perf(pt): reuse SO2 residual workspace during inference

[macqueen09]

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 new tests cover zero/partial edge tiles, multiple layer counts, narrow and
BMM gate-projection widths, FP32/FP16/BF16, input preservation, actual workspace
identity, and first/second derivatives plus two updates with trainable,
partially frozen and frozen weights. The allocation-identity test is designed
to fail on the unmodified base and pass with this change.

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.

Local test evidence

• 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.

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 to verify numerical results, input preservation, and consistency across workspace strategies and supported configurations.

[coderabbitai]

Important

Review skipped

Review was skipped as selected files did not have any reviewable changes.

⚙️ Run configuration

• Configuration used: Repository UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: 0f8b9764-71e0-495e-be6c-3623cd43bc38

📥 Commits

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

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

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: 8d0c7d14-4a10-476b-88aa-4d392e6c7c77

📥 Commits

Reviewing files that changed from the base of the PR and between dbca0b1 and fc19816.

📒 Files selected for processing (2)

• deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
• 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 compare results with reference implementations and cover workspace selection, input preservation, multiple dtypes, and force-loss updates.

Changes: SO(2) backward workspace reuse

Layer / File(s)	Summary
Scratch buffer selection and validation 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 cover the selection conditions, compare results with references across inputs and dtypes, check input tensors remain unchanged, and compare force-loss gradients and parameter updates with the dense reference.

Priority: ➖ Normal

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

Change: Refactor

Suggested reviewers: outisli

Merge Risk: ⚪ Minimal · up to fc198

No merge-blocking issue is identified. The workspace change appears ready for normal required checks.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 10 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	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly identifies the main change: reusing the SO2 residual workspace. The inference qualifier is somewhat narrow because the optimization also applies when weight gradients and saved state…

✨ Finishing Touches 💡 1

🧪 Generate unit tests (beta)

• Create a new PR

🛠️ Fix failing CI checks 💡

• Commit to this branch
• 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.55%. Comparing base (668b60c) to head (7cfd926).

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.55%   -0.25%     
==========================================
  Files        1171     1171              
  Lines      141562   141564       +2     
  Branches     5066     5065       -1     
==========================================
- Hits       110135   109784     -351     
- Misses      29550    29905     +355     
+ Partials     1877     1875       -2     

☔ 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.