From a3dcd972cfbbfb41947e29d96c58ee6bd02cc91f Mon Sep 17 00:00:00 2001 From: longjunhao Date: Thu, 8 Oct 2026 19:44:23 +0800 Subject: [PATCH 01/10] perf(pt): fuse the mixing-stack per-layer backward chain (DP_STACK_BWD_FUSED_INFER) Each gated layer's backward runs as recompute -> pointwise -> GEMM, materializing the sigmoid surface and the pre-activation gradient to HBM and reading both straight back. With the flag on (and the supported regime: fp32, narrow channels, no higher-order graph), one kernel walks the chain with both surfaces in registers, matching the relay bit-for-bit; the m = 0 and |m| = 1 column segments launch separately to keep the lighter segment off the heavy one's register profile. The 2.15 GB standing scratch never allocates; on the N=4096 eval workload it does not coincide with the frame's peak. --- .../kernels/triton/sezm/so2_value_path.py | 304 +++++++++++++++++- deepmd/pt_expt/kernels/utils.py | 25 ++ .../pt/model/test_descriptor_sezm_triton.py | 247 ++++++++++++++ 3 files changed, 573 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 7a8ac3fd43..678f600d2a 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py +++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py @@ -117,6 +117,7 @@ from .so2_rotation import ( _block_to_local_op, ) +from ...utils import use_stack_bwd_fused from .tile_configs import ( GATE_BMM_MIN_FOCUS_DIM, gate_config, @@ -2152,6 +2153,202 @@ def _stack_grad_alpha_kernel( mask=m_mask, ) + # Launch geometry for the fused per-layer backward below. BLOCK_M=64 / + # BLOCK_N=64 with 4 warps fills the 255-register budget and spills 226 + # slots; the spill traffic eats the fusion win, so the M dimension is + # halved instead. BLOCK_K is not free here: it must equal the channel + # group width so one K tile is exactly one row group (see the kernel). + _STACK_BWD_FUSED_CONFIG = (32, 64, 4, 3) + + @triton.jit + def _stack_bwd_fused_kernel( + g_ptr, # (F, E, ROW) upstream gradient of the layer output, focus-major + z_ptr, # z_all stack (NL, F, E, ROW), layer selected by ``layer`` + gw_ptr, # (NL, F, CF, L*CF) gate projections + gwt_ptr, # (NL, F, L*CF, CF) transposed gate projections + w0t_ptr, # (NL, F, M0, M0) stacked transposed weights (m = 0 block) + w1t_ptr, # (NL, F, M1, M1) stacked transposed weights (|m| = 1 block) + gu_ptr, # (F, E, ROW) layer-input gradient output + n_edge, + layer, + L: tl.constexpr, + CF: tl.constexpr, + PART: tl.constexpr, # 0: m = 0 block (with the scalar chain); 1: |m| = 1 + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + """Single-pass backward of one gated layer; ``gz`` lives only in registers. + + Replaces the three-kernel relay (gate-sigmoid recompute -> pointwise + chain writing ``gz`` to HBM -> block GEMM reading ``gz`` back) with two + launches over disjoint column segments of ``g_next``. The tiles + (BLOCK_M/N/K) and the ``tl.dot`` reduction order copy + ``_stack_gemm_bwd_kernel``; each row group's ``gz`` tile is recomputed + in place from ``g``/``z`` -- the scalar group (r = 0) replays the full + dot-accumulation chain of ``_stack_point_bwd_kernel``, the other row + groups multiply by their gate -- so the numerics match the relay + bit-for-bit. PART splits the launch because the m = 0 path has far + higher register pressure than m = 1; a shared kernel would drag the + m = 1 tiles down to its occupancy. Requires ``CF == BLOCK_K`` so each + K tile is exactly one row group. + """ + M0: tl.constexpr = (L + 1) * CF + M1: tl.constexpr = 2 * L * CF + ROW: tl.constexpr = (3 * L + 1) * CF + LG: tl.constexpr = L * CF + CP: tl.constexpr = triton.next_power_of_2(CF) + NT0: tl.constexpr = (M0 + BLOCK_N - 1) // BLOCK_N + NT1: tl.constexpr = (M1 + BLOCK_N - 1) // BLOCK_N + + if PART == 0: + NTile: tl.constexpr = NT0 + else: + NTile: tl.constexpr = NT1 + + pid = tl.program_id(0) + fid = tl.program_id(1).to(tl.int64) + n_focus = tl.num_programs(1) + pid_m = pid // NTile + pid_n = pid % NTile + + offs_m = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)).to(tl.int64) + m_mask = offs_m < n_edge + mm = m_mask[:, None] + nc = tl.arange(0, CP) + cm = mm & (nc < CF)[None, :] + wm = ((nc < CF)[:, None]) & ((nc < CF)[None, :]) + + g_row = g_ptr + fid * n_edge * ROW + offs_m * ROW + z_row = z_ptr + (layer * n_focus + fid) * n_edge * ROW + offs_m * ROW + gu_row = gu_ptr + fid * n_edge * ROW + offs_m * ROW + + # Gate sigmoid, recomputed in place with the same operator sequence + # as _stack_recompute_kernel. + z_s = tl.load(z_row[:, None] + nc[None, :], mask=cm, other=0.0).to(tl.float32) + sigs = () + for g_ in tl.static_range(L): + gw_g = tl.load( + gw_ptr + + (layer * n_focus + fid) * CF * LG + + nc[:, None] * LG + + (g_ * CF + nc)[None, :], + mask=wm, + other=0.0, + ).to(tl.float32) + sigs = sigs + (tl.sigmoid(tl.dot(z_s, gw_g, input_precision="ieee")),) + + offs_k = tl.arange(0, BLOCK_K) + if PART == 0: + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + n_mask = offs_n < M0 + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + w_ptrs = ( + w0t_ptr + + (layer * n_focus + fid) * M0 * M0 + + offs_k[:, None] * M0 + + offs_n[None, :] + ) + # K tiles advance one row group at a time: the scalar group walks + # the full gated chain, the gated groups scale by their gate. + for r in tl.static_range(M0 // CF): + if r == 0: + g_s = tl.load(g_row[:, None] + nc[None, :], mask=cm, other=0.0).to( + tl.float32 + ) + s0 = tl.sigmoid(z_s) + a = g_s * s0 * (1.0 + z_s * (1.0 - s0)) + for g_ in tl.static_range(L): + gr0 = tl.load( + g_row[:, None] + ((1 + g_) * CF + nc)[None, :], + mask=cm, + other=0.0, + ) + zr0 = tl.load( + z_row[:, None] + ((1 + g_) * CF + nc)[None, :], + mask=cm, + other=0.0, + ) + rn = (L + 1) + g_ + grn = tl.load( + g_row[:, None] + (rn * CF + nc)[None, :], + mask=cm, + other=0.0, + ) + zrn = tl.load( + z_row[:, None] + (rn * CF + nc)[None, :], + mask=cm, + other=0.0, + ) + rp = (2 * L + 1) + g_ + grp = tl.load( + g_row[:, None] + (rp * CF + nc)[None, :], + mask=cm, + other=0.0, + ) + zrp = tl.load( + z_row[:, None] + (rp * CF + nc)[None, :], + mask=cm, + other=0.0, + ) + sig_g = sigs[g_] + g_sig = gr0 * zr0 + grn * zrn + grp * zrp + g_logit = g_sig * sig_g * (1.0 - sig_g) + gwt_g = tl.load( + gwt_ptr + + (layer * n_focus + fid) * LG * CF + + (g_ * CF + nc)[:, None] * CF + + nc[None, :], + mask=wm, + other=0.0, + ).to(tl.float32) + a = tl.dot(g_logit, gwt_g, a, input_precision="ieee") + else: + g_r = tl.load( + g_row[:, None] + (r * CF + nc)[None, :], mask=cm, other=0.0 + ) + a = g_r * sigs[r - 1] + w = tl.load(w_ptrs, mask=n_mask[None, :], other=0.0) + acc = tl.dot(a, w, acc, input_precision="ieee") + w_ptrs += BLOCK_K * M0 + col0 = offs_n + col_mask = n_mask + else: + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + n_mask = offs_n < M1 + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + w_ptrs = ( + w1t_ptr + + (layer * n_focus + fid) * M1 * M1 + + offs_k[:, None] * M1 + + offs_n[None, :] + ) + # The m = 1 block is all |m| = 1 stripes: the two stripes of gate + # group g_ ((L+1)+g_ and (2L+1)+g_) share one gate. Stripes outer, + # gate groups inner, keeping the K-tile accumulation order of the + # relay; static indices use plain loop variables only. + for s_ in tl.static_range(2): + for g_ in tl.static_range(L): + r = (L + 1) + s_ * L + g_ + g_r = tl.load( + g_row[:, None] + (r * CF + nc)[None, :], mask=cm, other=0.0 + ) + a = g_r * sigs[g_] + w = tl.load(w_ptrs, mask=n_mask[None, :], other=0.0) + acc = tl.dot(a, w, acc, input_precision="ieee") + w_ptrs += BLOCK_K * M1 + col0 = M0 + offs_n + col_mask = n_mask + + # The residual is this layer's upstream gradient g, loaded the same + # way _stack_gemm_bwd_kernel loads it. + res = tl.load( + g_row[:, None] + col0[None, :], mask=mm & col_mask[None, :], other=0.0 + ) + tl.store( + gu_row[:, None] + col0[None, :], acc + res, mask=mm & col_mask[None, :] + ) + # ====================================================================== # Zero-edge guard and dispatch predicate @@ -2463,6 +2660,60 @@ def _launch_stack_point_backward( ) +def _launch_stack_bwd_fused( + g_cur: Tensor, + z_all: Tensor, + gw_all: Tensor, + gwt_all: Tensor, + w0t_all: Tensor, + w1t_all: Tensor, + g_next: Tensor, + n_edge, + layer: int, + lmax: int, + focus_dim: int, + block_k: int, + n_focus: int, +) -> None: + """Launch one gated layer's fused backward into ``g_next``. + + The two PART launches write disjoint column segments (the ``m = 0`` + block with the scalar chain, then the ``|m| = 1`` stripes); see + ``_stack_bwd_fused_kernel`` for the register-pressure rationale. + ``block_k`` comes from the shared backward GEMM configuration and the + dispatch has already verified it equals ``focus_dim``. + """ + lmax = int(lmax) + focus_dim = int(focus_dim) + m0 = (lmax + 1) * focus_dim + m1 = 2 * lmax * focus_dim + fbm, fbn, fwarps, fstages = _STACK_BWD_FUSED_CONFIG + nt0 = triton.cdiv(m0, fbn) + nt1 = triton.cdiv(m1, fbn) + for part, nt_part in ((0, nt0), (1, nt1)): + wrap_triton(_stack_bwd_fused_kernel)[ + (triton.cdiv(n_edge, fbm) * nt_part, n_focus) + ]( + g_cur, + z_all, + gw_all, + gwt_all, + w0t_all, + w1t_all, + g_next, + n_edge, + layer, + L=lmax, + CF=focus_dim, + PART=part, + BLOCK_M=fbm, + BLOCK_N=fbn, + BLOCK_K=block_k, + num_warps=fwarps, + num_stages=fstages, + ) + + # ====================================================================== # Operator implementations (Triton on CUDA fp32, eager reference otherwise) # ====================================================================== @@ -3276,7 +3527,31 @@ def _stack_backward_traversal( # Inference and first-order weight gradients consume each surface before # advancing to the next layer, so one scratch allocation serves the stack. gate_width = lmax * focus_dim - sig = torch.empty((n_focus, n_edge, gate_width), device=device, dtype=torch.float32) + use_bmm = focus_dim >= GATE_BMM_MIN_FOCUS_DIM + # The gate is read here (when the backward runs), matching the repository + # convention of call-time reads for these switches; a compiled graph bakes + # in the value observed at trace time. The fused kernel registers no + # next-derivative formula, so it engages only when no higher-order + # gradient graph is being built (torch.is_grad_enabled() is True inside + # this backward exactly then); the supported relay carries those + # derivatives. ``CF == block_k`` makes one K tile exactly one row group + # -- the bitwise-matching reduction order the fused kernel assumes. + fused = ( + use_stack_bwd_fused() + and not use_bmm + and focus_dim == block_k + and (focus_dim & (focus_dim - 1)) == 0 + and grad_out.dtype == torch.float32 + and not with_weights + and not keep + and grad_z_upstream is None + and not torch.is_grad_enabled() + ) + sig = ( + torch.empty((n_focus, n_edge, gate_width), device=device, dtype=torch.float32) + if not fused + else None + ) grad_z_all = ( torch.empty((n_gated, n_focus, n_edge, row), device=device, dtype=dtype) if keep @@ -3284,10 +3559,9 @@ def _stack_backward_traversal( ) grad_z_scratch = ( torch.empty((n_focus, n_edge, row), device=device, dtype=dtype) - if not keep + if not keep and not fused else None ) - use_bmm = focus_dim >= GATE_BMM_MIN_FOCUS_DIM store_logit = with_weights or use_bmm or (keep and need_logit) grad_logit_all = ( torch.empty( @@ -3311,6 +3585,30 @@ def _stack_backward_traversal( ) u_next = u_final for layer in range(n_gated - 1, -1, -1): + if fused: + # Single-pass per-layer backward: the gate sigmoids and the + # pre-activation gradient stay in registers, so neither the + # sigmoid surface nor the gz scratch buffer exists. + assert g_spare is not None + g_next = g_spare + _launch_stack_bwd_fused( + g_cur, + z_all, + gw_all, + gwt_all, + w0t_all, + w1t_all, + g_next, + n_edge, + layer, + lmax, + focus_dim, + block_k, + n_focus, + ) + g_spare = g_cur + g_cur = g_next + continue if keep: assert grad_z_all is not None gz = grad_z_all[layer] diff --git a/deepmd/pt_expt/kernels/utils.py b/deepmd/pt_expt/kernels/utils.py index c6e54bd17f..6c42b26440 100644 --- a/deepmd/pt_expt/kernels/utils.py +++ b/deepmd/pt_expt/kernels/utils.py @@ -380,3 +380,28 @@ 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_stack_bwd_fused() -> bool: + """Return whether the mixing-stack backward fuses its per-layer chain. + + The flag is controlled by the ``DP_STACK_BWD_FUSED_INFER`` environment + variable and is read when the backward traversal runs, so it takes effect + whenever it is set before evaluation; a compiled graph bakes in the value + observed at trace time. When enabled, each gated layer's backward chain + (gate-sigmoid recompute, pointwise backward, transposed-weight GEMM) runs + as one kernel and the intermediate sigmoid and pre-activation-gradient + surfaces never reach HBM. The fused kernel registers no next-derivative + formula, so the dispatch additionally requires that no higher-order + gradient graph is being built (``create_graph=True``, e.g. force-loss + training) and falls back to the unfused chain otherwise; the Triton stack + itself also binds for training, so the gate is not inference-only. + + Returns + ------- + bool + ``True`` when ``DP_STACK_BWD_FUSED_INFER`` is set to a truthy value. + """ + return ( + os.environ.get("DP_STACK_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..86b3f3ba6a 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -23,6 +23,7 @@ """ import math +import os import typing import unittest from unittest import ( @@ -2051,5 +2052,251 @@ def test_tune_missing_configs_runs_the_real_flash_bwd_sweep(self): self.assertTrue(tc.has_tile_config("flash_bwd_block", key)) +@_GPU_KERNELS +class TestSeZMStackBwdFused(unittest.TestCase): + """Fused per-layer fp32 mixing-stack backward against the unfused relay. + + The fused kernel keeps each layer's gate sigmoids and pre-activation + gradient in registers, so its results must match the unfused + recompute/pointwise/GEMM relay bit-for-bit at fp32; the remaining tests + pin the dispatch contract (env flag, higher-order fallback, wide-channel + and non-fp32 fallback). + """ + + N_EDGE = 20000 + N_FOCUS = 2 + N_LAYERS = 4 # one final linear layer plus three gated layers + CASES = ((1, 32, False), (3, 32, False), (3, 32, True)) # lmax, Cf, apply_alpha + + def _bwd_inputs(self, lmax, focus_dim, generator): + m0 = (lmax + 1) * focus_dim + half = lmax * focus_dim + row = (3 * lmax + 1) * focus_dim + n_gated = self.N_LAYERS - 1 + + def randn(*shape, scale=1.0): + return torch.randn(*shape, device="cuda", generator=generator) * scale + + grad_out = randn(self.N_EDGE, self.N_FOCUS, row, scale=0.05) + x_local = randn(self.N_EDGE, self.N_FOCUS, row, scale=0.05) + alpha = ( + torch.rand(self.N_EDGE, self.N_FOCUS, device="cuda", generator=generator) + + 0.5 + ) + z_all = randn(n_gated, self.N_FOCUS, self.N_EDGE, row, scale=0.5) + u_final = randn(self.N_EDGE, self.N_FOCUS, row, scale=0.05) + w0t_all = randn(self.N_LAYERS, self.N_FOCUS, m0, m0, scale=0.02) + block_u = randn(self.N_LAYERS, self.N_FOCUS, half, half, scale=0.02) + block_v = randn(self.N_LAYERS, self.N_FOCUS, half, half, scale=0.02) + # The |m| = 1 weight carries the [[U, V], [-V, U]] complex structure + # of SO2Linear so the synthetic stack matches the production operator. + w1t_all = torch.cat( + [ + torch.cat([block_u, block_v], dim=3), + torch.cat([-block_v, block_u], dim=3), + ], + dim=2, + ).contiguous() + gw_all = randn(n_gated, self.N_FOCUS, focus_dim, half, scale=0.2) + gwt_all = gw_all.transpose(-1, -2).contiguous() + return ( + grad_out, + x_local, + alpha, + z_all, + u_final, + w0t_all, + w1t_all, + gw_all, + gwt_all, + ) + + def _run_traversal(self, inputs, lmax, focus_dim, apply_alpha): + from deepmd.pt_expt.kernels.triton.sezm.so2_value_path import ( + _stack_backward_traversal, + ) + + with torch.no_grad(): + g, ga, _, _ = _stack_backward_traversal( + inputs[0], + inputs[1], + inputs[3], + inputs[4], + inputs[2], + inputs[5], + inputs[6], + inputs[7], + inputs[8], + None, + None, + lmax, + focus_dim, + apply_alpha, + with_weights=False, + keep=False, + ) + torch.cuda.synchronize() + return g, ga + + def test_fused_matches_unfused_bitwise(self): + generator = torch.Generator(device="cuda").manual_seed(42) + for lmax, focus_dim, apply_alpha in self.CASES: + with self.subTest(lmax=lmax, cf=focus_dim, alpha=apply_alpha): + inputs = self._bwd_inputs(lmax, focus_dim, generator) + old = os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) + try: + ref = self._run_traversal(inputs, lmax, focus_dim, apply_alpha) + os.environ["DP_STACK_BWD_FUSED_INFER"] = "1" + fused = self._run_traversal(inputs, lmax, focus_dim, apply_alpha) + finally: + os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) + if old is not None: + os.environ["DP_STACK_BWD_FUSED_INFER"] = old + self.assertTrue(torch.equal(ref[0], fused[0])) + # grad_alpha is computed only when alpha was applied; the + # traversal leaves it uninitialized otherwise. + if apply_alpha: + self.assertTrue(torch.equal(ref[1], fused[1])) + + def test_higher_order_backward_preserved_with_flag_on(self): + """``create_graph=True`` keeps the supported path with the flag on. + + The fused kernel registers no next-derivative formula; the backward + of the backward must route through the operator that carries the + hand-derived second order even when the flag is enabled. Both flag + settings take that same path here, so the second-order results + agree. + """ + from deepmd.pt_expt.kernels.triton.sezm.so2_value_path import ( + _mixing_stack_op, + ) + + lmax, focus_dim = 3, 32 + generator = torch.Generator(device="cuda").manual_seed(7) + m0 = (lmax + 1) * focus_dim + half = lmax * focus_dim + row = (3 * lmax + 1) * focus_dim + + def randn(*shape, scale=1.0): + return torch.randn(*shape, device="cuda", generator=generator) * scale + + u0 = randn(self.N_FOCUS, self.N_EDGE, row, scale=0.5) + alpha = ( + torch.rand(self.N_EDGE, self.N_FOCUS, device="cuda", generator=generator) + + 0.1 + ) + w0_all = randn(self.N_LAYERS, self.N_FOCUS, m0, m0, scale=0.2) + block_u = randn(self.N_LAYERS, self.N_FOCUS, half, half, scale=0.2) + block_v = randn(self.N_LAYERS, self.N_FOCUS, half, half, scale=0.2) + w1_all = torch.cat( + [ + torch.cat([block_u, block_v], dim=3), + torch.cat([-block_v, block_u], dim=3), + ], + dim=2, + ).contiguous() + gw_all = randn(self.N_LAYERS - 1, self.N_FOCUS, focus_dim, half, scale=0.3) + + # Gradients are requested for the gate weights as well: the + # double-backward replays the gate-logit surfaces only when the + # weight-gradient channel carries a cotangent. + results = {} + old = os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) + try: + for flag in (False, True): + os.environ["DP_STACK_BWD_FUSED_INFER"] = "1" if flag else "0" + # Reseed so both states draw identical cotangent samples. + torch.manual_seed(99) + u0r = u0.clone().requires_grad_(True) + gwr = gw_all.clone().requires_grad_(True) + out = _mixing_stack_op( + u0r, alpha, w0_all, w1_all, gwr, lmax, focus_dim, True + ) + seed = torch.randn_like(out[0]) + g_u0, g_gw = torch.autograd.grad( + (out[0] * seed).sum(), (u0r, gwr), create_graph=True + ) + seed_gw = torch.randn_like(g_gw) + (second,) = torch.autograd.grad((g_gw * seed_gw).sum(), u0r) + results[flag] = (g_u0.detach().clone(), second.detach().clone()) + finally: + os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) + if old is not None: + os.environ["DP_STACK_BWD_FUSED_INFER"] = old + torch.testing.assert_close(results[False][0], results[True][0]) + torch.testing.assert_close(results[False][1], results[True][1]) + + def _record_launches(self, lmax, focus_dim, apply_alpha=False, dtype=torch.float32): + """Run the traversal with a recording stand-in for the fused launch.""" + import deepmd.pt_expt.kernels.triton.sezm.so2_value_path as vp + + calls = [] + original = vp._launch_stack_bwd_fused + + def recorder(*args, **kwargs): + calls.append(args[8]) # the layer index + return original(*args, **kwargs) + + generator = torch.Generator(device="cuda").manual_seed(11) + inputs = self._bwd_inputs(lmax, focus_dim, generator) + inputs = tuple(t.to(dtype) for t in inputs) + with mock.patch.object(vp, "_launch_stack_bwd_fused", side_effect=recorder): + out = self._run_traversal(inputs, lmax, focus_dim, apply_alpha) + return calls, out + + def test_dispatch_follows_env_flag(self): + old = os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) + try: + for flag, expected in (("1", 3), ("0", 0)): + with self.subTest(flag=flag): + os.environ["DP_STACK_BWD_FUSED_INFER"] = flag + calls, _ = self._record_launches(3, 32) + self.assertEqual(len(calls), expected) + finally: + os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) + if old is not None: + os.environ["DP_STACK_BWD_FUSED_INFER"] = old + + def test_wide_channels_fall_back(self): + """``Cf >= GATE_BMM_MIN_FOCUS_DIM`` routes through cuBLAS, not fusion.""" + from deepmd.pt_expt.kernels.triton.sezm.tile_configs import ( + GATE_BMM_MIN_FOCUS_DIM, + ) + + calls, out = self._record_launches(2, GATE_BMM_MIN_FOCUS_DIM) + self.assertEqual(calls, []) + self.assertTrue(torch.isfinite(out[0]).all()) + + def test_non_fp32_fall_back(self): + """The fused kernel assumes fp32 rounding; other dtypes keep the relay.""" + calls, out = self._record_launches(3, 32, dtype=torch.bfloat16) + self.assertEqual(calls, []) + self.assertTrue(torch.isfinite(out[0]).all()) + + +class TestSeZMStackBwdFlagParser(unittest.TestCase): + def test_values(self): + from deepmd.pt_expt.kernels.utils import use_stack_bwd_fused + + old = os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) + try: + for value, expected in ( + ("1", True), + ("true", True), + ("ON", True), + ("0", False), + ("nope", False), + ): + with self.subTest(value=value): + os.environ["DP_STACK_BWD_FUSED_INFER"] = value + self.assertEqual(use_stack_bwd_fused(), expected) + os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) + self.assertFalse(use_stack_bwd_fused()) + finally: + os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) + if old is not None: + os.environ["DP_STACK_BWD_FUSED_INFER"] = old + + if __name__ == "__main__": unittest.main() From 4455e674cc2ddcaa05e57ebe23df595ad741dc19 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 8 Oct 2026 11:45:54 +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, 6 insertions(+), 2 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 678f600d2a..abb5fda9d9 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_stack_bwd_fused, +) from .gated_activation import ( gated_activation_second_order, gated_activation_second_order_reference, @@ -117,7 +120,6 @@ from .so2_rotation import ( _block_to_local_op, ) -from ...utils import use_stack_bwd_fused from .tile_configs import ( GATE_BMM_MIN_FOCUS_DIM, gate_config, diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index 86b3f3ba6a..cd1b8e98e4 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -2276,7 +2276,9 @@ def test_non_fp32_fall_back(self): class TestSeZMStackBwdFlagParser(unittest.TestCase): def test_values(self): - from deepmd.pt_expt.kernels.utils import use_stack_bwd_fused + from deepmd.pt_expt.kernels.utils import ( + use_stack_bwd_fused, + ) old = os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) try: From ad9af590248897b720b7401e03269f09bbb68181 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Fri, 16 Oct 2026 19:27:49 +0800 Subject: [PATCH 03/10] test(pt): exercise the stack-backward fallback guards with the flag enabled --- .../pt/model/test_descriptor_sezm_triton.py | 41 ++++++++++++++----- 1 file changed, 31 insertions(+), 10 deletions(-) diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index cd1b8e98e4..7b1f3bc767 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -2257,21 +2257,42 @@ def test_dispatch_follows_env_flag(self): if old is not None: os.environ["DP_STACK_BWD_FUSED_INFER"] = old - def test_wide_channels_fall_back(self): - """``Cf >= GATE_BMM_MIN_FOCUS_DIM`` routes through cuBLAS, not fusion.""" + def test_wide_channels_fall_back_with_flag_on(self): + """``Cf >= GATE_BMM_MIN_FOCUS_DIM`` routes through cuBLAS, not fusion. + + The flag is enabled so the fallback is the guard's doing, not the + flag's: zero fused launches with the flag on is the contract. + """ from deepmd.pt_expt.kernels.triton.sezm.tile_configs import ( GATE_BMM_MIN_FOCUS_DIM, ) - calls, out = self._record_launches(2, GATE_BMM_MIN_FOCUS_DIM) - self.assertEqual(calls, []) - self.assertTrue(torch.isfinite(out[0]).all()) + old = os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) + try: + os.environ["DP_STACK_BWD_FUSED_INFER"] = "1" + calls, out = self._record_launches(2, GATE_BMM_MIN_FOCUS_DIM) + self.assertEqual(calls, []) + self.assertTrue(torch.isfinite(out[0]).all()) + finally: + os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) + if old is not None: + os.environ["DP_STACK_BWD_FUSED_INFER"] = old - def test_non_fp32_fall_back(self): - """The fused kernel assumes fp32 rounding; other dtypes keep the relay.""" - calls, out = self._record_launches(3, 32, dtype=torch.bfloat16) - self.assertEqual(calls, []) - self.assertTrue(torch.isfinite(out[0]).all()) + def test_non_fp32_fall_back_with_flag_on(self): + """The fused kernel assumes fp32 rounding; other dtypes keep the relay. + + Scoped with the flag enabled for the same reason as above. + """ + old = os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) + try: + os.environ["DP_STACK_BWD_FUSED_INFER"] = "1" + calls, out = self._record_launches(3, 32, dtype=torch.bfloat16) + self.assertEqual(calls, []) + self.assertTrue(torch.isfinite(out[0]).all()) + finally: + os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) + if old is not None: + os.environ["DP_STACK_BWD_FUSED_INFER"] = old class TestSeZMStackBwdFlagParser(unittest.TestCase): From fa5e98e582fc49a4dbd6e928fd845875acf28900 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Sun, 18 Oct 2026 13:52:41 +0800 Subject: [PATCH 04/10] feat(pt): select the fused mixing-stack backward by tuned tile entry Replace the standalone DP_STACK_BWD_FUSED_INFER switch with the existing DP_TRITON_INFER >= 2 value-path selection: the fused per-layer backward engages only where the freeze-time sweep measured it as the winner for the (focus_dim, lmax) shape, through the new stack_bwd_fused tile-config family. The sweep drives the real backward traversal, accepts a candidate only on bitwise agreement with the relay, and records None (the relay) when fusion does not win by at least 2%; an untuned environment therefore never leaves the unfused relay. The SiLU and sigmoid cotangent chains the fused kernel repeated from the pointwise kernels now live in shared inlined Triton helpers with an identical instruction tree (bitwise agreement re-verified). --- .../kernels/triton/sezm/so2_value_path.py | 85 +++++--- .../kernels/triton/sezm/sweep_tile_configs.py | 186 ++++++++++++++++++ .../kernels/triton/sezm/tile_configs.py | 25 +++ deepmd/pt_expt/kernels/utils.py | 25 --- 4 files changed, 272 insertions(+), 49 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 abb5fda9d9..f43e7d8eb4 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_stack_bwd_fused, + triton_infer_level, + triton_train_level, ) from .gated_activation import ( gated_activation_second_order, @@ -129,6 +130,7 @@ recompute_config, rotate_mix_bwd_block_config, rotate_mix_fwd_config, + stack_bwd_fused_config, stack_fp16x3_configs, stack_fp32_configs, stack_m0_gate_config, @@ -503,6 +505,23 @@ def _mixing_stack_backward_reference( # ====================================================================== if SO2_VALUE_PATH_TRITON_AVAILABLE: + @triton.jit + def _silu_bwd(g, z): + """Cotangent of ``z * sigmoid(z)``: one shared SiLU backward chain. + + The mixing-stack kernels (pointwise backward, training backward and + the fused per-layer backward) contract the upstream gradient + through the same derivative, so the expression lives here once and + inlines into each kernel with an identical instruction tree. + """ + s = tl.sigmoid(z) + return g * s * (1.0 + z * (1.0 - s)) + + @triton.jit + def _sigmoid_bwd(g, sig): + """Cotangent through a precomputed sigmoid: ``g * sig * (1 - sig)``.""" + return g * sig * (1.0 - sig) + @triton.jit def _rotate_mix_fwd_kernel( x_ptr, # (N, D, C_wide), strides (x_sn, x_sd, 1) @@ -1471,11 +1490,14 @@ def _stack_point_bwd_kernel( # the output dtype. z_s = tl.load(z_row[:, None] + nc[None, :], mask=cm, other=0.0).to(tl.float32) g_s = tl.load(g_row[:, None] + nc[None, :], mask=cm, other=0.0).to(tl.float32) - s0 = tl.sigmoid(z_s) - gz_s = g_s * s0 * (1.0 + z_s * (1.0 - s0)) + gz_s = _silu_bwd(g_s, z_s) if RECOVER_INPUT: un_s = tl.load(un_row[:, None] + nc[None, :], mask=cm, other=0.0) - tl.store(up_row[:, None] + nc[None, :], un_s - z_s * s0, mask=cm) + tl.store( + up_row[:, None] + nc[None, :], + un_s - z_s * tl.sigmoid(z_s), + mask=cm, + ) for g in tl.static_range(L): if RECOMPUTE_SIG: @@ -1540,7 +1562,7 @@ def _stack_point_bwd_kernel( ) # Gate path: three value rows share gate group g. g_sig = gr0 * zr0 + grn * zrn + grp * zrp - g_logit = g_sig * sig_g * (1.0 - sig_g) + g_logit = _sigmoid_bwd(g_sig, sig_g) if GLOGIT_STORE: tl.store(gl_row[:, None] + (g * CF + nc)[None, :], g_logit, mask=cm) if not GLOGIT_OUT: @@ -1849,10 +1871,13 @@ def _stack_train_traversal_kernel( ) if KEEP: tl.store(up_row[:, None] + nc[None, :], g_s, mask=cm) - s0 = tl.sigmoid(z_s) - gz_s = g_s * s0 * (1.0 + z_s * (1.0 - s0)) + gz_s = _silu_bwd(g_s, z_s) un_s = tl.load(un_row[:, None] + nc[None, :], mask=cm, other=0.0) - tl.store(ul_row[:, None] + nc[None, :], un_s - z_s * s0, mask=cm) + tl.store( + ul_row[:, None] + nc[None, :], + un_s - z_s * tl.sigmoid(z_s), + mask=cm, + ) for grp in tl.static_range(L): gw_g = tl.load( @@ -2160,8 +2185,6 @@ def _stack_grad_alpha_kernel( # slots; the spill traffic eats the fusion win, so the M dimension is # halved instead. BLOCK_K is not free here: it must equal the channel # group width so one K tile is exactly one row group (see the kernel). - _STACK_BWD_FUSED_CONFIG = (32, 64, 4, 3) - @triton.jit def _stack_bwd_fused_kernel( g_ptr, # (F, E, ROW) upstream gradient of the layer output, focus-major @@ -2258,8 +2281,7 @@ def _stack_bwd_fused_kernel( g_s = tl.load(g_row[:, None] + nc[None, :], mask=cm, other=0.0).to( tl.float32 ) - s0 = tl.sigmoid(z_s) - a = g_s * s0 * (1.0 + z_s * (1.0 - s0)) + a = _silu_bwd(g_s, z_s) for g_ in tl.static_range(L): gr0 = tl.load( g_row[:, None] + ((1 + g_) * CF + nc)[None, :], @@ -2295,7 +2317,7 @@ def _stack_bwd_fused_kernel( ) sig_g = sigs[g_] g_sig = gr0 * zr0 + grn * zrn + grp * zrp - g_logit = g_sig * sig_g * (1.0 - sig_g) + g_logit = _sigmoid_bwd(g_sig, sig_g) gwt_g = tl.load( gwt_ptr + (layer * n_focus + fid) * LG * CF @@ -2676,6 +2698,7 @@ def _launch_stack_bwd_fused( focus_dim: int, block_k: int, n_focus: int, + config: tuple[int, int, int, int], ) -> None: """Launch one gated layer's fused backward into ``g_next``. @@ -2689,7 +2712,7 @@ def _launch_stack_bwd_fused( focus_dim = int(focus_dim) m0 = (lmax + 1) * focus_dim m1 = 2 * lmax * focus_dim - fbm, fbn, fwarps, fstages = _STACK_BWD_FUSED_CONFIG + fbm, fbn, fwarps, fstages = config nt0 = triton.cdiv(m0, fbn) nt1 = triton.cdiv(m1, fbn) for part, nt_part in ((0, nt0), (1, nt1)): @@ -3538,17 +3561,30 @@ def _stack_backward_traversal( # this backward exactly then); the supported relay carries those # derivatives. ``CF == block_k`` makes one K tile exactly one row group # -- the bitwise-matching reduction order the fused kernel assumes. - fused = ( - use_stack_bwd_fused() - and not use_bmm - and focus_dim == block_k - and (focus_dim & (focus_dim - 1)) == 0 - and grad_out.dtype == torch.float32 - and not with_weights - and not keep - and grad_z_upstream is None - and not torch.is_grad_enabled() + # The fused per-layer backward lives inside the table-tuned Triton + # value-path selection (the same >= 2 level as the other tuned kernels): + # it engages only where the freeze-time sweep measured it as the winner + # for this (focus_dim, lmax) shape. A missing or "relay won" table + # entry keeps the unfused relay, so untuned environments never take the + # fused kernel. The remaining conditions are the internal shape and + # higher-order gates: ``CF == block_k`` keeps one K tile exactly one row + # group -- the bitwise-matching reduction order the fused kernel assumes. + stack_bwd_fused_cfg = ( + stack_bwd_fused_config(focus_dim, lmax) + if ( + max(triton_infer_level(), triton_train_level()) >= 2 + and not use_bmm + and focus_dim == block_k + and (focus_dim & (focus_dim - 1)) == 0 + and grad_out.dtype == torch.float32 + and not with_weights + and not keep + and grad_z_upstream is None + and not torch.is_grad_enabled() + ) + else None ) + fused = stack_bwd_fused_cfg is not None sig = ( torch.empty((n_focus, n_edge, gate_width), device=device, dtype=torch.float32) if not fused @@ -3607,6 +3643,7 @@ def _stack_backward_traversal( focus_dim, block_k, n_focus, + stack_bwd_fused_cfg, ) g_spare = g_cur g_cur = g_next diff --git a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py index 44c0a91aac..cf111e55f9 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py +++ b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py @@ -2074,6 +2074,191 @@ def fp16_backward() -> tuple[torch.Tensor, torch.Tensor]: return {"stack_fp16x3": {key: conclusion}} +# ====================================================================== +# Fused mixing-stack backward (fused vs relay selection) +# ====================================================================== +def sweep_stack_bwd_fused( + cf: int, + lmax: int, + *, + n_focus: int = 2, + n_edge: int | None = None, + device: torch.device | str = "cuda", +) -> SweepResult: + """Select between the fused per-layer backward and the unfused relay. + + The comparison drives the real :func:`_stack_backward_traversal` over a + synthetic stack (three gated layers plus the final linear layer, alpha + enabled), so the baseline includes exactly the relay chain production + runs -- the gate-sigmoid recompute, the pointwise backward and the + tuned GEMM launches -- and each candidate replaces it with one fused + launch per layer. A candidate only qualifies when the traversal + outputs match the relay bit-for-bit; the fastest qualifying candidate + is registered, otherwise the key records that the relay won (``None``). + + Parameters + ---------- + cf : int + Per-focus channel width ``Cf``. + lmax : int + Maximum spherical harmonic degree. + n_focus : int + Focus count of the synthetic tensors. + n_edge : int or None + Edge count of the synthetic tensors; ``None`` selects a saturating + count. + device : torch.device or str + CUDA device to sweep on. + + Returns + ------- + SweepResult + The ``stack_bwd_fused`` entry under ``(cf, lmax)``: the winning + ``(BLOCK_M, BLOCK_N, num_warps, num_stages)`` or ``None``. + """ + from .so2_value_path import ( + _stack_backward_traversal, + ) + from .tile_configs import ( + _runtime_tile_configs, + register_tile_configs, + stack_fp32_configs, + ) + + device = torch.device(device) + if n_edge is None: + n_edge = _saturating_edges(cf) + key = (cf, lmax) + + # ``CF == block_k`` is the bitwise reduction-order contract; without it + # the dispatch can never take the fused kernel, so the relay wins by + # construction and no benchmark is needed. + _, _, bwd_cfg = stack_fp32_configs(cf, lmax) + if bwd_cfg[2] != cf or cf & (cf - 1) or cf >= GATE_BMM_MIN_FOCUS_DIM: + register_tile_configs("stack_bwd_fused", {key: None}) + print( + f"[stack_bwd_fused] ({cf},{lmax}) relay by dispatch contract, " + "no sweep needed" + ) + return {"stack_bwd_fused": {key: None}} + + n_layers = 4 + n_gated = n_layers - 1 + m0 = (lmax + 1) * cf + m1 = 2 * lmax * cf + row = (3 * lmax + 1) * cf + half = lmax * cf + + generator = torch.Generator(device=device).manual_seed(0) + + def randn(*shape: int, scale: float = 1.0) -> torch.Tensor: + return torch.randn(*shape, device=device, generator=generator) * scale + + grad_out = randn(n_edge, n_focus, row, scale=0.05) + z_all = randn(n_gated, n_focus, n_edge, row, scale=0.5) + u_final = randn(n_edge, n_focus, row, scale=0.05) + alpha = ( + torch.rand( + n_edge, n_focus, device=device, dtype=torch.float32, generator=generator + ) + + 0.5 + ) + w0t_all = randn(n_layers, n_focus, m0, m0, scale=0.02) + block_u = randn(n_layers, n_focus, half, half, scale=0.02) + block_v = randn(n_layers, n_focus, half, half, scale=0.02) + w1t_all = torch.cat( + [ + torch.cat([block_u, block_v], dim=3), + torch.cat([-block_v, block_u], dim=3), + ], + dim=2, + ) + gw_all = randn(n_gated, n_focus, cf, lmax * cf, scale=0.05) + gwt_all = gw_all.transpose(2, 3).contiguous() + + def run_traversal() -> tuple[torch.Tensor, torch.Tensor]: + g, ga, _, _ = _stack_backward_traversal( + grad_out, + grad_out, + z_all, + u_final, + alpha, + w0t_all, + w1t_all, + gw_all, + gwt_all, + None, + None, + lmax, + cf, + True, + with_weights=False, + keep=False, + ) + return g.clone(), ga.clone() + + candidates = ( + (32, 64, 4, 3), + (64, 64, 4, 3), + (32, 128, 4, 3), + (64, 128, 8, 3), + (16, 64, 4, 2), + (128, 64, 8, 4), + ) + runtime = _runtime_tile_configs("stack_bwd_fused") + saved = dict(runtime) + try: + # Baseline: an explicit None entry forces the relay whatever the + # built-in tables say, so the comparison reflects this GPU only. + register_tile_configs("stack_bwd_fused", {key: None}) + with torch.no_grad(): + ref_g, ref_ga = run_traversal() + relay_ms = _bench(run_traversal, iters=50) + + best_ms, best_cfg = float("inf"), None + for cfg in candidates: + register_tile_configs("stack_bwd_fused", {key: cfg}) + try: + with torch.no_grad(): + got_g, got_ga = run_traversal() + if not (torch.equal(got_g, ref_g) and torch.equal(got_ga, ref_ga)): + print( + f"[stack_bwd_fused] ({cf},{lmax}) candidate " + f"{cfg} rejected: outputs differ from the relay" + ) + continue + fused_ms = _bench(run_traversal, iters=50) + except Exception as err: + print( + f"[stack_bwd_fused] ({cf},{lmax}) candidate {cfg} " + f"rejected: {type(err).__name__}: {err}" + ) + continue + print( + f"[stack_bwd_fused] ({cf},{lmax}) candidate {cfg}: " + f"{fused_ms:.3f} ms (relay {relay_ms:.3f} ms)" + ) + if fused_ms < best_ms: + best_ms, best_cfg = fused_ms, cfg + finally: + runtime.clear() + runtime.update(saved) + + # A 2% margin keeps run-to-run noise from registering a coin-flip as + # a winner; anything below it records the relay. + if best_cfg is not None and best_ms < relay_ms * 0.98: + winner: tuple | None = best_cfg + print( + f"[stack_bwd_fused] ({cf},{lmax}) fused {winner} wins: " + f"{best_ms:.3f} vs {relay_ms:.3f} ms" + ) + else: + winner = None + print(f"[stack_bwd_fused] ({cf},{lmax}) relay wins: {relay_ms:.3f} ms") + register_tile_configs("stack_bwd_fused", {key: winner}) + return {"stack_bwd_fused": {key: winner}} + + @dataclass(frozen=True) class _SweepSpec: """Describe one independently covered launch-configuration sweep.""" @@ -2103,6 +2288,7 @@ class _SweepSpec: "fp32": _SweepSpec(sweep_fp32, "stack_fp32", "focus"), "m0_gate": _SweepSpec(sweep_m0_gate, "stack_m0_gate", "focus"), "fp16x3": _SweepSpec(sweep_fp16x3, "stack_fp16x3", "focus", min_level=3), + "stack_bwd_fused": _SweepSpec(sweep_stack_bwd_fused, "stack_bwd_fused", "focus"), } diff --git a/deepmd/pt_expt/kernels/triton/sezm/tile_configs.py b/deepmd/pt_expt/kernels/triton/sezm/tile_configs.py index 878f734099..44985b6e50 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/tile_configs.py +++ b/deepmd/pt_expt/kernels/triton/sezm/tile_configs.py @@ -138,6 +138,7 @@ "rotate_mix_bwd_block", "stack_fp32", "stack_m0_gate", + "stack_bwd_fused", "stack_fp16x3", ) @@ -488,6 +489,30 @@ def stack_fp32_configs( return _lookup("stack_fp32", (focus_dim, lmax)) or _STACK_FP32_DEFAULT +def stack_bwd_fused_config( + focus_dim: int, lmax: int +) -> tuple[int, int, int, int] | None: + """Return the tuned fused mixing-stack backward launch configuration. + + Parameters + ---------- + focus_dim + Per-focus channel width ``Cf``. + lmax + Maximum spherical harmonic degree. + + Returns + ------- + tuple or None + The ``(BLOCK_M, BLOCK_N, num_warps, num_stages)`` configuration the + sweep measured as the winner, or ``None`` to keep the unfused + relay. ``None`` covers both "the sweep ran and the relay won" and + "never swept on this GPU", so an untuned environment never engages + the fused kernel. + """ + return _lookup("stack_bwd_fused", (focus_dim, lmax)) + + def stack_fp16x3_configs( focus_dim: int, lmax: int ) -> ( diff --git a/deepmd/pt_expt/kernels/utils.py b/deepmd/pt_expt/kernels/utils.py index 6c42b26440..c6e54bd17f 100644 --- a/deepmd/pt_expt/kernels/utils.py +++ b/deepmd/pt_expt/kernels/utils.py @@ -380,28 +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_stack_bwd_fused() -> bool: - """Return whether the mixing-stack backward fuses its per-layer chain. - - The flag is controlled by the ``DP_STACK_BWD_FUSED_INFER`` environment - variable and is read when the backward traversal runs, so it takes effect - whenever it is set before evaluation; a compiled graph bakes in the value - observed at trace time. When enabled, each gated layer's backward chain - (gate-sigmoid recompute, pointwise backward, transposed-weight GEMM) runs - as one kernel and the intermediate sigmoid and pre-activation-gradient - surfaces never reach HBM. The fused kernel registers no next-derivative - formula, so the dispatch additionally requires that no higher-order - gradient graph is being built (``create_graph=True``, e.g. force-loss - training) and falls back to the unfused chain otherwise; the Triton stack - itself also binds for training, so the gate is not inference-only. - - Returns - ------- - bool - ``True`` when ``DP_STACK_BWD_FUSED_INFER`` is set to a truthy value. - """ - return ( - os.environ.get("DP_STACK_BWD_FUSED_INFER", "0").strip().lower() in _INFER_TRUE - ) From 8476318c7303de497ab69af5e474790f0cf617ed Mon Sep 17 00:00:00 2001 From: longjunhao Date: Sun, 18 Oct 2026 13:52:54 +0800 Subject: [PATCH 05/10] test(pt): retarget the fused stack-backward tests at the tuned selection --- .../pt/model/test_descriptor_sezm_triton.py | 202 +++++++++--------- 1 file changed, 102 insertions(+), 100 deletions(-) diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index 7b1f3bc767..89fd4140f0 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -1951,6 +1951,9 @@ def fake_flash_sweep(cf, lmax, **kwargs): "fp32": fake_sweep("fp32", "stack_fp32", (48, 2)), "m0_gate": fake_sweep("m0_gate", "stack_m0_gate", (48, 2)), "fp16x3": fake_sweep("fp16x3", "stack_fp16x3", (48, 2)), + "stack_bwd_fused": fake_sweep( + "stack_bwd_fused", "stack_bwd_fused", (48, 2) + ), } fake_specs = { name: replace(spec, sweep=fake_sweeps[name]) @@ -1973,6 +1976,7 @@ def fake_flash_sweep(cf, lmax, **kwargs): "point_recompute", "rotate_bwd", "rotate_fwd", + "stack_bwd_fused", ], ) self.assertEqual( @@ -1983,6 +1987,7 @@ def fake_flash_sweep(cf, lmax, **kwargs): "point_recompute", "rotate_mix_bwd_block", "rotate_mix_fwd", + "stack_bwd_fused", "stack_fp32", "stack_m0_gate", ], @@ -2058,11 +2063,32 @@ class TestSeZMStackBwdFused(unittest.TestCase): The fused kernel keeps each layer's gate sigmoids and pre-activation gradient in registers, so its results must match the unfused - recompute/pointwise/GEMM relay bit-for-bit at fp32; the remaining tests - pin the dispatch contract (env flag, higher-order fallback, wide-channel - and non-fp32 fallback). + recompute/pointwise/GEMM relay bit-for-bit at fp32. Engagement follows + the tuned table (a winning ``stack_bwd_fused`` entry) inside the + ``DP_TRITON_INFER >= 2`` level selection; the remaining tests pin the + internal guards (higher-order fallback, wide-channel and non-fp32 + fallback, level and table gating). """ + FUSED_CFG = (32, 64, 4, 3) + + def _register_fused(self, key, entry): + """Register a fused table entry and restore the prior state.""" + from deepmd.pt_expt.kernels.triton.sezm.tile_configs import ( + _runtime_tile_configs, + register_tile_configs, + ) + + runtime = _runtime_tile_configs("stack_bwd_fused") + saved = dict(runtime) + + def restore(): + runtime.clear() + runtime.update(saved) + + register_tile_configs("stack_bwd_fused", {key: entry}) + self.addCleanup(restore) + N_EDGE = 20000 N_FOCUS = 2 N_LAYERS = 4 # one final linear layer plus three gated layers @@ -2143,28 +2169,24 @@ def test_fused_matches_unfused_bitwise(self): for lmax, focus_dim, apply_alpha in self.CASES: with self.subTest(lmax=lmax, cf=focus_dim, alpha=apply_alpha): inputs = self._bwd_inputs(lmax, focus_dim, generator) - old = os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) - try: - ref = self._run_traversal(inputs, lmax, focus_dim, apply_alpha) - os.environ["DP_STACK_BWD_FUSED_INFER"] = "1" - fused = self._run_traversal(inputs, lmax, focus_dim, apply_alpha) - finally: - os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) - if old is not None: - os.environ["DP_STACK_BWD_FUSED_INFER"] = old + key = (focus_dim, lmax) + self._register_fused(key, None) + ref = self._run_traversal(inputs, lmax, focus_dim, apply_alpha) + self._register_fused(key, self.FUSED_CFG) + fused = self._run_traversal(inputs, lmax, focus_dim, apply_alpha) self.assertTrue(torch.equal(ref[0], fused[0])) # grad_alpha is computed only when alpha was applied; the # traversal leaves it uninitialized otherwise. if apply_alpha: self.assertTrue(torch.equal(ref[1], fused[1])) - def test_higher_order_backward_preserved_with_flag_on(self): - """``create_graph=True`` keeps the supported path with the flag on. + def test_higher_order_backward_preserved_with_winning_entry(self): + """``create_graph=True`` keeps the supported path with an entry present. The fused kernel registers no next-derivative formula; the backward of the backward must route through the operator that carries the - hand-derived second order even when the flag is enabled. Both flag - settings take that same path here, so the second-order results + hand-derived second order even with a winning table entry. Both + table states take that same path here, so the second-order results agree. """ from deepmd.pt_expt.kernels.triton.sezm.so2_value_path import ( @@ -2201,32 +2223,32 @@ def randn(*shape, scale=1.0): # double-backward replays the gate-logit surfaces only when the # weight-gradient channel carries a cotangent. results = {} - old = os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) - try: - for flag in (False, True): - os.environ["DP_STACK_BWD_FUSED_INFER"] = "1" if flag else "0" - # Reseed so both states draw identical cotangent samples. - torch.manual_seed(99) - u0r = u0.clone().requires_grad_(True) - gwr = gw_all.clone().requires_grad_(True) - out = _mixing_stack_op( - u0r, alpha, w0_all, w1_all, gwr, lmax, focus_dim, True - ) - seed = torch.randn_like(out[0]) - g_u0, g_gw = torch.autograd.grad( - (out[0] * seed).sum(), (u0r, gwr), create_graph=True - ) - seed_gw = torch.randn_like(g_gw) - (second,) = torch.autograd.grad((g_gw * seed_gw).sum(), u0r) - results[flag] = (g_u0.detach().clone(), second.detach().clone()) - finally: - os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) - if old is not None: - os.environ["DP_STACK_BWD_FUSED_INFER"] = old + key = (focus_dim, lmax) + for entry in (None, self.FUSED_CFG): + self._register_fused(key, entry) + # Reseed so both states draw identical cotangent samples. + torch.manual_seed(99) + u0r = u0.clone().requires_grad_(True) + gwr = gw_all.clone().requires_grad_(True) + out = _mixing_stack_op( + u0r, alpha, w0_all, w1_all, gwr, lmax, focus_dim, True + ) + seed = torch.randn_like(out[0]) + g_u0, g_gw = torch.autograd.grad( + (out[0] * seed).sum(), (u0r, gwr), create_graph=True + ) + seed_gw = torch.randn_like(g_gw) + (second,) = torch.autograd.grad((g_gw * seed_gw).sum(), u0r) + results[entry is not None] = ( + g_u0.detach().clone(), + second.detach().clone(), + ) torch.testing.assert_close(results[False][0], results[True][0]) torch.testing.assert_close(results[False][1], results[True][1]) - def _record_launches(self, lmax, focus_dim, apply_alpha=False, dtype=torch.float32): + def _record_launches( + self, lmax, focus_dim, apply_alpha=False, dtype=torch.float32, n_edge=None + ): """Run the traversal with a recording stand-in for the fused launch.""" import deepmd.pt_expt.kernels.triton.sezm.so2_value_path as vp @@ -2238,87 +2260,67 @@ def recorder(*args, **kwargs): return original(*args, **kwargs) generator = torch.Generator(device="cuda").manual_seed(11) - inputs = self._bwd_inputs(lmax, focus_dim, generator) + saved_edges = self.N_EDGE + if n_edge is not None: + type(self).N_EDGE = n_edge + try: + inputs = self._bwd_inputs(lmax, focus_dim, generator) + finally: + type(self).N_EDGE = saved_edges inputs = tuple(t.to(dtype) for t in inputs) with mock.patch.object(vp, "_launch_stack_bwd_fused", side_effect=recorder): out = self._run_traversal(inputs, lmax, focus_dim, apply_alpha) return calls, out - def test_dispatch_follows_env_flag(self): - old = os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) + def test_dispatch_follows_table_and_level(self): + """A winning table entry engages at level >= 2; anything else relays. + + Dispatch-only checks run on a small edge count (the guards are + scale-free), unlike the numerical fixtures above. + """ + old_level = os.environ.get("DP_TRITON_INFER") try: - for flag, expected in (("1", 3), ("0", 0)): - with self.subTest(flag=flag): - os.environ["DP_STACK_BWD_FUSED_INFER"] = flag - calls, _ = self._record_launches(3, 32) + for level, entry, expected in ( + ("2", self.FUSED_CFG, 3), + ("2", None, 0), + ("1", self.FUSED_CFG, 0), + ): + with self.subTest(level=level, entry=entry): + os.environ["DP_TRITON_INFER"] = level + self._register_fused((32, 3), entry) + calls, _ = self._record_launches(3, 32, n_edge=512) self.assertEqual(len(calls), expected) finally: - os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) - if old is not None: - os.environ["DP_STACK_BWD_FUSED_INFER"] = old + os.environ.pop("DP_TRITON_INFER", None) + if old_level is not None: + os.environ["DP_TRITON_INFER"] = old_level - def test_wide_channels_fall_back_with_flag_on(self): + def test_wide_channels_fall_back_with_winning_entry(self): """``Cf >= GATE_BMM_MIN_FOCUS_DIM`` routes through cuBLAS, not fusion. - The flag is enabled so the fallback is the guard's doing, not the - flag's: zero fused launches with the flag on is the contract. + A winning table entry is registered for the wide key so the + fallback is the internal guard's doing: zero fused launches with + the entry present is the contract. Dispatch-only, hence the small + edge count. """ from deepmd.pt_expt.kernels.triton.sezm.tile_configs import ( GATE_BMM_MIN_FOCUS_DIM, ) - old = os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) - try: - os.environ["DP_STACK_BWD_FUSED_INFER"] = "1" - calls, out = self._record_launches(2, GATE_BMM_MIN_FOCUS_DIM) - self.assertEqual(calls, []) - self.assertTrue(torch.isfinite(out[0]).all()) - finally: - os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) - if old is not None: - os.environ["DP_STACK_BWD_FUSED_INFER"] = old + self._register_fused((GATE_BMM_MIN_FOCUS_DIM, 2), self.FUSED_CFG) + calls, out = self._record_launches(2, GATE_BMM_MIN_FOCUS_DIM, n_edge=512) + self.assertEqual(calls, []) + self.assertTrue(torch.isfinite(out[0]).all()) - def test_non_fp32_fall_back_with_flag_on(self): + def test_non_fp32_fall_back_with_winning_entry(self): """The fused kernel assumes fp32 rounding; other dtypes keep the relay. - Scoped with the flag enabled for the same reason as above. + Scoped with a winning entry registered for the same reason as above. """ - old = os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) - try: - os.environ["DP_STACK_BWD_FUSED_INFER"] = "1" - calls, out = self._record_launches(3, 32, dtype=torch.bfloat16) - self.assertEqual(calls, []) - self.assertTrue(torch.isfinite(out[0]).all()) - finally: - os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) - if old is not None: - os.environ["DP_STACK_BWD_FUSED_INFER"] = old - - -class TestSeZMStackBwdFlagParser(unittest.TestCase): - def test_values(self): - from deepmd.pt_expt.kernels.utils import ( - use_stack_bwd_fused, - ) - - old = os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) - try: - for value, expected in ( - ("1", True), - ("true", True), - ("ON", True), - ("0", False), - ("nope", False), - ): - with self.subTest(value=value): - os.environ["DP_STACK_BWD_FUSED_INFER"] = value - self.assertEqual(use_stack_bwd_fused(), expected) - os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) - self.assertFalse(use_stack_bwd_fused()) - finally: - os.environ.pop("DP_STACK_BWD_FUSED_INFER", None) - if old is not None: - os.environ["DP_STACK_BWD_FUSED_INFER"] = old + self._register_fused((32, 3), self.FUSED_CFG) + calls, out = self._record_launches(3, 32, dtype=torch.bfloat16, n_edge=512) + self.assertEqual(calls, []) + self.assertTrue(torch.isfinite(out[0]).all()) if __name__ == "__main__": From 2e8e6f331889f0b8f6bb2af512d2694e8141ba73 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Sun, 18 Oct 2026 17:05:13 +0800 Subject: [PATCH 06/10] fix(pt): time fused sweep candidates under no_grad and assert the selection in tests The candidate benchmark ran outside the no_grad context, so the traversal's is-grad-enabled guard routed every timed call to the relay: the sweep compared relay against relay. Both the correctness call and the timing now run under no_grad. The retargeted tests ran at whatever the ambient Triton level was; at the default level 0 both table states take the relay and the comparison goes vacuous. The class now sets level 2 itself (the dispatch matrix test manages its own levels), and the bitwise test counts the fused launches -- 0 with a None entry, one per gated layer with a winning entry -- so output equality alone can no longer pass. --- .../kernels/triton/sezm/sweep_tile_configs.py | 18 ++++---- .../pt/model/test_descriptor_sezm_triton.py | 42 ++++++++++++++++++- 2 files changed, 51 insertions(+), 9 deletions(-) diff --git a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py index cf111e55f9..f6e4c6b362 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py +++ b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py @@ -2219,15 +2219,19 @@ def run_traversal() -> tuple[torch.Tensor, torch.Tensor]: for cfg in candidates: register_tile_configs("stack_bwd_fused", {key: cfg}) try: + # Both the correctness call and the timing must run under + # no_grad: the traversal selects the fused kernel only when + # grad mode is off, so timing outside this context would + # silently benchmark the relay instead of the candidate. with torch.no_grad(): got_g, got_ga = run_traversal() - if not (torch.equal(got_g, ref_g) and torch.equal(got_ga, ref_ga)): - print( - f"[stack_bwd_fused] ({cf},{lmax}) candidate " - f"{cfg} rejected: outputs differ from the relay" - ) - continue - fused_ms = _bench(run_traversal, iters=50) + if not (torch.equal(got_g, ref_g) and torch.equal(got_ga, ref_ga)): + print( + f"[stack_bwd_fused] ({cf},{lmax}) candidate " + f"{cfg} rejected: outputs differ from the relay" + ) + continue + fused_ms = _bench(run_traversal, iters=50) except Exception as err: print( f"[stack_bwd_fused] ({cf},{lmax}) candidate {cfg} " diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index 89fd4140f0..00bbfed6d6 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -2072,6 +2072,20 @@ class TestSeZMStackBwdFused(unittest.TestCase): FUSED_CFG = (32, 64, 4, 3) + def setUp(self): + # The numerical and fallback tests need the level-2 selection + # active: at the default level 0 the traversal silently takes the + # relay for both table states and the comparisons go vacuous. + # (The dispatch matrix test manages the level itself.) + self._saved_level = os.environ.get("DP_TRITON_INFER") + os.environ["DP_TRITON_INFER"] = "2" + + def tearDown(self): + if self._saved_level is None: + os.environ.pop("DP_TRITON_INFER", None) + else: + os.environ["DP_TRITON_INFER"] = self._saved_level + def _register_fused(self, key, entry): """Register a fused table entry and restore the prior state.""" from deepmd.pt_expt.kernels.triton.sezm.tile_configs import ( @@ -2164,6 +2178,21 @@ def _run_traversal(self, inputs, lmax, focus_dim, apply_alpha): torch.cuda.synchronize() return g, ga + def _recorded_traversal(self, inputs, lmax, focus_dim, apply_alpha): + """Run the traversal while counting fused launches; return outputs.""" + import deepmd.pt_expt.kernels.triton.sezm.so2_value_path as vp + + calls = [] + original = vp._launch_stack_bwd_fused + + def recorder(*args, **kwargs): + calls.append(args[8]) + return original(*args, **kwargs) + + with mock.patch.object(vp, "_launch_stack_bwd_fused", side_effect=recorder): + out = self._run_traversal(inputs, lmax, focus_dim, apply_alpha) + return out, len(calls) + def test_fused_matches_unfused_bitwise(self): generator = torch.Generator(device="cuda").manual_seed(42) for lmax, focus_dim, apply_alpha in self.CASES: @@ -2171,9 +2200,18 @@ def test_fused_matches_unfused_bitwise(self): inputs = self._bwd_inputs(lmax, focus_dim, generator) key = (focus_dim, lmax) self._register_fused(key, None) - ref = self._run_traversal(inputs, lmax, focus_dim, apply_alpha) + ref, relay_launches = self._recorded_traversal( + inputs, lmax, focus_dim, apply_alpha + ) + self.assertEqual(relay_launches, 0) self._register_fused(key, self.FUSED_CFG) - fused = self._run_traversal(inputs, lmax, focus_dim, apply_alpha) + fused, fused_launches = self._recorded_traversal( + inputs, lmax, focus_dim, apply_alpha + ) + # Counting the fused launch proves the selection happened; + # output equality alone could pass with both runs on the + # relay. + self.assertEqual(fused_launches, self.N_LAYERS - 1) self.assertTrue(torch.equal(ref[0], fused[0])) # grad_alpha is computed only when alpha was applied; the # traversal leaves it uninitialized otherwise. From 5df6813b833afd05a4d9c98dfffb798d9779f0a1 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Wed, 21 Oct 2026 20:53:07 +0800 Subject: [PATCH 07/10] fix(pt_expt): make stack_bwd_fused tuning trustworthy and dual-backend Register the SeZM tuning flow in the freeze-time autotuner registry (sezm_value_path) and share one implementation (autotune_sezm) with the PT freeze path, including the post-tuning value-path rebind, so a PT-expt export tunes the fused backward instead of silently keeping the relay. The sweep now refuses to run below Triton level 2 (the production dispatch would otherwise time the relay against itself), verifies each candidate actually engaged the fused launch before trusting its timing, and propagates unexpected internal errors instead of recording them as a measured relay win; only output mismatches and out-of-device-resource candidates count as rejections. Candidate timing follows the fp32 stack selection: short prefilter, repeated runs over rotating orders, medians, 2% win margin against the relay median. tune_missing_configs pins DP_TRITON_INFER to the requested level for the sweeps. Tests: 20001-edge and 129-edge partial-tile bitwise cross-checks, caller-option guard coverage with a winning entry, CPU policy tests for the level guard and unexpected-error handling, and routing tests proving both freeze paths reach the shared flow. The training kernel's gate-logit cotangent now uses the shared _sigmoid_bwd helper. --- deepmd/pt/entrypoints/freeze_pt2.py | 56 +--- .../kernels/triton/sezm/so2_value_path.py | 43 ++- .../kernels/triton/sezm/sweep_tile_configs.py | 213 ++++++++++++-- .../pt/model/test_descriptor_sezm_triton.py | 270 +++++++++++++++++- 4 files changed, 505 insertions(+), 77 deletions(-) diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index b3cd33798a..34ad73910a 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/deepmd/pt/entrypoints/freeze_pt2.py @@ -57,7 +57,6 @@ GeometricInitialEmbedding, ) from deepmd.pt.model.descriptor.sezm_nn.so2 import ( - SO2Convolution, SO2Linear, ) from deepmd.pt.model.model import ( @@ -331,55 +330,20 @@ def _tune_triton_configs(model: torch.nn.Module, target_device: torch.device) -> width/degree) are swept here on the local GPU -- the exact hardware the ``.pt2`` will run on, since AOTInductor artifacts are not portable across GPU models -- and registered for the current process before tracing. - Keys already covered cost nothing. - - The fused value-path entries are then rebound: the mixing-stack operator - selection (fp32 versus fp16x3) is fixed at construction time, which - predates the registrations made here. + Keys already covered cost nothing. The fused value-path entries are + then rebound (the mixing-stack operator selection, fp32 versus fp16x3, + is fixed at construction time, which predates the registrations). + + The whole flow lives in + :func:`deepmd.pt_expt.kernels.triton.sezm.sweep_tile_configs.autotune_sezm` + and is shared with the PT-expt export path, which reaches it through + the freeze-time autotuner registry. """ - if triton_infer_level() < 2: - return - if target_device.type != "cuda" or not torch.cuda.is_available(): - return - from deepmd.pt_expt.kernels.triton.sezm.so2_value_path import ( - SO2_VALUE_PATH_TRITON_AVAILABLE, - make_triton_value_path, - ) - - if not SO2_VALUE_PATH_TRITON_AVAILABLE: - return from deepmd.pt_expt.kernels.triton.sezm.sweep_tile_configs import ( - collect_model_shape_keys, - tune_missing_configs, + autotune_sezm, ) - # The built-in tables and the sweep both resolve against the current - # device; pin it to the AOTI target so a freeze aimed at a secondary GPU - # tunes and looks up the right hardware (mixed-model hosts). - if target_device.index is not None: - torch.cuda.set_device(target_device) - - shape_keys = collect_model_shape_keys(model) - registered = tune_missing_configs( - shape_keys, level=triton_infer_level(), device=target_device - ) - if registered: - log.info( - "Registered freshly tuned Triton launch configurations: %s", - {family: sorted(entries) for family, entries in registered.items()}, - ) - else: - log.info( - "Triton launch tables already cover this checkpoint's shapes on %s; " - "no tuning needed.", - torch.cuda.get_device_name(target_device), - ) - # Rebind unconditionally: the fp32-versus-fp16x3 stack selection was made - # at construction time, possibly against a different current device's - # tables, and must reflect the target device and any fresh registrations. - for module in model.modules(): - if isinstance(module, SO2Convolution) and module.triton_infer_level >= 2: - module._triton_value_path = make_triton_value_path(module) + autotune_sezm(model, triton_infer_level(), target_device) # The trace-time sendlist for the with-comm artifact embeds the address of a 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 f43e7d8eb4..a5d9ccbc2b 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py +++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py @@ -1950,8 +1950,10 @@ def _stack_train_traversal_kernel( ul_row[:, None] + (rp + nc)[None, :], un_rp - zrp * sig_g, mask=cm ) # Gate-logit gradient and its register-dot fold onto the - # scalar rows. - g_logit = (gr0 * zr0 + grn * zrn + grp_v * zrp) * sig_g * (1.0 - sig_g) + # scalar rows; the shared cotangent helper keeps the + # training kernel on the same derivative definition as the + # other stack kernels. + g_logit = _sigmoid_bwd(gr0 * zr0 + grn * zrn + grp_v * zrp, sig_g) tl.store(gq_row[:, None] + (grp * CF + nc)[None, :], g_logit, mask=cm) gz_s = tl.dot(g_logit, tl.trans(gw_g), gz_s, input_precision="ieee") if HAS_GZUP: @@ -2180,11 +2182,14 @@ def _stack_grad_alpha_kernel( mask=m_mask, ) - # Launch geometry for the fused per-layer backward below. BLOCK_M=64 / - # BLOCK_N=64 with 4 warps fills the 255-register budget and spills 226 - # slots; the spill traffic eats the fusion win, so the M dimension is - # halved instead. BLOCK_K is not free here: it must equal the channel - # group width so one K tile is exactly one row group (see the kernel). + # Launch geometry for the fused per-layer backward below. The baseline + # choice came from the register budget: BLOCK_M=64 / BLOCK_N=64 with 4 + # warps fills the 255-register budget and spills 226 slots, and the + # spill traffic eats the fusion win, so the M dimension is halved + # instead. BLOCK_M/BLOCK_N and the warp/stage counts are per-shape + # decisions of the ``stack_bwd_fused`` sweep table; BLOCK_K is not free + # here: it must equal the channel group width so one K tile is exactly + # one row group (see the kernel). @triton.jit def _stack_bwd_fused_kernel( g_ptr, # (F, E, ROW) upstream gradient of the layer output, focus-major @@ -5783,3 +5788,27 @@ def make_triton_value_path(conv: SO2Convolution) -> _TritonSO2ValuePath | None: if not SO2_VALUE_PATH_TRITON_AVAILABLE or not _is_supported(conv): return None return _TritonSO2ValuePath(conv) + + +def _freeze_autotune(model: torch.nn.Module, level: int, device: torch.device) -> None: + """Run the SeZM freeze-time launch-table tuning for ``model``. + + Thin seam for the freeze-time autotuner registry + (:mod:`deepmd.pt_expt.kernels.autotune`): the actual flow lives in + :func:`.sweep_tile_configs.autotune_sezm`, shared with the PT freeze + path. Imported lazily so the sweep module and this module stay + import-cycle-free. + """ + from deepmd.pt_expt.kernels.triton.sezm.sweep_tile_configs import ( + autotune_sezm, + ) + + autotune_sezm(model, level, device) + + +if SO2_VALUE_PATH_TRITON_AVAILABLE: + from deepmd.pt_expt.kernels.autotune import ( + register_autotuner, + ) + + register_autotuner("sezm_value_path", _freeze_autotune) diff --git a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py index f6e4c6b362..eefe31ed5e 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py +++ b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py @@ -121,6 +121,7 @@ import argparse import itertools import logging +import os import statistics from contextlib import ( contextmanager, @@ -139,6 +140,11 @@ wrap_triton, ) +from deepmd.pt_expt.kernels.utils import ( + triton_infer_level, + triton_train_level, +) + if TYPE_CHECKING: from collections.abc import ( Callable, @@ -190,6 +196,7 @@ ) __all__ = [ + "autotune_sezm", "collect_model_shape_keys", "tune_missing_configs", ] @@ -346,6 +353,27 @@ def install(entry: tuple | None) -> None: raise +@contextmanager +def _infer_level_scope(level: int) -> Iterator[None]: + """Pin ``DP_TRITON_INFER`` to the level being tuned for the sweep duration. + + The sweeps that drive a real dispatch (the fused per-layer backward) + resolve the Triton level from the ambient environment at call time, so + the tuning run must execute under the same level it tunes for -- the + level the caller already fixed through the existing ``DP_TRITON_INFER`` + mechanism. The prior value is restored on exit. + """ + saved = os.environ.get("DP_TRITON_INFER") + os.environ["DP_TRITON_INFER"] = str(level) + try: + yield + finally: + if saved is None: + os.environ.pop("DP_TRITON_INFER", None) + else: + os.environ["DP_TRITON_INFER"] = saved + + def _bench(fn: Callable[[], object], iters: int = 20) -> float: """Return the mean kernel time of ``fn`` in milliseconds.""" for _ in range(3): @@ -2142,10 +2170,22 @@ def sweep_stack_bwd_fused( ) return {"stack_bwd_fused": {key: None}} + # The fused candidate engages through the production dispatch, which + # selects it only at Triton level >= 2 (the table-tuned level); the + # sweep drives that same dispatch. At a lower ambient level every + # "candidate" timing would silently run the relay, so refuse to measure + # a comparison that never engaged the candidates. + if max(triton_infer_level(), triton_train_level()) < 2: + raise RuntimeError( + "sweep_stack_bwd_fused requires DP_TRITON_INFER>=2 (or " + "DP_TRITON_TRAIN>=2) in the environment: the production dispatch " + "engages the fused candidate only at that level, so a sweep at a " + "lower level would time the relay against itself" + ) + n_layers = 4 n_gated = n_layers - 1 m0 = (lmax + 1) * cf - m1 = 2 * lmax * cf row = (3 * lmax + 1) * cf half = lmax * cf @@ -2207,6 +2247,19 @@ def run_traversal() -> tuple[torch.Tensor, torch.Tensor]: ) runtime = _runtime_tile_configs("stack_bwd_fused") saved = dict(runtime) + # The sweep must observe each candidate actually engaging the fused + # launch (output equality alone passes with both runs on the relay), so + # a counting stand-in wraps the real launcher for the whole measurement. + from deepmd.pt_expt.kernels.triton.sezm import so2_value_path + + fused_launches = [0] + real_launch = so2_value_path._launch_stack_bwd_fused + + def counting_launch(*args: Any, **kwargs: Any) -> Any: + fused_launches[0] += 1 + return real_launch(*args, **kwargs) + + so2_value_path._launch_stack_bwd_fused = counting_launch try: # Baseline: an explicit None entry forces the relay whatever the # built-in tables say, so the comparison reflects this GPU only. @@ -2215,7 +2268,13 @@ def run_traversal() -> tuple[torch.Tensor, torch.Tensor]: ref_g, ref_ga = run_traversal() relay_ms = _bench(run_traversal, iters=50) - best_ms, best_cfg = float("inf"), None + # Prefilter: bitwise equality plus a short timing per candidate. + # Only known rejections disqualify a candidate -- an output + # mismatch, or the launch not fitting this device. Any other + # failure is an internal error of the sweep itself and propagates + # (with the key left uncovered) instead of being recorded as a + # measured "relay won" conclusion. + ranked: list[tuple[float, tuple[int, int, int, int]]] = [] for cfg in candidates: register_tile_configs("stack_bwd_fused", {key: cfg}) try: @@ -2224,41 +2283,73 @@ def run_traversal() -> tuple[torch.Tensor, torch.Tensor]: # grad mode is off, so timing outside this context would # silently benchmark the relay instead of the candidate. with torch.no_grad(): + fused_launches[0] = 0 got_g, got_ga = run_traversal() + if fused_launches[0] == 0: + raise RuntimeError( + f"stack_bwd_fused sweep: candidate {cfg} for " + f"({cf},{lmax}) never engaged the fused launch" + ) if not (torch.equal(got_g, ref_g) and torch.equal(got_ga, ref_ga)): print( f"[stack_bwd_fused] ({cf},{lmax}) candidate " f"{cfg} rejected: outputs differ from the relay" ) continue - fused_ms = _bench(run_traversal, iters=50) - except Exception as err: + fused_ms = _bench(run_traversal, iters=6) + except (torch.cuda.OutOfMemoryError, triton.runtime.errors.OutOfResources): print( f"[stack_bwd_fused] ({cf},{lmax}) candidate {cfg} " - f"rejected: {type(err).__name__}: {err}" + "rejected: not launchable on this device" ) continue print( f"[stack_bwd_fused] ({cf},{lmax}) candidate {cfg}: " f"{fused_ms:.3f} ms (relay {relay_ms:.3f} ms)" ) - if fused_ms < best_ms: - best_ms, best_cfg = fused_ms, cfg + ranked.append((fused_ms, cfg)) + + # Repeat timing over rotating orders; medians decide. This follows + # the fp32 stack selection (_select_fp32_config): a single + # relay-first sequence lets clock and thermal drift favor whatever + # runs last, which the 2% margin below cannot offset. + finalists = [cfg for _ms, cfg in sorted(ranked)[:3]] + samples: dict[tuple[int, int, int, int] | None, list[float]] = { + entry: [] for entry in [*finalists, None] + } + if finalists: + order: list[tuple[int, int, int, int] | None] = [*finalists, None] + for shift in range(3): + rotated = order[shift % len(order) :] + order[: shift % len(order)] + for entry in rotated: + register_tile_configs("stack_bwd_fused", {key: entry}) + with torch.no_grad(): + samples[entry].append(_bench(run_traversal, iters=20)) + else: + samples[None] = [relay_ms] finally: + so2_value_path._launch_stack_bwd_fused = real_launch runtime.clear() runtime.update(saved) - # A 2% margin keeps run-to-run noise from registering a coin-flip as - # a winner; anything below it records the relay. - if best_cfg is not None and best_ms < relay_ms * 0.98: + # A 2% margin over the relay median keeps run-to-run noise from + # registering a coin-flip as a winner; anything below it records the + # relay. + relay_med = statistics.median(samples[None]) + best_ms, best_cfg = float("inf"), None + for cfg in finalists: + cfg_ms = statistics.median(samples[cfg]) + if cfg_ms < best_ms: + best_ms, best_cfg = cfg_ms, cfg + if best_cfg is not None and best_ms < relay_med * 0.98: winner: tuple | None = best_cfg print( f"[stack_bwd_fused] ({cf},{lmax}) fused {winner} wins: " - f"{best_ms:.3f} vs {relay_ms:.3f} ms" + f"{best_ms:.3f} vs {relay_med:.3f} ms (relay median)" ) else: winner = None - print(f"[stack_bwd_fused] ({cf},{lmax}) relay wins: {relay_ms:.3f} ms") + print(f"[stack_bwd_fused] ({cf},{lmax}) relay wins: {relay_med:.3f} ms") register_tile_configs("stack_bwd_fused", {key: winner}) return {"stack_bwd_fused": {key: winner}} @@ -2423,6 +2514,9 @@ def tune_missing_configs( Shapes as returned by :func:`collect_model_shape_keys`. level : int The active Triton inference level; levels below 2 tune nothing. + Sweeps that drive a real dispatch execute with ``DP_TRITON_INFER`` + pinned to this level (see :func:`_infer_level_scope`), so candidate + selection during tuning matches the level being tuned for. device : torch.device or str CUDA device to sweep on. n_edge : int or None @@ -2466,15 +2560,16 @@ def tune_missing_configs( [group for group, _spec in pending], ) for _group, spec in pending: - result = _run_sweep( - spec, - cf, - lmax, - n_focus=n_focus, - n_head=n_head, - n_edge=n_edge, - device=device, - ) + with _infer_level_scope(level): + result = _run_sweep( + spec, + cf, + lmax, + n_focus=n_focus, + n_head=n_head, + n_edge=n_edge, + device=device, + ) for family, entries in result.items(): register_tile_configs(family, entries) registered.setdefault(family, {}).update(entries) @@ -2485,6 +2580,82 @@ def tune_missing_configs( return registered +def autotune_sezm( + model: torch.nn.Module, + level: int, + device: torch.device, +) -> SweepResult: + """Tune and register the SeZM launch tables for a model about to be frozen. + + Shared tuning flow behind both freeze paths: the PT ``.pt2`` freeze + calls this directly, and the PT-expt export reaches it through the + freeze-time autotuner registry (registered at the bottom of + :mod:`.so2_value_path`, driven by + :func:`deepmd.pt_expt.kernels.autotune.run_autotune`). Sweeps every + shape key the model's fused-value-path convolutions need that the + tables do not yet cover, then rebinds the per-module value paths: the + fp32-versus-fp16x3 stack selection is fixed at construction time, + which predates the registrations made here. + + Parameters + ---------- + model : torch.nn.Module + The model about to be frozen; its convolutions fix the shape keys. + level : int + The active Triton inference level; levels below 2 tune nothing. + device : torch.device + The target CUDA device the artifact will run on. + + Returns + ------- + SweepResult + The entries registered by this call, empty when every key was + already covered. + """ + if level < 2: + return {} + if device.type != "cuda" or not torch.cuda.is_available(): + return {} + from deepmd.pt_expt.kernels.triton.sezm.so2_value_path import ( + SO2_VALUE_PATH_TRITON_AVAILABLE, + make_triton_value_path, + ) + + if not SO2_VALUE_PATH_TRITON_AVAILABLE: + return {} + from deepmd.pt.model.descriptor.sezm_nn.so2 import ( + SO2Convolution, + ) + + # The built-in tables and the sweep both resolve against the current + # device; pin it to the freeze target so a freeze aimed at a secondary + # GPU tunes and looks up the right hardware (mixed-model hosts). + if device.index is not None: + torch.cuda.set_device(device) + + shape_keys = collect_model_shape_keys(model) + registered = tune_missing_configs(shape_keys, level=level, device=device) + if registered: + log.info( + "Registered freshly tuned Triton launch configurations: %s", + {family: sorted(entries) for family, entries in registered.items()}, + ) + else: + log.info( + "Triton launch tables already cover this checkpoint's shapes on %s; " + "no tuning needed.", + torch.cuda.get_device_name(device), + ) + # Rebind unconditionally: the fp32-versus-fp16x3 stack selection was + # made at construction time, possibly against a different current + # device's tables, and must reflect the target device and any fresh + # registrations. + for module in model.modules(): + if isinstance(module, SO2Convolution) and module.triton_infer_level >= 2: + module._triton_value_path = make_triton_value_path(module) + return registered + + def _format_data_fragment(entries: SweepResult) -> str: """Render registered entries as a :mod:`.tile_config_data` fragment.""" if not torch.cuda.is_available(): diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index 00bbfed6d6..fd630f3ada 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -2103,7 +2103,7 @@ def restore(): register_tile_configs("stack_bwd_fused", {key: entry}) self.addCleanup(restore) - N_EDGE = 20000 + N_EDGE = 20001 N_FOCUS = 2 N_LAYERS = 4 # one final linear layer plus three gated layers CASES = ((1, 32, False), (3, 32, False), (3, 32, True)) # lmax, Cf, apply_alpha @@ -2180,7 +2180,7 @@ def _run_traversal(self, inputs, lmax, focus_dim, apply_alpha): def _recorded_traversal(self, inputs, lmax, focus_dim, apply_alpha): """Run the traversal while counting fused launches; return outputs.""" - import deepmd.pt_expt.kernels.triton.sezm.so2_value_path as vp + from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp calls = [] original = vp._launch_stack_bwd_fused @@ -2218,6 +2218,99 @@ def test_fused_matches_unfused_bitwise(self): if apply_alpha: self.assertTrue(torch.equal(ref[1], fused[1])) + def test_fused_matches_unfused_partial_tile(self): + """A partial final row tile matches the relay bit-for-bit too. + + Both the large fixture above and the dispatch fixtures use edge + counts that are multiples of the candidate BLOCK_M = 32, so the + fused kernel's row-tail mask was never exercised numerically; + 129 = 4 * 32 + 1 leaves one partial row tile. + """ + saved_edges = self.N_EDGE + type(self).N_EDGE = 129 + try: + generator = torch.Generator(device="cuda").manual_seed(5) + inputs = self._bwd_inputs(1, 32, generator) + finally: + type(self).N_EDGE = saved_edges + key = (32, 1) + self._register_fused(key, None) + ref, relay_launches = self._recorded_traversal(inputs, 1, 32, False) + self.assertEqual(relay_launches, 0) + self._register_fused(key, self.FUSED_CFG) + fused, fused_launches = self._recorded_traversal(inputs, 1, 32, False) + self.assertEqual(fused_launches, self.N_LAYERS - 1) + self.assertTrue(torch.equal(ref[0], fused[0])) + + def test_caller_option_guards_fall_back_with_winning_entry(self): + """The caller-option guards keep the relay with an entry present. + + ``with_weights``, ``keep``, and an upstream cotangent request the + per-layer surfaces or weight contractions the fused kernel does not + produce, and a focus width of 48 is neither a power of two nor equal + to the shared backward GEMM's block width, breaking the row-group + tiling contract. Each keeps the relay with a winning entry + registered. Dispatch-only, hence the small edge count. + """ + from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp + + for name, lmax, cf, overrides, with_upstream, u_focus_major in ( + ("with_weights", 1, 32, {"with_weights": True}, False, True), + ("keep", 1, 32, {"keep": True}, False, False), + ("grad_z_upstream", 1, 32, {}, True, False), + ("non_pow2_cf", 1, 48, {}, False, False), + ): + with self.subTest(guard=name): + call_kwargs = {"with_weights": False, "keep": False} + call_kwargs.update(overrides) + self._register_fused((cf, lmax), self.FUSED_CFG) + saved_edges = self.N_EDGE + type(self).N_EDGE = 512 + try: + generator = torch.Generator(device="cuda").manual_seed(13) + inputs = self._bwd_inputs(lmax, cf, generator) + finally: + type(self).N_EDGE = saved_edges + # The relay's weight-contraction fallback expects the final + # activation focus-major, unlike the edge-major output the + # forward returns; this check exercises the dispatch guard, + # so the fixture follows that branch's contract. + u_final = inputs[4] + if u_focus_major: + u_final = u_final.permute(1, 0, 2).contiguous() + grad_z_upstream = torch.randn_like(inputs[3]) if with_upstream else None + calls = [] + original = vp._launch_stack_bwd_fused + + def recorder(*args, **kwargs_): + calls.append(args[8]) # the layer index + return original(*args, **kwargs_) + + with mock.patch.object( + vp, "_launch_stack_bwd_fused", side_effect=recorder + ): + with torch.no_grad(): + g, _ga, _weights, _state = vp._stack_backward_traversal( + inputs[0], + inputs[1], + inputs[3], + u_final, + inputs[2], + inputs[5], + inputs[6], + inputs[7], + inputs[8], + grad_z_upstream, + None, + lmax, + cf, + False, + **call_kwargs, + ) + torch.cuda.synchronize() + self.assertEqual(calls, []) + self.assertTrue(torch.isfinite(g).all()) + def test_higher_order_backward_preserved_with_winning_entry(self): """``create_graph=True`` keeps the supported path with an entry present. @@ -2288,7 +2381,7 @@ def _record_launches( self, lmax, focus_dim, apply_alpha=False, dtype=torch.float32, n_edge=None ): """Run the traversal with a recording stand-in for the fused launch.""" - import deepmd.pt_expt.kernels.triton.sezm.so2_value_path as vp + from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp calls = [] original = vp._launch_stack_bwd_fused @@ -2361,5 +2454,176 @@ def test_non_fp32_fall_back_with_winning_entry(self): self.assertTrue(torch.isfinite(out[0]).all()) +class TestSeZMStackBwdSweepPolicy(unittest.TestCase): + """CPU checks of the ``stack_bwd_fused`` sweep's measurement policy. + + The sweep drives the real traversal, so a misconfigured level or an + internal failure must be refused loudly rather than recorded as a + measured conclusion. The traversal and the timer are replaced by + doubles here; these checks verify the policy, not the kernel. + """ + + KEY = (32, 1) + + def setUp(self): + self._saved_level = os.environ.get("DP_TRITON_INFER") + os.environ["DP_TRITON_INFER"] = "2" + from deepmd.pt_expt.kernels.triton.sezm.tile_configs import ( + _runtime_tile_configs, + ) + + runtime = _runtime_tile_configs("stack_bwd_fused") + self._saved_table = dict(runtime) + runtime.pop(self.KEY, None) + + def restore_table(): + runtime.clear() + runtime.update(self._saved_table) + + self.addCleanup(restore_table) + + def tearDown(self): + if self._saved_level is None: + os.environ.pop("DP_TRITON_INFER", None) + else: + os.environ["DP_TRITON_INFER"] = self._saved_level + + def test_sweep_refuses_levels_below_two(self): + """A sweep at a lower level would time the relay against itself.""" + from deepmd.pt_expt.kernels.triton.sezm.sweep_tile_configs import ( + sweep_stack_bwd_fused, + ) + + os.environ["DP_TRITON_INFER"] = "0" + with self.assertRaisesRegex(RuntimeError, "DP_TRITON_INFER"): + sweep_stack_bwd_fused(32, 1, n_edge=129, device="cpu") + + def _stub_sweep(self, *, fail_candidates): + """Run the sweep against traversal/timer doubles; return its result.""" + from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp + from deepmd.pt_expt.kernels.triton.sezm import sweep_tile_configs as sweep_mod + from deepmd.pt_expt.kernels.triton.sezm.tile_configs import ( + _runtime_tile_configs, + ) + + outputs = ( + torch.zeros(2, 3, dtype=torch.float32, device="cpu"), + torch.zeros(2, 3, dtype=torch.float32, device="cpu"), + None, + None, + ) + + def stub_traversal(*args, **kwargs): + entry = _runtime_tile_configs("stack_bwd_fused").get(self.KEY, "missing") + if entry is None: + return tuple(t.clone() if torch.is_tensor(t) else t for t in outputs) + # The dispatch would launch the fused kernel here; calling the + # (doubled) launcher stands in for that engagement. + vp._launch_stack_bwd_fused() + if fail_candidates: + raise AssertionError("injected internal failure") + return tuple(t.clone() if torch.is_tensor(t) else t for t in outputs) + + def fake_bench(fn, iters=20): + entry = _runtime_tile_configs("stack_bwd_fused").get(self.KEY, "missing") + return 2.0 if entry is None else 1.0 + + with ( + mock.patch.object(vp, "_stack_backward_traversal", stub_traversal), + mock.patch.object(vp, "_launch_stack_bwd_fused", lambda *a, **k: None), + mock.patch.object(sweep_mod, "_bench", fake_bench), + ): + return sweep_mod.sweep_stack_bwd_fused(32, 1, n_edge=129, device="cpu") + + def test_unexpected_candidate_failure_aborts_and_leaves_key_uncovered(self): + """An internal error must not become a recorded relay win.""" + from deepmd.pt_expt.kernels.triton.sezm.tile_configs import ( + has_tile_config, + ) + + with self.assertRaisesRegex(AssertionError, "injected internal failure"): + self._stub_sweep(fail_candidates=True) + self.assertFalse(has_tile_config("stack_bwd_fused", self.KEY)) + + def test_engaged_fast_candidate_registers_a_winner(self): + """Positive control: a measured winner is registered and covered.""" + from deepmd.pt_expt.kernels.triton.sezm.tile_configs import ( + _runtime_tile_configs, + ) + + result = self._stub_sweep(fail_candidates=False) + # Every candidate times at 1.0 ms against the relay's 2.0 ms, so + # the fastest-by-tiebreak candidate wins and the key is covered. + self.assertEqual(result["stack_bwd_fused"][self.KEY], (16, 64, 4, 2)) + self.assertEqual( + _runtime_tile_configs("stack_bwd_fused")[self.KEY], (16, 64, 4, 2) + ) + + +class TestSeZMFreezeAutotuneRouting(unittest.TestCase): + """Both freeze paths must reach the shared SeZM tuning flow. + + The PT ``.pt2`` freeze calls ``autotune_sezm`` through its own seam and + the PT-expt export reaches it through the freeze-time autotuner + registry; routing only, so the shared flow itself is replaced by a + recording double and CUDA availability is mocked. + """ + + def _record_autotune(self): + from deepmd.pt_expt.kernels.triton.sezm import sweep_tile_configs + + calls = [] + + def fake_autotune(model, level, device): + calls.append((model, level, device)) + + return calls, mock.patch.object( + sweep_tile_configs, "autotune_sezm", fake_autotune + ) + + def test_pt_freeze_seam_reaches_shared_flow(self): + import deepmd.pt.entrypoints.freeze_pt2 as freeze_pt2 + + saved_level = os.environ.get("DP_TRITON_INFER") + os.environ["DP_TRITON_INFER"] = "2" + try: + calls, patched = self._record_autotune() + model = torch.nn.Module() + device = torch.device("cuda", 0) + with patched: + freeze_pt2._tune_triton_configs(model, device) + self.assertEqual(calls, [(model, 2, device)]) + finally: + if saved_level is None: + os.environ.pop("DP_TRITON_INFER", None) + else: + os.environ["DP_TRITON_INFER"] = saved_level + + def test_pt_expt_registry_reaches_shared_flow(self): + from deepmd.pt_expt.kernels.autotune import ( + _AUTOTUNERS, + run_autotune, + ) + + self.assertIn("sezm_value_path", _AUTOTUNERS) + saved_level = os.environ.get("DP_TRITON_INFER") + os.environ["DP_TRITON_INFER"] = "2" + try: + calls, patched = self._record_autotune() + model = torch.nn.Module() + device = torch.device("cuda", 0) + with ( + patched, + mock.patch("torch.cuda.is_available", return_value=True), + ): + run_autotune(model, device) + self.assertEqual(calls, [(model, 2, device)]) + finally: + if saved_level is None: + os.environ.pop("DP_TRITON_INFER", None) + else: + os.environ["DP_TRITON_INFER"] = saved_level + + if __name__ == "__main__": unittest.main() From ef2eaf71b4574994c67a0778799db61fa29a72d4 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 12:54:17 +0000 Subject: [PATCH 08/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/sweep_tile_configs.py | 4 +++- source/tests/pt/model/test_descriptor_sezm_triton.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py index eefe31ed5e..24a7d98974 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py +++ b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py @@ -2250,7 +2250,9 @@ def run_traversal() -> tuple[torch.Tensor, torch.Tensor]: # The sweep must observe each candidate actually engaging the fused # launch (output equality alone passes with both runs on the relay), so # a counting stand-in wraps the real launcher for the whole measurement. - from deepmd.pt_expt.kernels.triton.sezm import so2_value_path + from deepmd.pt_expt.kernels.triton.sezm import ( + so2_value_path, + ) fused_launches = [0] real_launch = so2_value_path._launch_stack_bwd_fused diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index fd630f3ada..828d519710 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -2570,7 +2570,9 @@ class TestSeZMFreezeAutotuneRouting(unittest.TestCase): """ def _record_autotune(self): - from deepmd.pt_expt.kernels.triton.sezm import sweep_tile_configs + from deepmd.pt_expt.kernels.triton.sezm import ( + sweep_tile_configs, + ) calls = [] From f5240628921b1a6d1fc4e01dd1ffc6236a123078 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Thu, 22 Oct 2026 19:24:00 +0800 Subject: [PATCH 09/10] fix(pt): keep the freeze tuning seam import-safe and backend-neutral The optional-Triton guards (inference level, CUDA target, value-path availability) return before the sweep-module import again: that module imports Triton unconditionally, and the previous delegation broke CPU-only / Triton-less freezes with a ModuleNotFoundError even at the default level. A regression test blocks the sweep module in sys.modules and covers all four skip cases through the real freeze seam. The post-tuning rebind probes the value-path interface instead of an isinstance against the PT-side SO2Convolution class, so the PT-expt descriptor modules rebind as well. The stack sweep's level-guard message no longer suggests the nonexistent DP_TRITON_TRAIN>=2 setting, and the new CPU test classes are gated on Triton availability like the rest of the file. --- deepmd/pt/entrypoints/freeze_pt2.py | 16 +++- .../kernels/triton/sezm/sweep_tile_configs.py | 24 +++--- .../pt/model/test_descriptor_sezm_triton.py | 73 ++++++++++++++++++- 3 files changed, 101 insertions(+), 12 deletions(-) diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index 34ad73910a..3e8e8c4de3 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/deepmd/pt/entrypoints/freeze_pt2.py @@ -337,8 +337,22 @@ def _tune_triton_configs(model: torch.nn.Module, target_device: torch.device) -> The whole flow lives in :func:`deepmd.pt_expt.kernels.triton.sezm.sweep_tile_configs.autotune_sezm` and is shared with the PT-expt export path, which reaches it through - the freeze-time autotuner registry. + the freeze-time autotuner registry. The early returns below must stay + ahead of the sweep-module import: that module imports Triton + unconditionally, and this helper runs on every freeze -- a CPU-only or + otherwise Triton-less installation must skip tuning (as it always did) + instead of failing the import. """ + if triton_infer_level() < 2: + return + if target_device.type != "cuda" or not torch.cuda.is_available(): + return + from deepmd.pt_expt.kernels.triton.sezm.so2_value_path import ( + SO2_VALUE_PATH_TRITON_AVAILABLE, + ) + + if not SO2_VALUE_PATH_TRITON_AVAILABLE: + return from deepmd.pt_expt.kernels.triton.sezm.sweep_tile_configs import ( autotune_sezm, ) diff --git a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py index 24a7d98974..2b4d926c55 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py +++ b/deepmd/pt_expt/kernels/triton/sezm/sweep_tile_configs.py @@ -2174,13 +2174,15 @@ def sweep_stack_bwd_fused( # selects it only at Triton level >= 2 (the table-tuned level); the # sweep drives that same dispatch. At a lower ambient level every # "candidate" timing would silently run the relay, so refuse to measure - # a comparison that never engaged the candidates. + # a comparison that never engaged the candidates. The training level + # tops out at 1, so in practice only DP_TRITON_INFER>=2 satisfies the + # dispatch predicate the sweep mirrors. if max(triton_infer_level(), triton_train_level()) < 2: raise RuntimeError( - "sweep_stack_bwd_fused requires DP_TRITON_INFER>=2 (or " - "DP_TRITON_TRAIN>=2) in the environment: the production dispatch " - "engages the fused candidate only at that level, so a sweep at a " - "lower level would time the relay against itself" + "sweep_stack_bwd_fused requires DP_TRITON_INFER>=2 in the " + "environment: the production dispatch engages the fused " + "candidate only at that level, so a sweep at a lower level " + "would time the relay against itself" ) n_layers = 4 @@ -2625,9 +2627,6 @@ def autotune_sezm( if not SO2_VALUE_PATH_TRITON_AVAILABLE: return {} - from deepmd.pt.model.descriptor.sezm_nn.so2 import ( - SO2Convolution, - ) # The built-in tables and the sweep both resolve against the current # device; pin it to the freeze target so a freeze aimed at a secondary @@ -2651,9 +2650,14 @@ def autotune_sezm( # Rebind unconditionally: the fp32-versus-fp16x3 stack selection was # made at construction time, possibly against a different current # device's tables, and must reflect the target device and any fresh - # registrations. + # registrations. The interface probe (not an isinstance against a + # backend-specific SO2Convolution class) keeps the rebind effective + # for both the PT and the PT-expt descriptor modules. for module in model.modules(): - if isinstance(module, SO2Convolution) and module.triton_infer_level >= 2: + if ( + _has_so2_value_path_interface(module) + and getattr(module, "triton_infer_level", 0) >= 2 + ): module._triton_value_path = make_triton_value_path(module) return registered diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index 828d519710..b5ded2d57c 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -22,6 +22,7 @@ lowering the resulting graph through AOTAutograd's forward-only compiler. """ +import contextlib import math import os import typing @@ -2454,6 +2455,7 @@ def test_non_fp32_fall_back_with_winning_entry(self): self.assertTrue(torch.isfinite(out[0]).all()) +@unittest.skipUnless(TRITON_AVAILABLE, "Triton is required for the sweep module") class TestSeZMStackBwdSweepPolicy(unittest.TestCase): """CPU checks of the ``stack_bwd_fused`` sweep's measurement policy. @@ -2560,6 +2562,7 @@ def test_engaged_fast_candidate_registers_a_winner(self): ) +@unittest.skipUnless(TRITON_AVAILABLE, "Triton is required for the sweep module") class TestSeZMFreezeAutotuneRouting(unittest.TestCase): """Both freeze paths must reach the shared SeZM tuning flow. @@ -2592,7 +2595,7 @@ def test_pt_freeze_seam_reaches_shared_flow(self): calls, patched = self._record_autotune() model = torch.nn.Module() device = torch.device("cuda", 0) - with patched: + with patched, mock.patch("torch.cuda.is_available", return_value=True): freeze_pt2._tune_triton_configs(model, device) self.assertEqual(calls, [(model, 2, device)]) finally: @@ -2601,6 +2604,74 @@ def test_pt_freeze_seam_reaches_shared_flow(self): else: os.environ["DP_TRITON_INFER"] = saved_level + def test_pt_freeze_seam_skips_before_sweep_import(self): + """Untuned/Triton-less freezes must skip, not import the sweep module. + + The sweep module imports Triton unconditionally, so the freeze seam + must return at its level/device/availability guards before that + import; blocking the module in ``sys.modules`` makes any premature + import fail loudly. + """ + import sys + + import deepmd.pt.entrypoints.freeze_pt2 as freeze_pt2 + from deepmd.pt_expt.kernels.triton.sezm import ( + so2_value_path, + ) + + blocked = { + "deepmd.pt_expt.kernels.triton.sezm.sweep_tile_configs": None, + } + saved_level = os.environ.get("DP_TRITON_INFER") + try: + cases = ( + ("default-level cpu", {}, torch.device("cpu"), None), + ("level-2 cpu", {"DP_TRITON_INFER": "2"}, torch.device("cpu"), None), + ( + "level-2 cuda-unavailable", + {"DP_TRITON_INFER": "2"}, + torch.device("cuda", 0), + False, + ), + ( + "level-2 triton-unavailable", + {"DP_TRITON_INFER": "2"}, + torch.device("cuda", 0), + True, + ), + ) + for name, env, device, cuda_available in cases: + with self.subTest(case=name): + if env: + os.environ["DP_TRITON_INFER"] = env["DP_TRITON_INFER"] + else: + os.environ.pop("DP_TRITON_INFER", None) + patches = [mock.patch.dict(sys.modules, blocked)] + if cuda_available is not None: + patches.append( + mock.patch( + "torch.cuda.is_available", + return_value=cuda_available, + ) + ) + if name == "level-2 triton-unavailable": + patches.append( + mock.patch.object( + so2_value_path, + "SO2_VALUE_PATH_TRITON_AVAILABLE", + False, + ) + ) + with contextlib.ExitStack() as stack: + for patch_ in patches: + stack.enter_context(patch_) + freeze_pt2._tune_triton_configs(torch.nn.Module(), device) + finally: + if saved_level is None: + os.environ.pop("DP_TRITON_INFER", None) + else: + os.environ["DP_TRITON_INFER"] = saved_level + def test_pt_expt_registry_reaches_shared_flow(self): from deepmd.pt_expt.kernels.autotune import ( _AUTOTUNERS, From 8a61fbb2cc55b01927279a2d6b1aa0caad3477ab Mon Sep 17 00:00:00 2001 From: longjunhao Date: Sun, 25 Oct 2026 13:11:44 +0800 Subject: [PATCH 10/10] fix(pt_expt): give the fused stack backward a non-aliasing output buffer The fused per-layer kernel reads the whole g row from every program (the scalar chain plus every gate row group) while writing only its own column tile, so unlike the relay GEMM it cannot run in place: after the relay-path buffer sharing lands on master, the shared spare buffer aliases g_cur for exactly the configurations this branch engages, and a textually clean rebase would silently turn every tuned lmax>=3 launch into a cross-program race. The fused branch now keeps its own ping-pong buffer, with an identity check as a belt against future refactors, and a recorder-based test asserts the launch input and output never alias. --- .../kernels/triton/sezm/so2_value_path.py | 26 ++++++++++++++-- .../pt/model/test_descriptor_sezm_triton.py | 31 +++++++++++++++++++ 2 files changed, 54 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 a5d9ccbc2b..e692312746 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py +++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py @@ -3590,6 +3590,18 @@ def _stack_backward_traversal( else None ) fused = stack_bwd_fused_cfg is not None + # The fused per-layer kernel reads the whole ``g`` row from every program + # (the scalar chain plus every gate row group) while writing only its own + # column tile, so unlike the relay GEMM it cannot run in place: its output + # buffer must never alias ``g_cur``. The shared spare buffer may alias + # ``g_cur`` on the buffer-shared relay path (whose GEMM reads its matrix + # input from the separate ``gz`` scratch and touches only its own output + # tile), so the fused branch keeps its own ping-pong buffer instead. + g_fused_spare: torch.Tensor | None = ( + torch.empty((n_focus, n_edge, row), device=device, dtype=dtype) + if fused and n_gated > 0 + else None + ) sig = ( torch.empty((n_focus, n_edge, gate_width), device=device, dtype=torch.float32) if not fused @@ -3632,8 +3644,16 @@ def _stack_backward_traversal( # Single-pass per-layer backward: the gate sigmoids and the # pre-activation gradient stay in registers, so neither the # sigmoid surface nor the gz scratch buffer exists. - assert g_spare is not None - g_next = g_spare + assert g_fused_spare is not None + # Anti-alias belt: the launch reads whole rows of g_cur while + # writing its own column tile of g_next, so a future + # buffer-sharing refactor must not silently reintroduce an + # in-place launch here (a cross-program race). + if g_fused_spare is g_cur: + g_fused_spare = torch.empty( + (n_focus, n_edge, row), device=device, dtype=dtype + ) + g_next = g_fused_spare _launch_stack_bwd_fused( g_cur, z_all, @@ -3650,7 +3670,7 @@ def _stack_backward_traversal( n_focus, stack_bwd_fused_cfg, ) - g_spare = g_cur + g_fused_spare = g_cur g_cur = g_next continue if keep: diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index b5ded2d57c..16b20104bd 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -2243,6 +2243,37 @@ def test_fused_matches_unfused_partial_tile(self): self.assertEqual(fused_launches, self.N_LAYERS - 1) self.assertTrue(torch.equal(ref[0], fused[0])) + def test_fused_launch_never_aliases_its_input(self): + """The fused kernel reads whole rows while writing column tiles. + + Each program reads the full ``g`` row (scalar chain plus every gate + row group) and writes only its own column tile, so an in-place + launch (``g_next is g_cur``, e.g. after a buffer-sharing + optimization on the relay path) would race across programs. A + recorder captures the input/output pair of every fused launch and + asserts they are distinct tensors throughout the traversal. + """ + from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp + + lmax, focus_dim = 3, 32 + generator = torch.Generator(device="cuda").manual_seed(31) + inputs = self._bwd_inputs(lmax, focus_dim, generator) + self._register_fused((focus_dim, lmax), self.FUSED_CFG) + pairs = [] + original = vp._launch_stack_bwd_fused + + def recorder(*args, **kwargs): + pairs.append((args[0], args[6])) # (g_cur, g_next) + return original(*args, **kwargs) + + with mock.patch.object(vp, "_launch_stack_bwd_fused", side_effect=recorder): + with torch.no_grad(): + self._run_traversal(inputs, lmax, focus_dim, True) + torch.cuda.synchronize() + self.assertEqual(len(pairs), self.N_LAYERS - 1) + for g_in, g_out in pairs: + self.assertIsNot(g_out, g_in) + def test_caller_option_guards_fall_back_with_winning_entry(self): """The caller-option guards keep the relay with an entry present.