chore(cute): remove development artifacts from Neo PR

Author: atulcthakur

deepmd/kernels/cute/neo/gie.py

@@ -11,8 +11,8 @@
 produces radial, zonal/Wigner, source-gate, and degree-normalization gradients
 without an edge-by-row-by-channel temporary.
 
-``DP_NEO_CUTE_INFER`` is the master opt-in. The SM80/SM86 path is the
-default there, while ``DP_CUTE_GIE=0`` remains a diagnostic kill switch.
+``DP_NEO_CUTE_INFER`` is the master opt-in. The SM80/SM86 path enables this
+fusion by default; ``DP_CUTE_GIE=0`` disables it explicitly.
 """
 
 from __future__ import annotations
@@ -90,8 +90,8 @@ def _standard_index_contract(module: Any, lmax: int, row_count: int) -> bool:
     if rows.dtype != torch.long or slots.dtype != torch.long:
         return False
 
-    # These buffers are constructor-owned and immutable in normal execution.
-    # Check values on CPU tests, but avoid a device synchronization in inference.
+    # These buffers are constructor-owned and immutable. Validate their values
+    # when host-resident without introducing a CUDA synchronization.
     if rows.device.type == "cpu" and slots.device.type == "cpu":
         expected_rows = torch.arange(
             1, row_count + 1, dtype=torch.long, device=torch.device("cpu")

deepmd/kernels/cute/neo/k1.py

@@ -275,7 +275,7 @@ def _dtypes_use_strict_fp32(dtypes: Any) -> bool:
 def is_supported_k1_compute_capability(
     compute_capability: tuple[int, int],
 ) -> bool:
-    """Return whether the production K1 path has been validated on this GPU."""
+    """Return whether K1 supports this compute capability."""
     return runtime_policy.is_supported_k1_capability(compute_capability)
 
 
@@ -295,8 +295,8 @@ def _tensor_compute_capability(tensor: Any) -> tuple[int, int] | None:
 
 
 def _device_is_supported_for_k1(device: torch.device) -> bool:
-    # CPU-only contract tests use a synthetic CUDA device. Real CUDA dispatch
-    # always has an available runtime and must match a validated architecture.
+    # Metadata-only eligibility checks may use a CUDA device on a host without
+    # a CUDA runtime. Concrete CUDA dispatch always validates the capability.
     return not torch.cuda.is_available() or is_supported_k1_compute_capability(
         _device_compute_capability(device)
     )
@@ -431,7 +431,6 @@ def __init__(
         try:
             self._block_ref = weakref.ref(block, on_collect)
         except TypeError:
-            # Lightweight contract tests use plain object() sentinels.
             self._block_ref = lambda: block
         self.config = config
 

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

@@ -136,7 +136,7 @@ def _compile_layout_boundary(
 
 
 class CuteNeoPhaseCBackwardLayout:
-    """Validated callable for the fused node-owned Phase-C boundary.
+    """Callable for the fused node-owned Phase-C boundary.
 
     One invocation replaces Phase-C backward, envelope-softmax backward plus
     its z reduction, and focus-source backward.  The Phase-C residual adjoint

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

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Strict-FP32 one-pass Neo Phase-C forward with output gating.
 
-This module specializes the production Neo shape ``D=16, Dm=10, F=2,
+This module specializes the supported Neo shape ``D=16, Dm=10, F=2,
 C=32``.  One 64-thread CTA owns one destination CSR row.  Lane
 ``focus * 32 + channel`` sweeps every edge in that row and retains its 16
 output-degree accumulators in a CuTe register fragment.  Wigner values are
@@ -54,7 +54,7 @@
 
 @dataclass(frozen=True)
 class NeoPhaseCOnePassParams:
-    """Runtime tensors for the validated warp-private packed kernel."""
+    """Runtime tensors for the warp-private packed kernel."""
 
     x_local: cute.Tensor
     wigner_dt: cute.Tensor
@@ -196,7 +196,7 @@ def neo_phase_c_onepass_output_gate_packed_direct_warp_private_kernel(
     accumulator_layout: cute.Layout,
     eps: cutlass.Constexpr[float],
 ):
-    """Triton-style final store with no Wigner shared-memory round trip."""
+    """Store the gated output without a Wigner shared-memory round trip."""
     tid, _, _ = cute.arch.thread_idx()
     node, _, _ = cute.arch.block_idx()
     focus = tid // CHANNELS
@@ -395,7 +395,7 @@ def _fake_common_tensors():
 def compile_neo_phase_c_onepass_output_gate(
     eps: float,
 ) -> Callable:
-    """Compile the validated packed, focus-major, warp-private specialization."""
+    """Compile the packed, focus-major, warp-private specialization."""
     if not (eps > 0.0 and eps < float("inf")):
         raise ValueError("output-gate RMSNorm eps must be finite and positive")
     common_args = _fake_common_tensors()

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

@@ -637,7 +637,7 @@ def _compile_combined_forward(
     compute_capability: tuple[int, int],
 ) -> Callable:
     if not _supports_combined_forward(compute_capability):
-        raise RuntimeError("combined forward requires a validated compute capability")
+        raise RuntimeError("combined forward requires a supported compute capability")
     with torch.cuda.device(device_index):
         fake_x = _fake_focus_tensor(FULL_WIDTH)
         fake_residual = _fake_focus_tensor(FULL_WIDTH)
@@ -709,7 +709,7 @@ def __init__(
         compute_capability = tuple(torch.cuda.get_device_capability(device_index))
         if not _supports_combined_forward(compute_capability):
             raise RuntimeError(
-                "combined forward requires a validated compute capability"
+                "combined forward requires a supported compute capability"
             )
 
         packed_w0, packed_wp, packed_gate = packed_weights

deepmd/kernels/cute/neo/k1_radial_phase_a_node.py

@@ -50,8 +50,8 @@ def build_source_csr(
 ) -> NeoSourceCSR:
     """Build indirect source CSR without changing the physical edge order.
 
-    ``validate_sources=True`` synchronizes when ``src`` is CUDA. Production
-    callers should build this once with the edge cache and retain both tensors.
+    ``validate_sources=True`` synchronizes when ``src`` is CUDA. Callers should
+    build this once with the edge cache and retain both tensors.
     """
     import torch
 

deepmd/kernels/cute/neo/k1_runner.py

@@ -29,7 +29,7 @@ def _validate_runtime_config(
     compute_capability: tuple[int, int] | None = None,
 ) -> None:
     if compute_capability not in SUPPORTED_K1_CAPABILITIES:
-        raise RuntimeError("Neo K1 requires a validated compute capability")
+        raise RuntimeError("Neo K1 requires a supported compute capability")
     if runtime_config.native_sm90_path != (compute_capability == SM90_CAPABILITY):
         raise RuntimeError("the native SM90 K1 path must be selected only on sm_90")
     if runtime_config.per_focus_so2_fwd_pair != (

deepmd/kernels/cute/neo/k1_wigner_layout.py

@@ -12,8 +12,6 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
-import torch
-
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
@@ -85,39 +83,3 @@ def iter_packed_entries() -> Iterator[PackedWignerEntry]:
                     full_row=full_row,
                     full_col=block_start + local_col,
                 )
-
-
-def pack_dense_reference(d_full: torch.Tensor) -> torch.Tensor:
-    """Gather a dense block-diagonal matrix for CPU differential tests only."""
-    entries = tuple(iter_packed_entries())
-    rows = [entry.full_row for entry in entries]
-    cols = [entry.full_col for entry in entries]
-    return d_full[..., rows, cols]
-
-
-def phase_a_reference(panel: torch.Tensor, x_wide: torch.Tensor) -> torch.Tensor:
-    """Reference ``D[coeff_index_m] @ x_wide`` without dense unpacking."""
-    outputs = []
-    for reduced, degree in enumerate(REDUCED_DEGREES):
-        block_start = FULL_BLOCK_OFFSETS[degree]
-        block_stop = FULL_BLOCK_OFFSETS[degree + 1]
-        row_start = REDUCED_PANEL_ROW_OFFSETS[reduced]
-        values = panel[..., row_start : row_start + block_stop - block_start]
-        x_block = x_wide[..., block_start:block_stop, :]
-        outputs.append((values.unsqueeze(-1) * x_block).sum(dim=-2))
-    return torch.stack(outputs, dim=-2)
-
-
-def phase_c_reference(panel: torch.Tensor, x_local: torch.Tensor) -> torch.Tensor:
-    """Reference ``Dt[:, :, coeff_index_m] @ x_local`` from the same panel."""
-    outputs = []
-    for full_row in range(FULL_BLOCK_OFFSETS[-1]):
-        value = torch.zeros_like(x_local[..., 0, :])
-        for reduced in range(len(COEFF_INDEX_M)):
-            offset = dt_offset(full_row, reduced)
-            if offset is not None:
-                value = (
-                    value + panel[..., offset].unsqueeze(-1) * x_local[..., reduced, :]
-                )
-        outputs.append(value)
-    return torch.stack(outputs, dim=-2)

deepmd/kernels/cute/neo/output_grid_product.py

@@ -1,6 +1,6 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""Strict-FP32 CuTe middle contractions for validated Neo grid MLPs."""
+"""Strict-FP32 CuTe middle contractions for supported Neo grid MLPs."""
 
 from __future__ import annotations
 

deepmd/kernels/cute/neo/runtime_policy.py

@@ -113,7 +113,7 @@ def _sm80_profile_feature(
     name: str,
     compute_capability: tuple[int, int] | None = None,
 ) -> bool:
-    """Apply an SM80-family default with an environment kill switch."""
+    """Apply an SM80-family default with an explicit disable override."""
     if not is_sm80_profile_enabled(compute_capability):
         return False
     return _env_override(name) is not False
@@ -208,14 +208,14 @@ def is_readout_input_fold_sm90_enabled(
 def is_readout_input_fold_enabled(
     compute_capability: tuple[int, int] | None = None,
 ) -> bool:
-    """Select the validated architecture-specific frozen readout fold."""
+    """Select the architecture-specific frozen readout fold."""
     return is_readout_input_fold_sm80_enabled(
         compute_capability
     ) or is_readout_input_fold_sm90_enabled(compute_capability)
 
 
 def is_supported_k1_capability(compute_capability: tuple[int, int]) -> bool:
-    """Return whether K1 has a validated production profile on this GPU."""
+    """Return whether K1 supports this compute capability."""
     return tuple(compute_capability) in SUPPORTED_K1_CAPABILITIES