fix(pt): address Neo CuTe review findings

Author: atulcthakur

deepmd/kernels/cute/neo/k1.py

@@ -482,6 +482,7 @@ class _RunnerState:
 
 _REGISTRY: dict[int, _RegistryEntry] = {}
 _NEXT_HANDLE = 1
+_REGISTRY_LOCK = threading.Lock()
 _PACKED_RUNNER_CACHE: dict[tuple[str, int | None, int], _RunnerState] = {}
 _PACKED_RUNNER_CACHE_LOCK = threading.Lock()
 
@@ -523,27 +524,30 @@ 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
-    handle = _NEXT_HANDLE
-    _NEXT_HANDLE += 1
 
     def remove_collected_entry(block_ref: weakref.ReferenceType[Any]) -> None:
-        entry = _REGISTRY.get(handle)
-        if entry is not None and entry._block_ref is block_ref:
-            _REGISTRY.pop(handle, None)
-
-    _REGISTRY[handle] = _RegistryEntry(
-        block=block,
-        config=config,
-        on_collect=remove_collected_entry,
-    )
+        with _REGISTRY_LOCK:
+            entry = _REGISTRY.get(handle)
+            if entry is not None and entry._block_ref is block_ref:
+                _REGISTRY.pop(handle, None)
+
+    with _REGISTRY_LOCK:
+        handle = _NEXT_HANDLE
+        _NEXT_HANDLE += 1
+        _REGISTRY[handle] = _RegistryEntry(
+            block=block,
+            config=config,
+            on_collect=remove_collected_entry,
+        )
     return handle
 
 
 def invalidate_cute_k1_state(block: Any) -> None:
     """Release a block's registered CuTe state after its modules change."""
     state = getattr(block, "_deepmd_cute_k1_state", None)
     if isinstance(state, _RegisteredK1State):
-        _REGISTRY.pop(state.handle, None)
+        with _REGISTRY_LOCK:
+            _REGISTRY.pop(state.handle, None)
     if hasattr(block, "_deepmd_cute_k1_state"):
         delattr(block, "_deepmd_cute_k1_state")
     if hasattr(block, "_deepmd_cute_gate_expand_contract"):
@@ -568,15 +572,16 @@ def _register_cute_k1_state(
     _validate_gate_expand_index(block)
     old_state = getattr(block, "_deepmd_cute_k1_state", None)
     if isinstance(old_state, _RegisteredK1State):
-        old_entry = _REGISTRY.get(old_state.handle)
-        if (
-            old_state.device_index == device_index
-            and old_state.config == config
-            and old_entry is not None
-            and old_entry.block is block
-        ):
-            return old_state
-        _REGISTRY.pop(old_state.handle, None)
+        with _REGISTRY_LOCK:
+            old_entry = _REGISTRY.get(old_state.handle)
+            if (
+                old_state.device_index == device_index
+                and old_state.config == config
+                and old_entry is not None
+                and old_entry.block is block
+            ):
+                return old_state
+            _REGISTRY.pop(old_state.handle, None)
     if not _module_state_is_aligned(block):
         # Module parameters and buffers are frozen for this inference path, so
         # cache a failed static contract until explicit state invalidation.
@@ -1074,7 +1079,6 @@ def _final_manual_backward(runner: Any, grad_out: Tensor) -> tuple[Tensor, Tenso
         grad_post_norm_in.squeeze(2).unsqueeze(2),
     ).squeeze(2)
 
-    grad_x_wide_down = None
     if so2.message_node_grid_product is not None:
         if runner.packed_message_grid:
             message_grid_product = runner.message_grid_product
@@ -1099,14 +1103,13 @@ def _final_manual_backward(runner: Any, grad_out: Tensor) -> tuple[Tensor, Tenso
                 grad_post_mix,
             )
         grad_out_gate_flat.add_(grad_post_mix)
-        if grad_x_wide_down is None:
-            grad_x_wide_down = torch.zeros(
-                n_node,
-                16 * 64,
-                device=x_wide.device,
-                dtype=x_wide.dtype,
-            ).view(n_node, 16, 64)
-            grad_x_wide_down.add_(grad_grid_context)
+        grad_x_wide_down = torch.zeros(
+            n_node,
+            16 * 64,
+            device=x_wide.device,
+            dtype=x_wide.dtype,
+        ).view(n_node, 16, 64)
+        grad_x_wide_down.add_(grad_grid_context)
     else:
         grad_out_gate_flat = grad_post_mix
         grad_x_wide_down = torch.zeros(
@@ -1354,7 +1357,8 @@ def _build_runner(
     here and repair contiguous offset views without introducing a Dynamo graph
     break above the op.
     """
-    entry = _REGISTRY[int(handle)]
+    with _REGISTRY_LOCK:
+        entry = _REGISTRY[int(handle)]
     config = entry.config
     if config.native_sm90_path:
         from .sm90_k1.runner import NeoSm90K1Runner as Runner
@@ -1611,7 +1615,6 @@ def _runner_backward_manual(
         runner.focus_alpha,
         runner.dst_ptr_i32,
         runner.rotate,
-        runner.attn_logits.detach().contiguous(),
         runner.edge_gate,
         so2.adamw_attn_z_bias_raw.detach().reshape(2).float().contiguous(),
         runner.group_max,
@@ -1992,9 +1995,10 @@ def _maybe_run_prepared_cute_k1(
     state = getattr(block, "_deepmd_cute_k1_state", None)
     if not isinstance(state, _RegisteredK1State):
         return None
-    entry = _REGISTRY.get(state.handle)
-    if entry is None or entry.block is not block:
-        return None
+    with _REGISTRY_LOCK:
+        entry = _REGISTRY.get(state.handle)
+        if entry is None or entry.block is not block:
+            return None
     d_full = edge_cache.D_packed
     dt_full = d_full
     destinations_sorted = bool(getattr(edge_cache, "destinations_sorted", False))

deepmd/kernels/cute/neo/k1_kernels/cute_envelope_gated_softmax.py

@@ -21,6 +21,7 @@
 
 import cutlass
 import cutlass.cute as cute
+import cutlass.utils
 from cuda.bindings.driver import (
     CUstream,
 )

deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py

@@ -12,6 +12,7 @@
 
 import cutlass
 import cutlass.cute as cute
+import cutlass.utils
 from cuda.bindings.driver import (
     CUstream,
 )

deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_a_radial_forward.py

@@ -265,11 +265,19 @@ def _combined_radial_weight(radial_hidden_proj, radial_degree_mixer):
     cache = getattr(radial_hidden_proj, "_deepmd_cute_neo_radial_combined", None)
     key = (
         hidden_weight.data_ptr(),
-        mixer_weight.data_ptr(),
         hidden_weight._version,
-        mixer_weight._version,
         hidden_weight.dtype,
         hidden_weight.device,
+        tuple(hidden_weight.shape),
+        tuple(hidden_weight.stride()),
+        hidden_weight.storage_offset(),
+        mixer_weight.data_ptr(),
+        mixer_weight._version,
+        mixer_weight.dtype,
+        mixer_weight.device,
+        tuple(mixer_weight.shape),
+        tuple(mixer_weight.stride()),
+        mixer_weight.storage_offset(),
     )
     if cache is not None and cache[0] == key:
         return cache[1]
@@ -308,19 +316,45 @@ def run_neo_phase_a_radial_forward_packed_direct(
             f"expected radial_feat_m0 shape {(edge_count, 4, 32)}, "
             f"got {tuple(radial_feat_m0.shape)}"
         )
+    device = x_wide.device
+    if device.type != "cuda":
+        raise ValueError("packed Phase-A/radial forward requires CUDA tensors")
+    if src.device != device or src.dtype not in (torch.int32, torch.int64):
+        raise ValueError("src must be an int32 or int64 tensor on the input device")
+    if src.data_ptr() % 16:
+        raise ValueError("src must be 16-byte aligned")
+    source_tensors = (
+        ("x_wide", x_wide),
+        ("D_full", D_full),
+        ("radial_feat_m0", radial_feat_m0),
+        ("radial_hidden_proj.weight", radial_hidden_proj.weight),
+        ("radial_degree_mixer.weight", radial_degree_mixer.weight),
+        ("radial_degree_mixer.channel_basis", radial_degree_mixer.channel_basis),
+    )
+    for name, tensor in source_tensors:
+        if tensor.device != device or tensor.dtype != torch.float32:
+            raise ValueError(f"{name} must be FP32 on {device}")
+        if tensor.data_ptr() % 16:
+            raise ValueError(f"{name} must be 16-byte aligned")
     if radial_hidden_proj.bias is not None:
         raise NotImplementedError("collapsed radial mixer expects no hidden bias")
+    if tuple(radial_hidden_proj.weight.shape) != (32, 64):
+        raise NotImplementedError("collapsed radial mixer expects a (32,64) projection")
     if radial_degree_mixer.mode != "degree_channel" or radial_degree_mixer.rank != 1:
         raise NotImplementedError(
             "collapsed radial mixer expects degree_channel rank=1"
         )
     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")
 
     combined_weight = _combined_radial_weight(
         radial_hidden_proj,
         radial_degree_mixer,
     )
+    if not combined_weight.is_contiguous() or combined_weight.data_ptr() % 16:
+        raise ValueError("combined radial weight must be contiguous and aligned")
     kernel = _compiled_neo_phase_a_radial_forward_packed_direct()
     out = torch.empty(
         edge_count,

deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout.py

@@ -70,7 +70,6 @@ class NeoPhaseCBackwardLayoutParams:
     focus_alpha: cute.Tensor
     dst_ptr: cute.Tensor
     rotate_inv_rescale: cute.Tensor
-    logits: cute.Tensor
     edge_gate: cute.Tensor
     z_bias_raw: cute.Tensor
     group_max: cute.Tensor
@@ -445,7 +444,6 @@ def neo_phase_c_backward_layout_jit(
     focus_alpha: cute.Tensor,
     dst_ptr: cute.Tensor,
     rotate_inv_rescale: cute.Tensor,
-    logits: cute.Tensor,
     edge_gate: cute.Tensor,
     z_bias_raw: cute.Tensor,
     group_max: cute.Tensor,
@@ -473,7 +471,6 @@ def neo_phase_c_backward_layout_jit(
         focus_alpha=focus_alpha,
         dst_ptr=dst_ptr,
         rotate_inv_rescale=rotate_inv_rescale,
-        logits=logits,
         edge_gate=edge_gate,
         z_bias_raw=z_bias_raw,
         group_max=group_max,
@@ -580,12 +577,6 @@ def compile_neo_phase_c_backward_layout(
         stride_order=(0,),
         **FAKE_TENSOR_KW,
     )
-    fake_logits = make_fake_compact_tensor(
-        cutlass.Float32,
-        (edge_count, N_FOCUS),
-        stride_order=(1, 0),
-        **FAKE_TENSOR_KW,
-    )
     fake_edge_gate = make_fake_compact_tensor(
         cutlass.Float32,
         (edge_count,),
@@ -680,7 +671,6 @@ def compile_neo_phase_c_backward_layout(
         fake_focus_alpha,
         fake_dst_ptr,
         fake_rotate,
-        fake_logits,
         fake_edge_gate,
         fake_z_bias,
         fake_group_max,

deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout_runner.py

@@ -175,7 +175,6 @@ def __call__(
         focus_alpha: torch.Tensor,
         dst_ptr: torch.Tensor,
         rotate_inv_rescale: torch.Tensor,
-        logits: torch.Tensor,
         edge_gate: torch.Tensor,
         z_bias_raw: torch.Tensor,
         group_max: torch.Tensor,
@@ -225,7 +224,6 @@ def __call__(
             (DEGREE_COUNT,),
             device=device,
         )
-        _require_tensor("logits", logits, (edge_count, N_FOCUS), device=device)
         _require_tensor("edge_gate", edge_gate, (edge_count,), device=device)
         _require_tensor("z_bias_raw", z_bias_raw, (N_FOCUS,), device=device)
         _require_tensor("group_max", group_max, (node_count, N_FOCUS), device=device)
@@ -296,7 +294,6 @@ def __call__(
                 ("focus_alpha", focus_alpha),
                 ("dst_ptr", dst_ptr),
                 ("rotate_inv_rescale", rotate_inv_rescale),
-                ("logits", logits),
                 ("edge_gate", edge_gate),
                 ("z_bias_raw", z_bias_raw),
                 ("group_max", group_max),
@@ -361,7 +358,6 @@ def __call__(
             focus_alpha,
             dst_ptr,
             rotate_inv_rescale,
-            logits,
             edge_gate,
             z_bias_raw,
             group_max,

deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_onepass.py

@@ -32,6 +32,7 @@
 
 import cutlass
 import cutlass.cute as cute
+import cutlass.utils
 from cuda.bindings.driver import (
     CUstream,
 )
@@ -58,6 +59,9 @@
 REDUCED_COUNT = 10
 FOCUS_COUNT = 2
 CHANNELS = 32
+
+if PACKED_WIGNER_VALUES != 46 or PACKED_WIGNER_VALUES > 2 * CHANNELS:
+    raise RuntimeError("one-pass Neo Phase-C requires the 46-value Wigner layout")
 HIDDEN = FOCUS_COUNT * CHANNELS
 PHASE_WIDTH = REDUCED_COUNT * HIDDEN
 OUTPUT_WIDTH = DEGREE_COUNT * HIDDEN

deepmd/kernels/cute/neo/k1_kernels/cute_neo_so2_gate_combined_fwd.py

@@ -46,6 +46,7 @@
 M0_WIDTH = 4 * 32
 PAIR_WIDTH = 6 * 32
 FULL_WIDTH = M0_WIDTH + PAIR_WIDTH
+DEFAULT_STREAM = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT)
 
 
 def _supports_combined_forward(compute_capability: tuple[int, int]) -> bool:
@@ -112,7 +113,7 @@ def __call__(
         mGate: cute.Tensor,
         mY: cute.Tensor,
         mOut: cute.Tensor,
-        stream: cuda.CUstream = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT),
+        stream: cuda.CUstream = DEFAULT_STREAM,
     ):
         sA_layout = cute.make_layout(
             (TILE_M, TILE_K, STAGES),

deepmd/kernels/cute/neo/k1_radial_phase_a_node.py

@@ -17,6 +17,9 @@
 from .compile_cache import (
     device_aware_lru_cache,
 )
+from .k1_wigner_layout import (
+    PACKED_VALUE_COUNT,
+)
 
 if TYPE_CHECKING:
     from torch import (
@@ -29,7 +32,7 @@
 HIDDEN = 64
 FOCUS_COUNT = 2
 FOCUS_HIDDEN = 32
-PACKED_WIGNER_VALUES = 46
+PACKED_WIGNER_VALUES = PACKED_VALUE_COUNT
 RADIAL_WIDTH = 4 * FOCUS_HIDDEN
 COMPACT_WIDTH = 25
 PROJECTION_INPUT_WIDTH = COMPACT_WIDTH + FOCUS_COUNT
@@ -239,6 +242,8 @@ def run_neo_radial_phase_a_backward_node_tiled(
     The kernel recomputes Phase A, fuses the focus-source adjoint, uses
     four-lane warp reductions with a 68-float shared row pitch, and packs the
     27-column radial projection input for one strict-FP32 matrix call.
+    ``grad_out_focus`` is repacked in place as projection workspace and must
+    not be reused after this function returns.
     """
     import torch
 

deepmd/kernels/cute/neo/k1_runner.py

@@ -614,9 +614,17 @@ def ensure_backward_workspace(
             radial_scratch = getattr(self, "structural_scratch", None)
             if radial_scratch is None:
                 radial_scratch = next(
-                    cache.logits
-                    for cache in self.stack_caches
-                    if cache.logits is not None
+                    (
+                        cache.logits
+                        for cache in self.stack_caches
+                        if cache.logits is not None
+                    ),
+                    None,
+                )
+            if radial_scratch is None:
+                raise RuntimeError(
+                    "Neo K1 backward requires a reusable radial scratch buffer; "
+                    "no stack layer stored gate logits"
                 )
             if self.phase_c_stack.device.type == "cuda":
                 stream = self.torch.cuda.current_stream(self.phase_c_stack.device)
@@ -843,6 +851,7 @@ def _build_forward_graph(self) -> None:
             self.group_max,
             self.denom,
         )
+        self.attn_logits = None
         x_wide_down = self.x_wide.detach()
         from .k1_kernels.cute_neo_phase_c_onepass import (
             run_neo_phase_c_onepass_output_gate,