perf(pt): reuse SO2 residual workspace during inference

[macqueen09]

Review follow-up: focused regression scope (2026-10-19)

Addressing OutisLi's requested changes.
The production implementation is unchanged from 7cfd926b; this follow-up changes only the test file. No new inference level, autotuning, performance study, production-size benchmark or full-model GPU test suite was added.

Coverage contract	Current cases	Mapping
Numerical/reference, shapes, caller-input preservation	7	The suggested representative list retains all five shape profiles, both alpha modes, empty/partial edges, zero/one/multiple gated layers and narrow/BMM paths
Workspace identity	4	Unchanged function: all with_weights/keep pairs; final out-of-place; gz/output always distinct
Dtype/alpha versus separate buffers	6	Unchanged function and tolerances: FP32/FP16/BF16 × alpha on/off
Actual autograd dispatch	6	Trainable/partially frozen/frozen weights × create_graph=False/True; mocked gradient operators test selection only
Saved-output cotangent dispatch	2	Frozen inputs with a cotangent on z or the final activation must use the protected training path
Unique frozen direct-stack derivative numerics	1	One small FP32 force-loss case checks output, first derivative and second-order objective gradients; no optimizer loop or width cross-product
Total	26	17 core + 8 dispatch guards + 1 numerical higher-order check

Existing TestSeZMModelCompile.test_forward_backward_double_backward_matches_compile covers ordinary model-level parameter gradients, force loss and short training; TestSeZMTrainPathParity.test_training_step_matches_the_dense_path covers descriptor/training binding parity. Those do not automatically establish the fully frozen direct-stack force-loss path, so one representative numerical check remains. The six new dispatch cases execute real torch.autograd.grad(..., create_graph=...); their mocked gradient operators are not presented as numerical training-equivalence evidence.

Fresh validation: 26 CUDA tests passed; with no visible GPU 26 skipped. The retained ownership guard still fails at its residual/output alias assertion when the exact baseline traversal is substituted in an isolated diagnostic process. Both retained guard/dtype functions are AST-identical to the previous head. Full-repository Ruff check/format and whitespace checks passed; production source is byte-for-byte unchanged.

Local CUDA coverage executes the three new statements (lines 3199–3201 at this revision), including both outcomes of the line-3200 guard. This is not a claim that upstream Codecov/required checks are green. No CI configuration or coverage exclusions were changed. Please verify GPU execution and the uploaded coverage in the new upstream runs.

The 106-case counts, full-model measurements and multi-step training evidence below are historical pre-review results, not newly rerun tests or performance claims for this follow-up.

---

Summary

Reuse the internally owned residual-gradient workspace in the SO2 mixing-stack
backward traversal when neither weight gradients nor saved per-layer state are
needed. This removes one edge-sized allocation without changing kernel
arithmetic, model parameters, precision settings or public APIs.

This is an isolated workspace change. It does not include compiler scheduling
passes, the inference-compiler flag from #6056, additional fusion, checkpoint
changes, or local benchmark infrastructure.

Why the alias is safe

The gated-layer backward GEMM computes residual + gz @ W.T. Its matrix input
gz remains separate from the output. Each output tile reads only its matching
residual elements before writing those same locations; distinct programs do not
read one another's residual tiles. The final-layer GEMM remains out-of-place.

Only buffers allocated inside this traversal are reused. Caller inputs and
forward-saved tensors are not mutated, so the operator's mutates_args=()
contract is unchanged. The weight-gradient and saved-state traversals keep their
original allocation paths. The guard uses actual gradient/state requirements,
not simply model.eval(), which matters for frozen-weight and force-loss cases.

The generated full-model graph was checked as well: four gated GEMM calls alias
residual/output, while both final-layer calls and every matrix-input/output pair
remain separate.

Scope

• deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py: local allocation change
  and the alias-safety explanation.
• source/tests/pt_expt/kernels/test_so2_workspace.py: generated-input CUDA
  regressions, with no checkpoint or machine-specific dependencies.

The current focused tests cover the contracts mapped in the review follow-up above.
They retain the allocation-identity baseline-fails/candidate-passes assertion,
unchanged dtype/alpha checks and tolerances, lightweight autograd-dispatch checks,
and one numerical frozen-weight force-loss case rather than repeated updates.

Full-model measurements

Base: dbca0b1cfabc9be166feffcbbc93e9c833603aca.
RTX 5090, driver 595.91.07; Python 3.11.12, PyTorch 2.11.0+cu128, Triton 3.6.0,
Warp 1.16.0. Same DPA4-Neo checkpoint and normal ASE/DP energy/forces/virial/stress
evaluation, including neighbor construction. DP_COMPILE_INFER=1,
DP_TRITON_INFER=2, TF32/AMP off, one host thread,
PYTORCH_ALLOC_CONF=expandable_segments:True.

Diamond C, lattice 3.567 Angstrom, seed 0, position noise sigma 0.02 Angstrom.
Each process: 8 warmup calls including the first, 30 synchronized wall-clock
samples, and a separate 30-call memory replay. Three independent processes for
4096 atoms; two for 1000. Timing values below are medians of process quantiles;
memory values are the maximum measured peaks. GB means 1e9 bytes.

Atoms	Metric	Base	Candidate	Change
4096	Steady P50, ms	155.339543	155.281764	−0.0372%
4096	Steady P90, ms	155.657487	155.602488	approximately unchanged
4096	Full measurement-window allocated peak, GB	24.061593	22.404837	−6.885%
4096	Steady allocated peak, GB	23.970400	22.313645	−6.912%
4096	Reserved peak, GB	25.612517	23.955767	−1.656755 GB
1000	Steady P50, ms	40.442513	40.511280	+0.1700%
1000	Full measurement-window allocated peak, GB	5.968265	5.563780	−6.777%
1000	Steady allocated peak, GB	5.869245	5.464760	−6.892%

This is a memory optimization, not a claim of significant speedup. The small
observed 1000-atom latency increase is disclosed rather than classified as a
speed improvement. The 4096-atom allocation reduction is exactly 1,656,755,200
bytes, matching one [2, 647170, 320] FP32 workspace. First-call compilation is
excluded from steady timing but included in the full memory window. GPU clocks
were not locked, and these short research runs are not a long-duration
performance acceptance test.

Numerical validation and limitations

Full-model comparison uses all energy, force, virial and stress elements with
finite-value checks and abs(actual-reference) <= 1e-5 + 1e-4*abs(reference).

• 1000 atoms: candidate versus same-round compiled base passes in both
  independent processes; all 10 three-frame self-repeatability rounds pass.
• 4096 atoms: all 3 candidate-versus-base comparisons pass, but strict
  self-repeatability is not universally satisfied: candidate 6/15 and base
  9/15 rounds pass. The failed checks are the third frame's virial. Baseline
  cross-process comparisons also fail that component in both repeat runs.
  No tolerance was relaxed and no output was excluded. These observations do
  not establish either a fix or a regression in global numerical determinism.

An additional independent dense eager GPU reference, with compilation and
Triton/CUDA fast paths disabled, also passes for both 1000-atom candidate runs
and both base runs at the same elementwise thresholds. This was one three-frame
reference evaluation, not an exhaustive determinism test.

Historical local test evidence (before the focused-test revision)

• New workspace regressions: 106 passed on CUDA.
• Full-repository Ruff check and format check: passed.
• Project-configured pre-commit hooks for both changed files: all applicable
  hooks passed.
• Patch whitespace check and clean application to the base: passed.
• Original model/descriptor/training-path regressions: 87 passed, 5 skipped,
  127 subtests passed. All nine TestSeZMModelCompile tests passed. The five
  skips require DPA4 native CUDA operators absent from the installed library.
• Red/green: the unmodified base fails the allocation-identity assertion as
  expected; the candidate passes.
• NVIDIA Compute Sanitizer memcheck: four alias/guard cases passed, 0 errors.
• CPU native tabulation forward/backward/second-order smoke: 6 passed.
• GPU native tabulation smoke: 2 failed on both base and candidate, because
  the installed customized-op library was built without GPU support. This is
  an environment limitation, not a passing test or a newly introduced failure.
• Without a visible GPU the new file collects normally and skips 106 tests.

The local source-selection harness used the new Python checkout with existing
native libraries from the measured environment. This does not claim a fresh
C++/wheel build or TensorFlow/JAX/Paddle validation. Other GPU architectures and
Torch 2.12/2.13 are not locally validated here. A CUDA-native rebuild is needed
to cover the missing native paths; this PR does not claim that every local
or required CI check is green.

Historical additional cross-checkout training equivalence

Separate BASE and candidate processes, identical generated inputs/initialization,
TF32 disabled, same GPU/environment:

Path	Cases / updates per case	Compared
Direct mixing stack, widths 32 and 96, all / partially frozen / frozen weights	6 cases, 5 AdamW steps each	output, force-loss second derivatives, gradients, updated parameters, step counters and moments
Actual SeZM model, dense eager	3 SGD + momentum steps	energy/force/virial loss, gradients, updated parameters and optimizer state
Actual SeZM model, eager Triton training level 1	3 SGD + momentum steps	same
Actual SeZM model, default dense compiled training	3 SGD + momentum steps	same

2516/2516 finite arrays pass the same atol=1e-5, rtol=1e-4
elementwise comparison. 1384 arrays are bitwise identical; global maximum
absolute difference is 3.725290298461914e-9; maximum error/tolerance ratio is
0.0003452008823. This is bounded multi-step equivalence evidence, not a
claim that arbitrary long training is bitwise identical.

The normal model.train() route is excluded from this inference stack by
the existing not self.training guard in so2.py. Direct-stack force-loss
tests additionally cover trainable, partially frozen and frozen operands:
observed backward calls requiring weight gradients retain separate workspace,
and calls saving second-order state retain the saved-state path. No public
training flags or APIs change.

An additional tiny-model configuration combining Triton training with compilation
triggered a Torch 2.11 Inductor replace_by_example pattern-matcher assertion
on the unmodified BASE. The final comparison therefore explicitly separates
eager Triton training from default compiled training, rather than claiming that
combined configuration passed. Candidate equivalence for that failing BASE
configuration is not established. The installed native library also lacks CUDA
operators; local CUDA toolkit 13.2 does not match this Torch cu128 build, and a
fresh matching-toolkit CUDA build remains for CI/a suitable environment.

Reproduction / reviewer requests

Run the new test file using the normal project installation. It generates all
operands and does not download a model. Also run the original
TestSeZMModelCompile, TestSeZMTritonValuePath, descriptor and training-path
regressions. The commands and standalone scripts below provide reproduction without
committing machine-specific benchmark infrastructure.

The exact checkpoint is DPA4-Neo-OMat24-v20260627.pt, SHA256
8d8b2b02b9575cb9e0af420cc1e8b1d9ffcbc04752c2203d29f48ccd0cd5f07d;
it is not bundled. Exact full-model numbers require an authorized copy.

Please trigger the upstream CUDA workflow and verify that relevant compile
tests actually run rather than skip under the workflow's Torch version. Local
results do not replace required status checks, maintainer review, or the merge
queue.

Reproducible commands

Install each revision following the project's CUDA/PyTorch build instructions
in separate environments, then verify deepmd.__file__ and
torch.__version__. Do not accidentally import another editable checkout.
Tests below need no checkpoint:

CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=1 python -m pytest -q \
  source/tests/pt_expt/kernels/test_so2_workspace.py
CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=1 python -m pytest -q \
  source/tests/pt/model/test_sezm_model.py::TestSeZMModelCompile \
  source/tests/pt/model/test_descriptor_sezm_triton.py::TestSeZMTritonValuePath \
  source/tests/pt/model/test_descriptor_sezm_train_paths.py \
  source/tests/pt/model/test_descriptor_sezm.py
ruff check .
ruff format --check .
prek run --files deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py \
  source/tests/pt_expt/kernels/test_so2_workspace.py
CUDA_VISIBLE_DEVICES=0 compute-sanitizer --tool memcheck --error-exitcode 86 \
  python -m pytest -q \
  source/tests/pt_expt/kernels/test_so2_workspace.py::test_reuse_only_without_weights_or_saved_state

Memcheck does not prove absence of all global-memory races; see the tile-level
read/write argument above.

For full-model reproduction, save the following script outside the repository.
Run serially on an otherwise idle GPU, with a fresh output directory and cache
for every invocation. With BASE installed/imported first:

CUDA_VISIBLE_DEVICES=0 MODEL=/path/to/authorized/model.pt OUTPUT=/scratch/base-r0 \
  TORCHINDUCTOR_CACHE_DIR=/scratch/cache-base-r0 NREP=8 \
  python /scratch/public_benchmark.py
# Switch to the candidate environment/checkout:
CUDA_VISIBLE_DEVICES=0 MODEL=/path/to/authorized/model.pt OUTPUT=/scratch/candidate-r0 \
  TORCHINDUCTOR_CACHE_DIR=/scratch/cache-candidate-r0 NREP=8 \
  REFERENCE=/scratch/base-r0/outputs.npz python /scratch/public_benchmark.py

Repeat 3 processes for NREP=8 (4096 atoms), 2 for NREP=5 (1000 atoms),
using different paths; interleave/reverse order when repeating. Inspect every
JSON comparison and finite check. This portable script reproduces normal API
timing/memory and cross-revision checks; the reported main experiment additionally
ran five three-frame self-repeatability rounds per process and audited generated
wrappers. Those extra diagnostics are not implemented by this short script.
Do not use torch.inference_mode(): conservative force evaluation needs
differentiation. No private compiler hooks are needed for this change.

Standalone public_benchmark.py (normal DP / ASE APIs)

"""Standalone normal-API reproduction; attach as text, not a production file.

Use a fresh process with the chosen checkout installed. MODEL is supplied by
the caller. No private compiler hooks or absolute machine paths are used.
"""
import hashlib
import json
import os
from pathlib import Path
import time

os.environ.update(
    DP_COMPILE_INFER="1", DP_COMPILE="0", DP_TRITON_INFER="2",
    DP_TF32_INFER="0", DP_AMP_INFER="0",
    DP_INTRA_OP_PARALLELISM_THREADS="1", DP_INTER_OP_PARALLELISM_THREADS="1",
    OMP_NUM_THREADS="1", PYTORCH_ALLOC_CONF="expandable_segments:True",
)
import numpy as np
import torch
import deepmd
from ase.build import bulk
from deepmd.calculator import DP

torch.set_num_threads(1)
torch.set_float32_matmul_precision("highest")
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
model = Path(os.environ["MODEL"])
out = Path(os.environ["OUTPUT"])
out.mkdir(parents=True, exist_ok=False)
nrep = int(os.environ.get("NREP", "8"))
warmup, samples = int(os.environ.get("WARMUP", "8")), int(os.environ.get("SAMPLES", "30"))
assert warmup >= 1 and samples >= 1
base = bulk("C", "diamond", a=3.567, cubic=True).repeat((nrep,) * 3)
rng = np.random.default_rng(0)
frames = [base.positions + 0.02*rng.standard_normal(base.positions.shape)
          for _ in range(max(3, warmup + samples))]
calc = DP(model=str(model))

def evaluate(pos):
    calc.reset()
    atoms = base.copy()
    atoms.positions = pos
    atoms.calc = calc
    force = atoms.get_forces().copy()
    return dict(energy=np.asarray(atoms.get_potential_energy()), forces=force,
                virial=np.asarray(calc.results["virial"]).copy(), stress=atoms.get_stress().copy())

torch.cuda.synchronize()
start = time.perf_counter()
evaluate(frames[0])
torch.cuda.synchronize()
first_ms = 1000*(time.perf_counter()-start)
for pos in frames[1:warmup]:
    evaluate(pos)
times = []
for pos in frames[warmup:warmup+samples]:
    torch.cuda.synchronize()
    start = time.perf_counter()
    evaluate(pos)
    torch.cuda.synchronize()
    times.append(1000*(time.perf_counter()-start))
full_allocated, full_reserved = torch.cuda.max_memory_allocated(), torch.cuda.max_memory_reserved()
torch.cuda.reset_peak_memory_stats()
for pos in frames[warmup:warmup+samples]:
    evaluate(pos)
torch.cuda.synchronize()
steady_allocated, steady_reserved = torch.cuda.max_memory_allocated(), torch.cuda.max_memory_reserved()
outputs = {f"{i}_{key}": value for i, pos in enumerate(frames[:3]) for key, value in evaluate(pos).items()}
np.savez_compressed(out/"outputs.npz", **outputs)
comparison = {}
if "REFERENCE" in os.environ:
    with np.load(os.environ["REFERENCE"]) as expected:
        for key, actual in outputs.items():
            ratio = np.abs(actual-expected[key])/(1e-5+1e-4*np.abs(expected[key]))
            comparison[key] = dict(passed=bool(np.isfinite(ratio).all() and (ratio <= 1).all()),
                                   max_ratio=float(ratio.max()))
result = dict(source=deepmd.__file__, torch=torch.__version__, atoms=len(base),
              model_sha256=hashlib.sha256(model.read_bytes()).hexdigest(),
              first_ms=first_ms, samples_ms=times, p50_ms=float(np.median(times)),
              p90_ms=float(np.percentile(times,90)), full_allocated=full_allocated,
              full_reserved=full_reserved, steady_allocated=steady_allocated,
              steady_reserved=steady_reserved, comparison=comparison,
              finite=all(bool(np.isfinite(x).all()) for x in outputs.values()))
(out/"result.json").write_text(json.dumps(result,indent=2,allow_nan=False))
print(json.dumps(result,indent=2,allow_nan=False))

Supplementary multi-step training comparison (outside the production branch)

Save this as training_equivalence.py outside the repository. Run once from
each installed checkout's root in a fresh process/cache:
CUDA_VISIBLE_DEVICES=0 PYTHONPATH="$PWD" TRAIN_OUTPUT=/scratch/train-base TORCHINDUCTOR_CACHE_DIR=/scratch/train-cache-base python /scratch/training_equivalence.py.
Use distinct candidate output/cache paths. The code below only makes GPU
selection configurable relative to the measured local script; computations are
the same. It uses existing repository test helpers and generated data; no
pretrained model is required.

"""Save actual optimizer trajectories for cross-checkout comparison; local only."""
import os
from pathlib import Path
import json

for key in list(os.environ):
    if key.startswith('DP_'):
        del os.environ[key]
os.environ.update(CUDA_VISIBLE_DEVICES=os.environ.get('CUDA_VISIBLE_DEVICES', '0'), OMP_NUM_THREADS='1',
                  DP_INTRA_OP_PARALLELISM_THREADS='1', DP_INTER_OP_PARALLELISM_THREADS='1',
                  DP_TRITON_TRAIN='1', DP_TRITON_INFER='2', DP_TF32_INFER='0',DP_AMP_INFER='0')
import numpy as np
import torch
import deepmd
from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
from source.tests.pt.model.test_sezm_model import TestSeZMModelCompile, get_sezm_model

# The repository test package deliberately installs an invalid default device;
# restore normal application semantics for Adam's CPU scalar step counters.
torch.set_default_device(None)
torch.set_num_threads(1)
torch.set_float32_matmul_precision('highest')
torch.backends.cuda.matmul.allow_tf32=False
torch.backends.cudnn.allow_tf32=False
out=Path(os.environ['TRAIN_OUTPUT'])
out.mkdir(parents=True,exist_ok=False)
arrays={}
counts={}
originals={}
for name in ('_stack_backward_traversal','_stack_train_traversal'):
    original=getattr(vp,name)
    originals[name]=original
    def observed(*args,_name=name,_fn=original,**kwargs):
        key=f'{_name}:weights={kwargs.get("with_weights")}:keep={kwargs.get("keep")}'
        counts[key]=counts.get(key,0)+1
        return _fn(*args,**kwargs)
    setattr(vp,name,observed)

def save(key,value):
    arrays[key]=value.detach().cpu().numpy().copy()
    assert np.isfinite(arrays[key]).all(),key

def state(prefix,optimizer,leaves):
    for i,p in enumerate(leaves):
        save(f'{prefix}/param/{i}',p)
        for name,value in optimizer.state[p].items():
            if isinstance(value,torch.Tensor):
                save(f'{prefix}/optimizer/{i}/{name}',value)

# Direct stack coverage includes narrow fused training and wide replay paths.
for lmax,width in ((2,32),(6,96)):
    for mode,requires in (('all',(True,True,True)),('partial',(True,False,False)),('frozen',(False,False,False))):
        torch.manual_seed(19)
        edges,focus,gates=7,2,2
        row,m0,m1=(3*lmax+1)*width,(lmax+1)*width,2*lmax*width
        operands=[torch.randn(focus,edges,row,device='cuda')*.1,
                  torch.rand(edges,focus,device='cuda')+.2,
                  torch.randn(gates+1,focus,m0,m0,device='cuda')*.02,
                  torch.randn(gates+1,focus,m1,m1,device='cuda')*.02,
                  torch.randn(gates,focus,width,lmax*width,device='cuda')*.02]
        for value,enabled in zip(operands,(True,True,*requires),strict=True):
            value.requires_grad_(enabled)
        leaves=[p for p in operands if p.requires_grad]
        optimizer=torch.optim.AdamW(leaves,lr=1e-4,weight_decay=.01)
        prefix=f'stack/{lmax}_{width}/{mode}'
        for step in range(5):
            optimizer.zero_grad(set_to_none=True)
            output=vp._mixing_stack_op(*operands,lmax,width,True)[0]
            energy=output.square().sum()/edges
            force=torch.autograd.grad(energy,operands[0],create_graph=True)[0]
            loss=energy+force.square().sum()
            save(f'{prefix}/{step}/output',output)
            save(f'{prefix}/{step}/force',force)
            save(f'{prefix}/{step}/loss',loss)
            loss.backward()
            for i,p in enumerate(leaves):
                assert p.grad is not None
                save(f'{prefix}/{step}/gradient/{i}',p.grad)
            optimizer.step()
            state(f'{prefix}/{step}',optimizer,leaves)
        print('STACK DONE',prefix,flush=True)

# Real SeZM energy/force/virial losses, not merely an isolated stack.
for name,original in originals.items():
    setattr(vp,name,original)
# Match the repository's model-test context. Keep the default compiled route
# separate from eager Triton training; neither is an inference benchmark.
torch.set_default_device('cuda:9999999')
helper=TestSeZMModelCompile()
helper.setUp()
coord,atype,box,energy_target,force_target,virial_target=helper._make_tiny_frame()
for mode,compiled,triton_level in (('dense_eager',False,0),('triton_eager',False,1),('dense_compiled',True,0)):
    os.environ['DP_TRITON_TRAIN']=str(triton_level)
    os.environ['DP_TRITON_INFER']='0'
    params=helper._build_model_params(use_compile=compiled)
    params['descriptor']['l_schedule']=[1,1]
    params['descriptor']['so3_readout']='mlp'
    if triton_level:
        params['descriptor'].update(channels=32,n_focus=2,l_schedule=[2,2],so2_layers=3,
                                    n_radial=8,ffn_neurons=64)
    model=get_sezm_model(params)
    helper._randomize_params(model)
    model.train()
    optimizer=torch.optim.SGD(model.parameters(),lr=1e-7,momentum=.9)
    leaves=list(model.parameters())
    prefix=f'model/{mode}'
    for step in range(3):
        optimizer.zero_grad(set_to_none=True)
        outputs=model(coord,atype,box=box)
        loss=0
        for key,target in (('energy',energy_target),('force',force_target),('virial',virial_target)):
            save(f'{prefix}/{step}/{key}',outputs[key])
            loss=loss+(outputs[key]-target.to(outputs[key].dtype)).square().mean()
        save(f'{prefix}/{step}/loss',loss)
        loss.backward()
        for i,p in enumerate(leaves):
            if p.grad is not None:
                save(f'{prefix}/{step}/gradient/{i}',p.grad)
        torch.nn.utils.clip_grad_norm_(leaves,1.)
        optimizer.step()
        state(f'{prefix}/{step}',optimizer,leaves)
    print('MODEL DONE',prefix,flush=True)
np.savez_compressed(out/'trajectory.npz',**arrays)
(out/'manifest.json').write_text(json.dumps(dict(source=deepmd.__file__,torch=torch.__version__,
        arrays=len(arrays),traversals=counts,finite=True),indent=2))
print('SAVED',len(arrays),counts,flush=True)

Compare the two generated archives with shape/key/dtype and finite checks:

import numpy as np
with np.load("/scratch/train-base/trajectory.npz") as base, \
     np.load("/scratch/train-candidate/trajectory.npz") as candidate:
    assert set(base.files) == set(candidate.files)
    for name in base.files:
        a, b = base[name], candidate[name]
        assert a.shape == b.shape and a.dtype == b.dtype, name
        assert np.isfinite(a).all() and np.isfinite(b).all(), name
        assert (np.abs(b-a) <= 1e-5 + 1e-4*np.abs(a)).all(), name

Contribution / merge checklist

• Small source change with generated-data regression tests; no checkpoints,
  benchmark caches, raw customer data, or machine-specific files committed.
• Conventional commit / PR title; Ruff and applicable project hooks pass.
• Numerical, memory, latency, training results and limitations disclosed.
• Upstream required status checks and actual CUDA job coverage.
• Maintainer approval and resolution of review threads.
• Maintainer-managed merge queue after approval/checks.

Prepared against the contributing guide
and linked coding conventions
and CI documentation.
Current default-branch rules require seven check contexts (C++ tests, C++ build,
C-library build, Python tests, wheel builds, CUDA aggregate, pre-commit.ci),
at least one approval, resolved threads, and the squash merge queue.
The CUDA workflow needs maintainer activation (e.g. the Test CUDA label);
a passing aggregate with skipped GPU jobs is not CUDA validation. Please also
check Torch-version gates on compile tests. No claim is made that local tests
replace these server-side requirements.

Summary by CodeRabbit

• Performance
  • Reduced temporary memory allocation during SO(2) backward computations when weight gradients and saved per-layer state are not needed. Computations that require weight gradients or retained state continue to use separate workspace as needed.
• Tests
  • Added coverage for numerical correctness, input preservation, and consistency across workspace strategies and supported configurations, including autograd behavior for trainable and frozen weights and force-loss derivatives.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: 5005469b-693f-410b-87a6-bd161f61274d

📥 Commits

Reviewing files that changed from the base of the PR and between 7cfd926 and 30c9fc9.

📒 Files selected for processing: 1

• source/tests/pt_expt/kernels/test_so2_workspace.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.

📝 Walkthrough

The SO(2) backward traversal now reuses its current gradient buffer as scratch when weights and saved state are not requested. CUDA-gated tests cover buffer selection, backward correctness, autograd dispatch, and force-loss gradients.

Changes: SO(2) backward workspace reuse

Layer / File(s)	Summary
Scratch buffer selection and alias checks deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py, source/tests/pt_expt/kernels/test_so2_workspace.py	The traversal reuses g_cur as g_spare when weights and saved state are not requested. Tests check launch-level aliasing, optional returned states, and unchanged inputs.
Traversal correctness and result parity source/tests/pt_expt/kernels/test_so2_workspace.py	Tests compare backward results with a reference and compare reused and separate workspace outputs across float32, float16, and bfloat16, with alpha enabled and disabled.
Autograd dispatch and force-loss gradients source/tests/pt_expt/kernels/test_so2_workspace.py	Tests check backward dispatch for trainable and frozen weights, pre-activation and final-activation cotangents, and force-loss gradients against a reference.

Priority: ➖ Normal

Estimated code review effort: 2 (Simple) | ~15 minutes

Change: Refactor

Suggested reviewers: outisli

Merge Risk: ⚪ Minimal · up to 30c9f

No identified issue currently blocks merging. Complete the pending upstream CUDA checks as part of normal validation.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 69.23% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 13 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly summarizes the main change: reusing the SO2 residual workspace to reduce allocation overhead during inference.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[njzjz-bot]

Agent: dot

Reviewed the exact two-file delta against dbca0b1; no actionable source defect found. Reuse is restricted to internally allocated g_cur without weight gradients or retained layer state. Gated launches keep gz separate, the final launch stays out-of-place, and the residual load/output store use the same uniquely owned (focus, edge, column) tile. Caller inputs and forward-saved tensors are not mutated.

Independent CPU checks using the actual traversal and four kernel ASTs: 40 baseline/candidate cases pass across weights/keep/alpha, zero-gated and narrow/wide-BMM paths; four reversed/shuffled multi-tile fp32/fp16 comparisons pass; caller buffers remain unchanged. An intentionally unsafe gz/output alias fails the negative control. Address enumeration covers 360 supported shape/edge partitions and 71,176,320 addresses without overlap or gaps. The actual autograd-dispatch AST passes all 64 combinations: trainable/partially frozen operands, create_graph and saved-output cotangents route through unchanged training handling; second-order replay uses keep=True. Both changed files compile syntactically and the new test matrix contains 106 CUDA cases.

These are CPU semantic/address probes, not CUDA execution or Inductor functionalization validation. Torch/GPU are unavailable here. Author-reported CUDA/memcheck, optimizer equivalence and measured memory reduction remain unverified locally. The disclosed 4096-atom self-repeatability failures on both base and candidate do not establish determinism or a causal regression; no timing speedup is inferred from the near-neutral latency measurements.

At review time all seven returned build/test workflows remain action_required. No workflow activation, code modification or merge action was taken; relevant actual CUDA/compile coverage must be distinguished from skipped aggregate checks.

[macqueen09]

Local inference and training validation, benchmark results, and reproduction commands are documented in the PR description.
The upstream workflows currently require approval.

Could a maintainer approve the pending runs and trigger the CUDA tests using the Test CUDA label? Please also verify that the relevant CUDA and compile tests execute rather than skip. Thanks!

[OutisLi]

Requesting changes to the test scope and organization, not to the production optimization.

The residual-workspace reuse is a small, coherent change. The matrix operand gz remains separate, each output tile reads only its own residual elements before overwriting them, the final-layer output stays out-of-place, and weight-gradient/saved-state traversals retain their original buffers. I did not find an actionable production correctness defect within this review. Please keep that implementation focused; it needs neither a new acceleration level nor a new autotuning mechanism.

Proposed test simplification

1. Replace the 90-case Cartesian product with a representative case list.

The main parameterization expands to 5 shapes × 3 edge counts × 3 layer counts × 2 alpha settings. Thirty cases have zero edges and return before the changed traversal; another twenty have nonzero edges but no gated layer, so the aliased gated update is never exercised. These boundaries deserve coverage, but not the full cross product. Even empty-edge cases still construct and snapshot weights.

For example, one parameterization over the following seven cases retains all five original shape profiles, both alpha modes, zero/partial-edge workloads, zero/single/multiple gated layers, and narrow/wide gate-projection paths:

# (lmax, focus, width, edges, gates, apply_alpha)
CASES = [
    (2, 2, 32, 0, 1, True),       # Empty-edge return.
    (2, 2, 32, 37, 0, True),      # Final layer only.
    (2, 2, 32, 37, 1, False),     # Single gated layer without alpha.
    (3, 2, 32, 129, 3, True),     # Multi-layer recurrence and partial tiles.
    (3, 1, 64, 37, 1, True),      # Single focus, wider narrow-channel path.
    (6, 2, 96, 37, 3, False),     # BMM gate path, non-power-of-two width.
    (6, 3, 128, 129, 1, True),    # BMM gate path and multiple focuses.
]

Keep the existing numerical, output-shape and caller-input-preservation assertions. This list is a concrete starting point, not a hard case-count target: add or adjust a case only when it covers a distinct ownership/branch/layout condition missing from the others.

2. Keep the four workspace-identity guard cases.

test_reuse_only_without_weights_or_saved_state directly tests the intended change and has a useful baseline-fails/candidate-passes assertion. Preserve all four with_weights/keep combinations, the out-of-place final layer, and the assertion that gz never aliases the output.

3. Keep the six dtype/alpha alias-versus-separate comparisons.

The three supported dtypes × two alpha settings in test_reuse_matches_separate_workspace cover a distinct contract. Preserve their numerical and input-preservation checks rather than weakening tolerances to reduce work.

4. Separate training dispatch protection from broad training equivalence.

The six frozen/trainable-weight cases each repeat first/second derivatives and two parameter updates. Much of this checks paths deliberately excluded from the changed reuse branch.

Use lightweight dispatch checks for trainable, partially frozen and fully frozen weights, especially create_graph=True, to establish that they select the protected training/saved-state path. Reuse existing model-level gradient/training regressions where they actually cover the same contract. Those existing tests do not automatically prove every frozen-weight direct-stack case, so identify any genuinely missing coverage before removing it. If a distinct numerical higher-order path still needs protection, retain the smallest representative single-step check for that path, rather than automatically crossing every weight-freezing mode with both widths and repeating optimizer updates.

This yields 17 focused core cases (7 numerical + 4 ownership guards + 6 dtype/alpha checks), plus any justified unique derivative coverage. Please explain that coverage mapping briefly in the updated PR. Do not add a new timing study, production-size benchmark, or full-model GPU test suite just to validate this test reorganization.

Verification and limits

The prior review's 11 small CPU-interpreter comparisons of the exact base/head traversal passed bitwise, preserved caller inputs, and confirmed the expected buffer identities. Ruff and whitespace checks passed. No GPU execution, test-runtime benchmark or independent reproduction of the author's GPU memory measurements was performed. The test-cost concern above is based on redundant paths and compilation/workload structure, not a claimed measured CI speedup.

Please retain the production change and make the regression suite proportionate to its contract before merging.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.82%. Comparing base (668b60c) to head (30c9fc9).
⚠️ Report is 2 commits behind head on master.

Files with missing lines	Patch %	Lines
...epmd/pt_expt/kernels/triton/sezm/so2_value_path.py	0.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #6058      +/-   ##
==========================================
+ Coverage   77.79%   77.82%   +0.02%     
==========================================
  Files        1171     1171              
  Lines      141562   141609      +47     
  Branches     5066     5072       +6     
==========================================
+ Hits       110135   110202      +67     
+ Misses      29550    29467      -83     
- Partials     1877     1940      +63     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[macqueen09]

@OutisLi Addressed in 30c9fc98. This follow-up changes only test_so2_workspace.py; the production implementation is unchanged.

• Replaced the 90-case Cartesian product with your seven representative configurations, keeping the original numerical tolerances, output-shape assertions and caller-input preservation.
• Kept the four workspace-identity cases and six dtype/alpha alias-versus-separate comparisons unchanged.
• Replaced the repeated width/freezing-mode optimizer loops with six lightweight real autograd dispatch checks (three freezing modes × create_graph off/on), plus two checks for saved-output cotangents. The gradient operators are mocked only to observe dispatch, not to claim numerical equivalence.
• Retained one small fully frozen direct-stack force-loss numerical case, comparing output, first derivative and second-order objective gradients against the dense reference. Existing model-level training/compile regressions do not automatically cover that distinct direct-stack case.
• Added the coverage mapping to the PR description: 17 core + 8 dispatch + 1 numerical higher-order = 26 cases. The old 106-case and multi-step results are now explicitly historical.

Validation on RTX 5090 / Torch 2.11.0+cu128: 26 passed; without a visible GPU 26 skipped. The retained identity guard fails at the expected alias assertion with the exact baseline traversal substituted in an isolated diagnostic process, and passes with the candidate. Full-repository Ruff check/format and whitespace checks passed.

Local CUDA coverage executes the three new production statements and both outcomes of the ownership guard. This does not assert that the upstream Codecov or required workflows are already green; no CI/coverage exclusions were changed. No timing study, production-size benchmark or full-model GPU suite was added for this revision.

Please re-review the focused scope when convenient.