From 27beb6c273af9128b23a5058acb2ba2868a1865d Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sun, 11 Oct 2026 18:01:34 +0800
Subject: [PATCH 01/10] 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 296205002a..ed2460b953 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,
         *,
@@ -292,6 +320,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 629eeba2c5809c28303954fa9a0a0842fbd8fb67 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sun, 11 Oct 2026 18:01:34 +0800
Subject: [PATCH 02/10] 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 603c5969d2325a4c2fc312a1774aa7eb4edf6175 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 03/10] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt_expt/kernels/triton/sezm/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 f2251c28c454da46b707f353611becae71539cd0 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Fri, 16 Oct 2026 19:33:39 +0800
Subject: [PATCH 04/10] 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 ed2460b953..6ce294c157 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/embedding.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/embedding.py
@@ -325,8 +325,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 fa4e1dfdd20c3e30dbc8fa7bda6e0aec08ae0c12 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Fri, 16 Oct 2026 21:55:20 +0800
Subject: [PATCH 05/10] 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 6ce294c157..084d139077 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)

From 43e64375fbedb61a1a2d58561240316b69f03eb9 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sun, 18 Oct 2026 14:19:48 +0800
Subject: [PATCH 06/10] fix(pt): channel-masked radial loads, level-1 routing
 and a unified scatter boundary for the GIE head

- The backward's three radial loads now carry the shared joint
  edge/channel mask: a padded channel lane reading past the tensor
  dragged NaN into valid zonal-gradient entries (0 * NaN is NaN).
  A NaN-sentinel regression test (radial as a slice of a NaN-filled
  parent, leaf on the parent so the backward walks the padding) fails
  on the old kernel and passes on the fixed one.
- DP_GIE_SCATTER_INFER is removed; the operator rides the level-1
  universal-kernel contract (no tunable configuration, strictly removes
  the message round-trip).
- One eligibility block and one destination-CSR acquisition now serve
  both fused backends; only the backend call differs at the boundary.
- The strided-input coverage is real now: the view case reaches the
  operator as a parent-backed slice with an asserted non-contiguous
  stride; small deterministic topologies cover isolated nodes, short
  segments and a BLOCK_E boundary; module dispatch asserts which path
  ran via a counting operator.
---
 .../pt/model/descriptor/sezm_nn/embedding.py  |  89 +++++-----
 .../kernels/triton/sezm/gie_scatter.py        |  16 +-
 deepmd/pt_expt/kernels/utils.py               |  20 ---
 .../pt/model/test_descriptor_sezm_triton.py   | 158 ++++++++++++++++--
 4 files changed, 193 insertions(+), 90 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm_nn/embedding.py b/deepmd/pt/model/descriptor/sezm_nn/embedding.py
index 084d139077..6c063721ab 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/embedding.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/embedding.py
@@ -37,7 +37,7 @@
 )
 from deepmd.pt_expt.kernels.utils import (
     cuda_infer_level,
-    use_gie_scatter_infer,
+    triton_infer_level,
 )
 from deepmd.utils.version import (
     check_version_compatibility,
@@ -233,7 +233,10 @@ def __init__(
         # 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:
+        # The fused Triton scatter rides the level-1 universal-kernel
+        # contract: no tunable configuration, strictly removes the
+        # (E, D-1, C) message round-trip, so no table selection is involved.
+        if triton_infer_level() >= 1 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,
@@ -310,53 +313,42 @@ def forward(
 
         # === Step 3. Broadcast radial features per row ===
         # Each non-scalar packed row reuses the radial feature of its degree l.
-        # The fused operator spans this broadcast and the scatter of Step 4, so
-        # it takes over whenever nothing else joins the message in between.
-        if (
-            self._can_fuse_scatter(zonal_coupling)
-            and spin_l1_message is None
-            and edge_cache.csr_cache is not None
-        ):
-            return self.forward_fused_scatter(
-                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.
-        # 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
-        ):
+        # One fused operator spans this broadcast and the scatter of Step 4,
+        # so it takes over whenever nothing else joins the message in
+        # between.  The eligibility conditions and the destination CSR are
+        # shared once here; only the backend selection differs (the CUDA
+        # operator folds the degree normalization into its kernel, the
+        # Triton one leaves the scale in the traced graph so its cotangent
+        # keeps flowing).  Training keeps the reference composition for the
+        # Triton backend: the fused backward registers no next-derivative
+        # formula, and force-loss training differentiates the backward
+        # itself.
+        if spin_l1_message 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
+            if self._can_fuse_scatter(zonal_coupling):
+                return self.forward_fused_scatter(
+                    n_nodes, edge_cache, radial_feat, zonal_coupling, order, row_ptr
+                )
+            if self._triton_scatter is not None and not self.training:
+                # 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
@@ -399,6 +391,8 @@ def forward_fused_scatter(
         edge_cache: EdgeFeatureCache,
         radial_feat: torch.Tensor,
         zonal_coupling: torch.Tensor,
+        order: torch.Tensor,
+        row_ptr: torch.Tensor,
     ) -> torch.Tensor:
         """
         Build and reduce the geometric message with the fused CUDA operator.
@@ -425,13 +419,6 @@ def forward_fused_scatter(
             zonal_scatter,
         )
 
-        from .edge_cache import (
-            cached_edge_csr,
-        )
-
-        # === Step 1. Destination CSR, shared with every other edge consumer ===
-        order, row_ptr = cached_edge_csr(edge_cache, "dst", n_nodes)
-
         # === Step 2. Fused message build, reduction, padding and normalization ===
         # The operator emits the packed node layout already normalized, so the
         # scalar row and the degree scaling cost no extra pass. The scaling is
diff --git a/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py b/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py
index 3aeb1a8820..ae796747ee 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py
@@ -241,16 +241,14 @@ def _gie_scatter_bwd_kernel(
         mask=mask_e[:, None] & m2[None, :],
         other=0.0,
     )
+    # Joint edge/channel validity shared by every radial load: masked-out
+    # lanes must read zeros, because a padded-channel lane that reads past
+    # the tensor can drag NaN into a valid reduction (0 * NaN is NaN).
+    rmask = mask_e[:, None] & cmask[None, :]
     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
-    )
+    rd0 = tl.load(radial_ptr + rad_base + 0 * r_stride1, mask=rmask, other=0.0)
+    rd1 = tl.load(radial_ptr + rad_base + 1 * r_stride1, mask=rmask, other=0.0)
+    rd2 = tl.load(radial_ptr + rad_base + 2 * r_stride1, mask=rmask, other=0.0)
     gz0 = tl.sum(g0 * rd0[:, None, :], 2)
     gz1 = tl.sum(g1 * rd1[:, None, :], 2)
     gz2 = tl.sum(g2 * rd2[:, None, :], 2)
diff --git a/deepmd/pt_expt/kernels/utils.py b/deepmd/pt_expt/kernels/utils.py
index 6da551f361..c6e54bd17f 100644
--- a/deepmd/pt_expt/kernels/utils.py
+++ b/deepmd/pt_expt/kernels/utils.py
@@ -380,23 +380,3 @@ def use_amp_infer() -> bool:
         ``True`` when ``DP_AMP_INFER`` is set to a truthy value.
     """
     return os.environ.get("DP_AMP_INFER", "0").strip().lower() in _INFER_TRUE
-
-
-def use_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
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index bcb4af4baf..a5822c03d5 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2061,7 +2061,12 @@ 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, channels=32):
-        """Random operator inputs; a sliced radial view when asked."""
+        """Random operator inputs; a sliced radial view when asked.
+
+        Returns ``(zonal, radial, dst, order, row_ptr, parent)`` where
+        ``parent`` is the unsliced backing tensor (the leaf for the view
+        case, so the strided operand survives to the operator boundary).
+        """
         zonal = torch.randn(n_edge, self._ROWS, 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()
@@ -2074,7 +2079,7 @@ def _make_case(self, n_edge, n_nodes, generator, view_radial=False, channels=32)
                 torch.cumsum(counts, 0).to(torch.int64),
             ]
         )
-        return zonal, radial, dst, order, row_ptr
+        return zonal, radial, dst, order, row_ptr, parent
 
     def _reference(self, zonal, radial, dst, n_nodes):
         """Composition reference: broadcast the message, then index_add."""
@@ -2099,12 +2104,28 @@ def test_forward_backward_matches_reference(self):
         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(
+                    zonal, radial, dst, order, row_ptr, parent = self._make_case(
                         n_edge, n_nodes, generator, view_radial, channels
                     )
+                    # The view case must reach the operator boundary as a
+                    # real slice: the leaf is the parent and the operand is
+                    # the strided view (a clone here would hand the kernel a
+                    # contiguous copy and hide the strided addressing).
                     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)
+                    if view_radial:
+                        self.assertFalse(radial.is_contiguous())
+                        leaf = parent.requires_grad_(True)
+                        operand = leaf[:, 1:, :]
+                        grad_pair = (z1, leaf)
+                        grad_slice = slice(1, None)
+                    else:
+                        leaf = None
+                        operand = radial.clone().requires_grad_(True)
+                        grad_pair = (z1, operand)
+                        grad_slice = slice(None)
+                    out = gie_zonal_scatter(
+                        z1, operand, 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)
@@ -2117,10 +2138,12 @@ def test_forward_backward_matches_reference(self):
                         device="cuda",
                         generator=torch.Generator(device="cuda").manual_seed(37),
                     )
-                    g1 = torch.autograd.grad(out, (z1, r1), cotangent)
+                    g1 = torch.autograd.grad(out, grad_pair, 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)
+                    torch.testing.assert_close(
+                        g1[1][:, grad_slice], g2[1], atol=1e-5, rtol=1e-5
+                    )
 
     def test_reference_cpu_matches_manual(self):
         """CPU reference equals a hand-written sequential accumulation."""
@@ -2143,6 +2166,105 @@ def test_reference_cpu_matches_manual(self):
         )
         torch.testing.assert_close(out, manual, atol=0.0, rtol=0.0)
 
+    @_GPU_KERNELS
+    def test_small_deterministic_topology(self):
+        """Hand-built edges cover isolated nodes, short segments and a tile
+        boundary crossing (BLOCK_E = 16) with padded channels.
+        """
+        from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import (
+            gie_zonal_scatter,
+        )
+
+        generator = torch.Generator(device="cuda").manual_seed(47)
+        n_nodes, channels = 6, 40
+        # node0: a long segment; node1: isolated (no edges); node2: one
+        # edge; nodes 3-5: segments straddling the BLOCK_E boundary
+        dst = torch.tensor(
+            [0] * 18 + [2] + [3] * 7 + [4] * 8 + [5] * 9,
+            device="cuda",
+            dtype=torch.int64,
+        )
+        n_edge = dst.numel()
+        zonal = torch.randn(n_edge, self._ROWS, device="cuda", generator=generator)
+        radial = torch.randn(n_edge, 3, channels, 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),
+            ]
+        )
+
+        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)
+        # the isolated node keeps a zero row block
+        self.assertTrue(torch.equal(out[1], torch.zeros_like(out[1])))
+        cot = torch.randn_like(out)
+        g1 = torch.autograd.grad(out, (z1, r1), cot)
+        g2 = torch.autograd.grad(ref, (z2, r2), cot)
+        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)
+
+    @_GPU_KERNELS
+    def test_backward_respects_channel_padding_boundary(self):
+        """Padded channel lanes must not read past the radial tensor.
+
+        The radial operand is a slice of a NaN-filled parent sized to the
+        padded channel block: any out-of-bounds lane load reads a NaN and,
+        through ``0 * NaN``, surfaces in a valid zonal-gradient entry.  A
+        correct joint edge/channel mask keeps every gradient finite.
+        """
+        from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import (
+            gie_zonal_scatter,
+        )
+
+        generator = torch.Generator(device="cuda").manual_seed(43)
+        n_edge, n_nodes, channels = 777, 33, 40
+        zonal = torch.randn(n_edge, self._ROWS, device="cuda", generator=generator)
+        parent = torch.full(
+            (n_edge, 3, 64), float("nan"), device="cuda", dtype=torch.float32
+        )
+        radial = parent[:, :, :channels]
+        radial.copy_(
+            torch.randn(n_edge, 3, channels, device="cuda", generator=generator)
+        )
+        self.assertFalse(radial.is_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),
+            ]
+        )
+
+        # The leaf is the parent: the operator receives the strided view
+        # itself, so the backward's radial loads walk the parent's NaN
+        # padding when the channel mask is missing (a clone here would
+        # hand the backward a contiguous copy and hide the overread).
+        parent_leaf = parent.requires_grad_(True)
+        radial_view = parent_leaf[:, :, :channels]
+        z1 = zonal.clone().requires_grad_(True)
+        out = gie_zonal_scatter(z1, radial_view, dst, order, row_ptr, n_nodes, 3)
+        cot = torch.randn_like(out)
+        gz, gr = torch.autograd.grad(out, (z1, parent_leaf), cot)
+        self.assertTrue(torch.isfinite(gz).all())
+        self.assertTrue(torch.isfinite(gr).all())
+        # The reference composition on the same inputs stays the oracle.
+        z2 = zonal.clone().requires_grad_(True)
+        r2 = radial.clone().requires_grad_(True)
+        ref = self._reference(z2, r2, dst, n_nodes)
+        ref_gz, ref_gr = torch.autograd.grad(ref, (z2, r2), cot)
+        torch.testing.assert_close(gz, ref_gz, atol=1e-5, rtol=1e-5)
+        torch.testing.assert_close(gr[:, :, :channels], ref_gr, atol=1e-5, rtol=1e-5)
+
     @_GPU_KERNELS
     def test_compile_matches_eager(self):
         """The compiled path must match eager (guards the launch-grid trap)."""
@@ -2151,7 +2273,7 @@ def test_compile_matches_eager(self):
         )
 
         generator = torch.Generator(device="cuda").manual_seed(11)
-        zonal, radial, dst, order, row_ptr = self._make_case(
+        zonal, radial, dst, order, row_ptr, _parent = self._make_case(
             4096, 128, generator, view_radial=True
         )
 
@@ -2185,7 +2307,9 @@ def test_embedding_dispatch_two_state(self):
 
         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)
+        zonal, radial, dst, order, row_ptr, _parent = 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
@@ -2210,9 +2334,19 @@ def make_cache(df):
         ).cuda()
         module.eval()
         cotangent = None
+        calls = {"n": 0}
+        real_op = gie_zonal_scatter
+
+        def counting_op(*args, **kwargs):
+            calls["n"] += 1
+            return real_op(*args, **kwargs)
 
         def run(bound):
-            module._triton_scatter = gie_zonal_scatter if bound else None
+            # Counting wrapper: the assertions check which path was
+            # actually selected, not just that outputs agree (a silently
+            # reference-path run would also produce equal outputs).
+            module._triton_scatter = counting_op if bound else None
+            calls["n"] = 0
             df = d_full.clone().requires_grad_(True)
             rad = radial.clone().requires_grad_(True)
             out = module(
@@ -2231,7 +2365,9 @@ def run(bound):
             return out.detach(), grads
 
         out_ref, (gd_ref, gr_ref) = run(bound=False)
+        self.assertEqual(calls["n"], 0)
         out_fused, (gd_fused, gr_fused) = run(bound=True)
+        self.assertEqual(calls["n"], 1)
         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)
@@ -2240,6 +2376,7 @@ def run(bound):
         # bound: its backward carries no next-derivative formula, so the
         # dispatch must stay out of the fused path.
         module.train()
+        calls["n"] = 0
         dp = d_full.clone().requires_grad_(True)
         rad = radial.clone().requires_grad_(True)
         out_train = module(
@@ -2248,6 +2385,7 @@ def run(bound):
             radial_feat=rad,
         )
         gdp, grad_train = torch.autograd.grad(out_train, (dp, rad), cotangent)
+        self.assertEqual(calls["n"], 0)  # training keeps the reference path
         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 71035233564c9338c746d7994114df18426eea56 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sun, 18 Oct 2026 17:06:48 +0800
Subject: [PATCH 07/10] fix(pt): guard the optional Triton import and defer CSR
 acquisition to a fused backend

---
 .../pt/model/descriptor/sezm_nn/embedding.py  | 32 +++++++++++++------
 1 file changed, 22 insertions(+), 10 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm_nn/embedding.py b/deepmd/pt/model/descriptor/sezm_nn/embedding.py
index 6c063721ab..dd1e7caae3 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/embedding.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/embedding.py
@@ -237,10 +237,17 @@ def __init__(
         # contract: no tunable configuration, strictly removes the
         # (E, D-1, C) message round-trip, so no table selection is involved.
         if triton_infer_level() >= 1 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,
-            )
+            try:
+                from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import (
+                    GIE_SCATTER_TRITON_AVAILABLE,
+                    gie_zonal_scatter,
+                )
+            except ImportError:
+                # The Triton stack is optional; a level-1 environment
+                # without it keeps the reference composition (and the CUDA
+                # operator above remains unaffected).
+                GIE_SCATTER_TRITON_AVAILABLE = False
+                gie_zonal_scatter = None
 
             index_device = self.non_scalar_row_index.device
             packed_rows = torch.arange(
@@ -324,16 +331,21 @@ def forward(
         # formula, and force-loss training differentiates the backward
         # itself.
         if spin_l1_message is None and edge_cache.csr_cache is not None:
-            from .edge_cache import (
-                cached_edge_csr,
-            )
+            cuda_selected = self._can_fuse_scatter(zonal_coupling)
+            triton_selected = self._triton_scatter is not None and not self.training
+            if cuda_selected or triton_selected:
+                from .edge_cache import (
+                    cached_edge_csr,
+                )
 
-            order, row_ptr = cached_edge_csr(edge_cache, "dst", n_nodes)
-            if self._can_fuse_scatter(zonal_coupling):
+                # The destination view is only needed by a fused backend;
+                # the reference path below must not pay for its sort.
+                order, row_ptr = cached_edge_csr(edge_cache, "dst", n_nodes)
+            if cuda_selected:
                 return self.forward_fused_scatter(
                     n_nodes, edge_cache, radial_feat, zonal_coupling, order, row_ptr
                 )
-            if self._triton_scatter is not None and not self.training:
+            if triton_selected:
                 # 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

From f687a614508442ea8c71b1896f304c04d75cff8d Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Fri, 30 Oct 2026 19:50:00 +0800
Subject: [PATCH 08/10] fix(sezm): guard the gie_scatter triton import at
 module level

The module now follows the package's optional-dependency convention: the
triton import sits in try/except, the kernel definitions and operator
registrations live under the availability flag, and the flag is False
whenever the stack is absent, so CPU-only environments import the module
instead of crashing at collection time.  The embedding binding reads the
flag directly (the consumer-side try/except and its unused fallback
assignment are gone) and keeps the CUDA requirement at the binding,
preserving the previous effective condition.  The package __init__
aggregates the new flag like every sibling's.

Two new tests: the degree-normalization cotangent guard (the scaling
deliberately stays in the traced graph; folding it into the kernel
without a gradient path now fails the suite, mirroring the CUDA
operator's guard) and a subprocess regression that blocks the triton
import and checks flag/ reference/ loud-error behavior.
---
 .../pt/model/descriptor/sezm_nn/embedding.py  |  19 +-
 .../pt_expt/kernels/triton/sezm/__init__.py   |   5 +
 .../kernels/triton/sezm/gie_scatter.py        | 739 +++++++++---------
 .../pt/model/test_descriptor_sezm_triton.py   | 135 ++++
 4 files changed, 520 insertions(+), 378 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm_nn/embedding.py b/deepmd/pt/model/descriptor/sezm_nn/embedding.py
index dd1e7caae3..32813ec22c 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/embedding.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/embedding.py
@@ -237,17 +237,13 @@ def __init__(
         # contract: no tunable configuration, strictly removes the
         # (E, D-1, C) message round-trip, so no table selection is involved.
         if triton_infer_level() >= 1 and self.dtype is torch.float32 and self.lmax == 3:
-            try:
-                from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import (
-                    GIE_SCATTER_TRITON_AVAILABLE,
-                    gie_zonal_scatter,
-                )
-            except ImportError:
-                # The Triton stack is optional; a level-1 environment
-                # without it keeps the reference composition (and the CUDA
-                # operator above remains unaffected).
-                GIE_SCATTER_TRITON_AVAILABLE = False
-                gie_zonal_scatter = None
+            # The module owns its optional-triton import guard, so this
+            # import succeeds everywhere and the exported flag is the
+            # single source of truth for the capability check below.
+            from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import (
+                GIE_SCATTER_TRITON_AVAILABLE,
+                gie_zonal_scatter,
+            )
 
             index_device = self.non_scalar_row_index.device
             packed_rows = torch.arange(
@@ -260,6 +256,7 @@ def __init__(
             )
             if (
                 GIE_SCATTER_TRITON_AVAILABLE
+                and torch.cuda.is_available()
                 and torch.equal(self.non_scalar_row_index, packed_rows)
                 and torch.equal(self.radial_slot_index_for_row, packed_slots)
             ):
diff --git a/deepmd/pt_expt/kernels/triton/sezm/__init__.py b/deepmd/pt_expt/kernels/triton/sezm/__init__.py
index a0c3b2c7a1..e8b04838f8 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/__init__.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/__init__.py
@@ -9,6 +9,9 @@
 from .force_assembly import (
     FORCE_ASSEMBLY_TRITON_AVAILABLE,
 )
+from .gie_scatter import (
+    GIE_SCATTER_TRITON_AVAILABLE,
+)
 from .radial_mix import (
     RADIAL_MIX_TRITON_AVAILABLE,
 )
@@ -39,8 +42,10 @@
     and STACK_FP16X3_TRITON_AVAILABLE
     and WIGNER_MONOMIALS_TRITON_AVAILABLE
     and FORCE_ASSEMBLY_TRITON_AVAILABLE
+    and GIE_SCATTER_TRITON_AVAILABLE
 )
 
 __all__ = [
+    "GIE_SCATTER_TRITON_AVAILABLE",
     "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 ae796747ee..46283fc349 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py
@@ -30,33 +30,29 @@
 )
 
 import torch
-import triton
-import triton.language as tl
 from torch import (
     Tensor,
 )
-from torch.library import (
-    wrap_triton,
-)
-from triton import (
-    cdiv,
-)
+
+try:
+    import triton
+    import triton.language as tl
+    from torch.library import (
+        wrap_triton,
+    )
+    from triton import (
+        cdiv,
+    )
+
+    GIE_SCATTER_TRITON_AVAILABLE = True
+except ImportError:  # pragma: no cover - exercised only without triton
+    GIE_SCATTER_TRITON_AVAILABLE = False
 
 __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,
@@ -78,321 +74,6 @@ def _reference(
     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,
-    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, 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)
-    r1 = tl.arange(0, P1)
-    r2 = tl.arange(0, P2)
-    m0 = r0 < R0
-    m1 = r1 < R1
-    m2 = r2 < R2
-    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
-        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, 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((BC,), dtype=tl.float32),
-        mask=cmask,
-    )
-    tl.store(
-        out_ptr + out_base + ((1 + B0 + r0) * C)[:, None] + offs_c[None, :],
-        acc0,
-        mask=m0[:, None] & cmask[None, :],
-    )
-    tl.store(
-        out_ptr + out_base + ((1 + B1 + r1) * C)[:, None] + offs_c[None, :],
-        acc1,
-        mask=m1[:, None] & cmask[None, :],
-    )
-    tl.store(
-        out_ptr + out_base + ((1 + B2 + r2) * C)[:, None] + offs_c[None, :],
-        acc2,
-        mask=m2[:, None] & cmask[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,
-    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, BC)
-    cmask = offs_c < 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] & 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] & 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] & cmask[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,
-    )
-    # Joint edge/channel validity shared by every radial load: masked-out
-    # lanes must read zeros, because a padded-channel lane that reads past
-    # the tensor can drag NaN into a valid reduction (0 * NaN is NaN).
-    rmask = mask_e[:, None] & cmask[None, :]
-    rad_base = (offs_e * r_stride0)[:, None] + offs_c[None, :]
-    rd0 = tl.load(radial_ptr + rad_base + 0 * r_stride1, mask=rmask, other=0.0)
-    rd1 = tl.load(radial_ptr + rad_base + 1 * r_stride1, mask=rmask, other=0.0)
-    rd2 = tl.load(radial_ptr + rad_base + 2 * r_stride1, mask=rmask, 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] & 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(
-    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,
-        BC=triton.next_power_of_2(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,
-        BC=triton.next_power_of_2(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,
@@ -421,44 +102,363 @@ def _backward_reference(
     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
+if GIE_SCATTER_TRITON_AVAILABLE:
+
+    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
+
+    @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,
+        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, 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)
+        r1 = tl.arange(0, P1)
+        r2 = tl.arange(0, P2)
+        m0 = r0 < R0
+        m1 = r1 < R1
+        m2 = r2 < R2
+        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
+            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, 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((BC,), dtype=tl.float32),
+            mask=cmask,
+        )
+        tl.store(
+            out_ptr + out_base + ((1 + B0 + r0) * C)[:, None] + offs_c[None, :],
+            acc0,
+            mask=m0[:, None] & cmask[None, :],
+        )
+        tl.store(
+            out_ptr + out_base + ((1 + B1 + r1) * C)[:, None] + offs_c[None, :],
+            acc1,
+            mask=m1[:, None] & cmask[None, :],
+        )
+        tl.store(
+            out_ptr + out_base + ((1 + B2 + r2) * C)[:, None] + offs_c[None, :],
+            acc2,
+            mask=m2[:, None] & cmask[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,
+        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, BC)
+        cmask = offs_c < 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] & 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] & 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] & cmask[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,
+        )
+        # Joint edge/channel validity shared by every radial load: masked-out
+        # lanes must read zeros, because a padded-channel lane that reads past
+        # the tensor can drag NaN into a valid reduction (0 * NaN is NaN).
+        rmask = mask_e[:, None] & cmask[None, :]
+        rad_base = (offs_e * r_stride0)[:, None] + offs_c[None, :]
+        rd0 = tl.load(radial_ptr + rad_base + 0 * r_stride1, mask=rmask, other=0.0)
+        rd1 = tl.load(radial_ptr + rad_base + 1 * r_stride1, mask=rmask, other=0.0)
+        rd2 = tl.load(radial_ptr + rad_base + 2 * r_stride1, mask=rmask, 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] & 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(
+        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,
+            BC=triton.next_power_of_2(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(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
+    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,
+            BC=triton.next_power_of_2(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
 
+    _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,
-)
+    torch.library.register_autograd(
+        "sezm_triton::gie_zonal_scatter",
+        _backward,
+        setup_context=_setup_context,
+    )
 
 
 def gie_zonal_scatter(
@@ -494,4 +494,9 @@ def gie_zonal_scatter(
     torch.Tensor
         ``(N, R+1, C)`` node features with a zero scalar row, unnormalized.
     """
+    if not GIE_SCATTER_TRITON_AVAILABLE:
+        raise RuntimeError(
+            "the fused GIE scatter needs the triton stack; the reference "
+            "composition covers environments without it"
+        )
     return _gie_scatter_op(zonal, radial, dst, order, row_ptr, n_nodes, lmax)
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index a5822c03d5..8dcd2023c6 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2145,6 +2145,73 @@ def test_forward_backward_matches_reference(self):
                         g1[1][:, grad_slice], g2[1], atol=1e-5, rtol=1e-5
                     )
 
+    @_GPU_KERNELS
+    def test_backward_reaches_the_degree_normalization(self):
+        """The degree normalization stays outside the operator on purpose.
+
+        The scaling descends from the cutoff envelope, so its cotangent is
+        part of the force.  The Triton path deliberately leaves the
+        ``out * inv_sqrt_deg`` multiply to the traced graph; this check
+        fails if a future refactor folds it into the kernel without a
+        gradient path for the scale (mirroring the CUDA operator's guard).
+        """
+        from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import (
+            gie_zonal_scatter,
+        )
+
+        generator = torch.Generator(device="cuda").manual_seed(11)
+        n_edge, n_nodes = 2048, 64
+        zonal, radial, dst, order, row_ptr, _ = self._make_case(
+            n_edge, n_nodes, generator
+        )
+        inv_sqrt_deg = torch.rand(
+            n_nodes, device="cuda", generator=generator
+        ).clamp_min(0.1)
+        cotangent = torch.randn(
+            n_nodes,
+            self._ROWS + 1,
+            radial.shape[2],
+            device="cuda",
+            generator=torch.Generator(device="cuda").manual_seed(13),
+        )
+
+        grads = {}
+        for name, compose in (
+            (
+                "fused",
+                lambda z, r, sc: (
+                    gie_zonal_scatter(z, r, dst, order, row_ptr, n_nodes, 3)
+                    * sc[:, None, None]
+                ),
+            ),
+            (
+                "reference",
+                lambda z, r, sc: (
+                    self._reference(z, r, dst, n_nodes) * sc[:, None, None]
+                ),
+            ),
+        ):
+            z = zonal.clone().requires_grad_(True)
+            r = radial.clone().requires_grad_(True)
+            sc = inv_sqrt_deg.clone().requires_grad_(True)
+            compose(z, r, sc).backward(cotangent)
+            grads[name] = (z.grad, r.grad, sc.grad)
+
+        self.assertIsNotNone(
+            grads["fused"][2],
+            "the scale gradient must flow; a fold without a gradient path "
+            "silently biases inference forces",
+        )
+        torch.testing.assert_close(
+            grads["fused"][2], grads["reference"][2], atol=1e-5, rtol=1e-5
+        )
+        torch.testing.assert_close(
+            grads["fused"][0], grads["reference"][0], atol=1e-5, rtol=1e-5
+        )
+        torch.testing.assert_close(
+            grads["fused"][1], grads["reference"][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 (
@@ -2166,6 +2233,74 @@ def test_reference_cpu_matches_manual(self):
         )
         torch.testing.assert_close(out, manual, atol=0.0, rtol=0.0)
 
+    def test_module_imports_and_falls_back_without_triton(self):
+        """The package contract: the module imports without the triton stack.
+
+        CPU-only torch wheels ship no triton; every sibling kernel module
+        guards its import and exports a meaningful availability flag.  A
+        child process blocks the triton import and checks the flag flips to
+        False, the reference composition still runs, and calling the fused
+        entry raises a loud error instead of an import-time crash.
+        """
+        import pathlib
+        import subprocess
+        import sys
+        import textwrap
+
+        import deepmd.pt_expt.kernels.triton.sezm.gie_scatter as gie_mod
+
+        tree = pathlib.Path(gie_mod.__file__).resolve().parents[5]
+        child = textwrap.dedent(
+            """
+            import pathlib
+            import sys
+
+            tree = pathlib.Path(sys.argv[1])
+            sys.meta_path = [
+                f
+                for f in sys.meta_path
+                if "editable"
+                not in (type(f).__module__ + "." + type(f).__name__).lower()
+                and "editable" not in getattr(f, "__name__", "").lower()
+            ]
+            sys.path.insert(0, str(tree))
+            sys.modules["triton"] = None  # any `import triton` now fails
+
+            import torch
+
+            import deepmd.pt_expt.kernels.triton.sezm.gie_scatter as gs
+
+            assert gs.GIE_SCATTER_TRITON_AVAILABLE is False
+
+            zonal = torch.ones(2, 3, device="cpu")
+            radial = torch.ones(2, 1, 4, device="cpu")
+            dst = torch.tensor([0, 1])
+            out = gs._reference(zonal, radial, dst, 2)
+            assert out.shape == (2, 4, 4)
+            assert float(out[0, 0, 0]) == 0.0
+
+            try:
+                gs.gie_zonal_scatter(zonal, radial, dst, None, None, 2, 1)
+            except RuntimeError:
+                pass
+            else:
+                raise AssertionError("expected a loud RuntimeError")
+            print("no-triton-ok")
+            """
+        )
+        result = subprocess.run(
+            [sys.executable, "-c", child, str(tree)],
+            capture_output=True,
+            text=True,
+            timeout=300,
+        )
+        self.assertEqual(
+            result.returncode,
+            0,
+            f"no-triton child failed:\n{result.stdout}\n{result.stderr}",
+        )
+        self.assertIn("no-triton-ok", result.stdout)
+
     @_GPU_KERNELS
     def test_small_deterministic_topology(self):
         """Hand-built edges cover isolated nodes, short segments and a tile

From f835ca0d8dda9855818b71ebf92c0ba854c1bbf1 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Sun, 1 Nov 2026 11:24:11 +0800
Subject: [PATCH 09/10] fix(sezm): widen the GIE scatter row base to 64 bits

tl.program_id returns int32, and n * (N_ROW + 1) * C kept that width, so
at node counts past 2**31/(rows*C) (~4.19M nodes at the lmax==3, C=32
gate) the offset wrapped and the row stores landed up to ~8.6 GB outside
the output buffer -- an illegal access or silent corruption of a
neighbouring allocation.  The node id is now cast to int64 before the
multiply.

The regression test pre-dirties the wrap-scale buffer with NaN through
the caching allocator (every node program writes its full row block, so
a surviving NaN marks a row whose stores left the buffer) and checks
the first wrapping node and the last node read back zeros; with the
int32 base restored the test fails.
---
 .../kernels/triton/sezm/gie_scatter.py        |  5 +-
 .../pt/model/test_descriptor_sezm_triton.py   | 49 +++++++++++++++++++
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py b/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py
index 46283fc349..0ff5032a23 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py
@@ -136,7 +136,10 @@ def _gie_scatter_fwd_kernel(
         BC: tl.constexpr,
     ):
         # tl.arange needs a power-of-two extent, so channels ride a padded block
-        n = tl.program_id(0)
+        # The node id must be 64-bit before the row-base multiply: program ids
+        # are int32, and n * (N_ROW + 1) * C wraps past ~2**31/512 nodes,
+        # sending the row stores outside the output buffer.
+        n = tl.program_id(0).to(tl.int64)
         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, BC)
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index 8dcd2023c6..b127c0fe4e 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -2212,6 +2212,55 @@ def test_backward_reaches_the_degree_normalization(self):
             grads["fused"][1], grads["reference"][1], atol=1e-5, rtol=1e-5
         )
 
+    @_GPU_KERNELS
+    def test_forward_row_base_survives_past_int32(self):
+        """The row base is 64-bit: nodes past 2**31/(rows*C) must not wrap.
+
+        At the wrap scale the output itself is ~8 GiB, so the check is
+        gated on free VRAM.  The buffer is pre-dirtied with NaN through the
+        caching allocator (fill, free, let the operator's allocation reuse
+        the block): every node program writes its full 16-row block, so any
+        surviving NaN marks a row whose stores landed outside the buffer --
+        exactly what an int32 row base does past the threshold.
+        """
+        from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import (
+            gie_zonal_scatter,
+        )
+
+        free, _total = torch.cuda.mem_get_info()
+        channels = 32
+        row_elems = 16 * channels  # (N_ROW + 1) * C at the lmax == 3 gate
+        n_nodes = 2**31 // row_elems + 16
+        need_bytes = n_nodes * row_elems * 4
+        if free < need_bytes + (2 << 30):
+            self.skipTest(f"needs ~{need_bytes >> 30} GiB free VRAM")
+
+        generator = torch.Generator(device="cuda").manual_seed(19)
+        n_edge = 8
+        zonal = torch.randn(n_edge, self._ROWS, device="cuda", generator=generator)
+        radial = torch.randn(n_edge, 3, channels, device="cuda", generator=generator)
+        dst = torch.zeros(n_edge, dtype=torch.int64, device="cuda")
+        order = torch.arange(n_edge, dtype=torch.int64, device="cuda")
+        row_ptr = torch.full((n_nodes + 1,), n_edge, dtype=torch.int64, device="cuda")
+        row_ptr[0] = 0
+
+        dirty = torch.full((n_nodes, 16, channels), float("nan"), device="cuda")
+        del dirty
+        out = gie_zonal_scatter(zonal, radial, dst, order, row_ptr, n_nodes, 3)
+
+        self.assertFalse(torch.isnan(out).any().item(), "unwritten row detected")
+        slot_of_row = torch.repeat_interleave(
+            torch.arange(3, device="cuda"),
+            torch.tensor([3, 5, 7], device="cuda"),
+        )
+        manual = (zonal.unsqueeze(-1) * radial[:, slot_of_row]).sum(0)
+        torch.testing.assert_close(out[0, 1:, :], manual, atol=1e-5, rtol=1e-5)
+        for probe in (2**31 // row_elems, n_nodes - 1):
+            with self.subTest(node=probe):
+                self.assertEqual(float(out[probe].abs().max()), 0.0)
+        del out
+        torch.cuda.empty_cache()
+
     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 (

From 1745fc2d92d4fb231007793972fd69af1da80010 Mon Sep 17 00:00:00 2001
From: longjunhao <longjunhao@localhost>
Date: Mon, 2 Nov 2026 14:12:36 +0800
Subject: [PATCH 10/10] fix(sezm): widen the GIE backward edge base and gate
 the op on device

The backward widened its edge offsets only after the block multiply:
tl.program_id is int32, so pid * BLOCK_E wrapped past 2**31 / BLOCK_E
blocks and the upper-bound-only edge mask let the negative offsets
through to loads and stores.  The program id is now widened first,
mirroring the forward.

The embedding call site now also requires CUDA operands (the CUDA
sibling's contract): the constructor binds on host availability, but a
GPU host running device-qualified CPU inference sent CPU tensors into
the operator, whose internal fallback builds data-dependent shapes that
break compiled CPU graphs.

Tests: the compile test now traces at one node count, re-runs the
compiled callable at a second one and asserts the operator node is
captured (fullgraph + a capture backend), so the launch-grid trap it
names can actually fire; the constructor binding conditions themselves
are exercised (level >= 1 + fp32 + lmax == 3 binds, level 0 / fp64 /
other ladders never do); and CPU operands never reach the operator
while producing the reference composition's values (dropping the
device term turns it red).
---
 .../pt/model/descriptor/sezm_nn/embedding.py  |  11 +-
 .../kernels/triton/sezm/gie_scatter.py        |   9 +-
 .../pt/model/test_descriptor_sezm_triton.py   | 157 ++++++++++++++++--
 3 files changed, 155 insertions(+), 22 deletions(-)

diff --git a/deepmd/pt/model/descriptor/sezm_nn/embedding.py b/deepmd/pt/model/descriptor/sezm_nn/embedding.py
index 32813ec22c..7b3c5cc776 100644
--- a/deepmd/pt/model/descriptor/sezm_nn/embedding.py
+++ b/deepmd/pt/model/descriptor/sezm_nn/embedding.py
@@ -329,7 +329,16 @@ def forward(
         # itself.
         if spin_l1_message is None and edge_cache.csr_cache is not None:
             cuda_selected = self._can_fuse_scatter(zonal_coupling)
-            triton_selected = self._triton_scatter is not None and not self.training
+            # The Triton operator serves the same CUDA-only contract: the
+            # constructor binds on host availability, but a GPU host may
+            # still run this forward on CPU tensors (device-qualified
+            # inference), and the operator's internal fallback builds
+            # data-dependent shapes that break compiled CPU graphs.
+            triton_selected = (
+                self._triton_scatter is not None
+                and not self.training
+                and zonal_coupling.is_cuda
+            )
             if cuda_selected or triton_selected:
                 from .edge_cache import (
                     cached_edge_csr,
diff --git a/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py b/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py
index 0ff5032a23..a946d86042 100644
--- a/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py
+++ b/deepmd/pt_expt/kernels/triton/sezm/gie_scatter.py
@@ -229,9 +229,12 @@ def _gie_scatter_bwd_kernel(
         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)
+        # tl.arange needs a power-of-two extent, so channels ride a padded
+        # block.  The program id is int32 and must be widened before the
+        # block multiply: at pid = 2**31 / BLOCK_E the product wraps negative,
+        # and the upper-bound-only edge mask lets the wrapped offsets through.
+        pid = tl.program_id(0).to(tl.int64)
+        offs_e = pid * BLOCK_E + tl.arange(0, BLOCK_E)
         mask_e = offs_e < n_edge
         offs_c = tl.arange(0, BC)
         cmask = offs_c < C
diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py
index b127c0fe4e..faefb36f5a 100644
--- a/source/tests/pt/model/test_descriptor_sezm_triton.py
+++ b/source/tests/pt/model/test_descriptor_sezm_triton.py
@@ -23,6 +23,7 @@
 """
 
 import math
+import os
 import typing
 import unittest
 from unittest import (
@@ -2450,34 +2451,154 @@ def test_backward_respects_channel_padding_boundary(self):
         torch.testing.assert_close(gr[:, :, :channels], ref_gr, atol=1e-5, rtol=1e-5)
 
     @_GPU_KERNELS
-    def test_compile_matches_eager(self):
-        """The compiled path must match eager (guards the launch-grid trap)."""
+    def test_compile_matches_eager_across_node_counts(self):
+        """The compiled path matches eager and re-runs at a second shape.
+
+        The launch grid is derived from ``row_ptr`` inside the operator's
+        Python wrapper, so the trap this guards on only appears when a
+        compiled graph is re-run at a different node count: the callable
+        is compiled at one shape and re-run at another.  A capture backend
+        first asserts the operator is actually captured (a silent graph
+        break would run eager and keep every value comparison green), and
+        the value runs use ``fullgraph=True`` for the same reason.
+        """
         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, _parent = self._make_case(
-            4096, 128, generator, view_radial=True
+        shapes = ((4096, 128), (5000, 131))
+        cases = []
+        for n_edge, n_nodes in shapes:
+            zonal, radial, dst, order, row_ptr, _parent = self._make_case(
+                n_edge, n_nodes, generator, view_radial=True
+            )
+            cases.append((zonal, radial, dst, order, row_ptr, n_nodes))
+
+        def fn(z, r, dst_, order_, row_ptr_, n_nodes_):
+            return gie_zonal_scatter(z, r, dst_, order_, row_ptr_, n_nodes_, 3)
+
+        graphs = []
+
+        def _capture(gm, example_inputs):
+            graphs.append(gm)
+            return gm.forward
+
+        first = cases[0]
+        torch.compile(fn, backend=_capture, dynamic=True)(*first[:2], *first[2:])
+        self.assertTrue(
+            any("gie_zonal_scatter" in g.code for g in graphs),
+            "the compiled graph must capture the operator node",
+        )
+
+        compiled = torch.compile(fn, dynamic=True, fullgraph=True)
+        for zonal, radial, dst, order, row_ptr, n_nodes in cases:
+            with self.subTest(n_nodes=n_nodes):
+                z_e = zonal.clone().requires_grad_(True)
+                r_e = radial.clone().requires_grad_(True)
+                out_e = fn(z_e, r_e, dst, order, row_ptr, n_nodes)
+                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 = compiled(z_c, r_c, dst, order, row_ptr, n_nodes)
+                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_constructor_binds_the_fused_scatter_only_on_the_supported_ladder(
+        self,
+    ):
+        """The binding conditions themselves, not the post-hoc override.
+
+        Level >= 1 with fp32 and lmax == 3 binds the fused entry; level 0,
+        a wider dtype or a different ladder degree never does.
+        """
+        from deepmd.pt.model.descriptor.sezm_nn.embedding import (
+            GeometricInitialEmbedding,
+        )
+        from deepmd.pt_expt.kernels.triton.sezm.gie_scatter import (
+            gie_zonal_scatter,
         )
 
-        def fn(z, r):
-            return gie_zonal_scatter(z, r, dst, order, row_ptr, 128, 3)
+        with mock.patch.dict(os.environ, {"DP_TRITON_INFER": "1"}, clear=False):
+            bound = GeometricInitialEmbedding(
+                lmax=3, channels=32, dtype=torch.float32
+            ).cuda()
+            self.assertIs(bound._triton_scatter, gie_zonal_scatter)
 
-        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)
+            wider = GeometricInitialEmbedding(
+                lmax=3, channels=32, dtype=torch.float64
+            ).cuda()
+            other_ladder = GeometricInitialEmbedding(
+                lmax=2, channels=32, dtype=torch.float32
+            ).cuda()
+        self.assertIsNone(wider._triton_scatter, "fp64 must not bind")
+        self.assertIsNone(other_ladder._triton_scatter, "lmax != 3 must not bind")
 
-        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)
+        with mock.patch.dict(os.environ, {"DP_TRITON_INFER": "0"}, clear=False):
+            off = GeometricInitialEmbedding(
+                lmax=3, channels=32, dtype=torch.float32
+            ).cuda()
+        self.assertIsNone(off._triton_scatter, "level 0 must not bind")
 
-        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_cpu_tensors_never_enter_the_triton_scatter(self):
+        """A CUDA host may still run device-qualified CPU inference.
+
+        The constructor binds on host availability, but CPU operands must
+        keep the module's own reference composition: the operator's
+        internal fallback builds data-dependent shapes that break compiled
+        CPU graphs, and the kernel cannot run on CPU tensors at all.
+        """
+        from deepmd.pt.model.descriptor.sezm_nn.embedding import (
+            GeometricInitialEmbedding,
+        )
+
+        with mock.patch.dict(os.environ, {"DP_TRITON_INFER": "1"}, clear=False):
+            module = GeometricInitialEmbedding(
+                lmax=3, channels=32, dtype=torch.float32
+            ).cpu()
+        module.eval()
+        self.assertIsNotNone(module._triton_scatter, "precondition: bound on host")
+
+        calls = []
+        real = module._triton_scatter
+        module._triton_scatter = lambda *a, **k: (calls.append(1), real(*a, **k))[1]
+
+        n_edge, n_nodes = 257, 9
+        generator = torch.Generator(device="cpu").manual_seed(17)
+        dst = torch.randint(0, n_nodes, (n_edge,), generator=generator, device="cpu")
+        d_full = torch.randn(n_edge, 16, 16, generator=generator, device="cpu")
+        radial = torch.randn(n_edge, 3, 32, generator=generator, device="cpu")
+
+        class _Cache:
+            pass
+
+        cache = _Cache()
+        cache.dst = dst
+        cache.edge_vec = torch.zeros(n_edge, 3, device="cpu")
+        cache.inv_sqrt_deg = torch.rand(
+            n_nodes, 1, 1, generator=generator, device="cpu"
+        )
+        cache.csr_cache = {}
+        cache.edge_src_gate = None
+        cache.Dt_full = d_full
+
+        out_bound = module.forward(
+            n_nodes=n_nodes, edge_cache=cache, radial_feat=radial
+        )
+        self.assertEqual(calls, [], "CPU operands must not reach the operator")
+
+        module._triton_scatter = None
+        out_reference = module.forward(
+            n_nodes=n_nodes, edge_cache=cache, radial_feat=radial
+        )
+        torch.testing.assert_close(out_bound, out_reference, atol=0.0, rtol=0.0)
 
     @_GPU_KERNELS
     def test_embedding_dispatch_two_state(self):