fix(pt-expt): harden Neo CuTe integration

Author: atulcthakur

deepmd/pt/model/descriptor/sezm.py

@@ -1802,6 +1802,8 @@ def _prepare_cute_k1_sorted_metadata(
         torch.Tensor | None,
     ]:
         """Build one per-forward CSR tensor bundle for eligible K1 blocks."""
+        if torch.jit.is_scripting():
+            return None, None, None
         if (
             self.training
             or not edge_cache.destinations_sorted

deepmd/pt_expt/kernels/cute/sezm/k1.py

@@ -15,6 +15,9 @@
 from functools import (
     lru_cache,
 )
+from itertools import (
+    count,
+)
 from types import (
     SimpleNamespace,
 )
@@ -481,7 +484,7 @@ class _RunnerState:
 
 
 _REGISTRY: dict[int, _RegistryEntry] = {}
-_NEXT_HANDLE = 1
+_HANDLE_COUNTER = count(1)
 _REGISTRY_LOCK = threading.Lock()
 _PACKED_RUNNER_CACHE: dict[tuple[str, int | None, int], _RunnerState] = {}
 _PACKED_RUNNER_CACHE_LOCK = threading.Lock()
@@ -523,7 +526,6 @@ def _compile_output_gate_backward(
 
 def register_cute_k1_block(block: Any, config: Any) -> int:
     """Register a DeePMD block/config pair and return a stable integer handle."""
-    global _NEXT_HANDLE
 
     def remove_collected_entry(block_ref: weakref.ReferenceType[Any]) -> None:
         with _REGISTRY_LOCK:
@@ -532,8 +534,7 @@ def remove_collected_entry(block_ref: weakref.ReferenceType[Any]) -> None:
                 _REGISTRY.pop(handle, None)
 
     with _REGISTRY_LOCK:
-        handle = _NEXT_HANDLE
-        _NEXT_HANDLE += 1
+        handle = next(_HANDLE_COUNTER)
         _REGISTRY[handle] = _RegistryEntry(
             block=block,
             config=config,
@@ -800,8 +801,7 @@ def _equivariant_rmsnorm_backward(norm: Any, x: Tensor, grad_out: Tensor) -> Ten
     return torch.cat([grad_x0_in, grad_xt_in], dim=1).to(dtype=in_dtype)
 
 
-def _so3_linear_backward_input(linear: Any, x: Tensor, grad_out: Tensor) -> Tensor:
-    del x
+def _so3_linear_backward_input(linear: Any, grad_out: Tensor) -> Tensor:
     weight = linear.weight.view(
         linear.lmax + 1,
         linear.in_channels,
@@ -994,14 +994,9 @@ def _so3_grid_cross_glu_flat_backward(
     coeff = torch.einsum("dkg,ngfc->ndfkc", from_grid, left_grid * right_grid)
     coeff_flat = coeff.reshape(n_batch, coeff_dim, n_focus, expanded)
 
-    scalar_out = _swiglu_forward(scalar_pair)
     scalar_logits = _focus_linear_forward(net.scalar_gate, scalar_pair)
     scalar_gate = torch.sigmoid(scalar_logits)
     coeff_view = coeff_flat.reshape(n_batch, coeff_dim, n_focus, n_frames, channels)
-    scalar_path = coeff_view * scalar_gate[:, None, :, None, :]
-    scalar_path = scalar_path.clone()
-    scalar_path[:, 0, :, net.frame_zero_index, :].add_(scalar_out)
-    scalar_path_flat = scalar_path.reshape(n_batch, coeff_dim, n_focus, expanded)
 
     grad = grad_out_flat.reshape(n_batch, coeff_dim, n_focus, channels).to(
         dtype=net.dtype
@@ -1046,7 +1041,6 @@ def _so3_grid_cross_glu_flat_backward(
     grad_context = _frame_expand_backward_input(net.frame_expand, grad_right)
     grad_query[:, 0, :, :].add_(grad_scalar_pair[:, :, :channels])
     grad_context[:, 0, :, :].add_(grad_scalar_pair[:, :, channels:])
-    del scalar_path_flat
     return (
         grad_query.reshape_as(query_flat).to(dtype=q_dtype),
         grad_context.reshape_as(context_flat).to(dtype=c_dtype),
@@ -1065,7 +1059,6 @@ def _final_manual_backward(runner: Any, grad_out: Tensor) -> tuple[Tensor, Tenso
     phase = runner.phase_c_out.detach()
     x_wide = runner.x_wide.detach()
     out_gate_flat = runner.out_gate_flat
-    post_in = runner.post_mix_input.unsqueeze(2)
     post_norm_in = runner.post_norm_input
 
     grad_post_norm_in = _equivariant_rmsnorm_backward(
@@ -1075,7 +1068,6 @@ def _final_manual_backward(runner: Any, grad_out: Tensor) -> tuple[Tensor, Tenso
     )
     grad_post_mix = _so3_linear_backward_input(
         so2.post_focus_mix,
-        post_in,
         grad_post_norm_in.squeeze(2).unsqueeze(2),
     ).squeeze(2)
 
@@ -1145,7 +1137,6 @@ def _qk_manual_backward(
 ) -> Tensor:
     so2 = runner.so2
     n_node = runner.node_count
-    n_edge = runner.edge_count
     x_wide = runner.x_wide.detach()
     x_l0 = x_wide[:, 0, :].reshape(n_node, 2, 32)
     q_node = runner.q_node
@@ -1285,20 +1276,16 @@ def _stack_backward_manual(runner: Any, grad_stack_out: Tensor) -> Tensor:
 def _x_wide_manual_backward(runner: Any, grad_x_wide_total: Tensor) -> Tensor:
     block = runner.block
     so2 = runner.so2
-    n_node = runner.node_count
     x_so2 = runner.x if runner.use_full_node else runner.x[:, : block.mp_ebed_dim, :, :]
-    x_pre = block.pre_so2_norm(x_so2)
-    x_pre_flat = x_pre.reshape(n_node, x_so2.shape[1], block.channels).unsqueeze(2)
+    if not _is_identity(block.pre_so2_norm):
+        raise NotImplementedError(
+            "manual x-wide backward currently expects Identity pre norm"
+        )
     grad_x_pre_flat = _so3_linear_backward_input(
         so2.pre_focus_mix,
-        x_pre_flat,
         grad_x_wide_total.unsqueeze(2),
     )
     grad_x_pre = grad_x_pre_flat.squeeze(2).reshape_as(x_so2)
-    if type(block.pre_so2_norm).__name__ != "Identity":
-        raise NotImplementedError(
-            "manual x-wide backward currently expects Identity pre norm"
-        )
     if runner.use_full_node:
         grad_x = grad_x_pre
     else:

deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_a_radial_forward.py

@@ -346,8 +346,10 @@ def run_neo_phase_a_radial_forward_packed_direct(
         )
     if tuple(radial_degree_mixer.weight.shape) != (4 * 64, 25):
         raise NotImplementedError("collapsed radial mixer expects lmax=3,mmax=1,C=64")
-    if tuple(radial_degree_mixer.channel_basis.shape) != (64,):
-        raise NotImplementedError("collapsed radial mixer expects 64 channel weights")
+    if tuple(radial_degree_mixer.channel_basis.shape) != (1, 64):
+        raise NotImplementedError(
+            "collapsed radial mixer expects a rank-1, 64-channel basis"
+        )
 
     combined_weight = _combined_radial_weight(
         radial_hidden_proj,

deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_onepass.py

@@ -512,6 +512,19 @@ def run_neo_phase_c_onepass_output_gate(
         raise ValueError("dst_ptr must be on the input CUDA device")
     if dst_ptr.dtype not in (torch.int32, torch.int64):
         raise TypeError(f"dst_ptr must be int32 or int64, got {dst_ptr.dtype}")
+    torch._assert_async(
+        dst_ptr[0] == 0,
+        "Phase-C forward requires dst_ptr[0] == 0",
+    )
+    torch._assert_async(
+        dst_ptr[-1] == edge_count,
+        "Phase-C forward requires dst_ptr[-1] == edge_count",
+    )
+    if dst_ptr.numel() > 1:
+        torch._assert_async(
+            torch.all(dst_ptr[1:] >= dst_ptr[:-1]),
+            "Phase-C forward requires nondecreasing dst_ptr",
+        )
 
     if out is None:
         out = torch.empty(

deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_radial_phase_a_backward_node.py

@@ -120,27 +120,25 @@ def _warp_owned_grad_compact(
     shared_row_pitch: cutlass.Constexpr[int],
 ):
     """Reduce one compact-kernel gradient inside a four-lane subgroup."""
-    in_coeff = cutlass.Int32(0)
-    out_coeff = cutlass.Int32(0)
+    value = cutlass.Float32(0.0)
     if compact_idx < 16:
         in_coeff = compact_idx // 4
         out_coeff = compact_idx - in_coeff * 4
-    else:
-        pair = compact_idx - 16
-        in_coeff = pair // 3
-        out_coeff = pair - in_coeff * 3
-
-    value = cutlass.Float32(0.0)
-    for channel_step in cutlass.range_constexpr(CHANNELS_PER_SUBGROUP_LANE):
-        hidden_channel = subgroup_lane + channel_step * WARP_REDUCTION_GROUP
-        basis = channel_basis[hidden_channel].to(cutlass.Float32)
-        if compact_idx < 16:
+        for channel_step in cutlass.range_constexpr(CHANNELS_PER_SUBGROUP_LANE):
+            hidden_channel = subgroup_lane + channel_step * WARP_REDUCTION_GROUP
+            basis = channel_basis[hidden_channel].to(cutlass.Float32)
             value += (
                 focus_grad[out_coeff * shared_row_pitch + hidden_channel]
                 * local_values[in_coeff * shared_row_pitch + hidden_channel]
                 * basis
             )
-        else:
+    else:
+        pair = compact_idx - 16
+        in_coeff = pair // 3
+        out_coeff = pair - in_coeff * 3
+        for channel_step in cutlass.range_constexpr(CHANNELS_PER_SUBGROUP_LANE):
+            hidden_channel = subgroup_lane + channel_step * WARP_REDUCTION_GROUP
+            basis = channel_basis[hidden_channel].to(cutlass.Float32)
             value += basis * (
                 focus_grad[(4 + out_coeff) * shared_row_pitch + hidden_channel]
                 * local_values[(4 + in_coeff) * shared_row_pitch + hidden_channel]

deepmd/pt_expt/kernels/cute/sezm/k1_message_grid_packed.py

@@ -7,24 +7,26 @@
 )
 
 from typing import (
-    TYPE_CHECKING,
     Any,
+    Protocol,
 )
 
 import torch
 from torch import (
     Tensor,
 )
 
-if TYPE_CHECKING:
-    from .message_grid_readout_sm90 import (
-        Sm90MessageGridState,
-    )
-
-
 COEFF_DIM = 16
 N_FOCUS = 2
 N_FRAMES = 3
+
+
+class _Sm90MessageGridState(Protocol):
+    """State subset consumed by the packed SM90 forward path."""
+
+    schedule: Tensor
+
+
 CHANNELS = 32
 HIDDEN_CHANNELS = N_FOCUS * CHANNELS
 
@@ -205,7 +207,7 @@ def run_packed_message_grid_forward(
     context_flat: Tensor,
     *,
     return_product: bool = False,
-    sm90_state: Sm90MessageGridState | None = None,
+    sm90_state: _Sm90MessageGridState | None = None,
 ) -> Tensor | tuple[Tensor, Tensor]:
     """Run only the message-grid module and return its canonical flat output."""
     query, context = _validate_contract(net, query_flat, context_flat)

deepmd/pt_expt/kernels/cute/sezm/k1_radial_phase_a_node.py

@@ -63,8 +63,11 @@ def build_source_csr(
 ) -> NeoSourceCSR:
     """Build indirect source CSR without changing the physical edge order.
 
-    ``validate_sources=True`` synchronizes when ``src`` is CUDA. Callers should
-    build this once with the edge cache and retain both tensors.
+    Source bounds are always checked before constructing the CSR. Setting
+    ``validate_sources=True`` reports an eager ``ValueError`` and therefore
+    synchronizes when ``src`` is CUDA; the default uses an asynchronous device
+    assertion. Callers should build this once with the edge cache and retain
+    both tensors.
     """
     import torch
 
@@ -78,11 +81,16 @@ def build_source_csr(
         raise ValueError("source CSR int32 indexing requires E <= 2**31 - 1")
 
     src = src.contiguous()
-    if validate_sources and src.numel() != 0:
+    if src.numel() != 0:
         valid = torch.all((src >= 0) & (src < node_count))
-        if not bool(valid):
-            raise ValueError(
-                "source-CSR backward requires source indices in [0, node_count)"
+        message = "source-CSR backward requires source indices in [0, node_count)"
+        if validate_sources:
+            if not bool(valid):
+                raise ValueError(message)
+        else:
+            torch._assert_async(
+                valid,
+                message,
             )
 
     source_order_i64 = torch.argsort(src, stable=True)
@@ -169,6 +177,10 @@ def _project_batched_radial_adjoint(
     """Project the 27-column adjoint packed in consumed edge scratch."""
     import torch
 
+    if torch.backends.cuda.matmul.allow_tf32:
+        raise RuntimeError("strict FP32 requires allow_tf32=False")
+    if torch.get_float32_matmul_precision() != "highest":
+        raise RuntimeError("strict FP32 requires float32 matmul precision 'highest'")
     edge_count = consumed_workspace.shape[0]
     device = consumed_workspace.device
     tensors = (
@@ -330,6 +342,24 @@ def run_neo_radial_phase_a_backward_node_tiled(
     )
     if node_count == 0 and edge_count != 0:
         raise ValueError("a non-empty edge list requires at least one source node")
+    torch._assert_async(
+        source_ptr[0] == 0,
+        "source-CSR backward requires source_ptr[0] == 0",
+    )
+    torch._assert_async(
+        source_ptr[-1] == edge_count,
+        "source-CSR backward requires source_ptr[-1] == edge_count",
+    )
+    if source_ptr.numel() > 1:
+        torch._assert_async(
+            torch.all(source_ptr[1:] >= source_ptr[:-1]),
+            "source-CSR backward requires nondecreasing source_ptr",
+        )
+    if source_order.numel() != 0:
+        torch._assert_async(
+            torch.all((source_order >= 0) & (source_order < edge_count)),
+            "source-CSR backward requires source_order entries in [0, E)",
+        )
     if validate_csr:
         _validate_csr_values(source_order, source_ptr, edge_count)
 

deepmd/pt_expt/kernels/cute/sezm/message_grid_gaunt_sm90.py

@@ -42,7 +42,6 @@
 CHANNELS = 64
 THREADS = 256
 GROUPS = 4
-THREADS_PER_GROUP = CHANNELS
 VALUES_PER_NODE = COEFF_DIM * CHANNELS
 EXPECTED_COMPACT_PATHS = 1968
 EXPECTED_ORDERED_PATHS = 3833

deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/cute_readout_l0.py

@@ -15,6 +15,7 @@
 
 import cutlass
 import cutlass.cute as cute
+import cutlass.pipeline as pipeline
 from cuda.bindings.driver import (
     CUstream,
 )
@@ -32,7 +33,6 @@
     TILE_K,
     TILE_M,
     TILE_N,
-    TiledOutputGridProductBackward,
 )
 
 # CuTe JIT functions use DSL-inferred argument and return types.
@@ -116,11 +116,16 @@ def kernel(
         out[node, channel] = value.to(out.element_type)
 
 
-class TiledReadoutL0GramBackward(TiledOutputGridProductBackward):
+class TiledReadoutL0GramBackward:
     """Apply the frozen 48x48 Gram matrix to one 64-channel tile."""
 
     def __init__(self) -> None:
-        super().__init__(HIDDEN_CHANNELS)
+        self.cta_tiler = (TILE_M, TILE_N, TILE_K)
+        self.channel_tiles = HIDDEN_CHANNELS // TILE_N
+        self.cta_sync_barrier = pipeline.NamedBarrier(
+            barrier_id=1,
+            num_threads=THREADS,
+        )
 
     @cute.jit
     def __call__(

deepmd/pt_expt/kernels/cute/sezm/sm90_k1/final_phase_c.py

@@ -37,6 +37,9 @@
 from ..compile_cache import (
     device_aware_lru_cache,
 )
+from ..k1_wigner_layout import (
+    PACKED_VALUE_COUNT,
+)
 
 if TYPE_CHECKING:
     from collections.abc import (
@@ -53,7 +56,7 @@
 DEGREE_COUNT = 16
 M0_WIDTH = 128
 M1_WIDTH = 96
-PACKED_WIGNER_VALUES = 46
+PACKED_WIGNER_VALUES = PACKED_VALUE_COUNT
 M0_THREADS = M0_WIDTH
 M1_THREADS = M1_WIDTH * 2
 THREADS = M0_THREADS + M1_THREADS