From 5caac25a9af8a032f5ce84091446c50e9869bb28 Mon Sep 17 00:00:00 2001
From: Macqueen Ma-hnu <16522286+macqueen09@users.noreply.github.com>
Date: Thu, 29 Oct 2026 10:01:37 +0800
Subject: [PATCH 1/9] perf(pt): save SO2 competition scalars during rotation

---
 .../kernels/triton/sezm/so2_value_path.py     | 133 +++++++-
 .../kernels/triton/sezm/sweep_tile_configs.py |   2 +
 doc/model/dpa4.md                             |  12 +
 .../pt_expt/benchmark_so2_scalar_capture.py   | 293 ++++++++++++++++++
 .../pt_expt/compare_so2_scalar_capture.py     |  76 +++++
 .../kernels/test_so2_scalar_capture.py        | 189 +++++++++++
 6 files changed, 703 insertions(+), 2 deletions(-)
 create mode 100644 source/tests/pt_expt/benchmark_so2_scalar_capture.py
 create mode 100644 source/tests/pt_expt/compare_so2_scalar_capture.py
 create mode 100644 source/tests/pt_expt/kernels/test_so2_scalar_capture.py

diff --git a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
index a02f428825..fefcaead40 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
@@ -91,6 +91,7 @@
     annotations,
 )
 
+import os
 from typing import (
     TYPE_CHECKING,
     NamedTuple,
@@ -508,6 +509,7 @@ def _rotate_mix_fwd_kernel(
         kc_ptr,  # (E, KSZ * RANK) compact kernel, or (E, L+1, CW) when RANK == 0
         cb_ptr,  # (RANK, CW) channel basis (unread when RANK == 0)
         u_ptr,  # (F, E, ROW) focus-major output
+        scalar_ptr,  # optional independent (E, F, CF), pitch 2*CF
         n_edge,
         x_sn,
         x_sd,
@@ -516,6 +518,7 @@ def _rotate_mix_fwd_kernel(
         CW: tl.constexpr,  # true C_wide; BC = next_power_of_2(CW) lanes with mask
         BC: tl.constexpr,
         RANK: tl.constexpr,
+        SAVE_SCALAR: tl.constexpr = False,
     ):
         """One program per edge, channels vectorized.
 
@@ -576,6 +579,13 @@ def _rotate_mix_fwd_kernel(
                     tl.float32
                 )
                 tl.store(u_ptr + f_off + o * CF, xl[o] * rad, mask=cmask)
+                if SAVE_SCALAR and o == 0:
+                    scalar_off = (
+                        (chan // CF).to(tl.int64) * n_edge * (2 * CF)
+                        + edge * (2 * CF)
+                        + chan % CF
+                    )
+                    tl.store(scalar_ptr + scalar_off, xl[o] * rad, mask=cmask)
             for o in tl.static_range(L):
                 rad = tl.load(rad_base + (o + 1) * CW + chan, mask=cmask, other=0.0).to(
                     tl.float32
@@ -606,6 +616,13 @@ def _rotate_mix_fwd_kernel(
             if RANK == 1:
                 acc = acc * cb[0]
             tl.store(u_ptr + f_off + o * CF, acc, mask=cmask)
+            if SAVE_SCALAR and o == 0:
+                scalar_off = (
+                    (chan // CF).to(tl.int64) * n_edge * (2 * CF)
+                    + edge * (2 * CF)
+                    + chan % CF
+                )
+                tl.store(scalar_ptr + scalar_off, acc, mask=cmask)
         for o in tl.static_range(L):
             accn = tl.zeros((BC,), dtype=tl.float32)
             accq = tl.zeros((BC,), dtype=tl.float32)
@@ -2498,6 +2515,7 @@ def _rotate_mix_impl(
         kc,
         cb,
         u,
+        u,  # unread unless SAVE_SCALAR is enabled
         n_edge,
         x.stride(0),
         x.stride(1),
@@ -2506,12 +2524,73 @@ def _rotate_mix_impl(
         CW=c_wide,
         BC=triton.next_power_of_2(c_wide),
         RANK=int(rank),
+        SAVE_SCALAR=False,
         num_warps=warps,
         num_stages=stages,
     )
     return u
 
 
+def _rotate_mix_scalar_impl(
+    x: Tensor,
+    src: Tensor,
+    src_order: Tensor,
+    src_rowptr: Tensor,
+    wigner: Tensor,
+    kc: Tensor,
+    cb: Tensor,
+    lmax: int,
+    n_focus: int,
+    rank: int,
+) -> tuple[Tensor, Tensor]:
+    """Emit the l=0 competition input without retaining a view of ``u``.
+
+    The extra store uses the existing rotation accumulator; no second gather
+    or rotation is performed.  The padded focus-major scalar allocation keeps
+    the competition reduction layout used by the original compiled path.
+    Its padding is never read.  Both outputs own independent storage.
+    """
+    del src_order, src_rowptr
+    n_edge = src.shape[0]
+    c_wide = int(x.shape[2])
+    focus_dim = c_wide // int(n_focus)
+    row = (3 * int(lmax) + 1) * focus_dim
+    if not _use_triton(x):
+        u = _rotate_mix_reference(x, src, wigner, kc, cb, lmax, n_focus, rank)
+        return u, u[:, :, :focus_dim].permute(1, 0, 2).clone()
+    u = torch.empty(n_focus, n_edge, row, device=x.device, dtype=x.dtype)
+    pitch = 2 * focus_dim
+    scalar = torch.empty_strided(
+        (n_edge, n_focus, focus_dim),
+        (pitch, n_edge * pitch, 1),
+        device=x.device,
+        dtype=x.dtype,
+    )
+    if not _has_no_edges(n_edge):
+        warps, stages = rotate_mix_fwd_config(c_wide, int(lmax))
+        wrap_triton(_rotate_mix_fwd_kernel)[(n_edge,)](
+            x,
+            src,
+            wigner,
+            kc,
+            cb,
+            u,
+            scalar,
+            n_edge,
+            x.stride(0),
+            x.stride(1),
+            L=int(lmax),
+            CF=focus_dim,
+            CW=c_wide,
+            BC=triton.next_power_of_2(c_wide),
+            RANK=int(rank),
+            SAVE_SCALAR=True,
+            num_warps=warps,
+            num_stages=stages,
+        )
+    return u, scalar
+
+
 def _rotate_mix_bwd_impl(
     grad_u: Tensor,
     x: Tensor,
@@ -3608,6 +3687,9 @@ def _mixing_stack_train_bwd_impl(
 _rotate_mix_op = torch.library.triton_op(
     "sezm_triton::so2_rotate_mix", mutates_args=()
 )(_rotate_mix_impl)
+_rotate_mix_scalar_op = torch.library.triton_op(
+    "sezm_triton::so2_rotate_mix_scalar", mutates_args=()
+)(_rotate_mix_scalar_impl)
 _rotate_mix_bwd_op = torch.library.triton_op(
     "sezm_triton::so2_rotate_mix_bwd", mutates_args=()
 )(_rotate_mix_bwd_impl)
@@ -3892,6 +3974,25 @@ def _rotate_mix_backward(ctx, grad_u):
     return grad_x, None, None, None, grad_wigner, grad_kc, grad_cb, None, None, None
 
 
+def _rotate_mix_scalar_setup_context(ctx, inputs, output):
+    _rotate_mix_setup_context(ctx, inputs, output)
+    ctx.set_materialize_grads(False)
+
+
+def _rotate_mix_scalar_backward(ctx, grad_u, grad_scalar):
+    # Exactly the slice/permutation adjoint of the old competition branch.
+    # Keep pad/add visible to Inductor so it can fuse the gradient merge.
+    if grad_scalar is not None:
+        focus_dim = grad_scalar.shape[-1]
+        extra = torch.nn.functional.pad(
+            grad_scalar.permute(1, 0, 2), (0, 3 * ctx.lmax * focus_dim)
+        )
+        grad_u = extra if grad_u is None else grad_u + extra
+    if grad_u is None:
+        return (None,) * 10
+    return _rotate_mix_backward(ctx, grad_u)
+
+
 def _rotate_mix_bwd_setup_context(ctx, inputs, output):
     grad_u, x, src, wigner, kc, cb, lmax, n_focus, rank = inputs
     ctx.save_for_backward(grad_u, x, src, wigner, kc, cb)
@@ -4010,6 +4111,9 @@ def _segment_sum_backward(ctx, grad_out):
 _rotate_mix_op.register_autograd(
     _rotate_mix_backward, setup_context=_rotate_mix_setup_context
 )
+_rotate_mix_scalar_op.register_autograd(
+    _rotate_mix_scalar_backward, setup_context=_rotate_mix_scalar_setup_context
+)
 _rotate_mix_bwd_op.register_autograd(
     _rotate_mix_bwd_backward, setup_context=_rotate_mix_bwd_setup_context
 )
@@ -5176,6 +5280,22 @@ def make_triton_rotate_mix(conv: SO2Convolution) -> _TritonRotateMix | None:
 )
 
 
+def _save_competition_scalar(conv: SO2Convolution, x: Tensor) -> bool:
+    """Opt in only to the validated frozen two-focus FP32 inference layout."""
+    return (
+        os.environ.get("DP_SO2_SCALAR_CAPTURE_INFER", "0").strip().lower()
+        in ("1", "true", "yes", "on")
+        and not conv.training
+        and conv.triton_infer_level == 2
+        and conv.lmax == 3
+        and conv.n_focus == 2
+        and conv.so2_focus_dim == 32
+        and conv.focus_compete
+        and x.dtype == torch.float32
+        and x.is_cuda
+    )
+
+
 class _TritonSO2ValuePath:
     """Per-convolution entry running the SO(2) value path through the fused ops.
 
@@ -5331,7 +5451,8 @@ def __call__(
             src_rowptr = torch.cat([counts.new_zeros(1), torch.cumsum(counts, 0)])
         else:
             src_order, src_rowptr = csr
-        u0 = _rotate_mix_op(
+        save_scalar = _save_competition_scalar(conv, x)
+        rotation = (_rotate_mix_scalar_op if save_scalar else _rotate_mix_op)(
             x.contiguous(),
             src,
             src_order,
@@ -5343,13 +5464,21 @@ def __call__(
             conv.n_focus,
             rank,
         )
+        if save_scalar:
+            u0, scalar = rotation
+        else:
+            u0 = rotation
 
         # === Step 3. Cross-focus competition weight from the l = 0 scalars ===
         apply_alpha = bool(conv.focus_compete and conv.n_focus > 1)
         if apply_alpha:
             # The small (E, F, Cf) copy keeps the softmax backward from
             # retaining a view of the whole focus-major activation.
-            gate_src = u0[:, :, : conv.so2_focus_dim].permute(1, 0, 2).contiguous()
+            gate_src = (
+                scalar.contiguous()
+                if save_scalar
+                else u0[:, :, : conv.so2_focus_dim].permute(1, 0, 2).contiguous()
+            )
             alpha = conv._focus_alpha(gate_src).to(u0.dtype).contiguous()
         else:
             alpha = torch.ones(
diff --git a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py
index 44c0a91aac..6696b2e9ce 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py
@@ -894,6 +894,7 @@ def launch(warps: int, stages: int) -> None:
             kc,
             cb,
             u,
+            u,  # scalar output is disabled during the original rotate sweep
             n_edge,
             x.stride(0),
             x.stride(1),
@@ -902,6 +903,7 @@ def launch(warps: int, stages: int) -> None:
             CW=c_wide,
             BC=triton.next_power_of_2(c_wide),
             RANK=1,
+            SAVE_SCALAR=False,
             num_warps=warps,
             num_stages=stages,
         )
diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md
index ad70c54b6f..8f06e7581a 100644
--- a/doc/model/dpa4.md
+++ b/doc/model/dpa4.md
@@ -568,6 +568,18 @@ shipping built in. On other GPUs the kernels fall back to conservative
 configurations, and `dp --pt freeze` tunes the missing entries on the local GPU
 before exporting, a one-off sweep of a few minutes baked into the `.pt2`.
 
+`DP_SO2_SCALAR_CAPTURE_INFER=1` optionally saves the focus-competition scalars
+into independent storage during the existing rotate-and-mix launch. This lets
+the full rotation output die after its mixing consumer, rather than remain
+live for competition backward. It does not change model weights or precision.
+The option defaults off and currently selects only evaluation with level `2`,
+CUDA float32, `lmax=3`, two focuses, 32 channels per focus and focus competition.
+Other layouts and training retain the original path. Set it before compilation
+or export; existing `.pt2` artifacts must be regenerated. This is a
+capacity optimization: the extra scalar store need not reduce latency.
+The manual full-API benchmark is
+`source/tests/pt_expt/benchmark_so2_scalar_capture.py` (`--help` lists inputs).
+
 `DP_CUTILE_INFER` replaces the whole SeZM edge pipeline — Wigner monomials,
 rotate-and-mix, the gated SO(2) mixing stack, the attention aggregation and the
 force / virial assembly — with kernels written in the `cuda.tile` DSL. On an
diff --git a/source/tests/pt_expt/benchmark_so2_scalar_capture.py b/source/tests/pt_expt/benchmark_so2_scalar_capture.py
new file mode 100644
index 0000000000..8c7f3b59f9
--- /dev/null
+++ b/source/tests/pt_expt/benchmark_so2_scalar_capture.py
@@ -0,0 +1,293 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Manual full-API CSC5 benchmark; run one route in each fresh process.
+
+Example: python source/tests/pt_expt/benchmark_so2_scalar_capture.py \
+    --model model.pt --route CSC5 --cell 8 8 8 --out results/csc5
+"""
+# ruff: noqa: T201 -- standalone benchmark CLI
+
+import argparse
+import copy
+import hashlib
+import json
+import os
+from pathlib import Path
+import site
+import subprocess
+import sys
+import time
+import traceback
+
+p = argparse.ArgumentParser(description=__doc__)
+p.add_argument("--model", type=Path, required=True)
+p.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[3])
+p.add_argument("--route", choices=("TORCH", "BASE", "CSC5"), required=True)
+p.add_argument("--cell", type=int, nargs=3, default=(5, 5, 5))
+p.add_argument("--out", type=Path, required=True)
+p.add_argument("--cache", type=Path)
+p.add_argument("--gpu", type=int, default=0)
+p.add_argument("--seed", type=int, default=0)
+p.add_argument("--warmup", type=int, default=12)
+p.add_argument("--samples", type=int, default=30)
+p.add_argument("--legacy-pro", action="store_true")
+a = p.parse_args()
+assert a.warmup >= 3 and a.samples >= 10
+a.out.mkdir(parents=True, exist_ok=False)
+cache = a.cache or a.out / "cache"
+
+
+def write(name, data):
+    (a.out / name).write_text(json.dumps(data, indent=2, allow_nan=False))
+
+
+def sha(path):
+    return hashlib.sha256(Path(path).read_bytes()).hexdigest()
+
+
+busy = subprocess.check_output(
+    [
+        "nvidia-smi",
+        "-i",
+        str(a.gpu),
+        "--query-compute-apps=pid",
+        "--format=csv,noheader",
+    ],
+    text=True,
+).strip()
+if busy:
+    write("result.json", {"status": "busy", "pids": busy})
+    raise SystemExit(3)
+for key in list(os.environ):
+    if key.startswith("DP_"):
+        del os.environ[key]
+os.environ.pop("CUBLAS_WORKSPACE_CONFIG", None)
+os.environ.update(
+    CUDA_VISIBLE_DEVICES=str(a.gpu),
+    DP_COMPILE_INFER="1",
+    DP_COMPILE="0",
+    DP_TRITON_INFER="0" if a.route == "TORCH" else "2",
+    DP_CUDA_INFER="0",
+    DP_TF32_INFER="0",
+    DP_AMP_INFER="0",
+    DP_SO2_SCALAR_CAPTURE_INFER="1" if a.route == "CSC5" else "0",
+    OMP_NUM_THREADS="1",
+    DP_INTRA_OP_PARALLELISM_THREADS="1",
+    DP_INTER_OP_PARALLELISM_THREADS="1",
+    PYTORCH_ALLOC_CONF="expandable_segments:True",
+    TORCHINDUCTOR_COMPILE_THREADS="1",
+    TORCHINDUCTOR_CACHE_DIR=str(cache / "inductor"),
+    TRITON_CACHE_DIR=str(cache / "triton"),
+    TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD="1",
+)
+sys.path.insert(0, str(a.root.resolve()))
+sys.meta_path[:] = [
+    f
+    for f in sys.meta_path
+    if not type(f).__module__.startswith("_editable_skbc_deepmd_kit")
+]
+import deepmd
+
+assert Path(deepmd.__file__).resolve().parent.parent == a.root.resolve()
+# Permit a source checkout to use its matching installed native extension.
+for installed in site.getsitepackages():
+    lib = Path(installed) / "deepmd"
+    if (lib / "lib").is_dir() and str(lib) not in deepmd.__path__:
+        deepmd.__path__.append(str(lib))
+import numpy as np
+import torch
+import triton
+from ase.build import bulk
+from deepmd.calculator import DP
+from deepmd.utils.errors import OutOfMemoryError as DPOutOfMemoryError
+
+torch.set_num_threads(1)
+torch.set_float32_matmul_precision("highest")
+torch.backends.cuda.matmul.allow_tf32 = False
+torch.backends.cudnn.allow_tf32 = False
+torch.use_deterministic_algorithms(False)
+torch.cuda.reset_peak_memory_stats()
+write(
+    "manifest.json",
+    {
+        "args": {k: str(v) for k, v in vars(a).items()},
+        "model_sha256": sha(a.model),
+        "torch": torch.__version__,
+        "triton": triton.__version__,
+        "cuda": torch.version.cuda,
+        "revision": subprocess.check_output(
+            ["git", "rev-parse", "HEAD"], cwd=a.root, text=True
+        ).strip(),
+        "source_sha256": sha(
+            a.root / "deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py"
+        ),
+        "benchmark_sha256": sha(__file__),
+        "flags": {
+            k: v
+            for k, v in os.environ.items()
+            if k.startswith(("DP_", "TORCHINDUCTOR", "TRITON_CACHE", "PYTORCH_ALLOC"))
+        },
+        "gpu": subprocess.check_output(
+            [
+                "nvidia-smi",
+                "-i",
+                str(a.gpu),
+                "--query-gpu=name,driver_version,uuid,clocks.sm,clocks.mem,temperature.gpu",
+                "--format=csv",
+            ],
+            text=True,
+        ),
+    },
+)
+result = {
+    "status": "running",
+    "route": a.route,
+    "n_atoms": int(8 * np.prod(a.cell)),
+    "atol": 1e-5,
+    "rtol": 1e-4,
+}
+arrays = {}
+started = time.time()
+try:
+    atoms = bulk("C", "diamond", a=3.567, cubic=True).repeat(tuple(a.cell))
+    rng = np.random.default_rng(a.seed)
+    frames = np.stack(
+        [
+            atoms.positions + 0.02 * rng.standard_normal(atoms.positions.shape)
+            for _ in range(3)
+        ]
+    )
+    np.savez_compressed(
+        a.out / "inputs.npz",
+        positions=frames,
+        cell=np.asarray(atoms.cell),
+        numbers=atoms.numbers,
+    )
+    result["input_sha256"] = sha(a.out / "inputs.npz")
+    load = torch.load
+    if a.legacy_pro:
+        assert (
+            sha(a.model)
+            == "446a6e893101fc52190cfccabb9d2328863d651aca36981fd8d3401c71a61221"
+        )
+
+        def legacy_load(file, *args, **kwargs):
+            obj = load(file, *args, **kwargs)
+            if (
+                isinstance(file, (str, Path))
+                and Path(file).resolve() == a.model.resolve()
+            ):
+                state = obj["model"] if "model" in obj else obj
+                state["_extra_state"] = copy.deepcopy(state["_extra_state"])
+                fitting = state["_extra_state"]["model_params"]["fitting_net"]
+                assert fitting["type"] == "ener"
+                fitting["type"] = "dpa4_ener"
+            return obj
+
+        torch.load = legacy_load
+        result["legacy_pro"] = (
+            "in-memory fitting type alias ener -> dpa4_ener; weights unchanged"
+        )
+    try:
+        calc = DP(model=str(a.model), freeze=False)
+    finally:
+        torch.load = load
+    assert not calc.dp.get_model().training
+    result["parameters"] = sum(t.numel() for t in calc.dp.get_model().parameters())
+    write("model_definition.json", calc.dp.get_model_def_script())
+
+    def evaluate(pos):
+        calc.reset()
+        sample = atoms.copy()
+        sample.positions = pos
+        sample.calc = calc
+        forces = sample.get_forces().copy()
+        return {
+            "energy": np.asarray(sample.get_potential_energy()),
+            "forces": forces,
+            "virial": np.asarray(calc.results["virial"]).copy(),
+            "stress": sample.get_stress().copy(),
+        }
+
+    def peak():
+        return {
+            "allocated": torch.cuda.max_memory_allocated(),
+            "reserved": torch.cuda.max_memory_reserved(),
+        }
+
+    def stats(values):
+        return {
+            "median": float(np.median(values)),
+            "p90": float(np.percentile(values, 90)),
+            "min": min(values),
+            "max": max(values),
+            "samples": values,
+        }
+
+    result["accuracy_wall_ms"] = []
+    for repeat in range(3):
+        for frame, pos in enumerate(frames):
+            tick = time.perf_counter()
+            out = evaluate(pos)
+            torch.cuda.synchronize()
+            assert all(np.isfinite(v).all() for v in out.values())
+            result["accuracy_wall_ms"].append(1000 * (time.perf_counter() - tick))
+            arrays.update({f"{repeat}_{frame}_{k}": v for k, v in out.items()})
+            np.savez_compressed(a.out / "outputs.npz", **arrays)
+            if repeat == frame == 0:
+                result["first_peak"] = peak()
+            print(a.route, len(atoms), repeat, frame, flush=True)
+    for i in range(a.warmup):
+        evaluate(frames[i % 3])
+    wall, events = [], []
+    for i in range(a.samples):
+        torch.cuda.synchronize()
+        begin, end = (
+            torch.cuda.Event(enable_timing=True),
+            torch.cuda.Event(enable_timing=True),
+        )
+        tick = time.perf_counter()
+        begin.record()
+        out = evaluate(frames[i % 3])
+        end.record()
+        torch.cuda.synchronize()
+        wall.append(1000 * (time.perf_counter() - tick))
+        events.append(begin.elapsed_time(end))
+        assert all(np.isfinite(v).all() for v in out.values())
+    full = peak()
+    torch.cuda.reset_peak_memory_stats()
+    for pos in frames:
+        evaluate(pos)
+    torch.cuda.synchronize()
+    result["steady_peak"] = peak()
+    result["complete_peak"] = {k: max(full[k], result["steady_peak"][k]) for k in full}
+    from torch._inductor.codecache import PyCodeCache
+
+    wrappers = [
+        str(m.__file__)
+        for m in PyCodeCache.modules
+        if "def call(" in Path(m.__file__).read_text()
+    ]
+    assert wrappers
+    result.update(
+        status="completed",
+        finite=True,
+        wall_ms=stats(wall),
+        event_ms=stats(events),
+        wrappers=wrappers,
+    )
+except Exception as exc:
+    result.update(
+        status="oom"
+        if isinstance(exc, (torch.OutOfMemoryError, DPOutOfMemoryError))
+        else "failed",
+        exception=traceback.format_exc(),
+        complete_peak={
+            "allocated": torch.cuda.max_memory_allocated(),
+            "reserved": torch.cuda.max_memory_reserved(),
+        },
+    )
+    print(result["exception"], flush=True)
+finally:
+    result["elapsed_s"] = time.time() - started
+    write("result.json", result)
+raise SystemExit(0 if result["status"] == "completed" else 1)
diff --git a/source/tests/pt_expt/compare_so2_scalar_capture.py b/source/tests/pt_expt/compare_so2_scalar_capture.py
new file mode 100644
index 0000000000..545e408544
--- /dev/null
+++ b/source/tests/pt_expt/compare_so2_scalar_capture.py
@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Compare same-frame repeats from benchmark_so2_scalar_capture.py.
+
+Run with --torch DIR --base DIR --candidate DIR. An OOM reference is unavailable,
+not a pass. Output is JSON; absolute units follow the ASE energy/force/virial/stress
+arrays. Relative errors use max(abs(reference), 1e-8) as the denominator.
+"""
+# ruff: noqa: T201 -- standalone comparison CLI
+
+import argparse
+import json
+from itertools import (
+    product,
+)
+from pathlib import (
+    Path,
+)
+
+import numpy as np
+
+
+def compare(candidate, reference):
+    """Compare all 3x3 repeat pairs for each of the three corresponding frames."""
+    fields = {}
+    for field in ("energy", "forces", "virial", "stress"):
+        rows = []
+        for frame, i, j in product(range(3), range(3), range(3)):
+            x = candidate[f"{i}_{frame}_{field}"]
+            y = reference[f"{j}_{frame}_{field}"]
+            assert x.shape == y.shape
+            delta = np.abs(x - y)
+            rows.append(
+                {
+                    "max_abs": float(delta.max()),
+                    "mean_abs": float(delta.mean()),
+                    "max_relative": float((delta / np.maximum(abs(y), 1e-8)).max()),
+                    "max_ratio": float((delta / (1e-5 + 1e-4 * abs(y))).max()),
+                    "finite": bool(np.isfinite(x).all() and np.isfinite(y).all()),
+                }
+            )
+        fields[field] = {
+            k: max(row[k] for row in rows)
+            for k in ("max_abs", "mean_abs", "max_relative", "max_ratio")
+        }
+        fields[field]["passed"] = all(
+            row["finite"] and row["max_ratio"] <= 1 for row in rows
+        )
+    return fields
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    for route in ("torch", "base", "candidate"):
+        parser.add_argument(f"--{route}", type=Path, required=True)
+    args = parser.parse_args()
+    results, arrays, out = {}, {}, {}
+    for route, path in vars(args).items():
+        results[route] = json.loads((path / "result.json").read_text())
+        out[route] = {"status": results[route]["status"]}
+        if results[route]["status"] == "completed":
+            arrays[route] = dict(np.load(path / "outputs.npz"))
+            out[route]["repeat"] = compare(arrays[route], arrays[route])
+    for route in arrays:
+        for reference in ("torch", "base"):
+            if route != reference and reference in arrays:
+                assert (
+                    results[route]["input_sha256"] == results[reference]["input_sha256"]
+                )
+                out[route][f"vs_{reference}"] = compare(
+                    arrays[route], arrays[reference]
+                )
+    print(json.dumps(out, indent=2, allow_nan=False))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/source/tests/pt_expt/kernels/test_so2_scalar_capture.py b/source/tests/pt_expt/kernels/test_so2_scalar_capture.py
new file mode 100644
index 0000000000..bd1d466ddb
--- /dev/null
+++ b/source/tests/pt_expt/kernels/test_so2_scalar_capture.py
@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Independent scalar storage emitted by the existing rotation producer."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or not vp.SO2_VALUE_PATH_TRITON_AVAILABLE,
+    reason="CUDA and Triton are required",
+)
+
+
+@pytest.fixture(autouse=True)
+def _strict_fp32():
+    old = torch.backends.cuda.matmul.allow_tf32
+    torch.backends.cuda.matmul.allow_tf32 = False
+    yield
+    torch.backends.cuda.matmul.allow_tf32 = old
+
+
+def _inputs(edges: int, rank: int, zero: bool = False) -> tuple:
+    torch.manual_seed(42)
+    # Noncontiguous node stride; the existing kernel requires channel stride 1.
+    x = (torch.randn(14, 16, 64, device="cuda") * 0.1)[::2].detach()
+    if zero:
+        x.zero_()
+    x.requires_grad_()
+    src = torch.arange(edges, device="cuda") % 7
+    order = src.argsort(stable=True)
+    counts = torch.bincount(src, minlength=7)
+    ptr = torch.cat((counts.new_zeros(1), counts.cumsum(0)))
+    w = torch.randn(edges, 16, 16, device="cuda") * 0.1
+    mask = torch.zeros(16, 16, device="cuda")
+    for degree in range(4):
+        mask[degree**2 : (degree + 1) ** 2, degree**2 : (degree + 1) ** 2] = 1
+    w = (w * mask).requires_grad_()
+    shape = (edges, 4, 64) if rank == 0 else (edges, 25 * rank)
+    kc = (torch.randn(shape, device="cuda") * 0.1).requires_grad_()
+    cb = torch.randn(1 if rank == 0 else rank * 64, device="cuda") * 0.1
+    cb.requires_grad_(rank > 0)
+    return x, src, order, ptr, w, kc, cb, 3, 2, rank
+
+
+@pytest.mark.parametrize("edges,rank", [(0, 1), (1, 0), (17, 1), (257, 2)])
+@pytest.mark.parametrize("branch", ["both", "main", "scalar"])
+@pytest.mark.parametrize("zero", [False, True])
+def test_capture_and_vjp(edges: int, rank: int, branch: str, zero: bool) -> None:
+    args = _inputs(edges, rank, zero)
+    original = vp._rotate_mix_op(*args)
+    expected = (original, original[:, :, :32].permute(1, 0, 2))
+    actual = vp._rotate_mix_scalar_op(*args)
+    for a, b in zip(actual, expected, strict=True):
+        torch.testing.assert_close(a, b, rtol=0, atol=0)
+    if edges:
+        assert (
+            actual[0].untyped_storage().data_ptr()
+            != actual[1].untyped_storage().data_ptr()
+        )
+    assert actual[1].stride() == (64, edges * 64, 1)
+    chosen = {"both": (0, 1), "main": (0,), "scalar": (1,)}[branch]
+    upstream = tuple(torch.randn_like(expected[i]) for i in chosen)
+    leaves = (args[0], args[4], args[5]) + ((args[6],) if rank else ())
+    ref_grad = torch.autograd.grad(tuple(expected[i] for i in chosen), leaves, upstream)
+    got_grad = torch.autograd.grad(tuple(actual[i] for i in chosen), leaves, upstream)
+    for a, b in zip(got_grad, ref_grad, strict=True):
+        assert torch.isfinite(a).all()
+        torch.testing.assert_close(a, b, rtol=1e-4, atol=1e-5)
+
+
+def test_vjp_can_be_differentiated() -> None:
+    """Preserve the existing rotation's differentiable backward contract."""
+    args = _inputs(17, 1, False)
+    original = vp._rotate_mix_op(*args)
+    reference = (original, original[:, :, :32].permute(1, 0, 2))
+    actual = vp._rotate_mix_scalar_op(*args)
+
+    def derivatives(outputs):
+        loss = sum(v.square().sum() for v in outputs)
+        first = torch.autograd.grad(loss, args[0], create_graph=True)[0]
+        second = torch.autograd.grad(first.square().sum(), args[0])[0]
+        return first, second
+
+    expected_grads = derivatives(reference)
+    actual_grads = derivatives(actual)
+    for actual_grad, expected_grad in zip(actual_grads, expected_grads, strict=True):
+        assert torch.isfinite(actual_grad).all()
+        torch.testing.assert_close(actual_grad, expected_grad, rtol=1e-4, atol=1e-5)
+
+
+def test_float64_reference() -> None:
+    args = _inputs(17, 2)
+    reference_args = tuple(
+        v.detach().double().requires_grad_(v.requires_grad)
+        if isinstance(v, torch.Tensor) and v.is_floating_point()
+        else v
+        for v in args
+    )
+    # Wigner's off-degree entries are structural zeros, not independent
+    # variables. Project the dense reference onto the same tangent space.
+    block_mask = torch.zeros(16, 16, device="cuda", dtype=torch.float64)
+    for degree in range(4):
+        block_mask[degree**2 : (degree + 1) ** 2, degree**2 : (degree + 1) ** 2] = 1
+    ref = vp._rotate_mix_reference(
+        reference_args[0],
+        reference_args[1],
+        reference_args[4] * block_mask,
+        reference_args[5],
+        reference_args[6],
+        3,
+        2,
+        2,
+    )
+    got, scalar = vp._rotate_mix_scalar_op(*args)
+    torch.testing.assert_close(got.double(), ref, rtol=1e-4, atol=1e-5)
+    g, gs = torch.randn_like(got), torch.randn_like(scalar)
+    actual = torch.autograd.grad(
+        (got, scalar), (args[0], args[4], args[5], args[6]), (g, gs)
+    )
+    expected = torch.autograd.grad(
+        (ref, ref[:, :, :32].permute(1, 0, 2)),
+        tuple(reference_args[i] for i in (0, 4, 5, 6)),
+        (g.double(), gs.double()),
+    )
+    for a, b in zip(actual, expected, strict=True):
+        torch.testing.assert_close(a.double(), b, rtol=1e-4, atol=1e-5)
+
+
+def test_schema_and_fake() -> None:
+    torch.library.opcheck(
+        vp._rotate_mix_scalar_op,
+        _inputs(17, 1),
+        test_utils=("test_schema", "test_faketensor", "test_autograd_registration"),
+    )
+
+
+@pytest.mark.parametrize("capture", [False, True])
+def test_dynamic_compile(capture: bool) -> None:
+    op = vp._rotate_mix_scalar_op if capture else vp._rotate_mix_op
+    compiled = torch.compile(op, fullgraph=True, dynamic=True)
+    for edges in (17, 33):
+        args = _inputs(edges, 1)
+        actual = compiled(*args)
+        expected = op(*args)
+        if not capture:
+            actual, expected = (actual,), (expected,)
+        for a, b in zip(actual, expected, strict=True):
+            torch.testing.assert_close(a, b, rtol=0, atol=0)
+        g = tuple(torch.randn_like(t) for t in expected)
+        leaves = tuple(args[i] for i in (0, 4, 5))
+        a = torch.autograd.grad(actual, leaves, g)
+        b = torch.autograd.grad(expected, leaves, g)
+        for x, y in zip(a, b, strict=True):
+            torch.testing.assert_close(x, y, rtol=1e-4, atol=1e-5)
+
+
+@pytest.mark.parametrize(
+    "field,value",
+    [
+        ("training", True),
+        ("triton_infer_level", 1),
+        ("triton_infer_level", 3),
+        ("lmax", 2),
+        ("n_focus", 1),
+        ("so2_focus_dim", 64),
+        ("focus_compete", False),
+    ],
+)
+def test_scope(monkeypatch: pytest.MonkeyPatch, field: str, value: object) -> None:
+    conv = SimpleNamespace(
+        training=False,
+        triton_infer_level=2,
+        lmax=3,
+        n_focus=2,
+        so2_focus_dim=32,
+        focus_compete=True,
+    )
+    x = torch.empty(1, 16, 64, device="cuda")
+    monkeypatch.delenv("DP_SO2_SCALAR_CAPTURE_INFER", raising=False)
+    assert not vp._save_competition_scalar(conv, x)
+    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", "1")
+    assert vp._save_competition_scalar(conv, x)
+    assert not vp._save_competition_scalar(conv, x.cpu())
+    assert not vp._save_competition_scalar(conv, x.double())
+    setattr(conv, field, value)
+    assert not vp._save_competition_scalar(conv, x)

From 00452250af60c95019cc3f8d2f4a7ba8cf46e209 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 02:45:55 +0000
Subject: [PATCH 2/9] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../pt_expt/benchmark_so2_scalar_capture.py     | 17 +++++++++++++----
 .../pt_expt/kernels/test_so2_scalar_capture.py  |  4 +++-
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/source/tests/pt_expt/benchmark_so2_scalar_capture.py b/source/tests/pt_expt/benchmark_so2_scalar_capture.py
index 8c7f3b59f9..dcbe26266f 100644
--- a/source/tests/pt_expt/benchmark_so2_scalar_capture.py
+++ b/source/tests/pt_expt/benchmark_so2_scalar_capture.py
@@ -11,12 +11,14 @@
 import hashlib
 import json
 import os
-from pathlib import Path
 import site
 import subprocess
 import sys
 import time
 import traceback
+from pathlib import (
+    Path,
+)
 
 p = argparse.ArgumentParser(description=__doc__)
 p.add_argument("--model", type=Path, required=True)
@@ -96,8 +98,13 @@ def sha(path):
 import numpy as np
 import torch
 import triton
-from ase.build import bulk
-from deepmd.calculator import DP
+from ase.build import (
+    bulk,
+)
+
+from deepmd.calculator import (
+    DP,
+)
 from deepmd.utils.errors import OutOfMemoryError as DPOutOfMemoryError
 
 torch.set_num_threads(1)
@@ -260,7 +267,9 @@ def stats(values):
     torch.cuda.synchronize()
     result["steady_peak"] = peak()
     result["complete_peak"] = {k: max(full[k], result["steady_peak"][k]) for k in full}
-    from torch._inductor.codecache import PyCodeCache
+    from torch._inductor.codecache import (
+        PyCodeCache,
+    )
 
     wrappers = [
         str(m.__file__)
diff --git a/source/tests/pt_expt/kernels/test_so2_scalar_capture.py b/source/tests/pt_expt/kernels/test_so2_scalar_capture.py
index bd1d466ddb..7f34a7964f 100644
--- a/source/tests/pt_expt/kernels/test_so2_scalar_capture.py
+++ b/source/tests/pt_expt/kernels/test_so2_scalar_capture.py
@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Independent scalar storage emitted by the existing rotation producer."""
 
-from types import SimpleNamespace
+from types import (
+    SimpleNamespace,
+)
 
 import pytest
 import torch

From 7161a051a1cab70c34369c150abdd7fd0aca0938 Mon Sep 17 00:00:00 2001
From: Macqueen Ma-hnu <16522286+macqueen09@users.noreply.github.com>
Date: Tue, 3 Nov 2026 11:27:57 +0800
Subject: [PATCH 3/9] fix(pt): preserve SO2 scalar capture in exported models

---
 deepmd/pt/entrypoints/freeze_pt2.py           |   5 +-
 .../kernels/triton/sezm/so2_value_path.py     |  32 ++-
 deepmd/pt_expt/utils/serialization.py         |   6 +-
 doc/model/dpa4.md                             |   4 +-
 .../pt_expt/compare_so2_scalar_capture.py     |  21 +-
 .../kernels/test_so2_scalar_dispatch.py       | 272 ++++++++++++++++++
 .../test_compare_so2_scalar_capture.py        |  26 ++
 7 files changed, 358 insertions(+), 8 deletions(-)
 create mode 100644 source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py
 create mode 100644 source/tests/pt_expt/test_compare_so2_scalar_capture.py

diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py
index 81d47e45ea..094166b1ff 100644
--- a/deepmd/pt/entrypoints/freeze_pt2.py
+++ b/deepmd/pt/entrypoints/freeze_pt2.py
@@ -1007,10 +1007,13 @@ def _freeze_sezm_to_pt2(
 
     # Device-dependent Python branches resolve on the CPU tracing inputs, so
     # pin them to the AOTI target. Non-CPU targets bake the block-diagonal SO(2)
-    # contraction, while CUDA targets also bake the fused GIE scatter.
+    # contraction, while CUDA targets also bake the fused GIE scatter and
+    # permit opt-in SO(2) scalar capture despite the CPU tracing inputs.
     force_block_diag = target_device.type != "cpu"
     force_fused_scatter = target_device.type == "cuda"
     for module in model.modules():
+        if isinstance(module, SO2Convolution):
+            module._force_scalar_capture_cuda_target = target_device.type == "cuda"
         if isinstance(module, SO2Linear):
             module._force_block_diag_matmul = force_block_diag
         if isinstance(module, GeometricInitialEmbedding):
diff --git a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
index fefcaead40..b5f11a0fce 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
@@ -2557,7 +2557,16 @@ def _rotate_mix_scalar_impl(
     row = (3 * int(lmax) + 1) * focus_dim
     if not _use_triton(x):
         u = _rotate_mix_reference(x, src, wigner, kc, cb, lmax, n_focus, rank)
-        return u, u[:, :, :focus_dim].permute(1, 0, 2).clone()
+        # CPU export tracing must describe the same independent, padded
+        # scalar storage as the CUDA implementation.
+        scalar = torch.empty_strided(
+            (n_edge, n_focus, focus_dim),
+            (2 * focus_dim, n_edge * 2 * focus_dim, 1),
+            device=x.device,
+            dtype=x.dtype,
+        )
+        scalar.copy_(u[:, :, :focus_dim].permute(1, 0, 2))
+        return u, scalar
     u = torch.empty(n_focus, n_edge, row, device=x.device, dtype=x.dtype)
     pitch = 2 * focus_dim
     scalar = torch.empty_strided(
@@ -3721,6 +3730,24 @@ def _mixing_stack_train_bwd_impl(
 )(_gated_act_bwd_impl)
 
 
+@_rotate_mix_scalar_op.register_fake
+def _rotate_mix_scalar_fake(
+    x, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank
+):
+    """Describe both outputs without specializing symbolic edges in einsum."""
+    focus_dim = x.shape[2] // n_focus
+    n_edge = src.shape[0]
+    return (
+        x.new_empty((n_focus, n_edge, (3 * lmax + 1) * focus_dim)),
+        torch.empty_strided(
+            (n_edge, n_focus, focus_dim),
+            (2 * focus_dim, n_edge * 2 * focus_dim, 1),
+            device=x.device,
+            dtype=x.dtype,
+        ),
+    )
+
+
 @_rotate_mix_op.register_fake
 def _(x, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank):
     focus_dim = x.shape[2] // n_focus
@@ -5292,7 +5319,8 @@ def _save_competition_scalar(conv: SO2Convolution, x: Tensor) -> bool:
         and conv.so2_focus_dim == 32
         and conv.focus_compete
         and x.dtype == torch.float32
-        and x.is_cuda
+        # Export traces CPU inputs before moving the graph to its target.
+        and (x.is_cuda or getattr(conv, "_force_scalar_capture_cuda_target", False))
     )
 
 
diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py
index 8104cce43b..8160718cc5 100644
--- a/deepmd/pt_expt/utils/serialization.py
+++ b/deepmd/pt_expt/utils/serialization.py
@@ -1933,17 +1933,21 @@ def _trace_and_export_impl(
 
     # Device-dependent Python branches resolve on the CPU tracing inputs, so
     # pin them to the AOTI target. Non-CPU targets bake the block-diagonal SO(2)
-    # contraction, while CUDA targets also bake the fused GIE scatter.
+    # contraction, while CUDA targets also bake the fused GIE scatter and
+    # permit opt-in SO(2) scalar capture despite the CPU tracing inputs.
     from deepmd.pt_expt.descriptor.dpa4_nn.embedding import (
         GeometricInitialEmbedding,
     )
     from deepmd.pt_expt.descriptor.dpa4_nn.so2 import (
+        SO2Convolution,
         SO2Linear,
     )
 
     force_block_diag = target_device.type != "cpu"
     force_fused_scatter = target_device.type == "cuda"
     for module in model.modules():
+        if isinstance(module, SO2Convolution):
+            module._force_scalar_capture_cuda_target = target_device.type == "cuda"
         if isinstance(module, SO2Linear):
             module._force_block_diag_matmul = force_block_diag
         if isinstance(module, GeometricInitialEmbedding):
diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md
index 8f06e7581a..60730dca4c 100644
--- a/doc/model/dpa4.md
+++ b/doc/model/dpa4.md
@@ -575,7 +575,9 @@ live for competition backward. It does not change model weights or precision.
 The option defaults off and currently selects only evaluation with level `2`,
 CUDA float32, `lmax=3`, two focuses, 32 channels per focus and focus competition.
 Other layouts and training retain the original path. Set it before compilation
-or export; existing `.pt2` artifacts must be regenerated. This is a
+or export; existing `.pt2` artifacts must be regenerated. Both PT freeze and
+the `pt_expt` exporter select this branch from the CUDA deployment target,
+even though tracing runs on CPU. CPU targets retain the original path. This is a
 capacity optimization: the extra scalar store need not reduce latency.
 The manual full-API benchmark is
 `source/tests/pt_expt/benchmark_so2_scalar_capture.py` (`--help` lists inputs).
diff --git a/source/tests/pt_expt/compare_so2_scalar_capture.py b/source/tests/pt_expt/compare_so2_scalar_capture.py
index 545e408544..91b020c223 100644
--- a/source/tests/pt_expt/compare_so2_scalar_capture.py
+++ b/source/tests/pt_expt/compare_so2_scalar_capture.py
@@ -48,23 +48,38 @@ def compare(candidate, reference):
     return fields
 
 
+def validate_identity(result, reference, manifest, reference_manifest):
+    """Reject cross-checkpoint or cross-input comparisons before scoring."""
+    for key, left, right in (
+        ("input_sha256", result, reference),
+        ("model_sha256", manifest, reference_manifest),
+    ):
+        if not left.get(key) or left[key] != right.get(key):
+            raise ValueError(f"Comparison requires matching {key}")
+
+
 def main():
+    """Load benchmark routes and report compatible repeated comparisons."""
     parser = argparse.ArgumentParser(description=__doc__)
     for route in ("torch", "base", "candidate"):
         parser.add_argument(f"--{route}", type=Path, required=True)
     args = parser.parse_args()
-    results, arrays, out = {}, {}, {}
+    results, arrays, out, manifests = {}, {}, {}, {}
     for route, path in vars(args).items():
         results[route] = json.loads((path / "result.json").read_text())
         out[route] = {"status": results[route]["status"]}
         if results[route]["status"] == "completed":
+            manifests[route] = json.loads((path / "manifest.json").read_text())
             arrays[route] = dict(np.load(path / "outputs.npz"))
             out[route]["repeat"] = compare(arrays[route], arrays[route])
     for route in arrays:
         for reference in ("torch", "base"):
             if route != reference and reference in arrays:
-                assert (
-                    results[route]["input_sha256"] == results[reference]["input_sha256"]
+                validate_identity(
+                    results[route],
+                    results[reference],
+                    manifests[route],
+                    manifests[reference],
                 )
                 out[route][f"vs_{reference}"] = compare(
                     arrays[route], arrays[reference]
diff --git a/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py b/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py
new file mode 100644
index 0000000000..abea9c3b43
--- /dev/null
+++ b/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py
@@ -0,0 +1,272 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Scalar capture must be selected and consumed, not merely numerically valid."""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
+
+GPU = pytest.mark.skipif(
+    not torch.cuda.is_available() or not vp.SO2_VALUE_PATH_TRITON_AVAILABLE,
+    reason="CUDA and Triton are required",
+)
+
+
+@pytest.fixture
+def strict_environment(monkeypatch):
+    """Keep the tested route independent of ambient acceleration flags."""
+    for name, value in {
+        "DP_TRITON_INFER": "2",
+        "DP_CUDA_INFER": "0",
+        "DP_TRITON_TRAIN": "0",
+        "DP_CUDA_TRAIN": "0",
+        "DP_CUTE_INFER": "0",
+        "DP_CUTILE_INFER": "0",
+    }.items():
+        monkeypatch.setenv(name, value)
+    old = torch.backends.cuda.matmul.allow_tf32
+    torch.backends.cuda.matmul.allow_tf32 = False
+    yield
+    torch.backends.cuda.matmul.allow_tf32 = old
+
+
+def test_export_target_scope(monkeypatch):
+    """CPU tracing may opt in for CUDA export, but ordinary CPU stays off."""
+    conv = SimpleNamespace(
+        training=False,
+        triton_infer_level=2,
+        lmax=3,
+        n_focus=2,
+        so2_focus_dim=32,
+        focus_compete=True,
+    )
+    x = torch.empty(7, 16, 64, device="cpu")
+    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", "1")
+    assert not vp._save_competition_scalar(conv, x)
+    conv._force_scalar_capture_cuda_target = True
+    assert vp._save_competition_scalar(conv, x)
+    conv._force_scalar_capture_cuda_target = False
+    assert not vp._save_competition_scalar(conv, x)
+    conv._force_scalar_capture_cuda_target = True
+    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", "0")
+    assert not vp._save_competition_scalar(conv, x)
+
+
+@pytest.mark.parametrize("edges", [0, 1, 17])
+def test_cpu_scalar_layout_and_vjp(edges):
+    """CPU reference and fake/CUDA layouts share the same output contract."""
+    torch.manual_seed(42)
+    x = torch.randn(7, 16, 64, device="cpu", requires_grad=True)
+    src = torch.arange(edges, device="cpu") % 7
+    order = src.argsort(stable=True)
+    ptr = torch.cat((src.new_zeros(1), torch.bincount(src, minlength=7).cumsum(0)))
+    w = torch.eye(16, device="cpu").repeat(edges, 1, 1).requires_grad_()
+    kc = torch.randn(edges, 25, device="cpu", requires_grad=True)
+    cb = torch.randn(64, device="cpu", requires_grad=True)
+    args = (x, src, order, ptr, w, kc, cb, 3, 2, 1)
+    u, scalar = vp._rotate_mix_scalar_op(*args)
+    ref = vp._rotate_mix_op(*args)
+    torch.testing.assert_close(u, ref, atol=0, rtol=0)
+    torch.testing.assert_close(scalar, ref[:, :, :32].permute(1, 0, 2), atol=0, rtol=0)
+    assert scalar.stride() == (64, edges * 64, 1)
+    if edges:
+        assert scalar.untyped_storage().data_ptr() != u.untyped_storage().data_ptr()
+    gu, gs = torch.randn_like(u), torch.randn_like(scalar)
+    actual = torch.autograd.grad((u, scalar), (x, w, kc, cb), (gu, gs))
+    expected = torch.autograd.grad(
+        (ref, ref[:, :, :32].permute(1, 0, 2)), (x, w, kc, cb), (gu, gs)
+    )
+    for a, b in zip(actual, expected, strict=True):
+        torch.testing.assert_close(a, b, atol=1e-5, rtol=1e-4)
+
+
+@GPU
+@pytest.mark.parametrize("backend", ["pt", "pt_expt"])
+@pytest.mark.parametrize("capture", [False, True])
+def test_value_path_dispatch(monkeypatch, strict_environment, backend, capture):
+    """Exercise the production callable and prove the competition data source."""
+    if backend == "pt":
+        from deepmd.pt.model.descriptor.sezm_nn.so2 import SO2Convolution
+
+        precision = {"dtype": torch.float32}
+    else:
+        from deepmd.pt_expt.descriptor.dpa4_nn.so2 import SO2Convolution
+
+        precision = {"precision": "float32"}
+    conv = (
+        SO2Convolution(
+            lmax=3,
+            mmax=1,
+            channels=32,
+            n_focus=2,
+            focus_dim=32,
+            mixing_layers=3,
+            radial_so2_mode="degree_channel",
+            radial_so2_rank=1,
+            n_atten_head=1,
+            seed=7,
+            trainable=False,
+            **precision,
+        )
+        .to("cuda")
+        .eval()
+    )
+    path = vp.make_triton_value_path(conv)
+    assert path is not None
+    x = torch.randn(7, 16, 64, device="cuda") * 0.1
+    radial = torch.randn(17, 4, 32, device="cuda") * 0.1
+    cache = SimpleNamespace(
+        src=torch.arange(17, device="cuda") % 7,
+        D_full=torch.eye(16, device="cuda").repeat(17, 1, 1),
+    )
+    seen = {}
+    original = vp._rotate_mix_scalar_op
+
+    def capture_spy(*args):
+        u, scalar = original(*args)
+        # A distinct sentinel proves provenance: a wrong slice has identical
+        # normal values and would otherwise pass an equality-only assertion.
+        scalar = scalar + 0.125
+        seen["scalar"] = scalar
+        return u, scalar
+
+    original_alpha = conv._focus_alpha
+
+    def alpha_spy(gate_src):
+        if capture:
+            torch.testing.assert_close(gate_src, seen["scalar"], rtol=0, atol=0)
+        return original_alpha(gate_src)
+
+    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", str(int(capture)))
+    with (
+        patch.object(
+            vp, "_rotate_mix_scalar_op", side_effect=capture_spy
+        ) as scalar_call,
+        patch.object(vp, "_rotate_mix_op", wraps=vp._rotate_mix_op) as default_call,
+        patch.object(conv, "_focus_alpha", side_effect=alpha_spy) as alpha_call,
+    ):
+        path(x, cache, radial)
+    assert scalar_call.call_count == int(capture)
+    assert default_call.call_count == int(not capture)
+    assert alpha_call.call_count == 1
+
+
+def _model_params():
+    """A small two-block model with the supported scalar-capture layout."""
+    return {
+        "type": "SeZM",
+        "type_map": ["C"],
+        "use_compile": False,
+        "descriptor": {
+            "type": "SeZM",
+            "sel": [32],
+            "rcut": 4.0,
+            "channels": 32,
+            "n_focus": 2,
+            "focus_dim": 32,
+            "n_radial": 4,
+            "radial_mlp": [32],
+            "use_env_seed": True,
+            "l_schedule": [3, 3, 0],
+            "mmax": 1,
+            "so2_layers": 3,
+            "n_atten_head": 1,
+            "ffn_neurons": 32,
+            "ffn_blocks": 1,
+            "use_amp": False,
+            "precision": "float32",
+            "seed": 7,
+        },
+        "fitting_net": {"neuron": [16], "precision": "float32", "seed": 7},
+    }
+
+
+@GPU
+def test_model_capture_on_off(monkeypatch, strict_environment):
+    """A real frozen model preserves outputs and executes scalar capture."""
+    from deepmd.pt.model.model import get_model
+
+    model = get_model(_model_params()).to("cuda").eval()
+    model.requires_grad_(False)
+    coord = torch.tensor(
+        [[[0.0, 0.0, 0.0], [1.2, 0.2, 0.1], [0.3, 1.1, 0.2], [0.1, 0.2, 1.3]]],
+        device="cuda",
+        dtype=torch.float64,
+    )
+    atype = torch.zeros(1, 4, device="cuda", dtype=torch.long)
+    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", "0")
+    reference = model(coord, atype, None)
+    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", "1")
+    with patch.object(
+        vp, "_rotate_mix_scalar_op", wraps=vp._rotate_mix_scalar_op
+    ) as spy:
+        actual = model(coord, atype, None)
+    assert spy.call_count > 0
+    assert actual.keys() == reference.keys()
+    for name in actual:
+        if isinstance(actual[name], torch.Tensor):
+            assert torch.isfinite(actual[name]).all()
+            torch.testing.assert_close(
+                actual[name], reference[name], rtol=1e-4, atol=1e-5
+            )
+
+
+@GPU
+@pytest.mark.parametrize("backend", ["pt", "pt_expt"])
+@pytest.mark.parametrize(
+    "capture,target", [(True, "cuda"), (False, "cuda"), (True, "cpu")]
+)
+def test_export_dispatch(
+    monkeypatch, strict_environment, tmp_path, backend, capture, target
+):
+    """Both official exporters must select capture before CPU tracing ends.
+
+    Stop at the export boundary, avoiding compiler/toolchain dependencies in
+    this routing regression. Full AOTI package inference is benchmarked separately.
+    """
+    from deepmd.pt.entrypoints import freeze_pt2
+    from deepmd.pt.model.model import get_model
+    from deepmd.pt.train.wrapper import ModelWrapper
+    from deepmd.pt_expt.kernels import autotune
+    from deepmd.pt_expt.utils.serialization import _trace_and_export_impl
+
+    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", str(int(capture)))
+    monkeypatch.setattr(freeze_pt2, "_tune_triton_configs", lambda *args: None)
+    monkeypatch.setattr(autotune, "run_autotune", lambda *args: None)
+    params = _model_params()
+    model = get_model(params).cpu().eval()
+    seen = []
+
+    class ExportBoundary(Exception):
+        pass
+
+    def inspect_graph(traced, *args, **kwargs):
+        targets = [str(node.target) for node in traced.graph.nodes]
+        count = sum("sezm_triton.so2_rotate_mix_scalar." in name for name in targets)
+        seen.append(count)
+        assert (count > 0) == (capture and target == "cuda")
+        raise ExportBoundary
+
+    monkeypatch.setattr(torch.export, "export", inspect_graph)
+    with pytest.raises(ExportBoundary):
+        if backend == "pt":
+            checkpoint = tmp_path / "model.pt"
+            wrapper = ModelWrapper(model, model_params=params)
+            torch.save({"model": wrapper.state_dict()}, checkpoint)
+            freeze_pt2.freeze_sezm_to_pt2(
+                str(checkpoint),
+                str(tmp_path / "model.pt2"),
+                device=torch.device(target),
+                with_comm=False,
+            )
+        else:
+            _trace_and_export_impl(
+                {"model": model.serialize(), "model_def_script": params},
+                lower_kind="graph",
+                do_atomic_virial=True,
+                target_device=torch.device(target),
+            )
+    assert len(seen) == 1
diff --git a/source/tests/pt_expt/test_compare_so2_scalar_capture.py b/source/tests/pt_expt/test_compare_so2_scalar_capture.py
new file mode 100644
index 0000000000..d79ff6d3ca
--- /dev/null
+++ b/source/tests/pt_expt/test_compare_so2_scalar_capture.py
@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Benchmark comparisons must use the same checkpoint and input frames."""
+
+import pytest
+
+from .compare_so2_scalar_capture import validate_identity
+
+
+@pytest.mark.parametrize("key", ["model_sha256", "input_sha256"])
+@pytest.mark.parametrize("missing", [False, True])
+def test_reject_mismatched_identity(key, missing):
+    """Reject a changed or absent identity instead of reporting accuracy."""
+    left = {"model_sha256": "model", "input_sha256": "input"}
+    right = dict(left)
+    if missing:
+        right.pop(key)
+    else:
+        right[key] = "different"
+    with pytest.raises(ValueError, match=key):
+        validate_identity(left, right, left, right)
+
+
+def test_matching_identity():
+    """An identical checkpoint and input pair remains comparable."""
+    identity = {"model_sha256": "model", "input_sha256": "input"}
+    validate_identity(identity, identity, identity, identity)

From 5e0e1c16a9ba27f962f196363ffc9dfe5b654d51 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 3 Nov 2026 03:34:11 +0000
Subject: [PATCH 4/9] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../kernels/test_so2_scalar_dispatch.py       | 40 ++++++++++++++-----
 .../test_compare_so2_scalar_capture.py        |  4 +-
 2 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py b/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py
index abea9c3b43..073d4c3f78 100644
--- a/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py
+++ b/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py
@@ -1,8 +1,12 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Scalar capture must be selected and consumed, not merely numerically valid."""
 
-from types import SimpleNamespace
-from unittest.mock import patch
+from types import (
+    SimpleNamespace,
+)
+from unittest.mock import (
+    patch,
+)
 
 import pytest
 import torch
@@ -89,11 +93,15 @@ def test_cpu_scalar_layout_and_vjp(edges):
 def test_value_path_dispatch(monkeypatch, strict_environment, backend, capture):
     """Exercise the production callable and prove the competition data source."""
     if backend == "pt":
-        from deepmd.pt.model.descriptor.sezm_nn.so2 import SO2Convolution
+        from deepmd.pt.model.descriptor.sezm_nn.so2 import (
+            SO2Convolution,
+        )
 
         precision = {"dtype": torch.float32}
     else:
-        from deepmd.pt_expt.descriptor.dpa4_nn.so2 import SO2Convolution
+        from deepmd.pt_expt.descriptor.dpa4_nn.so2 import (
+            SO2Convolution,
+        )
 
         precision = {"precision": "float32"}
     conv = (
@@ -187,7 +195,9 @@ def _model_params():
 @GPU
 def test_model_capture_on_off(monkeypatch, strict_environment):
     """A real frozen model preserves outputs and executes scalar capture."""
-    from deepmd.pt.model.model import get_model
+    from deepmd.pt.model.model import (
+        get_model,
+    )
 
     model = get_model(_model_params()).to("cuda").eval()
     model.requires_grad_(False)
@@ -227,11 +237,21 @@ def test_export_dispatch(
     Stop at the export boundary, avoiding compiler/toolchain dependencies in
     this routing regression. Full AOTI package inference is benchmarked separately.
     """
-    from deepmd.pt.entrypoints import freeze_pt2
-    from deepmd.pt.model.model import get_model
-    from deepmd.pt.train.wrapper import ModelWrapper
-    from deepmd.pt_expt.kernels import autotune
-    from deepmd.pt_expt.utils.serialization import _trace_and_export_impl
+    from deepmd.pt.entrypoints import (
+        freeze_pt2,
+    )
+    from deepmd.pt.model.model import (
+        get_model,
+    )
+    from deepmd.pt.train.wrapper import (
+        ModelWrapper,
+    )
+    from deepmd.pt_expt.kernels import (
+        autotune,
+    )
+    from deepmd.pt_expt.utils.serialization import (
+        _trace_and_export_impl,
+    )
 
     monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", str(int(capture)))
     monkeypatch.setattr(freeze_pt2, "_tune_triton_configs", lambda *args: None)
diff --git a/source/tests/pt_expt/test_compare_so2_scalar_capture.py b/source/tests/pt_expt/test_compare_so2_scalar_capture.py
index d79ff6d3ca..74400c6ad8 100644
--- a/source/tests/pt_expt/test_compare_so2_scalar_capture.py
+++ b/source/tests/pt_expt/test_compare_so2_scalar_capture.py
@@ -3,7 +3,9 @@
 
 import pytest
 
-from .compare_so2_scalar_capture import validate_identity
+from .compare_so2_scalar_capture import (
+    validate_identity,
+)
 
 
 @pytest.mark.parametrize("key", ["model_sha256", "input_sha256"])

From 39094625cb967ec2ea8a4157834bf86bb9eca487 Mon Sep 17 00:00:00 2001
From: Macqueen Ma-hnu <16522286+macqueen09@users.noreply.github.com>
Date: Tue, 3 Nov 2026 20:57:52 +0800
Subject: [PATCH 5/9] fix(pt): preserve SO2 gradients for empty CPU edges

---
 .../kernels/triton/sezm/so2_value_path.py     |  15 +++
 .../kernels/test_so2_scalar_dispatch.py       | 118 ++++++++++++++++--
 .../test_compare_so2_scalar_capture.py        |  29 ++++-
 3 files changed, 149 insertions(+), 13 deletions(-)

diff --git a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
index b5f11a0fce..235c142c15 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
@@ -189,6 +189,12 @@ def _rotate_mix_reference(
     dim = (lmax + 1) ** 2
     n_deg = lmax + 1
     reduced = 3 * lmax + 1
+    if _has_no_edges(n_edge):
+        # opt_einsum 3.3 cannot plan zero-cost contractions (0 / 0).
+        # Empty reductions retain the zero derivative to every operand,
+        # including the nonempty node features and shared channel basis.
+        dependency = x[:0].sum() + wigner.sum() + kc.sum() + cb[:0].sum()
+        return x.new_empty((n_focus, n_edge, reduced * focus_dim)) + dependency
     coeff = build_m_major_index(lmax, 1, device=x.device)
     d_to_m = wigner[:, :dim, :dim].index_select(1, coeff)
     x_local = torch.bmm(d_to_m, x.index_select(0, src))  # (E, reduced, C_wide)
@@ -249,6 +255,15 @@ def _rotate_mix_backward_reference(
     dim = (lmax + 1) ** 2
     n_deg = lmax + 1
     reduced = 3 * lmax + 1
+    if _has_no_edges(n_edge):
+        # All three edge cotangents are empty. Preserve their differentiable
+        # dependency without sending zero-size einsums to the path planner.
+        dependency = grad_u.sum() + x[:0].sum() + wigner.sum() + kc.sum() + cb[:0].sum()
+        return (
+            x.new_empty((n_edge, dim, c_wide)) + dependency,
+            wigner.new_empty(wigner.shape) + dependency,
+            kc.new_empty(kc.shape) + dependency,
+        )
     coeff = build_m_major_index(lmax, 1, device=x.device)
     d_to_m = wigner[:, :dim, :dim].index_select(1, coeff)
     x_src = x.index_select(0, src)
diff --git a/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py b/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py
index 073d4c3f78..8e5ce8c229 100644
--- a/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py
+++ b/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py
@@ -60,7 +60,8 @@ def test_export_target_scope(monkeypatch):
 
 
 @pytest.mark.parametrize("edges", [0, 1, 17])
-def test_cpu_scalar_layout_and_vjp(edges):
+@pytest.mark.parametrize("rank", [0, 1, 2])
+def test_cpu_scalar_layout_and_vjp(edges, rank):
     """CPU reference and fake/CUDA layouts share the same output contract."""
     torch.manual_seed(42)
     x = torch.randn(7, 16, 64, device="cpu", requires_grad=True)
@@ -68,9 +69,12 @@ def test_cpu_scalar_layout_and_vjp(edges):
     order = src.argsort(stable=True)
     ptr = torch.cat((src.new_zeros(1), torch.bincount(src, minlength=7).cumsum(0)))
     w = torch.eye(16, device="cpu").repeat(edges, 1, 1).requires_grad_()
-    kc = torch.randn(edges, 25, device="cpu", requires_grad=True)
-    cb = torch.randn(64, device="cpu", requires_grad=True)
-    args = (x, src, order, ptr, w, kc, cb, 3, 2, 1)
+    kc_shape = (edges, 4, 64) if rank == 0 else (edges, 25 * rank)
+    kc = torch.randn(kc_shape, device="cpu", requires_grad=True)
+    cb = torch.randn(
+        1 if rank == 0 else 64 * rank, device="cpu", requires_grad=rank > 0
+    )
+    args = (x, src, order, ptr, w, kc, cb, 3, 2, rank)
     u, scalar = vp._rotate_mix_scalar_op(*args)
     ref = vp._rotate_mix_op(*args)
     torch.testing.assert_close(u, ref, atol=0, rtol=0)
@@ -79,12 +83,106 @@ def test_cpu_scalar_layout_and_vjp(edges):
     if edges:
         assert scalar.untyped_storage().data_ptr() != u.untyped_storage().data_ptr()
     gu, gs = torch.randn_like(u), torch.randn_like(scalar)
-    actual = torch.autograd.grad((u, scalar), (x, w, kc, cb), (gu, gs))
+    leaves = (x, w, kc) + ((cb,) if rank else ())
+    actual = torch.autograd.grad((u, scalar), leaves, (gu, gs))
     expected = torch.autograd.grad(
-        (ref, ref[:, :, :32].permute(1, 0, 2)), (x, w, kc, cb), (gu, gs)
+        (ref, ref[:, :, :32].permute(1, 0, 2)), leaves, (gu, gs)
     )
     for a, b in zip(actual, expected, strict=True):
         torch.testing.assert_close(a, b, atol=1e-5, rtol=1e-4)
+        assert torch.isfinite(a).all()
+        if not edges:
+            torch.testing.assert_close(a, torch.zeros_like(a), atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("branch", ["both", "main", "scalar"])
+def test_empty_cpu_second_derivative(branch):
+    """Empty output branches retain zero first and second derivatives."""
+    x = torch.randn(3, 16, 64, device="cpu", requires_grad=True)
+    src = torch.empty(0, device="cpu", dtype=torch.long)
+    ptr = torch.zeros(4, device="cpu", dtype=torch.long)
+    w = torch.empty(0, 16, 16, device="cpu", requires_grad=True)
+    kc = torch.empty(0, 25, device="cpu", requires_grad=True)
+    cb = torch.randn(64, device="cpu", requires_grad=True)
+    outputs = vp._rotate_mix_scalar_op(x, src, src, ptr, w, kc, cb, 3, 2, 1)
+    chosen = {"both": outputs, "main": outputs[:1], "scalar": outputs[1:]}[branch]
+    leaves = (x, w, kc, cb)
+    first = torch.autograd.grad(
+        sum(v.square().sum() for v in chosen), leaves, create_graph=True
+    )
+    second = torch.autograd.grad(sum(v.square().sum() for v in first), leaves)
+    for gradient in (*first, *second):
+        torch.testing.assert_close(gradient, torch.zeros_like(gradient), rtol=0, atol=0)
+
+
+@GPU
+def test_capture_freeze_load(monkeypatch, strict_environment, tmp_path):
+    """Compile a real capture-enabled package and compare loaded predictions."""
+    import re
+    import zipfile
+
+    import deepmd.pt_expt.utils.env as export_env
+    from deepmd.infer import (
+        DeepPot,
+    )
+    from deepmd.pt.entrypoints import (
+        freeze_pt2,
+    )
+    from deepmd.pt.model.model import (
+        get_model,
+    )
+    from deepmd.pt.train.wrapper import (
+        ModelWrapper,
+    )
+
+    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", "1")
+    # The launch-table sweep is a performance concern; retain the real
+    # tracing, torch.export, AOTI compilation and package loader below.
+    monkeypatch.setattr(freeze_pt2, "_tune_triton_configs", lambda *args: None)
+    monkeypatch.setattr(export_env, "DEVICE", torch.device("cuda:0"))
+    params = _model_params()
+    params["descriptor"].update(radial_so2_mode="degree_channel", radial_so2_rank=1)
+    model = get_model(params).to("cuda").eval()
+    model.requires_grad_(False)
+    coord = torch.tensor(
+        [[[0.0, 0.0, 0.0], [1.2, 0.2, 0.1], [0.3, 1.1, 0.2], [0.1, 0.2, 1.3]]],
+        device="cuda",
+        dtype=torch.float64,
+    )
+    atype = torch.zeros(1, 4, device="cuda", dtype=torch.long)
+    reference = model(coord, atype, None)
+    checkpoint, package = tmp_path / "model.pt", tmp_path / "model.pt2"
+    wrapper = ModelWrapper(model, model_params=params)
+    torch.save({"model": wrapper.state_dict()}, checkpoint)
+    freeze_pt2.freeze_sezm_to_pt2(
+        str(checkpoint),
+        str(package),
+        device=torch.device("cuda:0"),
+        with_comm=False,
+    )
+    with zipfile.ZipFile(package) as archive:
+        wrappers = [
+            archive.read(name).decode()
+            for name in archive.namelist()
+            if name.endswith(".wrapper.cpp")
+        ]
+    assert any(
+        re.search(r"['\"]SAVE_SCALAR['\"]\s*:\s*True", text) for text in wrappers
+    )
+    evaluator = DeepPot(str(package))
+    energy, force, virial = evaluator.eval(
+        coord.detach().cpu().numpy(),
+        None,
+        [0] * 4,
+    )
+    for actual, name in ((energy, "energy"), (force, "force"), (virial, "virial")):
+        expected = reference[name].detach().cpu()
+        torch.testing.assert_close(
+            torch.as_tensor(actual).reshape(expected.shape),
+            expected,
+            rtol=1e-4,
+            atol=1e-5,
+        )
 
 
 @GPU
@@ -224,10 +322,14 @@ def test_model_capture_on_off(monkeypatch, strict_environment):
             )
 
 
-@GPU
 @pytest.mark.parametrize("backend", ["pt", "pt_expt"])
 @pytest.mark.parametrize(
-    "capture,target", [(True, "cuda"), (False, "cuda"), (True, "cpu")]
+    "capture,target",
+    [
+        pytest.param(True, "cuda", marks=GPU),
+        pytest.param(False, "cuda", marks=GPU),
+        (True, "cpu"),
+    ],
 )
 def test_export_dispatch(
     monkeypatch, strict_environment, tmp_path, backend, capture, target
diff --git a/source/tests/pt_expt/test_compare_so2_scalar_capture.py b/source/tests/pt_expt/test_compare_so2_scalar_capture.py
index 74400c6ad8..0e0e970e75 100644
--- a/source/tests/pt_expt/test_compare_so2_scalar_capture.py
+++ b/source/tests/pt_expt/test_compare_so2_scalar_capture.py
@@ -1,9 +1,11 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Benchmark comparisons must use the same checkpoint and input frames."""
 
+import numpy as np
 import pytest
 
 from .compare_so2_scalar_capture import (
+    compare,
     validate_identity,
 )
 
@@ -12,17 +14,34 @@
 @pytest.mark.parametrize("missing", [False, True])
 def test_reject_mismatched_identity(key, missing):
     """Reject a changed or absent identity instead of reporting accuracy."""
-    left = {"model_sha256": "model", "input_sha256": "input"}
-    right = dict(left)
+    result = {"input_sha256": "input"}
+    manifest = {"model_sha256": "model"}
+    right_result, right_manifest = dict(result), dict(manifest)
+    right = right_result if key == "input_sha256" else right_manifest
     if missing:
         right.pop(key)
     else:
         right[key] = "different"
     with pytest.raises(ValueError, match=key):
-        validate_identity(left, right, left, right)
+        validate_identity(result, right_result, manifest, right_manifest)
 
 
 def test_matching_identity():
     """An identical checkpoint and input pair remains comparable."""
-    identity = {"model_sha256": "model", "input_sha256": "input"}
-    validate_identity(identity, identity, identity, identity)
+    result = {"input_sha256": "input"}
+    manifest = {"model_sha256": "model"}
+    validate_identity(result, result, manifest, manifest)
+
+
+def test_compare_checks_cross_repeat_pairs():
+    """Equal matching repeats must not hide variability across repeats."""
+    arrays = {
+        f"{repeat}_{frame}_{field}": np.array([1.0 + repeat])
+        for repeat in range(3)
+        for frame in range(3)
+        for field in ("energy", "forces", "virial", "stress")
+    }
+    result = compare(arrays, arrays)
+    for field in result.values():
+        assert not field["passed"]
+        assert field["max_abs"] == 2.0

From 3724fe3104f63795e64da35ba0a206c2cdba99dd Mon Sep 17 00:00:00 2001
From: Macqueen Ma-hnu <16522286+macqueen09@users.noreply.github.com>
Date: Wed, 4 Nov 2026 23:40:54 +0800
Subject: [PATCH 6/9] fix(pt): select SO2 scalar capture by inference path

---
 deepmd/pt/entrypoints/freeze_pt2.py           |  3 +-
 .../kernels/triton/sezm/so2_value_path.py     | 11 +--
 deepmd/pt_expt/utils/serialization.py         |  5 +-
 doc/model/dpa4.md                             | 18 +++--
 .../pt_expt/benchmark_so2_scalar_capture.py   | 17 +++--
 .../kernels/test_so2_scalar_capture.py        |  4 +-
 .../kernels/test_so2_scalar_dispatch.py       | 70 ++++++++++++++-----
 .../test_compare_so2_scalar_capture.py        | 30 ++++++--
 8 files changed, 112 insertions(+), 46 deletions(-)

diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py
index 094166b1ff..896a348cfe 100644
--- a/deepmd/pt/entrypoints/freeze_pt2.py
+++ b/deepmd/pt/entrypoints/freeze_pt2.py
@@ -1008,11 +1008,12 @@ def _freeze_sezm_to_pt2(
     # Device-dependent Python branches resolve on the CPU tracing inputs, so
     # pin them to the AOTI target. Non-CPU targets bake the block-diagonal SO(2)
     # contraction, while CUDA targets also bake the fused GIE scatter and
-    # permit opt-in SO(2) scalar capture despite the CPU tracing inputs.
+    # permit level-2 SO(2) scalar capture despite the CPU tracing inputs.
     force_block_diag = target_device.type != "cpu"
     force_fused_scatter = target_device.type == "cuda"
     for module in model.modules():
         if isinstance(module, SO2Convolution):
+            module._scalar_capture_retains_rotation = True
             module._force_scalar_capture_cuda_target = target_device.type == "cuda"
         if isinstance(module, SO2Linear):
             module._force_block_diag_matmul = force_block_diag
diff --git a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
index 235c142c15..5377342d7d 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
@@ -91,7 +91,6 @@
     annotations,
 )
 
-import os
 from typing import (
     TYPE_CHECKING,
     NamedTuple,
@@ -5323,10 +5322,14 @@ def make_triton_rotate_mix(conv: SO2Convolution) -> _TritonRotateMix | None:
 
 
 def _save_competition_scalar(conv: SO2Convolution, x: Tensor) -> bool:
-    """Opt in only to the validated frozen two-focus FP32 inference layout."""
+    """Capture only where competition backward retains the rotation storage.
+
+    The PT-expt exporter can already release that storage without capture;
+    it pins this internal policy off before tracing. Runtime and PT export
+    retain the independent scalar buffer for the validated level-2 layout.
+    """
     return (
-        os.environ.get("DP_SO2_SCALAR_CAPTURE_INFER", "0").strip().lower()
-        in ("1", "true", "yes", "on")
+        getattr(conv, "_scalar_capture_retains_rotation", True)
         and not conv.training
         and conv.triton_infer_level == 2
         and conv.lmax == 3
diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py
index 8160718cc5..6572129c73 100644
--- a/deepmd/pt_expt/utils/serialization.py
+++ b/deepmd/pt_expt/utils/serialization.py
@@ -1934,7 +1934,9 @@ def _trace_and_export_impl(
     # Device-dependent Python branches resolve on the CPU tracing inputs, so
     # pin them to the AOTI target. Non-CPU targets bake the block-diagonal SO(2)
     # contraction, while CUDA targets also bake the fused GIE scatter and
-    # permit opt-in SO(2) scalar capture despite the CPU tracing inputs.
+    # identify the CUDA target despite the CPU tracing inputs. This exporter
+    # already releases the full rotation output: a separate scalar store adds
+    # latency without reducing its peak, so keep the non-capture route.
     from deepmd.pt_expt.descriptor.dpa4_nn.embedding import (
         GeometricInitialEmbedding,
     )
@@ -1947,6 +1949,7 @@ def _trace_and_export_impl(
     force_fused_scatter = target_device.type == "cuda"
     for module in model.modules():
         if isinstance(module, SO2Convolution):
+            module._scalar_capture_retains_rotation = False
             module._force_scalar_capture_cuda_target = target_device.type == "cuda"
         if isinstance(module, SO2Linear):
             module._force_block_diag_matmul = force_block_diag
diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md
index 60730dca4c..59283a147b 100644
--- a/doc/model/dpa4.md
+++ b/doc/model/dpa4.md
@@ -568,17 +568,21 @@ shipping built in. On other GPUs the kernels fall back to conservative
 configurations, and `dp --pt freeze` tunes the missing entries on the local GPU
 before exporting, a one-off sweep of a few minutes baked into the `.pt2`.
 
-`DP_SO2_SCALAR_CAPTURE_INFER=1` optionally saves the focus-competition scalars
+Level `2` saves the focus-competition scalars on supported inference paths
 into independent storage during the existing rotate-and-mix launch. This lets
 the full rotation output die after its mixing consumer, rather than remain
 live for competition backward. It does not change model weights or precision.
-The option defaults off and currently selects only evaluation with level `2`,
+Selection is internal and currently applies only to evaluation with level `2`,
 CUDA float32, `lmax=3`, two focuses, 32 channels per focus and focus competition.
-Other layouts and training retain the original path. Set it before compilation
-or export; existing `.pt2` artifacts must be regenerated. Both PT freeze and
-the `pt_expt` exporter select this branch from the CUDA deployment target,
-even though tracing runs on CPU. CPU targets retain the original path. This is a
-capacity optimization: the extra scalar store need not reduce latency.
+Other layouts and training retain the original path. PT freeze selects this
+branch from the CUDA deployment target even though tracing runs on CPU;
+existing `.pt2` artifacts must be regenerated. CPU targets and `pt_expt` export
+retain the original path. The latter already releases the large rotation
+storage, so capture adds no measured peak-memory benefit there. This is a
+capacity optimization: the extra scalar store can increase latency. Measurements
+on RTX 5090 with the Neo checkpoint found about 12.35%/12.38% lower peak allocated
+memory for PT export at 1000/4096 atoms, with about 0.65%/1.09% higher latency.
+These measurements do not imply the same benefit for other execution paths.
 The manual full-API benchmark is
 `source/tests/pt_expt/benchmark_so2_scalar_capture.py` (`--help` lists inputs).
 
diff --git a/source/tests/pt_expt/benchmark_so2_scalar_capture.py b/source/tests/pt_expt/benchmark_so2_scalar_capture.py
index dcbe26266f..d2095afe2b 100644
--- a/source/tests/pt_expt/benchmark_so2_scalar_capture.py
+++ b/source/tests/pt_expt/benchmark_so2_scalar_capture.py
@@ -71,7 +71,6 @@ def sha(path):
     DP_CUDA_INFER="0",
     DP_TF32_INFER="0",
     DP_AMP_INFER="0",
-    DP_SO2_SCALAR_CAPTURE_INFER="1" if a.route == "CSC5" else "0",
     OMP_NUM_THREADS="1",
     DP_INTRA_OP_PARALLELISM_THREADS="1",
     DP_INTER_OP_PARALLELISM_THREADS="1",
@@ -105,8 +104,14 @@ def sha(path):
 from deepmd.calculator import (
     DP,
 )
+from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as scalar_path
 from deepmd.utils.errors import OutOfMemoryError as DPOutOfMemoryError
 
+# BASE is an explicit benchmark-only ablation, not a public inference switch.
+# On a pre-feature checkout BASE naturally uses its original implementation.
+if a.route == "BASE" and hasattr(scalar_path, "_save_competition_scalar"):
+    scalar_path._save_competition_scalar = lambda conv, x: False
+
 torch.set_num_threads(1)
 torch.set_float32_matmul_precision("highest")
 torch.backends.cuda.matmul.allow_tf32 = False
@@ -157,12 +162,12 @@ def sha(path):
 try:
     atoms = bulk("C", "diamond", a=3.567, cubic=True).repeat(tuple(a.cell))
     rng = np.random.default_rng(a.seed)
-    frames = np.stack(
-        [
+    frames = []
+    for _ in range(3):
+        frames.append(
             atoms.positions + 0.02 * rng.standard_normal(atoms.positions.shape)
-            for _ in range(3)
-        ]
-    )
+        )
+    frames = np.stack(frames)
     np.savez_compressed(
         a.out / "inputs.npz",
         positions=frames,
diff --git a/source/tests/pt_expt/kernels/test_so2_scalar_capture.py b/source/tests/pt_expt/kernels/test_so2_scalar_capture.py
index 7f34a7964f..e549917039 100644
--- a/source/tests/pt_expt/kernels/test_so2_scalar_capture.py
+++ b/source/tests/pt_expt/kernels/test_so2_scalar_capture.py
@@ -169,6 +169,7 @@ def test_dynamic_compile(capture: bool) -> None:
         ("n_focus", 1),
         ("so2_focus_dim", 64),
         ("focus_compete", False),
+        ("_scalar_capture_retains_rotation", False),
     ],
 )
 def test_scope(monkeypatch: pytest.MonkeyPatch, field: str, value: object) -> None:
@@ -181,9 +182,6 @@ def test_scope(monkeypatch: pytest.MonkeyPatch, field: str, value: object) -> No
         focus_compete=True,
     )
     x = torch.empty(1, 16, 64, device="cuda")
-    monkeypatch.delenv("DP_SO2_SCALAR_CAPTURE_INFER", raising=False)
-    assert not vp._save_competition_scalar(conv, x)
-    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", "1")
     assert vp._save_competition_scalar(conv, x)
     assert not vp._save_competition_scalar(conv, x.cpu())
     assert not vp._save_competition_scalar(conv, x.double())
diff --git a/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py b/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py
index 8e5ce8c229..2abf9ffce9 100644
--- a/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py
+++ b/source/tests/pt_expt/kernels/test_so2_scalar_dispatch.py
@@ -37,8 +37,8 @@ def strict_environment(monkeypatch):
     torch.backends.cuda.matmul.allow_tf32 = old
 
 
-def test_export_target_scope(monkeypatch):
-    """CPU tracing may opt in for CUDA export, but ordinary CPU stays off."""
+def test_export_target_scope():
+    """Target and retained-storage policy both constrain CPU tracing."""
     conv = SimpleNamespace(
         training=False,
         triton_infer_level=2,
@@ -48,14 +48,44 @@ def test_export_target_scope(monkeypatch):
         focus_compete=True,
     )
     x = torch.empty(7, 16, 64, device="cpu")
-    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", "1")
     assert not vp._save_competition_scalar(conv, x)
     conv._force_scalar_capture_cuda_target = True
     assert vp._save_competition_scalar(conv, x)
     conv._force_scalar_capture_cuda_target = False
     assert not vp._save_competition_scalar(conv, x)
     conv._force_scalar_capture_cuda_target = True
-    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", "0")
+    conv._scalar_capture_retains_rotation = False
+    assert not vp._save_competition_scalar(conv, x)
+
+
+@pytest.mark.parametrize(
+    "field,value",
+    [
+        ("training", True),
+        ("triton_infer_level", 0),
+        ("triton_infer_level", 1),
+        ("triton_infer_level", 3),
+        ("lmax", 2),
+        ("n_focus", 1),
+        ("so2_focus_dim", 64),
+        ("focus_compete", False),
+    ],
+)
+def test_cpu_tracing_preserves_selection_constraints(field, value):
+    """A CUDA export target must not bypass inference/layout eligibility."""
+    conv = SimpleNamespace(
+        training=False,
+        triton_infer_level=2,
+        lmax=3,
+        n_focus=2,
+        so2_focus_dim=32,
+        focus_compete=True,
+        _force_scalar_capture_cuda_target=True,
+    )
+    x = torch.empty(7, 16, 64, device="cpu")
+    assert vp._save_competition_scalar(conv, x)
+    assert not vp._save_competition_scalar(conv, x.double())
+    setattr(conv, field, value)
     assert not vp._save_competition_scalar(conv, x)
 
 
@@ -135,7 +165,6 @@ def test_capture_freeze_load(monkeypatch, strict_environment, tmp_path):
         ModelWrapper,
     )
 
-    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", "1")
     # The launch-table sweep is a performance concern; retain the real
     # tracing, torch.export, AOTI compilation and package loader below.
     monkeypatch.setattr(freeze_pt2, "_tune_triton_configs", lambda *args: None)
@@ -187,8 +216,10 @@ def test_capture_freeze_load(monkeypatch, strict_environment, tmp_path):
 
 @GPU
 @pytest.mark.parametrize("backend", ["pt", "pt_expt"])
-@pytest.mark.parametrize("capture", [False, True])
-def test_value_path_dispatch(monkeypatch, strict_environment, backend, capture):
+@pytest.mark.parametrize("retains_rotation", [False, True])
+def test_value_path_dispatch(
+    monkeypatch, strict_environment, backend, retains_rotation
+):
     """Exercise the production callable and prove the competition data source."""
     if backend == "pt":
         from deepmd.pt.model.descriptor.sezm_nn.so2 import (
@@ -222,6 +253,8 @@ def test_value_path_dispatch(monkeypatch, strict_environment, backend, capture):
     )
     path = vp.make_triton_value_path(conv)
     assert path is not None
+    conv._scalar_capture_retains_rotation = retains_rotation
+    capture = retains_rotation
     x = torch.randn(7, 16, 64, device="cuda") * 0.1
     radial = torch.randn(17, 4, 32, device="cuda") * 0.1
     cache = SimpleNamespace(
@@ -246,7 +279,6 @@ def alpha_spy(gate_src):
             torch.testing.assert_close(gate_src, seen["scalar"], rtol=0, atol=0)
         return original_alpha(gate_src)
 
-    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", str(int(capture)))
     with (
         patch.object(
             vp, "_rotate_mix_scalar_op", side_effect=capture_spy
@@ -305,9 +337,9 @@ def test_model_capture_on_off(monkeypatch, strict_environment):
         dtype=torch.float64,
     )
     atype = torch.zeros(1, 4, device="cuda", dtype=torch.long)
-    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", "0")
-    reference = model(coord, atype, None)
-    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", "1")
+    # A test-only ablation leaves the production acceleration policy intact.
+    with patch.object(vp, "_save_competition_scalar", return_value=False):
+        reference = model(coord, atype, None)
     with patch.object(
         vp, "_rotate_mix_scalar_op", wraps=vp._rotate_mix_scalar_op
     ) as spy:
@@ -324,17 +356,17 @@ def test_model_capture_on_off(monkeypatch, strict_environment):
 
 @pytest.mark.parametrize("backend", ["pt", "pt_expt"])
 @pytest.mark.parametrize(
-    "capture,target",
+    "level,target",
     [
-        pytest.param(True, "cuda", marks=GPU),
-        pytest.param(False, "cuda", marks=GPU),
-        (True, "cpu"),
+        pytest.param(2, "cuda", marks=GPU),
+        pytest.param(1, "cuda", marks=GPU),
+        (2, "cpu"),
     ],
 )
 def test_export_dispatch(
-    monkeypatch, strict_environment, tmp_path, backend, capture, target
+    monkeypatch, strict_environment, tmp_path, backend, level, target
 ):
-    """Both official exporters must select capture before CPU tracing ends.
+    """Official exporters respect the level, target and storage-lifetime policy.
 
     Stop at the export boundary, avoiding compiler/toolchain dependencies in
     this routing regression. Full AOTI package inference is benchmarked separately.
@@ -355,7 +387,7 @@ def test_export_dispatch(
         _trace_and_export_impl,
     )
 
-    monkeypatch.setenv("DP_SO2_SCALAR_CAPTURE_INFER", str(int(capture)))
+    monkeypatch.setenv("DP_TRITON_INFER", str(level))
     monkeypatch.setattr(freeze_pt2, "_tune_triton_configs", lambda *args: None)
     monkeypatch.setattr(autotune, "run_autotune", lambda *args: None)
     params = _model_params()
@@ -369,7 +401,7 @@ def inspect_graph(traced, *args, **kwargs):
         targets = [str(node.target) for node in traced.graph.nodes]
         count = sum("sezm_triton.so2_rotate_mix_scalar." in name for name in targets)
         seen.append(count)
-        assert (count > 0) == (capture and target == "cuda")
+        assert (count > 0) == (backend == "pt" and level == 2 and target == "cuda")
         raise ExportBoundary
 
     monkeypatch.setattr(torch.export, "export", inspect_graph)
diff --git a/source/tests/pt_expt/test_compare_so2_scalar_capture.py b/source/tests/pt_expt/test_compare_so2_scalar_capture.py
index 0e0e970e75..533d29784f 100644
--- a/source/tests/pt_expt/test_compare_so2_scalar_capture.py
+++ b/source/tests/pt_expt/test_compare_so2_scalar_capture.py
@@ -35,13 +35,33 @@ def test_matching_identity():
 
 def test_compare_checks_cross_repeat_pairs():
     """Equal matching repeats must not hide variability across repeats."""
+    scales = {"energy": 1.0, "forces": 10.0, "virial": 100.0, "stress": 1000.0}
     arrays = {
-        f"{repeat}_{frame}_{field}": np.array([1.0 + repeat])
+        f"{repeat}_{frame}_{field}": np.array([(1.0 + repeat) * scales[field]])
         for repeat in range(3)
         for frame in range(3)
-        for field in ("energy", "forces", "virial", "stress")
+        for field in scales
     }
     result = compare(arrays, arrays)
-    for field in result.values():
-        assert not field["passed"]
-        assert field["max_abs"] == 2.0
+    for field, scale in scales.items():
+        assert not result[field]["passed"]
+        assert result[field]["max_abs"] == 2.0 * scale
+
+
+@pytest.mark.parametrize("changed_frame", [0, 1, 2])
+def test_compare_uses_reference_and_each_frame(changed_frame):
+    """Distinct routes must expose a discrepancy confined to any one frame."""
+    scales = {"energy": 1.0, "forces": 10.0, "virial": 100.0, "stress": 1000.0}
+    candidate = {
+        f"{repeat}_{frame}_{field}": np.array([scale])
+        for repeat in range(3)
+        for frame in range(3)
+        for field, scale in scales.items()
+    }
+    reference = {key: value.copy() for key, value in candidate.items()}
+    for field, scale in scales.items():
+        reference[f"0_{changed_frame}_{field}"] += 2 * scale
+    result = compare(candidate, reference)
+    for field, scale in scales.items():
+        assert not result[field]["passed"]
+        assert result[field]["max_abs"] == 2 * scale

From 6264f6fc35705c72cfc76703faa3ef84a8d205e3 Mon Sep 17 00:00:00 2001
From: Macqueen Ma-hnu <16522286+macqueen09@users.noreply.github.com>
Date: Fri, 6 Nov 2026 22:42:54 +0800
Subject: [PATCH 7/9] fix(test): validate SO2 scalar capture benchmark results

---
 doc/model/dpa4.md                             |  4 ++-
 .../pt_expt/benchmark_so2_scalar_capture.py   | 13 +++++--
 .../pt_expt/compare_so2_scalar_capture.py     | 25 ++++++++++++-
 .../test_compare_so2_scalar_capture.py        | 35 +++++++++++++++++++
 4 files changed, 73 insertions(+), 4 deletions(-)

diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md
index 59283a147b..e93eadee57 100644
--- a/doc/model/dpa4.md
+++ b/doc/model/dpa4.md
@@ -581,7 +581,9 @@ retain the original path. The latter already releases the large rotation
 storage, so capture adds no measured peak-memory benefit there. This is a
 capacity optimization: the extra scalar store can increase latency. Measurements
 on RTX 5090 with the Neo checkpoint found about 12.35%/12.38% lower peak allocated
-memory for PT export at 1000/4096 atoms, with about 0.65%/1.09% higher latency.
+memory for PT export at 1000/4096 atoms, with about 0.338%/0.453% higher latency
+in the [updated measurement round](https://github.com/deepmodeling/deepmd-kit/pull/6114#issuecomment-6635037675)
+(PyTorch 2.11.0+cu128, Triton 3.6.0, TF32/AMP off, 12 warmups and 30 samples).
 These measurements do not imply the same benefit for other execution paths.
 The manual full-API benchmark is
 `source/tests/pt_expt/benchmark_so2_scalar_capture.py` (`--help` lists inputs).
diff --git a/source/tests/pt_expt/benchmark_so2_scalar_capture.py b/source/tests/pt_expt/benchmark_so2_scalar_capture.py
index d2095afe2b..6b7e10ceac 100644
--- a/source/tests/pt_expt/benchmark_so2_scalar_capture.py
+++ b/source/tests/pt_expt/benchmark_so2_scalar_capture.py
@@ -107,10 +107,19 @@ def sha(path):
 from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as scalar_path
 from deepmd.utils.errors import OutOfMemoryError as DPOutOfMemoryError
 
+from compare_so2_scalar_capture import (
+    configure_capture_route,
+)
+
 # BASE is an explicit benchmark-only ablation, not a public inference switch.
 # On a pre-feature checkout BASE naturally uses its original implementation.
-if a.route == "BASE" and hasattr(scalar_path, "_save_competition_scalar"):
-    scalar_path._save_competition_scalar = lambda conv, x: False
+try:
+    configure_capture_route(a.route, scalar_path)
+except ValueError as error:
+    write(
+        "result.json", {"status": "unsupported", "route": a.route, "reason": str(error)}
+    )
+    raise SystemExit(2) from error
 
 torch.set_num_threads(1)
 torch.set_float32_matmul_precision("highest")
diff --git a/source/tests/pt_expt/compare_so2_scalar_capture.py b/source/tests/pt_expt/compare_so2_scalar_capture.py
index 91b020c223..9b0b6e24c1 100644
--- a/source/tests/pt_expt/compare_so2_scalar_capture.py
+++ b/source/tests/pt_expt/compare_so2_scalar_capture.py
@@ -4,6 +4,7 @@
 Run with --torch DIR --base DIR --candidate DIR. An OOM reference is unavailable,
 not a pass. Output is JSON; absolute units follow the ASE energy/force/virial/stress
 arrays. Relative errors use max(abs(reference), 1e-8) as the denominator.
+The worst tolerance ratio records its frame, repeat pair and component values.
 """
 # ruff: noqa: T201 -- standalone comparison CLI
 
@@ -19,6 +20,15 @@
 import numpy as np
 
 
+def configure_capture_route(route, scalar_path):
+    """Reject missing CSC5 support; retain supported-layout fallback behavior."""
+    available = hasattr(scalar_path, "_save_competition_scalar")
+    if route == "CSC5" and not available:
+        raise ValueError("CSC5 is unavailable in the selected source checkout")
+    if route == "BASE" and available:
+        scalar_path._save_competition_scalar = lambda conv, x: False
+
+
 def compare(candidate, reference):
     """Compare all 3x3 repeat pairs for each of the three corresponding frames."""
     fields = {}
@@ -29,12 +39,24 @@ def compare(candidate, reference):
             y = reference[f"{j}_{frame}_{field}"]
             assert x.shape == y.shape
             delta = np.abs(x - y)
+            ratio = delta / (1e-5 + 1e-4 * abs(y))
+            index = np.unravel_index(np.argmax(ratio), ratio.shape)
             rows.append(
                 {
                     "max_abs": float(delta.max()),
                     "mean_abs": float(delta.mean()),
                     "max_relative": float((delta / np.maximum(abs(y), 1e-8)).max()),
-                    "max_ratio": float((delta / (1e-5 + 1e-4 * abs(y))).max()),
+                    "max_ratio": float(ratio.max()),
+                    "worst": {
+                        "frame": frame,
+                        "candidate_repeat": i,
+                        "reference_repeat": j,
+                        "index": [int(k) for k in index],
+                        "candidate": float(x[index]),
+                        "reference": float(y[index]),
+                        "abs_error": float(delta[index]),
+                        "tolerance": float(1e-5 + 1e-4 * abs(y[index])),
+                    },
                     "finite": bool(np.isfinite(x).all() and np.isfinite(y).all()),
                 }
             )
@@ -45,6 +67,7 @@ def compare(candidate, reference):
         fields[field]["passed"] = all(
             row["finite"] and row["max_ratio"] <= 1 for row in rows
         )
+        fields[field]["worst"] = max(rows, key=lambda row: row["max_ratio"])["worst"]
     return fields
 
 
diff --git a/source/tests/pt_expt/test_compare_so2_scalar_capture.py b/source/tests/pt_expt/test_compare_so2_scalar_capture.py
index 533d29784f..6cc3aaaee0 100644
--- a/source/tests/pt_expt/test_compare_so2_scalar_capture.py
+++ b/source/tests/pt_expt/test_compare_so2_scalar_capture.py
@@ -1,15 +1,45 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Benchmark comparisons must use the same checkpoint and input frames."""
 
+from types import (
+    SimpleNamespace,
+)
+
 import numpy as np
 import pytest
 
 from .compare_so2_scalar_capture import (
     compare,
+    configure_capture_route,
     validate_identity,
 )
 
 
+@pytest.mark.parametrize("route", ["BASE", "CSC5", "TORCH"])
+@pytest.mark.parametrize("available", [False, True])
+@pytest.mark.parametrize("eligible", [False, True])
+def test_capture_route_support(route, available, eligible):
+    """A pre-feature checkout cannot masquerade as a completed CSC5 run."""
+    scalar_path = SimpleNamespace()
+    if available:
+
+        def original(conv, x):
+            return eligible
+
+        scalar_path._save_competition_scalar = original
+    if route == "CSC5" and not available:
+        with pytest.raises(ValueError, match="CSC5 is unavailable"):
+            configure_capture_route(route, scalar_path)
+        return
+    configure_capture_route(route, scalar_path)
+    if available:
+        assert scalar_path._save_competition_scalar(None, None) == (
+            eligible and route != "BASE"
+        )
+        if route != "BASE":
+            assert scalar_path._save_competition_scalar is original
+
+
 @pytest.mark.parametrize("key", ["model_sha256", "input_sha256"])
 @pytest.mark.parametrize("missing", [False, True])
 def test_reject_mismatched_identity(key, missing):
@@ -65,3 +95,8 @@ def test_compare_uses_reference_and_each_frame(changed_frame):
     for field, scale in scales.items():
         assert not result[field]["passed"]
         assert result[field]["max_abs"] == 2 * scale
+        assert result[field]["worst"]["frame"] == changed_frame
+        assert result[field]["worst"]["reference_repeat"] == 0
+        assert result[field]["worst"]["index"] == [0]
+        assert result[field]["worst"]["candidate"] == scale
+        assert result[field]["worst"]["reference"] == 3 * scale

From ae294158742f029f2ef8d185e49b3ba82e969580 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 6 Nov 2026 15:04:33 +0000
Subject: [PATCH 8/9] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/tests/pt_expt/benchmark_so2_scalar_capture.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/source/tests/pt_expt/benchmark_so2_scalar_capture.py b/source/tests/pt_expt/benchmark_so2_scalar_capture.py
index 6b7e10ceac..54f822a36f 100644
--- a/source/tests/pt_expt/benchmark_so2_scalar_capture.py
+++ b/source/tests/pt_expt/benchmark_so2_scalar_capture.py
@@ -100,6 +100,9 @@ def sha(path):
 from ase.build import (
     bulk,
 )
+from compare_so2_scalar_capture import (
+    configure_capture_route,
+)
 
 from deepmd.calculator import (
     DP,
@@ -107,10 +110,6 @@ def sha(path):
 from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as scalar_path
 from deepmd.utils.errors import OutOfMemoryError as DPOutOfMemoryError
 
-from compare_so2_scalar_capture import (
-    configure_capture_route,
-)
-
 # BASE is an explicit benchmark-only ablation, not a public inference switch.
 # On a pre-feature checkout BASE naturally uses its original implementation.
 try:

From ae36377d0e7e7ac072e001a5178c9d90298039a7 Mon Sep 17 00:00:00 2001
From: Macqueen Ma-hnu <16522286+macqueen09@users.noreply.github.com>
Date: Wed, 11 Nov 2026 00:24:40 +0800
Subject: [PATCH 9/9] docs(dpa4): keep scalar capture details in PR discussion

---
 doc/model/dpa4.md | 20 --------------------
 1 file changed, 20 deletions(-)

diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md
index e93eadee57..ad70c54b6f 100644
--- a/doc/model/dpa4.md
+++ b/doc/model/dpa4.md
@@ -568,26 +568,6 @@ shipping built in. On other GPUs the kernels fall back to conservative
 configurations, and `dp --pt freeze` tunes the missing entries on the local GPU
 before exporting, a one-off sweep of a few minutes baked into the `.pt2`.
 
-Level `2` saves the focus-competition scalars on supported inference paths
-into independent storage during the existing rotate-and-mix launch. This lets
-the full rotation output die after its mixing consumer, rather than remain
-live for competition backward. It does not change model weights or precision.
-Selection is internal and currently applies only to evaluation with level `2`,
-CUDA float32, `lmax=3`, two focuses, 32 channels per focus and focus competition.
-Other layouts and training retain the original path. PT freeze selects this
-branch from the CUDA deployment target even though tracing runs on CPU;
-existing `.pt2` artifacts must be regenerated. CPU targets and `pt_expt` export
-retain the original path. The latter already releases the large rotation
-storage, so capture adds no measured peak-memory benefit there. This is a
-capacity optimization: the extra scalar store can increase latency. Measurements
-on RTX 5090 with the Neo checkpoint found about 12.35%/12.38% lower peak allocated
-memory for PT export at 1000/4096 atoms, with about 0.338%/0.453% higher latency
-in the [updated measurement round](https://github.com/deepmodeling/deepmd-kit/pull/6114#issuecomment-6635037675)
-(PyTorch 2.11.0+cu128, Triton 3.6.0, TF32/AMP off, 12 warmups and 30 samples).
-These measurements do not imply the same benefit for other execution paths.
-The manual full-API benchmark is
-`source/tests/pt_expt/benchmark_so2_scalar_capture.py` (`--help` lists inputs).
-
 `DP_CUTILE_INFER` replaces the whole SeZM edge pipeline — Wigner monomials,
 rotate-and-mix, the gated SO(2) mixing stack, the attention aggregation and the
 force / virial assembly — with kernels written in the `cuda.tile` DSL. On an