From a2ebc42e234238d5a35f60cb7fb7f1b26257db69 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Sun, 11 Oct 2026 18:01:34 +0800 Subject: [PATCH 1/5] perf(pt): fused Triton zonal scatter for the GIE head (DP_GIE_SCATTER_INFER) --- .../pt/model/descriptor/sezm_nn/embedding.py | 60 +++ .../kernels/triton/sezm/gie_scatter.py | 481 ++++++++++++++++++ deepmd/pt_expt/kernels/utils.py | 19 + 3 files changed, 560 insertions(+) create mode 100644 deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py diff --git a/deepmd/pt/model/descriptor/sezm_nn/embedding.py b/deepmd/pt/model/descriptor/sezm_nn/embedding.py index f3215f3072..87a8506ec6 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/embedding.py +++ b/deepmd/pt/model/descriptor/sezm_nn/embedding.py @@ -37,6 +37,7 @@ ) from deepmd.pt_expt.kernels.utils import ( cuda_infer_level, + use_gie_scatter_infer, ) from deepmd.utils.version import ( check_version_compatibility, @@ -226,6 +227,33 @@ def __init__( self.lmax, self.ebed_dim - 1, self.channels ) + # === Fused Triton message-and-scatter operator === + # Same boundary as the CUDA operator above, served by the Triton + # stack instead: the kernel specializes on the packed contiguous row + # layout, so anything else (or a non-fp32 compute) keeps the reference + # composition. + self._triton_scatter = None + if use_gie_scatter_infer() and self.dtype is torch.float32 and self.lmax == 3: + from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import ( + GIE_SCATTER_TRITON_AVAILABLE, + gie_zonal_scatter, + ) + + packed_rows = torch.arange( + 1, self.ebed_dim, device=self.non_scalar_row_index.device + ) + counts = [2 * ell + 1 for ell in range(1, self.lmax + 1)] + packed_slots = torch.repeat_interleave( + torch.arange(self.lmax, device=self.non_scalar_row_index.device), + torch.tensor(counts, device=self.non_scalar_row_index.device), + ) + if ( + GIE_SCATTER_TRITON_AVAILABLE + and torch.equal(self.non_scalar_row_index, packed_rows) + and torch.equal(self.radial_slot_index_for_row, packed_slots) + ): + self._triton_scatter = gie_zonal_scatter + def forward( self, *, @@ -293,6 +321,38 @@ def forward( n_nodes, edge_cache, radial_feat, zonal_coupling ) + # === Step 3a. Fused Triton message-and-scatter === + # Same boundary as the CUDA operator above: the degree-normalized + # scale stays outside so its cotangent keeps flowing through the + # traced graph, while the (E, D-1, C) message itself never leaves the + # kernel's registers in either direction. + if ( + self._triton_scatter is not None + and spin_l1_message is None + and edge_cache.edge_src_gate is None + and edge_cache.csr_cache is not None + ): + from .edge_cache import ( + cached_edge_csr, + ) + + order, row_ptr = cached_edge_csr(edge_cache, "dst", n_nodes) + # The operator contract promises contiguous inputs: the + # radial slice and the gathered coupling are flattened in the + # traced graph before entering the op, like every other + # sezm_triton operator. + out = self._triton_scatter( + zonal_coupling.contiguous(), + radial_feat.contiguous(), + edge_cache.dst, + order, + row_ptr, + n_nodes, + self.lmax, + ) # (N, D, C), l=0 zero, unnormalized + out.mul_(edge_cache.inv_sqrt_deg) + return out + radial_value_for_row = radial_feat.index_select( 1, self.radial_slot_index_for_row ) # (E, D-1, C) diff --git a/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py b/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py new file mode 100644 index 0000000000..5c54d88956 --- /dev/null +++ b/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py @@ -0,0 +1,481 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +# pyright: reportMissingImports=false +# ruff: noqa: ANN001, ANN202 +"""Fused zonal message build and destination reduction for the GIE head. + +The geometric initial embedding forms, per edge, the message + + ``msg[e, r, c] = zonal[e, r] * radial[e, slot(r), c]`` + +and reduces it onto the destination node rows of the packed (l, m) layout. +The reference composition materializes the ``(E, R, C)`` message first -- at +the production shape that surface is 1.2 GB and it stays alive until the +backward has re-read it, so it also sits in the peak-residency set. The fused +forward keeps the message in registers: one node-parallel program walks the +destination CSR (stable-sorted, so the reduction order matches the eager +``index_add_`` on CPU exactly) and writes the already-padded ``(N, R+1, C)`` +output with a zero scalar row. The fused backward is edge-parallel and emits +both operand gradients in one pass, so neither direction ever forms the edge +message. + +The degree grouping (rows are contiguous per degree, ``2l+1`` rows each) is a +compile-time tuple like the monomial exponent table; the operator therefore +specializes per ``lmax`` at launch. The layout assumptions are checked by the +binding site, which falls back to the reference composition when the packed +row order or the degree slots are not the canonical ones. +""" + +from __future__ import ( + annotations, +) + +import torch +from torch import ( + Tensor, +) +from torch.library import ( + wrap_triton, +) +from triton import ( + cdiv, +) +import triton +import triton.language as tl + +__all__ = [ + "GIE_SCATTER_TRITON_AVAILABLE", + "gie_zonal_scatter", +] + +GIE_SCATTER_TRITON_AVAILABLE = torch.cuda.is_available() and triton is not None + + +def _slot_layout(lmax: int) -> tuple[list[int], list[int], list[int]]: + """Row counts, row bases and padded widths of the per-degree row groups.""" + counts = [2 * ell + 1 for ell in range(1, lmax + 1)] + bases = [sum(counts[:i]) for i in range(len(counts))] + pads = [triton.next_power_of_2(c) for c in counts] + return counts, bases, pads + + +def _reference( + zonal: Tensor, + radial: Tensor, + dst: Tensor, + n_nodes: int, +) -> Tensor: + """Eager fallback: same math, same stable within-node addition order.""" + n_row = zonal.shape[1] + n_slot = radial.shape[1] + counts = [2 * ell + 1 for ell in range(1, n_slot + 1)] + slot_of_row = torch.repeat_interleave( + torch.arange(n_slot, device=zonal.device), + torch.tensor(counts, device=zonal.device), + ) + msg = zonal.unsqueeze(-1) * radial[:, slot_of_row] # (E, R, C) + out = zonal.new_zeros(n_nodes, n_row + 1, radial.shape[2]) + out[:, 1:, :].index_add_(0, dst, msg) + return out + + +@triton.jit +def _gie_scatter_fwd_kernel( + zonal_ptr, + radial_ptr, + order_ptr, + row_ptr_ptr, + out_ptr, + z_stride, + r_stride0, + r_stride1, + C: tl.constexpr, + SLOTS: tl.constexpr, + R0: tl.constexpr, + R1: tl.constexpr, + R2: tl.constexpr, + P0: tl.constexpr, + P1: tl.constexpr, + P2: tl.constexpr, + B0: tl.constexpr, + B1: tl.constexpr, + B2: tl.constexpr, + N_ROW: tl.constexpr, +): + n = tl.program_id(0) + start = tl.load(row_ptr_ptr + n).to(tl.int64) + end = tl.load(row_ptr_ptr + n + 1).to(tl.int64) + offs_c = tl.arange(0, C) + # One register accumulator per degree group (row counts are + # compile-time constants and each group's rows are contiguous) + r0 = tl.arange(0, P0) + r1 = tl.arange(0, P1) + r2 = tl.arange(0, P2) + m0 = r0 < R0 + m1 = r1 < R1 + m2 = r2 < R2 + acc0 = tl.zeros((P0, C), dtype=tl.float32) + acc1 = tl.zeros((P1, C), dtype=tl.float32) + acc2 = tl.zeros((P2, C), dtype=tl.float32) + for e in range(start, end): + # Edges inside a CSR segment map back to the original edge + # order before zonal/radial can be read + eo = tl.load(order_ptr + e).to(tl.int64) + # zonal/radial may be strided views (radial_feat[:, 1:, :]), + # so address them through their strides + z_base = eo * z_stride + rad_base = eo * r_stride0 + z0 = tl.load(zonal_ptr + z_base + B0 + r0, mask=m0, other=0.0) + z1 = tl.load(zonal_ptr + z_base + B1 + r1, mask=m1, other=0.0) + z2 = tl.load(zonal_ptr + z_base + B2 + r2, mask=m2, other=0.0) + rd0 = tl.load(radial_ptr + rad_base + 0 * r_stride1 + offs_c) + rd1 = tl.load(radial_ptr + rad_base + 1 * r_stride1 + offs_c) + rd2 = tl.load(radial_ptr + rad_base + 2 * r_stride1 + offs_c) + acc0 += z0[:, None] * rd0[None, :] + acc1 += z1[:, None] * rd1[None, :] + acc2 += z2[:, None] * rd2[None, :] + out_base = n * (N_ROW + 1) * C + tl.store(out_ptr + out_base + 0 * C + offs_c, tl.zeros((C,), dtype=tl.float32)) + tl.store( + out_ptr + out_base + ((1 + B0 + r0) * C)[:, None] + offs_c[None, :], + acc0, + mask=m0[:, None], + ) + tl.store( + out_ptr + out_base + ((1 + B1 + r1) * C)[:, None] + offs_c[None, :], + acc1, + mask=m1[:, None], + ) + tl.store( + out_ptr + out_base + ((1 + B2 + r2) * C)[:, None] + offs_c[None, :], + acc2, + mask=m2[:, None], + ) + + +@triton.jit +def _gie_scatter_bwd_kernel( + grad_ptr, + zonal_ptr, + radial_ptr, + dst_ptr, + gz_ptr, + gr_ptr, + n_edge, + g_nstride, + g_rstride, + z_stride, + r_stride0, + r_stride1, + C: tl.constexpr, + SLOTS: tl.constexpr, + R0: tl.constexpr, + R1: tl.constexpr, + R2: tl.constexpr, + P0: tl.constexpr, + P1: tl.constexpr, + P2: tl.constexpr, + B0: tl.constexpr, + B1: tl.constexpr, + B2: tl.constexpr, + N_ROW: tl.constexpr, + BLOCK_E: tl.constexpr, +): + pid = tl.program_id(0) + offs_e = (pid * BLOCK_E + tl.arange(0, BLOCK_E)).to(tl.int64) + mask_e = offs_e < n_edge + offs_c = tl.arange(0, C) + r0 = tl.arange(0, P0) + r1 = tl.arange(0, P1) + r2 = tl.arange(0, P2) + m0 = r0 < R0 + m1 = r1 < R1 + m2 = r2 < R2 + d = tl.load(dst_ptr + offs_e, mask=mask_e, other=0).to(tl.int64) + # The node grad surface carries the zero scalar row at index 0 and + # the message rows from 1; honor its strides + g_base = (d * g_nstride + g_rstride)[:, None, None] + g0 = tl.load( + grad_ptr + g_base + ((B0 + r0) * C)[None, :, None] + offs_c[None, None, :], + mask=mask_e[:, None, None] & m0[None, :, None], + other=0.0, + ) + g1 = tl.load( + grad_ptr + g_base + ((B1 + r1) * C)[None, :, None] + offs_c[None, None, :], + mask=mask_e[:, None, None] & m1[None, :, None], + other=0.0, + ) + g2 = tl.load( + grad_ptr + g_base + ((B2 + r2) * C)[None, :, None] + offs_c[None, None, :], + mask=mask_e[:, None, None] & m2[None, :, None], + other=0.0, + ) + z_base = (offs_e * z_stride)[:, None] + z0 = tl.load( + zonal_ptr + z_base + B0 + r0[None, :], + mask=mask_e[:, None] & m0[None, :], + other=0.0, + ) + z1 = tl.load( + zonal_ptr + z_base + B1 + r1[None, :], + mask=mask_e[:, None] & m1[None, :], + other=0.0, + ) + z2 = tl.load( + zonal_ptr + z_base + B2 + r2[None, :], + mask=mask_e[:, None] & m2[None, :], + other=0.0, + ) + rad_base = (offs_e * r_stride0)[:, None] + offs_c[None, :] + rd0 = tl.load( + radial_ptr + rad_base + 0 * r_stride1, mask=mask_e[:, None], other=0.0 + ) + rd1 = tl.load( + radial_ptr + rad_base + 1 * r_stride1, mask=mask_e[:, None], other=0.0 + ) + rd2 = tl.load( + radial_ptr + rad_base + 2 * r_stride1, mask=mask_e[:, None], other=0.0 + ) + gz0 = tl.sum(g0 * rd0[:, None, :], 2) + gz1 = tl.sum(g1 * rd1[:, None, :], 2) + gz2 = tl.sum(g2 * rd2[:, None, :], 2) + gr0 = tl.sum(g0 * z0[:, :, None], 1) + gr1 = tl.sum(g1 * z1[:, :, None], 1) + gr2 = tl.sum(g2 * z2[:, :, None], 1) + gz_base = (offs_e * N_ROW)[:, None] + tl.store( + gz_ptr + gz_base + B0 + r0[None, :], gz0, mask=mask_e[:, None] & m0[None, :] + ) + tl.store( + gz_ptr + gz_base + B1 + r1[None, :], gz1, mask=mask_e[:, None] & m1[None, :] + ) + tl.store( + gz_ptr + gz_base + B2 + r2[None, :], gz2, mask=mask_e[:, None] & m2[None, :] + ) + gr_store = (offs_e * SLOTS * C)[:, None] + offs_c[None, :] + tl.store(gr_ptr + gr_store + 0 * C, gr0, mask=mask_e[:, None]) + tl.store(gr_ptr + gr_store + 1 * C, gr1, mask=mask_e[:, None]) + tl.store(gr_ptr + gr_store + 2 * C, gr2, mask=mask_e[:, None]) + + +def _forward_impl( + zonal: Tensor, + radial: Tensor, + dst: Tensor, + order: Tensor, + row_ptr: Tensor, + n_nodes: int, + lmax: int, +) -> Tensor: + if not ( + GIE_SCATTER_TRITON_AVAILABLE + and zonal.is_cuda + and zonal.dtype is torch.float32 + and lmax == 3 + ): + return _reference(zonal, radial, dst, n_nodes) + n_edge = zonal.shape[0] + n_row = zonal.shape[1] + c = radial.shape[2] + if zonal.stride(-1) != 1: + zonal = zonal.contiguous() + if radial.stride(-1) != 1: + radial = radial.contiguous() + out = torch.empty((n_nodes, n_row + 1, c), device=zonal.device, dtype=torch.float32) + if type(n_edge) is int and n_edge == 0: + out.zero_() + return out + counts, bases, pads = _slot_layout(int(lmax)) + # Derive the grid from the row-pointer shape: under tracing this + # keeps a symbolic expression, while a python-int argument would be + # baked into a wrong constant (measured: grid baked to 35) + wrap_triton(_gie_scatter_fwd_kernel)[(row_ptr.shape[0] - 1,)]( + zonal, + radial, + order, + row_ptr, + out, + zonal.stride(0), + radial.stride(0), + radial.stride(1), + C=c, + SLOTS=len(counts), + R0=counts[0], + R1=counts[1], + R2=counts[2], + P0=pads[0], + P1=pads[1], + P2=pads[2], + B0=bases[0], + B1=bases[1], + B2=bases[2], + N_ROW=n_row, + num_warps=4, + num_stages=2, + ) + return out + + +def _backward_impl( + grad_out: Tensor, + zonal: Tensor, + radial: Tensor, + dst: Tensor, + lmax: int, +) -> tuple[Tensor, Tensor]: + if not ( + GIE_SCATTER_TRITON_AVAILABLE + and zonal.is_cuda + and zonal.dtype is torch.float32 + and lmax == 3 + ): + return _backward_reference(grad_out, zonal, radial, dst, int(lmax)) + n_edge, n_row = zonal.shape + c = radial.shape[2] + if zonal.stride(-1) != 1: + zonal = zonal.contiguous() + if radial.stride(-1) != 1: + radial = radial.contiguous() + if grad_out.stride(-1) != 1 or grad_out.stride(1) != grad_out.shape[2]: + grad_out = grad_out.contiguous() + gz = torch.empty((n_edge, n_row), device=zonal.device, dtype=torch.float32) + gr = torch.empty( + (n_edge, radial.shape[1], c), device=zonal.device, dtype=torch.float32 + ) + if type(n_edge) is int and n_edge == 0: + return gz, gr + counts, bases, pads = _slot_layout(int(lmax)) + block_e = 16 + wrap_triton(_gie_scatter_bwd_kernel)[(cdiv(n_edge, block_e),)]( + grad_out, + zonal, + radial, + dst, + gz, + gr, + n_edge, + grad_out.stride(0), + grad_out.stride(1), + zonal.stride(0), + radial.stride(0), + radial.stride(1), + C=c, + SLOTS=len(counts), + R0=counts[0], + R1=counts[1], + R2=counts[2], + P0=pads[0], + P1=pads[1], + P2=pads[2], + B0=bases[0], + B1=bases[1], + B2=bases[2], + N_ROW=n_row, + BLOCK_E=block_e, + num_warps=4, + num_stages=2, + ) + return gz, gr + + +def _backward_reference( + grad_out: Tensor, + zonal: Tensor, + radial: Tensor, + dst: Tensor, + lmax: int, +) -> tuple[Tensor, Tensor]: + """Eager fallback of the backward: matches the kernel's grouping.""" + g = grad_out[:, 1:, :] # (N, R, C) 跳过零行 + g_msg = g[dst] # (E, R, C) + n_slot = radial.shape[1] + counts = [2 * ell + 1 for ell in range(1, n_slot + 1)] + slot_of_row = torch.repeat_interleave( + torch.arange(n_slot, device=zonal.device), + torch.tensor(counts, device=zonal.device), + ) + radial_rows = radial[:, slot_of_row] + gz = (g_msg * radial_rows).sum(-1) + gr = torch.zeros_like(radial) + base = 0 + for s, cnt in enumerate(counts): + gr[:, s, :] = ( + g_msg[:, base : base + cnt, :] * zonal[:, base : base + cnt, None] + ).sum(1) + base += cnt + return gz, gr + + +_gie_scatter_op = torch.library.triton_op( + "sezm_triton::gie_zonal_scatter", mutates_args=() +)(_forward_impl) + +_gie_scatter_bwd_op = torch.library.triton_op( + "sezm_triton::gie_zonal_scatter_bwd", mutates_args=() +)(_backward_impl) + + +@_gie_scatter_op.register_fake +def _(zonal, radial, dst, order, row_ptr, n_nodes, lmax): + return zonal.new_empty((n_nodes, zonal.shape[1] + 1, radial.shape[2])) + + +@_gie_scatter_bwd_op.register_fake +def _(grad_out, zonal, radial, dst, lmax): + return zonal.new_empty(zonal.shape), radial.new_empty(radial.shape) + + +def _setup_context(ctx, inputs, output): + zonal, radial, dst, order, row_ptr, n_nodes, lmax = inputs + # The backward reads the grad surface in the original edge order, + # so only the three small operands need saving + ctx.save_for_backward(zonal, radial, dst) + ctx.lmax = lmax + + +def _backward(ctx, grad_out): + zonal, radial, dst = ctx.saved_tensors + gz, gr = _gie_scatter_bwd_op(grad_out.contiguous(), zonal, radial, dst, ctx.lmax) + return gz, gr, None, None, None, None, None + + +torch.library.register_autograd( + "sezm_triton::gie_zonal_scatter", + _backward, + setup_context=_setup_context, +) + + +def gie_zonal_scatter( + zonal: Tensor, + radial: Tensor, + dst: Tensor, + order: Tensor, + row_ptr: Tensor, + n_nodes: int, + lmax: int, +) -> Tensor: + """Public entry: fused zonal message build + destination reduction. + + Parameters + ---------- + zonal + Per-edge zonal coupling with shape (E, R), R = D-1 for the packed + non-scalar rows. + radial + Per-edge radial features with shape (E, lmax, C), one slot per degree + ``l >= 1``. + dst + Original-order destination index with shape (E,). + order, row_ptr + Stable destination CSR view (int64, from ``cached_edge_csr``). + n_nodes + Node count of the reduction target. + lmax + Maximum degree; the packed contiguous row layout is required. + + Returns + ------- + torch.Tensor + ``(N, R+1, C)`` node features with a zero scalar row, unnormalized. + """ + return _gie_scatter_op(zonal, radial, dst, order, row_ptr, n_nodes, lmax) diff --git a/deepmd/pt_expt/kernels/utils.py b/deepmd/pt_expt/kernels/utils.py index c6e54bd17f..45ef849949 100644 --- a/deepmd/pt_expt/kernels/utils.py +++ b/deepmd/pt_expt/kernels/utils.py @@ -380,3 +380,22 @@ def use_amp_infer() -> bool: ``True`` when ``DP_AMP_INFER`` is set to a truthy value. """ return os.environ.get("DP_AMP_INFER", "0").strip().lower() in _INFER_TRUE +def use_gie_scatter_infer() -> bool: + """Return whether the GIE head message runs the fused scatter operator. + + The flag is controlled by the ``DP_GIE_SCATTER_INFER`` environment + variable and is read when the geometric initial embedding module is + constructed. When truthy (CUDA fp32 inference with the packed contiguous + row layout), the per-edge zonal message is built and destination-reduced + inside one Triton operator instead of the reference composition, so the + ``(E, D-1, C)`` edge message surface is never materialized in either + direction. The reduction walks the stable destination CSR, matching the + eager ``index_add_`` addition order. + + Returns + ------- + bool + ``True`` when ``DP_GIE_SCATTER_INFER`` is set to a truthy value. + """ + return os.environ.get("DP_GIE_SCATTER_INFER", "0").strip().lower() in _INFER_TRUE + From af902aa2d589b1c4ddbc9b4dee9e964892afb386 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Sun, 11 Oct 2026 18:01:34 +0800 Subject: [PATCH 2/5] test(pt): cover the fused GIE zonal scatter operator --- .../pt/model/test_descriptor_sezm_triton.py | 179 ++++++++++++++++++ 1 file changed, 179 insertions(+) diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index f359f073d6..a01097e4fd 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -2053,3 +2053,182 @@ def test_tune_missing_configs_runs_the_real_flash_bwd_sweep(self): if __name__ == "__main__": unittest.main() + + +class TestSeZMGieScatter(unittest.TestCase): + """Fused zonal message build + destination reduction for the GIE head.""" + + _ROWS = 15 # non-scalar rows at lmax=3: 3+5+7 + + def _make_case(self, n_edge, n_nodes, generator, view_radial=False): + """Random operator inputs; a sliced radial view when asked.""" + zonal = torch.randn(n_edge, self._ROWS, device="cuda", generator=generator) + parent = torch.randn(n_edge, 4, 32, device="cuda", generator=generator) + radial = parent[:, 1:, :] if view_radial else parent[:, 1:, :].contiguous() + dst = torch.randint(0, n_nodes, (n_edge,), device="cuda", generator=generator) + order = torch.argsort(dst, dim=0, stable=True) + counts = torch.bincount(dst, minlength=n_nodes) + row_ptr = torch.cat( + [ + torch.zeros(1, dtype=torch.int64, device="cuda"), + torch.cumsum(counts, 0).to(torch.int64), + ] + ) + return zonal, radial, dst, order, row_ptr + + def _reference(self, zonal, radial, dst, n_nodes): + """Composition reference: broadcast the message, then index_add.""" + counts = [3, 5, 7] + slot_of_row = torch.repeat_interleave( + torch.arange(3, device=zonal.device), + torch.tensor(counts, device=zonal.device), + ) + msg = zonal.unsqueeze(-1) * radial[:, slot_of_row] + out = zonal.new_zeros(n_nodes, self._ROWS + 1, 32) + out[:, 1:, :].index_add_(0, dst, msg) + return out + + @_GPU_KERNELS + def test_forward_backward_matches_reference(self): + from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import ( + gie_zonal_scatter, + ) + + generator = torch.Generator(device="cuda").manual_seed(5) + for n_edge, n_nodes in ((4096, 128), (777, 33)): + for view_radial in (False, True): + with self.subTest(n_edge=n_edge, n_nodes=n_nodes, view=view_radial): + zonal, radial, dst, order, row_ptr = self._make_case( + n_edge, n_nodes, generator, view_radial + ) + z1 = zonal.clone().requires_grad_(True) + r1 = radial.clone().requires_grad_(True) + out = gie_zonal_scatter(z1, r1, dst, order, row_ptr, n_nodes, 3) + z2 = zonal.clone().requires_grad_(True) + r2 = radial.clone().requires_grad_(True) + ref = self._reference(z2, r2, dst, n_nodes) + torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5) + self.assertTrue( + torch.equal(out[:, 0, :], torch.zeros_like(out[:, 0, :])) + ) + cotangent = torch.randn( + out.shape, + device="cuda", + generator=torch.Generator(device="cuda").manual_seed(37), + ) + g1 = torch.autograd.grad(out, (z1, r1), cotangent) + g2 = torch.autograd.grad(ref, (z2, r2), cotangent) + torch.testing.assert_close(g1[0], g2[0], atol=1e-5, rtol=1e-5) + torch.testing.assert_close(g1[1], g2[1], atol=1e-5, rtol=1e-5) + + def test_reference_cpu_matches_manual(self): + """CPU reference equals a hand-written sequential accumulation.""" + from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import _reference + + generator = torch.Generator().manual_seed(7) + n_edge, n_nodes = 513, 17 + zonal = torch.randn(n_edge, self._ROWS, generator=generator, device="cpu") + radial = torch.randn(n_edge, 3, 32, generator=generator, device="cpu") + dst = torch.randint(0, n_nodes, (n_edge,), generator=generator, device="cpu") + out = _reference(zonal, radial, dst, n_nodes) + slot_of_row = torch.repeat_interleave( + torch.arange(3, device="cpu"), torch.tensor([3, 5, 7], device="cpu") + ) + manual = torch.zeros(n_nodes, self._ROWS + 1, 32, device="cpu") + manual[:, 1:, :].index_add_( + 0, dst, zonal.unsqueeze(-1) * radial[:, slot_of_row] + ) + torch.testing.assert_close(out, manual, atol=0.0, rtol=0.0) + + @_GPU_KERNELS + def test_compile_matches_eager(self): + """The compiled path must match eager (guards the launch-grid trap).""" + from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import ( + gie_zonal_scatter, + ) + + generator = torch.Generator(device="cuda").manual_seed(11) + zonal, radial, dst, order, row_ptr = self._make_case( + 4096, 128, generator, view_radial=True + ) + + def fn(z, r): + return gie_zonal_scatter(z, r, dst, order, row_ptr, 128, 3) + + z_e = zonal.clone().requires_grad_(True) + r_e = radial.clone().requires_grad_(True) + out_e = fn(z_e, r_e) + cotangent = torch.randn_like(out_e) + g_e = torch.autograd.grad(out_e, (z_e, r_e), cotangent) + + z_c = zonal.clone().requires_grad_(True) + r_c = radial.clone().requires_grad_(True) + out_c = torch.compile(fn)(z_c, r_c) + g_c = torch.autograd.grad(out_c, (z_c, r_c), cotangent) + + torch.testing.assert_close(out_c, out_e, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(g_c[0], g_e[0], atol=1e-6, rtol=1e-6) + torch.testing.assert_close(g_c[1], g_e[1], atol=1e-6, rtol=1e-6) + + @_GPU_KERNELS + def test_embedding_dispatch_two_state(self): + """Module-level two states: bound op vs reference composition.""" + from deepmd.pt.model.descriptor.sezm_nn.embedding import ( + GeometricInitialEmbedding, + ) + from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import ( + gie_zonal_scatter, + ) + + generator = torch.Generator(device="cuda").manual_seed(13) + n_edge, n_nodes, channels = 2048, 64, 32 + zonal, radial, dst, order, row_ptr = self._make_case(n_edge, n_nodes, generator) + # Dense rotation surface the base layout gathers the coupling from + d_full = torch.randn(n_edge, 16, 16, device="cuda", generator=generator) + # Fixed per-test tensors so both states see identical inputs + edge_vec = torch.zeros(n_edge, 3, device="cuda") + inv_sqrt_deg = torch.rand(n_nodes, 1, 1, device="cuda", generator=generator) + + class _Cache: + pass + + def make_cache(df): + cache = _Cache() + cache.dst = dst + cache.edge_vec = edge_vec + cache.inv_sqrt_deg = inv_sqrt_deg + cache.csr_cache = {} + cache.edge_src_gate = None + cache.Dt_full = df + return cache + + module = GeometricInitialEmbedding( + lmax=3, channels=channels, dtype=torch.float32 + ).cuda() + module.eval() + cotangent = None + + def run(bound): + module._triton_scatter = gie_zonal_scatter if bound else None + df = d_full.clone().requires_grad_(True) + rad = radial.clone().requires_grad_(True) + out = module( + n_nodes=n_nodes, + edge_cache=make_cache(df), + radial_feat=rad, + ) + nonlocal cotangent + if cotangent is None: + cotangent = torch.randn( + out.shape, + device="cuda", + generator=torch.Generator(device="cuda").manual_seed(37), + ) + grads = torch.autograd.grad(out, (df, rad), cotangent) + return out.detach(), grads + + out_ref, (gd_ref, gr_ref) = run(bound=False) + out_fused, (gd_fused, gr_fused) = run(bound=True) + torch.testing.assert_close(out_fused, out_ref, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(gd_fused, gd_ref, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(gr_fused, gr_ref, atol=1e-5, rtol=1e-5) From 288e49c478975519aa55e4afd0bf8895e4476a9c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 11 Oct 2026 14:51:17 +0000 Subject: [PATCH 3/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py | 4 ++-- deepmd/pt_expt/kernels/utils.py | 3 ++- source/tests/pt/model/test_descriptor_sezm_triton.py | 4 +++- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py b/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py index 5c54d88956..e2b2fbe270 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py +++ b/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py @@ -30,6 +30,8 @@ ) import torch +import triton +import triton.language as tl from torch import ( Tensor, ) @@ -39,8 +41,6 @@ from triton import ( cdiv, ) -import triton -import triton.language as tl __all__ = [ "GIE_SCATTER_TRITON_AVAILABLE", diff --git a/deepmd/pt_expt/kernels/utils.py b/deepmd/pt_expt/kernels/utils.py index 45ef849949..6da551f361 100644 --- a/deepmd/pt_expt/kernels/utils.py +++ b/deepmd/pt_expt/kernels/utils.py @@ -380,6 +380,8 @@ def use_amp_infer() -> bool: ``True`` when ``DP_AMP_INFER`` is set to a truthy value. """ return os.environ.get("DP_AMP_INFER", "0").strip().lower() in _INFER_TRUE + + def use_gie_scatter_infer() -> bool: """Return whether the GIE head message runs the fused scatter operator. @@ -398,4 +400,3 @@ def use_gie_scatter_infer() -> bool: ``True`` when ``DP_GIE_SCATTER_INFER`` is set to a truthy value. """ return os.environ.get("DP_GIE_SCATTER_INFER", "0").strip().lower() in _INFER_TRUE - diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index a01097e4fd..4a5569b035 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -2123,7 +2123,9 @@ def test_forward_backward_matches_reference(self): def test_reference_cpu_matches_manual(self): """CPU reference equals a hand-written sequential accumulation.""" - from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import _reference + from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import ( + _reference, + ) generator = torch.Generator().manual_seed(7) n_edge, n_nodes = 513, 17 From 4f42e828a93ee09b2d75241c03eddfb926121dfb Mon Sep 17 00:00:00 2001 From: longjunhao Date: Fri, 16 Oct 2026 19:33:39 +0800 Subject: [PATCH 4/5] fix(pt): padded channel blocks and a training gate for the fused GIE scatter --- .../pt/model/descriptor/sezm_nn/embedding.py | 5 ++ .../kernels/triton/sezm/gie_scatter.py | 54 ++++++++++++------- .../pt/model/test_descriptor_sezm_triton.py | 27 ++++++++-- 3 files changed, 63 insertions(+), 23 deletions(-) diff --git a/deepmd/pt/model/descriptor/sezm_nn/embedding.py b/deepmd/pt/model/descriptor/sezm_nn/embedding.py index 87a8506ec6..b405bc2092 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/embedding.py +++ b/deepmd/pt/model/descriptor/sezm_nn/embedding.py @@ -326,8 +326,13 @@ def forward( # scale stays outside so its cotangent keeps flowing through the # traced graph, while the (E, D-1, C) message itself never leaves the # kernel's registers in either direction. + # Training keeps the reference composition: its autograd is exact + # to all orders, while the fused operator's backward registers no + # next-derivative formula (force-loss training differentiates the + # backward itself). if ( self._triton_scatter is not None + and not self.training and spin_l1_message is None and edge_cache.edge_src_gate is None and edge_cache.csr_cache is not None diff --git a/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py b/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py index e2b2fbe270..035de41cc7 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py +++ b/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py @@ -100,11 +100,14 @@ def _gie_scatter_fwd_kernel( B1: tl.constexpr, B2: tl.constexpr, N_ROW: tl.constexpr, + BC: tl.constexpr, ): + # tl.arange needs a power-of-two extent, so channels ride a padded block n = tl.program_id(0) start = tl.load(row_ptr_ptr + n).to(tl.int64) end = tl.load(row_ptr_ptr + n + 1).to(tl.int64) - offs_c = tl.arange(0, C) + offs_c = tl.arange(0, BC) + cmask = offs_c < C # One register accumulator per degree group (row counts are # compile-time constants and each group's rows are contiguous) r0 = tl.arange(0, P0) @@ -113,9 +116,9 @@ def _gie_scatter_fwd_kernel( m0 = r0 < R0 m1 = r1 < R1 m2 = r2 < R2 - acc0 = tl.zeros((P0, C), dtype=tl.float32) - acc1 = tl.zeros((P1, C), dtype=tl.float32) - acc2 = tl.zeros((P2, C), dtype=tl.float32) + acc0 = tl.zeros((P0, BC), dtype=tl.float32) + acc1 = tl.zeros((P1, BC), dtype=tl.float32) + acc2 = tl.zeros((P2, BC), dtype=tl.float32) for e in range(start, end): # Edges inside a CSR segment map back to the original edge # order before zonal/radial can be read @@ -127,28 +130,38 @@ def _gie_scatter_fwd_kernel( z0 = tl.load(zonal_ptr + z_base + B0 + r0, mask=m0, other=0.0) z1 = tl.load(zonal_ptr + z_base + B1 + r1, mask=m1, other=0.0) z2 = tl.load(zonal_ptr + z_base + B2 + r2, mask=m2, other=0.0) - rd0 = tl.load(radial_ptr + rad_base + 0 * r_stride1 + offs_c) - rd1 = tl.load(radial_ptr + rad_base + 1 * r_stride1 + offs_c) - rd2 = tl.load(radial_ptr + rad_base + 2 * r_stride1 + offs_c) + rd0 = tl.load( + radial_ptr + rad_base + 0 * r_stride1 + offs_c, mask=cmask, other=0.0 + ) + rd1 = tl.load( + radial_ptr + rad_base + 1 * r_stride1 + offs_c, mask=cmask, other=0.0 + ) + rd2 = tl.load( + radial_ptr + rad_base + 2 * r_stride1 + offs_c, mask=cmask, other=0.0 + ) acc0 += z0[:, None] * rd0[None, :] acc1 += z1[:, None] * rd1[None, :] acc2 += z2[:, None] * rd2[None, :] out_base = n * (N_ROW + 1) * C - tl.store(out_ptr + out_base + 0 * C + offs_c, tl.zeros((C,), dtype=tl.float32)) + tl.store( + out_ptr + out_base + 0 * C + offs_c, + tl.zeros((BC,), dtype=tl.float32), + mask=cmask, + ) tl.store( out_ptr + out_base + ((1 + B0 + r0) * C)[:, None] + offs_c[None, :], acc0, - mask=m0[:, None], + mask=m0[:, None] & cmask[None, :], ) tl.store( out_ptr + out_base + ((1 + B1 + r1) * C)[:, None] + offs_c[None, :], acc1, - mask=m1[:, None], + mask=m1[:, None] & cmask[None, :], ) tl.store( out_ptr + out_base + ((1 + B2 + r2) * C)[:, None] + offs_c[None, :], acc2, - mask=m2[:, None], + mask=m2[:, None] & cmask[None, :], ) @@ -179,11 +192,14 @@ def _gie_scatter_bwd_kernel( B2: tl.constexpr, N_ROW: tl.constexpr, BLOCK_E: tl.constexpr, + BC: tl.constexpr, ): + # tl.arange needs a power-of-two extent, so channels ride a padded block pid = tl.program_id(0) offs_e = (pid * BLOCK_E + tl.arange(0, BLOCK_E)).to(tl.int64) mask_e = offs_e < n_edge - offs_c = tl.arange(0, C) + offs_c = tl.arange(0, BC) + cmask = offs_c < C r0 = tl.arange(0, P0) r1 = tl.arange(0, P1) r2 = tl.arange(0, P2) @@ -196,17 +212,17 @@ def _gie_scatter_bwd_kernel( g_base = (d * g_nstride + g_rstride)[:, None, None] g0 = tl.load( grad_ptr + g_base + ((B0 + r0) * C)[None, :, None] + offs_c[None, None, :], - mask=mask_e[:, None, None] & m0[None, :, None], + mask=mask_e[:, None, None] & m0[None, :, None] & cmask[None, None, :], other=0.0, ) g1 = tl.load( grad_ptr + g_base + ((B1 + r1) * C)[None, :, None] + offs_c[None, None, :], - mask=mask_e[:, None, None] & m1[None, :, None], + mask=mask_e[:, None, None] & m1[None, :, None] & cmask[None, None, :], other=0.0, ) g2 = tl.load( grad_ptr + g_base + ((B2 + r2) * C)[None, :, None] + offs_c[None, None, :], - mask=mask_e[:, None, None] & m2[None, :, None], + mask=mask_e[:, None, None] & m2[None, :, None] & cmask[None, None, :], other=0.0, ) z_base = (offs_e * z_stride)[:, None] @@ -252,9 +268,9 @@ def _gie_scatter_bwd_kernel( gz_ptr + gz_base + B2 + r2[None, :], gz2, mask=mask_e[:, None] & m2[None, :] ) gr_store = (offs_e * SLOTS * C)[:, None] + offs_c[None, :] - tl.store(gr_ptr + gr_store + 0 * C, gr0, mask=mask_e[:, None]) - tl.store(gr_ptr + gr_store + 1 * C, gr1, mask=mask_e[:, None]) - tl.store(gr_ptr + gr_store + 2 * C, gr2, mask=mask_e[:, None]) + tl.store(gr_ptr + gr_store + 0 * C, gr0, mask=mask_e[:, None] & cmask[None, :]) + tl.store(gr_ptr + gr_store + 1 * C, gr1, mask=mask_e[:, None] & cmask[None, :]) + tl.store(gr_ptr + gr_store + 2 * C, gr2, mask=mask_e[:, None] & cmask[None, :]) def _forward_impl( @@ -298,6 +314,7 @@ def _forward_impl( radial.stride(0), radial.stride(1), C=c, + BC=triton.next_power_of_2(c), SLOTS=len(counts), R0=counts[0], R1=counts[1], @@ -359,6 +376,7 @@ def _backward_impl( radial.stride(0), radial.stride(1), C=c, + BC=triton.next_power_of_2(c), SLOTS=len(counts), R0=counts[0], R1=counts[1], diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index 4a5569b035..bcb4af4baf 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -2060,10 +2060,10 @@ class TestSeZMGieScatter(unittest.TestCase): _ROWS = 15 # non-scalar rows at lmax=3: 3+5+7 - def _make_case(self, n_edge, n_nodes, generator, view_radial=False): + def _make_case(self, n_edge, n_nodes, generator, view_radial=False, channels=32): """Random operator inputs; a sliced radial view when asked.""" zonal = torch.randn(n_edge, self._ROWS, device="cuda", generator=generator) - parent = torch.randn(n_edge, 4, 32, device="cuda", generator=generator) + parent = torch.randn(n_edge, 4, channels, device="cuda", generator=generator) radial = parent[:, 1:, :] if view_radial else parent[:, 1:, :].contiguous() dst = torch.randint(0, n_nodes, (n_edge,), device="cuda", generator=generator) order = torch.argsort(dst, dim=0, stable=True) @@ -2084,7 +2084,7 @@ def _reference(self, zonal, radial, dst, n_nodes): torch.tensor(counts, device=zonal.device), ) msg = zonal.unsqueeze(-1) * radial[:, slot_of_row] - out = zonal.new_zeros(n_nodes, self._ROWS + 1, 32) + out = zonal.new_zeros(n_nodes, self._ROWS + 1, radial.shape[2]) out[:, 1:, :].index_add_(0, dst, msg) return out @@ -2095,11 +2095,12 @@ def test_forward_backward_matches_reference(self): ) generator = torch.Generator(device="cuda").manual_seed(5) - for n_edge, n_nodes in ((4096, 128), (777, 33)): + # channels = 40 (non-power-of-two) rides the padded channel block + for n_edge, n_nodes, channels in ((4096, 128, 32), (777, 33, 40)): for view_radial in (False, True): with self.subTest(n_edge=n_edge, n_nodes=n_nodes, view=view_radial): zonal, radial, dst, order, row_ptr = self._make_case( - n_edge, n_nodes, generator, view_radial + n_edge, n_nodes, generator, view_radial, channels ) z1 = zonal.clone().requires_grad_(True) r1 = radial.clone().requires_grad_(True) @@ -2234,3 +2235,19 @@ def run(bound): torch.testing.assert_close(out_fused, out_ref, atol=1e-5, rtol=1e-5) torch.testing.assert_close(gd_fused, gd_ref, atol=1e-5, rtol=1e-5) torch.testing.assert_close(gr_fused, gr_ref, atol=1e-5, rtol=1e-5) + + # Training keeps the reference composition even with the operator + # bound: its backward carries no next-derivative formula, so the + # dispatch must stay out of the fused path. + module.train() + dp = d_full.clone().requires_grad_(True) + rad = radial.clone().requires_grad_(True) + out_train = module( + n_nodes=n_nodes, + edge_cache=make_cache(dp), + radial_feat=rad, + ) + gdp, grad_train = torch.autograd.grad(out_train, (dp, rad), cotangent) + torch.testing.assert_close(out_train, out_ref, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(gdp, gd_ref, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(grad_train, gr_ref, atol=1e-5, rtol=1e-5) From e26dbcc5c2700a601daf3f0cbf7c9b88a45acd33 Mon Sep 17 00:00:00 2001 From: longjunhao Date: Fri, 16 Oct 2026 21:55:20 +0800 Subject: [PATCH 5/5] style(pt): explicit index dtypes for the GIE scatter path --- deepmd/pt/model/descriptor/sezm_nn/embedding.py | 7 ++++--- deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py | 8 ++++---- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/deepmd/pt/model/descriptor/sezm_nn/embedding.py b/deepmd/pt/model/descriptor/sezm_nn/embedding.py index b405bc2092..bbe5b8e66e 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/embedding.py +++ b/deepmd/pt/model/descriptor/sezm_nn/embedding.py @@ -239,13 +239,14 @@ def __init__( gie_zonal_scatter, ) + index_device = self.non_scalar_row_index.device packed_rows = torch.arange( - 1, self.ebed_dim, device=self.non_scalar_row_index.device + 1, self.ebed_dim, device=index_device, dtype=torch.long ) counts = [2 * ell + 1 for ell in range(1, self.lmax + 1)] packed_slots = torch.repeat_interleave( - torch.arange(self.lmax, device=self.non_scalar_row_index.device), - torch.tensor(counts, device=self.non_scalar_row_index.device), + torch.arange(self.lmax, device=index_device, dtype=torch.long), + torch.tensor(counts, device=index_device, dtype=torch.long), ) if ( GIE_SCATTER_TRITON_AVAILABLE diff --git a/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py b/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py index 035de41cc7..3aeb1a8820 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py +++ b/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py @@ -69,8 +69,8 @@ def _reference( n_slot = radial.shape[1] counts = [2 * ell + 1 for ell in range(1, n_slot + 1)] slot_of_row = torch.repeat_interleave( - torch.arange(n_slot, device=zonal.device), - torch.tensor(counts, device=zonal.device), + torch.arange(n_slot, device=zonal.device, dtype=torch.long), + torch.tensor(counts, device=zonal.device, dtype=torch.long), ) msg = zonal.unsqueeze(-1) * radial[:, slot_of_row] # (E, R, C) out = zonal.new_zeros(n_nodes, n_row + 1, radial.shape[2]) @@ -408,8 +408,8 @@ def _backward_reference( n_slot = radial.shape[1] counts = [2 * ell + 1 for ell in range(1, n_slot + 1)] slot_of_row = torch.repeat_interleave( - torch.arange(n_slot, device=zonal.device), - torch.tensor(counts, device=zonal.device), + torch.arange(n_slot, device=zonal.device, dtype=torch.long), + torch.tensor(counts, device=zonal.device, dtype=torch.long), ) radial_rows = radial[:, slot_of_row] gz = (g_msg * radial_rows).sum(-1)