From ec97031b8817edcae909efd67d6c6ba860bee8d8 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Wed, 7 Oct 2026 19:25:36 +0800 Subject: [PATCH 01/10] 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/10] [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 Date: Wed, 7 Oct 2026 20:22:03 +0800 Subject: [PATCH 03/10] 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 Date: Wed, 7 Oct 2026 21:12:21 +0800 Subject: [PATCH 04/10] 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 Date: Thu, 8 Oct 2026 12:54:21 +0800 Subject: [PATCH 05/10] 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 Date: Thu, 8 Oct 2026 18:18:16 +0800 Subject: [PATCH 06/10] 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 Date: Thu, 15 Oct 2026 16:49:05 +0800 Subject: [PATCH 07/10] 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 Date: Fri, 16 Oct 2026 19:25:08 +0800 Subject: [PATCH 08/10] 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 Date: Sun, 18 Oct 2026 18:00:29 +0800 Subject: [PATCH 09/10] 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 Date: Sun, 18 Oct 2026 18:00:29 +0800 Subject: [PATCH 10/10] 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