diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index b3cd33798a..3e8e8c4de3 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,11 +330,18 @@ 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. 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 @@ -343,43 +349,15 @@ def _tune_triton_configs(model: torch.nn.Module, target_device: torch.device) -> 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 7a8ac3fd43..e692312746 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,10 @@ wrap_triton, ) +from ...utils import ( + triton_infer_level, + triton_train_level, +) from .gated_activation import ( gated_activation_second_order, gated_activation_second_order_reference, @@ -126,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, @@ -500,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) @@ -1468,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: @@ -1537,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: @@ -1846,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( @@ -1922,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: @@ -2152,6 +2182,202 @@ def _stack_grad_alpha_kernel( mask=m_mask, ) + # 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 + 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 + ) + 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, :], + 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 = _sigmoid_bwd(g_sig, 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 +2689,61 @@ 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, + config: tuple[int, int, int, 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 = 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 +3557,56 @@ 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. + # 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 + # 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 + else None + ) grad_z_all = ( torch.empty((n_gated, n_focus, n_edge, row), device=device, dtype=dtype) if keep @@ -3284,10 +3614,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 +3640,39 @@ 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_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, + gw_all, + gwt_all, + w0t_all, + w1t_all, + g_next, + n_edge, + layer, + lmax, + focus_dim, + block_k, + n_focus, + stack_bwd_fused_cfg, + ) + g_fused_spare = g_cur + g_cur = g_next + continue if keep: assert grad_z_all is not None gz = grad_z_all[layer] @@ -5446,3 +5808,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 44c0a91aac..2b4d926c55 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): @@ -2074,6 +2102,262 @@ 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}} + + # 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. 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 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 + 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) + # 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. + 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) + + # 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: + # 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(): + 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=6) + except (torch.cuda.OutOfMemoryError, triton.runtime.errors.OutOfResources): + print( + f"[stack_bwd_fused] ({cf},{lmax}) candidate {cfg} " + "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)" + ) + 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 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_med:.3f} ms (relay median)" + ) + else: + winner = None + 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}} + + @dataclass(frozen=True) class _SweepSpec: """Describe one independently covered launch-configuration sweep.""" @@ -2103,6 +2387,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"), } @@ -2233,6 +2518,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 @@ -2276,15 +2564,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) @@ -2295,6 +2584,84 @@ 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 {} + + # 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. 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 ( + _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 + + 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/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/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index f359f073d6..16b20104bd 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -22,7 +22,9 @@ lowering the resulting graph through AOTAutograd's forward-only compiler. """ +import contextlib import math +import os import typing import unittest from unittest import ( @@ -1950,6 +1952,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]) @@ -1972,6 +1977,7 @@ def fake_flash_sweep(cf, lmax, **kwargs): "point_recompute", "rotate_bwd", "rotate_fwd", + "stack_bwd_fused", ], ) self.assertEqual( @@ -1982,6 +1988,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", ], @@ -2051,5 +2058,676 @@ 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. 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 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 ( + _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 = 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 + + 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 _recorded_traversal(self, inputs, lmax, focus_dim, apply_alpha): + """Run the traversal while counting fused launches; return outputs.""" + from deepmd.pt_expt.kernels.triton.sezm import 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: + with self.subTest(lmax=lmax, cf=focus_dim, alpha=apply_alpha): + inputs = self._bwd_inputs(lmax, focus_dim, generator) + key = (focus_dim, lmax) + self._register_fused(key, None) + 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, 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. + 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_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. + + ``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. + + 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 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 ( + _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 = {} + 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, n_edge=None + ): + """Run the traversal with a recording stand-in for the fused launch.""" + from deepmd.pt_expt.kernels.triton.sezm import 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) + 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_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 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_TRITON_INFER", None) + if old_level is not None: + os.environ["DP_TRITON_INFER"] = old_level + + def test_wide_channels_fall_back_with_winning_entry(self): + """``Cf >= GATE_BMM_MIN_FOCUS_DIM`` routes through cuBLAS, not fusion. + + 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, + ) + + 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_winning_entry(self): + """The fused kernel assumes fp32 rounding; other dtypes keep the relay. + + Scoped with a winning entry registered for the same reason as above. + """ + 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()) + + +@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. + + 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) + ) + + +@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. + + 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, mock.patch("torch.cuda.is_available", return_value=True): + 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_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, + 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()