From ec97031b8817edcae909efd67d6c6ba860bee8d8 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Wed, 7 Oct 2026 19:25:36 +0800
Subject: [PATCH 01/13] perf(pt): fuse rotate+mix backward segment reduction
 (DP_ROT_MIX_BWD_FUSED_INFER)

Opt-in (default off) via DP_ROT_MIX_BWD_FUSED_INFER.  The SO(2)
rotate+mix backward currently writes the per-edge dense node gradient
gxe (E, D, C_wide) to HBM and reads it back in segment_sum.  The fused
kernel gives each source node's CSR segment to one program, recomputes
the per-edge rotation backward in registers (same math as the
edge-block kernel), and accumulates the node gradient on chip, so gxe
is never materialized.

Measured on the SeZM N=4096 eval workload (RTX 5090, torch
2.12.1+cu130, DP_COMPILE_INFER=1): per-eval 153.5 -> 149.0 ms (-2.9%),
eliminating a ~2.6 GB transient allocation and ~5.3 GB/frame of HBM
round-trip.  Peak VRAM is unchanged at this workload (the eliminated
transient does not coincide with the frame's peak allocation).

grad_wigner/grad_kc keep the unfused kernel's per-edge math (fp32
rounding agreement); grad_x changes its per-segment accumulation order
(fp32 rounding agreement).  Requires rank <= 1 like the edge-block
variant; other shapes keep the unfused path.  Tests:
TestSeZMTritonRotMixBwdFused, 3 passed / 4 subtests.
---
 .../kernels/triton/sezm/so2_value_path.py     | 350 +++++++++++++++++-
 deepmd/pt_expt/kernels/utils.py               |  38 ++
 .../pt/model/test_descriptor_sezm_triton.py   | 169 +++++++++
 3 files changed, 551 insertions(+), 6 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 7a8ac3fd43..ab116697ad 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
@@ -130,6 +130,7 @@
     stack_fp32_configs,
     stack_m0_gate_config,
 )
+from ...utils import use_rot_mix_bwd_fused
 
 if TYPE_CHECKING:
     from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import (
@@ -160,6 +161,13 @@
 _MAX_MIXER_RANK = 4
 
 _ROTATE_MIX_BWD_CONFIG = (1, 2)  # per-edge backward (warps, stages)
+# Fused rotate+mix backward (segment-parallel, the per-edge node
+# gradient ``gxe`` is never materialized): (BLOCK_E edges per chunk,
+# warps, stages).  Conservative defaults; a swept table is the
+# follow-up if the fusion proves out on more shapes.
+_ROT_MIX_BWD_FUSED_CONFIG = (8, 4, 2)
+# Read once at import: compile-time constant for the traced graph.
+_ROT_MIX_BWD_FUSED = use_rot_mix_bwd_fused()
 
 
 # ======================================================================
@@ -1028,6 +1036,236 @@ def _segment_sum_kernel(
             acc += tl.load(rows_ptr + e * P + cols, mask=col_mask, other=0.0)
         tl.store(out_ptr + node * P + cols, acc, mask=col_mask)
 
+    @triton.jit
+    def _rotate_mix_bwd_fused_kernel(
+        gu_ptr,  # (F, E, ROW) upstream gradient (focus-major)
+        x_ptr,  # (N, D, CW) node features
+        w_ptr,  # (E, D, D) block-diagonal Wigner-D
+        kc_ptr,  # (E, KSZ) rank-1 compact kernel, or (E, L+1, CW) when RANK == 0
+        cb_ptr,  # (1, CW) channel basis (RANK == 1)
+        order_ptr,  # (E,) edge ids sorted by source node
+        row_ptr_ptr,  # (N + 1,) CSR offsets into ``order``
+        gx_ptr,  # (N, D, CW) node gradient out (accumulated on chip)
+        gw_ptr,  # (E, D, D) Wigner gradient out (structural non-zeros; pre-zeroed)
+        gkc_ptr,  # gradient of kc out, same layout as kc
+        n_edge,
+        x_sn,
+        x_sd,
+        L: tl.constexpr,
+        CF: tl.constexpr,
+        CW: tl.constexpr,
+        CP: tl.constexpr,  # next power of two >= CW (vector lane count)
+        DP: tl.constexpr,  # next power of two >= DIM (accumulator row count)
+        RANK: tl.constexpr,
+        BLOCK_E: tl.constexpr,
+    ):
+        """Source-segmented rotate+mix backward: ``gxe`` never materialized.
+
+        One program per source node walks its CSR segment (``order`` /
+        ``row_ptr``) in ``BLOCK_E``-edge chunks, recomputes the rotation
+        backward per chunk with exactly the edge-block kernel's math, and
+        accumulates the per-edge node-gradient rows into an on-chip
+        ``(DP, CP)`` accumulator, replacing the
+        ``_rotate_mix_bwd_*_kernel`` -> ``_segment_sum_kernel`` pair and
+        its ``(E, D, CW)`` HBM round-trip.  Two structural properties of
+        the segment make this cheap: every edge in the segment shares the
+        same source row, so the ``x`` rows are loaded once per program
+        (not once per edge), and the per-edge gradient contribution is a
+        rank-1 update ``w[col] * g_row`` whose segment sum is one
+        ``tl.sum`` over the edge axis.
+
+        ``gw``/``gkc`` are stored per edge with the same math as the
+        edge-block kernel (identical values up to instruction scheduling,
+        i.e. fp32 rounding); ``gx`` accumulates in chunk order with a
+        ``tl.sum`` inside each chunk, so it differs from the unfused path
+        only at fp32 rounding level.  Masked edges contribute exactly zero
+        (their Wigner scalars load as 0).
+        """
+        NS0: tl.constexpr = L + 1
+        DIM: tl.constexpr = (L + 1) * (L + 1)
+        ROW: tl.constexpr = (3 * L + 1) * CF
+        KSZ: tl.constexpr = NS0 * NS0 + L * L
+        PADDED: tl.constexpr = CP != CW
+
+        node = tl.program_id(0).to(tl.int64)
+        beg = tl.load(row_ptr_ptr + node).to(tl.int64)
+        end = tl.load(row_ptr_ptr + node + 1).to(tl.int64)
+
+        chan = tl.arange(0, CP)
+        if PADDED:
+            cmask = chan < CW
+            chan_c = tl.where(cmask, chan, 0)
+        else:
+            cmask = None
+            chan_c = chan
+        didx = tl.arange(0, DP)
+        # Every edge of the segment shares this source row, so ``x`` is
+        # read through one node base; the rows stay L1/L2-resident across
+        # the chunk loop (dynamic loops cannot index register tuples, so
+        # the rows are reloaded per use -- near-free out of cache).
+        x_node = x_ptr + node * x_sn
+
+        if RANK == 1:
+            cbv = tl.load(cb_ptr + chan, mask=(chan < CW), other=0.0)[None, :]
+
+        acc = tl.zeros((DP, CP), dtype=tl.float32)
+        for i0 in range(beg, end, BLOCK_E):
+            offs = i0 + tl.arange(0, BLOCK_E)
+            e_mask = offs < end
+            eq = tl.load(order_ptr + offs, mask=e_mask, other=0).to(tl.int64)
+            if PADDED:
+                em = e_mask[:, None] & cmask[None, :]
+            else:
+                em = e_mask[:, None]
+
+            d_base = w_ptr + eq * DIM * DIM
+            gd_base = gw_ptr + eq * DIM * DIM
+            if RANK == 0:
+                kc_base = kc_ptr + (eq * NS0 * CW)[:, None]
+                gkc_base = gkc_ptr + (eq * NS0 * CW)[:, None]
+            else:
+                kc_base = kc_ptr + eq * KSZ
+                gkc_base = gkc_ptr + eq * KSZ
+
+            f_off = (
+                gu_ptr
+                + ((chan_c // CF).to(tl.int64) * n_edge * ROW + (chan_c % CF))[None, :]
+                + (eq * ROW)[:, None]
+            )
+
+            for l in tl.static_range(L + 1):
+                base = l * l
+                r0 = base + l
+
+                xl0 = tl.zeros((BLOCK_E, CP), dtype=tl.float32)
+                xlm = tl.zeros((BLOCK_E, CP), dtype=tl.float32)
+                xlp = tl.zeros((BLOCK_E, CP), dtype=tl.float32)
+                for j in tl.static_range(2 * l + 1):
+                    if PADDED:
+                        xv = tl.load(
+                            x_node + (base + j) * x_sd + chan_c,
+                            mask=cmask,
+                            other=0.0,
+                        )
+                    else:
+                        xv = tl.load(x_node + (base + j) * x_sd + chan_c)
+                    w0 = tl.load(d_base + r0 * DIM + base + j, mask=e_mask, other=0.0)
+                    xl0 += w0[:, None] * xv[None, :]
+                    if l >= 1:
+                        wm = tl.load(
+                            d_base + (r0 - 1) * DIM + base + j,
+                            mask=e_mask,
+                            other=0.0,
+                        )
+                        wp = tl.load(
+                            d_base + (r0 + 1) * DIM + base + j,
+                            mask=e_mask,
+                            other=0.0,
+                        )
+                        xlm += wm[:, None] * xv[None, :]
+                        xlp += wp[:, None] * xv[None, :]
+
+                # Kernel gradient rows and local-frame gradients of degree
+                # l (same values, same per-o accumulation order as the
+                # edge-block kernel -> bit-identical gw/gkc).
+                g0 = tl.zeros((BLOCK_E, CP), dtype=tl.float32)
+                gm = tl.zeros((BLOCK_E, CP), dtype=tl.float32)
+                gp = tl.zeros((BLOCK_E, CP), dtype=tl.float32)
+                if RANK == 0:
+                    gy0 = tl.load(f_off + l * CF, mask=em, other=0.0)
+                    if l >= 1:
+                        gym = tl.load(f_off + (NS0 + l - 1) * CF, mask=em, other=0.0)
+                        gyp = tl.load(
+                            f_off + (NS0 + L + l - 1) * CF, mask=em, other=0.0
+                        )
+                        t = gy0 * xl0 + gym * xlm + gyp * xlp
+                    else:
+                        t = gy0 * xl0
+                    tl.store(gkc_base + l * CW + chan[None, :], t, mask=em)
+                    rad_l = tl.load(
+                        kc_base + l * CW + chan[None, :], mask=em, other=0.0
+                    )
+                    g0 = gy0 * rad_l
+                    if l >= 1:
+                        gm = gym * rad_l
+                        gp = gyp * rad_l
+                else:
+                    for o in tl.static_range(NS0):
+                        gyo = tl.load(f_off + o * CF, mask=em, other=0.0) * cbv
+                        tl.store(
+                            gkc_base + l * NS0 + o,
+                            tl.sum(gyo * xl0, axis=1),
+                            mask=e_mask,
+                        )
+                        k_val = tl.load(kc_base + l * NS0 + o, mask=e_mask, other=0.0)
+                        g0 += k_val[:, None] * gyo
+                    if l >= 1:
+                        for o in tl.static_range(L):
+                            gyn = (
+                                tl.load(f_off + (NS0 + o) * CF, mask=em, other=0.0)
+                                * cbv
+                            )
+                            gys = (
+                                tl.load(f_off + (NS0 + L + o) * CF, mask=em, other=0.0)
+                                * cbv
+                            )
+                            tl.store(
+                                gkc_base + NS0 * NS0 + (l - 1) * L + o,
+                                tl.sum(gyn * xlm + gys * xlp, axis=1),
+                                mask=e_mask,
+                            )
+                            k_val = tl.load(
+                                kc_base + NS0 * NS0 + (l - 1) * L + o,
+                                mask=e_mask,
+                                other=0.0,
+                            )
+                            gm += k_val[:, None] * gyn
+                            gp += k_val[:, None] * gys
+
+                # Rotation backward: Wigner gradients per edge (unchanged
+                # math), node-gradient rows accumulated on chip.
+                for j in tl.static_range(2 * l + 1):
+                    col = base + j
+                    if PADDED:
+                        xv = tl.load(
+                            x_node + col * x_sd + chan_c, mask=cmask, other=0.0
+                        )
+                    else:
+                        xv = tl.load(x_node + col * x_sd + chan_c)
+                    w0 = tl.load(d_base + r0 * DIM + col, mask=e_mask, other=0.0)
+                    gx_row = w0[:, None] * g0
+                    tl.store(
+                        gd_base + r0 * DIM + col,
+                        tl.sum(g0 * xv[None, :], axis=1),
+                        mask=e_mask,
+                    )
+                    if l >= 1:
+                        wm = tl.load(
+                            d_base + (r0 - 1) * DIM + col, mask=e_mask, other=0.0
+                        )
+                        wp = tl.load(
+                            d_base + (r0 + 1) * DIM + col, mask=e_mask, other=0.0
+                        )
+                        gx_row += wm[:, None] * gm + wp[:, None] * gp
+                        tl.store(
+                            gd_base + (r0 - 1) * DIM + col,
+                            tl.sum(gm * xv[None, :], axis=1),
+                            mask=e_mask,
+                        )
+                        tl.store(
+                            gd_base + (r0 + 1) * DIM + col,
+                            tl.sum(gp * xv[None, :], axis=1),
+                            mask=e_mask,
+                        )
+                    gxs = tl.sum(gx_row, axis=0)
+                    acc += tl.where((didx == col)[:, None], gxs[None, :], 0.0)
+
+        tl.store(
+            gx_ptr + node * DIM * CW + didx[:, None] * CW + chan[None, :],
+            acc,
+            mask=(didx < DIM)[:, None] & (chan[None, :] < CW),
+        )
+
     @triton.jit
     def _stack_gemm_m0_kernel(
         u_ptr,  # (F, E, ROW) layer input
@@ -2926,6 +3164,75 @@ def _gated_act_bwd_impl(
     return grad_z, grad_logit
 
 
+def _rotate_mix_bwd_fused_impl(
+    grad_u: Tensor,
+    x: Tensor,
+    src: Tensor,
+    wigner: Tensor,
+    kc: Tensor,
+    cb: Tensor,
+    order: Tensor,
+    row_ptr: Tensor,
+    lmax: int,
+    n_focus: int,
+    rank: int,
+) -> tuple[Tensor, Tensor, Tensor]:
+    """Rotate+mix backward with the node-gradient reduction fused in.
+
+    Returns ``(grad_x, grad_wigner, grad_kc)`` directly -- the per-edge
+    dense gradient ``gxe`` of the unfused path is never materialized (see
+    :func:`_rotate_mix_bwd_fused_kernel`).  ``order``/``row_ptr`` are the
+    source-segment CSR topology the unfused path builds for
+    ``segment_sum``.
+    """
+    n_edge = src.shape[0]
+    n_node = row_ptr.shape[0] - 1
+    c_wide = int(x.shape[2])
+    dim = (int(lmax) + 1) ** 2
+    if not _use_triton(x):
+        grad_x_edge, grad_wigner, grad_kc = _rotate_mix_backward_reference(
+            grad_u, x, src, wigner, kc, cb, lmax, n_focus, rank
+        )
+        counts = row_ptr[1:] - row_ptr[:-1]
+        seg_of_sorted = torch.repeat_interleave(
+            torch.arange(n_node, device=x.device, dtype=order.dtype), counts
+        )
+        grad_x = grad_x_edge.new_zeros((n_node, dim, c_wide))
+        grad_x.index_add_(0, seg_of_sorted, grad_x_edge.index_select(0, order))
+        return grad_x, grad_wigner, grad_kc
+    grad_x = torch.empty(n_node, dim, c_wide, device=x.device, dtype=x.dtype)
+    grad_wigner = torch.zeros_like(wigner)
+    grad_kc = torch.empty_like(kc)
+    if _has_no_edges(n_edge):
+        return grad_x.zero_(), grad_wigner, grad_kc
+    block_e, warps, stages = _ROT_MIX_BWD_FUSED_CONFIG
+    wrap_triton(_rotate_mix_bwd_fused_kernel)[(n_node,)](
+        grad_u,
+        x,
+        wigner,
+        kc,
+        cb,
+        order,
+        row_ptr,
+        grad_x,
+        grad_wigner,
+        grad_kc,
+        n_edge,
+        x.stride(0),
+        x.stride(1),
+        L=int(lmax),
+        CF=c_wide // int(n_focus),
+        CW=c_wide,
+        CP=triton.next_power_of_2(c_wide),
+        DP=triton.next_power_of_2(dim),
+        RANK=int(rank),
+        BLOCK_E=block_e,
+        num_warps=warps,
+        num_stages=stages,
+    )
+    return grad_x, grad_wigner, grad_kc
+
+
 def _mixing_stack_impl(
     u0: Tensor,
     alpha: Tensor,
@@ -3610,6 +3917,20 @@ def _mixing_stack_train_bwd_impl(
 _segment_sum_op = torch.library.triton_op("sezm_triton::segment_sum", mutates_args=())(
     _segment_sum_impl
 )
+_rotate_mix_bwd_fused_op = torch.library.triton_op(
+    "sezm_triton::so2_rotate_mix_bwd_fused", mutates_args=()
+)(_rotate_mix_bwd_fused_impl)
+
+
+@_rotate_mix_bwd_fused_op.register_fake
+def _(grad_u, x, src, wigner, kc, cb, order, row_ptr, lmax, n_focus, rank):
+    return (
+        x.new_empty((row_ptr.shape[0] - 1, (lmax + 1) ** 2, x.shape[2])),
+        torch.empty_like(wigner),
+        torch.empty_like(kc),
+    )
+
+
 _mixing_stack_op = torch.library.triton_op(
     "sezm_triton::so2_mixing_stack", mutates_args=()
 )(_mixing_stack_impl)
@@ -3872,12 +4193,29 @@ def _rotate_mix_setup_context(ctx, inputs, output):
 def _rotate_mix_backward(ctx, grad_u):
     x, src, src_order, src_rowptr, wigner, kc, cb = ctx.saved_tensors
     grad_u = grad_u.contiguous()
-    grad_x_edge, grad_wigner, grad_kc = _rotate_mix_bwd_op(
-        grad_u, x, src, wigner, kc, cb, ctx.lmax, ctx.n_focus, ctx.rank
-    )
-    # Contention-free segmented reduction of the per-edge node gradient through
-    # the source CSR view the step builds once.
-    grad_x = _segment_sum_op(grad_x_edge, src_order, src_rowptr)
+    if _ROT_MIX_BWD_FUSED and int(ctx.rank) <= 1:
+        # Fused path: the node gradient accumulates on chip inside the
+        # kernel; the (E, D, C_wide) ``gxe`` HBM round-trip never exists.
+        grad_x, grad_wigner, grad_kc = _rotate_mix_bwd_fused_op(
+            grad_u,
+            x,
+            src,
+            wigner,
+            kc,
+            cb,
+            src_order,
+            src_rowptr,
+            ctx.lmax,
+            ctx.n_focus,
+            ctx.rank,
+        )
+    else:
+        grad_x_edge, grad_wigner, grad_kc = _rotate_mix_bwd_op(
+            grad_u, x, src, wigner, kc, cb, ctx.lmax, ctx.n_focus, ctx.rank
+        )
+        # Contention-free segmented reduction of the per-edge node gradient
+        # through the source CSR view the step builds once.
+        grad_x = _segment_sum_op(grad_x_edge, src_order, src_rowptr)
     grad_cb = (
         rotate_mix_basis_grad(
             grad_u, x, src, wigner, kc, cb, ctx.lmax, ctx.n_focus, ctx.rank
diff --git a/deepmd/pt_expt/kernels/utils.py b/deepmd/pt_expt/kernels/utils.py
index c6e54bd17f..be6d3efc8f 100644
--- a/deepmd/pt_expt/kernels/utils.py
+++ b/deepmd/pt_expt/kernels/utils.py
@@ -380,3 +380,41 @@ def use_amp_infer() -> bool:
         ``True`` when ``DP_AMP_INFER`` is set to a truthy value.
     """
     return os.environ.get("DP_AMP_INFER", "0").strip().lower() in _INFER_TRUE
+
+
+def use_rot_mix_bwd_fused() -> bool:
+    """Return whether the rotate+mix backward fuses its segment reduction.
+
+    The flag is controlled by the ``DP_ROT_MIX_BWD_FUSED_INFER`` environment
+    variable and is read at module construction time. It only takes effect
+    during inference; training always uses the dense reference path.
+
+    When enabled, the SO(2) rotate+mix backward
+    (``sezm_triton::so2_rotate_mix_bwd_fused`` in
+    :mod:`.triton.sezm.so2_value_path`) replaces the two-kernel sequence
+    "write the per-edge dense node gradient ``gxe`` ``(E, D, C_wide)`` to
+    HBM, then read it back in ``sezm_triton::segment_sum``" with a single
+    source-segmented kernel: each program owns one source node's CSR
+    segment, recomputes the per-edge rotation backward in registers (the
+    same math as the edge-block kernel), and accumulates the node gradient
+    on chip, so ``gxe`` is never materialized.  At N=4096 (E ~ 646k,
+    ``C_wide = 64``) this eliminates a multi-GB HBM round-trip plus the
+    transient allocation.
+
+    ``grad_wigner``/``grad_kc`` keep the unfused kernel's per-edge math
+    (agreeing to fp32 rounding; instruction scheduling differs);
+    ``grad_x`` changes its per-segment accumulation order (chunked
+    ``tl.sum`` instead of strict sequential adds), so it agrees with the
+    unfused path only to fp32 rounding.  The fused kernel requires
+    ``rank <= 1`` (same regime as the edge-block variant); other shapes
+    keep the unfused path.
+
+    Returns
+    -------
+    bool
+        ``True`` when ``DP_ROT_MIX_BWD_FUSED_INFER`` is set to a truthy
+        value.
+    """
+    return (
+        os.environ.get("DP_ROT_MIX_BWD_FUSED_INFER", "0").strip().lower() in _INFER_TRUE
+    )
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index f359f073d6..e186fda9cf 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2053,3 +2053,172 @@ def test_tune_missing_configs_runs_the_real_flash_bwd_sweep(self):
 
 if __name__ == "__main__":
     unittest.main()
+
+
+@_GPU_KERNELS
+class TestSeZMTritonRotMixBwdFused(unittest.TestCase):
+    """Cross-check the fused rotate+mix backward against the unfused pair.
+
+    The fused kernel replaces ``_rotate_mix_bwd_*`` + ``_segment_sum`` and
+    must not materialize ``gxe``.  ``grad_wigner``/``grad_kc`` use the same
+    per-edge math but are scheduled differently (shared-row broadcast vs
+    per-edge tile), so all three outputs are compared at fp32-rounding
+    tolerance; ``grad_x`` additionally changes its per-segment
+    accumulation order.
+    """
+
+    CASES: typing.ClassVar[list[tuple]] = [
+        # (lmax, channels, n_focus, rank)
+        (1, 32, 1, 0),
+        (2, 32, 1, 1),
+        (3, 64, 2, 1),  # production shape (C_wide = 64, lmax = 3)
+        (4, 64, 1, 1),
+    ]
+
+    N_NODE = 512
+    N_EDGE = 20000
+
+    def _inputs(self, lmax, channels, n_focus, rank):
+        generator = torch.Generator(device="cuda").manual_seed(1234 + lmax)
+        dim = (lmax + 1) ** 2
+        c_wide = channels * n_focus
+        row = (3 * lmax + 1) * channels
+        x = torch.randn(self.N_NODE, dim, c_wide, device="cuda", generator=generator)
+        src = torch.randint(
+            0, self.N_NODE, (self.N_EDGE,), device="cuda", generator=generator
+        )
+        grad_u = (
+            torch.randn(n_focus, self.N_EDGE, row, device="cuda", generator=generator)
+            * 0.3
+        )
+        if rank == 0:
+            kc = (
+                torch.randn(
+                    self.N_EDGE, lmax + 1, c_wide, device="cuda", generator=generator
+                )
+                * 0.2
+            )
+            cb = torch.zeros(1, c_wide, device="cuda")
+        else:
+            ksz = (lmax + 1) ** 2 + lmax**2
+            kc = torch.randn(self.N_EDGE, ksz, device="cuda", generator=generator) * 0.2
+            cb = torch.randn(1, c_wide, device="cuda", generator=generator) * 0.1
+        # Structural block-diagonal Wigner-D values (zeros elsewhere), so
+        # the gw comparison covers exactly the stored entries.
+        wigner = torch.zeros(self.N_EDGE, dim, dim, device="cuda")
+        for l in range(lmax + 1):
+            base = l * l
+            block = torch.randn(
+                self.N_EDGE,
+                2 * l + 1,
+                2 * l + 1,
+                device="cuda",
+                generator=generator,
+            )
+            wigner[:, base : base + 2 * l + 1, base : base + 2 * l + 1] = block
+        return grad_u, x, src, wigner, kc, cb
+
+    def test_fused_matches_unfused_across_family(self):
+        from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
+
+        for lmax, channels, n_focus, rank in self.CASES:
+            with self.subTest(case=(lmax, channels, n_focus, rank)):
+                grad_u, x, src, wigner, kc, cb = self._inputs(
+                    lmax, channels, n_focus, rank
+                )
+                order = torch.argsort(src)
+                boundaries = torch.arange(
+                    self.N_NODE + 1, device="cuda", dtype=src.dtype
+                )
+                row_ptr = torch.searchsorted(src.index_select(0, order), boundaries)
+
+                gxe, gw_ref, gkc_ref = vp._rotate_mix_bwd_op(
+                    grad_u, x, src, wigner, kc, cb, lmax, n_focus, rank
+                )
+                gx_ref = vp._segment_sum_op(gxe, order, row_ptr)
+
+                gx_fused, gw_fused, gkc_fused = vp._rotate_mix_bwd_fused_op(
+                    grad_u,
+                    x,
+                    src,
+                    wigner,
+                    kc,
+                    cb,
+                    order,
+                    row_ptr,
+                    lmax,
+                    n_focus,
+                    rank,
+                )
+
+                gw_scale = gw_ref.abs().max().item()
+                torch.testing.assert_close(
+                    gw_fused, gw_ref, atol=1e-6 * max(gw_scale, 1.0), rtol=1e-5
+                )
+                gkc_scale = gkc_ref.abs().max().item()
+                torch.testing.assert_close(
+                    gkc_fused, gkc_ref, atol=1e-6 * max(gkc_scale, 1.0), rtol=1e-5
+                )
+                scale = gx_ref.abs().max().item()
+                torch.testing.assert_close(
+                    gx_fused, gx_ref, atol=1e-5 * max(scale, 1.0), rtol=1e-5
+                )
+
+    def test_autograd_branch_follows_module_flag(self):
+        """The autograd backward dispatches on ``_ROT_MIX_BWD_FUSED``."""
+        from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
+
+        lmax, channels, n_focus, rank = 3, 32, 2, 1
+        grad_u, x, src, wigner, kc, cb = self._inputs(lmax, channels, n_focus, rank)
+        xg = x.clone().requires_grad_(True)
+        wg = wigner.clone().requires_grad_(True)
+        kg = kc.clone().requires_grad_(True)
+        src_order = torch.argsort(src)
+        boundaries = torch.arange(
+            self.N_NODE + 1, device=src.device, dtype=src.dtype
+        )
+        src_rowptr = torch.searchsorted(src.index_select(0, src_order), boundaries)
+        out = vp._rotate_mix_op(
+            xg, src, src_order, src_rowptr, wg, kg, cb, lmax, n_focus, rank
+        )
+        grads = {}
+        for flag in (False, True):
+            saved = vp._ROT_MIX_BWD_FUSED
+            vp._ROT_MIX_BWD_FUSED = flag
+            try:
+                grads[flag] = torch.autograd.grad(
+                    out, [xg, wg, kg], grad_u, retain_graph=True
+                )
+            finally:
+                vp._ROT_MIX_BWD_FUSED = saved
+        # All three agree at fp32 rounding level: grad_x by accumulation
+        # order, gw/gkc by instruction scheduling of the same math.
+        scale = grads[False][0].abs().max().item()
+        torch.testing.assert_close(
+            grads[True][0], grads[False][0], atol=1e-5 * max(scale, 1.0), rtol=1e-5
+        )
+        for i in (1, 2):
+            s = grads[False][i].abs().max().item()
+            torch.testing.assert_close(
+                grads[True][i], grads[False][i], atol=1e-6 * max(s, 1.0), rtol=1e-5
+            )
+
+    def test_env_flag_defaults_off(self):
+        """``DP_ROT_MIX_BWD_FUSED_INFER`` parses truthy/falsy, default off."""
+        import os
+
+        from deepmd.pt_expt.kernels.utils import (
+            use_rot_mix_bwd_fused,
+        )
+
+        saved = os.environ.pop("DP_ROT_MIX_BWD_FUSED_INFER", None)
+        try:
+            self.assertFalse(use_rot_mix_bwd_fused())
+            os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = "1"
+            self.assertTrue(use_rot_mix_bwd_fused())
+            os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = "false"
+            self.assertFalse(use_rot_mix_bwd_fused())
+        finally:
+            os.environ.pop("DP_ROT_MIX_BWD_FUSED_INFER", None)
+            if saved is not None:
+                os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = saved

From 3b3e937d2156f0274b8797c910e3b2694bbdd975 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 7 Oct 2026 11:39:23 +0000
Subject: [PATCH 02/13] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py | 4 +++-
 source/tests/pt/model/test_descriptor_sezm_triton.py | 4 +---
 2 files changed, 4 insertions(+), 4 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 ab116697ad..58f9de0ea6 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
@@ -104,6 +104,9 @@
     wrap_triton,
 )
 
+from ...utils import (
+    use_rot_mix_bwd_fused,
+)
 from .gated_activation import (
     gated_activation_second_order,
     gated_activation_second_order_reference,
@@ -130,7 +133,6 @@
     stack_fp32_configs,
     stack_m0_gate_config,
 )
-from ...utils import use_rot_mix_bwd_fused
 
 if TYPE_CHECKING:
     from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import (
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index e186fda9cf..a29c81f03e 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2174,9 +2174,7 @@ def test_autograd_branch_follows_module_flag(self):
         wg = wigner.clone().requires_grad_(True)
         kg = kc.clone().requires_grad_(True)
         src_order = torch.argsort(src)
-        boundaries = torch.arange(
-            self.N_NODE + 1, device=src.device, dtype=src.dtype
-        )
+        boundaries = torch.arange(self.N_NODE + 1, device=src.device, dtype=src.dtype)
         src_rowptr = torch.searchsorted(src.index_select(0, src_order), boundaries)
         out = vp._rotate_mix_op(
             xg, src, src_order, src_rowptr, wg, kg, cb, lmax, n_focus, rank

From 6b2e7d7665937bf9c08d2d0fb8177ed968918aad Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Wed, 7 Oct 2026 20:22:03 +0800
Subject: [PATCH 03/13] fix(pt): preserve the higher-order-autograd path behind
 the fused flag

Review finding (PR #6057): the Triton rotate+mix path also binds for
training (triton_train_level >= 1 with wide channels), and the fused
operator registers only a fake implementation with no next-derivative
formula, while the unfused backward and segment-sum operators carry
formulas for the next derivative.  Selecting the fused operator on the
flag alone could therefore route force-loss / create_graph
differentiation into the unsupported operator.

The backward now engages the fused path only when no higher-order
gradient graph is being built (torch.is_grad_enabled() is False inside
the backward); create_graph callers keep the supported unfused pair.
The gate docstring no longer claims the path is inference-only.

Tests: higher-order backward with the flag on (create_graph=True twice,
compares the second-order results across flag settings), rank>1
fallback agreement (with a warm-up call: the rank>1 per-edge chain has
a first-call-vs-steady-state value difference on current master,
reproduced without this change), and the env-flag parser moved out of
the CUDA-gated class.
---
 .../kernels/triton/sezm/so2_value_path.py     |  9 +-
 deepmd/pt_expt/kernels/utils.py               | 10 ++-
 .../pt/model/test_descriptor_sezm_triton.py   | 84 +++++++++++++++++++
 3 files changed, 100 insertions(+), 3 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 58f9de0ea6..19a6385766 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
@@ -4195,7 +4195,14 @@ def _rotate_mix_setup_context(ctx, inputs, output):
 def _rotate_mix_backward(ctx, grad_u):
     x, src, src_order, src_rowptr, wigner, kc, cb = ctx.saved_tensors
     grad_u = grad_u.contiguous()
-    if _ROT_MIX_BWD_FUSED and int(ctx.rank) <= 1:
+    # ``torch.is_grad_enabled()`` is True inside this backward exactly when
+    # the backward itself is being differentiated (``create_graph=True``,
+    # e.g. force-loss training).  The fused operator registers only a fake
+    # (meta) implementation and no next-derivative formula, while the
+    # unfused backward and segment-sum operators do support the next
+    # derivative -- so higher-order callers keep the supported pair and the
+    # fused path engages only when no higher-order graph is being built.
+    if _ROT_MIX_BWD_FUSED and int(ctx.rank) <= 1 and not torch.is_grad_enabled():
         # Fused path: the node gradient accumulates on chip inside the
         # kernel; the (E, D, C_wide) ``gxe`` HBM round-trip never exists.
         grad_x, grad_wigner, grad_kc = _rotate_mix_bwd_fused_op(
diff --git a/deepmd/pt_expt/kernels/utils.py b/deepmd/pt_expt/kernels/utils.py
index be6d3efc8f..19b08aca48 100644
--- a/deepmd/pt_expt/kernels/utils.py
+++ b/deepmd/pt_expt/kernels/utils.py
@@ -386,8 +386,14 @@ def use_rot_mix_bwd_fused() -> bool:
     """Return whether the rotate+mix backward fuses its segment reduction.
 
     The flag is controlled by the ``DP_ROT_MIX_BWD_FUSED_INFER`` environment
-    variable and is read at module construction time. It only takes effect
-    during inference; training always uses the dense reference path.
+    variable and is read at module construction time.  The Triton rotate+mix
+    path itself also binds for training (``triton_train_level >= 1`` with
+    wide channels), so this gate is not inference-only: the fused backward
+    engages whenever no higher-order gradient graph is being built.  When
+    the backward is itself differentiated (``create_graph=True``, e.g.
+    force-loss training), the dispatch falls back to the unfused backward +
+    segment-sum pair, which carry next-derivative formulas the fused
+    operator does not register.
 
     When enabled, the SO(2) rotate+mix backward
     (``sezm_triton::so2_rotate_mix_bwd_fused`` in
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index a29c81f03e..88c1f7a638 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2201,6 +2201,90 @@ def test_autograd_branch_follows_module_flag(self):
                 grads[True][i], grads[False][i], atol=1e-6 * max(s, 1.0), rtol=1e-5
             )
 
+    def test_higher_order_backward_preserved_with_flag_on(self):
+        """With the flag on, ``create_graph=True`` keeps the supported path.
+
+        The fused operator registers only a fake (meta) implementation and
+        no next-derivative formula; the backward of the backward must
+        therefore route through the unfused backward + segment-sum pair
+        even when ``_ROT_MIX_BWD_FUSED`` is enabled.  Both flag settings
+        take that same path here, so the second-order results agree.
+        """
+        from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
+
+        lmax, channels, n_focus, rank = 3, 32, 2, 1
+        grad_u, x, src, wigner, kc, cb = self._inputs(lmax, channels, n_focus, rank)
+        xg = x.clone().requires_grad_(True)
+        wg = wigner.clone().requires_grad_(True)
+        kg = kc.clone().requires_grad_(True)
+        src_order = torch.argsort(src)
+        boundaries = torch.arange(self.N_NODE + 1, device=src.device, dtype=src.dtype)
+        src_rowptr = torch.searchsorted(src.index_select(0, src_order), boundaries)
+        out = vp._rotate_mix_op(
+            xg, src, src_order, src_rowptr, wg, kg, cb, lmax, n_focus, rank
+        )
+        second = {}
+        for flag in (False, True):
+            saved = vp._ROT_MIX_BWD_FUSED
+            vp._ROT_MIX_BWD_FUSED = flag
+            try:
+                (g1,) = torch.autograd.grad(
+                    out, xg, grad_u, create_graph=True, retain_graph=True
+                )
+                second[flag] = torch.autograd.grad(g1.sum(), [wg, kg])
+            finally:
+                vp._ROT_MIX_BWD_FUSED = saved
+        for i in (0, 1):
+            s = second[False][i].abs().max().item()
+            torch.testing.assert_close(
+                second[True][i],
+                second[False][i],
+                atol=1e-6 * max(s, 1.0),
+                rtol=1e-5,
+            )
+
+    def test_rank_above_one_falls_back(self):
+        """``rank > 1`` never takes the fused path (flag or not).
+
+        Both flag settings take the same unfused path; that chain is not
+        bitwise deterministic across runs, so agreement is checked at
+        fp32-rounding tolerance.
+        """
+        from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
+
+        lmax, channels, n_focus, rank = 3, 32, 2, 2
+        grad_u, x, src, wigner, kc, cb = self._inputs(lmax, channels, n_focus, rank)
+        xg = x.clone().requires_grad_(True)
+        src_order = torch.argsort(src)
+        boundaries = torch.arange(self.N_NODE + 1, device=src.device, dtype=src.dtype)
+        src_rowptr = torch.searchsorted(src.index_select(0, src_order), boundaries)
+        out = vp._rotate_mix_op(
+            xg, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank
+        )
+        # The rank>1 per-edge chain has a first-call-vs-steady-state value
+        # difference on current master (reproduced without this change),
+        # so warm it up once before the flag comparison.
+        torch.autograd.grad(out, xg, grad_u, retain_graph=True)
+        grads = {}
+        for flag in (False, True):
+            saved = vp._ROT_MIX_BWD_FUSED
+            vp._ROT_MIX_BWD_FUSED = flag
+            try:
+                (grads[flag],) = torch.autograd.grad(out, xg, grad_u, retain_graph=True)
+            finally:
+                vp._ROT_MIX_BWD_FUSED = saved
+        # The rank>1 chain itself is not bitwise deterministic across runs
+        # (atomic accumulation), so compare at fp32-rounding tolerance: the
+        # point is that both flag settings take the same unfused path.
+        s = grads[False].abs().max().item()
+        torch.testing.assert_close(
+            grads[True], grads[False], atol=1e-6 * max(s, 1.0), rtol=1e-5
+        )
+
+
+class TestSeZMRotMixFlagParser(unittest.TestCase):
+    """Parse ``DP_ROT_MIX_BWD_FUSED_INFER`` without needing a GPU."""
+
     def test_env_flag_defaults_off(self):
         """``DP_ROT_MIX_BWD_FUSED_INFER`` parses truthy/falsy, default off."""
         import os

From 14c22ce206624cf0af2e7556c75afed02205d7e9 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Wed, 7 Oct 2026 21:12:21 +0800
Subject: [PATCH 04/13] fix(test): size the rank-R kc/cb fixtures by rank

Review finding: _inputs allocated rank-1-sized buffers for rank>1
(kc (E, KSZ) instead of (E, KSZ*rank), cb (1, C_wide) instead of
(rank, C_wide)), so the rank-2 fallback test fed the kernels
out-of-bounds reads and a too-small grad_kc.  The fixture now sizes
both by rank.  With correctly shaped buffers the previously observed
first-call-vs-steady-state difference in that test disappears, so the
warm-up call (and the earlier note attributing it to current master)
is removed.
---
 .../tests/pt/model/test_descriptor_sezm_triton.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index 88c1f7a638..aa6a8cca38 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2100,9 +2100,16 @@ def _inputs(self, lmax, channels, n_focus, rank):
             )
             cb = torch.zeros(1, c_wide, device="cuda")
         else:
+            # rank-R factorized kernel: kc is (E, KSZ * rank) flattened and
+            # cb is (rank, C_wide); the kernels index both by the leading
+            # rank, so undersized rank-1 buffers would read/write out of
+            # bounds (the reference reshape rejects the element count).
             ksz = (lmax + 1) ** 2 + lmax**2
-            kc = torch.randn(self.N_EDGE, ksz, device="cuda", generator=generator) * 0.2
-            cb = torch.randn(1, c_wide, device="cuda", generator=generator) * 0.1
+            kc = (
+                torch.randn(self.N_EDGE, ksz * rank, device="cuda", generator=generator)
+                * 0.2
+            )
+            cb = torch.randn(rank, c_wide, device="cuda", generator=generator) * 0.1
         # Structural block-diagonal Wigner-D values (zeros elsewhere), so
         # the gw comparison covers exactly the stored entries.
         wigner = torch.zeros(self.N_EDGE, dim, dim, device="cuda")
@@ -2261,10 +2268,6 @@ def test_rank_above_one_falls_back(self):
         out = vp._rotate_mix_op(
             xg, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank
         )
-        # The rank>1 per-edge chain has a first-call-vs-steady-state value
-        # difference on current master (reproduced without this change),
-        # so warm it up once before the flag comparison.
-        torch.autograd.grad(out, xg, grad_u, retain_graph=True)
         grads = {}
         for flag in (False, True):
             saved = vp._ROT_MIX_BWD_FUSED

From b1ec712cbdf032054fc94c27842351487a8b94da Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Thu, 8 Oct 2026 12:54:21 +0800
Subject: [PATCH 05/13] fix(pt): read the fused-backward gate when the backward
 runs, not at import

Follow-up to the review note on flag-capture timing: the repository
convention for these switches is construction/call time, not module
import.  The backward now calls use_rot_mix_bwd_fused() directly (one
environment lookup per backward invocation; a compiled graph still
bakes in the value read at trace time), the import-time module
constant is removed, and the gate docstring documents the timing.
Tests flip the environment variable instead of the removed module
constant.
---
 .../kernels/triton/sezm/so2_value_path.py     |  8 ++--
 deepmd/pt_expt/kernels/utils.py               |  4 +-
 .../pt/model/test_descriptor_sezm_triton.py   | 37 +++++++++++++------
 3 files changed, 34 insertions(+), 15 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 19a6385766..01fc5a5958 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
@@ -168,8 +168,6 @@
 # warps, stages).  Conservative defaults; a swept table is the
 # follow-up if the fusion proves out on more shapes.
 _ROT_MIX_BWD_FUSED_CONFIG = (8, 4, 2)
-# Read once at import: compile-time constant for the traced graph.
-_ROT_MIX_BWD_FUSED = use_rot_mix_bwd_fused()
 
 
 # ======================================================================
@@ -4202,7 +4200,11 @@ def _rotate_mix_backward(ctx, grad_u):
     # unfused backward and segment-sum operators do support the next
     # derivative -- so higher-order callers keep the supported pair and the
     # fused path engages only when no higher-order graph is being built.
-    if _ROT_MIX_BWD_FUSED and int(ctx.rank) <= 1 and not torch.is_grad_enabled():
+    # The gate is read here rather than at module import: the repository
+    # convention for these switches is construction/call time, not import
+    # time, and a read here is baked into the traced graph at trace time
+    # just the same (one environment lookup per backward invocation).
+    if use_rot_mix_bwd_fused() and int(ctx.rank) <= 1 and not torch.is_grad_enabled():
         # Fused path: the node gradient accumulates on chip inside the
         # kernel; the (E, D, C_wide) ``gxe`` HBM round-trip never exists.
         grad_x, grad_wigner, grad_kc = _rotate_mix_bwd_fused_op(
diff --git a/deepmd/pt_expt/kernels/utils.py b/deepmd/pt_expt/kernels/utils.py
index 19b08aca48..a75e166bb6 100644
--- a/deepmd/pt_expt/kernels/utils.py
+++ b/deepmd/pt_expt/kernels/utils.py
@@ -386,7 +386,9 @@ def use_rot_mix_bwd_fused() -> bool:
     """Return whether the rotate+mix backward fuses its segment reduction.
 
     The flag is controlled by the ``DP_ROT_MIX_BWD_FUSED_INFER`` environment
-    variable and is read at module construction time.  The Triton rotate+mix
+    variable and is read when the backward runs (not at import time), so it
+    takes effect whenever it is set before evaluation; a compiled graph
+    bakes in the value read at trace time.  The Triton rotate+mix
     path itself also binds for training (``triton_train_level >= 1`` with
     wide channels), so this gate is not inference-only: the fused backward
     engages whenever no higher-order gradient graph is being built.  When
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index aa6a8cca38..f80a062ecf 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2171,8 +2171,8 @@ def test_fused_matches_unfused_across_family(self):
                     gx_fused, gx_ref, atol=1e-5 * max(scale, 1.0), rtol=1e-5
                 )
 
-    def test_autograd_branch_follows_module_flag(self):
-        """The autograd backward dispatches on ``_ROT_MIX_BWD_FUSED``."""
+    def test_autograd_branch_follows_env_flag(self):
+        """The autograd backward dispatches on ``DP_ROT_MIX_BWD_FUSED_INFER``."""
         from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
 
         lmax, channels, n_focus, rank = 3, 32, 2, 1
@@ -2188,14 +2188,19 @@ def test_autograd_branch_follows_module_flag(self):
         )
         grads = {}
         for flag in (False, True):
-            saved = vp._ROT_MIX_BWD_FUSED
-            vp._ROT_MIX_BWD_FUSED = flag
+            import os
+
+            saved = os.environ.get("DP_ROT_MIX_BWD_FUSED_INFER")
+            os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = "1" if flag else "0"
             try:
                 grads[flag] = torch.autograd.grad(
                     out, [xg, wg, kg], grad_u, retain_graph=True
                 )
             finally:
-                vp._ROT_MIX_BWD_FUSED = saved
+                if saved is None:
+                    os.environ.pop("DP_ROT_MIX_BWD_FUSED_INFER", None)
+                else:
+                    os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = saved
         # All three agree at fp32 rounding level: grad_x by accumulation
         # order, gw/gkc by instruction scheduling of the same math.
         scale = grads[False][0].abs().max().item()
@@ -2232,15 +2237,20 @@ def test_higher_order_backward_preserved_with_flag_on(self):
         )
         second = {}
         for flag in (False, True):
-            saved = vp._ROT_MIX_BWD_FUSED
-            vp._ROT_MIX_BWD_FUSED = flag
+            import os
+
+            saved = os.environ.get("DP_ROT_MIX_BWD_FUSED_INFER")
+            os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = "1" if flag else "0"
             try:
                 (g1,) = torch.autograd.grad(
                     out, xg, grad_u, create_graph=True, retain_graph=True
                 )
                 second[flag] = torch.autograd.grad(g1.sum(), [wg, kg])
             finally:
-                vp._ROT_MIX_BWD_FUSED = saved
+                if saved is None:
+                    os.environ.pop("DP_ROT_MIX_BWD_FUSED_INFER", None)
+                else:
+                    os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = saved
         for i in (0, 1):
             s = second[False][i].abs().max().item()
             torch.testing.assert_close(
@@ -2270,12 +2280,17 @@ def test_rank_above_one_falls_back(self):
         )
         grads = {}
         for flag in (False, True):
-            saved = vp._ROT_MIX_BWD_FUSED
-            vp._ROT_MIX_BWD_FUSED = flag
+            import os
+
+            saved = os.environ.get("DP_ROT_MIX_BWD_FUSED_INFER")
+            os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = "1" if flag else "0"
             try:
                 (grads[flag],) = torch.autograd.grad(out, xg, grad_u, retain_graph=True)
             finally:
-                vp._ROT_MIX_BWD_FUSED = saved
+                if saved is None:
+                    os.environ.pop("DP_ROT_MIX_BWD_FUSED_INFER", None)
+                else:
+                    os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = saved
         # The rank>1 chain itself is not bitwise deterministic across runs
         # (atomic accumulation), so compare at fp32-rounding tolerance: the
         # point is that both flag settings take the same unfused path.

From 0060d6c206443799e3423294a3dac38447bda069 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Thu, 8 Oct 2026 18:18:16 +0800
Subject: [PATCH 06/13] docs(pt): rename the removed gate constant in a test
 docstring

The higher-order-backward test still referenced _ROT_MIX_BWD_FUSED,
a module constant that no longer exists; name the environment
variable instead.
---
 source/tests/pt/model/test_descriptor_sezm_triton.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index f80a062ecf..3586dddba5 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2219,8 +2219,9 @@ def test_higher_order_backward_preserved_with_flag_on(self):
         The fused operator registers only a fake (meta) implementation and
         no next-derivative formula; the backward of the backward must
         therefore route through the unfused backward + segment-sum pair
-        even when ``_ROT_MIX_BWD_FUSED`` is enabled.  Both flag settings
-        take that same path here, so the second-order results agree.
+        even when ``DP_ROT_MIX_BWD_FUSED_INFER`` is enabled.  Both flag
+        settings take that same path here, so the second-order results
+        agree.
         """
         from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
 

From 36e63b081a4b9dac5e740e9ee5c5eaf4be30a819 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Thu, 15 Oct 2026 16:49:05 +0800
Subject: [PATCH 07/13] test(pt): cover non-power-of-two C_wide and rank-0 in
 the fused backward test

---
 source/tests/pt/model/test_descriptor_sezm_triton.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index 3586dddba5..91bdcfd730 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2071,8 +2071,11 @@ class TestSeZMTritonRotMixBwdFused(unittest.TestCase):
         # (lmax, channels, n_focus, rank)
         (1, 32, 1, 0),
         (2, 32, 1, 1),
-        (3, 64, 2, 1),  # production shape (C_wide = 64, lmax = 3)
+        (3, 64, 2, 1),  # production shape (C_wide = 128, lmax = 3)
         (4, 64, 1, 1),
+        # non-power-of-two C_wide exercises the PADDED channel lanes
+        (2, 24, 2, 1),  # C_wide = 48
+        (1, 40, 1, 0),
     ]
 
     N_NODE = 512

From 39440fd6f1ea5988d5aba3edc57fd004ea5ea38e Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Fri, 16 Oct 2026 19:25:08 +0800
Subject: [PATCH 08/13] test(pt): rank-2 unfused coverage and a real
 always-true channel mask

---
 .../kernels/triton/sezm/so2_value_path.py     |  4 +++-
 .../pt/model/test_descriptor_sezm_triton.py   | 24 +++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

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 01fc5a5958..789148e817 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
@@ -1096,7 +1096,9 @@ def _rotate_mix_bwd_fused_kernel(
             cmask = chan < CW
             chan_c = tl.where(cmask, chan, 0)
         else:
-            cmask = None
+            # CP == CW here; a real always-true mask keeps the variable
+            # uniformly defined instead of a None placeholder.
+            cmask = chan < CP
             chan_c = chan
         didx = tl.arange(0, DP)
         # Every edge of the segment shares this source row, so ``x`` is
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index 91bdcfd730..12e08e876c 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2076,6 +2076,8 @@ class TestSeZMTritonRotMixBwdFused(unittest.TestCase):
         # non-power-of-two C_wide exercises the PADDED channel lanes
         (2, 24, 2, 1),  # C_wide = 48
         (1, 40, 1, 0),
+        # rank 2 exercises the rank-indexed cb/cb and flattened kc addressing
+        (2, 32, 1, 2),
     ]
 
     N_NODE = 512
@@ -2147,6 +2149,28 @@ def test_fused_matches_unfused_across_family(self):
                 )
                 gx_ref = vp._segment_sum_op(gxe, order, row_ptr)
 
+                # The fused kernel supports rank <= 1 (the dispatch contract
+                # keeps rank >= 2 on the unfused pair), so rank 2 checks the
+                # unfused kernels' rank-indexed addressing instead: gkc and
+                # the segment-summed gx agree with the eager reference, while
+                # gw differs from it by a pre-existing upstream discrepancy
+                # that reproduces bit-for-bit on plain master (tracked
+                # separately; not introduced by this PR).
+                if rank > 1:
+                    gxe_e, gw_e, gkc_e = vp._rotate_mix_backward_reference(
+                        grad_u, x, src, wigner, kc, cb, lmax, n_focus, rank
+                    )
+                    torch.testing.assert_close(gkc_ref, gkc_e, atol=1e-5, rtol=1e-5)
+                    self.assertTrue(gw_ref.isfinite().all())
+                    gx_e = torch.zeros(
+                        self.N_NODE,
+                        (lmax + 1) ** 2,
+                        channels * n_focus,
+                        device="cuda",
+                    ).index_add_(0, src, gxe_e)
+                    torch.testing.assert_close(gx_ref, gx_e, atol=1e-4, rtol=1e-4)
+                    continue
+
                 gx_fused, gw_fused, gkc_fused = vp._rotate_mix_bwd_fused_op(
                     grad_u,
                     x,

From 20acbc0bad69757bc64e71d1b80d7588163a2cd1 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sun, 18 Oct 2026 18:00:29 +0800
Subject: [PATCH 09/13] refactor(pt): route the fused rotate-mix backward
 through the level contract and share the Wigner row addressing

---
 .../kernels/triton/sezm/so2_value_path.py     | 44 ++++++++++++------
 deepmd/pt_expt/kernels/utils.py               | 46 -------------------
 2 files changed, 29 insertions(+), 61 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 789148e817..66fb3bfc10 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
@@ -105,7 +105,8 @@
 )
 
 from ...utils import (
-    use_rot_mix_bwd_fused,
+    triton_infer_level,
+    triton_train_level,
 )
 from .gated_activation import (
     gated_activation_second_order,
@@ -565,10 +566,10 @@ def _rotate_mix_fwd_kernel(
             accp = tl.zeros((BC,), dtype=tl.float32)
             for j in tl.static_range(2 * l + 1):
                 xv = xrows[l * l + j]
-                acc0 += tl.load(d_base + r0 * DIM + base + j) * xv
+                acc0 += tl.load(_wigner_row_ptr(d_base, r0, base + j, DIM)) * xv
                 if l >= 1:
-                    accm += tl.load(d_base + (r0 - 1) * DIM + base + j) * xv
-                    accp += tl.load(d_base + (r0 + 1) * DIM + base + j) * xv
+                    accm += tl.load(_wigner_row_ptr(d_base, r0 - 1, base + j, DIM)) * xv
+                    accp += tl.load(_wigner_row_ptr(d_base, r0 + 1, base + j, DIM)) * xv
             rows0 = rows0 + (acc0,)
             if l >= 1:
                 rows_m = rows_m + (accm,)
@@ -708,10 +709,10 @@ def _rotate_mix_bwd_kernel(
             accp = tl.zeros((BC,), dtype=tl.float32)
             for j in tl.static_range(2 * l + 1):
                 xv = xrows[l * l + j]
-                acc0 += tl.load(d_base + r0 * DIM + base + j) * xv
+                acc0 += tl.load(_wigner_row_ptr(d_base, r0, base + j, DIM)) * xv
                 if l >= 1:
-                    accm += tl.load(d_base + (r0 - 1) * DIM + base + j) * xv
-                    accp += tl.load(d_base + (r0 + 1) * DIM + base + j) * xv
+                    accm += tl.load(_wigner_row_ptr(d_base, r0 - 1, base + j, DIM)) * xv
+                    accp += tl.load(_wigner_row_ptr(d_base, r0 + 1, base + j, DIM)) * xv
             rows0 = rows0 + (acc0,)
             if l >= 1:
                 rows_m = rows_m + (accm,)
@@ -1036,6 +1037,11 @@ def _segment_sum_kernel(
             acc += tl.load(rows_ptr + e * P + cols, mask=col_mask, other=0.0)
         tl.store(out_ptr + node * P + cols, acc, mask=col_mask)
 
+    @triton.jit
+    def _wigner_row_ptr(d_base, row, col, DIM: tl.constexpr):
+        """Address of one Wigner entry; shared row/column arithmetic."""
+        return d_base + row * DIM + col
+
     @triton.jit
     def _rotate_mix_bwd_fused_kernel(
         gu_ptr,  # (F, E, ROW) upstream gradient (focus-major)
@@ -1151,7 +1157,11 @@ def _rotate_mix_bwd_fused_kernel(
                         )
                     else:
                         xv = tl.load(x_node + (base + j) * x_sd + chan_c)
-                    w0 = tl.load(d_base + r0 * DIM + base + j, mask=e_mask, other=0.0)
+                    w0 = tl.load(
+                        _wigner_row_ptr(d_base, r0, base + j, DIM),
+                        mask=e_mask,
+                        other=0.0,
+                    )
                     xl0 += w0[:, None] * xv[None, :]
                     if l >= 1:
                         wm = tl.load(
@@ -3195,12 +3205,7 @@ def _rotate_mix_bwd_fused_impl(
         grad_x_edge, grad_wigner, grad_kc = _rotate_mix_backward_reference(
             grad_u, x, src, wigner, kc, cb, lmax, n_focus, rank
         )
-        counts = row_ptr[1:] - row_ptr[:-1]
-        seg_of_sorted = torch.repeat_interleave(
-            torch.arange(n_node, device=x.device, dtype=order.dtype), counts
-        )
-        grad_x = grad_x_edge.new_zeros((n_node, dim, c_wide))
-        grad_x.index_add_(0, seg_of_sorted, grad_x_edge.index_select(0, order))
+        grad_x = _segment_sum_impl(grad_x_edge, order, row_ptr)
         return grad_x, grad_wigner, grad_kc
     grad_x = torch.empty(n_node, dim, c_wide, device=x.device, dtype=x.dtype)
     grad_wigner = torch.zeros_like(wigner)
@@ -4206,7 +4211,16 @@ def _rotate_mix_backward(ctx, grad_u):
     # convention for these switches is construction/call time, not import
     # time, and a read here is baked into the traced graph at trace time
     # just the same (one environment lookup per backward invocation).
-    if use_rot_mix_bwd_fused() and int(ctx.rank) <= 1 and not torch.is_grad_enabled():
+    # The fused backward rides the level-1 universal-kernel contract of
+    # the rotate-mix family (no tunable configuration; it strictly removes
+    # the (E, D, C_wide) node-gradient round-trip).  The existing
+    # inference/training boundary is preserved: whichever level selected
+    # this operator's forward governs its backward too.
+    if (
+        max(triton_infer_level(), triton_train_level()) >= 1
+        and int(ctx.rank) <= 1
+        and not torch.is_grad_enabled()
+    ):
         # Fused path: the node gradient accumulates on chip inside the
         # kernel; the (E, D, C_wide) ``gxe`` HBM round-trip never exists.
         grad_x, grad_wigner, grad_kc = _rotate_mix_bwd_fused_op(
diff --git a/deepmd/pt_expt/kernels/utils.py b/deepmd/pt_expt/kernels/utils.py
index a75e166bb6..c6e54bd17f 100644
--- a/deepmd/pt_expt/kernels/utils.py
+++ b/deepmd/pt_expt/kernels/utils.py
@@ -380,49 +380,3 @@ def use_amp_infer() -> bool:
         ``True`` when ``DP_AMP_INFER`` is set to a truthy value.
     """
     return os.environ.get("DP_AMP_INFER", "0").strip().lower() in _INFER_TRUE
-
-
-def use_rot_mix_bwd_fused() -> bool:
-    """Return whether the rotate+mix backward fuses its segment reduction.
-
-    The flag is controlled by the ``DP_ROT_MIX_BWD_FUSED_INFER`` environment
-    variable and is read when the backward runs (not at import time), so it
-    takes effect whenever it is set before evaluation; a compiled graph
-    bakes in the value read at trace time.  The Triton rotate+mix
-    path itself also binds for training (``triton_train_level >= 1`` with
-    wide channels), so this gate is not inference-only: the fused backward
-    engages whenever no higher-order gradient graph is being built.  When
-    the backward is itself differentiated (``create_graph=True``, e.g.
-    force-loss training), the dispatch falls back to the unfused backward +
-    segment-sum pair, which carry next-derivative formulas the fused
-    operator does not register.
-
-    When enabled, the SO(2) rotate+mix backward
-    (``sezm_triton::so2_rotate_mix_bwd_fused`` in
-    :mod:`.triton.sezm.so2_value_path`) replaces the two-kernel sequence
-    "write the per-edge dense node gradient ``gxe`` ``(E, D, C_wide)`` to
-    HBM, then read it back in ``sezm_triton::segment_sum``" with a single
-    source-segmented kernel: each program owns one source node's CSR
-    segment, recomputes the per-edge rotation backward in registers (the
-    same math as the edge-block kernel), and accumulates the node gradient
-    on chip, so ``gxe`` is never materialized.  At N=4096 (E ~ 646k,
-    ``C_wide = 64``) this eliminates a multi-GB HBM round-trip plus the
-    transient allocation.
-
-    ``grad_wigner``/``grad_kc`` keep the unfused kernel's per-edge math
-    (agreeing to fp32 rounding; instruction scheduling differs);
-    ``grad_x`` changes its per-segment accumulation order (chunked
-    ``tl.sum`` instead of strict sequential adds), so it agrees with the
-    unfused path only to fp32 rounding.  The fused kernel requires
-    ``rank <= 1`` (same regime as the edge-block variant); other shapes
-    keep the unfused path.
-
-    Returns
-    -------
-    bool
-        ``True`` when ``DP_ROT_MIX_BWD_FUSED_INFER`` is set to a truthy
-        value.
-    """
-    return (
-        os.environ.get("DP_ROT_MIX_BWD_FUSED_INFER", "0").strip().lower() in _INFER_TRUE
-    )

From 4a37e7f050b70c8f954cb31d30a90be11a885e87 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sun, 18 Oct 2026 18:00:29 +0800
Subject: [PATCH 10/13] test(pt): structural-block rank-2 Wigner comparison;
 drop the flag parser class

---
 .../pt/model/test_descriptor_sezm_triton.py   | 42 +++++--------------
 1 file changed, 10 insertions(+), 32 deletions(-)

diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index 12e08e876c..270360c047 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2149,19 +2149,21 @@ def test_fused_matches_unfused_across_family(self):
                 )
                 gx_ref = vp._segment_sum_op(gxe, order, row_ptr)
 
-                # The fused kernel supports rank <= 1 (the dispatch contract
-                # keeps rank >= 2 on the unfused pair), so rank 2 checks the
-                # unfused kernels' rank-indexed addressing instead: gkc and
-                # the segment-summed gx agree with the eager reference, while
-                # gw differs from it by a pre-existing upstream discrepancy
-                # that reproduces bit-for-bit on plain master (tracked
-                # separately; not introduced by this PR).
+                # The fused kernel supports rank <= 1 (the dispatch
+                # contract keeps rank >= 2 on the unfused pair), so rank 2
+                # checks the unfused kernels against the eager reference.
+                # The Wigner gradient is compared on the structural block
+                # diagonal: the kernel promises only those entries, while
+                # the dense reference also returns off-block ones.
                 if rank > 1:
                     gxe_e, gw_e, gkc_e = vp._rotate_mix_backward_reference(
                         grad_u, x, src, wigner, kc, cb, lmax, n_focus, rank
                     )
                     torch.testing.assert_close(gkc_ref, gkc_e, atol=1e-5, rtol=1e-5)
-                    self.assertTrue(gw_ref.isfinite().all())
+                    mask = _block_mask(lmax, "cuda")
+                    torch.testing.assert_close(
+                        gw_ref * mask, gw_e * mask, atol=1e-5, rtol=1e-5
+                    )
                     gx_e = torch.zeros(
                         self.N_NODE,
                         (lmax + 1) ** 2,
@@ -2326,27 +2328,3 @@ def test_rank_above_one_falls_back(self):
         torch.testing.assert_close(
             grads[True], grads[False], atol=1e-6 * max(s, 1.0), rtol=1e-5
         )
-
-
-class TestSeZMRotMixFlagParser(unittest.TestCase):
-    """Parse ``DP_ROT_MIX_BWD_FUSED_INFER`` without needing a GPU."""
-
-    def test_env_flag_defaults_off(self):
-        """``DP_ROT_MIX_BWD_FUSED_INFER`` parses truthy/falsy, default off."""
-        import os
-
-        from deepmd.pt_expt.kernels.utils import (
-            use_rot_mix_bwd_fused,
-        )
-
-        saved = os.environ.pop("DP_ROT_MIX_BWD_FUSED_INFER", None)
-        try:
-            self.assertFalse(use_rot_mix_bwd_fused())
-            os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = "1"
-            self.assertTrue(use_rot_mix_bwd_fused())
-            os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = "false"
-            self.assertFalse(use_rot_mix_bwd_fused())
-        finally:
-            os.environ.pop("DP_ROT_MIX_BWD_FUSED_INFER", None)
-            if saved is not None:
-                os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = saved

From b10b427c344a5ed9ec6edf34a1934f7953e46030 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Wed, 21 Oct 2026 22:22:54 +0800
Subject: [PATCH 11/13] fix(pt_expt): retain the level that selected the
 rotate-mix forward

The fused backward reread the ambient DP_TRITON_INFER/DP_TRITON_TRAIN,
so a module constructed under level 1 whose environment was later
restored kept its Triton forward but silently fell back to the unfused
backward plus segment sum. The governing level now travels as an
explicit operator argument from the construction-bound wrappers (the
per-block level captured at construction, following the module's
current train/eval mode) and reaches the backward through the autograd
context, retaining exactly the selection the forward made.

The three tests that still toggled the deleted DP_ROT_MIX_BWD_FUSED_INFER
switch are retargeted to the level argument with a counting stand-in
asserting which backward helper runs, and a new regression test covers
the construction-then-env-restore scenario end to end.
---
 .../kernels/triton/sezm/so2_value_path.py     |  82 +++++--
 .../pt/model/test_descriptor_sezm_triton.py   | 227 ++++++++++++------
 2 files changed, 206 insertions(+), 103 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 66fb3bfc10..3301c333d8 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py
@@ -104,10 +104,6 @@
     wrap_triton,
 )
 
-from ...utils import (
-    triton_infer_level,
-    triton_train_level,
-)
 from .gated_activation import (
     gated_activation_second_order,
     gated_activation_second_order_reference,
@@ -2727,10 +2723,14 @@ def _rotate_mix_impl(
     lmax: int,
     n_focus: int,
     rank: int,
+    level: int,
 ) -> Tensor:
     # The source CSR view rides through the forward untouched so the autograd
-    # context can hand it to the backward's segment reduction.
-    del src_order, src_rowptr
+    # context can hand it to the backward's segment reduction.  The level is
+    # likewise inert here: it reaches the backward through the autograd
+    # context (``ctx.level``), retaining the selection the forward's caller
+    # made instead of rereading the ambient environment.
+    del src_order, src_rowptr, level
     if not _use_triton(x):
         return _rotate_mix_reference(x, src, wigner, kc, cb, lmax, n_focus, rank)
     n_edge = src.shape[0]
@@ -3964,7 +3964,8 @@ def _(grad_u, x, src, wigner, kc, cb, order, row_ptr, lmax, n_focus, rank):
 
 
 @_rotate_mix_op.register_fake
-def _(x, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank):
+def _(x, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank, level):
+    del level
     focus_dim = x.shape[2] // n_focus
     return x.new_empty((n_focus, src.shape[0], (3 * lmax + 1) * focus_dim))
 
@@ -4190,11 +4191,12 @@ def rotate_mix_basis_grad(
 
 
 def _rotate_mix_setup_context(ctx, inputs, output):
-    x, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank = inputs
+    x, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank, level = inputs
     ctx.save_for_backward(x, src, src_order, src_rowptr, wigner, kc, cb)
     ctx.lmax = lmax
     ctx.n_focus = n_focus
     ctx.rank = rank
+    ctx.level = level
 
 
 def _rotate_mix_backward(ctx, grad_u):
@@ -4207,20 +4209,15 @@ def _rotate_mix_backward(ctx, grad_u):
     # unfused backward and segment-sum operators do support the next
     # derivative -- so higher-order callers keep the supported pair and the
     # fused path engages only when no higher-order graph is being built.
-    # The gate is read here rather than at module import: the repository
-    # convention for these switches is construction/call time, not import
-    # time, and a read here is baked into the traced graph at trace time
-    # just the same (one environment lookup per backward invocation).
-    # The fused backward rides the level-1 universal-kernel contract of
-    # the rotate-mix family (no tunable configuration; it strictly removes
-    # the (E, D, C_wide) node-gradient round-trip).  The existing
-    # inference/training boundary is preserved: whichever level selected
-    # this operator's forward governs its backward too.
-    if (
-        max(triton_infer_level(), triton_train_level()) >= 1
-        and int(ctx.rank) <= 1
-        and not torch.is_grad_enabled()
-    ):
+    # The fused backward rides the level-1 universal-kernel contract of the
+    # rotate-mix family (no tunable configuration; it strictly removes the
+    # (E, D, C_wide) node-gradient round-trip).  The governing level is the
+    # one the forward's caller passed as an explicit argument -- the
+    # per-block value captured at construction, following the module's
+    # current train/eval mode -- so the backward retains exactly the
+    # selection the forward made and never rereads the ambient environment
+    # (which may have been restored between construction and evaluation).
+    if int(ctx.level) >= 1 and int(ctx.rank) <= 1 and not torch.is_grad_enabled():
         # Fused path: the node gradient accumulates on chip inside the
         # kernel; the (E, D, C_wide) ``gxe`` HBM round-trip never exists.
         grad_x, grad_wigner, grad_kc = _rotate_mix_bwd_fused_op(
@@ -4250,7 +4247,20 @@ def _rotate_mix_backward(ctx, grad_u):
         if int(ctx.rank) > 0 and ctx.needs_input_grad[6]
         else None
     )
-    return grad_x, None, None, None, grad_wigner, grad_kc, grad_cb, None, None, None
+    # The trailing None is the level argument: integers carry no gradient.
+    return (
+        grad_x,
+        None,
+        None,
+        None,
+        grad_wigner,
+        grad_kc,
+        grad_cb,
+        None,
+        None,
+        None,
+        None,
+    )
 
 
 def _rotate_mix_bwd_setup_context(ctx, inputs, output):
@@ -4283,8 +4293,22 @@ def _rotate_mix_bwd_backward(ctx, grad_grad_x, grad_grad_wigner, grad_grad_kc):
     edge_src = torch.arange(src.shape[0], device=src.device, dtype=src.dtype)
 
     def forward(x_arg: Tensor, src_arg: Tensor, w_arg: Tensor, k_arg: Tensor) -> Tensor:
+        # The level argument is inert on this re-entry: grad mode stays
+        # enabled while the second order is being built, so the inner
+        # operator's own backward can never take the fused path whatever
+        # this value is; 1 documents the family's level-1 contract.
         return _rotate_mix_op(
-            x_arg, src_arg, edge_src, edge_src, w_arg, k_arg, cb, lmax, n_focus, rank
+            x_arg,
+            src_arg,
+            edge_src,
+            edge_src,
+            w_arg,
+            k_arg,
+            cb,
+            lmax,
+            n_focus,
+            rank,
+            1,
         )
 
     def backward(
@@ -5471,6 +5495,11 @@ def __call__(
             conv.lmax,
             conv.n_focus,
             rank,
+            # The per-block levels were captured at construction; the
+            # backward must retain the same selection this forward made
+            # (training level in training mode, inference level
+            # otherwise), not reread the ambient environment.
+            conv.triton_train_level if conv.training else conv.triton_infer_level,
         )
         return u0, rad_feat
 
@@ -5703,6 +5732,11 @@ def __call__(
             conv.lmax,
             conv.n_focus,
             rank,
+            # The per-block levels were captured at construction; the
+            # backward must retain the same selection this forward made
+            # (training level in training mode, inference level
+            # otherwise), not reread the ambient environment.
+            conv.triton_train_level if conv.training else conv.triton_infer_level,
         )
 
         # === Step 3. Cross-focus competition weight from the l = 0 scalars ===
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index 270360c047..4f33d10eef 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -853,6 +853,7 @@ def test_float64_fallback_channel_basis_gradient_preserves_precision(self) -> No
             lmax,
             n_focus,
             rank,
+            0,  # CPU tensors take the reference path; level 0 keeps it
         )
         (grad_fused,) = torch.autograd.grad(fused, basis_fused, grad_out)
 
@@ -2200,131 +2201,199 @@ def test_fused_matches_unfused_across_family(self):
                     gx_fused, gx_ref, atol=1e-5 * max(scale, 1.0), rtol=1e-5
                 )
 
-    def test_autograd_branch_follows_env_flag(self):
-        """The autograd backward dispatches on ``DP_ROT_MIX_BWD_FUSED_INFER``."""
+    def _record_fused_backward(self, run):
+        """Run ``run()`` while counting fused-backward launches; pass through."""
         from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
 
+        calls = []
+        original = vp._rotate_mix_bwd_fused_op
+
+        def recorder(*args, **kwargs):
+            calls.append(args)
+            return original(*args, **kwargs)
+
+        with mock.patch.object(vp, "_rotate_mix_bwd_fused_op", side_effect=recorder):
+            result = run()
+        torch.cuda.synchronize()
+        return len(calls), result
+
+    def test_autograd_branch_follows_forward_level(self):
+        """The autograd backward dispatches on the level the forward was given.
+
+        The level argument stands in for the per-block value captured at
+        construction; the ambient environment is cleared to zero throughout
+        so the check proves the argument (not the environment) governs.
+        """
+        import os
+
+        from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
+
+        for name in ("DP_TRITON_INFER", "DP_TRITON_TRAIN"):
+            os.environ.pop(name, None)
         lmax, channels, n_focus, rank = 3, 32, 2, 1
         grad_u, x, src, wigner, kc, cb = self._inputs(lmax, channels, n_focus, rank)
-        xg = x.clone().requires_grad_(True)
-        wg = wigner.clone().requires_grad_(True)
-        kg = kc.clone().requires_grad_(True)
         src_order = torch.argsort(src)
         boundaries = torch.arange(self.N_NODE + 1, device=src.device, dtype=src.dtype)
         src_rowptr = torch.searchsorted(src.index_select(0, src_order), boundaries)
-        out = vp._rotate_mix_op(
-            xg, src, src_order, src_rowptr, wg, kg, cb, lmax, n_focus, rank
-        )
-        grads = {}
-        for flag in (False, True):
-            import os
 
-            saved = os.environ.get("DP_ROT_MIX_BWD_FUSED_INFER")
-            os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = "1" if flag else "0"
-            try:
-                grads[flag] = torch.autograd.grad(
+        grads, launches = {}, {}
+        for level in (0, 1):
+            xg = x.clone().requires_grad_(True)
+            wg = wigner.clone().requires_grad_(True)
+            kg = kc.clone().requires_grad_(True)
+            out = vp._rotate_mix_op(
+                xg, src, src_order, src_rowptr, wg, kg, cb, lmax, n_focus, rank, level
+            )
+            launches[level], grads[level] = self._record_fused_backward(
+                lambda: torch.autograd.grad(
                     out, [xg, wg, kg], grad_u, retain_graph=True
                 )
-            finally:
-                if saved is None:
-                    os.environ.pop("DP_ROT_MIX_BWD_FUSED_INFER", None)
-                else:
-                    os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = saved
+            )
+        # The counting stand-in proves the selection happened; gradient
+        # equality alone could pass with both levels on the unfused path.
+        self.assertEqual(launches[0], 0)
+        self.assertEqual(launches[1], 1)
         # All three agree at fp32 rounding level: grad_x by accumulation
         # order, gw/gkc by instruction scheduling of the same math.
-        scale = grads[False][0].abs().max().item()
+        scale = grads[0][0].abs().max().item()
         torch.testing.assert_close(
-            grads[True][0], grads[False][0], atol=1e-5 * max(scale, 1.0), rtol=1e-5
+            grads[1][0], grads[0][0], atol=1e-5 * max(scale, 1.0), rtol=1e-5
         )
         for i in (1, 2):
-            s = grads[False][i].abs().max().item()
+            s_i = grads[0][i].abs().max().item()
             torch.testing.assert_close(
-                grads[True][i], grads[False][i], atol=1e-6 * max(s, 1.0), rtol=1e-5
+                grads[1][i], grads[0][i], atol=1e-6 * max(s_i, 1.0), rtol=1e-5
             )
 
-    def test_higher_order_backward_preserved_with_flag_on(self):
-        """With the flag on, ``create_graph=True`` keeps the supported path.
+    def test_higher_order_backward_preserved_at_level_one(self):
+        """At level 1, ``create_graph=True`` keeps the supported path.
 
         The fused operator registers only a fake (meta) implementation and
         no next-derivative formula; the backward of the backward must
         therefore route through the unfused backward + segment-sum pair
-        even when ``DP_ROT_MIX_BWD_FUSED_INFER`` is enabled.  Both flag
-        settings take that same path here, so the second-order results
-        agree.
+        even with the level-1 selection active.  The counting stand-in
+        asserts the fused backward never launches during either order, and
+        the second-order results agree with a level-0 run.
         """
         from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
 
         lmax, channels, n_focus, rank = 3, 32, 2, 1
         grad_u, x, src, wigner, kc, cb = self._inputs(lmax, channels, n_focus, rank)
-        xg = x.clone().requires_grad_(True)
-        wg = wigner.clone().requires_grad_(True)
-        kg = kc.clone().requires_grad_(True)
         src_order = torch.argsort(src)
         boundaries = torch.arange(self.N_NODE + 1, device=src.device, dtype=src.dtype)
         src_rowptr = torch.searchsorted(src.index_select(0, src_order), boundaries)
-        out = vp._rotate_mix_op(
-            xg, src, src_order, src_rowptr, wg, kg, cb, lmax, n_focus, rank
-        )
-        second = {}
-        for flag in (False, True):
-            import os
 
-            saved = os.environ.get("DP_ROT_MIX_BWD_FUSED_INFER")
-            os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = "1" if flag else "0"
-            try:
-                (g1,) = torch.autograd.grad(
+        second, launches = {}, {}
+        for level in (0, 1):
+            xg = x.clone().requires_grad_(True)
+            wg = wigner.clone().requires_grad_(True)
+            kg = kc.clone().requires_grad_(True)
+            out = vp._rotate_mix_op(
+                xg, src, src_order, src_rowptr, wg, kg, cb, lmax, n_focus, rank, level
+            )
+            launches[level], (g1,) = self._record_fused_backward(
+                lambda: torch.autograd.grad(
                     out, xg, grad_u, create_graph=True, retain_graph=True
                 )
-                second[flag] = torch.autograd.grad(g1.sum(), [wg, kg])
-            finally:
-                if saved is None:
-                    os.environ.pop("DP_ROT_MIX_BWD_FUSED_INFER", None)
-                else:
-                    os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = saved
+            )
+            second[level] = torch.autograd.grad(g1.sum(), [wg, kg])
+        self.assertEqual(launches[0], 0)
+        self.assertEqual(launches[1], 0)
         for i in (0, 1):
-            s = second[False][i].abs().max().item()
+            s_i = second[0][i].abs().max().item()
             torch.testing.assert_close(
-                second[True][i],
-                second[False][i],
-                atol=1e-6 * max(s, 1.0),
+                second[1][i],
+                second[0][i],
+                atol=1e-6 * max(s_i, 1.0),
                 rtol=1e-5,
             )
 
-    def test_rank_above_one_falls_back(self):
-        """``rank > 1`` never takes the fused path (flag or not).
+    def test_rank_above_one_falls_back_at_level_one(self):
+        """``rank > 1`` never takes the fused path, level or not.
 
-        Both flag settings take the same unfused path; that chain is not
-        bitwise deterministic across runs, so agreement is checked at
-        fp32-rounding tolerance.
+        The rank guard stands even with the level-1 selection active; the
+        counting stand-in asserts the unfused chain, and the rank>1 chain
+        is not bitwise deterministic across runs, so agreement between the
+        two levels is checked at fp32-rounding tolerance.
         """
         from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
 
         lmax, channels, n_focus, rank = 3, 32, 2, 2
         grad_u, x, src, wigner, kc, cb = self._inputs(lmax, channels, n_focus, rank)
-        xg = x.clone().requires_grad_(True)
         src_order = torch.argsort(src)
         boundaries = torch.arange(self.N_NODE + 1, device=src.device, dtype=src.dtype)
         src_rowptr = torch.searchsorted(src.index_select(0, src_order), boundaries)
-        out = vp._rotate_mix_op(
-            xg, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank
-        )
-        grads = {}
-        for flag in (False, True):
-            import os
-
-            saved = os.environ.get("DP_ROT_MIX_BWD_FUSED_INFER")
-            os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = "1" if flag else "0"
-            try:
-                (grads[flag],) = torch.autograd.grad(out, xg, grad_u, retain_graph=True)
-            finally:
-                if saved is None:
-                    os.environ.pop("DP_ROT_MIX_BWD_FUSED_INFER", None)
-                else:
-                    os.environ["DP_ROT_MIX_BWD_FUSED_INFER"] = saved
-        # The rank>1 chain itself is not bitwise deterministic across runs
-        # (atomic accumulation), so compare at fp32-rounding tolerance: the
-        # point is that both flag settings take the same unfused path.
-        s = grads[False].abs().max().item()
+
+        grads, launches = {}, {}
+        for level in (0, 1):
+            xg = x.clone().requires_grad_(True)
+            out = vp._rotate_mix_op(
+                xg,
+                src,
+                src_order,
+                src_rowptr,
+                wigner,
+                kc,
+                cb,
+                lmax,
+                n_focus,
+                rank,
+                level,
+            )
+            launches[level], (grads[level],) = self._record_fused_backward(
+                lambda: torch.autograd.grad(out, xg, grad_u, retain_graph=True)
+            )
+        self.assertEqual(launches[0], 0)
+        self.assertEqual(launches[1], 0)
+        s_i = grads[0].abs().max().item()
         torch.testing.assert_close(
-            grads[True], grads[False], atol=1e-6 * max(s, 1.0), rtol=1e-5
+            grads[1], grads[0], atol=1e-6 * max(s_i, 1.0), rtol=1e-5
         )
+
+    def test_backward_retains_construction_level_across_env_restore(self):
+        """The wrapper's captured level governs the backward, not ambient env.
+
+        A module constructed with ``DP_TRITON_INFER=1`` keeps its Triton
+        forward after the environment is restored to the default; its
+        backward must retain the same selection instead of rereading the
+        (now zero) ambient level and silently falling back.
+        """
+        import os
+        from types import SimpleNamespace
+
+        from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
+
+        lmax, channels, n_focus, rank = 3, 32, 2, 0
+        grad_u, x, src, wigner, kc, cb = self._inputs(lmax, channels, n_focus, rank)
+        src_order = torch.argsort(src)
+        boundaries = torch.arange(self.N_NODE + 1, device=src.device, dtype=src.dtype)
+        src_rowptr = torch.searchsorted(src.index_select(0, src_order), boundaries)
+        edge_cache = SimpleNamespace(
+            src=src,
+            D_full=wigner,
+            csr_cache={"src": (src_order, src_rowptr)},
+        )
+        radial_feat = kc if rank == 0 else None
+
+        for captured, expected_launches in ((1, 1), (0, 0)):
+            with self.subTest(captured_infer_level=captured):
+                conv = SimpleNamespace(
+                    radial_hidden_proj=None,
+                    radial_degree_mixer=None,
+                    lmax=lmax,
+                    n_focus=n_focus,
+                    training=False,
+                    triton_infer_level=captured,
+                    triton_train_level=0,
+                )
+                wrapper = vp._TritonRotateMix(conv)
+                for name in ("DP_TRITON_INFER", "DP_TRITON_TRAIN"):
+                    os.environ.pop(name, None)
+                xg = x.clone().requires_grad_(True)
+                out, _rad = wrapper(xg, edge_cache, radial_feat)
+                launches, _grads = self._record_fused_backward(
+                    lambda: torch.autograd.grad(
+                        out, xg, grad_u.reshape(out.shape), retain_graph=True
+                    )
+                )
+                self.assertEqual(launches, expected_launches)

From 158e77e57b62be1cf0d5e9f73dca4bad6457a2e6 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 21 Oct 2026 14:24:09 +0000
Subject: [PATCH 12/13] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/tests/pt/model/test_descriptor_sezm_triton.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index 4f33d10eef..f721a876b1 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2359,7 +2359,9 @@ def test_backward_retains_construction_level_across_env_restore(self):
         (now zero) ambient level and silently falling back.
         """
         import os
-        from types import SimpleNamespace
+        from types import (
+            SimpleNamespace,
+        )
 
         from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
 

From 0befdc56cdcf1ddb4097affa266de27d931a69d4 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Thu, 22 Oct 2026 19:26:44 +0800
Subject: [PATCH 13/13] test(pt): scope the level-override environment in the
 rotate-mix tests

The two new checks deleted both configured Triton level variables from
the test process without restoring them, leaking the cleared environment
to every later test. Both now clear through a shared helper that
registers the restore as a cleanup, so the configured levels survive
even when an assertion fails mid-test.
---
 .../pt/model/test_descriptor_sezm_triton.py   | 37 +++++++++++++++----
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index f721a876b1..00d0770e42 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2217,6 +2217,30 @@ def recorder(*args, **kwargs):
         torch.cuda.synchronize()
         return len(calls), result
 
+    def _clear_and_restore_levels(self):
+        """Clear both Triton level variables; restore the prior values on exit.
+
+        The override is registered as a cleanup, so the configured
+        environment survives even when an assertion fails mid-test.
+        """
+        import os
+
+        saved = {
+            name: os.environ.get(name)
+            for name in ("DP_TRITON_INFER", "DP_TRITON_TRAIN")
+        }
+
+        def restore():
+            for name, value in saved.items():
+                if value is None:
+                    os.environ.pop(name, None)
+                else:
+                    os.environ[name] = value
+
+        for name in saved:
+            os.environ.pop(name, None)
+        self.addCleanup(restore)
+
     def test_autograd_branch_follows_forward_level(self):
         """The autograd backward dispatches on the level the forward was given.
 
@@ -2224,12 +2248,11 @@ def test_autograd_branch_follows_forward_level(self):
         construction; the ambient environment is cleared to zero throughout
         so the check proves the argument (not the environment) governs.
         """
-        import os
-
         from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
 
-        for name in ("DP_TRITON_INFER", "DP_TRITON_TRAIN"):
-            os.environ.pop(name, None)
+        # Scope the override: the cleared levels are restored even when an
+        # assertion fails, so later tests keep the configured environment.
+        self._clear_and_restore_levels()
         lmax, channels, n_focus, rank = 3, 32, 2, 1
         grad_u, x, src, wigner, kc, cb = self._inputs(lmax, channels, n_focus, rank)
         src_order = torch.argsort(src)
@@ -2358,13 +2381,15 @@ def test_backward_retains_construction_level_across_env_restore(self):
         backward must retain the same selection instead of rereading the
         (now zero) ambient level and silently falling back.
         """
-        import os
         from types import (
             SimpleNamespace,
         )
 
         from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp
 
+        # Scope the ambient-env restore once: the cleared levels come back
+        # even when a subTest assertion fails.
+        self._clear_and_restore_levels()
         lmax, channels, n_focus, rank = 3, 32, 2, 0
         grad_u, x, src, wigner, kc, cb = self._inputs(lmax, channels, n_focus, rank)
         src_order = torch.argsort(src)
@@ -2389,8 +2414,6 @@ def test_backward_retains_construction_level_across_env_restore(self):
                     triton_train_level=0,
                 )
                 wrapper = vp._TritonRotateMix(conv)
-                for name in ("DP_TRITON_INFER", "DP_TRITON_TRAIN"):
-                    os.environ.pop(name, None)
                 xg = x.clone().requires_grad_(True)
                 out, _rad = wrapper(xg, edge_cache, radial_feat)
                 launches, _grads = self._record_fused_backward(