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..bf44aed708 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py +++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py @@ -160,6 +160,11 @@ _MAX_MIXER_RANK = 4 _ROTATE_MIX_BWD_CONFIG = (1, 2) # per-edge backward (warps, stages) +# Fused rotate+mix backward (segment-parallel, the per-edge node +# gradient ``gxe`` is never materialized): (BLOCK_E edges per chunk, +# warps, stages). Conservative defaults; a swept table is the +# follow-up if the fusion proves out on more shapes. +_ROT_MIX_BWD_FUSED_CONFIG = (8, 4, 2) # ====================================================================== @@ -557,10 +562,10 @@ def _rotate_mix_fwd_kernel( accp = tl.zeros((BC,), dtype=tl.float32) for j in tl.static_range(2 * l + 1): xv = xrows[l * l + j] - acc0 += tl.load(d_base + r0 * DIM + base + j) * xv + acc0 += tl.load(_wigner_row_ptr(d_base, r0, base + j, DIM)) * xv if l >= 1: - accm += tl.load(d_base + (r0 - 1) * DIM + base + j) * xv - accp += tl.load(d_base + (r0 + 1) * DIM + base + j) * xv + accm += tl.load(_wigner_row_ptr(d_base, r0 - 1, base + j, DIM)) * xv + accp += tl.load(_wigner_row_ptr(d_base, r0 + 1, base + j, DIM)) * xv rows0 = rows0 + (acc0,) if l >= 1: rows_m = rows_m + (accm,) @@ -700,10 +705,10 @@ def _rotate_mix_bwd_kernel( accp = tl.zeros((BC,), dtype=tl.float32) for j in tl.static_range(2 * l + 1): xv = xrows[l * l + j] - acc0 += tl.load(d_base + r0 * DIM + base + j) * xv + acc0 += tl.load(_wigner_row_ptr(d_base, r0, base + j, DIM)) * xv if l >= 1: - accm += tl.load(d_base + (r0 - 1) * DIM + base + j) * xv - accp += tl.load(d_base + (r0 + 1) * DIM + base + j) * xv + accm += tl.load(_wigner_row_ptr(d_base, r0 - 1, base + j, DIM)) * xv + accp += tl.load(_wigner_row_ptr(d_base, r0 + 1, base + j, DIM)) * xv rows0 = rows0 + (acc0,) if l >= 1: rows_m = rows_m + (accm,) @@ -1028,6 +1033,247 @@ def _segment_sum_kernel( acc += tl.load(rows_ptr + e * P + cols, mask=col_mask, other=0.0) tl.store(out_ptr + node * P + cols, acc, mask=col_mask) + @triton.jit + def _wigner_row_ptr(d_base, row, col, DIM: tl.constexpr): + """Address of one Wigner entry; shared row/column arithmetic.""" + return d_base + row * DIM + col + + @triton.jit + def _rotate_mix_bwd_fused_kernel( + gu_ptr, # (F, E, ROW) upstream gradient (focus-major) + x_ptr, # (N, D, CW) node features + w_ptr, # (E, D, D) block-diagonal Wigner-D + kc_ptr, # (E, KSZ) rank-1 compact kernel, or (E, L+1, CW) when RANK == 0 + cb_ptr, # (1, CW) channel basis (RANK == 1) + order_ptr, # (E,) edge ids sorted by source node + row_ptr_ptr, # (N + 1,) CSR offsets into ``order`` + gx_ptr, # (N, D, CW) node gradient out (accumulated on chip) + gw_ptr, # (E, D, D) Wigner gradient out (structural non-zeros; pre-zeroed) + gkc_ptr, # gradient of kc out, same layout as kc + n_edge, + x_sn, + x_sd, + L: tl.constexpr, + CF: tl.constexpr, + CW: tl.constexpr, + CP: tl.constexpr, # next power of two >= CW (vector lane count) + DP: tl.constexpr, # next power of two >= DIM (accumulator row count) + RANK: tl.constexpr, + BLOCK_E: tl.constexpr, + ): + """Source-segmented rotate+mix backward: ``gxe`` never materialized. + + One program per source node walks its CSR segment (``order`` / + ``row_ptr``) in ``BLOCK_E``-edge chunks, recomputes the rotation + backward per chunk with exactly the edge-block kernel's math, and + accumulates the per-edge node-gradient rows into an on-chip + ``(DP, CP)`` accumulator, replacing the + ``_rotate_mix_bwd_*_kernel`` -> ``_segment_sum_kernel`` pair and + its ``(E, D, CW)`` HBM round-trip. Two structural properties of + the segment make this cheap: every edge in the segment shares the + same source row, so the ``x`` rows are loaded once per program + (not once per edge), and the per-edge gradient contribution is a + rank-1 update ``w[col] * g_row`` whose segment sum is one + ``tl.sum`` over the edge axis. + + ``gw``/``gkc`` are stored per edge with the same math as the + edge-block kernel (identical values up to instruction scheduling, + i.e. fp32 rounding); ``gx`` accumulates in chunk order with a + ``tl.sum`` inside each chunk, so it differs from the unfused path + only at fp32 rounding level. Masked edges contribute exactly zero + (their Wigner scalars load as 0). + """ + NS0: tl.constexpr = L + 1 + DIM: tl.constexpr = (L + 1) * (L + 1) + ROW: tl.constexpr = (3 * L + 1) * CF + KSZ: tl.constexpr = NS0 * NS0 + L * L + PADDED: tl.constexpr = CP != CW + + node = tl.program_id(0).to(tl.int64) + beg = tl.load(row_ptr_ptr + node).to(tl.int64) + end = tl.load(row_ptr_ptr + node + 1).to(tl.int64) + + chan = tl.arange(0, CP) + if PADDED: + cmask = chan < CW + chan_c = tl.where(cmask, chan, 0) + else: + # CP == CW here; a real always-true mask keeps the variable + # uniformly defined instead of a None placeholder. + cmask = chan < CP + chan_c = chan + didx = tl.arange(0, DP) + # Every edge of the segment shares this source row, so ``x`` is + # read through one node base; the rows stay L1/L2-resident across + # the chunk loop (dynamic loops cannot index register tuples, so + # the rows are reloaded per use -- near-free out of cache). + x_node = x_ptr + node * x_sn + + if RANK == 1: + cbv = tl.load(cb_ptr + chan, mask=(chan < CW), other=0.0)[None, :] + + acc = tl.zeros((DP, CP), dtype=tl.float32) + for i0 in range(beg, end, BLOCK_E): + offs = i0 + tl.arange(0, BLOCK_E) + e_mask = offs < end + eq = tl.load(order_ptr + offs, mask=e_mask, other=0).to(tl.int64) + if PADDED: + em = e_mask[:, None] & cmask[None, :] + else: + em = e_mask[:, None] + + d_base = w_ptr + eq * DIM * DIM + gd_base = gw_ptr + eq * DIM * DIM + if RANK == 0: + kc_base = kc_ptr + (eq * NS0 * CW)[:, None] + gkc_base = gkc_ptr + (eq * NS0 * CW)[:, None] + else: + kc_base = kc_ptr + eq * KSZ + gkc_base = gkc_ptr + eq * KSZ + + f_off = ( + gu_ptr + + ((chan_c // CF).to(tl.int64) * n_edge * ROW + (chan_c % CF))[None, :] + + (eq * ROW)[:, None] + ) + + for l in tl.static_range(L + 1): + base = l * l + r0 = base + l + + xl0 = tl.zeros((BLOCK_E, CP), dtype=tl.float32) + xlm = tl.zeros((BLOCK_E, CP), dtype=tl.float32) + xlp = tl.zeros((BLOCK_E, CP), dtype=tl.float32) + for j in tl.static_range(2 * l + 1): + if PADDED: + xv = tl.load( + x_node + (base + j) * x_sd + chan_c, + mask=cmask, + other=0.0, + ) + else: + xv = tl.load(x_node + (base + j) * x_sd + chan_c) + w0 = tl.load( + _wigner_row_ptr(d_base, r0, base + j, DIM), + mask=e_mask, + other=0.0, + ) + xl0 += w0[:, None] * xv[None, :] + if l >= 1: + wm = tl.load( + d_base + (r0 - 1) * DIM + base + j, + mask=e_mask, + other=0.0, + ) + wp = tl.load( + d_base + (r0 + 1) * DIM + base + j, + mask=e_mask, + other=0.0, + ) + xlm += wm[:, None] * xv[None, :] + xlp += wp[:, None] * xv[None, :] + + # Kernel gradient rows and local-frame gradients of degree + # l (same values, same per-o accumulation order as the + # edge-block kernel -> bit-identical gw/gkc). + g0 = tl.zeros((BLOCK_E, CP), dtype=tl.float32) + gm = tl.zeros((BLOCK_E, CP), dtype=tl.float32) + gp = tl.zeros((BLOCK_E, CP), dtype=tl.float32) + if RANK == 0: + gy0 = tl.load(f_off + l * CF, mask=em, other=0.0) + if l >= 1: + gym = tl.load(f_off + (NS0 + l - 1) * CF, mask=em, other=0.0) + gyp = tl.load( + f_off + (NS0 + L + l - 1) * CF, mask=em, other=0.0 + ) + t = gy0 * xl0 + gym * xlm + gyp * xlp + else: + t = gy0 * xl0 + tl.store(gkc_base + l * CW + chan[None, :], t, mask=em) + rad_l = tl.load( + kc_base + l * CW + chan[None, :], mask=em, other=0.0 + ) + g0 = gy0 * rad_l + if l >= 1: + gm = gym * rad_l + gp = gyp * rad_l + else: + for o in tl.static_range(NS0): + gyo = tl.load(f_off + o * CF, mask=em, other=0.0) * cbv + tl.store( + gkc_base + l * NS0 + o, + tl.sum(gyo * xl0, axis=1), + mask=e_mask, + ) + k_val = tl.load(kc_base + l * NS0 + o, mask=e_mask, other=0.0) + g0 += k_val[:, None] * gyo + if l >= 1: + for o in tl.static_range(L): + gyn = ( + tl.load(f_off + (NS0 + o) * CF, mask=em, other=0.0) + * cbv + ) + gys = ( + tl.load(f_off + (NS0 + L + o) * CF, mask=em, other=0.0) + * cbv + ) + tl.store( + gkc_base + NS0 * NS0 + (l - 1) * L + o, + tl.sum(gyn * xlm + gys * xlp, axis=1), + mask=e_mask, + ) + k_val = tl.load( + kc_base + NS0 * NS0 + (l - 1) * L + o, + mask=e_mask, + other=0.0, + ) + gm += k_val[:, None] * gyn + gp += k_val[:, None] * gys + + # Rotation backward: Wigner gradients per edge (unchanged + # math), node-gradient rows accumulated on chip. + for j in tl.static_range(2 * l + 1): + col = base + j + if PADDED: + xv = tl.load( + x_node + col * x_sd + chan_c, mask=cmask, other=0.0 + ) + else: + xv = tl.load(x_node + col * x_sd + chan_c) + w0 = tl.load(d_base + r0 * DIM + col, mask=e_mask, other=0.0) + gx_row = w0[:, None] * g0 + tl.store( + gd_base + r0 * DIM + col, + tl.sum(g0 * xv[None, :], axis=1), + mask=e_mask, + ) + if l >= 1: + wm = tl.load( + d_base + (r0 - 1) * DIM + col, mask=e_mask, other=0.0 + ) + wp = tl.load( + d_base + (r0 + 1) * DIM + col, mask=e_mask, other=0.0 + ) + gx_row += wm[:, None] * gm + wp[:, None] * gp + tl.store( + gd_base + (r0 - 1) * DIM + col, + tl.sum(gm * xv[None, :], axis=1), + mask=e_mask, + ) + tl.store( + gd_base + (r0 + 1) * DIM + col, + tl.sum(gp * xv[None, :], axis=1), + mask=e_mask, + ) + gxs = tl.sum(gx_row, axis=0) + acc += tl.where((didx == col)[:, None], gxs[None, :], 0.0) + + tl.store( + gx_ptr + node * DIM * CW + didx[:, None] * CW + chan[None, :], + acc, + mask=(didx < DIM)[:, None] & (chan[None, :] < CW), + ) + @triton.jit def _stack_gemm_m0_kernel( u_ptr, # (F, E, ROW) layer input @@ -2477,10 +2723,14 @@ def _rotate_mix_impl( lmax: int, n_focus: int, rank: int, + level: int, ) -> Tensor: # The source CSR view rides through the forward untouched so the autograd - # context can hand it to the backward's segment reduction. - del src_order, src_rowptr + # context can hand it to the backward's segment reduction. The level is + # likewise inert here: it reaches the backward through the autograd + # context (``ctx.level``), retaining the selection the forward's caller + # made instead of rereading the ambient environment. + del src_order, src_rowptr, level if not _use_triton(x): return _rotate_mix_reference(x, src, wigner, kc, cb, lmax, n_focus, rank) n_edge = src.shape[0] @@ -2926,6 +3176,70 @@ def _gated_act_bwd_impl( return grad_z, grad_logit +def _rotate_mix_bwd_fused_impl( + grad_u: Tensor, + x: Tensor, + src: Tensor, + wigner: Tensor, + kc: Tensor, + cb: Tensor, + order: Tensor, + row_ptr: Tensor, + lmax: int, + n_focus: int, + rank: int, +) -> tuple[Tensor, Tensor, Tensor]: + """Rotate+mix backward with the node-gradient reduction fused in. + + Returns ``(grad_x, grad_wigner, grad_kc)`` directly -- the per-edge + dense gradient ``gxe`` of the unfused path is never materialized (see + :func:`_rotate_mix_bwd_fused_kernel`). ``order``/``row_ptr`` are the + source-segment CSR topology the unfused path builds for + ``segment_sum``. + """ + n_edge = src.shape[0] + n_node = row_ptr.shape[0] - 1 + c_wide = int(x.shape[2]) + dim = (int(lmax) + 1) ** 2 + if not _use_triton(x): + grad_x_edge, grad_wigner, grad_kc = _rotate_mix_backward_reference( + grad_u, x, src, wigner, kc, cb, lmax, n_focus, rank + ) + grad_x = _segment_sum_impl(grad_x_edge, order, row_ptr) + return grad_x, grad_wigner, grad_kc + grad_x = torch.empty(n_node, dim, c_wide, device=x.device, dtype=x.dtype) + grad_wigner = torch.zeros_like(wigner) + grad_kc = torch.empty_like(kc) + if _has_no_edges(n_edge): + return grad_x.zero_(), grad_wigner, grad_kc + block_e, warps, stages = _ROT_MIX_BWD_FUSED_CONFIG + wrap_triton(_rotate_mix_bwd_fused_kernel)[(n_node,)]( + grad_u, + x, + wigner, + kc, + cb, + order, + row_ptr, + grad_x, + grad_wigner, + grad_kc, + n_edge, + x.stride(0), + x.stride(1), + L=int(lmax), + CF=c_wide // int(n_focus), + CW=c_wide, + CP=triton.next_power_of_2(c_wide), + DP=triton.next_power_of_2(dim), + RANK=int(rank), + BLOCK_E=block_e, + num_warps=warps, + num_stages=stages, + ) + return grad_x, grad_wigner, grad_kc + + def _mixing_stack_impl( u0: Tensor, alpha: Tensor, @@ -3610,6 +3924,20 @@ def _mixing_stack_train_bwd_impl( _segment_sum_op = torch.library.triton_op("sezm_triton::segment_sum", mutates_args=())( _segment_sum_impl ) +_rotate_mix_bwd_fused_op = torch.library.triton_op( + "sezm_triton::so2_rotate_mix_bwd_fused", mutates_args=() +)(_rotate_mix_bwd_fused_impl) + + +@_rotate_mix_bwd_fused_op.register_fake +def _(grad_u, x, src, wigner, kc, cb, order, row_ptr, lmax, n_focus, rank): + return ( + x.new_empty((row_ptr.shape[0] - 1, (lmax + 1) ** 2, x.shape[2])), + torch.empty_like(wigner), + torch.empty_like(kc), + ) + + _mixing_stack_op = torch.library.triton_op( "sezm_triton::so2_mixing_stack", mutates_args=() )(_mixing_stack_impl) @@ -3636,7 +3964,8 @@ def _mixing_stack_train_bwd_impl( @_rotate_mix_op.register_fake -def _(x, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank): +def _(x, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank, level): + del level focus_dim = x.shape[2] // n_focus return x.new_empty((n_focus, src.shape[0], (3 * lmax + 1) * focus_dim)) @@ -3862,22 +4191,65 @@ def rotate_mix_basis_grad( def _rotate_mix_setup_context(ctx, inputs, output): - x, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank = inputs + x, src, src_order, src_rowptr, wigner, kc, cb, lmax, n_focus, rank, level = inputs ctx.save_for_backward(x, src, src_order, src_rowptr, wigner, kc, cb) ctx.lmax = lmax ctx.n_focus = n_focus ctx.rank = rank + ctx.level = level def _rotate_mix_backward(ctx, grad_u): x, src, src_order, src_rowptr, wigner, kc, cb = ctx.saved_tensors grad_u = grad_u.contiguous() - grad_x_edge, grad_wigner, grad_kc = _rotate_mix_bwd_op( - grad_u, x, src, wigner, kc, cb, ctx.lmax, ctx.n_focus, ctx.rank - ) - # Contention-free segmented reduction of the per-edge node gradient through - # the source CSR view the step builds once. - grad_x = _segment_sum_op(grad_x_edge, src_order, src_rowptr) + # ``torch.is_grad_enabled()`` is True inside this backward exactly when + # the backward itself is being differentiated eagerly + # (``create_graph=True``). Traced force-loss training instead compiles + # the backward with grad mode disabled while handing it a cotangent + # that is still part of the outer graph (the tracers communicate the + # training requirement through the incoming gradient, not ambient + # mode), so a requiring cotangent equally announces a next derivative. + # The fused operator registers only a fake (meta) implementation and + # no next-derivative formula, while the unfused backward and + # segment-sum operators do support the next derivative -- so both + # higher-order signals keep the supported pair and the fused path + # engages only when neither is present. + # The fused backward rides the level-1 universal-kernel contract of the + # rotate-mix family (no tunable configuration; it strictly removes the + # (E, D, C_wide) node-gradient round-trip). The governing level is the + # one the forward's caller passed as an explicit argument -- the + # per-block value captured at construction, following the module's + # current train/eval mode -- so the backward retains exactly the + # selection the forward made and never rereads the ambient environment + # (which may have been restored between construction and evaluation). + if ( + int(ctx.level) >= 1 + and int(ctx.rank) <= 1 + and not torch.is_grad_enabled() + and not grad_u.requires_grad + ): + # Fused path: the node gradient accumulates on chip inside the + # kernel; the (E, D, C_wide) ``gxe`` HBM round-trip never exists. + grad_x, grad_wigner, grad_kc = _rotate_mix_bwd_fused_op( + grad_u, + x, + src, + wigner, + kc, + cb, + src_order, + src_rowptr, + ctx.lmax, + ctx.n_focus, + ctx.rank, + ) + else: + grad_x_edge, grad_wigner, grad_kc = _rotate_mix_bwd_op( + grad_u, x, src, wigner, kc, cb, ctx.lmax, ctx.n_focus, ctx.rank + ) + # Contention-free segmented reduction of the per-edge node gradient + # through the source CSR view the step builds once. + grad_x = _segment_sum_op(grad_x_edge, src_order, src_rowptr) grad_cb = ( rotate_mix_basis_grad( grad_u, x, src, wigner, kc, cb, ctx.lmax, ctx.n_focus, ctx.rank @@ -3885,7 +4257,20 @@ def _rotate_mix_backward(ctx, grad_u): if int(ctx.rank) > 0 and ctx.needs_input_grad[6] else None ) - return grad_x, None, None, None, grad_wigner, grad_kc, grad_cb, None, None, None + # The trailing None is the level argument: integers carry no gradient. + return ( + grad_x, + None, + None, + None, + grad_wigner, + grad_kc, + grad_cb, + None, + None, + None, + None, + ) def _rotate_mix_bwd_setup_context(ctx, inputs, output): @@ -3918,8 +4303,22 @@ def _rotate_mix_bwd_backward(ctx, grad_grad_x, grad_grad_wigner, grad_grad_kc): edge_src = torch.arange(src.shape[0], device=src.device, dtype=src.dtype) def forward(x_arg: Tensor, src_arg: Tensor, w_arg: Tensor, k_arg: Tensor) -> Tensor: + # The level argument is inert on this re-entry: grad mode stays + # enabled while the second order is being built, so the inner + # operator's own backward can never take the fused path whatever + # this value is; 1 documents the family's level-1 contract. return _rotate_mix_op( - x_arg, src_arg, edge_src, edge_src, w_arg, k_arg, cb, lmax, n_focus, rank + x_arg, + src_arg, + edge_src, + edge_src, + w_arg, + k_arg, + cb, + lmax, + n_focus, + rank, + 1, ) def backward( @@ -5106,6 +5505,11 @@ def __call__( conv.lmax, conv.n_focus, rank, + # The per-block levels were captured at construction; the + # backward must retain the same selection this forward made + # (training level in training mode, inference level + # otherwise), not reread the ambient environment. + conv.triton_train_level if conv.training else conv.triton_infer_level, ) return u0, rad_feat @@ -5338,6 +5742,11 @@ def __call__( conv.lmax, conv.n_focus, rank, + # The per-block levels were captured at construction; the + # backward must retain the same selection this forward made + # (training level in training mode, inference level + # otherwise), not reread the ambient environment. + conv.triton_train_level if conv.training else conv.triton_infer_level, ) # === Step 3. Cross-focus competition weight from the l = 0 scalars === diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index f359f073d6..8c3a8b715f 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -853,6 +853,7 @@ def test_float64_fallback_channel_basis_gradient_preserves_precision(self) -> No lmax, n_focus, rank, + 0, # CPU tensors take the reference path; level 0 keeps it ) (grad_fused,) = torch.autograd.grad(fused, basis_fused, grad_out) @@ -2053,3 +2054,450 @@ def test_tune_missing_configs_runs_the_real_flash_bwd_sweep(self): if __name__ == "__main__": unittest.main() + + +@_GPU_KERNELS +class TestSeZMTritonRotMixBwdFused(unittest.TestCase): + """Cross-check the fused rotate+mix backward against the unfused pair. + + The fused kernel replaces ``_rotate_mix_bwd_*`` + ``_segment_sum`` and + must not materialize ``gxe``. ``grad_wigner``/``grad_kc`` use the same + per-edge math but are scheduled differently (shared-row broadcast vs + per-edge tile), so all three outputs are compared at fp32-rounding + tolerance; ``grad_x`` additionally changes its per-segment + accumulation order. + """ + + CASES: typing.ClassVar[list[tuple]] = [ + # (lmax, channels, n_focus, rank) + (1, 32, 1, 0), + (2, 32, 1, 1), + (3, 64, 2, 1), # production shape (C_wide = 128, lmax = 3) + (4, 64, 1, 1), + # non-power-of-two C_wide exercises the PADDED channel lanes + (2, 24, 2, 1), # C_wide = 48 + (1, 40, 1, 0), + # rank 2 exercises the rank-indexed cb/cb and flattened kc addressing + (2, 32, 1, 2), + ] + + N_NODE = 512 + N_EDGE = 20000 + + def _inputs(self, lmax, channels, n_focus, rank): + generator = torch.Generator(device="cuda").manual_seed(1234 + lmax) + dim = (lmax + 1) ** 2 + c_wide = channels * n_focus + row = (3 * lmax + 1) * channels + x = torch.randn(self.N_NODE, dim, c_wide, device="cuda", generator=generator) + src = torch.randint( + 0, self.N_NODE, (self.N_EDGE,), device="cuda", generator=generator + ) + grad_u = ( + torch.randn(n_focus, self.N_EDGE, row, device="cuda", generator=generator) + * 0.3 + ) + if rank == 0: + kc = ( + torch.randn( + self.N_EDGE, lmax + 1, c_wide, device="cuda", generator=generator + ) + * 0.2 + ) + cb = torch.zeros(1, c_wide, device="cuda") + else: + # rank-R factorized kernel: kc is (E, KSZ * rank) flattened and + # cb is (rank, C_wide); the kernels index both by the leading + # rank, so undersized rank-1 buffers would read/write out of + # bounds (the reference reshape rejects the element count). + ksz = (lmax + 1) ** 2 + lmax**2 + kc = ( + torch.randn(self.N_EDGE, ksz * rank, device="cuda", generator=generator) + * 0.2 + ) + cb = torch.randn(rank, c_wide, device="cuda", generator=generator) * 0.1 + # Structural block-diagonal Wigner-D values (zeros elsewhere), so + # the gw comparison covers exactly the stored entries. + wigner = torch.zeros(self.N_EDGE, dim, dim, device="cuda") + for l in range(lmax + 1): + base = l * l + block = torch.randn( + self.N_EDGE, + 2 * l + 1, + 2 * l + 1, + device="cuda", + generator=generator, + ) + wigner[:, base : base + 2 * l + 1, base : base + 2 * l + 1] = block + return grad_u, x, src, wigner, kc, cb + + def test_fused_matches_unfused_across_family(self): + from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp + + for lmax, channels, n_focus, rank in self.CASES: + with self.subTest(case=(lmax, channels, n_focus, rank)): + grad_u, x, src, wigner, kc, cb = self._inputs( + lmax, channels, n_focus, rank + ) + order = torch.argsort(src) + boundaries = torch.arange( + self.N_NODE + 1, device="cuda", dtype=src.dtype + ) + row_ptr = torch.searchsorted(src.index_select(0, order), boundaries) + + gxe, gw_ref, gkc_ref = vp._rotate_mix_bwd_op( + grad_u, x, src, wigner, kc, cb, lmax, n_focus, rank + ) + gx_ref = vp._segment_sum_op(gxe, order, row_ptr) + + # The fused kernel supports rank <= 1 (the dispatch + # contract keeps rank >= 2 on the unfused pair), so rank 2 + # checks the unfused kernels against the eager reference. + # The Wigner gradient is compared on the structural block + # diagonal: the kernel promises only those entries, while + # the dense reference also returns off-block ones. + if rank > 1: + gxe_e, gw_e, gkc_e = vp._rotate_mix_backward_reference( + grad_u, x, src, wigner, kc, cb, lmax, n_focus, rank + ) + torch.testing.assert_close(gkc_ref, gkc_e, atol=1e-5, rtol=1e-5) + mask = _block_mask(lmax, "cuda") + torch.testing.assert_close( + gw_ref * mask, gw_e * mask, atol=1e-5, rtol=1e-5 + ) + gx_e = torch.zeros( + self.N_NODE, + (lmax + 1) ** 2, + channels * n_focus, + device="cuda", + ).index_add_(0, src, gxe_e) + torch.testing.assert_close(gx_ref, gx_e, atol=1e-4, rtol=1e-4) + continue + + gx_fused, gw_fused, gkc_fused = vp._rotate_mix_bwd_fused_op( + grad_u, + x, + src, + wigner, + kc, + cb, + order, + row_ptr, + lmax, + n_focus, + rank, + ) + + gw_scale = gw_ref.abs().max().item() + torch.testing.assert_close( + gw_fused, gw_ref, atol=1e-6 * max(gw_scale, 1.0), rtol=1e-5 + ) + gkc_scale = gkc_ref.abs().max().item() + torch.testing.assert_close( + gkc_fused, gkc_ref, atol=1e-6 * max(gkc_scale, 1.0), rtol=1e-5 + ) + scale = gx_ref.abs().max().item() + torch.testing.assert_close( + gx_fused, gx_ref, atol=1e-5 * max(scale, 1.0), rtol=1e-5 + ) + + def _record_fused_backward(self, run): + """Run ``run()`` while counting fused-backward launches; pass through.""" + from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp + + calls = [] + original = vp._rotate_mix_bwd_fused_op + + def recorder(*args, **kwargs): + calls.append(args) + return original(*args, **kwargs) + + with mock.patch.object(vp, "_rotate_mix_bwd_fused_op", side_effect=recorder): + result = run() + torch.cuda.synchronize() + return len(calls), result + + def _clear_and_restore_levels(self): + """Clear both Triton level variables; restore the prior values on exit. + + The override is registered as a cleanup, so the configured + environment survives even when an assertion fails mid-test. + """ + import os + + saved = { + name: os.environ.get(name) + for name in ("DP_TRITON_INFER", "DP_TRITON_TRAIN") + } + + def restore(): + for name, value in saved.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + + for name in saved: + os.environ.pop(name, None) + self.addCleanup(restore) + + def test_autograd_branch_follows_forward_level(self): + """The autograd backward dispatches on the level the forward was given. + + The level argument stands in for the per-block value captured at + construction; the ambient environment is cleared to zero throughout + so the check proves the argument (not the environment) governs. + """ + from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp + + # Scope the override: the cleared levels are restored even when an + # assertion fails, so later tests keep the configured environment. + self._clear_and_restore_levels() + lmax, channels, n_focus, rank = 3, 32, 2, 1 + grad_u, x, src, wigner, kc, cb = self._inputs(lmax, channels, n_focus, rank) + src_order = torch.argsort(src) + boundaries = torch.arange(self.N_NODE + 1, device=src.device, dtype=src.dtype) + src_rowptr = torch.searchsorted(src.index_select(0, src_order), boundaries) + + grads, launches = {}, {} + for level in (0, 1): + xg = x.clone().requires_grad_(True) + wg = wigner.clone().requires_grad_(True) + kg = kc.clone().requires_grad_(True) + out = vp._rotate_mix_op( + xg, src, src_order, src_rowptr, wg, kg, cb, lmax, n_focus, rank, level + ) + launches[level], grads[level] = self._record_fused_backward( + lambda: torch.autograd.grad( + out, [xg, wg, kg], grad_u, retain_graph=True + ) + ) + # The counting stand-in proves the selection happened; gradient + # equality alone could pass with both levels on the unfused path. + self.assertEqual(launches[0], 0) + self.assertEqual(launches[1], 1) + # All three agree at fp32 rounding level: grad_x by accumulation + # order, gw/gkc by instruction scheduling of the same math. + scale = grads[0][0].abs().max().item() + torch.testing.assert_close( + grads[1][0], grads[0][0], atol=1e-5 * max(scale, 1.0), rtol=1e-5 + ) + for i in (1, 2): + s_i = grads[0][i].abs().max().item() + torch.testing.assert_close( + grads[1][i], grads[0][i], atol=1e-6 * max(s_i, 1.0), rtol=1e-5 + ) + + @_GPU_KERNELS + def test_traced_force_loss_keeps_the_unfused_path(self): + """Traced force loss: the backward runs as a plain compiled function. + + The tracers compile the backward with grad mode disabled and hand + it a cotangent that is still part of the outer graph -- + ``autograd.grad(create_graph=True)`` cannot emulate this (the + engine re-enables grad inside an eager double backward), so the + check drives the backward function directly under ``no_grad`` + with a requiring cotangent, exactly the shape the tracer + executes. The fused operator registers no next-derivative + formula, so a requiring cotangent must keep the unfused pair. + """ + from types import ( + SimpleNamespace, + ) + + from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp + + self._clear_and_restore_levels() + lmax, channels, n_focus, rank = 3, 32, 2, 1 + grad_u, x, src, wigner, kc, cb = self._inputs(lmax, channels, n_focus, rank) + src_order = torch.argsort(src) + boundaries = torch.arange(self.N_NODE + 1, device=src.device, dtype=src.dtype) + src_rowptr = torch.searchsorted(src.index_select(0, src_order), boundaries) + + xg = x.clone().requires_grad_(True) + wg = wigner.clone().requires_grad_(True) + kg = kc.clone().requires_grad_(True) + out = vp._rotate_mix_op( + xg, src, src_order, src_rowptr, wg, kg, cb, lmax, n_focus, rank, 1 + ) + + holder = [] + ctx = SimpleNamespace( + save_for_backward=lambda *ts: holder.append(ts), + # The unfused path reads the tracer's training requirements + # from the context (all tensor inputs require grad here). + needs_input_grad=( + True, + True, + False, + False, + True, + True, + True, + False, + False, + False, + False, + ), + ) + vp._rotate_mix_setup_context( + ctx, + (xg, src, src_order, src_rowptr, wg, kg, cb, lmax, n_focus, rank, 1), + out, + ) + ctx.saved_tensors = holder[0] + cot = grad_u.clone().requires_grad_(True) + + def _backward(): + with torch.no_grad(): + return vp._rotate_mix_backward(ctx, cot) + + launches, grads = self._record_fused_backward(_backward) + self.assertEqual( + launches, 0, "a requiring cotangent must keep the unfused pair" + ) + ref = torch.autograd.grad(out, [xg, wg, kg], grad_u, retain_graph=True) + # The backward returns one gradient per forward input: x, wigner + # and kc live at positions 0, 4 and 5 of the 11-tuple. + scale = ref[0].abs().max().item() + torch.testing.assert_close( + grads[0], ref[0], atol=1e-5 * max(scale, 1.0), rtol=1e-5 + ) + for got, want in ((grads[4], ref[1]), (grads[5], ref[2])): + s_i = want.abs().max().item() + torch.testing.assert_close(got, want, atol=1e-6 * max(s_i, 1.0), rtol=1e-5) + + def test_higher_order_backward_preserved_at_level_one(self): + """At level 1, ``create_graph=True`` keeps the supported path. + + The fused operator registers only a fake (meta) implementation and + no next-derivative formula; the backward of the backward must + therefore route through the unfused backward + segment-sum pair + even with the level-1 selection active. The counting stand-in + asserts the fused backward never launches during either order, and + the second-order results agree with a level-0 run. + """ + from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp + + lmax, channels, n_focus, rank = 3, 32, 2, 1 + grad_u, x, src, wigner, kc, cb = self._inputs(lmax, channels, n_focus, rank) + src_order = torch.argsort(src) + boundaries = torch.arange(self.N_NODE + 1, device=src.device, dtype=src.dtype) + src_rowptr = torch.searchsorted(src.index_select(0, src_order), boundaries) + + second, launches = {}, {} + for level in (0, 1): + xg = x.clone().requires_grad_(True) + wg = wigner.clone().requires_grad_(True) + kg = kc.clone().requires_grad_(True) + out = vp._rotate_mix_op( + xg, src, src_order, src_rowptr, wg, kg, cb, lmax, n_focus, rank, level + ) + launches[level], (g1,) = self._record_fused_backward( + lambda: torch.autograd.grad( + out, xg, grad_u, create_graph=True, retain_graph=True + ) + ) + second[level] = torch.autograd.grad(g1.sum(), [wg, kg]) + self.assertEqual(launches[0], 0) + self.assertEqual(launches[1], 0) + for i in (0, 1): + s_i = second[0][i].abs().max().item() + torch.testing.assert_close( + second[1][i], + second[0][i], + atol=1e-6 * max(s_i, 1.0), + rtol=1e-5, + ) + + def test_rank_above_one_falls_back_at_level_one(self): + """``rank > 1`` never takes the fused path, level or not. + + The rank guard stands even with the level-1 selection active; the + counting stand-in asserts the unfused chain, and the rank>1 chain + is not bitwise deterministic across runs, so agreement between the + two levels is checked at fp32-rounding tolerance. + """ + from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp + + lmax, channels, n_focus, rank = 3, 32, 2, 2 + grad_u, x, src, wigner, kc, cb = self._inputs(lmax, channels, n_focus, rank) + src_order = torch.argsort(src) + boundaries = torch.arange(self.N_NODE + 1, device=src.device, dtype=src.dtype) + src_rowptr = torch.searchsorted(src.index_select(0, src_order), boundaries) + + grads, launches = {}, {} + for level in (0, 1): + xg = x.clone().requires_grad_(True) + out = vp._rotate_mix_op( + xg, + src, + src_order, + src_rowptr, + wigner, + kc, + cb, + lmax, + n_focus, + rank, + level, + ) + launches[level], (grads[level],) = self._record_fused_backward( + lambda: torch.autograd.grad(out, xg, grad_u, retain_graph=True) + ) + self.assertEqual(launches[0], 0) + self.assertEqual(launches[1], 0) + s_i = grads[0].abs().max().item() + torch.testing.assert_close( + grads[1], grads[0], atol=1e-6 * max(s_i, 1.0), rtol=1e-5 + ) + + def test_backward_retains_construction_level_across_env_restore(self): + """The wrapper's captured level governs the backward, not ambient env. + + A module constructed with ``DP_TRITON_INFER=1`` keeps its Triton + forward after the environment is restored to the default; its + backward must retain the same selection instead of rereading the + (now zero) ambient level and silently falling back. + """ + from types import ( + SimpleNamespace, + ) + + from deepmd.pt_expt.kernels.triton.sezm import so2_value_path as vp + + # Scope the ambient-env restore once: the cleared levels come back + # even when a subTest assertion fails. + self._clear_and_restore_levels() + lmax, channels, n_focus, rank = 3, 32, 2, 0 + grad_u, x, src, wigner, kc, cb = self._inputs(lmax, channels, n_focus, rank) + src_order = torch.argsort(src) + boundaries = torch.arange(self.N_NODE + 1, device=src.device, dtype=src.dtype) + src_rowptr = torch.searchsorted(src.index_select(0, src_order), boundaries) + edge_cache = SimpleNamespace( + src=src, + D_full=wigner, + csr_cache={"src": (src_order, src_rowptr)}, + ) + radial_feat = kc if rank == 0 else None + + for captured, expected_launches in ((1, 1), (0, 0)): + with self.subTest(captured_infer_level=captured): + conv = SimpleNamespace( + radial_hidden_proj=None, + radial_degree_mixer=None, + lmax=lmax, + n_focus=n_focus, + training=False, + triton_infer_level=captured, + triton_train_level=0, + ) + wrapper = vp._TritonRotateMix(conv) + xg = x.clone().requires_grad_(True) + out, _rad = wrapper(xg, edge_cache, radial_feat) + launches, _grads = self._record_fused_backward( + lambda: torch.autograd.grad( + out, xg, grad_u.reshape(out.shape), retain_graph=True + ) + ) + self.assertEqual(launches, expected_launches)