From 8d11cd06696948fa2727915107557f569635a472 Mon Sep 17 00:00:00 2001 From: atulcthakur Date: Mon, 17 Aug 2026 15:54:32 -0700 Subject: [PATCH 01/11] feat(pt): add optimized Neo CuTe inference path Add the production Neo CuTe inference integration, architecture-specific kernels, correctness coverage, runtime dependencies, licensing, and final review cleanup. --- .github/workflows/test_cuda.yml | 3 +- deepmd/kernels/cute/neo/__init__.py | 3 + deepmd/kernels/cute/neo/compile_cache.py | 78 + deepmd/kernels/cute/neo/gie.py | 740 +++++ deepmd/kernels/cute/neo/k1.py | 2220 +++++++++++++ deepmd/kernels/cute/neo/k1_gate_structural.py | 175 + .../kernels/cute/neo/k1_kernels/__init__.py | 3 + .../k1_kernels/cute_envelope_gated_softmax.py | 237 ++ .../k1_kernels/cute_neo_focus_src_backward.py | 265 ++ ...neo_gate_linear_residual_backward_fused.py | 168 + ...ute_neo_gate_split_structural_vec4_sm80.py | 497 +++ .../cute_neo_message_grid_product.py | 275 ++ .../cute_neo_output_gate_backward.py | 204 ++ .../cute_neo_phase_a_radial_forward.py | 358 ++ .../cute_neo_phase_c_backward_layout.py | 703 ++++ ...cute_neo_phase_c_backward_layout_runner.py | 380 +++ .../k1_kernels/cute_neo_phase_c_onepass.py | 548 ++++ .../cute/neo/k1_kernels/cute_neo_qk_edge.py | 364 ++ .../cute_neo_radial_phase_a_backward_node.py | 720 ++++ .../cute_neo_so2_gate_combined_fwd.py | 779 +++++ .../cute/neo/k1_message_grid_packed.py | 363 ++ .../cute/neo/k1_radial_phase_a_node.py | 395 +++ deepmd/kernels/cute/neo/k1_runner.py | 895 +++++ deepmd/kernels/cute/neo/k1_so2linear.py | 261 ++ deepmd/kernels/cute/neo/k1_wigner_layout.py | 93 + deepmd/kernels/cute/neo/k4_wignerd.py | 1004 ++++++ .../cute/neo/message_grid_gaunt_sm90.py | 323 ++ .../cute/neo/message_grid_readout_sm90.py | 849 +++++ .../cute/neo/output_grid_kernels/__init__.py | 3 + .../output_grid_kernels/cute_readout_l0.py | 834 +++++ .../cute_tiled_grid_product.py | 2915 +++++++++++++++++ .../kernels/cute/neo/output_grid_product.py | 329 ++ deepmd/kernels/cute/neo/readout_l0.py | 774 +++++ deepmd/kernels/cute/neo/runtime_policy.py | 290 ++ deepmd/kernels/cute/neo/sm90_k1/__init__.py | 3 + .../kernels/cute/neo/sm90_k1/final_phase_c.py | 483 +++ .../kernels/cute/neo/sm90_k1/output_gate.py | 204 ++ deepmd/kernels/cute/neo/sm90_k1/persistent.py | 593 ++++ deepmd/kernels/cute/neo/sm90_k1/phase_a.py | 554 ++++ .../cute/neo/sm90_k1/phase_a_backward.py | 918 ++++++ .../neo/sm90_k1/phase_c_attention_backward.py | 838 +++++ deepmd/kernels/cute/neo/sm90_k1/prefix.py | 92 + deepmd/kernels/cute/neo/sm90_k1/radial.py | 306 ++ deepmd/kernels/cute/neo/sm90_k1/runner.py | 661 ++++ deepmd/pt/infer/deep_eval.py | 46 +- deepmd/pt/model/descriptor/sezm.py | 188 +- deepmd/pt/model/descriptor/sezm_nn/block.py | 91 +- .../pt/model/descriptor/sezm_nn/edge_cache.py | 202 +- .../pt/model/descriptor/sezm_nn/embedding.py | 75 +- .../pt/model/descriptor/sezm_nn/grid_net.py | 79 +- deepmd/pt/model/descriptor/sezm_nn/lora.py | 22 + deepmd/pt/model/model/sezm_model.py | 257 +- deepmd/pt/model/network/mlp.py | 59 +- deepmd/pt/utils/nv_nlist.py | 21 + doc/install/easy-install.md | 14 + pyproject.toml | 4 + ...t_descriptor_sezm_block_ffn_grid_fusion.py | 377 +++ ...test_descriptor_sezm_cute_compile_cache.py | 115 + ...t_descriptor_sezm_cute_envelope_softmax.py | 83 + .../pt/model/test_descriptor_sezm_cute_gie.py | 249 ++ .../pt/model/test_descriptor_sezm_cute_k1.py | 555 ++++ .../test_descriptor_sezm_cute_k1_dst_ptr.py | 646 ++++ ...descriptor_sezm_cute_k1_gate_structural.py | 117 + ...escriptor_sezm_cute_k1_inplace_residual.py | 387 +++ ...descriptor_sezm_cute_k1_manual_adjoints.py | 360 ++ ...est_descriptor_sezm_cute_k1_packed_impl.py | 106 + ...or_sezm_cute_k1_phase_c_inplace_adjoint.py | 323 ++ ...riptor_sezm_cute_k1_radial_focus_fusion.py | 113 + ...riptor_sezm_cute_k1_radial_pitch68_safe.py | 148 + ...scriptor_sezm_cute_k1_radial_projection.py | 234 ++ ..._sezm_cute_k1_structural_gate_vec4_sm80.py | 516 +++ ...or_sezm_cute_k1_structural_memory_reuse.py | 236 ++ ...escriptor_sezm_cute_output_grid_product.py | 539 +++ .../test_descriptor_sezm_cute_readout_l0.py | 860 +++++ ...est_descriptor_sezm_cute_runtime_policy.py | 249 ++ ...descriptor_sezm_cute_tiled_grid_product.py | 387 +++ ...test_descriptor_sezm_readout_l0_algebra.py | 246 ++ source/tests/pt/model/test_mlp.py | 98 + source/tests/pt/model/test_nlist_backend.py | 52 + source/tests/pt/model/test_sezm_cute_efs.py | 452 +++ .../model/test_sezm_cute_multigpu_dispatch.py | 279 ++ source/tests/pt/model/test_sezm_model.py | 271 ++ 82 files changed, 31677 insertions(+), 79 deletions(-) create mode 100644 deepmd/kernels/cute/neo/__init__.py create mode 100644 deepmd/kernels/cute/neo/compile_cache.py create mode 100644 deepmd/kernels/cute/neo/gie.py create mode 100644 deepmd/kernels/cute/neo/k1.py create mode 100644 deepmd/kernels/cute/neo/k1_gate_structural.py create mode 100644 deepmd/kernels/cute/neo/k1_kernels/__init__.py create mode 100644 deepmd/kernels/cute/neo/k1_kernels/cute_envelope_gated_softmax.py create mode 100644 deepmd/kernels/cute/neo/k1_kernels/cute_neo_focus_src_backward.py create mode 100644 deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py create mode 100644 deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py create mode 100644 deepmd/kernels/cute/neo/k1_kernels/cute_neo_message_grid_product.py create mode 100644 deepmd/kernels/cute/neo/k1_kernels/cute_neo_output_gate_backward.py create mode 100644 deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_a_radial_forward.py create mode 100644 deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout.py create mode 100644 deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout_runner.py create mode 100644 deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_onepass.py create mode 100644 deepmd/kernels/cute/neo/k1_kernels/cute_neo_qk_edge.py create mode 100644 deepmd/kernels/cute/neo/k1_kernels/cute_neo_radial_phase_a_backward_node.py create mode 100644 deepmd/kernels/cute/neo/k1_kernels/cute_neo_so2_gate_combined_fwd.py create mode 100644 deepmd/kernels/cute/neo/k1_message_grid_packed.py create mode 100644 deepmd/kernels/cute/neo/k1_radial_phase_a_node.py create mode 100644 deepmd/kernels/cute/neo/k1_runner.py create mode 100644 deepmd/kernels/cute/neo/k1_so2linear.py create mode 100644 deepmd/kernels/cute/neo/k1_wigner_layout.py create mode 100644 deepmd/kernels/cute/neo/k4_wignerd.py create mode 100644 deepmd/kernels/cute/neo/message_grid_gaunt_sm90.py create mode 100644 deepmd/kernels/cute/neo/message_grid_readout_sm90.py create mode 100644 deepmd/kernels/cute/neo/output_grid_kernels/__init__.py create mode 100644 deepmd/kernels/cute/neo/output_grid_kernels/cute_readout_l0.py create mode 100644 deepmd/kernels/cute/neo/output_grid_kernels/cute_tiled_grid_product.py create mode 100644 deepmd/kernels/cute/neo/output_grid_product.py create mode 100644 deepmd/kernels/cute/neo/readout_l0.py create mode 100644 deepmd/kernels/cute/neo/runtime_policy.py create mode 100644 deepmd/kernels/cute/neo/sm90_k1/__init__.py create mode 100644 deepmd/kernels/cute/neo/sm90_k1/final_phase_c.py create mode 100644 deepmd/kernels/cute/neo/sm90_k1/output_gate.py create mode 100644 deepmd/kernels/cute/neo/sm90_k1/persistent.py create mode 100644 deepmd/kernels/cute/neo/sm90_k1/phase_a.py create mode 100644 deepmd/kernels/cute/neo/sm90_k1/phase_a_backward.py create mode 100644 deepmd/kernels/cute/neo/sm90_k1/phase_c_attention_backward.py create mode 100644 deepmd/kernels/cute/neo/sm90_k1/prefix.py create mode 100644 deepmd/kernels/cute/neo/sm90_k1/radial.py create mode 100644 deepmd/kernels/cute/neo/sm90_k1/runner.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_block_ffn_grid_fusion.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_compile_cache.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_envelope_softmax.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_gie.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_k1.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_k1_dst_ptr.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_k1_manual_adjoints.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_k1_packed_impl.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_focus_fusion.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_pitch68_safe.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_memory_reuse.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_output_grid_product.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_readout_l0.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_runtime_policy.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_cute_tiled_grid_product.py create mode 100644 source/tests/pt/model/test_descriptor_sezm_readout_l0_algebra.py create mode 100644 source/tests/pt/model/test_sezm_cute_efs.py create mode 100644 source/tests/pt/model/test_sezm_cute_multigpu_dispatch.py diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index 866eca155e..0f3532fb45 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -58,7 +58,7 @@ jobs: - run: | export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])') export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') - source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit + source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,cute,torch,jax] mpi4py --reinstall-package deepmd-kit # See https://github.com/jax-ml/jax/issues/29042 source/install/uv_with_retry.sh pip install --system -U 'nvidia-cublas-cu12>=12.9.0.13' env: @@ -66,6 +66,7 @@ jobs: DP_ENABLE_NATIVE_OPTIMIZATION: 1 DP_ENABLE_PYTORCH: 1 - run: dp --version + - run: python -c "import cutlass.cute" - run: python -m pytest source/tests --ignore=source/tests/pd env: NUM_WORKERS: 0 diff --git a/deepmd/kernels/cute/neo/__init__.py b/deepmd/kernels/cute/neo/__init__.py new file mode 100644 index 0000000000..e98fa6b6db --- /dev/null +++ b/deepmd/kernels/cute/neo/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Neo-specialized CuTe kernels and PyTorch integration.""" diff --git a/deepmd/kernels/cute/neo/compile_cache.py b/deepmd/kernels/cute/neo/compile_cache.py new file mode 100644 index 0000000000..0a9b9b89b5 --- /dev/null +++ b/deepmd/kernels/cute/neo/compile_cache.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Device-aware caching for architecture-specific CuTe compilation.""" + +from __future__ import ( + annotations, +) + +from collections.abc import ( + Callable, +) +from contextlib import ( + nullcontext, +) +from functools import ( + lru_cache, + wraps, +) +from typing import ( + Any, + TypeVar, + cast, +) + +_T = TypeVar("_T", bound=Callable[..., Any]) + + +def current_cuda_compile_identity() -> tuple[int, int, int]: + """Return the current CUDA device and its compute capability.""" + import torch + + device_index = torch.cuda.current_device() + major, minor = torch.cuda.get_device_capability(device_index) + return device_index, major, minor + + +def device_aware_lru_cache( + *, + maxsize: int, + identity_getter: Callable[[], tuple[int, int, int]] = current_cuda_compile_identity, +) -> Callable[[_T], _T]: + """Cache a compile factory separately for each CUDA device architecture.""" + + def decorate(function: _T) -> _T: + @lru_cache(maxsize=maxsize) + def cached( + identity: tuple[int, int, int], + args: tuple[Any, ...], + kwargs: tuple[tuple[str, Any], ...], + ) -> Any: + import torch + + device_index = identity[0] + device_count = getattr(torch.cuda, "device_count", None) + device_is_visible = device_count is None or device_index < device_count() + compile_device = ( + torch.cuda.device(device_index) + if torch.cuda.is_available() and device_is_visible + else nullcontext() + ) + with compile_device: + return function(*args, **dict(kwargs)) + + @wraps(function) + def wrapper(*args: Any, **kwargs: Any) -> Any: + return cached( + identity_getter(), + args, + tuple(sorted(kwargs.items())), + ) + + wrapper.cache_clear = cached.cache_clear + wrapper.cache_info = cached.cache_info + wrapper.cache_parameters = cached.cache_parameters + wrapper._deepmd_cute_cached = True + return cast("_T", wrapper) + + return decorate diff --git a/deepmd/kernels/cute/neo/gie.py b/deepmd/kernels/cute/neo/gie.py new file mode 100644 index 0000000000..70e747b8c5 --- /dev/null +++ b/deepmd/kernels/cute/neo/gie.py @@ -0,0 +1,740 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +# pyright: reportMissingImports=false +# ruff: noqa: ANN001 +"""Opt-in CuTe fusion for the DPA4 geometric initial embedding. + +The eager implementation materializes both ``radial_value_for_row`` and +``non_scalar_message`` with shape ``(E, D - 1, C)`` before reducing by +destination. This module computes the same strict-FP32 expression directly +into ``(N, D, C)`` using the destination-sorted edge list. Its first backward +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 enables this +fusion by default; ``DP_CUTE_GIE=0`` disables it explicitly. +""" + +from __future__ import ( + annotations, +) + +import threading +from typing import ( + TYPE_CHECKING, + Any, +) + +import torch +from torch import ( + Tensor, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + +try: + import cutlass + import cutlass.cute as cute + import cutlass.torch as cutlass_torch + from cuda.bindings.driver import CUstream # noqa: TC002 + from cutlass.cute.runtime import ( + from_dlpack, + ) + + SEZM_CUTE_GIE_AVAILABLE = True +except Exception: # pragma: no cover - import guard for non-CuTe environments + SEZM_CUTE_GIE_AVAILABLE = False + + +def is_cute_gie_enabled(device: torch.device | None = None) -> bool: + """Return whether the architecture-selected GIE path is enabled.""" + from .runtime_policy import ( + is_gie_enabled, + ) + + if device is None: + if not torch.cuda.is_available(): + return False + device_index = torch.cuda.current_device() + else: + if device.type != "cuda": + return False + device_index = device.index + if device_index is None: + device_index = torch.cuda.current_device() + return is_gie_enabled(tuple(torch.cuda.get_device_capability(device_index))) + + +def _backward_compile_key( + device_identity: tuple[int, int, int], + lmax: int, + channels: int, + has_gate: bool, + radial_stride: tuple[int, ...], + dst_dtype: torch.dtype, +) -> tuple[Any, ...]: + """Build an ABI-complete GIE backward compilation key.""" + return ( + "gie_bwd", + *device_identity, + lmax, + channels, + has_gate, + tuple(radial_stride), + dst_dtype, + ) + + +def _degree_slots(lmax: int, *, device: torch.device) -> Tensor: + degrees = torch.arange(1, lmax + 1, device=device, dtype=torch.long) + return torch.repeat_interleave(degrees - 1, 2 * degrees + 1) + + +def _standard_index_contract(module: Any, lmax: int, row_count: int) -> bool: + rows = getattr(module, "non_scalar_row_index", None) + slots = getattr(module, "radial_slot_index_for_row", None) + if not isinstance(rows, Tensor) or not isinstance(slots, Tensor): + return False + if rows.numel() != row_count or slots.numel() != row_count: + return False + if rows.dtype != torch.long or slots.dtype != torch.long: + return False + + # 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") + ) + expected_slots = _degree_slots(lmax, device=torch.device("cpu")) + return bool( + torch.equal(rows, expected_rows) and torch.equal(slots, expected_slots) + ) + return True + + +def validate_gie_contract( + module: Any, + n_nodes: int, + edge_cache: Any, + radial: Tensor, + zonal: Tensor, +) -> bool: + """Validate the shape/layout contract without inspecting dynamic edge values.""" + lmax = int(getattr(module, "lmax", -1)) + channels = int(getattr(module, "channels", -1)) + if lmax <= 0 or channels <= 0 or n_nodes <= 0: + return False + row_count = (lmax + 1) ** 2 - 1 + dst = getattr(edge_cache, "dst", None) + inv_sqrt_deg = getattr(edge_cache, "inv_sqrt_deg", None) + gate = getattr(edge_cache, "edge_src_gate", None) + if not bool(getattr(edge_cache, "destinations_sorted", False)): + return False + if not isinstance(dst, Tensor) or not isinstance(inv_sqrt_deg, Tensor): + return False + if radial.dim() != 3 or zonal.dim() != 2 or dst.dim() != 1: + return False + edge_count = radial.shape[0] + if edge_count == 0: + return False + if ( + zonal.shape != (edge_count, row_count) + or radial.shape[1:] != (lmax, channels) + or dst.shape[0] != edge_count + or inv_sqrt_deg.shape != (n_nodes, 1, 1) + ): + return False + if radial.dtype != torch.float32 or zonal.dtype != torch.float32: + return False + if inv_sqrt_deg.dtype != torch.float32: + return False + if dst.dtype not in (torch.int32, torch.int64): + return False + if not (radial.device == zonal.device == dst.device == inv_sqrt_deg.device): + return False + if radial.stride(-1) != 1 or radial.stride(-2) != channels: + return False + if ( + not zonal.is_contiguous() + or not dst.is_contiguous() + or not inv_sqrt_deg.is_contiguous() + ): + return False + if gate is not None: + if not isinstance(gate, Tensor): + return False + if gate.shape not in ((edge_count,), (edge_count, 1)): + return False + if gate.dtype != torch.float32 or gate.device != radial.device: + return False + if not gate.is_contiguous(): + return False + return _standard_index_contract(module, lmax, row_count) + + +if SEZM_CUTE_GIE_AVAILABLE: + _F32 = cutlass.Float32 + _I32 = cutlass.Int32 + _WARPS_PER_BLOCK = 4 + _LANES = 32 + + def _build_forward(lmax: int, channels: int, has_gate: bool) -> Callable: + row_count = (lmax + 1) ** 2 - 1 + + @cute.kernel + def kernel(m_radial, m_zonal, m_inv, m_dst_ptr, m_gate, m_out) -> None: + node, _, _ = cute.arch.block_idx() + lane, warp, _ = cute.arch.thread_idx() + + if warp == 0: + for channel in cutlass.range(lane, channels, _LANES, unroll=1): + m_out[node, channel] = _F32(0.0) + + lo = m_dst_ptr[node].to(_I32) + hi = m_dst_ptr[node + 1].to(_I32) + for row in cutlass.range(warp, row_count, _WARPS_PER_BLOCK, unroll=1): + radial_slot = _I32(0) + for degree in cutlass.range_constexpr(lmax): + start = (degree + 1) * (degree + 1) - 1 + stop = (degree + 2) * (degree + 2) - 1 + if row >= start and row < stop: + radial_slot = degree + for channel in cutlass.range(lane, channels, _LANES, unroll=1): + acc = _F32(0.0) + for edge in cutlass.range(lo, hi, 1, unroll=1): + scale = _F32(1.0) + if has_gate: + scale = m_gate[edge].to(_F32) + acc += ( + m_zonal[edge, row].to(_F32) + * m_radial[edge, radial_slot * channels + channel].to(_F32) + * scale + ) + m_out[node, (row + 1) * channels + channel] = acc * m_inv[ + node, 0 + ].to(_F32) + + @cute.jit + def host( + m_radial, + m_zonal, + m_inv, + m_dst_ptr, + m_gate, + m_out, + stream: CUstream, + ) -> None: + nodes, _ = m_out.shape + kernel(m_radial, m_zonal, m_inv, m_dst_ptr, m_gate, m_out).launch( + grid=[nodes, 1, 1], + block=[_LANES, _WARPS_PER_BLOCK, 1], + stream=stream, + ) + + return host + + def _build_backward(lmax: int, channels: int, has_gate: bool) -> Callable: + row_count = (lmax + 1) ** 2 - 1 + full_width = (row_count + 1) * channels + + @cute.kernel + def edge_kernel( + m_grad_out, + m_radial, + m_zonal, + m_inv, + m_dst, + m_gate, + m_grad_radial, + m_grad_zonal, + m_grad_gate, + ) -> None: + block_edge, _, _ = cute.arch.block_idx() + lane, warp, _ = cute.arch.thread_idx() + edge = block_edge * _WARPS_PER_BLOCK + warp + edge_count, _ = m_radial.shape + load_edge = edge + if edge >= edge_count: + load_edge = 0 + node = m_dst[load_edge].to(_I32) + norm = m_inv[node, 0].to(_F32) + gate_value = _F32(1.0) + if has_gate: + gate_value = m_gate[load_edge].to(_F32) + grad_gate_acc = _F32(0.0) + + if channels <= _LANES: + active_channel = lane < channels + for degree in cutlass.range_constexpr(lmax): + start = (degree + 1) * (degree + 1) - 1 + stop = (degree + 2) * (degree + 2) - 1 + grad_radial_acc = _F32(0.0) + for row in cutlass.range_constexpr(start, stop, 1): + grad_value = _F32(0.0) + radial_value = _F32(0.0) + zonal_value = _F32(0.0) + if active_channel: + grad_value = ( + m_grad_out[node, (row + 1) * channels + lane].to(_F32) + * norm + ) + radial_value = m_radial[ + load_edge, degree * channels + lane + ].to(_F32) + zonal_value = m_zonal[load_edge, row].to(_F32) + grad_radial_acc += grad_value * zonal_value * gate_value + grad_gate_acc += grad_value * zonal_value * radial_value + grad_zonal_value = cute.arch.warp_reduction_sum( + grad_value * radial_value * gate_value + ) + if lane == 0 and edge < edge_count: + m_grad_zonal[edge, row] = grad_zonal_value + if active_channel and edge < edge_count: + m_grad_radial[edge, degree * channels + lane] = grad_radial_acc + else: + for degree in cutlass.range_constexpr(lmax): + start = (degree + 1) * (degree + 1) - 1 + stop = (degree + 2) * (degree + 2) - 1 + for channel in cutlass.range(lane, channels, _LANES, unroll=1): + grad_radial_acc = _F32(0.0) + for row in cutlass.range_constexpr(start, stop, 1): + grad_value = ( + m_grad_out[node, (row + 1) * channels + channel].to( + _F32 + ) + * norm + ) + radial_value = m_radial[ + load_edge, degree * channels + channel + ].to(_F32) + zonal_value = m_zonal[load_edge, row].to(_F32) + grad_radial_acc += grad_value * zonal_value * gate_value + grad_gate_acc += grad_value * zonal_value * radial_value + if edge < edge_count: + m_grad_radial[edge, degree * channels + channel] = ( + grad_radial_acc + ) + for row in cutlass.range_constexpr(start, stop, 1): + grad_zonal_acc = _F32(0.0) + for channel in cutlass.range(lane, channels, _LANES, unroll=1): + grad_value = ( + m_grad_out[node, (row + 1) * channels + channel].to( + _F32 + ) + * norm + ) + radial_value = m_radial[ + load_edge, degree * channels + channel + ].to(_F32) + grad_zonal_acc += grad_value * radial_value * gate_value + grad_zonal_value = cute.arch.warp_reduction_sum(grad_zonal_acc) + if lane == 0 and edge < edge_count: + m_grad_zonal[edge, row] = grad_zonal_value + + if has_gate: + grad_gate_value = cute.arch.warp_reduction_sum(grad_gate_acc) + if lane == 0 and edge < edge_count: + m_grad_gate[edge] = grad_gate_value + + @cute.kernel + def inv_kernel(m_grad_out, m_out, m_inv, m_grad_inv) -> None: + block_node, _, _ = cute.arch.block_idx() + lane, warp, _ = cute.arch.thread_idx() + node = block_node * _WARPS_PER_BLOCK + warp + node_count, _ = m_out.shape + load_node = node + if node >= node_count: + load_node = 0 + acc = _F32(0.0) + for idx in cutlass.range(channels + lane, full_width, _LANES, unroll=1): + acc += m_grad_out[load_node, idx].to(_F32) * m_out[load_node, idx].to( + _F32 + ) + acc = cute.arch.warp_reduction_sum(acc) + if lane == 0 and node < node_count: + m_grad_inv[node, 0] = acc / m_inv[node, 0].to(_F32) + + @cute.jit + def host( + m_grad_out, + m_radial, + m_zonal, + m_inv, + m_dst, + m_gate, + m_out, + m_grad_radial, + m_grad_zonal, + m_grad_inv, + m_grad_gate, + stream: CUstream, + ) -> None: + edge_count, _ = m_radial.shape + node_count, _ = m_out.shape + edge_kernel( + m_grad_out, + m_radial, + m_zonal, + m_inv, + m_dst, + m_gate, + m_grad_radial, + m_grad_zonal, + m_grad_gate, + ).launch( + grid=[cute.ceil_div(edge_count, _WARPS_PER_BLOCK), 1, 1], + block=[_LANES, _WARPS_PER_BLOCK, 1], + stream=stream, + ) + inv_kernel(m_grad_out, m_out, m_inv, m_grad_inv).launch( + grid=[cute.ceil_div(node_count, _WARPS_PER_BLOCK), 1, 1], + block=[_LANES, _WARPS_PER_BLOCK, 1], + stream=stream, + ) + + return host + + _compile_lock = threading.Lock() + _compiled: dict[tuple[Any, ...], Any] = {} + + def _as_cute(tensor: Tensor) -> Any: + value = from_dlpack(tensor) + if tensor.dim() <= 1: + return value.mark_layout_dynamic() + return value.mark_layout_dynamic(leading_dim=tensor.dim() - 1) + + def _device_key(tensor: Tensor) -> tuple[int, int, int]: + index = tensor.device.index + if index is None: + index = torch.cuda.current_device() + major, minor = torch.cuda.get_device_capability(index) + return index, major, minor + + def _get_compiled( + key: tuple[Any, ...], + builder: Callable[[], Callable], + example_args: tuple[Any, ...], + ) -> Any: + compiled = _compiled.get(key) + if compiled is not None: + return compiled + with _compile_lock: + compiled = _compiled.get(key) + if compiled is None: + compiled = cute.compile(builder(), *example_args) + _compiled[key] = compiled + return compiled + + def _flat_radial(radial: Tensor) -> Tensor: + return radial.view(radial.shape[0], radial.shape[1] * radial.shape[2]) + + def _flat_node(tensor: Tensor) -> Tensor: + return tensor.view(tensor.shape[0], -1) + + def _launch_forward( + radial: Tensor, + zonal: Tensor, + inv_sqrt_deg: Tensor, + dst_ptr: Tensor, + gate: Tensor, + lmax: int, + has_gate: bool, + ) -> Tensor: + channels = radial.shape[2] + out = torch.empty( + inv_sqrt_deg.shape[0], + (lmax + 1) ** 2, + channels, + device=radial.device, + dtype=radial.dtype, + ) + radial_flat = _flat_radial(radial.detach()) + inv_flat = _flat_node(inv_sqrt_deg.detach()) + gate_flat = gate.detach().view(-1) + out_flat = _flat_node(out) + args = tuple( + _as_cute(value) + for value in ( + radial_flat, + zonal.detach(), + inv_flat, + dst_ptr.detach(), + gate_flat, + out_flat, + ) + ) + device_identity = _device_key(radial) + with torch.cuda.device(device_identity[0]): + stream = cutlass_torch.current_stream() + key = ( + "gie_fwd", + *device_identity, + lmax, + channels, + has_gate, + tuple(radial_flat.stride()), + dst_ptr.dtype, + ) + compiled = _get_compiled( + key, + lambda: _build_forward(lmax, channels, has_gate), + (*args, stream), + ) + compiled(*args, stream) + return out + + def _launch_backward( + grad_out: Tensor, + radial: Tensor, + zonal: Tensor, + inv_sqrt_deg: Tensor, + dst: Tensor, + gate: Tensor, + out: Tensor, + lmax: int, + has_gate: bool, + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: + channels = radial.shape[2] + grad_out_flat = _flat_node(grad_out.detach().contiguous()) + radial_flat = _flat_radial(radial.detach()) + inv_flat = _flat_node(inv_sqrt_deg.detach()) + gate_flat = gate.detach().view(-1) + out_flat = _flat_node(out.detach()) + grad_radial = torch.empty( + radial.shape, + device=radial.device, + dtype=radial.dtype, + memory_format=torch.contiguous_format, + ) + grad_zonal = torch.empty_like(zonal, memory_format=torch.contiguous_format) + grad_inv = torch.empty_like(inv_sqrt_deg, memory_format=torch.contiguous_format) + grad_gate = torch.empty_like(gate, memory_format=torch.contiguous_format) + grad_radial_flat = _flat_radial(grad_radial) + grad_inv_flat = _flat_node(grad_inv) + grad_gate_flat = grad_gate.view(-1) + args = tuple( + _as_cute(value) + for value in ( + grad_out_flat, + radial_flat, + zonal.detach(), + inv_flat, + dst.detach(), + gate_flat, + out_flat, + grad_radial_flat, + grad_zonal, + grad_inv_flat, + grad_gate_flat, + ) + ) + device_identity = _device_key(radial) + with torch.cuda.device(device_identity[0]): + stream = cutlass_torch.current_stream() + key = _backward_compile_key( + device_identity, + lmax, + channels, + has_gate, + tuple(radial_flat.stride()), + dst.dtype, + ) + compiled = _get_compiled( + key, + lambda: _build_backward(lmax, channels, has_gate), + (*args, stream), + ) + compiled(*args, stream) + return grad_radial, grad_zonal, grad_inv, grad_gate + + @torch.library.custom_op( + "sezm_cute::gie_fused", mutates_args=(), device_types="cuda" + ) + def _gie_op( + radial: Tensor, + zonal: Tensor, + inv_sqrt_deg: Tensor, + dst: Tensor, + dst_ptr: Tensor, + gate: Tensor, + lmax: int, + has_gate: bool, + ) -> Tensor: + del dst + return _launch_forward( + radial, + zonal, + inv_sqrt_deg, + dst_ptr, + gate, + int(lmax), + bool(has_gate), + ) + + @_gie_op.register_fake + def _( + radial: Tensor, + zonal: Tensor, + inv_sqrt_deg: Tensor, + dst: Tensor, + dst_ptr: Tensor, + gate: Tensor, + lmax: int, + has_gate: bool, + ) -> Tensor: + del zonal, dst, dst_ptr, gate, has_gate + return radial.new_empty( + (inv_sqrt_deg.shape[0], (int(lmax) + 1) ** 2, radial.shape[2]) + ) + + @torch.library.custom_op( + "sezm_cute::gie_fused_bwd", mutates_args=(), device_types="cuda" + ) + def _gie_bwd_op( + grad_out: Tensor, + radial: Tensor, + zonal: Tensor, + inv_sqrt_deg: Tensor, + dst: Tensor, + gate: Tensor, + out: Tensor, + lmax: int, + has_gate: bool, + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: + return _launch_backward( + grad_out, + radial, + zonal, + inv_sqrt_deg, + dst, + gate, + out, + int(lmax), + bool(has_gate), + ) + + @_gie_bwd_op.register_fake + def _( + grad_out: Tensor, + radial: Tensor, + zonal: Tensor, + inv_sqrt_deg: Tensor, + dst: Tensor, + gate: Tensor, + out: Tensor, + lmax: int, + has_gate: bool, + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: + del grad_out, dst, out, lmax, has_gate + return ( + torch.empty_like(radial, memory_format=torch.contiguous_format), + torch.empty_like(zonal, memory_format=torch.contiguous_format), + torch.empty_like(inv_sqrt_deg, memory_format=torch.contiguous_format), + torch.empty_like(gate, memory_format=torch.contiguous_format), + ) + + def _gie_setup_context( + ctx: Any, + inputs: tuple[Any, ...], + output: Tensor, + ) -> None: + radial, zonal, inv_sqrt_deg, dst, _dst_ptr, gate, lmax, has_gate = inputs + ctx.save_for_backward(radial, zonal, inv_sqrt_deg, dst, gate, output) + ctx.lmax = int(lmax) + ctx.has_gate = bool(has_gate) + + def _gie_backward(ctx: Any, grad_out: Tensor) -> tuple[Any, ...]: + radial, zonal, inv_sqrt_deg, dst, gate, out = ctx.saved_tensors + grad_radial, grad_zonal, grad_inv, grad_gate = _gie_bwd_op( + grad_out, + radial, + zonal, + inv_sqrt_deg, + dst, + gate, + out, + ctx.lmax, + ctx.has_gate, + ) + return grad_radial, grad_zonal, grad_inv, None, None, grad_gate, None, None + + _gie_op.register_autograd(_gie_backward, setup_context=_gie_setup_context) + + +def gie_fused_cuda( + radial: Tensor, + zonal: Tensor, + inv_sqrt_deg: Tensor, + dst: Tensor, + gate: Tensor, + *, + n_nodes: int, + lmax: int, +) -> Tensor: + """Run the fused CUDA path after the caller has validated its contract.""" + if not SEZM_CUTE_GIE_AVAILABLE: + raise RuntimeError("CuTe DSL is unavailable") + boundaries = torch.arange( + n_nodes + 1, + device=dst.device, + dtype=dst.dtype, + ) + dst_ptr = torch.searchsorted(dst, boundaries) + has_gate = gate.numel() != 0 + kernel_gate = gate if has_gate else radial.new_ones((1,)) + return _gie_op( + radial, + zonal, + inv_sqrt_deg, + dst, + dst_ptr, + kernel_gate, + int(lmax), + has_gate, + ) + + +def maybe_run_cute_gie( + module: Any, + *, + n_nodes: int, + edge_cache: Any, + radial_feat: Tensor, + zonal_coupling: Tensor, +) -> Tensor | None: + """Run the opt-in path or return ``None`` for the eager fallback.""" + if ( + not is_cute_gie_enabled(radial_feat.device) + or not SEZM_CUTE_GIE_AVAILABLE + or bool(getattr(module, "training", True)) + or not radial_feat.is_cuda + or not validate_gie_contract( + module, n_nodes, edge_cache, radial_feat, zonal_coupling + ) + ): + return None + gate = getattr(edge_cache, "edge_src_gate", None) + if gate is None: + gate = radial_feat.new_empty((0,)) + return gie_fused_cuda( + radial_feat, + zonal_coupling, + edge_cache.inv_sqrt_deg, + edge_cache.dst, + gate, + n_nodes=n_nodes, + lmax=int(module.lmax), + ) + + +__all__ = [ + "SEZM_CUTE_GIE_AVAILABLE", + "gie_fused_cuda", + "is_cute_gie_enabled", + "maybe_run_cute_gie", + "validate_gie_contract", +] diff --git a/deepmd/kernels/cute/neo/k1.py b/deepmd/kernels/cute/neo/k1.py new file mode 100644 index 0000000000..279fd62ece --- /dev/null +++ b/deepmd/kernels/cute/neo/k1.py @@ -0,0 +1,2220 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Opt-in CuTe Neo K1 custom op for SeZM/DPA4 inference.""" + +from __future__ import ( + annotations, +) + +import threading +import weakref +from dataclasses import ( + dataclass, + field, +) +from functools import ( + lru_cache, +) +from types import ( + SimpleNamespace, +) +from typing import ( + Any, +) + +import torch +from torch import ( + Tensor, +) + +from deepmd.pt.model.descriptor.sezm_nn.norm import ( + EquivariantRMSNorm, +) + +from . import ( + runtime_policy, +) +from .k1_message_grid_packed import ( + is_supported_message_grid, +) + + +@dataclass(frozen=True) +class NeoK1RuntimeConfig: + """Architecture-specific choices for the compact Neo K1 path.""" + + native_sm90_path: bool = False + per_focus_so2_fwd_pair: bool = False + combined_so2_gate: bool = False + + +@dataclass(frozen=True) +class NeoK1Spec: + """Shape and branch contract for one Neo K1 replacement.""" + + lmax: int + node_lmax: int + mmax: int + full_dim: int + reduced_dim: int + channels: int + n_focus: int + focus_dim: int + hidden_channels: int + so2_layers: int + n_atten_head: int + radial_so2_mode: str + radial_so2_rank: int + so2_norm: bool + focus_compete: bool + message_node_so3: bool + atten_f_mix: bool + atten_v_proj: bool + atten_o_proj: bool + mlp_bias: bool + layer_scale: bool + use_so2_attn_res: bool + has_pre_so2_norm: bool + has_post_so2_norm: bool + + @property + def is_current_neo_target(self) -> bool: + return ( + self.lmax == 3 + and self.node_lmax == self.lmax + and self.mmax == 1 + and self.full_dim == 16 + and self.reduced_dim == 10 + and self.channels == 32 + and self.n_focus == 2 + and self.focus_dim == 32 + and self.hidden_channels == 64 + and self.so2_layers == 3 + and self.n_atten_head == 1 + and self.radial_so2_mode == "degree_channel" + and self.radial_so2_rank == 1 + and not self.so2_norm + and self.focus_compete + and self.message_node_so3 + and not self.atten_f_mix + and not self.atten_v_proj + and not self.atten_o_proj + and not self.mlp_bias + and not self.layer_scale + and not self.use_so2_attn_res + and not self.has_pre_so2_norm + and self.has_post_so2_norm + ) + + +def is_cute_k1_enabled() -> bool: + """Return whether the Neo CuTe K1 path is enabled by environment.""" + return runtime_policy.is_cute_infer_enabled() + + +def _is_identity(module: Any) -> bool: + return ( + module.__class__.__name__ == "Identity" + or module.__class__.__name__ == "_Identity" + ) + + +def get_neo_k1_spec(block: Any) -> NeoK1Spec: + """Extract the K1 contract from a SeZM interaction block.""" + so2 = block.so2_conv + return NeoK1Spec( + lmax=int(so2.lmax), + node_lmax=int(block.node_lmax), + mmax=int(so2.mmax), + full_dim=int(so2.ebed_dim_full), + reduced_dim=int(so2.reduced_dim), + channels=int(so2.channels), + n_focus=int(so2.n_focus), + focus_dim=int(so2.so2_focus_dim), + hidden_channels=int(so2.hidden_channels), + so2_layers=int(so2.mixing_layers), + n_atten_head=int(so2.n_atten_head), + radial_so2_mode=str(so2.radial_so2_mode), + radial_so2_rank=int(so2.radial_so2_rank), + so2_norm=bool(so2.so2_norm), + focus_compete=bool(so2.focus_compete), + message_node_so3=getattr(so2, "message_node_grid_product", None) is not None, + atten_f_mix=bool(so2.atten_f_mix), + atten_v_proj=getattr(so2, "attn_v_proj", None) is not None, + atten_o_proj=getattr(so2, "attn_o_proj", None) is not None, + mlp_bias=bool(so2.mlp_bias), + layer_scale=bool(so2.layer_scale), + use_so2_attn_res=bool(so2.use_so2_attn_res), + has_pre_so2_norm=not _is_identity(block.pre_so2_norm), + has_post_so2_norm=not _is_identity(block.post_so2_norm), + ) + + +def is_supported_neo_k1_block(block: Any) -> bool: + """Return whether this block can use the current Neo CuTe K1 path.""" + so2 = block.so2_conv + return ( + get_neo_k1_spec(block).is_current_neo_target + and bool(so2.focus_norm) + and not bool(so2.edge_cartesian) + and getattr(so2, "node_cartesian_tp", None) is None + and type(block.post_so2_norm) is EquivariantRMSNorm + and is_supported_message_grid(so2.message_node_grid_product) + ) + + +def _module_floating_state_uses_strict_fp32( + module: Any, + *, + require_floating_tensor: bool, +) -> bool: + """Check live floating parameters and buffers without reading device data.""" + saw_floating_tensor = False + for getter_name in ("parameters", "buffers"): + getter = getattr(module, getter_name, None) + if getter is None: + return False + for tensor in getter(): + if tensor.is_floating_point(): + saw_floating_tensor = True + if tensor.dtype != torch.float32: + return False + return saw_floating_tensor or not require_floating_tensor + + +def _module_uses_strict_fp32(module: Any) -> bool: + return _module_floating_state_uses_strict_fp32( + module, + require_floating_tensor=True, + ) + + +def _module_is_frozen(module: Any) -> bool: + """Return whether autograd cannot request gradients for module state.""" + parameters = getattr(module, "parameters", None) + return parameters is not None and not any( + parameter.requires_grad for parameter in parameters() + ) + + +def _tensor_is_aligned(tensor: Tensor, alignment: int = 16) -> bool: + return tensor.data_ptr() % alignment == 0 + + +def _module_state_is_aligned(module: Any, alignment: int = 16) -> bool: + """Check CuTe's declared alignment contract for parameters and buffers.""" + return all( + not tensor.is_floating_point() or _tensor_is_aligned(tensor, alignment) + for getter_name in ("parameters", "buffers") + for tensor in getattr(module, getter_name)() + ) + + +def _gate_expand_index_is_supported(block: Any) -> bool: + """Check the degree-to-gate map assumed by fused Neo gate kernels.""" + non_linearities = getattr( + getattr(block, "so2_conv", None), + "non_linearities", + None, + ) + if non_linearities is None: + return True + buffers = tuple( + expand_index + for non_linear in non_linearities + if (expand_index := getattr(non_linear, "expand_index", None)) is not None + and expand_index.numel() > 0 + ) + signature = tuple( + ( + tensor.data_ptr(), + tensor._version, + tensor.dtype, + tensor.device, + tuple(tensor.shape), + ) + for tensor in buffers + ) + cached = getattr(block, "_deepmd_cute_gate_expand_contract", None) + if cached is not None and cached[0] == signature: + return bool(cached[1]) + expected = torch.tensor( + [0, 1, 2, 0, 1, 2, 0, 1, 2], + dtype=torch.long, + device="cpu", + ) + supported = all( + torch.equal( + expand_index.detach().to(device="cpu", dtype=torch.long), + expected, + ) + for expand_index in buffers + ) + block._deepmd_cute_gate_expand_contract = (signature, supported) + return supported + + +def _gate_expand_index_structure_is_supported(block: Any) -> bool: + """Check graph-visible gate-index metadata without reading tensor values.""" + non_linearities = getattr( + getattr(block, "so2_conv", None), + "non_linearities", + None, + ) + if non_linearities is None: + return True + return all( + expand_index.dtype == torch.long and tuple(expand_index.shape) == (9,) + for non_linear in non_linearities + if (expand_index := getattr(non_linear, "expand_index", None)) is not None + and expand_index.numel() > 0 + ) + + +def _aligned_contiguous(tensor: Tensor, alignment: int = 16) -> Tensor: + """Return canonical storage satisfying CuTe's assumed alignment.""" + if tensor.is_contiguous() and _tensor_is_aligned(tensor, alignment): + return tensor + return tensor.clone(memory_format=torch.contiguous_format) + + +def _producer_modules_use_strict_fp32(modules: Any) -> bool: + return all( + _module_floating_state_uses_strict_fp32( + module, + require_floating_tensor=False, + ) + for module in modules + ) + + +def _dtypes_use_strict_fp32(dtypes: Any) -> bool: + return all(dtype == torch.float32 for dtype in dtypes) + + +def is_supported_k1_compute_capability( + compute_capability: tuple[int, int], +) -> bool: + """Return whether K1 supports this compute capability.""" + return runtime_policy.is_supported_k1_capability(compute_capability) + + +def _device_compute_capability(device: torch.device) -> tuple[int, int]: + device_index = device.index + if device_index is None: + device_index = torch.cuda.current_device() + return _cuda_compute_capability(device_index) + + +def _tensor_compute_capability(tensor: Any) -> tuple[int, int] | None: + """Resolve dispatch capability from an operand instead of global CUDA state.""" + device = getattr(tensor, "device", None) + if device is None or device.type != "cuda": + return None + return _device_compute_capability(device) + + +def _device_is_supported_for_k1(device: torch.device) -> bool: + # 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) + ) + + +def packed_wigner_edges_eligible( + *, + candidate: bool, + edge_count: int, + node_count: int, + destinations_sorted: bool, + runtime_dtypes: Any = (), +) -> bool: + """Finish packed eligibility from host-side edge-order provenance.""" + return ( + candidate + and edge_count > 0 + and destinations_sorted + and _dtypes_use_strict_fp32(runtime_dtypes) + and runtime_policy.k1_int32_indexing_is_safe( + edge_count=edge_count, + node_count=node_count, + ) + ) + + +def is_neo_k1_static_eligible( + block: Any, + *, + training: bool, + device: torch.device, + dtype: torch.dtype, +) -> bool: + """Check K1 conditions that are stable during one descriptor forward.""" + return ( + is_cute_k1_enabled() + and not training + and not bool(getattr(block, "training", False)) + and device.type == "cuda" + and _device_is_supported_for_k1(device) + and dtype == torch.float32 + and not torch.is_autocast_enabled(device.type) + and runtime_policy.uses_strict_fp32_matmul() + and _module_uses_strict_fp32(block) + and _module_is_frozen(block) + and _module_state_is_aligned(block) + and _gate_expand_index_structure_is_supported(block) + and getattr(block, "_deepmd_cute_k1_state", None) is not False + and is_supported_neo_k1_block(block) + ) + + +def is_neo_k1_runtime_eligible( + block: Any, + *, + training: bool, + device: torch.device, + dtype: torch.dtype, + edge_count: int, + node_count: int, + destinations_sorted: bool, +) -> bool: + """Return the exact strict-FP32 inference contract for K1 dispatch.""" + return ( + edge_count > 0 + and destinations_sorted + and is_neo_k1_static_eligible( + block, + training=training, + device=device, + dtype=dtype, + ) + ) + + +def is_packed_wigner_candidate( + *, + blocks: Any, + training: bool, + device: torch.device, + dtype: torch.dtype, + producer_modules: Any = (), + producer_dtypes: Any = (), + has_edge_src_gate: bool = False, +) -> bool: + """Check packed-Wigner conditions known before edge construction.""" + if has_edge_src_gate: + # K1 currently falls back for SFPG so eager must receive dense Wigner data. + return False + block_tuple = tuple(blocks) + if device.type == "cuda": + try: + compute_capability = _device_compute_capability(device) + except (AssertionError, RuntimeError): + packed_wigner_enabled = False + else: + packed_wigner_enabled = ( + is_cute_k1_enabled() + and runtime_policy.is_supported_k1_capability(compute_capability) + ) + else: + packed_wigner_enabled = False + if ( + not block_tuple + or not packed_wigner_enabled + or not _producer_modules_use_strict_fp32(producer_modules) + or not _dtypes_use_strict_fp32(producer_dtypes) + or torch.is_autocast_enabled(device.type) + ): + return False + return all( + is_neo_k1_static_eligible( + block, + training=training, + device=device, + dtype=dtype, + ) + for block in block_tuple + ) + + +class _RegistryEntry: + """Weakly retain a block so discarded models do not leak the registry.""" + + def __init__( + self, + block: Any, + config: Any, + *, + on_collect: Any | None = None, + ) -> None: + try: + self._block_ref = weakref.ref(block, on_collect) + except TypeError: + self._block_ref = lambda: block + self.config = config + + @property + def block(self) -> Any: + block = self._block_ref() + if block is None: + raise RuntimeError("the registered Neo K1 block has been released") + return block + + +@dataclass(frozen=True) +class _RegisteredK1State: + device_index: int + handle: int + config: NeoK1RuntimeConfig + + +@dataclass +class _RunnerState: + runner: Any | None + backward_calls: int = 0 + reservation_lock: threading.Lock = field( + default_factory=threading.Lock, + repr=False, + compare=False, + ) + + +_REGISTRY: dict[int, _RegistryEntry] = {} +_NEXT_HANDLE = 1 +_PACKED_RUNNER_CACHE: dict[tuple[str, int | None, int], _RunnerState] = {} +_PACKED_RUNNER_CACHE_LOCK = threading.Lock() + + +def _runner_compile_identity(runner: Any) -> tuple[int, int, int]: + """Return the runner's device-specific CuTe compilation identity.""" + return getattr(runner, "compile_identity", (-1, 0, 0)) + + +def _compile_on_runner_device( + compile_identity: tuple[int, int, int], + compiler: Any, + *args: Any, + **kwargs: Any, +) -> Any: + """Compile under the CUDA device represented by the cache key.""" + device_index = int(compile_identity[0]) + if device_index < 0: + return compiler(*args, **kwargs) + with torch.cuda.device(device_index): + return compiler(*args, **kwargs) + + +@lru_cache(maxsize=8) +def _compile_output_gate_backward( + compile_identity: tuple[int, int, int], eps: float +) -> Any: + from .k1_kernels.cute_neo_output_gate_backward import ( + compile_neo_output_gate_backward, + ) + + return _compile_on_runner_device( + compile_identity, + compile_neo_output_gate_backward, + eps, + ) + + +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, + ) + 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) + if hasattr(block, "_deepmd_cute_k1_state"): + delattr(block, "_deepmd_cute_k1_state") + if hasattr(block, "_deepmd_cute_gate_expand_contract"): + delattr(block, "_deepmd_cute_gate_expand_contract") + + +def _validate_gate_expand_index(block: Any) -> None: + """Pin the degree-to-gate map assumed by fused Neo gate kernels.""" + if not _gate_expand_index_is_supported(block): + raise ValueError( + "Neo K1 fused gate kernels require expand_index=[0,1,2,0,1,2,0,1,2]" + ) + + +@torch.compiler.disable +def _register_cute_k1_state( + block: Any, + device_index: int, + config: NeoK1RuntimeConfig, +) -> _RegisteredK1State | None: + """Publish cold K1 state eagerly before its handle enters a custom op.""" + _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) + 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. + block._deepmd_cute_k1_state = False + return None + state = _RegisteredK1State( + device_index=device_index, + handle=register_cute_k1_block(block, config), + config=config, + ) + block._deepmd_cute_k1_state = state + return state + + +@torch.compiler.disable +def prepare_cute_k1_blocks( + blocks: Any, + *, + training: bool, + device: torch.device, + dtype: torch.dtype, +) -> bool: + """Validate and register K1 state before model graph capture begins.""" + if training or device.type != "cuda" or dtype != torch.float32: + return False + block_tuple = tuple(blocks) + if not block_tuple: + return False + device_index = device.index + if device_index is None: + device_index = torch.cuda.current_device() + compute_capability = _cuda_compute_capability(device_index) + if not is_supported_k1_compute_capability(compute_capability): + return False + config = _architecture_default_config(compute_capability) + if not all( + is_neo_k1_static_eligible( + block, + training=training, + device=device, + dtype=dtype, + ) + for block in block_tuple + ): + return False + return all( + _register_cute_k1_state(block, device_index, config) is not None + for block in block_tuple + ) + + +def _runner_token_key( + runner_token: Tensor, *, path: str +) -> tuple[str, int | None, int]: + if runner_token.dtype != torch.uint8 or runner_token.numel() != 1: + raise ValueError(f"{path} Neo K1 runner token must be one uint8 value") + return ( + runner_token.device.type, + runner_token.device.index, + int(runner_token.data_ptr()), + ) + + +def _packed_runner_key(runner_token: Tensor) -> tuple[str, int | None, int]: + return _runner_token_key(runner_token, path="packed") + + +def _release_packed_runner( + key: tuple[str, int | None, int], + state: _RunnerState, +) -> None: + with _PACKED_RUNNER_CACHE_LOCK: + if _PACKED_RUNNER_CACHE.get(key) is state: + _PACKED_RUNNER_CACHE.pop(key, None) + + +def _store_packed_runner(runner_token: Tensor, runner: Any) -> None: + key = _packed_runner_key(runner_token) + state = _RunnerState(runner=runner) + with _PACKED_RUNNER_CACHE_LOCK: + old_state = _PACKED_RUNNER_CACHE.get(key) + if old_state is not None: + raise RuntimeError("packed Neo K1 runner token is already outstanding") + else: + _PACKED_RUNNER_CACHE[key] = state + weakref.finalize(runner_token, _release_packed_runner, key, state) + + +def _borrow_packed_runner(runner_token: Tensor) -> Any | None: + key = _packed_runner_key(runner_token) + with _PACKED_RUNNER_CACHE_LOCK: + state = _PACKED_RUNNER_CACHE.get(key) + if state is None: + raise RuntimeError("packed Neo K1 runner token is not outstanding") + with state.reservation_lock: + # Transfer the forward runner once; retained VJPs rebuild isolated workspaces. + runner = state.runner + state.runner = None + state.backward_calls += 1 + return runner + + +def _edge_src_gate_arg(edge_src_gate: Tensor) -> Tensor | None: + return None if edge_src_gate.numel() == 0 else edge_src_gate + + +def _layout_like(tensor: Tensor, like: Tensor) -> Tensor: + if tensor.shape == like.shape and tensor.stride() == like.stride(): + return tensor + out = torch.empty_strided( + like.shape, + like.stride(), + device=like.device, + dtype=like.dtype, + ) + out.copy_(tensor) + return out + + +def _grad_layout_like(tensor: Tensor, like: Tensor, *, skip: bool) -> Tensor: + if skip: + return tensor + return _layout_like(tensor, like) + + +def _assert_grad_meta_contract( + actual: Tensor, + expected_like: Tensor, + *, + name: str, + expected_stride: tuple[int, ...] | None = None, +) -> None: + """Fail before AOT consumes a gradient that contradicts register_fake.""" + stride = expected_like.stride() if expected_stride is None else expected_stride + if ( + actual.shape != expected_like.shape + or actual.dtype != expected_like.dtype + or actual.device != expected_like.device + or actual.stride() != stride + ): + raise RuntimeError( + f"Neo K1 {name} gradient violates the custom-op meta contract: " + f"got shape={tuple(actual.shape)} stride={actual.stride()}, expected " + f"shape={tuple(expected_like.shape)} stride={stride}" + ) + + +def _assert_not_cuda_graph_capturing(path: str, tensor: Tensor) -> None: + """Reject stateful paths that cannot be safely captured and replayed.""" + if tensor.is_cuda and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"Neo K1 {path} cannot run inside direct CUDA graph capture; " + "use torch.compile without wrapping this stateful custom op in " + "torch.cuda.graph" + ) + + +def _fake_x_wide_grad_like(x: Tensor, *, skip: bool) -> Tensor: + """Fake the native K1 grad-x layout when layout restore is skipped. + + The manual SO3 pre-mix backward returns a degree-major view with shape + ``(N, D, 1, C)`` and stride ``(C, N*C, C, 1)``. If the fake kernel + promises ``empty_like(x)`` instead, AOTAutograd emits stride assertions for + a contiguous ``(N, D, 1, C)`` gradient and rejects the runtime result. + """ + if not skip: + return torch.empty_like(x) + if x.ndim != 4: + return torch.empty_like(x) + n_node = x.shape[0] + channels = x.shape[-1] + return torch.empty_strided( + x.shape, + (channels, n_node * channels, channels, 1), + device=x.device, + dtype=x.dtype, + ) + + +def _equivariant_rmsnorm_backward(norm: Any, x: Tensor, grad_out: Tensor) -> Tensor: + # Matches EquivariantRMSNorm.forward for the inference case. Parameters are + # frozen; only the input gradient is needed for forces/stress. + in_dtype = x.dtype + xf = x.to(dtype=norm.dtype) + gf = grad_out.to(dtype=norm.dtype) + x0_in = xf[:, :1, :, :] + xt = xf[:, 1:, :, :] + x0 = x0_in - x0_in.mean(dim=-1, keepdim=True) + + mean_variance = x0.square().sum(dim=(1, 3)) * norm.balance_weight[0] + if xt.numel() > 0: + mean_variance = mean_variance + torch.einsum( + "ndfc,d->nf", xt * xt, norm.balance_weight[1:] + ) + inv = torch.rsqrt(mean_variance + norm.eps_tensor).unsqueeze(1).unsqueeze(-1) + expanded_scale = torch.index_select( + norm.adam_scale, dim=0, index=norm.expand_index + ).unsqueeze(0) + + grad_pre = gf * expanded_scale + grad_x0 = grad_pre[:, :1, :, :] + grad_xt = grad_pre[:, 1:, :, :] + + dinv = (grad_x0 * x0).sum(dim=(1, 3), keepdim=True) + if xt.numel() > 0: + dinv = dinv + (grad_xt * xt).sum(dim=(1, 3), keepdim=True) + dvar = -0.5 * dinv * inv.pow(3) + + grad_centered = grad_x0 * inv + dvar * (2.0 * norm.balance_weight[0] * x0) + grad_x0_in = grad_centered - grad_centered.mean(dim=-1, keepdim=True) + if xt.numel() == 0: + return grad_x0_in.to(dtype=in_dtype) + grad_xt_in = grad_xt * inv + dvar * ( + 2.0 * norm.balance_weight[1:].view(1, -1, 1, 1) * xt + ) + 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 + weight = linear.weight.view( + linear.lmax + 1, + linear.in_channels, + linear.n_focus, + linear.out_channels, + ) + weight_expanded = torch.index_select(weight, dim=0, index=linear.expand_index) + return torch.einsum("ndfo,difo->ndfi", grad_out, weight_expanded) + + +def _focus_linear_forward(linear: Any, x: Tensor) -> Tensor: + weight = linear.weight.view(linear.in_channels, linear.n_focus, linear.out_channels) + out = torch.einsum("bfi,ifo->bfo", x, weight) + if linear.use_bias: + out = out + linear.bias.view(linear.n_focus, linear.out_channels).unsqueeze(0) + return out + + +def _focus_linear_backward_input(linear: Any, grad_out: Tensor) -> Tensor: + weight = linear.weight.view(linear.in_channels, linear.n_focus, linear.out_channels) + return torch.einsum("bfo,ifo->bfi", grad_out, weight) + + +def _swiglu_forward(x: Tensor) -> Tensor: + gate, value = torch.chunk(x, chunks=2, dim=-1) + return gate * torch.sigmoid(gate) * value + + +def _swiglu_backward_input(x: Tensor, grad_out: Tensor) -> Tensor: + gate, value = torch.chunk(x, chunks=2, dim=-1) + sig = torch.sigmoid(gate) + grad_gate = grad_out * value * (sig + gate * sig * (1.0 - sig)) + grad_value = grad_out * gate * sig + return torch.cat([grad_gate, grad_value], dim=-1) + + +def _frame_expand_forward(module: Any, coeff: Tensor) -> Tensor: + weight = module.weight.index_select(0, module.degree_index) + return torch.einsum("ndfi,dio->ndfo", coeff, weight) + + +def _frame_expand_backward_input(module: Any, grad_out: Tensor) -> Tensor: + weight = module.weight.index_select(0, module.degree_index) + return torch.einsum("ndfo,dio->ndfi", grad_out, weight) + + +def _frame_contract_backward_input(module: Any, grad_out: Tensor) -> Tensor: + weight = module.weight.index_select(0, module.degree_index) + return torch.einsum("ndfo,dio->ndfi", grad_out, weight) + + +def _neo_so2_linear_backward_input_with_residual( + so2_linear: Any, + grad_out: Tensor, + residual: Tensor, + *, + inplace_residual: bool = False, + out: Tensor | None = None, +) -> Tensor: + from .k1_so2linear import ( + cached_neo_so2_linear_weights, + ) + + w0, wpair = cached_neo_so2_linear_weights(so2_linear) + cache = getattr(so2_linear, "_deepmd_cute_neo_manual_weights_t", None) + cache_key = (w0.data_ptr(), wpair.data_ptr(), w0.dtype, w0.device) + if ( + cache is None + or cache[0] is not w0 + or cache[1] is not wpair + or cache[2] != cache_key + ): + w0_t = w0.transpose(1, 2).contiguous() + wpair_t = wpair.transpose(1, 2).contiguous() + so2_linear._deepmd_cute_neo_manual_weights_t = ( + w0, + wpair, + cache_key, + w0_t, + wpair_t, + ) + else: + w0_t, wpair_t = cache[3], cache[4] + if inplace_residual: + if out is not None: + raise ValueError("SO2 backward cannot select both in-place and out storage") + from .k1_so2linear import ( + neo_so2_linear_backward_residual_inplace, + ) + + return neo_so2_linear_backward_residual_inplace( + residual, + grad_out, + w0_t, + wpair_t, + ) + if grad_out is residual: + from .k1_gate_structural import ( + focus_major_so2_backward_with_folded_residual_out, + ) + + folded_cache = getattr( + so2_linear, + "_deepmd_cute_neo_manual_folded_weights_t", + None, + ) + folded_key = ( + w0_t.data_ptr(), + wpair_t.data_ptr(), + w0_t.dtype, + w0_t.device, + ) + if ( + folded_cache is None + or folded_cache[0] is not w0_t + or folded_cache[1] is not wpair_t + or folded_cache[2] != folded_key + ): + w0_folded_t = w0_t.clone() + wpair_folded_t = wpair_t.clone() + w0_folded_t.diagonal(dim1=-2, dim2=-1).add_(1.0) + wpair_folded_t.diagonal(dim1=-2, dim2=-1).add_(1.0) + folded_cache = ( + w0_t, + wpair_t, + folded_key, + w0_folded_t, + wpair_folded_t, + ) + so2_linear._deepmd_cute_neo_manual_folded_weights_t = folded_cache + if out is None: + out = grad_out.new_empty(grad_out.shape) + return focus_major_so2_backward_with_folded_residual_out( + grad_out, + folded_cache[3], + folded_cache[4], + out=out, + ) + + raise RuntimeError("Neo K1 SO2 backward reached an unsupported storage pattern") + + +def _so3_grid_cross_glu_flat_backward( + net: Any, + query_flat: Tensor, + context_flat: Tensor, + grad_out_flat: Tensor, +) -> tuple[Tensor, Tensor]: + """Input-gradient for Neo's flat cross SO3GridNet GLU branch.""" + if ( + net.layout != "flat" + or net.mode != "cross" + or net.op_type != "glu" + or net.frame_expand is None + or net.frame_contract is None + ): + raise NotImplementedError("manual grid backward supports Neo cross/flat/glu") + + q_dtype = query_flat.dtype + c_dtype = context_flat.dtype + n_batch, coeff_dim, _ = query_flat.shape + n_focus = net.n_focus + channels = net.channels + n_frames = net.n_frames + expanded = net.expanded_channels + + query = query_flat.reshape(n_batch, coeff_dim, n_focus, channels) + context = context_flat.reshape(n_batch, coeff_dim, n_focus, channels) + scalar_pair = torch.cat( + [query[:, 0, :, :], context[:, 0, :, :]], + dim=-1, + ).to(dtype=net.dtype) + + left = _frame_expand_forward(net.frame_expand, query).to(dtype=net.dtype) + right = _frame_expand_forward(net.frame_expand, context).to(dtype=net.dtype) + left_view = left.reshape(n_batch, coeff_dim, n_focus, n_frames, channels) + right_view = right.reshape(n_batch, coeff_dim, n_focus, n_frames, channels) + to_grid = net.projector.to_grid_mat.reshape( + net.projector.grid_size, + coeff_dim, + n_frames, + ) + from_grid = net.projector.from_grid_mat.reshape( + coeff_dim, + n_frames, + net.projector.grid_size, + ) + left_grid = torch.einsum("gdk,ndfkc->ngfc", to_grid, left_view) + right_grid = torch.einsum("gdk,ndfkc->ngfc", to_grid, right_view) + 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 + ) + if net.residual_scale is not None: + grad = grad * net.residual_scale.reshape(1, 1, n_focus, channels) + grad_scalar_flat = _frame_contract_backward_input(net.frame_contract, grad) + grad_scalar_view = grad_scalar_flat.reshape( + n_batch, + coeff_dim, + n_focus, + n_frames, + channels, + ) + + grad_coeff = grad_scalar_view * scalar_gate[:, None, :, None, :] + grad_scalar_gate = (grad_scalar_view * coeff_view).sum(dim=(1, 3)) + grad_scalar_out = grad_scalar_view[:, 0, :, net.frame_zero_index, :] + grad_scalar_logits = grad_scalar_gate * scalar_gate * (1.0 - scalar_gate) + grad_scalar_pair = _focus_linear_backward_input(net.scalar_gate, grad_scalar_logits) + grad_scalar_pair = grad_scalar_pair + _swiglu_backward_input( + scalar_pair, + grad_scalar_out, + ) + + grad_grid = torch.einsum("dkg,ndfkc->ngfc", from_grid, grad_coeff) + grad_left_grid = grad_grid * right_grid + grad_right_grid = grad_grid * left_grid + grad_left = torch.einsum("gdk,ngfc->ndfkc", to_grid, grad_left_grid).reshape( + n_batch, + coeff_dim, + n_focus, + expanded, + ) + grad_right = torch.einsum("gdk,ngfc->ndfkc", to_grid, grad_right_grid).reshape( + n_batch, + coeff_dim, + n_focus, + expanded, + ) + grad_query = _frame_expand_backward_input(net.frame_expand, grad_left) + 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), + ) + + +def _final_manual_backward(runner: Any, grad_out: Tensor) -> tuple[Tensor, Tensor]: + so2 = runner.so2 + block = runner.block + n_node = runner.node_count + if runner.use_full_node: + grad_so2_out = grad_out + else: + grad_so2_out = grad_out[:, : block.mp_ebed_dim, :, :] + + 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( + block.post_so2_norm, + post_norm_in, + grad_so2_out, + ) + grad_post_mix = _so3_linear_backward_input( + so2.post_focus_mix, + post_in, + 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 + runner.message_grid_product = None + from .k1_message_grid_packed import ( + run_packed_message_grid_backward, + ) + + grad_out_gate_flat, grad_grid_context = run_packed_message_grid_backward( + so2.message_node_grid_product, + out_gate_flat, + x_wide, + grad_post_mix, + product_flat=message_grid_product, + ) + del message_grid_product + else: + grad_out_gate_flat, grad_grid_context = _so3_grid_cross_glu_flat_backward( + so2.message_node_grid_product, + out_gate_flat, + x_wide, + 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) + else: + grad_out_gate_flat = grad_post_mix + grad_x_wide_down = torch.zeros( + n_node, + 16 * 64, + device=x_wide.device, + dtype=x_wide.dtype, + ).view(n_node, 16, 64) + + grad_phase = grad_out_gate_flat.contiguous() + output_gate_backward = _compile_output_gate_backward( + _runner_compile_identity(runner), + float(so2.attn_output_gate_norm.eps), + ) + output_gate_backward( + grad_phase.view(n_node, 16 * 64), + phase.contiguous().view(n_node, 16 * 64), + x_wide.contiguous().view(n_node, 16 * 64), + so2.attn_output_gate_norm.adam_scale.detach() + .float() + .reshape(2, 32) + .contiguous(), + so2.adamw_attn_gate_w.detach().float().reshape(32, 2, 1).contiguous(), + grad_phase.view(n_node, 16 * 64), + grad_x_wide_down.view(n_node, 16 * 64), + ) + return grad_phase.reshape_as(phase), grad_x_wide_down + + +def _qk_manual_backward( + runner: Any, + grad_logits: Tensor, +) -> 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 + k_node = runner.k_node + grad_q_node = getattr(runner, "grad_q_node", None) + grad_k_node = getattr(runner, "grad_k_node", None) + if grad_q_node is None: + grad_q_node = torch.empty_like( + q_node, + memory_format=torch.contiguous_format, + ) + grad_k_node = torch.empty_like( + k_node, + memory_format=torch.contiguous_format, + ) + runner.grad_q_node = grad_q_node + runner.grad_k_node = grad_k_node + if not grad_q_node.is_contiguous() or not grad_k_node.is_contiguous(): + raise RuntimeError("Neo Q/K backward buffers must be compact N x 2 x 32") + grad_q_node.zero_() + grad_k_node.zero_() + runner.qk_edge_backward( + grad_logits.contiguous(), + q_node, + k_node, + runner.src_i32, + runner.dst_i32, + grad_q_node, + grad_k_node, + ) + grad_x_wide = getattr(runner, "grad_x_wide_qk", None) + if grad_x_wide is None or grad_x_wide.shape != x_wide.shape: + grad_x_wide = torch.empty( + x_wide.shape, + device=x_wide.device, + dtype=x_wide.dtype, + ) + runner.grad_x_wide_qk = grad_x_wide + runner.qk_node_input_adjoint( + x_l0.contiguous(), + grad_q_node, + grad_k_node, + so2.attn_q_proj.weight.detach().float().view(32, 2, 32).contiguous(), + so2.attn_k_proj.weight.detach().float().view(32, 2, 32).contiguous(), + so2.attn_qk_norm.adam_scale.detach().float().contiguous(), + grad_x_wide.view(n_node, 16 * 64), + ) + return grad_x_wide + + +def _native_edge_major_stack_grad(grad_stack_out: Tensor) -> Tensor: + """Validate the exact in-place Phase-C adjoint consumed by final SO2.""" + edge_count = grad_stack_out.shape[0] + expected_shape = (edge_count, 2, 10, 32) + expected_stride = (2 * 10 * 32, 10 * 32, 32, 1) + if ( + grad_stack_out.shape != expected_shape + or grad_stack_out.stride() != expected_stride + or grad_stack_out.dtype != torch.float32 + ): + raise RuntimeError( + "in-place Phase-C adjoint requires compact edge-major storage: " + f"got shape={tuple(grad_stack_out.shape)} " + f"stride={grad_stack_out.stride()} dtype={grad_stack_out.dtype}, " + f"expected shape={expected_shape} stride={expected_stride}" + ) + return grad_stack_out + + +def _phase_c_layout_grad_stack(runner: Any) -> Tensor: + """Select aliased edge-major or ordinary Phase-C output storage.""" + if runner.phase_c_y is not None: + raise RuntimeError( + "in-place Phase-C adjoint requires the folded single-input stack" + ) + grad_mixed = runner.grad_mixed_slab + if grad_mixed is None or ( + grad_mixed.untyped_storage()._cdata + == runner.phase_c_stack.untyped_storage()._cdata + ): + raise RuntimeError( + "in-place Phase-C adjoint must retain a distinct grad_mixed_slab" + ) + # Phase C owns every destination edge and stores G at the same + # edge/focus address only after its complete 10x32 input fragment is + # register-resident. Final SO2 consumes G before gate scratch reuses y2. + return _native_edge_major_stack_grad(runner.phase_c_stack) + + +def _stack_backward_manual(runner: Any, grad_stack_out: Tensor) -> Tensor: + cur_grad = _native_edge_major_stack_grad(grad_stack_out) + for cache_index in range(len(runner.stack_caches) - 1, -1, -1): + cache = runner.stack_caches[cache_index] + residual_grad = cur_grad + if cache.final: + grad_y = cur_grad + else: + if runner.config.combined_so2_gate: + assert runner.combined_gate_backward is not None + runner.combined_gate_backward( + cur_grad.view(-1, 10 * 32), + cache.y.detach().view(-1, 10 * 32), + cache.non_linear.gate_linear.weight.detach() + .view(32, 2, 3 * 32) + .contiguous(), + runner.grad_y, + ) + grad_y = runner.grad_y.view_as(cache.y) + else: + grad_gate_logits = runner._run_structural_gate_backward( + runner.structural_gate_backward, + cur_grad, + cache.y, + cache.logits, + runner.grad_y.view_as(cache.y), + grad_logits=runner.grad_gate_logits, + overwrite_logits=True, + ) + grad_y = runner.grad_y.view_as(cache.y) + runner._focus_major_gate_linear_backward_add( + grad_y, + grad_gate_logits, + cache.non_linear.gate_linear.weight.detach(), + ) + linear = runner.so2.so2_linears[cache_index] + so2_out = runner.grad_mixed_slab if cache.final else None + cur_grad = _neo_so2_linear_backward_input_with_residual( + linear, + grad_y, + residual_grad, + inplace_residual=not cache.final, + out=so2_out, + ) + return cur_grad + + +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) + 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: + grad_x = torch.zeros_like(runner.x) + grad_x[:, : block.mp_ebed_dim, :, :] = grad_x_pre + expected_stride = ( + grad_x.shape[-1], + grad_x.shape[0] * grad_x.shape[-1], + grad_x.shape[-1], + 1, + ) + if grad_x.stride() != expected_stride: + grad_x = grad_x.permute(1, 0, 2, 3).contiguous().permute(1, 0, 2, 3) + return grad_x + + +def _make_edge_cache( + *, + src: Tensor, + dst: Tensor, + d_full: Tensor, + dt_full: Tensor, + edge_env: Tensor, + edge_src_gate: Tensor, +) -> Any: + return SimpleNamespace( + src=src, + dst=dst, + D_full=d_full, + Dt_full=dt_full, + edge_env=edge_env, + edge_src_gate=_edge_src_gate_arg(edge_src_gate), + D_to_m_cache={}, + Dt_from_m_cache={}, + ) + + +def _build_runner( + handle: int, + x: Tensor, + d_full: Tensor, + dt_full: Tensor, + radial_feat: Tensor, + edge_env: Tensor, + src: Tensor, + dst: Tensor, + dst_ptr: Tensor, + source_order: Tensor, + source_ptr: Tensor, + edge_src_gate: Tensor, +) -> Any: + """Build a runner after enforcing CuTe's runtime pointer contract. + + This function executes below the custom-op boundary, including for the + compile-visible thin path. Exact ``data_ptr`` checks are therefore safe + here and repair contiguous offset views without introducing a Dynamo graph + break above the op. + """ + entry = _REGISTRY[int(handle)] + config = entry.config + if config.native_sm90_path: + from .sm90_k1.runner import NeoSm90K1Runner as Runner + else: + from .k1_runner import NeoFullCuteBackward as Runner + + x = _aligned_contiguous(x) + shared_wigner_storage = d_full.data_ptr() == dt_full.data_ptr() + d_full = _aligned_contiguous(d_full) + dt_full = d_full if shared_wigner_storage else _aligned_contiguous(dt_full) + radial_feat = _aligned_contiguous(radial_feat) + edge_env = _aligned_contiguous(edge_env) + src = _aligned_contiguous(src) + dst = _aligned_contiguous(dst) + dst_ptr = _aligned_contiguous(dst_ptr) + source_order = _aligned_contiguous(source_order) + source_ptr = _aligned_contiguous(source_ptr) + edge_src_gate = _aligned_contiguous(edge_src_gate) + edge_cache = _make_edge_cache( + src=src, + dst=dst, + d_full=d_full, + dt_full=dt_full, + edge_env=edge_env, + edge_src_gate=edge_src_gate, + ) + record = SimpleNamespace(edge_cache=edge_cache) + with torch.cuda.device(x.device), torch.no_grad(): + runner = Runner( + torch, + entry.block, + record, + x, + d_full, + dt_full, + radial_feat, + dst_ptr, + source_order, + source_ptr, + runtime_config=config, + ) + return runner + + +def _k1_packed_direct_forward_impl( + handle: int, + x: Tensor, + d_full: Tensor, + dt_full: Tensor, + radial_feat: Tensor, + edge_env: Tensor, + src: Tensor, + dst: Tensor, + dst_ptr: Tensor, + source_order: Tensor, + source_ptr: Tensor, + edge_src_gate: Tensor, +) -> tuple[Tensor, Tensor]: + _assert_not_cuda_graph_capturing("packed forward", x) + runner = _build_runner( + handle, + x, + d_full, + dt_full, + radial_feat, + edge_env, + src, + dst, + dst_ptr, + source_order, + source_ptr, + edge_src_gate, + ) + runner_token = torch.empty( + (1,), + device=x.device, + dtype=torch.uint8, + ) + _store_packed_runner(runner_token, runner) + return _layout_like(runner.final.detach(), x), runner_token + + +def _k1_backward_from_runner_current_device( + handle: int, + grad_out: Tensor, + x: Tensor, + d_full: Tensor, + dt_full: Tensor, + radial_feat: Tensor, + edge_env: Tensor, + edge_src_gate: Tensor, + runner: Any, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + del handle + grad_x, grad_d, grad_dt, grad_radial = _runner_backward_manual( + runner, + grad_out, + ) + + grad_edge_env = runner.grad_edge.view_as(edge_env) + if edge_src_gate.numel() != 0: + factor = edge_src_gate.reshape_as(edge_env).float().clamp_min(0.0).sqrt() + grad_edge_env = grad_edge_env * factor.to(dtype=grad_edge_env.dtype) + grad_edge_env = ( + grad_edge_env.clone() if grad_edge_env._base is not None else grad_edge_env + ) + grad_edge_env.masked_fill_(edge_env <= 0, 0) + grad_x_out = _grad_layout_like(grad_x, x, skip=True) + grad_d_out = _grad_layout_like(grad_d, d_full, skip=True) + grad_dt_out = _grad_layout_like(grad_dt, dt_full, skip=True) + grad_radial_out = _grad_layout_like(grad_radial, radial_feat, skip=True) + grad_edge_env_out = _grad_layout_like(grad_edge_env, edge_env, skip=True) + if x.ndim == 4: + n_node = x.shape[0] + channels = x.shape[-1] + grad_x_stride = (channels, n_node * channels, channels, 1) + else: + grad_x_stride = x.stride() + for name, actual, expected_like, expected_stride in ( + ("x", grad_x_out, x, grad_x_stride), + ("D", grad_d_out, d_full, d_full.stride()), + ("Dt", grad_dt_out, dt_full, dt_full.stride()), + ("radial", grad_radial_out, radial_feat, radial_feat.stride()), + ("edge_env", grad_edge_env_out, edge_env, edge_env.stride()), + ): + _assert_grad_meta_contract( + actual, + expected_like, + name=name, + expected_stride=expected_stride, + ) + return ( + grad_x_out, + grad_d_out, + grad_dt_out, + grad_radial_out, + grad_edge_env_out, + ) + + +def _k1_backward_from_runner( + handle: int, + grad_out: Tensor, + x: Tensor, + d_full: Tensor, + dt_full: Tensor, + radial_feat: Tensor, + edge_env: Tensor, + edge_src_gate: Tensor, + runner: Any, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + """Run all compilation and launches on the operand's CUDA device.""" + with torch.cuda.device(grad_out.device): + return _k1_backward_from_runner_current_device( + handle, + grad_out, + x, + d_full, + dt_full, + radial_feat, + edge_env, + edge_src_gate, + runner, + ) + + +def _k1_packed_direct_backward_impl( + handle: int, + grad_out: Tensor, + x: Tensor, + d_full: Tensor, + dt_full: Tensor, + radial_feat: Tensor, + edge_env: Tensor, + src: Tensor, + dst: Tensor, + dst_ptr: Tensor, + source_order: Tensor, + source_ptr: Tensor, + edge_src_gate: Tensor, + runner_token: Tensor, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + _assert_not_cuda_graph_capturing("packed backward", grad_out) + runner = _borrow_packed_runner(runner_token) + if runner is None: + runner = _build_runner( + handle, + x, + d_full, + dt_full, + radial_feat, + edge_env, + src, + dst, + dst_ptr, + source_order, + source_ptr, + edge_src_gate, + ) + return _k1_backward_from_runner( + handle, + grad_out, + x, + d_full, + dt_full, + radial_feat, + edge_env, + edge_src_gate, + runner, + ) + + +def _runner_backward_manual( + runner: Any, + grad_out: Tensor, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: + if getattr(runner, "uses_native_sm90_path", False): + return runner.input_adjoint(grad_out) + + so2 = runner.so2 + n_edge = runner.edge_count + + grad_phase_c_out, grad_x_wide_down = _final_manual_backward(runner, grad_out) + # Keep the large K1 slabs out of the readout-backward allocation crest. + runner.ensure_backward_workspace() + + from .k1_kernels.cute_neo_phase_c_backward_layout_runner import ( + NeoPhaseCBackwardLayoutOutputs, + ) + + layout_outputs = getattr(runner, "phase_c_layout_outputs", None) + if layout_outputs is None: + layout_outputs = NeoPhaseCBackwardLayoutOutputs( + grad_stack=_phase_c_layout_grad_stack(runner), + grad_wigner_dt=runner.grad_dt, + grad_logits=runner.grad_logits, + grad_edge=runner.grad_edge, + grad_z_partial=runner.grad_z_partial, + grad_z=runner.grad_z, + grad_focus_src=torch.empty( + 2, + runner.edge_count, + 32, + device=runner.x.device, + dtype=torch.float32, + ), + ) + runner.phase_c_layout_outputs = layout_outputs + runner.phase_c_layout_backward( + grad_phase_c_out.contiguous(), + runner.phase_c_stack.detach().contiguous(), + runner.dt.detach().contiguous(), + runner.alpha, + 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, + runner.denom, + runner.focus_gate_src.detach().contiguous(), + so2.adamw_focus_compete_w.detach().float().contiguous(), + so2.focus_compete_norm.adam_scale.detach().float().reshape(2, 32).contiguous(), + layout_outputs, + ) + grad_stack_out = _native_edge_major_stack_grad(layout_outputs.grad_stack) + grad_focus_src_focus = layout_outputs.grad_focus_src + + grad_x_wide_qk = _qk_manual_backward( + runner, + runner.grad_logits.view(n_edge, 2), + ) + grad_mixed = _stack_backward_manual(runner, grad_stack_out) + + grad_mixed_focus = grad_mixed + if not grad_mixed_focus.is_contiguous(): + grad_mixed_focus = grad_mixed_focus.contiguous() + x_wide_flat = runner.x_wide.detach().contiguous().view(runner.node_count, 16 * 64) + radial_state = runner.radial_compact + from .k1_radial_phase_a_node import ( + run_neo_radial_phase_a_backward_node_tiled, + ) + + run_neo_radial_phase_a_backward_node_tiled( + grad_mixed_focus.view(runner.edge_count, 2 * 10 * 32), + runner.grad_logits, + radial_state, + runner.so2.radial_degree_mixer.channel_basis.detach().view(64).contiguous(), + x_wide_flat, + runner.source_order_i32, + runner.source_ptr_i32, + runner.d.detach(), + grad_focus_src_focus=grad_focus_src_focus, + batched_radial_projection_weight=runner.batched_radial_projection_weight, + grad_x_wide=runner.grad_x_wide_phase_a, + grad_d_full=runner.grad_d, + grad_radial_m0=runner.grad_radial_flat, + validate_csr=runtime_policy.is_cute_strict_enabled(), + ) + grad_radial = runner.grad_radial_flat.view_as(runner.radial) + grad_x_wide_phase_a = runner.grad_x_wide_phase_a.view_as(runner.x_wide) + grad_d = runner.grad_d + + grad_x_wide_total = grad_x_wide_phase_a + grad_x_wide_total.add_(grad_x_wide_qk) + grad_x_wide_total.add_(grad_x_wide_down) + grad_x = _x_wide_manual_backward(runner, grad_x_wide_total) + return grad_x, grad_d, runner.grad_dt, grad_radial + + +def _stateful_custom_op_tags() -> tuple[Any, ...] | None: + """Tag hidden-state runner ops as unsafe for direct CUDA graph capture.""" + tag_type = getattr(getattr(torch, "_C", None), "Tag", None) + cudagraph_unsafe = getattr(tag_type, "cudagraph_unsafe", None) + if cudagraph_unsafe is None: + return None + return (cudagraph_unsafe,) + + +_K1_CUSTOM_OP_TAGS = _stateful_custom_op_tags() +_k1_packed_direct_op = torch.library.custom_op( + "sezm_cute::k1_packed_direct", mutates_args=(), tags=_K1_CUSTOM_OP_TAGS +)(_k1_packed_direct_forward_impl) +_k1_packed_direct_bwd_op = torch.library.custom_op( + "sezm_cute::k1_packed_direct_bwd", mutates_args=(), tags=_K1_CUSTOM_OP_TAGS +)(_k1_packed_direct_backward_impl) + + +@_k1_packed_direct_op.register_fake +def _( + handle: int, + x: Tensor, + d_full: Tensor, + dt_full: Tensor, + radial_feat: Tensor, + edge_env: Tensor, + src: Tensor, + dst: Tensor, + dst_ptr: Tensor, + source_order: Tensor, + source_ptr: Tensor, + edge_src_gate: Tensor, +) -> tuple[Tensor, Tensor]: + del handle + del d_full, dt_full, radial_feat, edge_env, src, dst, dst_ptr + del source_order, source_ptr, edge_src_gate + return ( + torch.empty_like(x), + torch.empty((1,), device=x.device, dtype=torch.uint8), + ) + + +@_k1_packed_direct_bwd_op.register_fake +def _( + handle: int, + grad_out: Tensor, + x: Tensor, + d_full: Tensor, + dt_full: Tensor, + radial_feat: Tensor, + edge_env: Tensor, + src: Tensor, + dst: Tensor, + dst_ptr: Tensor, + source_order: Tensor, + source_ptr: Tensor, + edge_src_gate: Tensor, + runner_token: Tensor, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + del handle, grad_out, src, dst, dst_ptr, source_order, source_ptr + del edge_src_gate, runner_token + return ( + _fake_x_wide_grad_like(x, skip=True), + torch.empty_like(d_full), + torch.empty_like(dt_full), + torch.empty_like(radial_feat), + torch.empty_like(edge_env), + ) + + +def _k1_packed_direct_setup_context( + ctx: Any, + inputs: tuple[Any, ...], + output: tuple[Tensor, Tensor], +) -> None: + _, runner_token = output + ( + handle, + x, + d_full, + dt_full, + radial_feat, + edge_env, + src, + dst, + dst_ptr, + source_order, + source_ptr, + edge_src_gate, + ) = inputs + ctx.handle = int(handle) + ctx.save_for_backward( + x, + d_full, + dt_full, + radial_feat, + edge_env, + src, + dst, + dst_ptr, + source_order, + source_ptr, + edge_src_gate, + runner_token, + ) + + +def _k1_packed_direct_registered_backward_impl( + ctx: Any, + grad_out: Tensor, +) -> tuple[Any, ...]: + ( + x, + d_full, + dt_full, + radial_feat, + edge_env, + src, + dst, + dst_ptr, + source_order, + source_ptr, + edge_src_gate, + runner_token, + ) = ctx.saved_tensors + grad_x, grad_d, grad_dt, grad_radial, grad_edge_env = _k1_packed_direct_bwd_op( + ctx.handle, + grad_out, + x, + d_full, + dt_full, + radial_feat, + edge_env, + src, + dst, + dst_ptr, + source_order, + source_ptr, + edge_src_gate, + runner_token, + ) + return ( + None, + grad_x, + grad_d, + grad_dt, + grad_radial, + grad_edge_env, + None, + None, + None, + None, + None, + None, + ) + + +def _k1_packed_direct_backward( + ctx: Any, + grad_out: Tensor, + grad_runner_token: Tensor | None, +) -> tuple[Any, ...]: + """Run packed-direct backward with its custom op visible to compilation.""" + del grad_runner_token + return _k1_packed_direct_registered_backward_impl(ctx, grad_out) + + +_k1_packed_direct_op.register_autograd( + _k1_packed_direct_backward, + setup_context=_k1_packed_direct_setup_context, +) + + +def _cute_k1_impl( + handle: int, + x: Tensor, + d_full: Tensor, + dt_full: Tensor, + radial_feat: Tensor, + edge_env: Tensor, + src: Tensor, + dst: Tensor, + dst_ptr: Tensor, + source_order: Tensor, + source_ptr: Tensor, + edge_src_gate: Tensor, +) -> Tensor: + output, _runner_token = _k1_packed_direct_op( + int(handle), + x, + d_full, + dt_full, + radial_feat, + edge_env, + src, + dst, + dst_ptr, + source_order, + source_ptr, + edge_src_gate, + ) + return output + + +def cute_k1( + handle: int, + x: Tensor, + d_full: Tensor, + dt_full: Tensor, + radial_feat: Tensor, + edge_env: Tensor, + src: Tensor, + dst: Tensor, + dst_ptr: Tensor, + edge_src_gate: Tensor, + source_order: Tensor | None = None, + source_ptr: Tensor | None = None, +) -> Tensor: + """Run K1 through its registered custom-op boundary.""" + if source_order is None: + source_order = src.new_empty((0,), dtype=torch.int32) + if source_ptr is None: + source_ptr = src.new_empty((0,), dtype=torch.int32) + return _cute_k1_impl( + handle, + x, + d_full, + dt_full, + radial_feat, + edge_env, + src, + dst, + dst_ptr, + source_order, + source_ptr, + edge_src_gate, + ) + + +def _dst_ptr_from_sorted( + torch_module: Any, + dst: Tensor, + n_node: int, + *, + destinations_sorted: bool, +) -> Tensor | None: + if not destinations_sorted: + return None + if runtime_policy.is_cute_strict_enabled() and dst.numel() > 1: + torch_module._assert_async( + torch_module.all(dst[1:] >= dst[:-1]), + "Neo K1 destinations_sorted=True requires monotonically " + "nondecreasing destination indices", + ) + boundaries = torch_module.arange( + n_node + 1, + device=dst.device, + dtype=torch_module.int64, + ) + return torch_module.searchsorted(dst.contiguous(), boundaries) + + +def _validated_sorted_edge_metadata_args( + edge_cache: Any, + *, + node_count: int, + dst_ptr: Tensor | None, + source_order: Tensor | None, + source_ptr: Tensor | None, +) -> tuple[Tensor, Tensor, Tensor] | None: + """Return validated invocation-local CSR tensors for this edge cache.""" + if not getattr(edge_cache, "destinations_sorted", False): + return None + if dst_ptr is None or source_order is None or source_ptr is None: + return None + device = edge_cache.src.device + if ( + dst_ptr.device != device + or source_order.device != device + or source_ptr.device != device + or dst_ptr.dtype != torch.int32 + or source_order.dtype != torch.int32 + or source_ptr.dtype != torch.int32 + or dst_ptr.numel() != node_count + 1 + or source_order.numel() != edge_cache.src.numel() + or source_ptr.numel() != node_count + 1 + ): + return None + return ( + dst_ptr.contiguous(), + source_order.contiguous(), + source_ptr.contiguous(), + ) + + +def _cuda_compute_capability(device_index: int) -> tuple[int, int]: + return tuple(torch.cuda.get_device_capability(device_index)) + + +def _architecture_default_config( + compute_capability: tuple[int, int], +) -> NeoK1RuntimeConfig: + return NeoK1RuntimeConfig( + native_sm90_path=compute_capability == runtime_policy.SM90_CAPABILITY, + per_focus_so2_fwd_pair=( + compute_capability in runtime_policy.SM80_PROFILE_CAPABILITIES + ), + combined_so2_gate=( + compute_capability in runtime_policy.FUSED_SO2_GATE_CAPABILITIES + ), + ) + + +def _maybe_run_prepared_cute_k1( + block: Any, + x: Tensor, + edge_cache: Any, + radial_feat: Tensor, + dst_ptr: Tensor | None = None, + source_order: Tensor | None = None, + source_ptr: Tensor | None = None, +) -> Tensor | None: + """Dispatch prevalidated packed K1 state without a Python graph break.""" + 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 + d_full = edge_cache.D_packed + dt_full = d_full + destinations_sorted = bool(getattr(edge_cache, "destinations_sorted", False)) + device_index = x.device.index + if ( + block.training + or x.device.type != "cuda" + or device_index is None + or device_index != state.device_index + or x.dtype != torch.float32 + or torch.is_autocast_enabled(x.device.type) + or d_full is None + or dt_full is None + or d_full is not dt_full + or d_full.dim() != 2 + or dt_full.dim() != 2 + or edge_cache.edge_src_gate is not None + or not destinations_sorted + or not _dtypes_use_strict_fp32( + ( + d_full.dtype, + dt_full.dtype, + radial_feat.dtype, + edge_cache.edge_env.dtype, + ) + ) + ): + return None + edge_count = edge_cache.src.numel() + if not runtime_policy.k1_int32_indexing_is_safe( + edge_count=edge_count, + node_count=x.shape[0], + ): + return None + metadata_args = _validated_sorted_edge_metadata_args( + edge_cache, + node_count=x.shape[0], + dst_ptr=dst_ptr, + source_order=source_order, + source_ptr=source_ptr, + ) + if metadata_args is None: + dst_ptr = _dst_ptr_from_sorted( + torch, + edge_cache.dst, + x.shape[0], + destinations_sorted=destinations_sorted, + ) + if dst_ptr is None: + return None + source_order = edge_cache.src.new_empty((0,), dtype=torch.int32) + source_ptr = edge_cache.src.new_empty((0,), dtype=torch.int32) + else: + dst_ptr, source_order, source_ptr = metadata_args + # The opaque custom-op implementation performs exact pointer-alignment + # canonicalization in ``_build_runner``. Keep this wrapper graph-visible. + empty_edge_src_gate = edge_cache.edge_env.new_empty((0,)) + d_arg = d_full.contiguous() + output = cute_k1( + state.handle, + x.contiguous(), + d_arg, + d_arg, + radial_feat.contiguous(), + edge_cache.edge_env.contiguous(), + edge_cache.src.contiguous(), + edge_cache.dst.contiguous(), + dst_ptr.contiguous(), + empty_edge_src_gate, + source_order=source_order.contiguous(), + source_ptr=source_ptr.contiguous(), + ) + return _layout_like(output, x) + + +@torch.compiler.disable +def _maybe_run_cute_k1_fallback( + block: Any, + x: Tensor, + edge_cache: Any, + radial_feat: Tensor, + dst_ptr: Tensor | None = None, + source_order: Tensor | None = None, + source_ptr: Tensor | None = None, +) -> Tensor | None: + """Run the opt-in Neo CuTe K1 path, or return ``None`` for fallback.""" + if edge_cache.D_packed is None: + return None + # The optimized backward does not expose the differentiable SFPG + # source-gate adjoint. Preserve force/stress correctness via eager fallback. + if edge_cache.edge_src_gate is not None: + return None + destinations_sorted = bool(getattr(edge_cache, "destinations_sorted", False)) + if not is_neo_k1_runtime_eligible( + block, + training=bool(block.training), + device=x.device, + dtype=x.dtype, + edge_count=edge_cache.src.numel(), + node_count=x.shape[0], + destinations_sorted=destinations_sorted, + ): + return None + if not _dtypes_use_strict_fp32( + ( + edge_cache.D_packed.dtype, + radial_feat.dtype, + edge_cache.edge_env.dtype, + ) + ): + return None + + state = getattr(block, "_deepmd_cute_k1_state", None) + if state is False: + return None + if state is not None and not isinstance(state, _RegisteredK1State): + state = None + device_index = x.device.index + if device_index is None: + device_index = torch.cuda.current_device() + compute_capability = _cuda_compute_capability(device_index) + if not is_supported_k1_compute_capability(compute_capability): + return None + config = _architecture_default_config(compute_capability) + if not runtime_policy.k1_int32_indexing_is_safe( + edge_count=edge_cache.src.numel(), + node_count=x.shape[0], + ): + return None + if edge_cache.D_packed.dim() != 2: + return None + if state is None or state.device_index != device_index or state.config != config: + state = _register_cute_k1_state(block, device_index, config) + if state is None: + return None + + metadata_args = _validated_sorted_edge_metadata_args( + edge_cache, + node_count=x.shape[0], + dst_ptr=dst_ptr, + source_order=source_order, + source_ptr=source_ptr, + ) + if metadata_args is None: + dst_ptr = _dst_ptr_from_sorted( + torch, + edge_cache.dst, + x.shape[0], + destinations_sorted=destinations_sorted, + ) + if dst_ptr is None: + return None + source_order = edge_cache.src.new_empty((0,), dtype=torch.int32) + source_ptr = edge_cache.src.new_empty((0,), dtype=torch.int32) + else: + dst_ptr, source_order, source_ptr = metadata_args + x_arg = _aligned_contiguous(x) + d_arg = _aligned_contiguous(edge_cache.D_packed) + dt_arg = d_arg + radial_arg = _aligned_contiguous(radial_feat) + edge_env_arg = _aligned_contiguous(edge_cache.edge_env) + src_arg = _aligned_contiguous(edge_cache.src) + dst_arg = _aligned_contiguous(edge_cache.dst) + dst_ptr_arg = _aligned_contiguous(dst_ptr) + source_order_arg = _aligned_contiguous(source_order) + source_ptr_arg = _aligned_contiguous(source_ptr) + edge_src_gate = edge_cache.edge_src_gate + if edge_src_gate is None: + edge_src_gate = edge_cache.edge_env.new_empty((0,)) + edge_src_gate_arg = _aligned_contiguous(edge_src_gate) + output = cute_k1( + state.handle, + x_arg, + d_arg, + dt_arg, + radial_arg, + edge_env_arg, + src_arg, + dst_arg, + dst_ptr_arg, + edge_src_gate_arg, + source_order=source_order_arg, + source_ptr=source_ptr_arg, + ) + return _layout_like(output, x) + + +def maybe_run_cute_k1( + block: Any, + x: Tensor, + edge_cache: Any, + radial_feat: Tensor, + dst_ptr: Tensor | None = None, + source_order: Tensor | None = None, + source_ptr: Tensor | None = None, +) -> Tensor | None: + """Use prevalidated opaque dispatch, or the conservative eager fallback.""" + state = getattr(block, "_deepmd_cute_k1_state", None) + use_prepared = isinstance( + state, + _RegisteredK1State, + ) or runtime_policy.is_k1_thin_wrapper_enabled(_tensor_compute_capability(x)) + if use_prepared: + output = _maybe_run_prepared_cute_k1( + block, + x, + edge_cache, + radial_feat, + dst_ptr, + source_order, + source_ptr, + ) + if output is not None: + return output + return _maybe_run_cute_k1_fallback( + block, + x, + edge_cache, + radial_feat, + dst_ptr, + source_order, + source_ptr, + ) diff --git a/deepmd/kernels/cute/neo/k1_gate_structural.py b/deepmd/kernels/cute/neo/k1_gate_structural.py new file mode 100644 index 0000000000..caafbf17ff --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_gate_structural.py @@ -0,0 +1,175 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""PyTorch/cuBLAS glue for the structural Neo K1 split-gate path. + +The two focus panels stay contiguous across the gate-linear boundary. Forward +uses direct cuBLAS matrix products, while backward accumulates into the scalar +slice of ``grad_y`` through the GEMM beta epilogue. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import torch +from torch import ( + Tensor, +) + +FOCUS_COUNT = 2 +CHANNELS = 32 +GATE_WIDTH = 3 * CHANNELS +VEC4_ALIGNMENT_BYTES = 16 +VEC4_STORAGE_OFFSET_MULTIPLE = 4 + + +def _dispatch_aligned_vec4_kernel( + kernel: Any, + tensor_names: tuple[str, ...], + *tensors: Tensor, +) -> Any: + """Validate the float4 load/store contract, then dispatch the kernel.""" + if len(tensor_names) != len(tensors): + raise ValueError("vec4 dispatch tensor names and arguments must match") + if not tensors or tensors[0].numel() == 0: + return None + + for name, tensor in zip(tensor_names, tensors, strict=True): + if tensor.dtype != torch.float32: + raise TypeError( + "SM80 vec4 structural gate requires float32 tensors; " + f"{name} has dtype={tensor.dtype}" + ) + if not tensor.is_contiguous(): + raise ValueError( + "SM80 vec4 structural gate requires compact tensors; " + f"{name} has shape={tuple(tensor.shape)} and stride={tensor.stride()}" + ) + if tensor.storage_offset() % VEC4_STORAGE_OFFSET_MULTIPLE: + raise ValueError( + "SM80 vec4 structural gate requires storage offsets divisible " + f"by {VEC4_STORAGE_OFFSET_MULTIPLE} float32 elements; " + f"{name} has storage_offset={tensor.storage_offset()}" + ) + pointer_remainder = tensor.data_ptr() % VEC4_ALIGNMENT_BYTES + if pointer_remainder: + raise ValueError( + "SM80 vec4 structural gate requires 16-byte-aligned tensors; " + f"{name} has data_ptr modulo 16={pointer_remainder}" + ) + + return kernel(*tensors) + + +def focus_major_gate_linear_forward( + gate_src: Tensor, + gate_weight: Tensor, +) -> Tensor: + """Project ``(E, 2, 32)`` into contiguous ``(2, E, 96)`` panels.""" + edge_count = gate_src.shape[0] + weight = gate_weight.view(CHANNELS, FOCUS_COUNT, GATE_WIDTH) + logits = torch.empty( + FOCUS_COUNT, + edge_count, + GATE_WIDTH, + dtype=gate_src.dtype, + device=gate_src.device, + ) + for focus in range(FOCUS_COUNT): + torch.mm(gate_src[:, focus, :], weight[:, focus, :], out=logits[focus]) + return logits + + +def focus_major_gate_linear_backward_add_( + grad_y: Tensor, + grad_logits: Tensor, + gate_weight: Tensor, +) -> Tensor: + """Accumulate the gate-linear adjoint into ``grad_y`` in place.""" + weight = gate_weight.view(CHANNELS, FOCUS_COUNT, GATE_WIDTH) + for focus in range(FOCUS_COUNT): + grad_y[:, focus, 0, :].addmm_( + grad_logits[focus], + weight[:, focus, :].T, + ) + return grad_y + + +def focus_major_so2_backward_with_folded_residual_out( + grad_out: Tensor, + w0_folded_t: Tensor, + wpair_folded_t: Tensor, + *, + out: Tensor, +) -> Tensor: + """Write ``grad_out @ (W.T + I)`` without seeding output copies.""" + if out.shape != grad_out.shape: + raise ValueError("SO2 backward tensors must have identical shapes") + if not out.is_contiguous(): + raise ValueError("SO2 backward output storage must be contiguous") + if out.untyped_storage()._cdata == grad_out.untyped_storage()._cdata: + raise ValueError("SO2 backward output must not alias its input") + edge_count = grad_out.shape[0] + grad_flat = grad_out.view(edge_count, FOCUS_COUNT, 10 * CHANNELS) + out_flat = out.view(edge_count, FOCUS_COUNT, 10 * CHANNELS) + split = 4 * CHANNELS + for focus in range(FOCUS_COUNT): + torch.mm( + grad_flat[:, focus, :split], + w0_folded_t[focus], + out=out_flat[:, focus, :split], + ) + torch.mm( + grad_flat[:, focus, split:], + wpair_folded_t[focus], + out=out_flat[:, focus, split:], + ) + return out + + +def run_structural_gate_forward( + kernel: Any, + residual: Tensor, + y: Tensor, + logits: Tensor, + *, + out: Tensor, +) -> Tensor: + """Run the alias-safe CuTe forward into caller-owned storage.""" + rows = residual.shape[0] * residual.shape[1] + kernel( + residual.view(rows, 10 * CHANNELS), + y.view(rows, 10 * CHANNELS), + logits, + out.view(rows, 10 * CHANNELS), + ) + return out + + +def run_structural_gate_backward( + kernel: Any, + grad_out: Tensor, + y: Tensor, + logits: Tensor, + grad_y: Tensor, + *, + grad_logits: Tensor | None, + overwrite_logits: bool, +) -> Tensor: + """Run gate backward, optionally replacing consumed logits with adjoints.""" + grad_logits_out = logits if overwrite_logits else grad_logits + if grad_logits_out is None: + raise ValueError("grad_logits storage is required when logits are preserved") + rows = y.shape[0] * y.shape[1] + kernel( + grad_out.view(rows, 10 * CHANNELS), + y.view(rows, 10 * CHANNELS), + logits, + grad_y.view(rows, 10 * CHANNELS), + grad_logits_out, + ) + return grad_logits_out diff --git a/deepmd/kernels/cute/neo/k1_kernels/__init__.py b/deepmd/kernels/cute/neo/k1_kernels/__init__.py new file mode 100644 index 0000000000..77198c0277 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Private CuTe kernels for the Neo SO2Conv/K1 inference path.""" diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_envelope_gated_softmax.py b/deepmd/kernels/cute/neo/k1_kernels/cute_envelope_gated_softmax.py new file mode 100644 index 0000000000..8a6512689b --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_envelope_gated_softmax.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""CuTe DSL kernel for envelope-gated segmented softmax. + +The input contract is intentionally strict: destination edges must already be +sorted and represented by CSR row pointers before calling the CuTe kernel. +""" + +# ruff: noqa: ANN001, ANN201, ANN204, TC002 + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, +) + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from ..compile_cache import ( + device_aware_lru_cache, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + +@dataclass(frozen=True) +class EnvelopeSoftmaxFwdParams: + threads: int + logits: cute.Tensor + edge_gate: cute.Tensor + dst_ptr: cute.Tensor + z_bias_raw: cute.Tensor + out: cute.Tensor + group_max: cute.Tensor + denom: cute.Tensor + + +class EnvelopeSoftmaxForward: + def __init__(self, threads: int): + if threads % 32 != 0: + raise ValueError("threads must be a multiple of 32") + self.threads = threads + self.warps = threads // 32 + self.dtype = cutlass.Float32 + + @cute.jit + def warp_sum(self, value): + return cute.arch.warp_reduction_sum(value) + + @cute.jit + def warp_max(self, value): + return cute.arch.warp_reduction_max(value) + + @cute.jit + def cta_sum(self, value, scratch, tidx): + lane = tidx % 32 + warp = tidx // 32 + value = self.warp_sum(value) + if lane == 0: + scratch[warp] = value + cute.arch.sync_threads() + + total = self.dtype(0.0) + if tidx < self.warps: + total = scratch[tidx] + total = self.warp_sum(total) + if tidx == 0: + scratch[0] = total + cute.arch.sync_threads() + return scratch[0] + + @cute.jit + def cta_max(self, value, scratch, tidx): + lane = tidx % 32 + warp = tidx // 32 + value = self.warp_max(value) + if lane == 0: + scratch[warp] = value + cute.arch.sync_threads() + + neg_large = self.dtype(-3.4028234663852886e38) + total = neg_large + if tidx < self.warps: + total = scratch[tidx] + total = self.warp_max(total) + if tidx == 0: + scratch[0] = total + cute.arch.sync_threads() + return scratch[0] + + @cute.jit + def softplus(self, value): + zero = self.dtype(0.0) + positive = cute.arch.fmax(value, zero) + magnitude = cute.arch.fmax(value, -value) + return positive + cute.log(self.dtype(1.0) + cute.exp(-magnitude)) + + @cute.kernel + def kernel(self, params: EnvelopeSoftmaxFwdParams, eps: cutlass.Constexpr[float]): + tidx, _, _ = cute.arch.thread_idx() + node, group, _ = cute.arch.block_idx() + + smem = cutlass.utils.SmemAllocator() + scratch = smem.allocate_tensor(self.dtype, self.warps) + + lo = params.dst_ptr[node] + hi = params.dst_ptr[node + 1] + null_mass = self.softplus(params.z_bias_raw[group].to(self.dtype)) + self.dtype( + eps + ) + local_max = cute.log(null_mass) + for edge in cutlass.range(lo + tidx, hi, self.threads, unroll=1): + gate = params.edge_gate[edge].to(self.dtype) + if gate < self.dtype(0.0): + gate = self.dtype(0.0) + if gate > self.dtype(0.0): + value = params.logits[edge, group].to(self.dtype) + self.dtype( + 2.0 + ) * cute.log(gate) + if value > local_max: + local_max = value + + group_max = self.cta_max(local_max, scratch, tidx) + # Every warp must consume scratch[0] before cta_sum reuses it. + cute.arch.sync_threads() + + local_sum = self.dtype(0.0) + for edge in cutlass.range(lo + tidx, hi, self.threads, unroll=1): + gate = params.edge_gate[edge].to(self.dtype) + if gate < self.dtype(0.0): + gate = self.dtype(0.0) + if gate > self.dtype(0.0): + effective_logit = params.logits[edge, group].to( + self.dtype + ) + self.dtype(2.0) * cute.log(gate) + local_sum += cute.exp(effective_logit - group_max) + + denom_sum = self.cta_sum(local_sum, scratch, tidx) + denom = denom_sum + null_mass * cute.exp(-group_max) + + if tidx == 0: + params.group_max[node, group] = group_max.to(params.group_max.element_type) + params.denom[node, group] = denom.to(params.denom.element_type) + cute.arch.sync_threads() + + for edge in cutlass.range(lo + tidx, hi, self.threads, unroll=1): + gate = params.edge_gate[edge].to(self.dtype) + if gate < self.dtype(0.0): + gate = self.dtype(0.0) + alpha = self.dtype(0.0) + if gate > self.dtype(0.0): + effective_logit = params.logits[edge, group].to( + self.dtype + ) + self.dtype(2.0) * cute.log(gate) + num = cute.exp(effective_logit - group_max) + alpha = num / denom + params.out[edge, group] = alpha.to(params.out.element_type) + + +@cute.jit +def envelope_softmax_forward_jit( + logits: cute.Tensor, + edge_gate: cute.Tensor, + dst_ptr: cute.Tensor, + z_bias_raw: cute.Tensor, + out: cute.Tensor, + group_max: cute.Tensor, + denom: cute.Tensor, + threads: cutlass.Constexpr[int], + eps: cutlass.Constexpr[float], + stream: CUstream, +): + params = EnvelopeSoftmaxFwdParams( + threads=threads, + logits=logits, + edge_gate=edge_gate, + dst_ptr=dst_ptr, + z_bias_raw=z_bias_raw, + out=out, + group_max=group_max, + denom=denom, + ) + n_nodes, groups = denom.shape + EnvelopeSoftmaxForward(threads).kernel(params, eps).launch( + grid=[n_nodes, groups, 1], + block=[threads, 1, 1], + stream=stream, + ) + + +@device_aware_lru_cache(maxsize=16) +def compile_envelope_softmax_forward(threads: int, eps: float = 1.0e-7) -> Callable: + e = cute.sym_int64() + n = cute.sym_int64() + g = cute.sym_int64() + fake_logits = make_fake_compact_tensor(cutlass.Float32, (e, g), stride_order=(1, 0)) + fake_gate = make_fake_compact_tensor(cutlass.Float32, (e,), stride_order=(0,)) + fake_dst_ptr = make_fake_compact_tensor( + cutlass.Int32, (cute.sym_int64(),), stride_order=(0,) + ) + fake_z = make_fake_compact_tensor(cutlass.Float32, (g,), stride_order=(0,)) + fake_out = make_fake_compact_tensor(cutlass.Float32, (e, g), stride_order=(1, 0)) + fake_group_max = make_fake_compact_tensor( + cutlass.Float32, (n, g), stride_order=(1, 0) + ) + fake_denom = make_fake_compact_tensor(cutlass.Float32, (n, g), stride_order=(1, 0)) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + envelope_softmax_forward_jit, + fake_logits, + fake_gate, + fake_dst_ptr, + fake_z, + fake_out, + fake_group_max, + fake_denom, + threads, + eps, + fake_stream, + options="--enable-tvm-ffi", + ) diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_focus_src_backward.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_focus_src_backward.py new file mode 100644 index 0000000000..8aee494844 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_focus_src_backward.py @@ -0,0 +1,265 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +# ruff: noqa: ANN001, ANN201, ANN202, TC002, UP035 +"""CuTe forward for Neo attention-prelude source features. + +The forward kernel fuses two independent PyTorch producer chains in one launch: + + focus RMSNorm -> two-focus logits -> softmax -> label smoothing + scalar Q/K RMSNorm -> Q projection + K projection + +The kernel is specialized to the Neo K1 shape `(E, F=2, C=32)`. +""" + +from __future__ import ( + annotations, +) + +from functools import ( + lru_cache, +) +from typing import ( + Callable, +) + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + + +@cute.jit +def _warp_sum(value): + return cute.arch.warp_reduction_sum(value) + + +@cute.jit +def neo_attention_prelude_forward_jit( + focus: cute.Tensor, + x_l0: cute.Tensor, + focus_weight: cute.Tensor, + focus_scale: cute.Tensor, + q_weight: cute.Tensor, + k_weight: cute.Tensor, + qk_scale: cute.Tensor, + focus_alpha: cute.Tensor, + q_node: cute.Tensor, + k_node: cute.Tensor, + stream: CUstream, + focus_eps: cutlass.Float32, + qk_eps: cutlass.Float32, + tau: cutlass.Float32, + label_smoothing: cutlass.Float32, +): + edges, _ = focus.shape + nodes, _, _ = x_l0.shape + neo_attention_prelude_forward_kernel( + focus, + x_l0, + focus_weight, + focus_scale, + q_weight, + k_weight, + qk_scale, + focus_alpha, + q_node, + k_node, + focus_eps, + qk_eps, + tau, + label_smoothing, + ).launch( + # The launch covers two independent domains. Using their sum avoids a + # host-side branch on symbolic E/N while guaranteeing coverage of both. + grid=[cute.ceil_div(edges + nodes, 8), 1, 1], + block=[256, 1, 1], + stream=stream, + ) + + +@cute.kernel +def neo_attention_prelude_forward_kernel( + focus: cute.Tensor, + x_l0: cute.Tensor, + focus_weight: cute.Tensor, + focus_scale: cute.Tensor, + q_weight: cute.Tensor, + k_weight: cute.Tensor, + qk_scale: cute.Tensor, + focus_alpha: cute.Tensor, + q_node: cute.Tensor, + k_node: cute.Tensor, + focus_eps: cutlass.Float32, + qk_eps: cutlass.Float32, + tau: cutlass.Float32, + label_smoothing: cutlass.Float32, +): + tid, _, _ = cute.arch.thread_idx() + block, _, _ = cute.arch.block_idx() + lane = tid % 32 + edge = block * 8 + tid // 32 + edges, _ = focus.shape + + if edge < edges: + x0 = focus[edge, lane].to(cutlass.Float32) + x1 = focus[edge, 32 + lane].to(cutlass.Float32) + inv0 = cute.rsqrt(_warp_sum(x0 * x0) / cutlass.Float32(32.0) + focus_eps) + inv1 = cute.rsqrt(_warp_sum(x1 * x1) / cutlass.Float32(32.0) + focus_eps) + norm0 = x0 * inv0 * focus_scale[0, lane].to(cutlass.Float32) + norm1 = x1 * inv1 * focus_scale[1, lane].to(cutlass.Float32) + logit0 = _warp_sum(norm0 * focus_weight[lane, 0].to(cutlass.Float32)) + logit1 = _warp_sum(norm1 * focus_weight[lane, 1].to(cutlass.Float32)) + + if lane == 0: + z0 = logit0 / tau + z1 = logit1 / tau + zmax = z0 + if z1 > zmax: + zmax = z1 + e0 = cute.exp(z0 - zmax) + e1 = cute.exp(z1 - zmax) + denom = e0 + e1 + keep = cutlass.Float32(1.0) - label_smoothing + smooth = label_smoothing / cutlass.Float32(2.0) + focus_alpha[edge, 0] = (e0 / denom * keep + smooth).to( + focus_alpha.element_type + ) + focus_alpha[edge, 1] = (e1 / denom * keep + smooth).to( + focus_alpha.element_type + ) + + # Q/K is a per-node chain, not a child of the per-edge focus chain. Keep + # this guard independent so every node row is initialized even when E < N. + nodes, _, _ = x_l0.shape + node = block * 8 + tid // 32 + if node < nodes: + qk_x0 = x_l0[node, 0, lane].to(cutlass.Float32) + qk_x1 = x_l0[node, 1, lane].to(cutlass.Float32) + qk_norm0 = ( + qk_x0 + * cute.rsqrt(_warp_sum(qk_x0 * qk_x0) / cutlass.Float32(32.0) + qk_eps) + * qk_scale[0, lane].to(cutlass.Float32) + ) + qk_norm1 = ( + qk_x1 + * cute.rsqrt(_warp_sum(qk_x1 * qk_x1) / cutlass.Float32(32.0) + qk_eps) + * qk_scale[1, lane].to(cutlass.Float32) + ) + q0 = cutlass.Float32(0.0) + k0 = cutlass.Float32(0.0) + q1 = cutlass.Float32(0.0) + k1 = cutlass.Float32(0.0) + for input_channel in cutlass.range_constexpr(32): + value0 = cute.arch.shuffle_sync(qk_norm0, input_channel) + value1 = cute.arch.shuffle_sync(qk_norm1, input_channel) + q0 += value0 * q_weight[input_channel, 0, lane].to(cutlass.Float32) + k0 += value0 * k_weight[input_channel, 0, lane].to(cutlass.Float32) + q1 += value1 * q_weight[input_channel, 1, lane].to(cutlass.Float32) + k1 += value1 * k_weight[input_channel, 1, lane].to(cutlass.Float32) + q_node[node, 0, lane] = q0.to(q_node.element_type) + k_node[node, 0, lane] = k0.to(k_node.element_type) + q_node[node, 1, lane] = q1.to(q_node.element_type) + k_node[node, 1, lane] = k1.to(k_node.element_type) + + +@lru_cache(maxsize=8) +def compile_neo_attention_prelude_forward( + focus_eps: float, + qk_eps: float, + tau: float, + label_smoothing: float, + compile_identity: tuple[int, int, int] | None = None, +) -> Callable: + # The identity keeps independently compiled device/architecture binaries in + # distinct cache entries. Compilation itself runs under the runner's device. + del compile_identity + edges = cute.sym_int64() + nodes = cute.sym_int64() + fake_focus = make_fake_compact_tensor( + cutlass.Float32, (edges, 64), stride_order=(1, 0), **FAKE_TENSOR_KW + ) + fake_x_l0 = make_fake_compact_tensor( + cutlass.Float32, (nodes, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + fake_focus_weight = make_fake_compact_tensor( + cutlass.Float32, (32, 2), stride_order=(1, 0), **FAKE_TENSOR_KW + ) + fake_focus_scale = make_fake_compact_tensor( + cutlass.Float32, (2, 32), stride_order=(1, 0), **FAKE_TENSOR_KW + ) + fake_q_weight = make_fake_compact_tensor( + cutlass.Float32, (32, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + fake_k_weight = make_fake_compact_tensor( + cutlass.Float32, (32, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + fake_qk_scale = make_fake_compact_tensor( + cutlass.Float32, (2, 32), stride_order=(1, 0), **FAKE_TENSOR_KW + ) + fake_focus_alpha = make_fake_compact_tensor( + cutlass.Float32, (edges, 2), stride_order=(1, 0), **FAKE_TENSOR_KW + ) + fake_q_node = make_fake_compact_tensor( + cutlass.Float32, (nodes, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + fake_k_node = make_fake_compact_tensor( + cutlass.Float32, (nodes, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True) + compiled = cute.compile( + neo_attention_prelude_forward_jit, + fake_focus, + fake_x_l0, + fake_focus_weight, + fake_focus_scale, + fake_q_weight, + fake_k_weight, + fake_qk_scale, + fake_focus_alpha, + fake_q_node, + fake_k_node, + fake_stream, + cutlass.Float32(focus_eps), + cutlass.Float32(qk_eps), + cutlass.Float32(tau), + cutlass.Float32(label_smoothing), + options="--enable-tvm-ffi", + ) + + def run( + focus, + x_l0, + focus_weight, + focus_scale, + q_weight, + k_weight, + qk_scale, + focus_alpha, + q_node, + k_node, + ): + return compiled( + focus, + x_l0, + focus_weight, + focus_scale, + q_weight, + k_weight, + qk_scale, + focus_alpha, + q_node, + k_node, + cutlass.Float32(focus_eps), + cutlass.Float32(qk_eps), + cutlass.Float32(tau), + cutlass.Float32(label_smoothing), + ) + + return run diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py new file mode 100644 index 0000000000..6c8658b887 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""CuTe Neo gate-linear + gate/residual backward without saved logits.""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, +) + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from ..compile_cache import ( + device_aware_lru_cache, +) + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN001, ANN201, ANN202, TC002 + + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} +ROWS_PER_BLOCK = 2 + + +@cute.jit +def _sigmoid(value): + return cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-value)) + + +@cute.jit +def neo_gate_linear_residual_backward_fused_jit( + grad_out: cute.Tensor, + y: cute.Tensor, + gate_weight: cute.Tensor, + grad_y: cute.Tensor, + stream: CUstream, + rows_per_block: cutlass.Constexpr[int], +): + rows, _ = y.shape + neo_gate_linear_residual_backward_fused_kernel( + grad_out, + y, + gate_weight, + grad_y, + rows_per_block, + ).launch( + grid=[cute.ceil_div(rows, rows_per_block), 1, 1], + block=[32 * rows_per_block, 1, 1], + stream=stream, + ) + + +@cute.kernel +def neo_gate_linear_residual_backward_fused_kernel( + grad_out: cute.Tensor, + y: cute.Tensor, + gate_weight: cute.Tensor, + grad_y: cute.Tensor, + rows_per_block: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + block_row, _, _ = cute.arch.block_idx() + row_slot = tidx // 32 + channel = tidx - row_slot * 32 + row = block_row * rows_per_block + row_slot + rows, _ = y.shape + + smem = cutlass.utils.SmemAllocator() + grad_logits = smem.allocate_tensor(cutlass.Float32, rows_per_block * 3 * 32) + smem_base = row_slot * 3 * 32 + + for gate_degree in cutlass.range_constexpr(3): + grad_logits[smem_base + gate_degree * 32 + channel] = cutlass.Float32(0.0) + cute.arch.sync_threads() + + grad_l0 = cutlass.Float32(0.0) + if row < rows: + focus = row - (row // 2) * 2 + y0 = y[row, channel].to(cutlass.Float32) + sig0 = _sigmoid(y0) + grad0 = grad_out[row, channel].to(cutlass.Float32) + grad_l0 = ( + grad0 * sig0 * (cutlass.Float32(1.0) + y0 * (cutlass.Float32(1.0) - sig0)) + ) + + gate0_logit = cutlass.Float32(0.0) + gate1_logit = cutlass.Float32(0.0) + gate2_logit = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(32): + src = y[row, k].to(cutlass.Float32) + gate0_logit += src * gate_weight[k, focus, channel].to(cutlass.Float32) + gate1_logit += src * gate_weight[k, focus, 32 + channel].to(cutlass.Float32) + gate2_logit += src * gate_weight[k, focus, 64 + channel].to(cutlass.Float32) + + gate0 = _sigmoid(gate0_logit) + gate1 = _sigmoid(gate1_logit) + gate2 = _sigmoid(gate2_logit) + for d in cutlass.range_constexpr(1, 10, 1): + gate_idx = channel + gate = gate0 + if cutlass.const_expr((d - 1) % 3 == 1): + gate_idx = 32 + channel + gate = gate1 + if cutlass.const_expr((d - 1) % 3 == 2): + gate_idx = 64 + channel + gate = gate2 + idx = d * 32 + channel + gout = grad_out[row, idx].to(cutlass.Float32) + yv = y[row, idx].to(cutlass.Float32) + grad_y[row, idx] = (gout * gate).to(grad_y.element_type) + old = grad_logits[smem_base + gate_idx] + grad_logits[smem_base + gate_idx] = old + gout * yv * gate * ( + cutlass.Float32(1.0) - gate + ) + cute.arch.sync_threads() + + if row < rows: + focus = row - (row // 2) * 2 + grad_gate_src = cutlass.Float32(0.0) + for out_idx in cutlass.range_constexpr(3 * 32): + grad_gate_src += grad_logits[smem_base + out_idx] * gate_weight[ + channel, focus, out_idx + ].to(cutlass.Float32) + grad_y[row, channel] = (grad_l0 + grad_gate_src).to(grad_y.element_type) + + +@device_aware_lru_cache(maxsize=2) +def compile_neo_gate_linear_residual_backward_fused() -> Callable: + rows = cute.sym_int64() + fake_grad_out = make_fake_compact_tensor( + cutlass.Float32, (rows, 10 * 32), stride_order=(1, 0), **FAKE_TENSOR_KW + ) + fake_y = make_fake_compact_tensor( + cutlass.Float32, (rows, 10 * 32), stride_order=(1, 0), **FAKE_TENSOR_KW + ) + fake_gate_weight = make_fake_compact_tensor( + cutlass.Float32, (32, 2, 3 * 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + fake_grad_y = make_fake_compact_tensor( + cutlass.Float32, (rows, 10 * 32), stride_order=(1, 0), **FAKE_TENSOR_KW + ) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + neo_gate_linear_residual_backward_fused_jit, + fake_grad_out, + fake_y, + fake_gate_weight, + fake_grad_y, + fake_stream, + ROWS_PER_BLOCK, + options="--enable-tvm-ffi", + ) diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py new file mode 100644 index 0000000000..c0a4fc058e --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py @@ -0,0 +1,497 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +# ruff: noqa: ANN001, ANN201, ANN202, TC002, TC003 +"""Strict-FP32 vectorized structural-gate forward/backward for Neo K1. + +The split path keeps both neighboring SO2 products and the two gate projections +in cuBLAS. Its elementwise consumer assigns eight threads to each 32-channel +row, with each thread moving an aligned float4. Both directions preserve the +split-gate tensor contract and leave dense projections in PyTorch/cuBLAS. +""" + +from __future__ import ( + annotations, +) + +from collections.abc import ( + Callable, +) +from functools import ( + lru_cache, +) + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from .. import ( + runtime_policy, +) + +FOCUS_COUNT = 2 +REDUCED_COUNT = 10 +CHANNELS = 32 +GATE_COUNT = 3 +VECTOR_WIDTH = 4 +CHANNEL_GROUPS = CHANNELS // VECTOR_WIDTH +ROWS_PER_BLOCK = 16 +THREADS = ROWS_PER_BLOCK * CHANNEL_GROUPS +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} +FORWARD_TENSOR_NAMES = ("residual", "y", "logits", "out") +BACKWARD_TENSOR_NAMES = ("grad_out", "y", "logits", "grad_y", "grad_logits") + + +def _guard_vec4_dispatch( + kernel: Callable, + tensor_names: tuple[str, ...], +) -> Callable: + from ..k1_gate_structural import ( + _dispatch_aligned_vec4_kernel, + ) + + def dispatch(*tensors: object): + return _dispatch_aligned_vec4_kernel(kernel, tensor_names, *tensors) + + return dispatch + + +@cute.jit +def _sigmoid(value): + return cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-value)) + + +@cute.jit +def neo_gate_split_structural_vec4_sm80_forward_jit( + residual: cute.Tensor, + y: cute.Tensor, + logits: cute.Tensor, + out: cute.Tensor, + stream: CUstream, +): + copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + residual.element_type, + num_bits_per_copy=residual.element_type.width * VECTOR_WIDTH, + ) + channel_thread_layout = cute.make_ordered_layout( + (1, CHANNEL_GROUPS), + order=(1, 0), + ) + channel_value_layout = cute.make_ordered_layout( + (1, VECTOR_WIDTH), + order=(1, 0), + ) + vector_copy = cute.make_tiled_copy_tv( + copy_atom, + channel_thread_layout, + channel_value_layout, + ) + channel_layout = cute.make_layout((1, CHANNELS), stride=(CHANNELS, 1)) + rows, _ = y.shape + neo_gate_split_structural_vec4_sm80_forward_kernel( + residual, + y, + logits, + out, + channel_layout, + vector_copy, + ).launch( + grid=[cute.ceil_div(rows, ROWS_PER_BLOCK), 1, 1], + block=[THREADS, 1, 1], + stream=stream, + ) + + +@cute.kernel +def neo_gate_split_structural_vec4_sm80_forward_kernel( + residual: cute.Tensor, + y: cute.Tensor, + logits: cute.Tensor, + out: cute.Tensor, + channel_layout: cute.Layout, + vector_copy: cute.TiledCopy, +): + tidx, _, _ = cute.arch.thread_idx() + block_row, _, _ = cute.arch.block_idx() + row_slot = tidx // CHANNEL_GROUPS + channel_group = tidx - row_slot * CHANNEL_GROUPS + row = block_row * ROWS_PER_BLOCK + row_slot + rows, _ = y.shape + + if row < rows: + edge = row // FOCUS_COUNT + focus = row - edge * FOCUS_COUNT + thread_copy = vector_copy.get_slice(channel_group) + + y0_tile = cute.local_tile( + y, + tiler=(1, CHANNELS), + coord=(row, 0), + ) + residual0_tile = cute.local_tile( + residual, + tiler=(1, CHANNELS), + coord=(row, 0), + ) + out0_tile = cute.local_tile( + out, + tiler=(1, CHANNELS), + coord=(row, 0), + ) + thread_y = thread_copy.partition_S(y0_tile) + thread_residual = thread_copy.partition_S(residual0_tile) + thread_out = thread_copy.partition_D(out0_tile) + y_fragment = cute.make_fragment_like(thread_y, cutlass.Float32) + residual_fragment = cute.make_fragment_like( + thread_residual, + cutlass.Float32, + ) + cute.copy(vector_copy, thread_y, y_fragment) + cute.copy(vector_copy, thread_residual, residual_fragment) + for value_idx in cutlass.range_constexpr(VECTOR_WIDTH): + y0 = y_fragment[value_idx].to(cutlass.Float32) + value = y0 * _sigmoid(y0) + value += residual_fragment[value_idx].to(cutlass.Float32) + residual_fragment[value_idx] = value + cute.copy(vector_copy, residual_fragment, thread_out) + + for gate_index in cutlass.range_constexpr(GATE_COUNT): + logits_tile = cute.local_tile( + logits, + tiler=(1, 1, CHANNELS), + coord=(focus, edge, gate_index), + ) + logits_panel = cute.make_tensor(logits_tile.iterator, channel_layout) + thread_logits = thread_copy.partition_S(logits_panel) + gate_fragment = cute.make_fragment_like( + thread_logits, + cutlass.Float32, + ) + cute.copy(vector_copy, thread_logits, gate_fragment) + for value_idx in cutlass.range_constexpr(VECTOR_WIDTH): + gate_fragment[value_idx] = _sigmoid( + gate_fragment[value_idx].to(cutlass.Float32) + ) + + for repeat in cutlass.range_constexpr(3): + degree = 1 + gate_index + repeat * GATE_COUNT + y_tile = cute.local_tile( + y, + tiler=(1, CHANNELS), + coord=(row, degree), + ) + residual_tile = cute.local_tile( + residual, + tiler=(1, CHANNELS), + coord=(row, degree), + ) + out_tile = cute.local_tile( + out, + tiler=(1, CHANNELS), + coord=(row, degree), + ) + thread_y = thread_copy.partition_S(y_tile) + thread_residual = thread_copy.partition_S(residual_tile) + thread_out = thread_copy.partition_D(out_tile) + y_fragment = cute.make_fragment_like(thread_y, cutlass.Float32) + residual_fragment = cute.make_fragment_like( + thread_residual, + cutlass.Float32, + ) + cute.copy(vector_copy, thread_y, y_fragment) + cute.copy(vector_copy, thread_residual, residual_fragment) + for value_idx in cutlass.range_constexpr(VECTOR_WIDTH): + value = y_fragment[value_idx].to(cutlass.Float32) * gate_fragment[ + value_idx + ].to(cutlass.Float32) + value += residual_fragment[value_idx].to(cutlass.Float32) + residual_fragment[value_idx] = value + cute.copy(vector_copy, residual_fragment, thread_out) + + +@cute.jit +def neo_gate_split_structural_vec4_sm80_backward_jit( + grad_out: cute.Tensor, + y: cute.Tensor, + logits: cute.Tensor, + grad_y: cute.Tensor, + grad_logits: cute.Tensor, + stream: CUstream, +): + copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + y.element_type, + num_bits_per_copy=y.element_type.width * VECTOR_WIDTH, + ) + channel_thread_layout = cute.make_ordered_layout( + (1, CHANNEL_GROUPS), + order=(1, 0), + ) + channel_value_layout = cute.make_ordered_layout( + (1, VECTOR_WIDTH), + order=(1, 0), + ) + vector_copy = cute.make_tiled_copy_tv( + copy_atom, + channel_thread_layout, + channel_value_layout, + ) + channel_layout = cute.make_layout((1, CHANNELS), stride=(CHANNELS, 1)) + rows, _ = y.shape + neo_gate_split_structural_vec4_sm80_backward_kernel( + grad_out, + y, + logits, + grad_y, + grad_logits, + channel_layout, + vector_copy, + ).launch( + grid=[cute.ceil_div(rows, ROWS_PER_BLOCK), 1, 1], + block=[THREADS, 1, 1], + stream=stream, + ) + + +@cute.kernel +def neo_gate_split_structural_vec4_sm80_backward_kernel( + grad_out: cute.Tensor, + y: cute.Tensor, + logits: cute.Tensor, + grad_y: cute.Tensor, + grad_logits: cute.Tensor, + channel_layout: cute.Layout, + vector_copy: cute.TiledCopy, +): + tidx, _, _ = cute.arch.thread_idx() + block_row, _, _ = cute.arch.block_idx() + row_slot = tidx // CHANNEL_GROUPS + channel_group = tidx - row_slot * CHANNEL_GROUPS + row = block_row * ROWS_PER_BLOCK + row_slot + rows, _ = y.shape + + if row < rows: + edge = row // FOCUS_COUNT + focus = row - edge * FOCUS_COUNT + thread_copy = vector_copy.get_slice(channel_group) + + y0_tile = cute.local_tile(y, tiler=(1, CHANNELS), coord=(row, 0)) + grad_y0_tile = cute.local_tile( + grad_y, + tiler=(1, CHANNELS), + coord=(row, 0), + ) + grad_out0_panel = cute.local_tile( + grad_out, + tiler=(1, CHANNELS), + coord=(row, 0), + ) + thread_y0 = thread_copy.partition_S(y0_tile) + thread_grad_out0 = thread_copy.partition_S(grad_out0_panel) + thread_grad_y0 = thread_copy.partition_D(grad_y0_tile) + y0_fragment = cute.make_fragment_like(thread_y0, cutlass.Float32) + grad_out0_fragment = cute.make_fragment_like( + thread_grad_out0, + cutlass.Float32, + ) + cute.copy(vector_copy, thread_y0, y0_fragment) + cute.copy(vector_copy, thread_grad_out0, grad_out0_fragment) + for value_idx in cutlass.range_constexpr(VECTOR_WIDTH): + y0 = y0_fragment[value_idx].to(cutlass.Float32) + sig0 = _sigmoid(y0) + grad0 = grad_out0_fragment[value_idx].to(cutlass.Float32) + y0_fragment[value_idx] = ( + grad0 + * sig0 + * (cutlass.Float32(1.0) + y0 * (cutlass.Float32(1.0) - sig0)) + ) + cute.copy(vector_copy, y0_fragment, thread_grad_y0) + + for gate_index in cutlass.range_constexpr(GATE_COUNT): + logits_tile = cute.local_tile( + logits, + tiler=(1, 1, CHANNELS), + coord=(focus, edge, gate_index), + ) + logits_panel = cute.make_tensor(logits_tile.iterator, channel_layout) + thread_logits = thread_copy.partition_S(logits_panel) + gate_fragment = cute.make_fragment_like( + thread_logits, + cutlass.Float32, + ) + grad_logit_fragment = cute.make_fragment_like( + thread_logits, + cutlass.Float32, + ) + cute.copy(vector_copy, thread_logits, gate_fragment) + for value_idx in cutlass.range_constexpr(VECTOR_WIDTH): + gate_fragment[value_idx] = _sigmoid( + gate_fragment[value_idx].to(cutlass.Float32) + ) + grad_logit_fragment[value_idx] = cutlass.Float32(0.0) + + for repeat in cutlass.range_constexpr(3): + degree = 1 + gate_index + repeat * GATE_COUNT + y_tile = cute.local_tile( + y, + tiler=(1, CHANNELS), + coord=(row, degree), + ) + grad_y_tile = cute.local_tile( + grad_y, + tiler=(1, CHANNELS), + coord=(row, degree), + ) + grad_out_panel = cute.local_tile( + grad_out, + tiler=(1, CHANNELS), + coord=(row, degree), + ) + thread_y = thread_copy.partition_S(y_tile) + thread_grad_out = thread_copy.partition_S(grad_out_panel) + thread_grad_y = thread_copy.partition_D(grad_y_tile) + y_fragment = cute.make_fragment_like(thread_y, cutlass.Float32) + grad_out_fragment = cute.make_fragment_like( + thread_grad_out, + cutlass.Float32, + ) + cute.copy(vector_copy, thread_y, y_fragment) + cute.copy(vector_copy, thread_grad_out, grad_out_fragment) + for value_idx in cutlass.range_constexpr(VECTOR_WIDTH): + gate = gate_fragment[value_idx].to(cutlass.Float32) + gout = grad_out_fragment[value_idx].to(cutlass.Float32) + y_value = y_fragment[value_idx].to(cutlass.Float32) + y_fragment[value_idx] = gout * gate + grad_logit_fragment[value_idx] += ( + gout * y_value * gate * (cutlass.Float32(1.0) - gate) + ) + cute.copy(vector_copy, y_fragment, thread_grad_y) + + grad_logits_tile = cute.local_tile( + grad_logits, + tiler=(1, 1, CHANNELS), + coord=(focus, edge, gate_index), + ) + grad_logits_panel = cute.make_tensor( + grad_logits_tile.iterator, + channel_layout, + ) + thread_grad_logits = thread_copy.partition_D(grad_logits_panel) + cute.copy(vector_copy, grad_logit_fragment, thread_grad_logits) + + +@lru_cache(maxsize=8) +def compile_neo_gate_split_structural_vec4_sm80_forward( + compile_identity: tuple[int, int, int] | None = None, +) -> Callable: + if ( + compile_identity is not None + and compile_identity[1:] not in runtime_policy.SUPPORTED_K1_CAPABILITIES + ): + raise ValueError( + "vectorized structural gate forward requires a supported K1 device" + ) + rows = cute.sym_int64() + edges = cute.sym_int64() + fake_residual = make_fake_compact_tensor( + cutlass.Float32, + (rows, REDUCED_COUNT * CHANNELS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_y = make_fake_compact_tensor( + cutlass.Float32, + (rows, REDUCED_COUNT * CHANNELS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_logits = make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, edges, GATE_COUNT * CHANNELS), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + fake_out = make_fake_compact_tensor( + cutlass.Float32, + (rows, REDUCED_COUNT * CHANNELS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return _guard_vec4_dispatch( + cute.compile( + neo_gate_split_structural_vec4_sm80_forward_jit, + fake_residual, + fake_y, + fake_logits, + fake_out, + fake_stream, + options="--enable-tvm-ffi", + ), + FORWARD_TENSOR_NAMES, + ) + + +@lru_cache(maxsize=8) +def compile_neo_gate_split_structural_vec4_sm80_backward( + compile_identity: tuple[int, int, int] | None = None, +) -> Callable: + if ( + compile_identity is not None + and compile_identity[1:] not in runtime_policy.SUPPORTED_K1_CAPABILITIES + ): + raise ValueError( + "vectorized structural gate backward requires a supported K1 device" + ) + rows = cute.sym_int64() + edges = cute.sym_int64() + fake_grad_out = make_fake_compact_tensor( + cutlass.Float32, + (rows, REDUCED_COUNT * CHANNELS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_y = make_fake_compact_tensor( + cutlass.Float32, + (rows, REDUCED_COUNT * CHANNELS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_logits = make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, edges, GATE_COUNT * CHANNELS), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + fake_grad_y = make_fake_compact_tensor( + cutlass.Float32, + (rows, REDUCED_COUNT * CHANNELS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_grad_logits = make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, edges, GATE_COUNT * CHANNELS), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return _guard_vec4_dispatch( + cute.compile( + neo_gate_split_structural_vec4_sm80_backward_jit, + fake_grad_out, + fake_y, + fake_logits, + fake_grad_y, + fake_grad_logits, + fake_stream, + options="--enable-tvm-ffi", + ), + BACKWARD_TENSOR_NAMES, + ) diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_message_grid_product.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_message_grid_product.py new file mode 100644 index 0000000000..54d66b9d99 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_message_grid_product.py @@ -0,0 +1,275 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Packed strict-FP32 CuTe grid product for Neo's F=2 K1 branch.""" + +from __future__ import ( + annotations, +) + +from collections.abc import ( + Callable, +) +from functools import ( + lru_cache, +) +from typing import ( + Any, +) + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from .. import ( + runtime_policy, +) + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN001, ANN201, TC003 + + +PACKED_COEFF_DIM = 48 +HIDDEN_CHANNELS = 64 +GRID_SIZE = 152 +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} +_SUPPORTED_CAPABILITIES = runtime_policy.SM80_PROFILE_CAPABILITIES | { + runtime_policy.SM90_CAPABILITY +} + + +def _make_grid_operation( + operation_type: type[Any], + *, + panel_adjoint: bool = False, +) -> Any: + from ..output_grid_kernels import cute_tiled_grid_product as tiled + + # F=2 and C=32 form one complete 64-channel panel. Bypass only the + # shared readout policy; the tiled implementation itself is unchanged. + operation = operation_type.__new__(operation_type) + operation.hidden_channels = HIDDEN_CHANNELS + operation.tile_k = tiled.TILE_K + operation.sm80_c96_n48_panel = bool(panel_adjoint) + operation.cta_tiler = (tiled.TILE_M, tiled.TILE_N, operation.tile_k) + operation.channel_tile_start = 0 + operation.channel_tiles = 1 + operation.has_channel_residue = False + operation.cta_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=tiled.THREADS, + ) + return operation + + +def _validate_compile_target( + device_index: int, + compute_capability: tuple[int, int], +) -> None: + import torch + + actual = tuple(torch.cuda.get_device_capability(device_index)) + if actual != tuple(compute_capability): + raise ValueError("compile target does not match the selected CUDA device") + if actual not in _SUPPORTED_CAPABILITIES: + raise ValueError( + "packed Neo message-grid product requires the SM80-family profile or sm90" + ) + + +def _fake_inputs() -> tuple[Any, Any, Any]: + nodes = cute.sym_int64() + coeff = make_fake_compact_tensor( + cutlass.Float32, + (nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + to_grid = make_fake_compact_tensor( + cutlass.Float32, + (GRID_SIZE, PACKED_COEFF_DIM), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + from_grid = make_fake_compact_tensor( + cutlass.Float32, + (PACKED_COEFF_DIM, GRID_SIZE), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + return coeff, to_grid, from_grid + + +def compile_message_grid_product_forward( + device_index: int, + compute_capability: tuple[int, int], +) -> Callable: + """Compile the F=2 packed product with a symbolic node count.""" + import torch + + _validate_compile_target(device_index, compute_capability) + from ..output_grid_kernels.cute_tiled_grid_product import ( + TiledOutputGridProductForward, + ) + + coeff, to_grid, from_grid = _fake_inputs() + stream = make_fake_stream(use_tvm_ffi_env_stream=True) + with torch.cuda.device(device_index): + return cute.compile( + _make_grid_operation(TiledOutputGridProductForward), + coeff, + coeff, + to_grid, + from_grid, + coeff, + stream, + options="--enable-tvm-ffi", + ) + + +def compile_message_grid_product_backward( + device_index: int, + compute_capability: tuple[int, int], +) -> Callable: + """Compile both packed input adjoints with a symbolic node count.""" + import torch + + _validate_compile_target(device_index, compute_capability) + from ..output_grid_kernels.cute_tiled_grid_product import ( + TiledOutputGridProductBackward, + ) + + coeff, to_grid, from_grid = _fake_inputs() + stream = make_fake_stream(use_tvm_ffi_env_stream=True) + with torch.cuda.device(device_index): + return cute.compile( + _make_grid_operation( + TiledOutputGridProductBackward, + panel_adjoint=( + compute_capability in runtime_policy.SM80_PROFILE_CAPABILITIES + ), + ), + coeff, + coeff, + coeff, + to_grid, + from_grid, + coeff, + coeff, + stream, + options="--enable-tvm-ffi", + ) + + +@lru_cache(maxsize=8) +def _compiled_forward( + device_index: int, + compute_capability: tuple[int, int], +) -> Callable: + return compile_message_grid_product_forward(device_index, compute_capability) + + +@lru_cache(maxsize=8) +def _compiled_backward( + device_index: int, + compute_capability: tuple[int, int], +) -> Callable: + return compile_message_grid_product_backward(device_index, compute_capability) + + +def _compile_identity(tensor) -> tuple[int, tuple[int, int]]: + import torch + + device_index = tensor.device.index + if device_index is None: + device_index = torch.cuda.current_device() + return int(device_index), tuple(torch.cuda.get_device_capability(device_index)) + + +def _validate_tensors(left, right, to_grid, from_grid, grad_out=None) -> None: + import torch + + floating = (left, right, to_grid, from_grid) + if ( + tuple(left.shape[1:]) != (PACKED_COEFF_DIM, HIDDEN_CHANNELS) + or left.shape[0] <= 0 + or right.shape != left.shape + or tuple(to_grid.shape) != (GRID_SIZE, PACKED_COEFF_DIM) + or tuple(from_grid.shape) != (PACKED_COEFF_DIM, GRID_SIZE) + or any(not tensor.is_cuda for tensor in floating) + or any(tensor.device != left.device for tensor in floating) + or any(tensor.dtype != torch.float32 for tensor in floating) + or any(not tensor.is_contiguous() for tensor in floating) + or any(tensor.data_ptr() % 16 != 0 for tensor in floating) + or tuple(torch.cuda.get_device_capability(left.device)) + not in _SUPPORTED_CAPABILITIES + or not runtime_policy.uses_strict_fp32_matmul() + ): + raise ValueError( + "packed message-grid product requires contiguous SM80-family/SM90 FP32 " + "left/right=(N,48,64), to_grid=(152,48), and from_grid=(48,152)" + ) + if grad_out is not None and ( + grad_out.shape != left.shape + or grad_out.device != left.device + or grad_out.dtype != torch.float32 + or not grad_out.is_contiguous() + or grad_out.data_ptr() % 16 != 0 + ): + raise ValueError( + "packed message-grid grad_out must be contiguous and match left" + ) + + +def run_message_grid_product(left, right, to_grid, from_grid): + """Run the F=2 product without materializing projection-layout clones.""" + import torch + + _validate_tensors(left, right, to_grid, from_grid) + out = torch.empty_like(left) + with torch.cuda.device(left.device): + _compiled_forward(*_compile_identity(left))( + left, + right, + to_grid, + from_grid, + out, + ) + return out + + +def run_message_grid_product_backward( + grad_out, + left, + right, + to_grid, + from_grid, +): + """Run both F=2 product input adjoints in the packed layout.""" + import torch + + _validate_tensors(left, right, to_grid, from_grid, grad_out) + grad_left = torch.empty_like(left) + grad_right = torch.empty_like(right) + with torch.cuda.device(left.device): + _compiled_backward(*_compile_identity(left))( + grad_out, + left, + right, + to_grid, + from_grid, + grad_left, + grad_right, + ) + return grad_left, grad_right + + +__all__ = [ + "compile_message_grid_product_backward", + "compile_message_grid_product_forward", + "run_message_grid_product", + "run_message_grid_product_backward", +] diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_output_gate_backward.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_output_gate_backward.py new file mode 100644 index 0000000000..e368e32d40 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_output_gate_backward.py @@ -0,0 +1,204 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Manual backward for the fused Neo attention output gate. + +Only input gradients are produced. The gate and RMSNorm parameter gradients +are intentionally omitted for the E/F/S path. ``grad_phase`` may alias +``grad_gated``; ``grad_x_wide`` is an existing accumulation buffer and only its +first 64 values per node are updated. + +The forward gate is recomputed. Its logit gradient uses +``sum(grad_gated * gated_out) * (1 - gate)``, so backward does not need either +an ungated Phase-C aggregate or a saved gate tensor. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + Callable, +) + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN001, ANN201, ANN202, TC002, UP035 + +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} +DEGREE_COUNT = 16 +FOCUS_COUNT = 2 +CHANNELS = 32 +HIDDEN = FOCUS_COUNT * CHANNELS +OUTPUT_WIDTH = DEGREE_COUNT * HIDDEN + + +@cute.jit +def _warp_sum(value): + return cute.arch.warp_reduction_sum(value) + + +@cute.jit +def _sigmoid(value): + one = cutlass.Float32(1.0) + return one / (one + cute.exp(-value)) + + +@cute.jit +def compute_neo_inv_rms(x, eps: cutlass.Constexpr[float]): + square_sum = _warp_sum(x * x) + return cute.rsqrt(square_sum / cutlass.Float32(CHANNELS) + cutlass.Float32(eps)) + + +@cute.jit +def compute_neo_output_gate( + x, + norm_scale: cute.Tensor, + gate_weight: cute.Tensor, + inv_rms, + focus, + channel, +): + logit_part = ( + x + * inv_rms + * norm_scale[focus, channel].to(cutlass.Float32) + * gate_weight[channel, focus, 0].to(cutlass.Float32) + ) + return _sigmoid(_warp_sum(logit_part)) + + +@cute.jit +def neo_output_gate_backward_jit( + grad_gated: cute.Tensor, + gated_out: cute.Tensor, + x_wide: cute.Tensor, + norm_scale: cute.Tensor, + gate_weight: cute.Tensor, + grad_phase: cute.Tensor, + grad_x_wide: cute.Tensor, + stream: CUstream, + eps: cutlass.Constexpr[float], +): + nodes, _ = gated_out.shape + neo_output_gate_backward_kernel( + grad_gated, + gated_out, + x_wide, + norm_scale, + gate_weight, + grad_phase, + grad_x_wide, + eps, + ).launch( + grid=[nodes, 1, 1], + block=[HIDDEN, 1, 1], + stream=stream, + ) + + +@cute.kernel +def neo_output_gate_backward_kernel( + grad_gated: cute.Tensor, + gated_out: cute.Tensor, + x_wide: cute.Tensor, + norm_scale: cute.Tensor, + gate_weight: cute.Tensor, + grad_phase: cute.Tensor, + grad_x_wide: cute.Tensor, + eps: cutlass.Constexpr[float], +): + tid, _, _ = cute.arch.thread_idx() + node, _, _ = cute.arch.block_idx() + focus = tid // CHANNELS + channel = tid - focus * CHANNELS + + x = x_wide[node, tid].to(cutlass.Float32) + inv_rms = compute_neo_inv_rms(x, eps) + gate_value = compute_neo_output_gate( + x, + norm_scale, + gate_weight, + inv_rms, + focus, + channel, + ) + + gate_dot = cutlass.Float32(0.0) + for degree in cutlass.range_constexpr(DEGREE_COUNT): + idx = degree * HIDDEN + tid + grad = grad_gated[node, idx].to(cutlass.Float32) + gated = gated_out[node, idx].to(cutlass.Float32) + grad_phase[node, idx] = (grad * gate_value).to(grad_phase.element_type) + gate_dot += grad * gated + + gate_dot = _warp_sum(gate_dot) + grad_logit = gate_dot * (cutlass.Float32(1.0) - gate_value) + + scale = norm_scale[focus, channel].to(cutlass.Float32) + weight = gate_weight[channel, focus, 0].to(cutlass.Float32) + grad_scaled = grad_logit * weight * scale + rms_coeff = _warp_sum(grad_scaled * x) / cutlass.Float32(CHANNELS) + + grad_x = grad_scaled * inv_rms + grad_x -= x * inv_rms * inv_rms * inv_rms * rms_coeff + previous = grad_x_wide[node, tid].to(cutlass.Float32) + grad_x_wide[node, tid] = (previous + grad_x).to(grad_x_wide.element_type) + + +def compile_neo_output_gate_backward(eps: float) -> Callable: + """Compile input-only backward. + + Runtime order is ``grad_gated, gated_out, x_wide, norm_scale, gate_weight, + grad_phase, grad_x_wide``. ``grad_phase`` may alias ``grad_gated`` and the + kernel adds only the scalar row into the preinitialized ``grad_x_wide``. + """ + nodes = cute.sym_int64() + + def fake_node_tensor(): + return make_fake_compact_tensor( + cutlass.Float32, + (nodes, OUTPUT_WIDTH), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + + grad_gated = fake_node_tensor() + gated_out = fake_node_tensor() + x_wide = fake_node_tensor() + norm_scale = make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, CHANNELS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + gate_weight = make_fake_compact_tensor( + cutlass.Float32, + (CHANNELS, FOCUS_COUNT, 1), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + grad_phase = fake_node_tensor() + grad_x_wide = fake_node_tensor() + stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + neo_output_gate_backward_jit, + grad_gated, + gated_out, + x_wide, + norm_scale, + gate_weight, + grad_phase, + grad_x_wide, + stream, + eps, + options="--enable-tvm-ffi", + ) diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_a_radial_forward.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_a_radial_forward.py new file mode 100644 index 0000000000..0855d3b426 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_a_radial_forward.py @@ -0,0 +1,358 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Packed-direct Phase-A/radial forward without an ``x_rot`` boundary.""" + +# ruff: noqa: ANN001, ANN201, ANN202, TC002, UP035 + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + Callable, +) + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from ..compile_cache import ( + device_aware_lru_cache, +) +from ..k1_wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES + +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + + +@dataclass(frozen=True) +class NeoPhaseARadialForwardParams: + x_wide: cute.Tensor + src: cute.Tensor + d_full: cute.Tensor + radial_m0: cute.Tensor + combined_weight: cute.Tensor + hidden_weight: cute.Tensor + channel_basis: cute.Tensor + out: cute.Tensor + rad_l0: cute.Tensor + + +@cute.jit +def _store_packed_phase_a_value( + params: NeoPhaseARadialForwardParams, + x_local: cute.Tensor, + edge, + src_node, + channel, + reduced: cutlass.Constexpr[int], + panel_start: cutlass.Constexpr[int], + full_start: cutlass.Constexpr[int], + width: cutlass.Constexpr[int], +): + acc = cutlass.Float32(0.0) + for local_col in cutlass.range_constexpr(width): + d_val = params.d_full[edge, panel_start + local_col].to(cutlass.Float32) + x_val = params.x_wide[ + src_node, + (full_start + local_col) * 64 + channel, + ].to(cutlass.Float32) + acc += d_val * x_val + x_local[reduced * 64 + channel] = acc + + +@cute.jit +def neo_phase_a_radial_forward_packed_direct_saved_jit( + x_wide: cute.Tensor, + src: cute.Tensor, + d_full: cute.Tensor, + radial_m0: cute.Tensor, + combined_weight: cute.Tensor, + hidden_weight: cute.Tensor, + channel_basis: cute.Tensor, + out: cute.Tensor, + rad_l0: cute.Tensor, + compact_out: cute.Tensor, + stream: CUstream, +): + params = NeoPhaseARadialForwardParams( + x_wide=x_wide, + src=src, + d_full=d_full, + radial_m0=radial_m0, + combined_weight=combined_weight, + hidden_weight=hidden_weight, + channel_basis=channel_basis, + out=out, + rad_l0=rad_l0, + ) + edges, _ = out.shape + neo_phase_a_radial_forward_packed_direct_saved_kernel( + params, + compact_out, + ).launch( + grid=[edges, 1, 1], + block=[64, 1, 1], + stream=stream, + ) + + +@cute.kernel +def neo_phase_a_radial_forward_packed_direct_saved_kernel( + params: NeoPhaseARadialForwardParams, + compact_out: cute.Tensor, +): + channel, _, _ = cute.arch.thread_idx() + edge, _, _ = cute.arch.block_idx() + + smem = cutlass.utils.SmemAllocator() + x_local = smem.allocate_tensor(cutlass.Float32, 10 * 64) + compact = smem.allocate_tensor(cutlass.Float32, 25) + src_node = params.src[edge] + + _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 0, 0, 0, 1) + _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 1, 1, 1, 3) + _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 2, 10, 4, 5) + _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 3, 25, 9, 7) + _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 4, 4, 1, 3) + _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 5, 15, 4, 5) + _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 6, 32, 9, 7) + _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 7, 7, 1, 3) + _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 8, 20, 4, 5) + _store_packed_phase_a_value(params, x_local, edge, src_node, channel, 9, 39, 9, 7) + + if channel < 25: + acc = cutlass.Float32(0.0) + for radial_idx in cutlass.range_constexpr(4 * 32): + radial_value = params.radial_m0[edge, radial_idx].to(cutlass.Float32) + weight = params.combined_weight[radial_idx, channel].to(cutlass.Float32) + acc += radial_value * weight + compact[channel] = acc + compact_out[edge, channel] = acc + + acc_l0 = cutlass.Float32(0.0) + for radial_channel in cutlass.range_constexpr(32): + radial_value = params.radial_m0[edge, radial_channel].to(cutlass.Float32) + weight = params.hidden_weight[radial_channel, channel].to(cutlass.Float32) + acc_l0 += radial_value * weight + params.rad_l0[edge, channel] = acc_l0 + + cute.arch.sync_threads() + + for coeff in cutlass.range_constexpr(10): + acc = cutlass.Float32(0.0) + if coeff < 4: + out_coeff = coeff + for in_coeff in cutlass.range_constexpr(4): + kval = compact[in_coeff * 4 + out_coeff] + acc += kval * x_local[in_coeff * 64 + channel] + elif coeff < 7: + out_coeff = coeff - 4 + for in_coeff in cutlass.range_constexpr(3): + kval = compact[16 + in_coeff * 3 + out_coeff] + acc += kval * x_local[(4 + in_coeff) * 64 + channel] + else: + out_coeff = coeff - 7 + for in_coeff in cutlass.range_constexpr(3): + kval = compact[16 + in_coeff * 3 + out_coeff] + acc += kval * x_local[(7 + in_coeff) * 64 + channel] + acc *= params.channel_basis[channel].to(cutlass.Float32) + focus = channel // 32 + focus_channel = channel - focus * 32 + out_idx = focus * 10 * 32 + coeff * 32 + focus_channel + params.out[edge, out_idx] = acc + + +def compile_neo_phase_a_radial_forward_packed_direct() -> Callable: + edge_count = cute.sym_int64() + node_count = cute.sym_int64() + fake_x_wide = make_fake_compact_tensor( + cutlass.Float32, + (node_count, 16 * 64), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_src = make_fake_compact_tensor( + cutlass.Int32, + (edge_count,), + stride_order=(0,), + **FAKE_TENSOR_KW, + ) + fake_d = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, PACKED_WIGNER_VALUES), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_radial = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, 4 * 32), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_combined = make_fake_compact_tensor( + cutlass.Float32, + (4 * 32, 25), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_hidden = make_fake_compact_tensor( + cutlass.Float32, + (32, 64), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_basis = make_fake_compact_tensor( + cutlass.Float32, + (64,), + stride_order=(0,), + **FAKE_TENSOR_KW, + ) + fake_out = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, 10 * 64), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_rad_l0 = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, 64), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True) + fake_compact = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, 25), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + return cute.compile( + neo_phase_a_radial_forward_packed_direct_saved_jit, + fake_x_wide, + fake_src, + fake_d, + fake_radial, + fake_combined, + fake_hidden, + fake_basis, + fake_out, + fake_rad_l0, + fake_compact, + fake_stream, + options="--enable-tvm-ffi", + ) + + +@device_aware_lru_cache(maxsize=4) +def _compiled_neo_phase_a_radial_forward_packed_direct() -> Callable: + return compile_neo_phase_a_radial_forward_packed_direct() + + +def _combined_radial_weight(radial_hidden_proj, radial_degree_mixer): + import torch + + hidden_weight = radial_hidden_proj.weight + mixer_weight = radial_degree_mixer.weight + 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, + ) + if cache is not None and cache[0] == key: + return cache[1] + + blocks = [] + for degree in range(4): + mixer_block = mixer_weight.detach()[degree * 64 : (degree + 1) * 64, :] + blocks.append(torch.mm(hidden_weight.detach(), mixer_block)) + combined = torch.cat(blocks, dim=0).contiguous() + radial_hidden_proj._deepmd_cute_neo_radial_combined = (key, combined) + return combined + + +def run_neo_phase_a_radial_forward_packed_direct( + *, + radial_hidden_proj, + radial_degree_mixer, + x_wide, + src, + D_full, + radial_feat_m0, +): + """Validate and launch the packed Phase-A/radial forward kernel.""" + import torch + + if x_wide.shape[1:] != (16, 64): + raise ValueError(f"expected x_wide shape (N,16,64), got {x_wide.shape}") + edge_count = src.numel() + if tuple(D_full.shape) != (edge_count, PACKED_WIGNER_VALUES): + raise ValueError( + "expected packed Wigner shape " + f"{(edge_count, PACKED_WIGNER_VALUES)}, got {tuple(D_full.shape)}" + ) + if radial_feat_m0.shape != (edge_count, 4, 32): + raise ValueError( + f"expected radial_feat_m0 shape {(edge_count, 4, 32)}, " + f"got {tuple(radial_feat_m0.shape)}" + ) + if radial_hidden_proj.bias is not None: + raise NotImplementedError("collapsed radial mixer expects no hidden bias") + 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") + + combined_weight = _combined_radial_weight( + radial_hidden_proj, + radial_degree_mixer, + ) + kernel = _compiled_neo_phase_a_radial_forward_packed_direct() + out = torch.empty( + edge_count, + 10 * 64, + device=x_wide.device, + dtype=x_wide.dtype, + ) + rad_l0 = torch.empty( + edge_count, + 64, + device=radial_feat_m0.device, + dtype=radial_feat_m0.dtype, + ) + compact_out = torch.empty( + (edge_count, 25), + device=radial_feat_m0.device, + dtype=torch.float32, + ) + kernel( + x_wide.contiguous().view(x_wide.shape[0], 16 * 64), + src.to(torch.int32).contiguous(), + D_full.contiguous(), + radial_feat_m0.contiguous().view(edge_count, 4 * 32), + combined_weight, + radial_hidden_proj.weight.detach().contiguous(), + radial_degree_mixer.channel_basis.detach().view(64).contiguous(), + out, + rad_l0, + compact_out, + ) + return ( + out.view(edge_count, 2, 10, 32), + rad_l0.view(edge_count, 2, 32), + compact_out, + ) diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout.py new file mode 100644 index 0000000000..e4a55bfd2f --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout.py @@ -0,0 +1,703 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Exact-shape Neo Phase-C backward with fused layout-boundary reductions. + +This kernel specializes ``D=16``, ``Dm=10``, ``F=2``, ``C=32`` and the +46-value packed Wigner panel. It keeps the destination adjoint in shared +memory once per node, but streams each degree once per edge into a lane-local +ten-value input-adjoint fragment. ``grad_Dt`` is reduced by the two focus +warps, so the effective stack input does not need a shared-memory slab. + +The kernel also owns the two reductions immediately downstream of Phase C: + +* attention ``grad_alpha`` is consumed in-place to produce envelope-softmax + ``grad_logits``, ``grad_edge`` and ``grad_z``; +* focus ``grad_alpha`` is consumed in-place by the two-focus softmax/RMSNorm + backward to produce ``grad_focus_src``. + +The stack adjoint is written edge-major into the fully consumed final saved +activation, avoiding a separate edge-sized allocation. The focus-source +gradient remains focus-major. +""" + +# ruff: noqa: ANN001, ANN201, ANN202, TC002 + +from __future__ import ( + annotations, +) + +import operator +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, +) + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from ..k1_wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + +DEGREE_COUNT = 16 +REDUCED_COUNT = 10 +N_FOCUS = 2 +FOCUS_CHANNELS = 32 +HIDDEN = N_FOCUS * FOCUS_CHANNELS + +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + + +@dataclass(frozen=True) +class NeoPhaseCBackwardLayoutParams: + grad_out: cute.Tensor + stack: cute.Tensor + wigner_dt: cute.Tensor + alpha: cute.Tensor + 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 + denom: cute.Tensor + focus_src: cute.Tensor + focus_weight: cute.Tensor + focus_scale: cute.Tensor + grad_stack: cute.Tensor + grad_wigner_dt: cute.Tensor + grad_logits: cute.Tensor + grad_edge: cute.Tensor + grad_z_partial: cute.Tensor + grad_z: cute.Tensor + grad_focus_src: cute.Tensor + + +@cute.jit +def _warp_sum(value): + return cute.arch.warp_reduction(value, operator.add) + + +@cute.jit +def _sigmoid(value): + one = cutlass.Float32(1.0) + return one / (one + cute.exp(-value)) + + +@cute.jit +def _record_panel_term( + panel_values, + panel_offset, + dt_partial, + raw_fragment, + grad_fragment, + transformed, + head, + lane, + reduced: cutlass.Constexpr[int], + panel_index: cutlass.Constexpr[int], + alpha_value, + focus_value, +): + """Consume one structural Wigner entry for grad-stack and grad-Dt.""" + panel = panel_values[panel_offset + panel_index] + grad_fragment[reduced] = ( + grad_fragment[reduced].to(cutlass.Float32) + panel * transformed + ) + + effective_stack = raw_fragment[reduced].to(cutlass.Float32) * focus_value + partial = _warp_sum(transformed * effective_stack * alpha_value) + if lane == 0: + dt_partial[head * PACKED_WIGNER_VALUES + panel_index] = partial + + +@cute.jit +def _focus_source_backward( + params: NeoPhaseCBackwardLayoutParams, + focus_grad, + focus_inv, + focus_grad_logits, + focus_coeff, + edge, + head, + lane, + tidx, + eps: cutlass.Constexpr[float], + tau: cutlass.Constexpr[float], + label_smoothing: cutlass.Constexpr[float], +): + """Consume the Phase-C focus-alpha reduction without a global temporary.""" + value = params.focus_src[edge, head, lane].to(cutlass.Float32) + scale = params.focus_scale[head, lane].to(cutlass.Float32) + weight = params.focus_weight[lane, head].to(cutlass.Float32) + + square_sum = _warp_sum(value * value) + if lane == 0: + focus_inv[head] = cute.rsqrt( + square_sum / cutlass.Float32(FOCUS_CHANNELS) + cutlass.Float32(eps) + ) + + if tidx == 0: + probability_keep = cutlass.Float32(1.0 - label_smoothing) + smooth = cutlass.Float32(label_smoothing / N_FOCUS) + probability0 = ( + params.focus_alpha[edge, 0].to(cutlass.Float32) - smooth + ) / probability_keep + probability1 = ( + params.focus_alpha[edge, 1].to(cutlass.Float32) - smooth + ) / probability_keep + keep = cutlass.Float32(1.0 - label_smoothing) + grad0 = focus_grad[0] * keep + grad1 = focus_grad[1] * keep + dot = grad0 * probability0 + grad1 * probability1 + inv_tau = cutlass.Float32(1.0 / tau) + focus_grad_logits[0] = probability0 * (grad0 - dot) * inv_tau + focus_grad_logits[1] = probability1 * (grad1 - dot) * inv_tau + cute.arch.sync_threads() + + grad_scaled = focus_grad_logits[head] * weight * scale + coeff_sum = _warp_sum(grad_scaled * value) + if lane == 0: + focus_coeff[head] = coeff_sum / cutlass.Float32(FOCUS_CHANNELS) + cute.arch.sync_warp() + + inv = focus_inv[head] + grad_value = grad_scaled * inv + grad_value -= value * inv * inv * inv * focus_coeff[head] + params.grad_focus_src[head, edge, lane] = grad_value.to( + params.grad_focus_src.element_type + ) + + +@cute.kernel +def neo_phase_c_backward_layout_kernel( + params: NeoPhaseCBackwardLayoutParams, + raw_layout: cute.Layout, + raw_tiled_copy: cute.TiledCopy, + focus_eps: cutlass.Constexpr[float], + focus_tau: cutlass.Constexpr[float], + focus_label_smoothing: cutlass.Constexpr[float], +): + tidx, _, _ = cute.arch.thread_idx() + node, _, _ = cute.arch.block_idx() + head = tidx // FOCUS_CHANNELS + lane = tidx - head * FOCUS_CHANNELS + lo = params.dst_ptr[node] + hi = params.dst_ptr[node + 1] + + smem = cutlass.utils.SmemAllocator() + t_values = smem.allocate_tensor(cutlass.Float32, DEGREE_COUNT * HIDDEN) + panel_values = smem.allocate_tensor( + cutlass.Float32, + PACKED_WIGNER_VALUES, + ) + dt_partial = smem.allocate_tensor(cutlass.Float32, N_FOCUS * PACKED_WIGNER_VALUES) + focus_grad = smem.allocate_tensor(cutlass.Float32, N_FOCUS) + focus_inv = smem.allocate_tensor(cutlass.Float32, N_FOCUS) + focus_grad_logits = smem.allocate_tensor(cutlass.Float32, N_FOCUS) + focus_coeff = smem.allocate_tensor(cutlass.Float32, N_FOCUS) + softmax_dot_by_focus = smem.allocate_tensor(cutlass.Float32, N_FOCUS) + gate_tile = smem.allocate_tensor(cutlass.Float32, HIDDEN) + + for degree in cutlass.range_constexpr(DEGREE_COUNT): + index = degree * HIDDEN + tidx + upstream = params.grad_out[node, degree, tidx].to(cutlass.Float32) + rotate = params.rotate_inv_rescale[degree].to(cutlass.Float32) + t_values[index] = upstream * rotate + cute.arch.sync_threads() + + raw_thread_copy = raw_tiled_copy.get_slice(lane) + softmax_dot = cutlass.Float32(0.0) + + for edge in cutlass.range(lo, hi, 1, unroll=1): + stack_tile = cute.local_tile( + params.stack, + tiler=(1, 1, REDUCED_COUNT, FOCUS_CHANNELS), + coord=(edge, head, 0, 0), + ) + stack_head = cute.make_tensor(stack_tile.iterator, raw_layout) + thread_stack = raw_thread_copy.partition_S(stack_head) + raw_fragment = cute.make_fragment_like(thread_stack, cutlass.Float32) + cute.copy(raw_tiled_copy, thread_stack, raw_fragment) + if tidx < PACKED_WIGNER_VALUES: + panel_values[tidx] = params.wigner_dt[edge, tidx].to(cutlass.Float32) + cute.arch.sync_threads() + + alpha_value = params.alpha[edge, head].to(cutlass.Float32) + focus_value = params.focus_alpha[edge, head].to(cutlass.Float32) + grad_fragment = cute.make_fragment_like(raw_fragment, cutlass.Float32) + grad_fragment.fill(0.0) + + transformed = t_values[tidx] + _record_panel_term( + panel_values, + 0, + dt_partial, + raw_fragment, + grad_fragment, + transformed, + head, + lane, + 0, + 0, + alpha_value, + focus_value, + ) + for local_col in cutlass.range_constexpr(3): + transformed = t_values[(1 + local_col) * HIDDEN + tidx] + for row_slot in cutlass.range_constexpr(3): + _record_panel_term( + panel_values, + 0, + dt_partial, + raw_fragment, + grad_fragment, + transformed, + head, + lane, + 1 + row_slot * 3, + 1 + row_slot * 3 + local_col, + alpha_value, + focus_value, + ) + for local_col in cutlass.range_constexpr(5): + transformed = t_values[(4 + local_col) * HIDDEN + tidx] + for row_slot in cutlass.range_constexpr(3): + _record_panel_term( + panel_values, + 0, + dt_partial, + raw_fragment, + grad_fragment, + transformed, + head, + lane, + 2 + row_slot * 3, + 10 + row_slot * 5 + local_col, + alpha_value, + focus_value, + ) + for local_col in cutlass.range_constexpr(7): + transformed = t_values[(9 + local_col) * HIDDEN + tidx] + for row_slot in cutlass.range_constexpr(3): + _record_panel_term( + panel_values, + 0, + dt_partial, + raw_fragment, + grad_fragment, + transformed, + head, + lane, + 3 + row_slot * 3, + 25 + row_slot * 7 + local_col, + alpha_value, + focus_value, + ) + + grad_focus_part = cutlass.Float32(0.0) + grad_alpha_part = cutlass.Float32(0.0) + for reduced in cutlass.range_constexpr(REDUCED_COUNT): + raw = raw_fragment[reduced].to(cutlass.Float32) + grad_raw = grad_fragment[reduced].to(cutlass.Float32) + grad_focus_part += grad_raw * raw * alpha_value + grad_alpha_part += grad_raw * raw * focus_value + grad_fragment[reduced] = grad_raw * focus_value * alpha_value + + # Every edge belongs to exactly one node CTA. Both focus warps have + # loaded their complete source fragment, so this exact-address store + # may reuse the fully consumed stack allocation. + grad_stack_tile = cute.local_tile( + params.grad_stack, + tiler=(1, 1, REDUCED_COUNT, FOCUS_CHANNELS), + coord=(edge, head, 0, 0), + ) + grad_stack_head = cute.make_tensor(grad_stack_tile.iterator, raw_layout) + thread_grad_stack = raw_thread_copy.partition_D(grad_stack_head) + cute.copy(raw_tiled_copy, grad_fragment, thread_grad_stack) + + grad_focus_value = _warp_sum(grad_focus_part) + grad_alpha_value = _warp_sum(grad_alpha_part) + if lane == 0: + focus_grad[head] = grad_focus_value + # grad_logits is the dead grad-alpha slab during the first node pass. + params.grad_logits[edge, head] = grad_alpha_value.to( + params.grad_logits.element_type + ) + softmax_dot += grad_alpha_value * alpha_value + + cute.arch.sync_threads() + if tidx < PACKED_WIGNER_VALUES: + value = dt_partial[tidx] + dt_partial[PACKED_WIGNER_VALUES + tidx] + params.grad_wigner_dt[edge, tidx] = value.to( + params.grad_wigner_dt.element_type + ) + + _focus_source_backward( + params, + focus_grad, + focus_inv, + focus_grad_logits, + focus_coeff, + edge, + head, + lane, + tidx, + focus_eps, + focus_tau, + focus_label_smoothing, + ) + cute.arch.sync_threads() + + if lane == 0: + softmax_dot_by_focus[head] = softmax_dot + max_value = params.group_max[node, head].to(cutlass.Float32) + denom = params.denom[node, head].to(cutlass.Float32) + z_sigmoid = _sigmoid(params.z_bias_raw[head].to(cutlass.Float32)) + params.grad_z_partial[node, head] = ( + -softmax_dot * cute.exp(-max_value) / denom * z_sigmoid + ).to(params.grad_z_partial.element_type) + cute.arch.sync_threads() + # The edge loop below is lane-striped, so every lane needs the node dot. + softmax_dot = softmax_dot_by_focus[head] + + for edge_base in cutlass.range(lo, hi, FOCUS_CHANNELS, unroll=1): + edge = edge_base + lane + gate_contribution = cutlass.Float32(0.0) + if edge < hi: + upstream = params.grad_logits[edge, head].to(cutlass.Float32) + centered = upstream - softmax_dot + alpha_value = params.alpha[edge, head].to(cutlass.Float32) + params.grad_logits[edge, head] = (alpha_value * centered).to( + params.grad_logits.element_type + ) + + gate = params.edge_gate[edge].to(cutlass.Float32) + if gate < cutlass.Float32(0.0): + gate = cutlass.Float32(0.0) + if gate > cutlass.Float32(0.0): + gate_contribution = alpha_value * centered * cutlass.Float32(2.0) / gate + gate_tile[tidx] = gate_contribution + cute.arch.sync_threads() + if head == 0 and edge < hi: + params.grad_edge[edge] = (gate_tile[lane] + gate_tile[lane + 32]).to( + params.grad_edge.element_type + ) + cute.arch.sync_threads() + + +@cute.jit +def _cta_sum(value, scratch, threads: cutlass.Constexpr[int]): + lane = cute.arch.lane_idx() + warp = cute.arch.warp_idx() + warps = threads // 32 + value = _warp_sum(value) + if lane == 0: + scratch[warp] = value + cute.arch.barrier() + + total = cutlass.Float32(0.0) + if lane < warps: + total = scratch[lane] + return _warp_sum(total) + + +@cute.kernel +def neo_phase_c_backward_z_reduce_kernel( + params: NeoPhaseCBackwardLayoutParams, + threads: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + head, _, _ = cute.arch.block_idx() + node_count, _ = params.grad_z_partial.shape + warps = threads // 32 + + smem = cutlass.utils.SmemAllocator() + scratch = smem.allocate_tensor(cutlass.Float32, warps) + local = cutlass.Float32(0.0) + for node in cutlass.range(tidx, node_count, threads, unroll=1): + local += params.grad_z_partial[node, head].to(cutlass.Float32) + total = _cta_sum(local, scratch, threads) + if tidx == 0: + params.grad_z[head] = total.to(params.grad_z.element_type) + + +@cute.jit +def neo_phase_c_backward_layout_jit( + grad_out: cute.Tensor, + stack: cute.Tensor, + wigner_dt: cute.Tensor, + alpha: cute.Tensor, + 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, + denom: cute.Tensor, + focus_src: cute.Tensor, + focus_weight: cute.Tensor, + focus_scale: cute.Tensor, + grad_stack: cute.Tensor, + grad_wigner_dt: cute.Tensor, + grad_logits: cute.Tensor, + grad_edge: cute.Tensor, + grad_z_partial: cute.Tensor, + grad_z: cute.Tensor, + grad_focus_src: cute.Tensor, + stream: CUstream, + focus_eps: cutlass.Constexpr[float], + focus_tau: cutlass.Constexpr[float], + focus_label_smoothing: cutlass.Constexpr[float], +): + params = NeoPhaseCBackwardLayoutParams( + grad_out=grad_out, + stack=stack, + wigner_dt=wigner_dt, + alpha=alpha, + 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, + denom=denom, + focus_src=focus_src, + focus_weight=focus_weight, + focus_scale=focus_scale, + grad_stack=grad_stack, + grad_wigner_dt=grad_wigner_dt, + grad_logits=grad_logits, + grad_edge=grad_edge, + grad_z_partial=grad_z_partial, + grad_z=grad_z, + grad_focus_src=grad_focus_src, + ) + copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + stack.element_type, + num_bits_per_copy=32, + ) + channel_thread_layout = cute.make_ordered_layout((1, 32), order=(1, 0)) + reduced_value_layout = cute.make_ordered_layout((10, 1), order=(1, 0)) + raw_tiled_copy = cute.make_tiled_copy_tv( + copy_atom, + channel_thread_layout, + reduced_value_layout, + ) + raw_layout = cute.make_layout((REDUCED_COUNT, FOCUS_CHANNELS), stride=(32, 1)) + node_count, _, _ = grad_out.shape + neo_phase_c_backward_layout_kernel( + params, + raw_layout, + raw_tiled_copy, + focus_eps, + focus_tau, + focus_label_smoothing, + ).launch( + grid=[node_count, 1, 1], + block=[HIDDEN, 1, 1], + stream=stream, + ) + neo_phase_c_backward_z_reduce_kernel(params, 128).launch( + grid=[N_FOCUS, 1, 1], + block=[128, 1, 1], + stream=stream, + ) + + +def compile_neo_phase_c_backward_layout( + *, + focus_eps: float, + focus_tau: float, + focus_label_smoothing: float, +) -> Callable: + """Compile the exact Neo Phase-C layout-boundary backward callable.""" + if focus_eps <= 0.0: + raise ValueError("focus_eps must be positive") + if focus_tau <= 0.0: + raise ValueError("focus_tau must be positive") + if not 0.0 <= focus_label_smoothing < 1.0: + raise ValueError("focus_label_smoothing must be in [0, 1)") + + edge_count = cute.sym_int64() + node_count = cute.sym_int64() + fake_grad_out = make_fake_compact_tensor( + cutlass.Float32, + (node_count, DEGREE_COUNT, HIDDEN), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + fake_stack = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, N_FOCUS, REDUCED_COUNT, FOCUS_CHANNELS), + stride_order=(3, 2, 1, 0), + **FAKE_TENSOR_KW, + ) + fake_wigner_dt = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, PACKED_WIGNER_VALUES), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_alpha = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, N_FOCUS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_focus_alpha = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, N_FOCUS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_dst_ptr = make_fake_compact_tensor( + cutlass.Int32, + (cute.sym_int64(),), + stride_order=(0,), + **FAKE_TENSOR_KW, + ) + fake_rotate = make_fake_compact_tensor( + cutlass.Float32, + (DEGREE_COUNT,), + 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,), + stride_order=(0,), + **FAKE_TENSOR_KW, + ) + fake_z_bias = make_fake_compact_tensor( + cutlass.Float32, + (N_FOCUS,), + stride_order=(0,), + **FAKE_TENSOR_KW, + ) + fake_group_max = make_fake_compact_tensor( + cutlass.Float32, + (node_count, N_FOCUS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_denom = make_fake_compact_tensor( + cutlass.Float32, + (node_count, N_FOCUS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_focus_src = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, N_FOCUS, FOCUS_CHANNELS), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + fake_focus_weight = make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_CHANNELS, N_FOCUS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_focus_scale = make_fake_compact_tensor( + cutlass.Float32, + (N_FOCUS, FOCUS_CHANNELS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_grad_stack = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, N_FOCUS, REDUCED_COUNT, FOCUS_CHANNELS), + stride_order=(3, 2, 1, 0), + **FAKE_TENSOR_KW, + ) + fake_grad_wigner_dt = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, PACKED_WIGNER_VALUES), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_grad_logits = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, N_FOCUS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_grad_edge = make_fake_compact_tensor( + cutlass.Float32, + (edge_count,), + stride_order=(0,), + **FAKE_TENSOR_KW, + ) + fake_grad_z_partial = make_fake_compact_tensor( + cutlass.Float32, + (node_count, N_FOCUS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_grad_z = make_fake_compact_tensor( + cutlass.Float32, + (N_FOCUS,), + stride_order=(0,), + **FAKE_TENSOR_KW, + ) + fake_grad_focus_src = make_fake_compact_tensor( + cutlass.Float32, + (N_FOCUS, edge_count, FOCUS_CHANNELS), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + neo_phase_c_backward_layout_jit, + fake_grad_out, + fake_stack, + fake_wigner_dt, + fake_alpha, + fake_focus_alpha, + fake_dst_ptr, + fake_rotate, + fake_logits, + fake_edge_gate, + fake_z_bias, + fake_group_max, + fake_denom, + fake_focus_src, + fake_focus_weight, + fake_focus_scale, + fake_grad_stack, + fake_grad_wigner_dt, + fake_grad_logits, + fake_grad_edge, + fake_grad_z_partial, + fake_grad_z, + fake_grad_focus_src, + fake_stream, + focus_eps, + focus_tau, + focus_label_smoothing, + options="--enable-tvm-ffi", + ) diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout_runner.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout_runner.py new file mode 100644 index 0000000000..5f016886b5 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout_runner.py @@ -0,0 +1,380 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Runtime contract for exact-shape Neo Phase-C backward.""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + Any, +) + +import torch + +from ..compile_cache import ( + device_aware_lru_cache, +) +from ..k1_wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES +from .cute_neo_phase_c_backward_layout import ( + DEGREE_COUNT, + FOCUS_CHANNELS, + HIDDEN, + N_FOCUS, + REDUCED_COUNT, + compile_neo_phase_c_backward_layout, +) + +REQUIRED_ALIGNMENT = 16 + + +@dataclass(frozen=True) +class NeoPhaseCBackwardLayoutOutputs: + """Caller-owned outputs and scratch for one fused Phase-C invocation.""" + + grad_stack: torch.Tensor + grad_wigner_dt: torch.Tensor + grad_logits: torch.Tensor + grad_edge: torch.Tensor + grad_z_partial: torch.Tensor + grad_z: torch.Tensor + grad_focus_src: torch.Tensor + + +def _storage_id(tensor: torch.Tensor) -> int: + return tensor.untyped_storage()._cdata + + +def _is_exact_view(lhs: torch.Tensor, rhs: torch.Tensor) -> bool: + """Return whether two tensors name the same logical and physical view.""" + return ( + _storage_id(lhs) == _storage_id(rhs) + and lhs.data_ptr() == rhs.data_ptr() + and lhs.storage_offset() == rhs.storage_offset() + and lhs.shape == rhs.shape + and lhs.stride() == rhs.stride() + and lhs.dtype == rhs.dtype + and lhs.device == rhs.device + ) + + +def _tensor_byte_region(tensor: torch.Tensor) -> tuple[int, int] | None: + """Return the physical byte range of a compact runtime tensor.""" + if tensor.numel() == 0 or tensor.device.type == "meta": + return None + start = tensor.data_ptr() + return start, start + tensor.numel() * tensor.element_size() + + +def _tensor_regions_overlap( + lhs: torch.Tensor, + lhs_region: tuple[int, int] | None, + rhs: torch.Tensor, + rhs_region: tuple[int, int] | None, +) -> bool: + """Compare precomputed physical regions, including external storage views.""" + if lhs.device != rhs.device: + return False + if lhs.device.type == "meta": + return torch._C._overlaps(lhs, rhs) + if lhs_region is None or rhs_region is None: + return False + lhs_start, lhs_stop = lhs_region + rhs_start, rhs_stop = rhs_region + return lhs_start < rhs_stop and rhs_start < lhs_stop + + +def _require_alignment( + name: str, + tensor: torch.Tensor, + alignment: int = REQUIRED_ALIGNMENT, +) -> None: + """Enforce the alignment promised to CuTe by ``assumed_align``.""" + if tensor.device.type == "meta": + return + + byte_offset = tensor.storage_offset() * tensor.element_size() + pointer_remainder = tensor.data_ptr() % alignment + storage_remainder = tensor.untyped_storage().data_ptr() % alignment + offset_remainder = byte_offset % alignment + if pointer_remainder or storage_remainder or offset_remainder: + raise ValueError( + f"{name} must be {alignment}-byte aligned; got data pointer " + f"remainder {pointer_remainder}, storage pointer remainder " + f"{storage_remainder}, and byte storage-offset remainder " + f"{offset_remainder}" + ) + + +def _require_tensor( + name: str, + tensor: torch.Tensor, + shape: tuple[int, ...], + *, + device: torch.device, + dtype: torch.dtype = torch.float32, +) -> None: + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}") + if tensor.dtype != dtype: + raise ValueError(f"{name} must have dtype {dtype}, got {tensor.dtype}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be compact") + _require_alignment(name, tensor) + + +@device_aware_lru_cache(maxsize=32) +def _compile_layout_boundary( + focus_eps: float, + focus_tau: float, + focus_label_smoothing: float, +) -> Any: + return compile_neo_phase_c_backward_layout( + focus_eps=focus_eps, + focus_tau=focus_tau, + focus_label_smoothing=focus_label_smoothing, + ) + + +class CuteNeoPhaseCBackwardLayout: + """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 + is the unmodified ``grad_out`` input and remains owned by the caller. + + Inputs retain their model-native ranks. ``grad_stack`` aliases compact + edge-major ``stack`` storage and is written only after the complete source + fragment has been consumed. ``grad_focus_src`` is compact ``(F,E,C)``. + """ + + def __init__( + self, + *, + focus_eps: float, + focus_tau: float, + focus_label_smoothing: float, + ) -> None: + self._compiled = _compile_layout_boundary( + float(focus_eps), + float(focus_tau), + float(focus_label_smoothing), + ) + + def __call__( + self, + grad_out: torch.Tensor, + stack: torch.Tensor, + wigner_dt: torch.Tensor, + alpha: torch.Tensor, + 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, + denom: torch.Tensor, + focus_src: torch.Tensor, + focus_weight: torch.Tensor, + focus_scale: torch.Tensor, + outputs: NeoPhaseCBackwardLayoutOutputs, + ) -> NeoPhaseCBackwardLayoutOutputs: + edge_count = stack.shape[0] + node_count = grad_out.shape[0] + device = stack.device + stack_shape = ( + edge_count, + N_FOCUS, + REDUCED_COUNT, + FOCUS_CHANNELS, + ) + + _require_tensor( + "grad_out", + grad_out, + (node_count, DEGREE_COUNT, HIDDEN), + device=device, + ) + _require_tensor("stack", stack, stack_shape, device=device) + _require_tensor( + "wigner_dt", + wigner_dt, + (edge_count, PACKED_WIGNER_VALUES), + device=device, + ) + _require_tensor("alpha", alpha, (edge_count, N_FOCUS), device=device) + _require_tensor( + "focus_alpha", focus_alpha, (edge_count, N_FOCUS), device=device + ) + _require_tensor( + "dst_ptr", + dst_ptr, + (node_count + 1,), + device=device, + dtype=torch.int32, + ) + _require_tensor( + "rotate_inv_rescale", + rotate_inv_rescale, + (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) + _require_tensor("denom", denom, (node_count, N_FOCUS), device=device) + _require_tensor( + "focus_src", + focus_src, + (edge_count, N_FOCUS, FOCUS_CHANNELS), + device=device, + ) + _require_tensor( + "focus_weight", + focus_weight, + (FOCUS_CHANNELS, N_FOCUS), + device=device, + ) + _require_tensor( + "focus_scale", + focus_scale, + (N_FOCUS, FOCUS_CHANNELS), + device=device, + ) + + _require_tensor( + "outputs.grad_stack", + outputs.grad_stack, + stack_shape, + device=device, + ) + if not _is_exact_view(outputs.grad_stack, stack): + raise ValueError("outputs.grad_stack must be the exact in-place stack view") + _require_tensor( + "outputs.grad_wigner_dt", + outputs.grad_wigner_dt, + (edge_count, PACKED_WIGNER_VALUES), + device=device, + ) + _require_tensor( + "outputs.grad_logits", + outputs.grad_logits, + (edge_count, N_FOCUS), + device=device, + ) + _require_tensor( + "outputs.grad_edge", outputs.grad_edge, (edge_count,), device=device + ) + _require_tensor( + "outputs.grad_z_partial", + outputs.grad_z_partial, + (node_count, N_FOCUS), + device=device, + ) + _require_tensor("outputs.grad_z", outputs.grad_z, (N_FOCUS,), device=device) + _require_tensor( + "outputs.grad_focus_src", + outputs.grad_focus_src, + (N_FOCUS, edge_count, FOCUS_CHANNELS), + device=device, + ) + + input_tensors = tuple( + (name, tensor, _tensor_byte_region(tensor)) + for name, tensor in ( + ("grad_out", grad_out), + ("stack", stack), + ("wigner_dt", wigner_dt), + ("alpha", alpha), + ("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), + ("denom", denom), + ("focus_src", focus_src), + ("focus_weight", focus_weight), + ("focus_scale", focus_scale), + ) + ) + output_tensors = tuple( + ( + f"outputs.{field_name}", + getattr(outputs, field_name), + _tensor_byte_region(getattr(outputs, field_name)), + ) + for field_name in ( + "grad_stack", + "grad_wigner_dt", + "grad_logits", + "grad_edge", + "grad_z_partial", + "grad_z", + "grad_focus_src", + ) + ) + + for output_index, ( + output_name, + output, + output_region, + ) in enumerate(output_tensors): + for other_name, other, other_region in output_tensors[output_index + 1 :]: + if _tensor_regions_overlap( + output, + output_region, + other, + other_region, + ): + raise ValueError( + f"{output_name} must not overlap output {other_name}" + ) + for input_name, input_tensor, input_region in input_tensors: + if not _tensor_regions_overlap( + output, + output_region, + input_tensor, + input_region, + ): + continue + is_stack_adjoint = ( + output_name == "outputs.grad_stack" and input_name == "stack" + ) + if is_stack_adjoint: + continue + raise ValueError(f"{output_name} must not overlap input {input_name}") + + self._compiled( + grad_out, + stack, + wigner_dt, + alpha, + focus_alpha, + dst_ptr, + rotate_inv_rescale, + logits, + edge_gate, + z_bias_raw, + group_max, + denom, + focus_src, + focus_weight, + focus_scale, + outputs.grad_stack, + outputs.grad_wigner_dt, + outputs.grad_logits, + outputs.grad_edge, + outputs.grad_z_partial, + outputs.grad_z, + outputs.grad_focus_src, + ) + return outputs diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_onepass.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_onepass.py new file mode 100644 index 0000000000..d366e3d1ab --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_onepass.py @@ -0,0 +1,548 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Strict-FP32 one-pass Neo Phase-C forward with output gating. + +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 +cooperatively staged through a double-buffered shared-memory panel, so no +``node * chunks`` partial tensor or second reduction launch is required. + +The final stores fuse the output-side attention gate: + +``sigmoid(project(RMSNorm(x_wide[:, 0]))) * rotate_inv_rescale * aggregate``. + +Dense Wigner, packed Wigner through the generic loader, and the 46-value +packed-direct path share the same runtime tensor signature. Every floating +tensor and every arithmetic operation is FP32 by construction. +""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, + Any, +) + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from ..compile_cache import ( + device_aware_lru_cache, +) +from ..k1_wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN001, ANN201, ANN202, TC002 + +DEGREE_COUNT = 16 +REDUCED_COUNT = 10 +FOCUS_COUNT = 2 +CHANNELS = 32 +HIDDEN = FOCUS_COUNT * CHANNELS +PHASE_WIDTH = REDUCED_COUNT * HIDDEN +OUTPUT_WIDTH = DEGREE_COUNT * HIDDEN +THREADS = HIDDEN +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + + +@dataclass(frozen=True) +class NeoPhaseCOnePassParams: + """Runtime tensors for the warp-private packed kernel.""" + + x_local: cute.Tensor + wigner_dt: cute.Tensor + alpha: cute.Tensor + focus_alpha: cute.Tensor + out: cute.Tensor + x_wide: cute.Tensor + norm_scale: cute.Tensor + gate_weight: cute.Tensor + dst_ptr: cute.Tensor + rotate_inv_rescale: cute.Tensor + + +@cute.jit +def _sigmoid(value): + one = cutlass.Float32(1.0) + return one / (one + cute.exp(-value)) + + +@cute.jit +def _weighted_input( + params: NeoPhaseCOnePassParams, + edge, + tid, + focus, + channel, + reduced: cutlass.Constexpr[int], + alpha, + focus_scale, +): + """Load one reduced coefficient and apply both edge attention weights.""" + load_idx = focus * REDUCED_COUNT * CHANNELS + reduced * CHANNELS + channel + value = params.x_local[edge, load_idx].to(cutlass.Float32) + # Match the two-launch path's FP32 association exactly. + return value * alpha * focus_scale + + +@cute.jit +def _focus_scale( + params: NeoPhaseCOnePassParams, + edge, + focus, +): + return params.focus_alpha[edge, focus].to(cutlass.Float32) + + +@cute.jit +def _output_gate( + params: NeoPhaseCOnePassParams, + node, + tid, + focus, + channel, + eps: cutlass.Constexpr[float], +): + """Compute one gate per focus; each warp owns one 32-channel focus.""" + x = params.x_wide[node, tid].to(cutlass.Float32) + square_sum = cute.arch.warp_reduction_sum(x * x) + inv_rms = cute.rsqrt(square_sum / cutlass.Float32(CHANNELS) + cutlass.Float32(eps)) + logit_part = ( + x + * inv_rms + * params.norm_scale[focus, channel].to(cutlass.Float32) + * params.gate_weight[channel, focus, 0].to(cutlass.Float32) + ) + logit = cute.arch.warp_reduction_sum(logit_part) + return _sigmoid(logit) + + +@cute.jit +def _store_gated_output( + params: NeoPhaseCOnePassParams, + accumulator: cute.Tensor, + node, + tid, + gate, +): + for degree in cutlass.range_constexpr(DEGREE_COUNT): + value = accumulator[degree] + value *= params.rotate_inv_rescale[degree].to(cutlass.Float32) + value *= gate + params.out[node, degree * HIDDEN + tid] = value + + +@cute.jit +def _warp_private_packed_wigner( + panel, + focus, + panel_index: cutlass.Constexpr[int], +): + """Load one scalar from a focus warp's private shared Wigner panel.""" + return panel[focus, panel_index] + + +@cute.jit +def neo_phase_c_onepass_output_gate_packed_direct_warp_private_jit( + x_local: cute.Tensor, + wigner_dt: cute.Tensor, + alpha: cute.Tensor, + focus_alpha: cute.Tensor, + out: cute.Tensor, + x_wide: cute.Tensor, + norm_scale: cute.Tensor, + gate_weight: cute.Tensor, + dst_ptr: cute.Tensor, + rotate_inv_rescale: cute.Tensor, + eps: cutlass.Constexpr[float], + stream: CUstream, +): + """Launch the packed-direct kernel with a warp-local Wigner epilogue.""" + params = NeoPhaseCOnePassParams( + x_local=x_local, + wigner_dt=wigner_dt, + alpha=alpha, + focus_alpha=focus_alpha, + out=out, + x_wide=x_wide, + norm_scale=norm_scale, + gate_weight=gate_weight, + dst_ptr=dst_ptr, + rotate_inv_rescale=rotate_inv_rescale, + ) + accumulator_layout = cute.make_layout((DEGREE_COUNT,), stride=(1,)) + node_count, _ = out.shape + neo_phase_c_onepass_output_gate_packed_direct_warp_private_kernel( + params, + accumulator_layout, + eps, + ).launch( + grid=[node_count, 1, 1], + block=[THREADS, 1, 1], + stream=stream, + ) + + +@cute.kernel +def neo_phase_c_onepass_output_gate_packed_direct_warp_private_kernel( + params: NeoPhaseCOnePassParams, + accumulator_layout: cute.Layout, + eps: cutlass.Constexpr[float], +): + """Store the gated output without a Wigner shared-memory round trip.""" + tid, _, _ = cute.arch.thread_idx() + node, _, _ = cute.arch.block_idx() + focus = tid // CHANNELS + channel = tid - focus * CHANNELS + lane = cute.arch.lane_idx() + + # Each focus is exactly one warp. The gate reduction already broadcasts + # its result within that warp, so it does not need shared memory either. + smem = cutlass.utils.SmemAllocator() + panel_storage = smem.allocate_tensor( + cutlass.Float32, + FOCUS_COUNT * PACKED_WIGNER_VALUES, + ) + panel_layout = cute.make_layout( + (FOCUS_COUNT, PACKED_WIGNER_VALUES), + stride=(PACKED_WIGNER_VALUES, 1), + ) + panel = cute.make_tensor(panel_storage.iterator, panel_layout) + gate = _output_gate(params, node, tid, focus, channel, eps) + accumulator = cute.make_rmem_tensor(accumulator_layout, cutlass.Float32) + accumulator.fill(0.0) + + lo = params.dst_ptr[node] + hi = params.dst_ptr[node + 1] + for edge in cutlass.range(lo, hi, 1, unroll=1): + panel[focus, lane] = params.wigner_dt[edge, lane].to(cutlass.Float32) + if lane < PACKED_WIGNER_VALUES - CHANNELS: + panel[focus, lane + CHANNELS] = params.wigner_dt[edge, lane + CHANNELS].to( + cutlass.Float32 + ) + cute.arch.sync_warp() + + alpha = params.alpha[edge, focus].to(cutlass.Float32) + focus_scale = _focus_scale(params, edge, focus) + value0 = _weighted_input( + params, + edge, + tid, + focus, + channel, + 0, + alpha, + focus_scale, + ) + accumulator[0] += _warp_private_packed_wigner(panel, focus, 0) * value0 + + for row_slot in cutlass.range_constexpr(3): + value1 = _weighted_input( + params, + edge, + tid, + focus, + channel, + 1 + row_slot * 3, + alpha, + focus_scale, + ) + panel_start1 = 1 + row_slot * 3 + for local_row in cutlass.range_constexpr(3): + accumulator[1 + local_row] += ( + _warp_private_packed_wigner( + panel, + focus, + panel_start1 + local_row, + ) + * value1 + ) + + for row_slot in cutlass.range_constexpr(3): + value2 = _weighted_input( + params, + edge, + tid, + focus, + channel, + 2 + row_slot * 3, + alpha, + focus_scale, + ) + panel_start2 = 10 + row_slot * 5 + for local_row in cutlass.range_constexpr(5): + accumulator[4 + local_row] += ( + _warp_private_packed_wigner( + panel, + focus, + panel_start2 + local_row, + ) + * value2 + ) + + for row_slot in cutlass.range_constexpr(3): + value3 = _weighted_input( + params, + edge, + tid, + focus, + channel, + 3 + row_slot * 3, + alpha, + focus_scale, + ) + panel_start3 = 25 + row_slot * 7 + for local_row in cutlass.range_constexpr(7): + accumulator[9 + local_row] += ( + _warp_private_packed_wigner( + panel, + focus, + panel_start3 + local_row, + ) + * value3 + ) + + # Do not overwrite this warp's panel until every lane has consumed it. + cute.arch.sync_warp() + + _store_gated_output(params, accumulator, node, tid, gate) + + +def _fake_common_tensors(): + edges = cute.sym_int64() + nodes = cute.sym_int64() + x_local = make_fake_compact_tensor( + cutlass.Float32, + (edges, PHASE_WIDTH), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + wigner_dt = make_fake_compact_tensor( + cutlass.Float32, + (edges, PACKED_WIGNER_VALUES), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + alpha = make_fake_compact_tensor( + cutlass.Float32, + (edges, FOCUS_COUNT), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + focus_alpha = make_fake_compact_tensor( + cutlass.Float32, + (edges, FOCUS_COUNT), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + out = make_fake_compact_tensor( + cutlass.Float32, + (nodes, OUTPUT_WIDTH), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + x_wide = make_fake_compact_tensor( + cutlass.Float32, + (nodes, OUTPUT_WIDTH), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + norm_scale = make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, CHANNELS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + gate_weight = make_fake_compact_tensor( + cutlass.Float32, + (CHANNELS, FOCUS_COUNT, 1), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + dst_ptr = make_fake_compact_tensor( + cutlass.Int32, + (cute.sym_int64(),), + stride_order=(0,), + **FAKE_TENSOR_KW, + ) + rotate_inv_rescale = make_fake_compact_tensor( + cutlass.Float32, + (DEGREE_COUNT,), + stride_order=(0,), + **FAKE_TENSOR_KW, + ) + return ( + x_local, + wigner_dt, + alpha, + focus_alpha, + out, + x_wide, + norm_scale, + gate_weight, + dst_ptr, + rotate_inv_rescale, + ) + + +def compile_neo_phase_c_onepass_output_gate( + eps: float, +) -> Callable: + """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() + stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + neo_phase_c_onepass_output_gate_packed_direct_warp_private_jit, + *common_args, + eps, + stream, + options="--enable-tvm-ffi", + ) + + +@device_aware_lru_cache(maxsize=8) +def _compiled_neo_phase_c_onepass_output_gate(eps: float) -> Callable: + return compile_neo_phase_c_onepass_output_gate(eps) + + +def _expect_shape(name: str, tensor: Any, expected: tuple[int, ...]) -> None: + actual = tuple(tensor.shape) + if actual != expected: + raise ValueError(f"expected {name} shape {expected}, got {actual}") + + +def _expect_fp32_cuda(name: str, tensor: Any, *, torch: Any, device: Any) -> None: + if tensor.dtype != torch.float32: + raise TypeError(f"{name} must be strict float32, got {tensor.dtype}") + if not tensor.is_cuda: + raise ValueError(f"{name} must be a CUDA tensor") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}") + + +def run_neo_phase_c_onepass_output_gate( + *, + x_local_flat: Any, + Dt_full: Any, + alpha_focus: Any, + focus_compete_alpha: Any, + dst_ptr: Any, + rotate_inv_rescale: Any, + x_wide: Any, + output_gate_norm_scale: Any, + output_gate_weight: Any, + output_gate_eps: float, + out: Any | None = None, +) -> Any: + """Validate and launch fused Phase C and output gating.""" + import torch + + if not x_local_flat.is_cuda: + raise ValueError("x_local_flat must be a CUDA tensor") + + device = x_local_flat.device + edge_count = x_local_flat.shape[0] + node_count = x_wide.shape[0] + _expect_shape( + "x_local_flat", + x_local_flat, + (edge_count, FOCUS_COUNT, REDUCED_COUNT, CHANNELS), + ) + _expect_shape("x_wide", x_wide, (node_count, DEGREE_COUNT, HIDDEN)) + _expect_shape("alpha_focus", alpha_focus, (edge_count, FOCUS_COUNT)) + _expect_shape( + "output_gate_norm_scale", + output_gate_norm_scale, + (FOCUS_COUNT, CHANNELS), + ) + _expect_shape( + "output_gate_weight", + output_gate_weight, + (CHANNELS, FOCUS_COUNT, 1), + ) + _expect_shape("rotate_inv_rescale", rotate_inv_rescale, (DEGREE_COUNT,)) + _expect_shape("dst_ptr", dst_ptr, (node_count + 1,)) + _expect_shape( + "focus_compete_alpha", + focus_compete_alpha, + (edge_count, FOCUS_COUNT), + ) + _expect_shape("Dt_full", Dt_full, (edge_count, PACKED_WIGNER_VALUES)) + + floating_tensors = { + "x_local_flat": x_local_flat, + "Dt_full": Dt_full, + "alpha_focus": alpha_focus, + "focus_compete_alpha": focus_compete_alpha, + "rotate_inv_rescale": rotate_inv_rescale, + "x_wide": x_wide, + "output_gate_norm_scale": output_gate_norm_scale, + "output_gate_weight": output_gate_weight, + } + for name, tensor in floating_tensors.items(): + _expect_fp32_cuda(name, tensor, torch=torch, device=device) + + if dst_ptr.device != device: + 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}") + + if out is None: + out = torch.empty( + node_count, + DEGREE_COUNT, + HIDDEN, + device=device, + dtype=torch.float32, + ) + else: + _expect_shape("out", out, (node_count, DEGREE_COUNT, HIDDEN)) + _expect_fp32_cuda("out", out, torch=torch, device=device) + if not out.is_contiguous(): + raise ValueError("out must be contiguous") + + with torch.cuda.device(device): + kernel = _compiled_neo_phase_c_onepass_output_gate(float(output_gate_eps)) + kernel( + x_local_flat.contiguous().view(edge_count, REDUCED_COUNT * HIDDEN), + Dt_full.contiguous(), + alpha_focus.contiguous(), + focus_compete_alpha.contiguous(), + out.view(node_count, DEGREE_COUNT * HIDDEN), + x_wide.contiguous().view(node_count, DEGREE_COUNT * HIDDEN), + output_gate_norm_scale.contiguous(), + output_gate_weight.contiguous(), + dst_ptr.to(dtype=torch.int32).contiguous(), + rotate_inv_rescale.contiguous(), + ) + return out + + +__all__ = [ + "compile_neo_phase_c_onepass_output_gate", + "neo_phase_c_onepass_output_gate_packed_direct_warp_private_jit", + "neo_phase_c_onepass_output_gate_packed_direct_warp_private_kernel", + "run_neo_phase_c_onepass_output_gate", +] diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_qk_edge.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_qk_edge.py new file mode 100644 index 0000000000..d3ccbe8ad7 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_qk_edge.py @@ -0,0 +1,364 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +# ruff: noqa: ANN001, ANN201, ANN202, TC002, UP035 +"""Fused Neo Q/K edge logits and first-backward input adjoints.""" + +from __future__ import ( + annotations, +) + +from functools import ( + lru_cache, +) +from typing import ( + Callable, +) + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + + +@cute.kernel +def neo_qk_edge_forward_kernel( + q_node: cute.Tensor, + k_node: cute.Tensor, + radial_l0: cute.Tensor, + attention_weight: cute.Tensor, + src: cute.Tensor, + dst: cute.Tensor, + logits: cute.Tensor, + scale: cutlass.Constexpr[float], +): + tidx, _, _ = cute.arch.thread_idx() + block, focus, _ = cute.arch.block_idx() + edge = block * 128 + tidx + edges, _ = logits.shape + if edge < edges: + src_node = src[edge] + dst_node = dst[edge] + qk_acc = cutlass.Float32(0.0) + radial_acc = cutlass.Float32(0.0) + for channel in cutlass.range_constexpr(32): + qk_acc += q_node[dst_node, focus, channel].to(cutlass.Float32) * k_node[ + src_node, focus, channel + ].to(cutlass.Float32) + radial_acc += radial_l0[edge, focus, channel].to( + cutlass.Float32 + ) * attention_weight[channel, focus, 0].to(cutlass.Float32) + logits[edge, focus] = (qk_acc * cutlass.Float32(scale) + radial_acc).to( + logits.element_type + ) + + +@cute.jit +def neo_qk_edge_forward_jit( + q_node: cute.Tensor, + k_node: cute.Tensor, + radial_l0: cute.Tensor, + attention_weight: cute.Tensor, + src: cute.Tensor, + dst: cute.Tensor, + logits: cute.Tensor, + stream: CUstream, + scale: cutlass.Constexpr[float], +): + edges, _ = logits.shape + neo_qk_edge_forward_kernel( + q_node, + k_node, + radial_l0, + attention_weight, + src, + dst, + logits, + scale, + ).launch( + grid=[cute.ceil_div(edges, 128), 2, 1], + block=[128, 1, 1], + stream=stream, + ) + + +@cute.kernel +def neo_qk_edge_backward_kernel( + grad_logits: cute.Tensor, + q_node: cute.Tensor, + k_node: cute.Tensor, + src: cute.Tensor, + dst: cute.Tensor, + grad_q_node: cute.Tensor, + grad_k_node: cute.Tensor, + scale: cutlass.Constexpr[float], +): + tidx, _, _ = cute.arch.thread_idx() + block, focus, _ = cute.arch.block_idx() + edge = block * 8 + tidx // 32 + channel = tidx % 32 + edges, _ = grad_logits.shape + if edge < edges: + src_node = src[edge] + dst_node = dst[edge] + grad = grad_logits[edge, focus].to(cutlass.Float32) * cutlass.Float32(scale) + grad_q = grad * k_node[src_node, focus, channel].to(cutlass.Float32) + grad_k = grad * q_node[dst_node, focus, channel].to(cutlass.Float32) + q_offset = (dst_node * 2 + focus) * 32 + channel + k_offset = (src_node * 2 + focus) * 32 + channel + q_ptr = grad_q_node.iterator + q_offset + k_ptr = grad_k_node.iterator + k_offset + cute.arch.atomic_add(q_ptr.llvm_ptr, grad_q, sem="relaxed", scope="gpu") + cute.arch.atomic_add(k_ptr.llvm_ptr, grad_k, sem="relaxed", scope="gpu") + + +@cute.jit +def neo_qk_edge_backward_jit( + grad_logits: cute.Tensor, + q_node: cute.Tensor, + k_node: cute.Tensor, + src: cute.Tensor, + dst: cute.Tensor, + grad_q_node: cute.Tensor, + grad_k_node: cute.Tensor, + stream: CUstream, + scale: cutlass.Constexpr[float], +): + edges, _ = grad_logits.shape + neo_qk_edge_backward_kernel( + grad_logits, + q_node, + k_node, + src, + dst, + grad_q_node, + grad_k_node, + scale, + ).launch( + grid=[cute.ceil_div(edges, 8), 2, 1], + block=[256, 1, 1], + stream=stream, + ) + + +@cute.kernel +def neo_qk_node_input_adjoint_kernel( + x_l0: cute.Tensor, + grad_q_node: cute.Tensor, + grad_k_node: cute.Tensor, + q_weight: cute.Tensor, + k_weight: cute.Tensor, + norm_scale: cute.Tensor, + grad_x_wide: cute.Tensor, + eps: cutlass.Float32, +): + tid, _, _ = cute.arch.thread_idx() + block, _, _ = cute.arch.block_idx() + node_in_block = tid // 64 + local_tid = tid % 64 + focus = local_tid // 32 + channel = local_tid % 32 + node = block * 4 + node_in_block + nodes, _, _ = x_l0.shape + + if node < nodes: + for flat_index in cutlass.range(local_tid, 16 * 64, 64): + grad_x_wide[node, flat_index] = cutlass.Float32(0.0) + + grad_norm = cutlass.Float32(0.0) + for output_channel in cutlass.range_constexpr(32): + grad_norm += grad_q_node[node, focus, output_channel].to( + cutlass.Float32 + ) * q_weight[channel, focus, output_channel].to(cutlass.Float32) + grad_norm += grad_k_node[node, focus, output_channel].to( + cutlass.Float32 + ) * k_weight[channel, focus, output_channel].to(cutlass.Float32) + + x = x_l0[node, focus, channel].to(cutlass.Float32) + grad_scaled = grad_norm * norm_scale[focus, channel].to(cutlass.Float32) + inv = cute.rsqrt( + cute.arch.warp_reduction_sum(x * x) / cutlass.Float32(32.0) + eps + ) + coeff = cute.arch.warp_reduction_sum(grad_scaled * x) / cutlass.Float32(32.0) + grad_x = grad_scaled * inv - x * inv * inv * inv * coeff + grad_x_wide[node, focus * 32 + channel] = grad_x.to(grad_x_wide.element_type) + + +@cute.jit +def neo_qk_node_input_adjoint_jit( + x_l0: cute.Tensor, + grad_q_node: cute.Tensor, + grad_k_node: cute.Tensor, + q_weight: cute.Tensor, + k_weight: cute.Tensor, + norm_scale: cute.Tensor, + grad_x_wide: cute.Tensor, + stream: CUstream, + eps: cutlass.Float32, +): + nodes, _, _ = x_l0.shape + neo_qk_node_input_adjoint_kernel( + x_l0, + grad_q_node, + grad_k_node, + q_weight, + k_weight, + norm_scale, + grad_x_wide, + eps, + ).launch( + grid=[cute.ceil_div(nodes, 4), 1, 1], + block=[256, 1, 1], + stream=stream, + ) + + +def _fake_inputs(): + edge_count = cute.sym_int64() + node_count = cute.sym_int64() + q_node = make_fake_compact_tensor( + cutlass.Float32, (node_count, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + k_node = make_fake_compact_tensor( + cutlass.Float32, (node_count, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + radial = make_fake_compact_tensor( + cutlass.Float32, (edge_count, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + weight = make_fake_compact_tensor( + cutlass.Float32, (32, 2, 1), stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + src = make_fake_compact_tensor( + cutlass.Int32, (edge_count,), stride_order=(0,), **FAKE_TENSOR_KW + ) + dst = make_fake_compact_tensor( + cutlass.Int32, (edge_count,), stride_order=(0,), **FAKE_TENSOR_KW + ) + logits = make_fake_compact_tensor( + cutlass.Float32, (edge_count, 2), stride_order=(1, 0), **FAKE_TENSOR_KW + ) + return q_node, k_node, radial, weight, src, dst, logits + + +@lru_cache(maxsize=8) +def compile_neo_qk_edge_forward( + scale: float, + compile_identity: tuple[int, int, int] | None = None, +) -> Callable: + del compile_identity + q_node, k_node, radial, weight, src, dst, logits = _fake_inputs() + stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + neo_qk_edge_forward_jit, + q_node, + k_node, + radial, + weight, + src, + dst, + logits, + stream, + scale, + options="--enable-tvm-ffi", + ) + + +@lru_cache(maxsize=8) +def compile_neo_qk_edge_backward( + scale: float, + compile_identity: tuple[int, int, int] | None = None, +) -> Callable: + del compile_identity + q_node, k_node, _radial, _weight, src, dst, logits = _fake_inputs() + grad_q = make_fake_compact_tensor( + cutlass.Float32, q_node.shape, stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + grad_k = make_fake_compact_tensor( + cutlass.Float32, k_node.shape, stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + neo_qk_edge_backward_jit, + logits, + q_node, + k_node, + src, + dst, + grad_q, + grad_k, + stream, + scale, + options="--enable-tvm-ffi", + ) + + +@lru_cache(maxsize=8) +def compile_neo_qk_node_input_adjoint( + eps: float, + compile_identity: tuple[int, int, int] | None = None, +) -> Callable: + del compile_identity + node_count = cute.sym_int64() + x_l0 = make_fake_compact_tensor( + cutlass.Float32, (node_count, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + grad_q = make_fake_compact_tensor( + cutlass.Float32, x_l0.shape, stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + grad_k = make_fake_compact_tensor( + cutlass.Float32, x_l0.shape, stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + q_weight = make_fake_compact_tensor( + cutlass.Float32, (32, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + k_weight = make_fake_compact_tensor( + cutlass.Float32, (32, 2, 32), stride_order=(2, 1, 0), **FAKE_TENSOR_KW + ) + norm_scale = make_fake_compact_tensor( + cutlass.Float32, (2, 32), stride_order=(1, 0), **FAKE_TENSOR_KW + ) + grad_x_wide = make_fake_compact_tensor( + cutlass.Float32, (node_count, 16 * 64), stride_order=(1, 0), **FAKE_TENSOR_KW + ) + stream = make_fake_stream(use_tvm_ffi_env_stream=True) + compiled = cute.compile( + neo_qk_node_input_adjoint_jit, + x_l0, + grad_q, + grad_k, + q_weight, + k_weight, + norm_scale, + grad_x_wide, + stream, + cutlass.Float32(eps), + options="--enable-tvm-ffi", + ) + + def run( + x_l0_tensor, + grad_q_tensor, + grad_k_tensor, + q_weight_tensor, + k_weight_tensor, + norm_scale_tensor, + grad_x_wide_tensor, + ): + return compiled( + x_l0_tensor, + grad_q_tensor, + grad_k_tensor, + q_weight_tensor, + k_weight_tensor, + norm_scale_tensor, + grad_x_wide_tensor, + cutlass.Float32(eps), + ) + + return run diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_radial_phase_a_backward_node.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_radial_phase_a_backward_node.py new file mode 100644 index 0000000000..603e330fb2 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_radial_phase_a_backward_node.py @@ -0,0 +1,720 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Source-CSR node-tiled Neo radial and Phase-A backward. + +This is the contention-free node-tiled implementation for the exact Neo K1 shape +``lmax=3, D=16, Dm=10, Cwide=64, F=2``. Edge tensors retain their physical +destination-sorted order. ``source_ptr`` delimits intervals in +``source_order``, whose slots hold the corresponding physical edge ids. + +One 64-thread CTA owns one source node. It keeps the node feature row in shared +memory and 16 adjoint values per thread in registers, processes all incident +edges, writes the per-edge radial/Wigner adjoints directly, then writes the +node adjoint once. There are no global atomics and no ``(E, 16, 64)`` +reduction intermediate. +""" + +# ruff: noqa: ANN001, ANN201, ANN202, TC002, UP035 + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + Callable, +) + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from ..k1_wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES + +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + +DEGREE_COUNT = 16 +REDUCED_COUNT = 10 +HIDDEN = 64 +FOCUS_COUNT = 2 +FOCUS_HIDDEN = 32 +FOCUS_ROW = REDUCED_COUNT * FOCUS_HIDDEN +RADIAL_WIDTH = 4 * FOCUS_HIDDEN +COMPACT_WIDTH = 25 +SHARED_ROW_PITCH = HIDDEN + 4 +WARP_REDUCTION_GROUP = 4 +GROUPS_PER_WARP = 32 // WARP_REDUCTION_GROUP +GROUPS_PER_CTA = 2 * GROUPS_PER_WARP +CHANNELS_PER_SUBGROUP_LANE = HIDDEN // WARP_REDUCTION_GROUP + + +@dataclass(frozen=True) +class NeoRadialPhaseABackwardNodeParams: + grad_out_focus: cute.Tensor + grad_focus_src: cute.Tensor + grad_logits: cute.Tensor + radial_state: cute.Tensor + channel_basis: cute.Tensor + x_wide: cute.Tensor + source_order: cute.Tensor + source_ptr: cute.Tensor + d_full: cute.Tensor + grad_x_wide: cute.Tensor + grad_d_full: cute.Tensor + + +@cute.jit +def _focus_grad_value( + grad_out_focus: cute.Tensor, + grad_focus_src: cute.Tensor, + edge, + coeff: cutlass.Constexpr[int], + channel, +): + focus = channel // FOCUS_HIDDEN + focus_channel = channel - focus * FOCUS_HIDDEN + offset = focus * FOCUS_ROW + coeff * FOCUS_HIDDEN + focus_channel + value = grad_out_focus[edge, offset].to(cutlass.Float32) + if cutlass.const_expr(coeff == 0): + value += grad_focus_src[focus, edge, focus_channel].to(cutlass.Float32) + return value + + +@cute.jit +def _recompute_local_value( + local_values, + x_values, + d_values, + channel, + reduced: cutlass.Constexpr[int], + panel_start: cutlass.Constexpr[int], + full_start: cutlass.Constexpr[int], + width: cutlass.Constexpr[int], + shared_row_pitch: cutlass.Constexpr[int], + x_row_pitch: cutlass.Constexpr[int], +): + acc = cutlass.Float32(0.0) + for local_col in cutlass.range_constexpr(width): + acc += ( + d_values[panel_start + local_col] + * x_values[(full_start + local_col) * x_row_pitch + channel] + ) + local_values[reduced * shared_row_pitch + channel] = acc + + +@cute.jit +def _warp_owned_grad_compact( + focus_grad, + local_values, + channel_basis: cute.Tensor, + compact_idx, + subgroup_lane, + 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) + 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: + value += ( + focus_grad[out_coeff * shared_row_pitch + hidden_channel] + * local_values[in_coeff * shared_row_pitch + hidden_channel] + * basis + ) + else: + value += basis * ( + focus_grad[(4 + out_coeff) * shared_row_pitch + hidden_channel] + * local_values[(4 + in_coeff) * shared_row_pitch + hidden_channel] + + focus_grad[(7 + out_coeff) * shared_row_pitch + hidden_channel] + * local_values[(7 + in_coeff) * shared_row_pitch + hidden_channel] + ) + return cute.arch.warp_reduction_sum( + value, + threads_in_group=WARP_REDUCTION_GROUP, + ) + + +@cute.jit +def _warp_owned_grad_d( + local_values, + x_values, + panel_idx, + subgroup_lane, + shared_row_pitch: cutlass.Constexpr[int], + x_row_pitch: cutlass.Constexpr[int], +): + """Reduce one packed Wigner adjoint inside a four-lane subgroup.""" + reduced = cutlass.Int32(0) + full_col = cutlass.Int32(0) + if panel_idx >= 25: + local_idx = panel_idx - 25 + row_slot = local_idx // 7 + reduced = 3 + row_slot * 3 + full_col = 9 + local_idx - row_slot * 7 + elif panel_idx >= 10: + local_idx = panel_idx - 10 + row_slot = local_idx // 5 + reduced = 2 + row_slot * 3 + full_col = 4 + local_idx - row_slot * 5 + elif panel_idx >= 1: + local_idx = panel_idx - 1 + row_slot = local_idx // 3 + reduced = 1 + row_slot * 3 + full_col = 1 + local_idx - row_slot * 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 + value += ( + local_values[reduced * shared_row_pitch + hidden_channel] + * x_values[full_col * x_row_pitch + hidden_channel] + ) + return cute.arch.warp_reduction_sum( + value, + threads_in_group=WARP_REDUCTION_GROUP, + ) + + +@cute.jit +def _grad_x_value( + local_values, + d_values, + channel, + degree: cutlass.Constexpr[int], + local_col: cutlass.Constexpr[int], + panel_start: cutlass.Constexpr[int], + width: cutlass.Constexpr[int], + rows: cutlass.Constexpr[int], + shared_row_pitch: cutlass.Constexpr[int], +): + acc = cutlass.Float32(0.0) + for row_slot in cutlass.range_constexpr(rows): + reduced = degree + row_slot * 3 + panel_offset = panel_start + row_slot * width + local_col + acc += ( + d_values[panel_offset] * local_values[reduced * shared_row_pitch + channel] + ) + return acc + + +@cute.jit +def neo_radial_phase_a_backward_node_jit( + grad_out_focus: cute.Tensor, + grad_focus_src: cute.Tensor, + grad_logits: cute.Tensor, + radial_state: cute.Tensor, + channel_basis: cute.Tensor, + x_wide: cute.Tensor, + source_order: cute.Tensor, + source_ptr: cute.Tensor, + d_full: cute.Tensor, + grad_x_wide: cute.Tensor, + grad_d_full: cute.Tensor, + stream: CUstream, +): + params = NeoRadialPhaseABackwardNodeParams( + grad_out_focus=grad_out_focus, + grad_focus_src=grad_focus_src, + grad_logits=grad_logits, + radial_state=radial_state, + channel_basis=channel_basis, + x_wide=x_wide, + source_order=source_order, + source_ptr=source_ptr, + d_full=d_full, + grad_x_wide=grad_x_wide, + grad_d_full=grad_d_full, + ) + node_count, _ = grad_x_wide.shape + neo_radial_phase_a_backward_node_kernel(params).launch( + grid=[node_count, 1, 1], + block=[HIDDEN, 1, 1], + stream=stream, + ) + + +@cute.kernel +def neo_radial_phase_a_backward_node_kernel( + params: NeoRadialPhaseABackwardNodeParams, +): + channel, _, _ = cute.arch.thread_idx() + node, _, _ = cute.arch.block_idx() + shared_row_pitch = SHARED_ROW_PITCH + + smem = cutlass.utils.SmemAllocator() + x_row_pitch = SHARED_ROW_PITCH + x_values = smem.allocate_tensor( + cutlass.Float32, + DEGREE_COUNT * x_row_pitch, + ) + focus_grad = smem.allocate_tensor( + cutlass.Float32, + REDUCED_COUNT * SHARED_ROW_PITCH, + ) + # The primal local rows are dead after grad_compact. Reuse this panel for + # their adjoints instead of reserving another 2.5 KiB per CTA. + local_values = smem.allocate_tensor( + cutlass.Float32, + REDUCED_COUNT * SHARED_ROW_PITCH, + ) + d_values = smem.allocate_tensor(cutlass.Float32, PACKED_WIGNER_VALUES) + compact = smem.allocate_tensor(cutlass.Float32, COMPACT_WIDTH) + grad_compact = smem.allocate_tensor(cutlass.Float32, COMPACT_WIDTH) + + for full_row in cutlass.range_constexpr(DEGREE_COUNT): + x_values[full_row * x_row_pitch + channel] = params.x_wide[ + node, + full_row * HIDDEN + channel, + ].to(cutlass.Float32) + cute.arch.sync_threads() + + grad_x_0 = cutlass.Float32(0.0) + grad_x_1 = cutlass.Float32(0.0) + grad_x_2 = cutlass.Float32(0.0) + grad_x_3 = cutlass.Float32(0.0) + grad_x_4 = cutlass.Float32(0.0) + grad_x_5 = cutlass.Float32(0.0) + grad_x_6 = cutlass.Float32(0.0) + grad_x_7 = cutlass.Float32(0.0) + grad_x_8 = cutlass.Float32(0.0) + grad_x_9 = cutlass.Float32(0.0) + grad_x_10 = cutlass.Float32(0.0) + grad_x_11 = cutlass.Float32(0.0) + grad_x_12 = cutlass.Float32(0.0) + grad_x_13 = cutlass.Float32(0.0) + grad_x_14 = cutlass.Float32(0.0) + grad_x_15 = cutlass.Float32(0.0) + + lo = params.source_ptr[node] + hi = params.source_ptr[node + 1] + for slot in cutlass.range(lo, hi, 1, unroll=1): + edge = params.source_order[slot] + for coeff in cutlass.range_constexpr(REDUCED_COUNT): + row_offset = coeff * shared_row_pitch + channel + focus_grad[row_offset] = _focus_grad_value( + params.grad_out_focus, + params.grad_focus_src, + edge, + coeff, + channel, + ) + if channel < COMPACT_WIDTH: + compact[channel] = params.radial_state[edge, channel].to(cutlass.Float32) + + if channel < PACKED_WIGNER_VALUES: + d_values[channel] = params.d_full[edge, channel].to(cutlass.Float32) + cute.arch.sync_threads() + + if cutlass.const_expr(True): + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 0, + 0, + 0, + 1, + shared_row_pitch, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 1, + 1, + 1, + 3, + shared_row_pitch, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 2, + 10, + 4, + 5, + shared_row_pitch, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 3, + 25, + 9, + 7, + shared_row_pitch, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 4, + 4, + 1, + 3, + shared_row_pitch, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 5, + 15, + 4, + 5, + shared_row_pitch, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 6, + 32, + 9, + 7, + shared_row_pitch, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 7, + 7, + 1, + 3, + shared_row_pitch, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 8, + 20, + 4, + 5, + shared_row_pitch, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 9, + 39, + 9, + 7, + shared_row_pitch, + x_row_pitch, + ) + cute.arch.sync_threads() + + lane = channel % 32 + warp = channel // 32 + subgroup = lane // WARP_REDUCTION_GROUP + subgroup_lane = lane % WARP_REDUCTION_GROUP + group = warp * GROUPS_PER_WARP + subgroup + for batch in cutlass.range_constexpr( + (COMPACT_WIDTH + GROUPS_PER_CTA - 1) // GROUPS_PER_CTA + ): + compact_idx = batch * GROUPS_PER_CTA + group + safe_compact_idx = compact_idx + if compact_idx >= COMPACT_WIDTH: + safe_compact_idx = cutlass.Int32(0) + grad_value = _warp_owned_grad_compact( + focus_grad, + local_values, + params.channel_basis, + safe_compact_idx, + subgroup_lane, + shared_row_pitch, + ) + if compact_idx < COMPACT_WIDTH: + if subgroup_lane == 0: + grad_compact[compact_idx] = grad_value + cute.arch.sync_threads() + + basis = params.channel_basis[channel].to(cutlass.Float32) + for coeff in cutlass.range_constexpr(REDUCED_COUNT): + grad_value = cutlass.Float32(0.0) + if coeff < 4: + for out_coeff in cutlass.range_constexpr(4): + grad_value += ( + focus_grad[out_coeff * shared_row_pitch + channel] + * compact[coeff * 4 + out_coeff] + ) + elif coeff < 7: + in_coeff = coeff - 4 + for out_coeff in cutlass.range_constexpr(3): + grad_value += ( + focus_grad[(4 + out_coeff) * shared_row_pitch + channel] + * compact[16 + in_coeff * 3 + out_coeff] + ) + else: + in_coeff = coeff - 7 + for out_coeff in cutlass.range_constexpr(3): + grad_value += ( + focus_grad[(7 + out_coeff) * shared_row_pitch + channel] + * compact[16 + in_coeff * 3 + out_coeff] + ) + local_values[coeff * shared_row_pitch + channel] = grad_value * basis + + if channel < COMPACT_WIDTH: + # grad_out_focus is dead once this edge has been reduced. Pack + # the 27-column GEMM operand in-place to avoid a separate cat. + params.grad_out_focus[edge, channel] = grad_compact[channel].to( + params.grad_out_focus.element_type + ) + if channel < FOCUS_COUNT: + params.grad_out_focus[edge, COMPACT_WIDTH + channel] = params.grad_logits[ + edge, channel + ].to(params.grad_out_focus.element_type) + cute.arch.sync_threads() + + for batch in cutlass.range_constexpr( + (PACKED_WIGNER_VALUES + GROUPS_PER_CTA - 1) // GROUPS_PER_CTA + ): + panel_idx = batch * GROUPS_PER_CTA + group + safe_panel_idx = panel_idx + if panel_idx >= PACKED_WIGNER_VALUES: + safe_panel_idx = cutlass.Int32(0) + grad_d_value = _warp_owned_grad_d( + local_values, + x_values, + safe_panel_idx, + subgroup_lane, + shared_row_pitch, + x_row_pitch, + ) + if panel_idx < PACKED_WIGNER_VALUES: + if subgroup_lane == 0: + params.grad_d_full[edge, panel_idx] = grad_d_value.to( + params.grad_d_full.element_type + ) + + grad_x_0 += _grad_x_value( + local_values, d_values, channel, 0, 0, 0, 1, 1, shared_row_pitch + ) + grad_x_1 += _grad_x_value( + local_values, d_values, channel, 1, 0, 1, 3, 3, shared_row_pitch + ) + grad_x_2 += _grad_x_value( + local_values, d_values, channel, 1, 1, 1, 3, 3, shared_row_pitch + ) + grad_x_3 += _grad_x_value( + local_values, d_values, channel, 1, 2, 1, 3, 3, shared_row_pitch + ) + grad_x_4 += _grad_x_value( + local_values, d_values, channel, 2, 0, 10, 5, 3, shared_row_pitch + ) + grad_x_5 += _grad_x_value( + local_values, d_values, channel, 2, 1, 10, 5, 3, shared_row_pitch + ) + grad_x_6 += _grad_x_value( + local_values, d_values, channel, 2, 2, 10, 5, 3, shared_row_pitch + ) + grad_x_7 += _grad_x_value( + local_values, d_values, channel, 2, 3, 10, 5, 3, shared_row_pitch + ) + grad_x_8 += _grad_x_value( + local_values, d_values, channel, 2, 4, 10, 5, 3, shared_row_pitch + ) + grad_x_9 += _grad_x_value( + local_values, d_values, channel, 3, 0, 25, 7, 3, shared_row_pitch + ) + grad_x_10 += _grad_x_value( + local_values, d_values, channel, 3, 1, 25, 7, 3, shared_row_pitch + ) + grad_x_11 += _grad_x_value( + local_values, d_values, channel, 3, 2, 25, 7, 3, shared_row_pitch + ) + grad_x_12 += _grad_x_value( + local_values, d_values, channel, 3, 3, 25, 7, 3, shared_row_pitch + ) + grad_x_13 += _grad_x_value( + local_values, d_values, channel, 3, 4, 25, 7, 3, shared_row_pitch + ) + grad_x_14 += _grad_x_value( + local_values, d_values, channel, 3, 5, 25, 7, 3, shared_row_pitch + ) + grad_x_15 += _grad_x_value( + local_values, d_values, channel, 3, 6, 25, 7, 3, shared_row_pitch + ) + cute.arch.sync_threads() + + params.grad_x_wide[node, 0 * HIDDEN + channel] = grad_x_0.to( + params.grad_x_wide.element_type + ) + params.grad_x_wide[node, 1 * HIDDEN + channel] = grad_x_1.to( + params.grad_x_wide.element_type + ) + params.grad_x_wide[node, 2 * HIDDEN + channel] = grad_x_2.to( + params.grad_x_wide.element_type + ) + params.grad_x_wide[node, 3 * HIDDEN + channel] = grad_x_3.to( + params.grad_x_wide.element_type + ) + params.grad_x_wide[node, 4 * HIDDEN + channel] = grad_x_4.to( + params.grad_x_wide.element_type + ) + params.grad_x_wide[node, 5 * HIDDEN + channel] = grad_x_5.to( + params.grad_x_wide.element_type + ) + params.grad_x_wide[node, 6 * HIDDEN + channel] = grad_x_6.to( + params.grad_x_wide.element_type + ) + params.grad_x_wide[node, 7 * HIDDEN + channel] = grad_x_7.to( + params.grad_x_wide.element_type + ) + params.grad_x_wide[node, 8 * HIDDEN + channel] = grad_x_8.to( + params.grad_x_wide.element_type + ) + params.grad_x_wide[node, 9 * HIDDEN + channel] = grad_x_9.to( + params.grad_x_wide.element_type + ) + params.grad_x_wide[node, 10 * HIDDEN + channel] = grad_x_10.to( + params.grad_x_wide.element_type + ) + params.grad_x_wide[node, 11 * HIDDEN + channel] = grad_x_11.to( + params.grad_x_wide.element_type + ) + params.grad_x_wide[node, 12 * HIDDEN + channel] = grad_x_12.to( + params.grad_x_wide.element_type + ) + params.grad_x_wide[node, 13 * HIDDEN + channel] = grad_x_13.to( + params.grad_x_wide.element_type + ) + params.grad_x_wide[node, 14 * HIDDEN + channel] = grad_x_14.to( + params.grad_x_wide.element_type + ) + params.grad_x_wide[node, 15 * HIDDEN + channel] = grad_x_15.to( + params.grad_x_wide.element_type + ) + + +def compile_neo_radial_phase_a_backward_node_tiled() -> Callable: + """Compile the exact-shape source-CSR node-tiled backward specialization.""" + edge_count = cute.sym_int64() + node_count = cute.sym_int64() + source_ptr_count = cute.sym_int64() + fake_grad_out = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, FOCUS_COUNT * REDUCED_COUNT * FOCUS_HIDDEN), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_grad_focus_src = make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, edge_count, FOCUS_HIDDEN), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + fake_grad_logits = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, FOCUS_COUNT), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_radial_state = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, COMPACT_WIDTH), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_basis = make_fake_compact_tensor( + cutlass.Float32, + (HIDDEN,), + stride_order=(0,), + **FAKE_TENSOR_KW, + ) + fake_x_wide = make_fake_compact_tensor( + cutlass.Float32, + (node_count, DEGREE_COUNT * HIDDEN), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_source_order = make_fake_compact_tensor( + cutlass.Int32, + (edge_count,), + stride_order=(0,), + **FAKE_TENSOR_KW, + ) + fake_source_ptr = make_fake_compact_tensor( + cutlass.Int32, + (source_ptr_count,), + stride_order=(0,), + **FAKE_TENSOR_KW, + ) + fake_d = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, PACKED_WIGNER_VALUES), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_grad_x = make_fake_compact_tensor( + cutlass.Float32, + (node_count, DEGREE_COUNT * HIDDEN), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_grad_d = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, PACKED_WIGNER_VALUES), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + neo_radial_phase_a_backward_node_jit, + fake_grad_out, + fake_grad_focus_src, + fake_grad_logits, + fake_radial_state, + fake_basis, + fake_x_wide, + fake_source_order, + fake_source_ptr, + fake_d, + fake_grad_x, + fake_grad_d, + fake_stream, + options="--enable-tvm-ffi", + ) diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_so2_gate_combined_fwd.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_so2_gate_combined_fwd.py new file mode 100644 index 0000000000..4dec70c29a --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_so2_gate_combined_fwd.py @@ -0,0 +1,779 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""One-launch strict-FP32 Neo SO2 gate/residual forward.""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, +) + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as cute_utils +import torch +from cutlass.cute.runtime import ( + make_fake_stream, + make_fake_tensor, +) + +from ..compile_cache import ( + device_aware_lru_cache, +) +from ..runtime_policy import ( + FUSED_SO2_GATE_CAPABILITIES, +) + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN001, ANN201, ANN202, ANN204 + + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + +TILE_M = 64 +TILE_K = 16 +THREADS = 256 +STAGES = 3 +FOCUS_COUNT = 2 +M0_WIDTH = 4 * 32 +PAIR_WIDTH = 6 * 32 +FULL_WIDTH = M0_WIDTH + PAIR_WIDTH + + +def _supports_combined_forward(compute_capability: tuple[int, int]) -> bool: + return compute_capability in FUSED_SO2_GATE_CAPABILITIES + + +def _require_16_byte_alignment(tensors: tuple[torch.Tensor, ...]) -> None: + if any(tensor.data_ptr() % 16 for tensor in tensors): + raise ValueError("combined Neo SO2 gate tensors must be 16-byte aligned") + + +@cute.jit +def _sigmoid(value): + return cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-value)) + + +def _fake_focus_tensor(width: int): + return make_fake_tensor( + cutlass.Float32, + (cute.sym_int32(), FOCUS_COUNT, width), + (FOCUS_COUNT * width, width, 1), + assumed_align=16, + ) + + +def _fake_focus_weight(width: int): + return make_fake_tensor( + cutlass.Float32, + (FOCUS_COUNT, width, width), + (width * width, width, 1), + assumed_align=16, + ) + + +def prepare_neo_so2_gate_combined_weights( + w0: torch.Tensor, + wp: torch.Tensor, + gate_weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pack immutable weights into contiguous focus-major tensors.""" + return ( + w0.transpose(1, 2).contiguous(), + wp.transpose(1, 2).contiguous(), + gate_weight.permute(1, 0, 2).contiguous(), + ) + + +class CuteNeoSO2GateCombined: + """M64/K16/T256/S3 SIMT SGEMMs with one CTA-resident gate epilogue.""" + + def __init__(self) -> None: + self.cta_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=THREADS, + ) + + @cute.jit + def __call__( + self, + mX: cute.Tensor, + mResidual: cute.Tensor, + mW0: cute.Tensor, + mWP: cute.Tensor, + mGate: cute.Tensor, + mY: cute.Tensor, + mOut: cute.Tensor, + stream: cuda.CUstream = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT), + ): + sA_layout = cute.make_layout( + (TILE_M, TILE_K, STAGES), + stride=(1, TILE_M + 4, TILE_K * (TILE_M + 4)), + ) + sB_pair_layout = cute.make_layout( + (PAIR_WIDTH, TILE_K, STAGES), + stride=(1, PAIR_WIDTH + 4, TILE_K * (PAIR_WIDTH + 4)), + ) + sB_m0_layout = cute.make_layout( + (M0_WIDTH, TILE_K, STAGES), + stride=(1, PAIR_WIDTH + 4, TILE_K * (PAIR_WIDTH + 4)), + ) + sY0_layout = cute.make_layout( + (TILE_M, 32), + stride=(32, 1), + ) + sGate_layout = cute.make_layout( + (TILE_M, 3 * 32), + stride=(3 * 32, 1), + ) + + copy_layout = cute.make_layout( + (THREADS // TILE_K, TILE_K), + stride=(TILE_K, 1), + ) + copy_value = cute.make_layout((1, 1)) + copy_a = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + mX.element_type, + num_bits_per_copy=mX.element_type.width, + ) + copy_b = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + mW0.element_type, + num_bits_per_copy=mW0.element_type.width, + ) + tiled_copy_A = cute.make_tiled_copy_tv(copy_a, copy_layout, copy_value) + tiled_copy_B = cute.make_tiled_copy_tv(copy_b, copy_layout, copy_value) + + atoms_layout = cute.make_layout( + (THREADS // 16, 16, 1), + stride=(16, 1, 0), + ) + permutation_m = cute.make_layout( + (atoms_layout.shape[0], 4), + stride=(4, 1), + ) + permutation_n = cute.make_layout( + (atoms_layout.shape[1], 4), + stride=(4, 1), + ) + m0_op = cute.nvgpu.MmaUniversalOp(cutlass.Float32) + pair_op = cute.nvgpu.MmaUniversalOp(cutlass.Float32) + tiled_mma_m0 = cute.make_tiled_mma( + m0_op, + atoms_layout, + permutation_mnk=(permutation_m, permutation_n, None), + ) + tiled_mma_pair = cute.make_tiled_mma( + pair_op, + atoms_layout, + permutation_mnk=(permutation_m, permutation_n, None), + ) + + self.kernel( + mX, + mResidual, + mW0, + mWP, + mGate, + mY, + mOut, + sA_layout, + sB_m0_layout, + sB_pair_layout, + sY0_layout, + sGate_layout, + tiled_copy_A, + tiled_copy_B, + tiled_mma_m0, + tiled_mma_pair, + ).launch( + grid=(cute.ceil_div(mY.shape[0], TILE_M), FOCUS_COUNT, 1), + block=[THREADS, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mX: cute.Tensor, + mResidual: cute.Tensor, + mW0: cute.Tensor, + mWP: cute.Tensor, + mGate: cute.Tensor, + mY: cute.Tensor, + mOut: cute.Tensor, + sA_layout: cute.Layout, + sB_m0_layout: cute.Layout, + sB_pair_layout: cute.Layout, + sY0_layout: cute.Layout, + sGate_layout: cute.Layout, + tiled_copy_A: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_mma_m0: cute.TiledMma, + tiled_mma_pair: cute.TiledMma, + ): + tidx, _, _ = cute.arch.thread_idx() + edge_tile, focus, _ = cute.arch.block_idx() + + x_focus = mX[None, focus, None] + residual_focus = mResidual[None, focus, None] + y_focus = mY[None, focus, None] + out_focus = mOut[None, focus, None] + matrix_layout_m0 = cute.make_layout( + (mY.shape[0], M0_WIDTH), + stride=(FOCUS_COUNT * FULL_WIDTH, 1), + ) + matrix_layout_pair = cute.make_layout( + (mY.shape[0], PAIR_WIDTH), + stride=(FOCUS_COUNT * FULL_WIDTH, 1), + ) + mA0 = cute.make_tensor(x_focus.iterator, matrix_layout_m0) + mAPair = cute.make_tensor( + x_focus.iterator + M0_WIDTH, + matrix_layout_pair, + ) + mR0 = cute.make_tensor(residual_focus.iterator, matrix_layout_m0) + mRPair = cute.make_tensor( + residual_focus.iterator + M0_WIDTH, + matrix_layout_pair, + ) + mY0 = cute.make_tensor(y_focus.iterator, matrix_layout_m0) + mYPair = cute.make_tensor( + y_focus.iterator + M0_WIDTH, + matrix_layout_pair, + ) + mOut0 = cute.make_tensor(out_focus.iterator, matrix_layout_m0) + mOutPair = cute.make_tensor( + out_focus.iterator + M0_WIDTH, + matrix_layout_pair, + ) + w0_focus = mW0[focus, None, None] + wp_focus = mWP[focus, None, None] + gate_focus = mGate[focus, None, None] + + smem = cute_utils.SmemAllocator() + sA = smem.allocate_tensor(cutlass.Float32, sA_layout, 16) + sB = smem.allocate_tensor(cutlass.Float32, sB_pair_layout, 16) + sY0 = smem.allocate_tensor(cutlass.Float32, sY0_layout, 16) + sGate = smem.allocate_tensor(cutlass.Float32, sGate_layout, 16) + + self._run_m0( + mA0, + w0_focus, + mR0, + gate_focus, + mY0, + mOut0, + sA, + sB, + sY0, + sGate, + sB_m0_layout, + tiled_copy_A, + tiled_copy_B, + tiled_mma_m0, + tidx, + edge_tile, + ) + cute.arch.sync_threads() + self._run_pair( + mAPair, + wp_focus, + mRPair, + mYPair, + mOutPair, + sA, + sB, + sGate, + tiled_copy_A, + tiled_copy_B, + tiled_mma_pair, + tidx, + edge_tile, + ) + + @cute.jit + def _run_m0( + self, + mA: cute.Tensor, + mB: cute.Tensor, + mR: cute.Tensor, + mGate: cute.Tensor, + mY: cute.Tensor, + mOut: cute.Tensor, + sA: cute.Tensor, + sB: cute.Tensor, + sY0: cute.Tensor, + sGate: cute.Tensor, + sB_m0_layout: cute.Layout, + tiled_copy_A: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + edge_tile: cutlass.Int32, + ): + sB_m0 = cute.make_tensor(sB.iterator, sB_m0_layout) + self._run_gemm( + mA, + mB, + mR, + mGate, + mY, + mOut, + sA, + sB_m0, + sY0, + sGate, + tiled_copy_A, + tiled_copy_B, + tiled_mma, + tidx, + edge_tile, + m0_block=True, + ) + + @cute.jit + def _run_pair( + self, + mA: cute.Tensor, + mB: cute.Tensor, + mR: cute.Tensor, + mY: cute.Tensor, + mOut: cute.Tensor, + sA: cute.Tensor, + sB: cute.Tensor, + sGate: cute.Tensor, + tiled_copy_A: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + edge_tile: cutlass.Int32, + ): + self._run_gemm( + mA, + mB, + mR, + mB, + mY, + mOut, + sA, + sB, + sA, + sGate, + tiled_copy_A, + tiled_copy_B, + tiled_mma, + tidx, + edge_tile, + m0_block=False, + ) + + @cute.jit + def _run_gemm( + self, + mA: cute.Tensor, + mB: cute.Tensor, + mR: cute.Tensor, + mGate: cute.Tensor, + mY: cute.Tensor, + mOut: cute.Tensor, + sA: cute.Tensor, + sB: cute.Tensor, + sY0: cute.Tensor, + sGate: cute.Tensor, + tiled_copy_A: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + edge_tile: cutlass.Int32, + m0_block: cutlass.Constexpr[bool], + ): + width = M0_WIDTH if cutlass.const_expr(m0_block) else PAIR_WIDTH + cta_tiler = (TILE_M, width, TILE_K) + tiler_coord = (edge_tile, 0, None) + thr_mma = tiled_mma.get_slice(tidx) + + gA = cute.local_tile( + mA, + tiler=cta_tiler, + coord=tiler_coord, + proj=(1, None, 1), + ) + gB = cute.local_tile( + mB, + tiler=cta_tiler, + coord=tiler_coord, + proj=(None, 1, 1), + ) + gR = cute.local_tile( + mR, + tiler=cta_tiler, + coord=tiler_coord, + proj=(1, 1, None), + ) + gY = cute.local_tile( + mY, + tiler=cta_tiler, + coord=tiler_coord, + proj=(1, 1, None), + ) + gOut = cute.local_tile( + mOut, + tiler=cta_tiler, + coord=tiler_coord, + proj=(1, 1, None), + ) + + thr_copy_A = tiled_copy_A.get_slice(tidx) + thr_copy_B = tiled_copy_B.get_slice(tidx) + tAgA = thr_copy_A.partition_S(gA) + tAsA = thr_copy_A.partition_D(sA) + tBgB = thr_copy_B.partition_S(gB) + tBsB = thr_copy_B.partition_D(sB) + + mcA = cute.make_identity_tensor(mA.shape) + cA = cute.local_tile( + mcA, + tiler=cta_tiler, + coord=tiler_coord, + proj=(1, None, 1), + ) + tAcA = thr_copy_A.partition_S(cA) + tApA = cute.make_rmem_tensor( + cute.make_layout( + ( + tAsA.shape[0][1], + cute.size(tAsA, mode=[1]), + cute.size(tAsA, mode=[2]), + ), + stride=(cute.size(tAsA, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in range(tApA.shape[0]): + for row in range(tApA.shape[1]): + tApA[rest_v, row, 0] = cute.elem_less( + tAcA[(0, rest_v), row, 0, 0][0], + mA.shape[0], + ) + + k_pipe_max = cute.size(tAsA, mode=[3]) + k_tile_count = cute.size(tAgA, mode=[3]) + gmem_pipe_read = cutlass.Int32(0) + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, 0], + pred=tApA, + ) + cute.copy( + tiled_copy_B, + tBgB[None, None, None, gmem_pipe_read], + tBsB[None, None, None, 0], + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = ( + gmem_pipe_read + 1 + if gmem_pipe_read + 1 < k_tile_count + else cutlass.Int32(0) + ) + for k_tile in range(1, k_pipe_max - 1): + if k_tile < k_tile_count: + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, k_tile], + pred=tApA, + ) + cute.copy( + tiled_copy_B, + tBgB[None, None, None, gmem_pipe_read], + tBsB[None, None, None, k_tile], + ) + gmem_pipe_read = ( + gmem_pipe_read + 1 + if gmem_pipe_read + 1 < k_tile_count + else cutlass.Int32(0) + ) + cute.arch.cp_async_commit_group() + + tCsA = thr_mma.partition_A(sA) + tCsB = thr_mma.partition_B(sB) + tCgR = thr_mma.partition_C(gR) + tCgY = thr_mma.partition_C(gY) + tCgOut = thr_mma.partition_C(gOut) + tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) + tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0]) + tCrC = tiled_mma.make_fragment_C(tCgOut) + tCrC.fill(0.0) + + smem_pipe_read = cutlass.Int32(0) + smem_pipe_write = cutlass.Int32(k_pipe_max - 1) + tiles_issued = cutlass.Int32(k_pipe_max - 1) + tCsA_p = tCsA[None, None, None, smem_pipe_read] + tCsB_p = tCsB[None, None, None, smem_pipe_read] + k_block_max = cute.size(tCrA, mode=[2]) + + if k_block_max > 1: + cute.arch.cp_async_wait_group(k_pipe_max - 2) + self.cta_sync_barrier.arrive_and_wait() + cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0]) + cute.autovec_copy(tCsB_p[None, None, 0], tCrB[None, None, 0]) + + for _ in range(k_tile_count): + for k_block in range(k_block_max, unroll_full=True): + if k_block == k_block_max - 1: + tCsA_p = tCsA[None, None, None, smem_pipe_read] + tCsB_p = tCsB[None, None, None, smem_pipe_read] + cute.arch.cp_async_wait_group(k_pipe_max - 2) + self.cta_sync_barrier.arrive_and_wait() + + k_block_next = (k_block + 1) % k_block_max + cute.autovec_copy( + tCsA_p[None, None, k_block_next], + tCrA[None, None, k_block_next], + ) + cute.autovec_copy( + tCsB_p[None, None, k_block_next], + tCrB[None, None, k_block_next], + ) + if k_block == 0: + if tiles_issued < k_tile_count: + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, smem_pipe_write], + pred=tApA, + ) + cute.gemm( + tiled_mma, + tCrC, + tCrA[None, None, k_block], + tCrB[None, None, k_block], + tCrC, + ) + if k_block == 0: + if tiles_issued < k_tile_count: + cute.copy( + tiled_copy_B, + tBgB[None, None, None, gmem_pipe_read], + tBsB[None, None, None, smem_pipe_write], + ) + cute.arch.cp_async_commit_group() + tiles_issued = tiles_issued + 1 + smem_pipe_write = smem_pipe_read + smem_pipe_read = smem_pipe_read + 1 + if smem_pipe_read == k_pipe_max: + smem_pipe_read = cutlass.Int32(0) + gmem_pipe_read = ( + gmem_pipe_read + 1 + if gmem_pipe_read + 1 < k_tile_count + else cutlass.Int32(1) + ) + + cute.arch.cp_async_wait_group(0) + self.cta_sync_barrier.arrive_and_wait() + tCrC.store(tCrC.load()) + + cC = cute.make_identity_tensor(gOut.shape) + tCpC = thr_mma.partition_C(cC) + predC = cute.make_rmem_tensor(tCrC.layout, cutlass.Boolean) + residue_m = mOut.shape[0] - cutlass.Int32(TILE_M) * edge_tile + for idx in range(cute.size(tCrC.shape)): + predC[idx] = cute.elem_less(tCpC[idx], (residue_m, width)) + + atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mOut.element_type, + ) + cute.copy(atom, tCrC, tCgY, pred=predC) + tCrR = tiled_mma.make_fragment_C(tCgR) + tCrR.fill(0.0) + cute.copy(atom, tCgR, tCrR, pred=predC) + + if cutlass.const_expr(m0_block): + for idx in range(cute.size(tCrC.shape)): + if predC[idx]: + local_row = tCpC[idx][0] + local_col = tCpC[idx][1] + if local_col < 32: + sY0[local_row, local_col] = tCrC[idx].to(cutlass.Float32) + cute.arch.sync_threads() + + gate_slots = (TILE_M * 3 * 32 + THREADS - 1) // THREADS + for slot in cutlass.range_constexpr(gate_slots): + linear_idx = tidx + slot * THREADS + if linear_idx < TILE_M * 3 * 32: + local_row = linear_idx // (3 * 32) + gate_idx = linear_idx - local_row * (3 * 32) + global_row = edge_tile * TILE_M + local_row + if global_row < mOut.shape[0]: + gate_logit = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(32): + gate_logit += sY0[local_row, k] * mGate[k, gate_idx] + sGate[local_row, gate_idx] = _sigmoid(gate_logit) + cute.arch.sync_threads() + + for idx in range(cute.size(tCrC.shape)): + if predC[idx]: + local_row = tCpC[idx][0] + local_col = tCpC[idx][1] + value = tCrC[idx].to(cutlass.Float32) + if cutlass.const_expr(m0_block): + if local_col < 32: + value = value * _sigmoid(value) + else: + value = value * sGate[local_row, local_col - 32] + else: + gate_idx = local_col + if gate_idx >= 3 * 32: + gate_idx = gate_idx - 3 * 32 + value = value * sGate[local_row, gate_idx] + tCrC[idx] = value + tCrR[idx].to(cutlass.Float32) + + cute.copy(atom, tCrC, tCgOut, pred=predC) + + +@device_aware_lru_cache(maxsize=8) +def _compile_combined_forward( + device_index: int, + compute_capability: tuple[int, int], +) -> Callable: + if not _supports_combined_forward(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) + fake_w0 = _fake_focus_weight(M0_WIDTH) + fake_wp = _fake_focus_weight(PAIR_WIDTH) + fake_gate = make_fake_tensor( + cutlass.Float32, + (FOCUS_COUNT, 32, 3 * 32), + (32 * 3 * 32, 3 * 32, 1), + assumed_align=16, + ) + fake_y = _fake_focus_tensor(FULL_WIDTH) + fake_out = _fake_focus_tensor(FULL_WIDTH) + operation = CuteNeoSO2GateCombined() + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) + return cute.compile( + operation, + fake_x, + fake_residual, + fake_w0, + fake_wp, + fake_gate, + fake_y, + fake_out, + stream=fake_stream, + options="--enable-tvm-ffi", + ) + + +class CuteNeoSO2GateCombinedFwdRunner: + """Prebuilt one-launch forward that writes full y and no global aux.""" + + def __init__( + self, + x: torch.Tensor, + residual: torch.Tensor, + y: torch.Tensor, + out: torch.Tensor, + *, + packed_weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + packed_weights_ready: tuple[torch.cuda.Event, int], + ) -> None: + expected = (x.shape[0], FOCUS_COUNT, 10, 32) + if tuple(x.shape) != expected: + raise ValueError(f"x must have shape {expected}, got {tuple(x.shape)}") + if residual.shape != x.shape or y.shape != x.shape or out.shape != x.shape: + raise ValueError("residual, y, and out must match x") + if x.shape[0] <= 0: + raise ValueError("combined Neo SO2 gate forward requires E > 0") + tensors = (x, residual, y, out) + if any( + tensor.dtype != torch.float32 or not tensor.is_cuda for tensor in tensors + ): + raise TypeError( + "combined Neo SO2 gate forward requires CUDA float32 tensors" + ) + if any(tensor.device != x.device for tensor in tensors): + raise ValueError( + "all combined Neo SO2 gate forward tensors must share x.device" + ) + if any(not tensor.is_contiguous() for tensor in tensors): + raise ValueError( + "combined Neo SO2 gate forward requires canonical contiguous tensors" + ) + _require_16_byte_alignment(tensors) + device_index = x.device.index + if device_index is None: + raise RuntimeError("combined Neo SO2 gate forward requires a CUDA index") + compute_capability = tuple(torch.cuda.get_device_capability(device_index)) + if not _supports_combined_forward(compute_capability): + raise RuntimeError( + "combined forward requires a supported compute capability" + ) + + packed_w0, packed_wp, packed_gate = packed_weights + packed = (packed_w0, packed_wp, packed_gate) + if any( + tensor.dtype != torch.float32 or not tensor.is_cuda for tensor in packed + ): + raise TypeError( + "packed combined forward weights must be CUDA float32 tensors" + ) + if any(tensor.device != x.device for tensor in packed): + raise ValueError("all packed combined forward weights must share x.device") + if any(not tensor.is_contiguous() for tensor in packed): + raise ValueError("packed combined forward weights must be contiguous") + _require_16_byte_alignment(packed) + if tuple(packed_w0.shape) != (FOCUS_COUNT, M0_WIDTH, M0_WIDTH): + raise ValueError("packed w0 must have shape (2,128,128)") + if tuple(packed_wp.shape) != (FOCUS_COUNT, PAIR_WIDTH, PAIR_WIDTH): + raise ValueError("packed wp must have shape (2,192,192)") + if tuple(packed_gate.shape) != (FOCUS_COUNT, 32, 3 * 32): + raise ValueError("packed gate weight must have shape (2,32,96)") + + with torch.cuda.device(x.device): + self._compiled = _compile_combined_forward( + device_index, + compute_capability, + ) + self._args = ( + x.reshape(x.shape[0], FOCUS_COUNT, FULL_WIDTH), + residual.reshape(x.shape[0], FOCUS_COUNT, FULL_WIDTH), + packed_w0, + packed_wp, + packed_gate, + y.reshape(x.shape[0], FOCUS_COUNT, FULL_WIDTH), + out.reshape(x.shape[0], FOCUS_COUNT, FULL_WIDTH), + ) + self._device = x.device + self._packed_weights_ready = packed_weights_ready + self._packed_weights_waited_streams: set[int] = set() + self.y = y + self.out = out + + def __call__(self) -> torch.Tensor: + with torch.cuda.device(self._device): + torch_stream = torch.cuda.current_stream(self._device) + ready_event, producer_stream = self._packed_weights_ready + if ( + torch_stream.cuda_stream != producer_stream + and torch_stream.cuda_stream not in self._packed_weights_waited_streams + ): + torch_stream.wait_event(ready_event) + self._packed_weights_waited_streams.add(torch_stream.cuda_stream) + stream = cuda.CUstream(torch_stream.cuda_stream) + self._compiled(*self._args, stream=stream) + return self.out diff --git a/deepmd/kernels/cute/neo/k1_message_grid_packed.py b/deepmd/kernels/cute/neo/k1_message_grid_packed.py new file mode 100644 index 0000000000..c7d9a6ce2c --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_message_grid_packed.py @@ -0,0 +1,363 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Packed-layout Neo message-grid forward and first input adjoint.""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, + Any, +) + +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 +CHANNELS = 32 +HIDDEN_CHANNELS = N_FOCUS * CHANNELS + + +def _validate_module_contract(net: Any) -> None: + expected = { + "layout": "flat", + "mode": "cross", + "op_type": "glu", + "n_focus": N_FOCUS, + "n_frames": N_FRAMES, + "channels": CHANNELS, + "dtype": torch.float32, + } + mismatches = { + name: (getattr(net, name, None), value) + for name, value in expected.items() + if getattr(net, name, None) != value + } + frame_expand = getattr(net, "frame_expand", None) + frame_contract = getattr(net, "frame_contract", None) + if frame_expand is None or frame_contract is None: + mismatches["frame_modules"] = ( + (frame_expand is not None, frame_contract is not None), + (True, True), + ) + else: + expected_frame = ("packed", N_FRAMES, CHANNELS) + for name, module in ( + ("frame_expand", frame_expand), + ("frame_contract", frame_contract), + ): + actual_frame = ( + getattr(module, "coefficient_layout", None), + getattr(module, "n_frames", None), + getattr(module, "channels", None), + ) + if actual_frame != expected_frame: + mismatches[name] = (actual_frame, expected_frame) + if mismatches: + details = ", ".join( + f"{name}={actual!r} (expected {wanted!r})" + for name, (actual, wanted) in mismatches.items() + ) + raise ValueError(f"packed message-grid module contract mismatch: {details}") + + +def is_supported_message_grid(net: Any) -> bool: + """Return whether forward and manual backward share the same contract.""" + try: + _validate_module_contract(net) + except ValueError: + return False + return True + + +def _as_flat_ndfc( + net: Any, + name: str, + value: Tensor, + *, + like: Tensor | None = None, +) -> Tensor: + """Adapt GridNet's flat layout without materializing a contiguous copy.""" + expected_shape = (COEFF_DIM, HIDDEN_CHANNELS) + if value.ndim != 3 or tuple(value.shape[1:]) != expected_shape: + raise ValueError( + f"packed message-grid {name} must have shape (N, {COEFF_DIM}, " + f"{HIDDEN_CHANNELS}), got {tuple(value.shape)}" + ) + if value.dtype != torch.float32: + raise ValueError(f"packed message-grid {name} must be FP32, got {value.dtype}") + if like is not None and (value.shape != like.shape or value.device != like.device): + raise ValueError( + f"packed message-grid {name} must match query shape/device; got " + f"shape={tuple(value.shape)}, device={value.device}" + ) + # SO3Linear's einsum returns the valid flat GridNet layout with stride + # (64, N*64, 1), while Phase C returns compact (1024, 64, 1). Both split + # their unit-stride final axis into (F=2, C=32) without a copy. + if value.stride(-1) != 1: + raise ValueError( + f"packed message-grid {name} requires a unit-stride folded F*C " + f"axis, got stride={tuple(value.stride())}" + ) + value_ndfc, shape_info = net._to_ndfc(value) + expected_ndfc = (value.shape[0], COEFF_DIM, N_FOCUS, CHANNELS) + if ( + tuple(shape_info) != tuple(value.shape) + or tuple(value_ndfc.shape) != expected_ndfc + ): + raise ValueError( + f"packed message-grid {name} did not adapt to {expected_ndfc}; got " + f"shape={tuple(value_ndfc.shape)}, stride={tuple(value_ndfc.stride())}" + ) + return value_ndfc + + +def _validate_contract( + net: Any, + query: Tensor, + context: Tensor, +) -> tuple[Tensor, Tensor]: + _validate_module_contract(net) + query_ndfc = _as_flat_ndfc(net, "query", query) + context_ndfc = _as_flat_ndfc(net, "context", context, like=query) + return query_ndfc, context_ndfc + + +def _expanded_frame_weight(module: Any) -> Tensor: + return module.weight.index_select(0, module.degree_index) + + +def _frame_expand_packed(module: Any, coeff: Tensor) -> Tensor: + weight = _expanded_frame_weight(module).view( + COEFF_DIM, + CHANNELS, + N_FRAMES, + CHANNELS, + ) + # Output order D,K,F,C is the native CuTe grid-product contract. The + # trailing F,C panel remains unit-stride and therefore coalesced. + # PyTorch's degree-batched einsum naturally returns a degree-major stride. + # Normalize once at this producer because the CuTe consumer's packed + # contract is compact (N,D,K,F,C), not because the flat input was strided. + return torch.einsum("ndfi,dikc->ndkfc", coeff, weight).contiguous() + + +def _frame_expand_packed_backward( + module: Any, + grad_packed: Tensor, +) -> Tensor: + weight = _expanded_frame_weight(module).view( + COEFF_DIM, + CHANNELS, + N_FRAMES, + CHANNELS, + ) + return torch.einsum("ndkfc,dikc->ndfi", grad_packed, weight) + + +def _frame_contract_packed(module: Any, coeff_packed: Tensor) -> Tensor: + weight = _expanded_frame_weight(module).view( + COEFF_DIM, + N_FRAMES, + CHANNELS, + CHANNELS, + ) + return torch.einsum("ndkfc,dkco->ndfo", coeff_packed, weight) + + +def _frame_contract_packed_backward(module: Any, grad_out: Tensor) -> Tensor: + weight = _expanded_frame_weight(module).view( + COEFF_DIM, + N_FRAMES, + CHANNELS, + CHANNELS, + ) + return torch.einsum("ndfo,dkco->ndkfc", grad_out, weight) + + +def _focus_linear_backward_input(linear: Any, grad_out: Tensor) -> Tensor: + weight = linear.weight.view(linear.in_channels, linear.n_focus, linear.out_channels) + return torch.einsum("bfo,ifo->bfi", grad_out, weight) + + +def _swiglu_backward_input(x: Tensor, grad_out: Tensor) -> Tensor: + gate, value = torch.chunk(x, chunks=2, dim=-1) + sigmoid = torch.sigmoid(gate) + grad_gate = grad_out * value * (sigmoid + gate * sigmoid * (1.0 - sigmoid)) + grad_value = grad_out * gate * sigmoid + return torch.cat([grad_gate, grad_value], dim=-1) + + +def run_packed_message_grid_forward( + net: Any, + query_flat: Tensor, + context_flat: Tensor, + *, + return_product: bool = False, + 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) + from .k1_kernels.cute_neo_message_grid_product import ( + run_message_grid_product, + ) + + nodes = query_flat.shape[0] + scalar_pair = torch.cat([query[:, 0], context[:, 0]], dim=-1).to(net.dtype) + + left_packed = _frame_expand_packed(net.frame_expand, query) + right_packed = _frame_expand_packed(net.frame_expand, context) + left = left_packed.view(nodes, COEFF_DIM * N_FRAMES, HIDDEN_CHANNELS) + right = right_packed.view(nodes, COEFF_DIM * N_FRAMES, HIDDEN_CHANNELS) + if sm90_state is None: + product_flat = run_message_grid_product( + left, + right, + net.projector.to_grid_mat, + net.projector.from_grid_mat, + ) + else: + from .message_grid_gaunt_sm90 import ( + run_sm90_gaunt_forward, + ) + + product_flat = run_sm90_gaunt_forward( + left, + right, + sm90_state.schedule, + ) + product = product_flat.view( + nodes, + COEFF_DIM, + N_FRAMES, + N_FOCUS, + CHANNELS, + ) + + scalar_out = net.scalar_act(scalar_pair) + scalar_gate = torch.sigmoid(net.scalar_gate(scalar_pair)) + coeff_packed = product * scalar_gate[:, None, None, :, :] + coeff_packed[:, 0, net.frame_zero_index].add_(scalar_out) + coeff = _frame_contract_packed(net.frame_contract, coeff_packed) + if net.residual_scale is not None: + coeff = coeff * net.residual_scale.view(1, 1, N_FOCUS, CHANNELS) + output = coeff.reshape_as(query_flat) + if return_product: + return output, product_flat + return output + + +def run_packed_message_grid_backward( + net: Any, + query_flat: Tensor, + context_flat: Tensor, + grad_out_flat: Tensor, + *, + product_flat: Tensor | None = None, +) -> tuple[Tensor, Tensor]: + """Return query/context adjoints while retaining packed grid intermediates.""" + query, context = _validate_contract(net, query_flat, context_flat) + from .k1_kernels.cute_neo_message_grid_product import ( + run_message_grid_product, + run_message_grid_product_backward, + ) + + nodes = query_flat.shape[0] + scalar_pair = torch.cat([query[:, 0], context[:, 0]], dim=-1).to(net.dtype) + + left_packed = _frame_expand_packed(net.frame_expand, query) + right_packed = _frame_expand_packed(net.frame_expand, context) + left = left_packed.view(nodes, COEFF_DIM * N_FRAMES, HIDDEN_CHANNELS) + right = right_packed.view(nodes, COEFF_DIM * N_FRAMES, HIDDEN_CHANNELS) + if product_flat is None: + product_flat = run_message_grid_product( + left, + right, + net.projector.to_grid_mat, + net.projector.from_grid_mat, + ) + elif ( + tuple(product_flat.shape) != (nodes, COEFF_DIM * N_FRAMES, HIDDEN_CHANNELS) + or product_flat.dtype != torch.float32 + or product_flat.device != query_flat.device + or not product_flat.is_contiguous() + ): + raise ValueError( + "saved packed message-grid product must be contiguous FP32 with shape " + f"({nodes}, {COEFF_DIM * N_FRAMES}, {HIDDEN_CHANNELS}) on " + f"{query_flat.device}" + ) + product = product_flat.view( + nodes, + COEFF_DIM, + N_FRAMES, + N_FOCUS, + CHANNELS, + ) + + scalar_gate = torch.sigmoid(net.scalar_gate(scalar_pair)) + + grad = _as_flat_ndfc( + net, + "grad_out", + grad_out_flat, + like=query_flat, + ).to(net.dtype) + if net.residual_scale is not None: + grad = grad * net.residual_scale.view(1, 1, N_FOCUS, CHANNELS) + grad_scalar_packed = _frame_contract_packed_backward(net.frame_contract, grad) + + grad_product = grad_scalar_packed * scalar_gate[:, None, None, :, :] + grad_scalar_gate = (grad_scalar_packed * product).sum(dim=(1, 2)) + grad_scalar_out = grad_scalar_packed[:, 0, net.frame_zero_index] + grad_scalar_logits = grad_scalar_gate * scalar_gate * (1.0 - scalar_gate) + grad_scalar_pair = _focus_linear_backward_input( + net.scalar_gate, + grad_scalar_logits, + ) + _swiglu_backward_input(scalar_pair, grad_scalar_out) + + # Broadcast multiplication preserves the degree-major einsum layout. The + # packed CuTe adjoint requires coefficient-major compact storage. + grad_product_flat = grad_product.contiguous().view( + nodes, + COEFF_DIM * N_FRAMES, + HIDDEN_CHANNELS, + ) + grad_left, grad_right = run_message_grid_product_backward( + grad_product_flat, + left, + right, + net.projector.to_grid_mat, + net.projector.from_grid_mat, + ) + grad_query = _frame_expand_packed_backward( + net.frame_expand, + grad_left.view(nodes, COEFF_DIM, N_FRAMES, N_FOCUS, CHANNELS), + ) + grad_context = _frame_expand_packed_backward( + net.frame_expand, + grad_right.view(nodes, COEFF_DIM, N_FRAMES, N_FOCUS, CHANNELS), + ) + grad_query[:, 0].add_(grad_scalar_pair[:, :, :CHANNELS]) + grad_context[:, 0].add_(grad_scalar_pair[:, :, CHANNELS:]) + return ( + grad_query.reshape_as(query_flat).to(dtype=query_flat.dtype), + grad_context.reshape_as(context_flat).to(dtype=context_flat.dtype), + ) + + +__all__ = [ + "run_packed_message_grid_backward", + "run_packed_message_grid_forward", +] diff --git a/deepmd/kernels/cute/neo/k1_radial_phase_a_node.py b/deepmd/kernels/cute/neo/k1_radial_phase_a_node.py new file mode 100644 index 0000000000..f9384fbdfa --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_radial_phase_a_node.py @@ -0,0 +1,395 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Runtime wrapper for the source-CSR node-tiled radial/Phase-A backward.""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, + Any, +) + +from .compile_cache import ( + device_aware_lru_cache, +) + +if TYPE_CHECKING: + from torch import ( + Tensor, + ) + + +DEGREE_COUNT = 16 +REDUCED_COUNT = 10 +HIDDEN = 64 +FOCUS_COUNT = 2 +FOCUS_HIDDEN = 32 +PACKED_WIGNER_VALUES = 46 +RADIAL_WIDTH = 4 * FOCUS_HIDDEN +COMPACT_WIDTH = 25 +PROJECTION_INPUT_WIDTH = COMPACT_WIDTH + FOCUS_COUNT + + +@dataclass(frozen=True) +class NeoSourceCSR: + """Indirect source CSR over the unchanged physical edge order.""" + + source_order: Tensor + source_ptr: Tensor + + +@dataclass(frozen=True) +class NeoRadialPhaseABackwardNodeResult: + """Output buffers populated by the node-tiled backward.""" + + grad_x_wide: Tensor + grad_d_full: Tensor + grad_radial_m0: Tensor + + +def build_source_csr( + src: Tensor, + node_count: int, + *, + validate_sources: bool = False, +) -> 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. + """ + import torch + + if src.dim() != 1: + raise ValueError("src must be one-dimensional") + if src.dtype not in (torch.int32, torch.int64): + raise TypeError("src must have dtype int32 or int64") + if node_count < 0: + raise ValueError("node_count must be non-negative") + if src.numel() > 2**31 - 1: + raise ValueError("source CSR int32 indexing requires E <= 2**31 - 1") + + src = src.contiguous() + if validate_sources and 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)" + ) + + source_order_i64 = torch.argsort(src, stable=True) + sorted_src = src.index_select(0, source_order_i64) + boundaries = torch.arange( + node_count + 1, + device=src.device, + dtype=src.dtype, + ) + source_ptr = torch.searchsorted( + sorted_src, + boundaries, + out_int32=True, + ).contiguous() + source_order = source_order_i64.to(dtype=torch.int32).contiguous() + return NeoSourceCSR(source_order=source_order, source_ptr=source_ptr) + + +@device_aware_lru_cache(maxsize=4) +def _compile_node_tiled() -> Any: + from .k1_kernels.cute_neo_radial_phase_a_backward_node import ( + compile_neo_radial_phase_a_backward_node_tiled, + ) + + return compile_neo_radial_phase_a_backward_node_tiled() + + +def _expect_tensor( + name: str, + tensor: Tensor, + shape: tuple[int, ...], + *, + device: Any, + dtype: Any, +) -> None: + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}") + if tensor.dtype != dtype: + raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + + +def prepare_batched_radial_projection_weight( + combined_weight: Tensor, + attention_radial_weight: Tensor, +) -> Tensor: + """Precombine the compact and attention adjoints for one FP32 GEMM.""" + import torch + + device = combined_weight.device + _expect_tensor( + "combined_weight", + combined_weight, + (RADIAL_WIDTH, COMPACT_WIDTH), + device=device, + dtype=torch.float32, + ) + _expect_tensor( + "attention_radial_weight", + attention_radial_weight, + (FOCUS_HIDDEN, FOCUS_COUNT), + device=device, + dtype=torch.float32, + ) + + projection_weight = combined_weight.new_zeros( + (PROJECTION_INPUT_WIDTH, RADIAL_WIDTH) + ) + projection_weight[:COMPACT_WIDTH].copy_(combined_weight.transpose(0, 1)) + projection_weight[COMPACT_WIDTH:, :FOCUS_HIDDEN].copy_( + attention_radial_weight.transpose(0, 1) + ) + return projection_weight + + +def _project_batched_radial_adjoint( + projection_weight: Tensor, + grad_radial_m0: Tensor, + consumed_workspace: Tensor, +) -> None: + """Project the 27-column adjoint packed in consumed edge scratch.""" + import torch + + edge_count = consumed_workspace.shape[0] + device = consumed_workspace.device + tensors = ( + ( + "projection_weight", + projection_weight, + (PROJECTION_INPUT_WIDTH, RADIAL_WIDTH), + ), + ("grad_radial_m0", grad_radial_m0, (edge_count, RADIAL_WIDTH)), + ( + "consumed_workspace", + consumed_workspace, + (edge_count, FOCUS_COUNT * REDUCED_COUNT * FOCUS_HIDDEN), + ), + ) + for name, tensor, shape in tensors: + _expect_tensor( + name, + tensor, + shape, + device=device, + dtype=torch.float32, + ) + + projection_input = consumed_workspace[:, :PROJECTION_INPUT_WIDTH] + torch.mm(projection_input, projection_weight, out=grad_radial_m0) + + +def _validate_csr_values( + source_order: Tensor, + source_ptr: Tensor, + edge_count: int, +) -> None: + import torch + + valid = (source_ptr[0] == 0) & (source_ptr[-1] == edge_count) + if source_ptr.numel() > 1: + valid = valid & torch.all(source_ptr[1:] >= source_ptr[:-1]) + if not bool(valid): + raise ValueError( + "source_ptr must be nondecreasing, begin at zero, and end at E" + ) + expected = torch.arange( + edge_count, + device=source_order.device, + dtype=source_order.dtype, + ) + if not torch.equal(torch.sort(source_order).values, expected): + raise ValueError("source_order must be a permutation of [0, E)") + + +def run_neo_radial_phase_a_backward_node_tiled( + grad_out_focus: Tensor, + grad_logits: Tensor, + radial_state: Tensor, + channel_basis: Tensor, + x_wide: Tensor, + source_order: Tensor, + source_ptr: Tensor, + d_full: Tensor, + *, + grad_focus_src_focus: Tensor, + batched_radial_projection_weight: Tensor, + grad_x_wide: Tensor | None = None, + grad_d_full: Tensor | None = None, + grad_radial_m0: Tensor | None = None, + validate_csr: bool = False, +) -> NeoRadialPhaseABackwardNodeResult: + """Run the node-owned backward over indirect source CSR. + + 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. + """ + import torch + + if not x_wide.is_cuda: + raise ValueError("node-tiled radial Phase-A backward requires CUDA tensors") + if x_wide.dtype != torch.float32: + raise TypeError("node-tiled radial Phase-A backward specializes float32") + if not x_wide.is_contiguous() or x_wide.dim() != 2: + raise ValueError("x_wide must be a contiguous two-dimensional tensor") + + device = x_wide.device + dtype = x_wide.dtype + node_count = x_wide.shape[0] + edge_count = grad_out_focus.shape[0] + _expect_tensor( + "x_wide", + x_wide, + (node_count, DEGREE_COUNT * HIDDEN), + device=device, + dtype=dtype, + ) + _expect_tensor( + "grad_out_focus", + grad_out_focus, + (edge_count, FOCUS_COUNT * REDUCED_COUNT * FOCUS_HIDDEN), + device=device, + dtype=dtype, + ) + _expect_tensor( + "grad_focus_src_focus", + grad_focus_src_focus, + (FOCUS_COUNT, edge_count, FOCUS_HIDDEN), + device=device, + dtype=dtype, + ) + _expect_tensor( + "grad_logits", + grad_logits, + (edge_count, FOCUS_COUNT), + device=device, + dtype=dtype, + ) + _expect_tensor( + "radial_state", + radial_state, + (edge_count, COMPACT_WIDTH), + device=device, + dtype=dtype, + ) + _expect_tensor( + "batched_radial_projection_weight", + batched_radial_projection_weight, + (PROJECTION_INPUT_WIDTH, RADIAL_WIDTH), + device=device, + dtype=dtype, + ) + _expect_tensor( + "channel_basis", + channel_basis, + (HIDDEN,), + device=device, + dtype=dtype, + ) + _expect_tensor( + "source_order", + source_order, + (edge_count,), + device=device, + dtype=torch.int32, + ) + _expect_tensor( + "source_ptr", + source_ptr, + (node_count + 1,), + device=device, + dtype=torch.int32, + ) + _expect_tensor( + "d_full", + d_full, + (edge_count, PACKED_WIGNER_VALUES), + device=device, + dtype=dtype, + ) + if node_count == 0 and edge_count != 0: + raise ValueError("a non-empty edge list requires at least one source node") + if validate_csr: + _validate_csr_values(source_order, source_ptr, edge_count) + + if grad_x_wide is None: + grad_x_wide = torch.empty_like(x_wide) + else: + _expect_tensor( + "grad_x_wide", + grad_x_wide, + (node_count, DEGREE_COUNT * HIDDEN), + device=device, + dtype=dtype, + ) + if grad_d_full is None: + grad_d_full = torch.empty_like(d_full) + else: + _expect_tensor( + "grad_d_full", + grad_d_full, + (edge_count, PACKED_WIGNER_VALUES), + device=device, + dtype=dtype, + ) + if grad_radial_m0 is None: + grad_radial_m0 = torch.empty( + (edge_count, RADIAL_WIDTH), + device=device, + dtype=dtype, + ) + else: + _expect_tensor( + "grad_radial_m0", + grad_radial_m0, + (edge_count, RADIAL_WIDTH), + device=device, + dtype=dtype, + ) + + result = NeoRadialPhaseABackwardNodeResult( + grad_x_wide=grad_x_wide, + grad_d_full=grad_d_full, + grad_radial_m0=grad_radial_m0, + ) + if edge_count == 0: + grad_x_wide.zero_() + return result + + with torch.cuda.device(device): + kernel = _compile_node_tiled() + kernel( + grad_out_focus, + grad_focus_src_focus, + grad_logits, + radial_state, + channel_basis, + x_wide, + source_order, + source_ptr, + d_full, + grad_x_wide, + grad_d_full, + ) + _project_batched_radial_adjoint( + batched_radial_projection_weight, + grad_radial_m0, + grad_out_focus, + ) + return result diff --git a/deepmd/kernels/cute/neo/k1_runner.py b/deepmd/kernels/cute/neo/k1_runner.py new file mode 100644 index 0000000000..c74256b972 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_runner.py @@ -0,0 +1,895 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Runtime runner for the Neo CuTe K1 SO2 unit.""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + Any, +) + +from .runtime_policy import ( + FUSED_SO2_GATE_CAPABILITIES, + SM80_PROFILE_CAPABILITIES, + SM90_CAPABILITY, + SUPPORTED_K1_CAPABILITIES, +) + + +def _uses_packed_message_grid(compute_capability: tuple[int, int]) -> bool: + """Return whether the packed message-grid kernels support this GPU.""" + return ( + compute_capability in SM80_PROFILE_CAPABILITIES + or compute_capability == SM90_CAPABILITY + ) + + +def _validate_runtime_config( + runtime_config: Any, + *, + compute_capability: tuple[int, int] | None = None, +) -> None: + if compute_capability not in SUPPORTED_K1_CAPABILITIES: + 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 != ( + compute_capability in SM80_PROFILE_CAPABILITIES + ): + raise RuntimeError("the per-focus SO2 path must be selected on sm_80/sm_86") + if runtime_config.combined_so2_gate != ( + compute_capability in FUSED_SO2_GATE_CAPABILITIES + ): + raise RuntimeError( + "the combined SO2/gate path must be selected on sm_89/sm_120" + ) + + +def _combined_radial_weight( + torch: Any, + radial_hidden_proj: Any, + radial_degree_mixer: Any, +) -> Any: + hidden_weight = radial_hidden_proj.weight.detach() + mixer_weight = radial_degree_mixer.weight.detach() + cache_key = ( + hidden_weight.data_ptr(), + hidden_weight._version, + mixer_weight.data_ptr(), + mixer_weight._version, + hidden_weight.dtype, + hidden_weight.device, + ) + cache = getattr(radial_degree_mixer, "_deepmd_cute_combined_weight", None) + if cache is not None and cache[0] == cache_key: + return cache[1] + blocks = [] + for degree in range(4): + mixer_block = mixer_weight[degree * 64 : (degree + 1) * 64, :] + blocks.append(torch.mm(hidden_weight, mixer_block)) + combined = torch.cat(blocks, dim=0).contiguous() + radial_degree_mixer._deepmd_cute_combined_weight = (cache_key, combined) + return combined + + +def _combined_attention_radial_weight( + torch: Any, + radial_hidden_proj: Any, + attention_weight: Any, +) -> Any: + hidden_weight = radial_hidden_proj.weight.detach() + attention_weight = attention_weight.detach() + cache_key = ( + hidden_weight.data_ptr(), + hidden_weight._version, + attention_weight.data_ptr(), + attention_weight._version, + hidden_weight.dtype, + hidden_weight.device, + ) + cache = getattr(radial_hidden_proj, "_deepmd_cute_attention_radial_weight", None) + if cache is not None and cache[0] == cache_key: + return cache[1] + blocks = [] + for focus in range(2): + hidden_block = hidden_weight[:, focus * 32 : (focus + 1) * 32] + blocks.append(torch.mv(hidden_block, attention_weight[:, focus, 0])) + combined = torch.stack(blocks, dim=1).contiguous() + radial_hidden_proj._deepmd_cute_attention_radial_weight = (cache_key, combined) + return combined + + +def _batched_radial_projection_weight( + radial_degree_mixer: Any, + combined_weight: Any, + attention_radial_weight: Any, +) -> Any: + cache_key = ( + combined_weight.data_ptr(), + combined_weight._version, + attention_radial_weight.data_ptr(), + attention_radial_weight._version, + combined_weight.dtype, + combined_weight.device, + ) + cache = getattr( + radial_degree_mixer, + "_deepmd_cute_batched_radial_projection_weight", + None, + ) + if cache is not None and cache[0] == cache_key: + return cache[1] + + from .k1_radial_phase_a_node import ( + prepare_batched_radial_projection_weight, + ) + + projection_weight = prepare_batched_radial_projection_weight( + combined_weight, + attention_radial_weight, + ) + # Keep the derived operands alive so allocator pointer reuse cannot spoof + # the versioned cache key after a parameter update. + radial_degree_mixer._deepmd_cute_batched_radial_projection_weight = ( + cache_key, + projection_weight, + combined_weight, + attention_radial_weight, + ) + return projection_weight + + +def _edge_gate(torch: Any, edge_cache: Any) -> Any: + gate = edge_cache.edge_env.reshape(-1).float().clamp_min(0.0) + if edge_cache.edge_src_gate is not None: + gate = gate * edge_cache.edge_src_gate.reshape(-1).float().clamp_min(0.0).sqrt() + return gate.contiguous() + + +@dataclass +class StackCache: + y: Any + logits: Any | None + non_linear: Any + final: bool + + def __post_init__(self) -> None: + self.y = self.y.detach() + if self.logits is not None: + self.logits = self.logits.detach() + + +class _NeoK1BackwardWorkspaceBase: + _EXPORTED_NAMES = ( + "grad_stack_focus", + "grad_focus_alpha", + "grad_dt", + "grad_alpha", + "grad_logits", + "grad_edge", + "grad_z_partial", + "grad_z", + "grad_x_rot", + "grad_radial_flat", + "grad_x_wide_phase_a", + "grad_d", + "grad_y", + "grad_gate_logits", + "grad_mixed_slab", + ) + + def attach_to(self, runner: Any) -> None: + for name in self._EXPORTED_NAMES: + setattr(runner, name, getattr(self, name)) + + +def _structural_memory_views( + torch: Any, + *, + edge_count: int, + like: Any, + phase_c_stack: Any, + phase_c_y: Any | None, + radial_scratch: Any, + radial_values_per_edge: int, + phase_c_single_input_reuse: bool, +) -> tuple[Any, Any | None, Any]: + """Validate and expose saved stack buffers that become backward scratch.""" + expected_phase_shape = (edge_count, 2, 10, 32) + if phase_c_single_input_reuse: + if phase_c_y is not None: + raise ValueError( + "structural memory reuse requires phase_c_y to be absent in " + "single-input mode" + ) + phase_tensors = (("phase_c_stack", phase_c_stack, expected_phase_shape),) + else: + phase_tensors = ( + ("phase_c_stack", phase_c_stack, expected_phase_shape), + ("phase_c_y", phase_c_y, expected_phase_shape), + ) + for name, tensor, expected_shape in phase_tensors: + if tensor is None or tuple(tensor.shape) != expected_shape: + raise ValueError( + f"structural memory reuse requires {name} shape {expected_shape}" + ) + if tensor.dtype != torch.float32 or tensor.dtype != like.dtype: + raise ValueError(f"structural memory reuse requires FP32 {name} storage") + if tensor.device != like.device or not tensor.is_contiguous(): + raise ValueError( + f"structural memory reuse requires contiguous {name} on {like.device}" + ) + if tensor.storage_offset() != 0: + raise ValueError( + f"structural memory reuse requires zero-offset {name} storage" + ) + required_radial_values = edge_count * radial_values_per_edge + if radial_scratch is None or radial_scratch.numel() < required_radial_values: + raise ValueError( + "structural memory reuse requires radial_scratch capacity of at least " + f"{required_radial_values} values" + ) + if radial_scratch.dtype != torch.float32 or radial_scratch.dtype != like.dtype: + raise ValueError("structural memory reuse requires FP32 radial_scratch storage") + if radial_scratch.device != like.device or not radial_scratch.is_contiguous(): + raise ValueError( + "structural memory reuse requires contiguous radial_scratch on " + f"{like.device}" + ) + if radial_scratch.storage_offset() != 0: + raise ValueError( + "structural memory reuse requires zero-offset radial_scratch storage" + ) + tensors = (*[tensor for _, tensor, _ in phase_tensors], radial_scratch) + storages = {tensor.untyped_storage()._cdata for tensor in tensors} + if len(storages) != len(tensors): + raise ValueError("structural memory reuse requires distinct storages") + return ( + phase_c_stack.view(edge_count, 10 * 64), + None if phase_c_y is None else phase_c_y.view(edge_count, 10 * 64), + radial_scratch.flatten(), + ) + + +class NeoK1BackwardWorkspace(_NeoK1BackwardWorkspaceBase): + """Lazily allocated K1 backward scratch with lifetime-based slab reuse.""" + + def __init__( + self, + torch: Any, + *, + edge_count: int, + node_count: int, + d_full: Any, + dt_full: Any, + radial: Any, + phase_c_stack: Any | None = None, + phase_c_y: Any | None = None, + radial_scratch: Any | None = None, + structural_memory_reuse: bool = False, + phase_c_single_input_reuse: bool = False, + ) -> None: + opts = {"device": d_full.device, "dtype": d_full.dtype} + d_values_per_edge = 1 + for size in d_full.shape[1:]: + d_values_per_edge *= size + radial_values_per_edge = 1 + for size in radial.shape[1:]: + radial_values_per_edge *= size + if d_values_per_edge > 10 * 64: + raise ValueError("Neo K1 grad_D does not fit the Phase-C scratch slab") + + if phase_c_single_input_reuse and not structural_memory_reuse: + raise ValueError( + "Phase-C single-input reuse requires structural memory reuse" + ) + + phase_c_stack_flat = None + phase_c_y_flat = None + radial_scratch_flat = None + if structural_memory_reuse: + phase_c_stack_flat, phase_c_y_flat, radial_scratch_flat = ( + _structural_memory_views( + torch, + edge_count=edge_count, + like=d_full, + phase_c_stack=phase_c_stack, + phase_c_y=phase_c_y, + radial_scratch=radial_scratch, + radial_values_per_edge=radial_values_per_edge, + phase_c_single_input_reuse=phase_c_single_input_reuse, + ) + ) + + self._phase_c_or_d = ( + phase_c_stack_flat + if phase_c_stack_flat is not None + else torch.empty( + edge_count, + 10 * 64, + device=opts["device"], + dtype=opts["dtype"], + ) + ) + self.grad_stack_focus = self._phase_c_or_d.view(edge_count, 2, 10, 32) + self.grad_d = self._phase_c_or_d.flatten()[ + : edge_count * d_values_per_edge + ].view_as(d_full) + self.grad_dt = torch.empty_like(dt_full) + + # Gate backward completes before radial backward writes the final radial grad. + if radial_scratch_flat is not None: + self._gate_or_radial = radial_scratch_flat + else: + self._gate_or_radial = torch.empty( + edge_count, + radial_values_per_edge, + device=opts["device"], + dtype=opts["dtype"], + ) + self.grad_gate_logits = None + self.grad_radial_flat = self._gate_or_radial.flatten()[ + : edge_count * radial_values_per_edge + ].view(edge_count, radial_values_per_edge) + + # The per-layer gate output is dead before radial backward writes grad_x_rot. + self._gate_or_x_rot = ( + phase_c_stack_flat + if phase_c_stack_flat is not None + else torch.empty( + edge_count, + 10 * 64, + device=opts["device"], + dtype=opts["dtype"], + ) + ) + self.grad_y = self._gate_or_x_rot.view(edge_count * 2, 10 * 32) + self.grad_x_rot = self._gate_or_x_rot + if phase_c_stack_flat is None: + self.grad_mixed_slab = None + elif phase_c_y_flat is not None: + self.grad_mixed_slab = phase_c_y_flat.view(edge_count, 2, 10, 32) + else: + self.grad_mixed_slab = torch.empty( + edge_count, + 2, + 10, + 32, + device=opts["device"], + dtype=opts["dtype"], + ) + + self.grad_focus_alpha = torch.empty( + edge_count, 2, device=opts["device"], dtype=opts["dtype"] + ) + self.grad_alpha = torch.empty( + edge_count, 2, device=opts["device"], dtype=opts["dtype"] + ) + self.grad_logits = torch.empty( + edge_count, 2, device=opts["device"], dtype=opts["dtype"] + ) + self.grad_edge = torch.empty( + edge_count, device=opts["device"], dtype=opts["dtype"] + ) + self.grad_z_partial = torch.empty( + node_count, 2, device=opts["device"], dtype=opts["dtype"] + ) + self.grad_z = torch.empty(2, device=opts["device"], dtype=opts["dtype"]) + self.grad_x_wide_phase_a = torch.empty( + node_count, + 16 * 64, + device=opts["device"], + dtype=opts["dtype"], + ) + + +class NeoFullCuteBackward: + def __init__( + self, + torch: Any, + block: Any, + record: Any, + x: Any, + d_full: Any, + dt_full: Any, + radial_feat: Any, + dst_ptr: Any, + source_order: Any, + source_ptr: Any, + *, + runtime_config: Any, + ) -> None: + from .k1_kernels.cute_envelope_gated_softmax import ( + compile_envelope_softmax_forward, + ) + from .k1_kernels.cute_neo_phase_a_radial_forward import ( + run_neo_phase_a_radial_forward_packed_direct, + ) + from .k1_so2linear import ( + run_neo_so2_linear_manual, + ) + from .k1_wigner_layout import ( + PACKED_VALUE_COUNT, + ) + + self.torch = torch + self.block = block + self.so2 = block.so2_conv + self.record = record + self.config = runtime_config + device_index = x.device.index + if device_index is None: + device_index = torch.cuda.current_device() + compute_capability = tuple(torch.cuda.get_device_capability(device_index)) + self.device_index = device_index + self.compute_capability = compute_capability + self.packed_message_grid = _uses_packed_message_grid(compute_capability) + self.compile_identity = (device_index, *compute_capability) + _validate_runtime_config( + runtime_config, + compute_capability=compute_capability, + ) + self.x = x + self.d = d_full + self.dt = dt_full + expected_shape = ( + record.edge_cache.src.numel(), + PACKED_VALUE_COUNT, + ) + if ( + tuple(d_full.shape) != expected_shape + or tuple(dt_full.shape) != expected_shape + ): + raise ValueError( + f"packed Neo K1 Wigner tensors must have shape {expected_shape}" + ) + if d_full.data_ptr() != dt_full.data_ptr(): + raise ValueError("packed Neo K1 D and Dt must share one storage") + self.radial = radial_feat + self.dst_ptr_i32 = dst_ptr.to(device=x.device, dtype=torch.int32).contiguous() + self.edge_count = record.edge_cache.src.numel() + self.node_count = x.shape[0] + self.src_i32 = record.edge_cache.src.to(torch.int32).contiguous() + self.src_i64 = record.edge_cache.src.contiguous() + self.dst_i32 = record.edge_cache.dst.to(torch.int32).contiguous() + self.dst_i64 = record.edge_cache.dst.contiguous() + if ( + source_order.numel() == self.edge_count + and source_ptr.numel() == self.node_count + 1 + ): + self.source_order_i32 = source_order.to( + device=x.device, + dtype=torch.int32, + ).contiguous() + self.source_ptr_i32 = source_ptr.to( + device=x.device, + dtype=torch.int32, + ).contiguous() + else: + from .k1_radial_phase_a_node import ( + build_source_csr, + ) + + source_csr = build_source_csr(self.src_i32, self.node_count) + self.source_order_i32 = source_csr.source_order + self.source_ptr_i32 = source_csr.source_ptr + self.rotate = self.so2.rotate_inv_rescale_full.contiguous() + self.edge_gate = _edge_gate(torch, record.edge_cache) + from .k1_kernels.cute_neo_focus_src_backward import ( + compile_neo_attention_prelude_forward, + ) + + with torch.cuda.device(self.device_index): + self.attention_prelude_forward = compile_neo_attention_prelude_forward( + float(self.so2.focus_compete_norm.eps), + float(self.so2.attn_qk_norm.eps), + float(self.so2.focus_softmax_tau), + float(self.so2.focus_label_smoothing), + self.compile_identity, + ) + from .k1_kernels.cute_neo_qk_edge import ( + compile_neo_qk_edge_backward, + compile_neo_qk_edge_forward, + ) + + qk_scale = float(self.so2.head_dim**-0.5) + with torch.cuda.device(self.device_index): + self.qk_edge_forward = compile_neo_qk_edge_forward( + qk_scale, + self.compile_identity, + ) + self.qk_edge_backward = compile_neo_qk_edge_backward( + qk_scale, + self.compile_identity, + ) + from .k1_kernels.cute_neo_qk_edge import ( + compile_neo_qk_node_input_adjoint, + ) + + with torch.cuda.device(self.device_index): + self.qk_node_input_adjoint = compile_neo_qk_node_input_adjoint( + float(self.so2.attn_qk_norm.eps), + self.compile_identity, + ) + from .k1_kernels.cute_neo_phase_c_backward_layout_runner import ( + CuteNeoPhaseCBackwardLayout, + ) + + self.phase_c_layout_backward = CuteNeoPhaseCBackwardLayout( + focus_eps=float(self.so2.focus_compete_norm.eps), + focus_tau=float(self.so2.focus_softmax_tau), + focus_label_smoothing=float(self.so2.focus_label_smoothing), + ) + with torch.cuda.device(self.device_index): + self.softmax_fwd = compile_envelope_softmax_forward( + 128, + float(self.so2.eps), + ) + self.structural_gate_forward = None + self.structural_gate_backward = None + self.combined_gate_backward = None + self._focus_major_gate_linear_forward = None + self._focus_major_gate_linear_backward_add = None + self._run_structural_gate_forward = None + self._run_structural_gate_backward = None + if runtime_config.combined_so2_gate: + from .k1_kernels.cute_neo_gate_linear_residual_backward_fused import ( + compile_neo_gate_linear_residual_backward_fused, + ) + + with torch.cuda.device(self.device_index): + self.combined_gate_backward = ( + compile_neo_gate_linear_residual_backward_fused() + ) + elif not runtime_config.native_sm90_path: + from .k1_gate_structural import ( + focus_major_gate_linear_backward_add_, + focus_major_gate_linear_forward, + run_structural_gate_backward, + run_structural_gate_forward, + ) + from .k1_kernels.cute_neo_gate_split_structural_vec4_sm80 import ( + compile_neo_gate_split_structural_vec4_sm80_backward, + compile_neo_gate_split_structural_vec4_sm80_forward, + ) + + self.structural_gate_forward = ( + compile_neo_gate_split_structural_vec4_sm80_forward( + self.compile_identity, + ) + ) + self.structural_gate_backward = ( + compile_neo_gate_split_structural_vec4_sm80_backward( + self.compile_identity, + ) + ) + self._focus_major_gate_linear_forward = focus_major_gate_linear_forward + self._focus_major_gate_linear_backward_add = ( + focus_major_gate_linear_backward_add_ + ) + self._run_structural_gate_forward = run_structural_gate_forward + self._run_structural_gate_backward = run_structural_gate_backward + opts = {"device": x.device, "dtype": x.dtype} + self._backward_workspace = None + self.alpha = torch.empty( + self.edge_count, 2, device=opts["device"], dtype=opts["dtype"] + ) + self.group_max = torch.empty( + self.node_count, 2, device=opts["device"], dtype=opts["dtype"] + ) + self.denom = torch.empty( + self.node_count, 2, device=opts["device"], dtype=opts["dtype"] + ) + + self._run_cute_phase_a_radial_forward = ( + run_neo_phase_a_radial_forward_packed_direct + ) + self._run_neo_so2_linear_manual = run_neo_so2_linear_manual + + self.combined_radial = _combined_radial_weight( + torch, self.so2.radial_hidden_proj, self.so2.radial_degree_mixer + ) + self.combined_attention_radial = _combined_attention_radial_weight( + torch, + self.so2.radial_hidden_proj, + self.so2.adamw_attn_logit_w, + ) + self.batched_radial_projection_weight = _batched_radial_projection_weight( + self.so2.radial_degree_mixer, + self.combined_radial, + self.combined_attention_radial, + ) + + self._build_forward_graph() + + def ensure_backward_workspace( + self, + ) -> NeoK1BackwardWorkspace: + if self._backward_workspace is None: + 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 + ) + if self.phase_c_stack.device.type == "cuda": + stream = self.torch.cuda.current_stream(self.phase_c_stack.device) + self.phase_c_stack.record_stream(stream) + radial_scratch.record_stream(stream) + workspace = NeoK1BackwardWorkspace( + self.torch, + edge_count=self.edge_count, + node_count=self.node_count, + d_full=self.d, + dt_full=self.dt, + radial=self.radial, + phase_c_stack=self.phase_c_stack, + phase_c_y=None, + radial_scratch=radial_scratch, + structural_memory_reuse=True, + phase_c_single_input_reuse=True, + ) + workspace.attach_to(self) + self._backward_workspace = workspace + return self._backward_workspace + + def _build_forward_graph(self) -> None: + torch = self.torch + so2 = self.so2 + block = self.block + n_node = self.node_count + n_edge = self.edge_count + self.structural_scratch = None + if self.config.combined_so2_gate: + self.structural_scratch = self.x.new_empty(self.radial.shape) + + use_full_node = block.node_lmax == block.lmax + self.use_full_node = use_full_node + x_so2 = self.x if use_full_node else self.x[:, : block.mp_ebed_dim, :, :] + x_pre = block.pre_so2_norm(x_so2) + self.x_wide = so2.pre_focus_mix( + x_pre.reshape(n_node, x_so2.shape[1], block.channels).unsqueeze(2) + ).squeeze(2) + + cur, rad_l0, radial_compact = self._run_cute_phase_a_radial_forward( + radial_hidden_proj=so2.radial_hidden_proj, + radial_degree_mixer=so2.radial_degree_mixer, + x_wide=self.x_wide.detach(), + src=self.src_i32, + D_full=self.d.detach(), + radial_feat_m0=self.radial.detach(), + ) + self.radial_compact = radial_compact.detach() + + self.focus_gate_src = cur[:, :, 0, :].detach().contiguous() + self.stack_caches: list[StackCache] = [] + + if self.config.combined_so2_gate: + from .k1_kernels.cute_neo_so2_gate_combined_fwd import ( + CuteNeoSO2GateCombinedFwdRunner, + prepare_neo_so2_gate_combined_weights, + ) + from .k1_so2linear import ( + cached_neo_so2_linear_weights, + ) + + stack_layers = zip( + so2.so2_linears, + so2.so2_inter_norms, + so2.non_linearities, + strict=True, + ) + for layer_idx, (so2_linear, _inter_norm, non_linear) in enumerate(stack_layers): + x_layer = cur.detach() + final = layer_idx == so2.mixing_layers - 1 + logits = None + if not final and self.config.combined_so2_gate: + w0, wpair = cached_neo_so2_linear_weights(so2_linear) + gate_parameter = non_linear.gate_linear.weight + gate_weight = gate_parameter.detach().view(32, 2, 3 * 32).contiguous() + pack_key = ( + so2_linear.weight_m0.data_ptr(), + so2_linear.weight_m0._version, + so2_linear.weight_m[0].data_ptr(), + so2_linear.weight_m[0]._version, + gate_parameter.data_ptr(), + gate_parameter._version, + x_layer.device, + ) + pack_cache = getattr( + so2_linear, + "_deepmd_cute_neo_combined_gate_weights", + None, + ) + if ( + not isinstance(pack_cache, tuple) + or len(pack_cache) != 6 + or pack_cache[0] != pack_key + ): + with torch.cuda.device(x_layer.device): + pack_stream = torch.cuda.current_stream(x_layer.device) + packed_weights = prepare_neo_so2_gate_combined_weights( + w0, + wpair, + gate_weight, + ) + ready_event = torch.cuda.Event() + ready_event.record(pack_stream) + packed_weights_ready = ( + ready_event, + pack_stream.cuda_stream, + ) + pack_cache = ( + pack_key, + packed_weights, + packed_weights_ready, + # Retain source parameters so allocator pointer reuse + # cannot spoof the versioned cache key. + so2_linear.weight_m0, + so2_linear.weight_m[0], + gate_parameter, + ) + so2_linear._deepmd_cute_neo_combined_gate_weights = pack_cache + else: + packed_weights = pack_cache[1] + packed_weights_ready = pack_cache[2] + y = torch.empty_like(x_layer) + # Each CTA reads its residual tile before storing it, and the + # m=0 and m>0 regions are disjoint. Reuse x_layer for output to + # retain the optimized two-buffer stack footprint. + combined_forward = CuteNeoSO2GateCombinedFwdRunner( + x_layer, + x_layer, + y, + x_layer, + packed_weights=packed_weights, + packed_weights_ready=packed_weights_ready, + ) + cur = combined_forward() + else: + y = self._run_neo_so2_linear_manual( + so2_linear, + x_layer, + add_residual=final, + per_focus_pair=self.config.per_focus_so2_fwd_pair, + ) + if not final and not self.config.combined_so2_gate: + gate_src = y[:, :, 0, :] + logits = self._focus_major_gate_linear_forward( + gate_src, + non_linear.gate_linear.weight.detach(), + ) + self._run_structural_gate_forward( + self.structural_gate_forward, + x_layer, + y, + logits, + out=x_layer, + ) + cur = x_layer + elif final: + self.phase_c_stack = y.detach() + self.phase_c_y = None + cur = None + self.stack_caches.append( + StackCache( + y=y, + logits=logits, + non_linear=non_linear, + final=final, + ) + ) + + x_wide_qk = self.x_wide.detach() + rad_l0_qk = rad_l0.detach().view(n_edge, 2, 32) + x_l0_node = x_wide_qk[:, 0, :].reshape(n_node, 2, 32) + focus_alpha = torch.empty( + n_edge, + 2, + device=self.focus_gate_src.device, + dtype=torch.float32, + ) + q_node = torch.empty_like( + x_l0_node, + memory_format=torch.contiguous_format, + ) + k_node = torch.empty_like( + x_l0_node, + memory_format=torch.contiguous_format, + ) + self.attention_prelude_forward( + self.focus_gate_src.view(n_edge, 64), + x_l0_node.contiguous(), + so2.adamw_focus_compete_w.detach().float().contiguous(), + so2.focus_compete_norm.adam_scale.detach().float().contiguous(), + so2.attn_q_proj.weight.detach().float().view(32, 2, 32).contiguous(), + so2.attn_k_proj.weight.detach().float().view(32, 2, 32).contiguous(), + so2.attn_qk_norm.adam_scale.detach().float().contiguous(), + focus_alpha, + q_node, + k_node, + ) + self.focus_alpha = focus_alpha.detach() + self.q_node = q_node.detach() + self.k_node = k_node.detach() + self.attn_logits = torch.empty( + n_edge, + 2, + device=self.q_node.device, + dtype=self.q_node.dtype, + ) + self.qk_edge_forward( + self.q_node, + self.k_node, + rad_l0_qk.contiguous(), + so2.adamw_attn_logit_w.detach().contiguous(), + self.src_i32, + self.dst_i32, + self.attn_logits, + ) + + self.softmax_fwd( + self.attn_logits.detach().contiguous(), + self.edge_gate, + self.dst_ptr_i32, + so2.adamw_attn_z_bias_raw.detach().reshape(2).float().contiguous(), + self.alpha, + self.group_max, + self.denom, + ) + x_wide_down = self.x_wide.detach() + from .k1_kernels.cute_neo_phase_c_onepass import ( + run_neo_phase_c_onepass_output_gate, + ) + + self.phase_c_out = run_neo_phase_c_onepass_output_gate( + x_local_flat=self.phase_c_stack, + Dt_full=self.dt.detach(), + alpha_focus=self.alpha, + focus_compete_alpha=self.focus_alpha, + dst_ptr=self.dst_ptr_i32, + rotate_inv_rescale=so2.rotate_inv_rescale_full, + x_wide=x_wide_down, + output_gate_norm_scale=so2.attn_output_gate_norm.adam_scale.detach() + .float() + .reshape(2, 32) + .contiguous(), + output_gate_weight=so2.adamw_attn_gate_w.detach() + .float() + .reshape(32, 2, 1) + .contiguous(), + output_gate_eps=float(so2.attn_output_gate_norm.eps), + ).to(dtype=so2.compute_dtype) + out = self.phase_c_out.detach().to(dtype=so2.dtype) + self.out_gate_flat = out.detach() + self.message_grid_product = None + if so2.message_node_grid_product is not None: + if self.packed_message_grid: + from .k1_message_grid_packed import ( + run_packed_message_grid_forward, + ) + + grid_out = run_packed_message_grid_forward( + so2.message_node_grid_product, + out, + x_wide_down, + ) + else: + grid_out = so2.message_node_grid_product(out, x_wide_down) + out = out + grid_out + self.post_mix_input = out.detach() + out = so2.post_focus_mix(out.unsqueeze(2)).squeeze(2) + self.post_norm_input = out.unsqueeze(2).detach() + so2_out = block.post_so2_norm(self.post_norm_input) + if use_full_node: + self.final = so2_out + else: + final = self.x.new_zeros(self.x.shape) + final[:, : block.mp_ebed_dim, :, :] = so2_out + self.final = final diff --git a/deepmd/kernels/cute/neo/k1_so2linear.py b/deepmd/kernels/cute/neo/k1_so2linear.py new file mode 100644 index 0000000000..8f35d4afd2 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_so2linear.py @@ -0,0 +1,261 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""SO2Linear helpers for the Neo CuTe K1 path. + +The in-place residual adjoint enforces FP32 operand dtypes. Strict FP32 GEMM +also requires the caller to select highest float32 matmul precision and disable +TF32; these helpers do not mutate process-wide backend settings. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import torch +from torch import ( + Tensor, +) + +FOCUS_COUNT = 2 +REDUCED_COUNT = 10 +CHANNELS = 32 +M0_WIDTH = 4 * 32 +PAIR_WIDTH = 6 * CHANNELS +FULL_WIDTH = REDUCED_COUNT * CHANNELS + + +def _validate_neo_so2_linear(so2_linear: Any) -> None: + if ( + so2_linear.lmax != 3 + or so2_linear.mmax != 1 + or so2_linear.in_channels != 32 + or so2_linear.out_channels != 32 + or so2_linear.n_focus != 2 + or so2_linear.mlp_bias + ): + raise NotImplementedError( + "Neo CuTe K1 SO2Linear expects lmax=3,mmax=1,F=2,C=32" + ) + + +def run_neo_so2_linear_manual( + so2_linear: Any, + x_local: Any, + *, + add_residual: bool = False, + per_focus_pair: bool = False, +) -> Any: + """Run the fixed Neo SO2Linear block with two dense focus batched GEMMs.""" + import torch + + _validate_neo_so2_linear(so2_linear) + if add_residual: + w0, wpair = cached_neo_so2_linear_residual_weights(so2_linear) + else: + w0, wpair = cached_neo_so2_linear_weights(so2_linear) + x_flat = x_local.reshape(x_local.shape[0], 2, 10 * 32).transpose(0, 1) + out = x_local.new_empty(x_local.shape[0], 2, 10 * 32) + out_t = out.transpose(0, 1) + torch.bmm(x_flat[:, :, : 4 * 32], w0, out=out_t[:, :, : 4 * 32]) + if per_focus_pair: + for focus in range(FOCUS_COUNT): + torch.mm( + x_flat[focus, :, M0_WIDTH:], + wpair[focus], + out=out_t[focus, :, M0_WIDTH:], + ) + else: + torch.bmm(x_flat[:, :, M0_WIDTH:], wpair, out=out_t[:, :, M0_WIDTH:]) + return out.reshape(x_local.shape[0], 2, 10, 32) + + +def cached_neo_so2_linear_residual_weights(so2_linear: Any) -> tuple[Any, Any]: + """Return dense weights with the fixed SO2 residual folded in.""" + w0, wpair = cached_neo_so2_linear_weights(so2_linear) + cache = getattr(so2_linear, "_deepmd_cute_neo_manual_residual_weights", None) + cache_key = (w0.data_ptr(), wpair.data_ptr(), w0.dtype, w0.device) + if ( + cache is not None + and cache[0] is w0 + and cache[1] is wpair + and cache[2] == cache_key + ): + return cache[3], cache[4] + + w0_residual = w0.clone() + wpair_residual = wpair.clone() + w0_residual.diagonal(dim1=-2, dim2=-1).add_(1.0) + wpair_residual.diagonal(dim1=-2, dim2=-1).add_(1.0) + so2_linear._deepmd_cute_neo_manual_residual_weights = ( + w0, + wpair, + cache_key, + w0_residual, + wpair_residual, + ) + return w0_residual, wpair_residual + + +def cached_neo_so2_linear_weights(so2_linear: Any) -> tuple[Any, Any]: + """Return cached dense block weights for Neo's fixed SO2Linear layout.""" + import torch + + cache = getattr(so2_linear, "_deepmd_cute_neo_manual_weights", None) + cache_key = ( + so2_linear.weight_m0.data_ptr(), + so2_linear.weight_m[0].data_ptr(), + so2_linear.weight_m0._version, + so2_linear.weight_m[0]._version, + so2_linear.weight_m0.dtype, + so2_linear.weight_m0.device, + ) + if ( + cache is not None + and cache[0] is so2_linear.weight_m0 + and cache[1] is so2_linear.weight_m[0] + and cache[2] == cache_key + ): + return cache[3], cache[4] + + w0 = so2_linear.weight_m0.detach().view(4 * 32, 2, 4 * 32) + w0 = w0.permute(1, 0, 2).contiguous() + raw_pair = so2_linear.weight_m[0].detach().view(3 * 32, 2, 2 * 3 * 32) + w_u = raw_pair[:, :, : 3 * 32] + w_v = raw_pair[:, :, 3 * 32 :] + wpair = torch.empty( + 2, + 2 * 3 * 32, + 2 * 3 * 32, + device=raw_pair.device, + dtype=raw_pair.dtype, + ) + wpair[:, : 3 * 32, : 3 * 32] = w_u.permute(1, 0, 2) + wpair[:, : 3 * 32, 3 * 32 :] = w_v.permute(1, 0, 2) + wpair[:, 3 * 32 :, : 3 * 32] = -w_v.permute(1, 0, 2) + wpair[:, 3 * 32 :, 3 * 32 :] = w_u.permute(1, 0, 2) + wpair = wpair.contiguous() + so2_linear._deepmd_cute_neo_manual_weights = ( + so2_linear.weight_m0, + so2_linear.weight_m[0], + cache_key, + w0, + wpair, + ) + return w0, wpair + + +def _has_direct_cublas_layout(tensor: Tensor) -> bool: + """Check only the direct-cuBLAS stride layouts used by PyTorch 2.10. + + This is a stride-layout predicate, not a complete cuBLAS eligibility test; + dtype/device requirements are validated separately. + """ + if tensor.ndim != 2: + return False + rows, columns = tensor.shape + stride0, stride1 = tensor.stride() + return (stride0 == 1 and stride1 >= max(1, rows)) or ( + stride1 == 1 and stride0 >= max(1, columns) + ) + + +def _contiguous_tensors_overlap(lhs: Tensor, rhs: Tensor) -> bool: + """Return exact byte overlap for validated, nonempty contiguous tensors.""" + if lhs.device != rhs.device: + return False + if lhs.device.type == "meta": + return torch._C._overlaps(lhs, rhs) + + lhs_start = lhs.data_ptr() + rhs_start = rhs.data_ptr() + lhs_stop = lhs_start + lhs.numel() * lhs.element_size() + rhs_stop = rhs_start + rhs.numel() * rhs.element_size() + return lhs_start < rhs_stop and rhs_start < lhs_stop + + +def _validate_edge_focus(tensor: Tensor, name: str) -> None: + if tensor.dtype != torch.float32: + raise TypeError(f"{name} must be float32, got {tensor.dtype}") + if tensor.ndim != 4 or tuple(tensor.shape[1:]) != ( + FOCUS_COUNT, + REDUCED_COUNT, + CHANNELS, + ): + raise ValueError( + f"{name} must have shape (E,2,10,32), got {tuple(tensor.shape)}" + ) + if tensor.shape[0] <= 0: + raise ValueError(f"{name} requires E > 0") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must use the canonical contiguous K1 layout") + + +def _validate_weight( + weight: Tensor, + name: str, + width: int, + device: torch.device, +) -> None: + if weight.dtype != torch.float32: + raise TypeError(f"{name} must be float32, got {weight.dtype}") + expected_shape = (FOCUS_COUNT, width, width) + if tuple(weight.shape) != expected_shape: + raise ValueError( + f"{name} must have shape {expected_shape}, got {tuple(weight.shape)}" + ) + if weight.device != device: + raise ValueError(f"{name} must be on {device}, got {weight.device}") + if not weight.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + + +def neo_so2_linear_backward_residual_inplace( + residual: Tensor, + grad_out: Tensor, + w0_t: Tensor, + wpair_t: Tensor, +) -> Tensor: + """Overwrite a dead residual with ``grad_out @ W.T + residual``.""" + _validate_edge_focus(residual, "residual") + _validate_edge_focus(grad_out, "grad_out") + if residual.shape != grad_out.shape: + raise ValueError("residual and grad_out shapes must match") + if residual.device != grad_out.device: + raise ValueError("residual and grad_out devices must match") + if _contiguous_tensors_overlap(residual, grad_out): + raise ValueError( + "residual and grad_out must not alias; keep the final layer out-of-place" + ) + + _validate_weight(w0_t, "w0_t", M0_WIDTH, residual.device) + _validate_weight(wpair_t, "wpair_t", PAIR_WIDTH, residual.device) + for name, weight in (("w0_t", w0_t), ("wpair_t", wpair_t)): + if _contiguous_tensors_overlap(residual, weight): + raise ValueError(f"residual and {name} must not alias") + + edge_count = residual.shape[0] + residual_flat = residual.view(edge_count, FOCUS_COUNT, FULL_WIDTH) + grad_flat = grad_out.view(edge_count, FOCUS_COUNT, FULL_WIDTH) + for focus in range(FOCUS_COUNT): + for start, stop, weight in ( + (0, M0_WIDTH, w0_t[focus]), + (M0_WIDTH, FULL_WIDTH, wpair_t[focus]), + ): + residual_block = residual_flat[:, focus, start:stop] + grad_block = grad_flat[:, focus, start:stop] + if not _has_direct_cublas_layout(residual_block) or not ( + _has_direct_cublas_layout(grad_block) + and _has_direct_cublas_layout(weight) + ): + raise ValueError("SO2 block layout would require cuBLAS staging") + residual_block.addmm_( + grad_block, + weight, + beta=1.0, + alpha=1.0, + ) + return residual diff --git a/deepmd/kernels/cute/neo/k1_wigner_layout.py b/deepmd/kernels/cute/neo/k1_wigner_layout.py new file mode 100644 index 0000000000..0283f23a2f --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_wigner_layout.py @@ -0,0 +1,93 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Fixed packed Wigner layout for the Neo ``lmax=3, mmax=1`` K1 path. + +The panel stores only rows selected by ``coeff_index_m`` and only columns from +the matching Wigner block. Phase A reads ``D[coeff, degree]`` while Phase C +reads ``Dt[degree, coeff]``; both expressions therefore address the same slot. +""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, +) + +if TYPE_CHECKING: + from collections.abc import ( + Iterator, + ) + + +BLOCK_WIDTHS = (1, 3, 5, 7) +FULL_BLOCK_OFFSETS = (0, 1, 4, 9, 16) + +# Rows are ordered m=0, m=-1, m=+1 within each non-scalar block. +SELECTED_LOCAL_ROWS = ((0,), (1, 0, 2), (2, 1, 3), (3, 2, 4)) +PANEL_BLOCK_OFFSETS = (0, 1, 10, 25, 46) +PACKED_VALUE_COUNT = PANEL_BLOCK_OFFSETS[-1] + +# DeePMD's m-major reduced ordering for lmax=3, mmax=1. +COEFF_INDEX_M = (0, 2, 6, 12, 1, 5, 11, 3, 7, 13) +REDUCED_DEGREES = (0, 1, 2, 3, 1, 2, 3, 1, 2, 3) +REDUCED_PANEL_ROW_OFFSETS = (0, 1, 10, 25, 4, 15, 32, 7, 20, 39) +ZONAL_PANEL_OFFSETS = tuple(range(1, 4)) + tuple(range(10, 15)) + tuple(range(25, 32)) + + +@dataclass(frozen=True) +class PackedWignerEntry: + offset: int + degree: int + reduced: int + full_row: int + full_col: int + + +_REDUCED_BY_FULL_ROW = { + full_row: reduced for reduced, full_row in enumerate(COEFF_INDEX_M) +} + + +def d_offset(reduced: int, full_col: int) -> int | None: + """Map ``D[coeff_index_m[reduced], full_col]`` into the packed panel.""" + degree = REDUCED_DEGREES[reduced] + block_start = FULL_BLOCK_OFFSETS[degree] + block_stop = FULL_BLOCK_OFFSETS[degree + 1] + if full_col < block_start or full_col >= block_stop: + return None + return REDUCED_PANEL_ROW_OFFSETS[reduced] + full_col - block_start + + +def dt_offset(full_row: int, reduced: int) -> int | None: + """Map ``Dt[full_row, coeff_index_m[reduced]]`` to its shared D slot.""" + return d_offset(reduced, full_row) + + +def iter_packed_entries() -> Iterator[PackedWignerEntry]: + """Yield the 46 stored entries in contiguous panel order.""" + for degree, (width, block_start, panel_start, local_rows) in enumerate( + zip( + BLOCK_WIDTHS, + FULL_BLOCK_OFFSETS[:-1], + PANEL_BLOCK_OFFSETS[:-1], + SELECTED_LOCAL_ROWS, + strict=True, + ) + ): + for row_slot, local_row in enumerate(local_rows): + full_row = block_start + local_row + reduced = _REDUCED_BY_FULL_ROW[full_row] + row_start = panel_start + row_slot * width + for local_col in range(width): + yield PackedWignerEntry( + offset=row_start + local_col, + degree=degree, + reduced=reduced, + full_row=full_row, + full_col=block_start + local_col, + ) diff --git a/deepmd/kernels/cute/neo/k4_wignerd.py b/deepmd/kernels/cute/neo/k4_wignerd.py new file mode 100644 index 0000000000..49ee535200 --- /dev/null +++ b/deepmd/kernels/cute/neo/k4_wignerd.py @@ -0,0 +1,1004 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +# pyright: reportMissingImports=false +# ruff: noqa: ANN001, ANN201, ANN202, ANN204, TC002, UP035 +"""Packed CuTe Wigner-D panel for the Neo K1 inference path.""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + Any, + Callable, +) + +import cutlass +import cutlass.cute as cute +import torch +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from . import ( + runtime_policy, +) +from .compile_cache import ( + device_aware_lru_cache, +) +from .k1_wigner_layout import PACKED_VALUE_COUNT as K1_PANEL_VALUES + +L2_SPARSE_TERMS = 10 +L3_SPARSE_TERMS = 20 + + +def _load_dpa4_wignerd_calculator(): + """Load DeePMD's WignerDCalculator to reuse its coefficient tables.""" + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + WignerDCalculator, + ) + + return WignerDCalculator + + +def _sparsify_rows( + coeffs: torch.Tensor, max_terms: int, threshold: float = 1.0e-12 +) -> tuple[torch.Tensor, torch.Tensor]: + values = torch.zeros( + coeffs.shape[0], max_terms, device=coeffs.device, dtype=coeffs.dtype + ) + indices = torch.zeros( + coeffs.shape[0], max_terms, device=coeffs.device, dtype=torch.int32 + ) + for row in range(coeffs.shape[0]): + nz = torch.nonzero(coeffs[row].abs() > threshold, as_tuple=False).flatten() + if nz.numel() > max_terms: + raise RuntimeError( + f"sparse row {row} has {nz.numel()} terms, max_terms={max_terms}" + ) + values[row, : nz.numel()] = coeffs[row, nz] + indices[row, : nz.numel()] = nz.to(torch.int32) + return values.contiguous(), indices.contiguous() + + +def _build_l2_l3_tables( + dtype: torch.dtype, device: torch.device +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + calc_cls = _load_dpa4_wignerd_calculator() + try: + cache = calc_cls._get_small_order_cache_cpu_fp64(3) + except TypeError: + cache = calc_cls._get_small_order_cache_cpu_fp64() + c_l2 = cache["C_l2"] + monomials_l2 = calc_cls._generate_monomials(4, 4) + c_l2_flat = torch.zeros( + 25, + len(monomials_l2), + device=c_l2.device, + dtype=c_l2.dtype, + ) + for mono_idx, exponents in enumerate(monomials_l2): + for a in range(4): + for b in range(4): + for c in range(4): + for d in range(4): + counts = ( + int(a == 0) + int(b == 0) + int(c == 0) + int(d == 0), + int(a == 1) + int(b == 1) + int(c == 1) + int(d == 1), + int(a == 2) + int(b == 2) + int(c == 2) + int(d == 2), + int(a == 3) + int(b == 3) + int(c == 3) + int(d == 3), + ) + if counts == exponents: + c_l2_flat[:, mono_idx] += c_l2[:, :, a, b, c, d].reshape(25) + exp_l2 = torch.tensor( + monomials_l2, + device=c_l2.device, + dtype=torch.int32, + ) + c_l2_sparse, c_l2_sparse_idx = _sparsify_rows(c_l2_flat, L2_SPARSE_TERMS) + c_l3_sparse, c_l3_sparse_idx = _sparsify_rows(cache["C_l3"], L3_SPARSE_TERMS) + return ( + exp_l2.to(device=device, dtype=torch.int32).contiguous(), + c_l2_sparse.to(device=device, dtype=dtype).contiguous(), + c_l2_sparse_idx.to(device=device, dtype=torch.int32).contiguous(), + cache["exp_l3"].to(device=device, dtype=torch.int32).contiguous(), + c_l3_sparse.to(device=device, dtype=dtype).contiguous(), + c_l3_sparse_idx.to(device=device, dtype=torch.int32).contiguous(), + ) + + +@dataclass(frozen=True) +class WignerDParams: + q: cute.Tensor + panel: cute.Tensor + exp_l2: cute.Tensor + c_l2_sparse: cute.Tensor + c_l2_sparse_idx: cute.Tensor + exp_l3: cute.Tensor + c_l3_sparse: cute.Tensor + c_l3_sparse_idx: cute.Tensor + + +class WignerDForward: + def __init__(self, threads: int, dtype): + if threads % 32 != 0: + raise ValueError("threads must be a multiple of 32") + self.threads = int(threads) + self.warps = threads // 32 + self.dtype = dtype + + @cute.jit + def rotmat(self, w, x, y, z, row, col): + two = self.dtype(2.0) + one = self.dtype(1.0) + value = self.dtype(0.0) + if row == 0 and col == 0: + value = one - two * (y * y + z * z) + elif row == 0 and col == 1: + value = two * (x * y - w * z) + elif row == 0 and col == 2: + value = two * (x * z + w * y) + elif row == 1 and col == 0: + value = two * (x * y + w * z) + elif row == 1 and col == 1: + value = one - two * (x * x + z * z) + elif row == 1 and col == 2: + value = two * (y * z - w * x) + elif row == 2 and col == 0: + value = two * (x * z - w * y) + elif row == 2 and col == 1: + value = two * (y * z + w * x) + elif row == 2 and col == 2: + value = one - two * (x * x + y * y) + return value + + @cute.jit + def perm(self, idx): + out = idx + if idx == 0: + out = 1 + elif idx == 1: + out = 2 + elif idx == 2: + out = 0 + return out + + @cute.jit + def sign(self, idx): + value = self.dtype(-1.0) + if idx == 2: + value = self.dtype(1.0) + return value + + @cute.jit + def d1(self, w, x, y, z, row, col): + r = self.rotmat(w, x, y, z, self.perm(row), self.perm(col)) + return r * self.sign(row) * self.sign(col) + + @cute.jit + def component(self, w, x, y, z, comp): + value = w + if comp == 1: + value = x + elif comp == 2: + value = y + elif comp == 3: + value = z + return value + + @cute.jit + def pow_small(self, base, exponent): + value = self.dtype(1.0) + if exponent >= 1: + value *= base + if exponent >= 2: + value *= base + if exponent >= 3: + value *= base + if exponent >= 4: + value *= base + if exponent >= 5: + value *= base + if exponent >= 6: + value *= base + return value + + @cute.jit + def monomial_l3(self, exp_l3, mono, w, x, y, z): + value = self.dtype(1.0) + for comp in cutlass.range_constexpr(4): + value *= self.pow_small( + self.component(w, x, y, z, comp), + exp_l3[mono, comp], + ) + return value + + @cute.kernel + def kernel_warp_edges_panel(self, params: WignerDParams): + tidx, _, _ = cute.arch.thread_idx() + edge_block, _, _ = cute.arch.block_idx() + lane = tidx % 32 + warp = tidx // 32 + edge_count, _ = params.q.shape + edge = edge_block * self.warps + warp + + smem = cutlass.utils.SmemAllocator() + l2_monomials = smem.allocate_tensor(self.dtype, self.warps * 35) + l3_monomials = smem.allocate_tensor(self.dtype, self.warps * 84) + + if edge < edge_count: + qw = params.q[edge, 0].to(self.dtype) + qx = params.q[edge, 1].to(self.dtype) + qy = params.q[edge, 2].to(self.dtype) + qz = params.q[edge, 3].to(self.dtype) + inv_norm = cute.rsqrt( + qw * qw + qx * qx + qy * qy + qz * qz + self.dtype(1.0e-14) + ) + qw = qw * inv_norm + qx = qx * inv_norm + qy = qy * inv_norm + qz = qz * inv_norm + + if lane == 0: + params.panel[edge, 0] = self.dtype(1.0).to(params.panel.element_type) + + for flat in cutlass.range(lane, 9, 32, unroll=1): + row_slot = flat // 3 + col = flat - row_slot * 3 + row = cutlass.Int32(1) + if row_slot == 1: + row = cutlass.Int32(0) + elif row_slot == 2: + row = cutlass.Int32(2) + value = self.d1(qw, qx, qy, qz, row, col) + params.panel[edge, 1 + flat] = value.to(params.panel.element_type) + + l2_base = warp * 35 + for mono in cutlass.range(lane, 35, 32, unroll=1): + l2_monomials[l2_base + mono] = self.monomial_l3( + params.exp_l2, mono, qw, qx, qy, qz + ) + cute.arch.sync_warp() + for flat in cutlass.range(lane, 15, 32, unroll=1): + row_slot = flat // 5 + col = flat - row_slot * 5 + row = cutlass.Int32(2) + if row_slot == 1: + row = cutlass.Int32(1) + elif row_slot == 2: + row = cutlass.Int32(3) + block_flat = row * 5 + col + value = self.dtype(0.0) + for term in cutlass.range_constexpr(L2_SPARSE_TERMS): + mono = params.c_l2_sparse_idx[block_flat, term] + value += ( + params.c_l2_sparse[block_flat, term].to(self.dtype) + * l2_monomials[l2_base + mono] + ) + params.panel[edge, 10 + flat] = value.to(params.panel.element_type) + + l3_base = warp * 84 + for mono in cutlass.range(lane, 84, 32, unroll=1): + l3_monomials[l3_base + mono] = self.monomial_l3( + params.exp_l3, mono, qw, qx, qy, qz + ) + cute.arch.sync_warp() + for flat in cutlass.range(lane, 21, 32, unroll=1): + row_slot = flat // 7 + col = flat - row_slot * 7 + row = cutlass.Int32(3) + if row_slot == 1: + row = cutlass.Int32(2) + elif row_slot == 2: + row = cutlass.Int32(4) + block_flat = row * 7 + col + value = self.dtype(0.0) + for term in cutlass.range_constexpr(L3_SPARSE_TERMS): + mono = params.c_l3_sparse_idx[block_flat, term] + value += ( + params.c_l3_sparse[block_flat, term].to(self.dtype) + * l3_monomials[l3_base + mono] + ) + params.panel[edge, 25 + flat] = value.to(params.panel.element_type) + + +@dataclass(frozen=True) +class WignerDBwdParams: + q: cute.Tensor + grad_panel: cute.Tensor + grad_q: cute.Tensor + exp_l2: cute.Tensor + c_l2_sparse: cute.Tensor + c_l2_sparse_idx: cute.Tensor + exp_l3: cute.Tensor + c_l3_sparse: cute.Tensor + c_l3_sparse_idx: cute.Tensor + + +class WignerDBackward: + def __init__(self, threads: int, dtype): + if threads % 32 != 0: + raise ValueError("threads must be a multiple of 32") + self.threads = int(threads) + self.warps = threads // 32 + self.dtype = dtype + + @cute.jit + def warp_sum(self, value): + return cute.arch.warp_reduction_sum(value) + + @cute.jit + def cta_sum(self, value, scratch, tidx): + lane = tidx % 32 + warp = tidx // 32 + value = self.warp_sum(value) + if lane == 0: + scratch[warp] = value + cute.arch.sync_threads() + + total = self.dtype(0.0) + if tidx < self.warps: + total = scratch[tidx] + total = self.warp_sum(total) + if tidx == 0: + scratch[0] = total + cute.arch.sync_threads() + return scratch[0] + + @cute.jit + def rotmat_grad(self, w, x, y, z, row, col, comp): + two = self.dtype(2.0) + four = self.dtype(4.0) + value = self.dtype(0.0) + if row == 0 and col == 0: + if comp == 2: + value = -four * y + elif comp == 3: + value = -four * z + elif row == 0 and col == 1: + if comp == 0: + value = -two * z + elif comp == 1: + value = two * y + elif comp == 2: + value = two * x + elif comp == 3: + value = -two * w + elif row == 0 and col == 2: + if comp == 0: + value = two * y + elif comp == 1: + value = two * z + elif comp == 2: + value = two * w + elif comp == 3: + value = two * x + elif row == 1 and col == 0: + if comp == 0: + value = two * z + elif comp == 1: + value = two * y + elif comp == 2: + value = two * x + elif comp == 3: + value = two * w + elif row == 1 and col == 1: + if comp == 1: + value = -four * x + elif comp == 3: + value = -four * z + elif row == 1 and col == 2: + if comp == 0: + value = -two * x + elif comp == 1: + value = -two * w + elif comp == 2: + value = two * z + elif comp == 3: + value = two * y + elif row == 2 and col == 0: + if comp == 0: + value = -two * y + elif comp == 1: + value = two * z + elif comp == 2: + value = -two * w + elif comp == 3: + value = two * x + elif row == 2 and col == 1: + if comp == 0: + value = two * x + elif comp == 1: + value = two * w + elif comp == 2: + value = two * z + elif comp == 3: + value = two * y + elif row == 2 and col == 2: + if comp == 1: + value = -four * x + elif comp == 2: + value = -four * y + return value + + @cute.jit + def perm(self, idx): + out = idx + if idx == 0: + out = 1 + elif idx == 1: + out = 2 + elif idx == 2: + out = 0 + return out + + @cute.jit + def sign(self, idx): + value = self.dtype(-1.0) + if idx == 2: + value = self.dtype(1.0) + return value + + @cute.jit + def d1_grad(self, w, x, y, z, row, col, comp): + r = self.rotmat_grad(w, x, y, z, self.perm(row), self.perm(col), comp) + return r * self.sign(row) * self.sign(col) + + @cute.jit + def component(self, w, x, y, z, comp): + value = w + if comp == 1: + value = x + elif comp == 2: + value = y + elif comp == 3: + value = z + return value + + @cute.jit + def pow_small(self, base, exponent): + value = self.dtype(1.0) + if exponent >= 1: + value *= base + if exponent >= 2: + value *= base + if exponent >= 3: + value *= base + if exponent >= 4: + value *= base + if exponent >= 5: + value *= base + if exponent >= 6: + value *= base + return value + + @cute.jit + def monomial_l3_grad(self, exp_l3, mono, w, x, y, z, comp): + exp_comp = exp_l3[mono, comp] + value = self.dtype(0.0) + if exp_comp > 0: + value = exp_comp.to(self.dtype) + for c in cutlass.range_constexpr(4): + exp_c = exp_l3[mono, c] + if c == comp: + exp_c = exp_c - 1 + value *= self.pow_small(self.component(w, x, y, z, c), exp_c) + return value + + @cute.kernel + def kernel_sparse_panel(self, params: WignerDBwdParams): + tidx, _, _ = cute.arch.thread_idx() + edge, _, _ = cute.arch.block_idx() + + smem = cutlass.utils.SmemAllocator() + scratch = smem.allocate_tensor(self.dtype, self.warps) + dmono_l2 = smem.allocate_tensor(self.dtype, 4 * 35) + dmono_l3 = smem.allocate_tensor(self.dtype, 4 * 84) + + raw_w = params.q[edge, 0].to(self.dtype) + raw_x = params.q[edge, 1].to(self.dtype) + raw_y = params.q[edge, 2].to(self.dtype) + raw_z = params.q[edge, 3].to(self.dtype) + inv_norm = cute.rsqrt( + raw_w * raw_w + + raw_x * raw_x + + raw_y * raw_y + + raw_z * raw_z + + self.dtype(1.0e-14) + ) + w = raw_w * inv_norm + x = raw_x * inv_norm + y = raw_y * inv_norm + z = raw_z * inv_norm + + for idx in cutlass.range(tidx, 4 * 35, self.threads, unroll=1): + comp = idx // 35 + mono = idx - comp * 35 + dmono_l2[idx] = self.monomial_l3_grad(params.exp_l2, mono, w, x, y, z, comp) + for idx in cutlass.range(tidx, 4 * 84, self.threads, unroll=1): + comp = idx // 84 + mono = idx - comp * 84 + dmono_l3[idx] = self.monomial_l3_grad(params.exp_l3, mono, w, x, y, z, comp) + cute.arch.sync_threads() + + gw_local = self.dtype(0.0) + gx_local = self.dtype(0.0) + gy_local = self.dtype(0.0) + gz_local = self.dtype(0.0) + + for flat in cutlass.range(tidx, 9, self.threads, unroll=1): + row_slot = flat // 3 + col = flat - row_slot * 3 + row = cutlass.Int32(1) + if row_slot == 1: + row = cutlass.Int32(0) + elif row_slot == 2: + row = cutlass.Int32(2) + grad = params.grad_panel[edge, 1 + flat].to(self.dtype) + gw_local += grad * self.d1_grad(w, x, y, z, row, col, 0) + gx_local += grad * self.d1_grad(w, x, y, z, row, col, 1) + gy_local += grad * self.d1_grad(w, x, y, z, row, col, 2) + gz_local += grad * self.d1_grad(w, x, y, z, row, col, 3) + + for flat in cutlass.range(tidx, 15, self.threads, unroll=1): + row_slot = flat // 5 + col = flat - row_slot * 5 + row = cutlass.Int32(2) + if row_slot == 1: + row = cutlass.Int32(1) + elif row_slot == 2: + row = cutlass.Int32(3) + block_flat = row * 5 + col + grad = params.grad_panel[edge, 10 + flat].to(self.dtype) + for term in cutlass.range_constexpr(L2_SPARSE_TERMS): + mono = params.c_l2_sparse_idx[block_flat, term] + coeff = params.c_l2_sparse[block_flat, term].to(self.dtype) + gw_local += grad * coeff * dmono_l2[mono] + gx_local += grad * coeff * dmono_l2[35 + mono] + gy_local += grad * coeff * dmono_l2[70 + mono] + gz_local += grad * coeff * dmono_l2[105 + mono] + + for flat in cutlass.range(tidx, 21, self.threads, unroll=1): + row_slot = flat // 7 + col = flat - row_slot * 7 + row = cutlass.Int32(3) + if row_slot == 1: + row = cutlass.Int32(2) + elif row_slot == 2: + row = cutlass.Int32(4) + block_flat = row * 7 + col + grad = params.grad_panel[edge, 25 + flat].to(self.dtype) + for term in cutlass.range_constexpr(L3_SPARSE_TERMS): + mono = params.c_l3_sparse_idx[block_flat, term] + coeff = params.c_l3_sparse[block_flat, term].to(self.dtype) + gw_local += grad * coeff * dmono_l3[mono] + gx_local += grad * coeff * dmono_l3[84 + mono] + gy_local += grad * coeff * dmono_l3[168 + mono] + gz_local += grad * coeff * dmono_l3[252 + mono] + + gw = self.cta_sum(gw_local, scratch, tidx) + cute.arch.sync_threads() + gx = self.cta_sum(gx_local, scratch, tidx) + cute.arch.sync_threads() + gy = self.cta_sum(gy_local, scratch, tidx) + cute.arch.sync_threads() + gz = self.cta_sum(gz_local, scratch, tidx) + + if tidx == 0: + dot = gw * raw_w + gx * raw_x + gy * raw_y + gz * raw_z + inv3 = inv_norm * inv_norm * inv_norm + params.grad_q[edge, 0] = (inv_norm * gw - raw_w * inv3 * dot).to( + params.grad_q.element_type + ) + params.grad_q[edge, 1] = (inv_norm * gx - raw_x * inv3 * dot).to( + params.grad_q.element_type + ) + params.grad_q[edge, 2] = (inv_norm * gy - raw_y * inv3 * dot).to( + params.grad_q.element_type + ) + params.grad_q[edge, 3] = (inv_norm * gz - raw_z * inv3 * dot).to( + params.grad_q.element_type + ) + + +@cute.jit +def wignerd_panel_forward_warp_edges_jit( + q: cute.Tensor, + panel: cute.Tensor, + exp_l2: cute.Tensor, + c_l2_sparse: cute.Tensor, + c_l2_sparse_idx: cute.Tensor, + exp_l3: cute.Tensor, + c_l3_sparse: cute.Tensor, + c_l3_sparse_idx: cute.Tensor, + threads: cutlass.Constexpr[int], + stream: CUstream, +): + params = WignerDParams( + q=q, + panel=panel, + exp_l2=exp_l2, + c_l2_sparse=c_l2_sparse, + c_l2_sparse_idx=c_l2_sparse_idx, + exp_l3=exp_l3, + c_l3_sparse=c_l3_sparse, + c_l3_sparse_idx=c_l3_sparse_idx, + ) + edge_count, _ = q.shape + warps = threads // 32 + edge_blocks = cute.ceil_div(edge_count, warps) + WignerDForward(threads, cutlass.Float32).kernel_warp_edges_panel(params).launch( + grid=[edge_blocks, 1, 1], + block=[threads, 1, 1], + stream=stream, + ) + + +@cute.jit +def wignerd_panel_backward_jit( + q: cute.Tensor, + grad_panel: cute.Tensor, + grad_q: cute.Tensor, + exp_l2: cute.Tensor, + c_l2_sparse: cute.Tensor, + c_l2_sparse_idx: cute.Tensor, + exp_l3: cute.Tensor, + c_l3_sparse: cute.Tensor, + c_l3_sparse_idx: cute.Tensor, + threads: cutlass.Constexpr[int], + stream: CUstream, +): + params = WignerDBwdParams( + q=q, + grad_panel=grad_panel, + grad_q=grad_q, + exp_l2=exp_l2, + c_l2_sparse=c_l2_sparse, + c_l2_sparse_idx=c_l2_sparse_idx, + exp_l3=exp_l3, + c_l3_sparse=c_l3_sparse, + c_l3_sparse_idx=c_l3_sparse_idx, + ) + edge_count, _ = q.shape + WignerDBackward(threads, cutlass.Float32).kernel_sparse_panel(params).launch( + grid=[edge_count, 1, 1], + block=[threads, 1, 1], + stream=stream, + ) + + +def compile_wignerd_panel_forward( + threads: int, +) -> Callable: + if threads != 32: + raise ValueError("packed Wigner forward requires one warp per edge") + e = cute.sym_int64() + fake_q = make_fake_compact_tensor(cutlass.Float32, (e, 4), stride_order=(1, 0)) + fake_panel = make_fake_compact_tensor( + cutlass.Float32, + (e, K1_PANEL_VALUES), + stride_order=(1, 0), + ) + fake_exp_l2 = make_fake_compact_tensor(cutlass.Int32, (35, 4), stride_order=(1, 0)) + fake_c_l2_sparse = make_fake_compact_tensor( + cutlass.Float32, + (25, L2_SPARSE_TERMS), + stride_order=(1, 0), + ) + fake_c_l2_sparse_idx = make_fake_compact_tensor( + cutlass.Int32, + (25, L2_SPARSE_TERMS), + stride_order=(1, 0), + ) + fake_exp_l3 = make_fake_compact_tensor(cutlass.Int32, (84, 4), stride_order=(1, 0)) + fake_c_l3_sparse = make_fake_compact_tensor( + cutlass.Float32, + (49, L3_SPARSE_TERMS), + stride_order=(1, 0), + ) + fake_c_l3_sparse_idx = make_fake_compact_tensor( + cutlass.Int32, + (49, L3_SPARSE_TERMS), + stride_order=(1, 0), + ) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + wignerd_panel_forward_warp_edges_jit, + fake_q, + fake_panel, + fake_exp_l2, + fake_c_l2_sparse, + fake_c_l2_sparse_idx, + fake_exp_l3, + fake_c_l3_sparse, + fake_c_l3_sparse_idx, + threads, + fake_stream, + options="--enable-tvm-ffi", + ) + + +def compile_wignerd_panel_backward( + threads: int, +) -> Callable: + if threads % 32 != 0: + raise ValueError("packed Wigner backward threads must be warp-aligned") + e = cute.sym_int64() + fake_q = make_fake_compact_tensor(cutlass.Float32, (e, 4), stride_order=(1, 0)) + fake_grad_panel = make_fake_compact_tensor( + cutlass.Float32, + (e, K1_PANEL_VALUES), + stride_order=(1, 0), + ) + fake_grad_q = make_fake_compact_tensor(cutlass.Float32, (e, 4), stride_order=(1, 0)) + fake_exp_l2 = make_fake_compact_tensor(cutlass.Int32, (35, 4), stride_order=(1, 0)) + fake_c_l2_sparse = make_fake_compact_tensor( + cutlass.Float32, + (25, L2_SPARSE_TERMS), + stride_order=(1, 0), + ) + fake_c_l2_sparse_idx = make_fake_compact_tensor( + cutlass.Int32, + (25, L2_SPARSE_TERMS), + stride_order=(1, 0), + ) + fake_exp_l3 = make_fake_compact_tensor(cutlass.Int32, (84, 4), stride_order=(1, 0)) + fake_c_l3_sparse = make_fake_compact_tensor( + cutlass.Float32, + (49, L3_SPARSE_TERMS), + stride_order=(1, 0), + ) + fake_c_l3_sparse_idx = make_fake_compact_tensor( + cutlass.Int32, + (49, L3_SPARSE_TERMS), + stride_order=(1, 0), + ) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + wignerd_panel_backward_jit, + fake_q, + fake_grad_panel, + fake_grad_q, + fake_exp_l2, + fake_c_l2_sparse, + fake_c_l2_sparse_idx, + fake_exp_l3, + fake_c_l3_sparse, + fake_c_l3_sparse_idx, + threads, + fake_stream, + options="--enable-tvm-ffi", + ) + + +@device_aware_lru_cache(maxsize=16) +def _cached_wignerd_panel_forward(threads: int) -> Callable: + return compile_wignerd_panel_forward(threads) + + +@device_aware_lru_cache(maxsize=16) +def _cached_wignerd_panel_backward(threads: int) -> Callable: + return compile_wignerd_panel_backward(threads) + + +_TABLE_CACHE: dict[tuple[str, int], tuple[torch.Tensor, ...]] = {} + + +def _device_cache_key(device: torch.device) -> tuple[str, int]: + index = -1 if device.index is None else int(device.index) + return (device.type, index) + + +def _get_lmax3_tables(device: torch.device) -> tuple[torch.Tensor, ...]: + key = _device_cache_key(device) + tables = _TABLE_CACHE.get(key) + if tables is None: + tables = _build_l2_l3_tables(torch.float32, device) + _TABLE_CACHE[key] = tables + return tables + + +def _wignerd_panel_impl(edge_quat: torch.Tensor) -> torch.Tensor: + q = edge_quat.detach().contiguous() + if q.dtype != torch.float32: + raise TypeError(f"packed WignerD requires float32, got {q.dtype}") + + panel = torch.empty( + q.shape[0], + K1_PANEL_VALUES, + device=q.device, + dtype=q.dtype, + ) + if q.shape[0] == 0: + return panel + + ( + exp_l2, + c_l2_sparse, + c_l2_sparse_idx, + exp_l3, + c_l3_sparse, + c_l3_sparse_idx, + ) = _get_lmax3_tables(q.device) + with torch.cuda.device(q.device): + compiled = _cached_wignerd_panel_forward(32) + compiled( + q, + panel, + exp_l2, + c_l2_sparse, + c_l2_sparse_idx, + exp_l3, + c_l3_sparse, + c_l3_sparse_idx, + ) + return panel + + +def _wignerd_panel_bwd_impl( + grad_panel: torch.Tensor, + edge_quat: torch.Tensor, +) -> torch.Tensor: + q = edge_quat.detach().contiguous() + if tuple(grad_panel.shape) != (q.shape[0], K1_PANEL_VALUES): + raise ValueError( + f"packed Wigner gradient must have shape ({q.shape[0]}, {K1_PANEL_VALUES})" + ) + if q.dtype != torch.float32: + raise TypeError(f"packed WignerD requires float32, got {q.dtype}") + if q.shape[0] == 0: + return torch.empty_like(q) + + ( + exp_l2, + c_l2_sparse, + c_l2_sparse_idx, + exp_l3, + c_l3_sparse, + c_l3_sparse_idx, + ) = _get_lmax3_tables(q.device) + grad_q = torch.empty_like(q) + with torch.cuda.device(q.device): + compiled = _cached_wignerd_panel_backward(128) + compiled( + q, + grad_panel.detach().contiguous(), + grad_q, + exp_l2, + c_l2_sparse, + c_l2_sparse_idx, + exp_l3, + c_l3_sparse, + c_l3_sparse_idx, + ) + return grad_q + + +def _stateful_custom_op_tags() -> tuple[Any, ...] | None: + """Tag hidden-state runner ops as unsafe for direct CUDA graph capture.""" + tag_type = getattr(getattr(torch, "_C", None), "Tag", None) + cudagraph_unsafe = getattr(tag_type, "cudagraph_unsafe", None) + if cudagraph_unsafe is None: + return None + return (cudagraph_unsafe,) + + +_K4_CUSTOM_OP_TAGS = _stateful_custom_op_tags() +_wignerd_panel_op = torch.library.custom_op( + "sezm_cute::wignerd_k1_panel", mutates_args=(), tags=_K4_CUSTOM_OP_TAGS +)(_wignerd_panel_impl) +_wignerd_panel_bwd_op = torch.library.custom_op( + "sezm_cute::wignerd_k1_panel_bwd", mutates_args=(), tags=_K4_CUSTOM_OP_TAGS +)(_wignerd_panel_bwd_impl) + + +@_wignerd_panel_op.register_fake +def _(edge_quat: torch.Tensor) -> torch.Tensor: + return edge_quat.new_empty((edge_quat.shape[0], K1_PANEL_VALUES)) + + +@_wignerd_panel_bwd_op.register_fake +def _(grad_panel: torch.Tensor, edge_quat: torch.Tensor) -> torch.Tensor: + del grad_panel + return torch.empty_like(edge_quat) + + +def _tensor_compute_capability( + tensor: torch.Tensor, +) -> tuple[int, int] | None: + """Resolve dispatch capability from the Wigner operand's CUDA device.""" + if tensor.device.type != "cuda": + return None + return tuple(torch.cuda.get_device_capability(tensor.device)) + + +def _wignerd_panel_setup_context( + ctx: Any, + inputs: tuple, + output: torch.Tensor, +) -> None: + del output + (edge_quat,) = inputs + ctx.save_for_backward(edge_quat) + + +def _wignerd_panel_registered_backward_impl( + ctx: Any, + grad_panel: torch.Tensor, +) -> tuple[torch.Tensor]: + (edge_quat,) = ctx.saved_tensors + return (_wignerd_panel_bwd_op(grad_panel, edge_quat),) + + +def _wignerd_panel_backward( + ctx: Any, + grad_panel: torch.Tensor, +) -> tuple[torch.Tensor]: + """Run packed K4 backward with its custom op visible to compilation.""" + return _wignerd_panel_registered_backward_impl(ctx, grad_panel) + + +_wignerd_panel_op.register_autograd( + _wignerd_panel_backward, + setup_context=_wignerd_panel_setup_context, +) + + +def _run_cute_wignerd_impl( + edge_quat: torch.Tensor, + wigner_calc: Any, + *, + packed_wigner: bool = False, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Return CuTe Wigner data for the Neo ``lmax=3`` CUDA path. + + Unsupported shapes/devices return ``None`` so the caller can use DeePMD's + original implementation unchanged. A prevalidated strict-FP32 packed + request returns the same ``(E,46)`` panel object in both tuple positions. + """ + lmax = int(getattr(wigner_calc, "lmax", -1)) + if ( + not packed_wigner + or lmax != 3 + or edge_quat.dim() != 2 + or edge_quat.shape[-1] != 4 + or not edge_quat.is_cuda + or edge_quat.dtype != torch.float32 + or edge_quat.shape[0] == 0 + ): + return None + compute_capability = _tensor_compute_capability(edge_quat) + if compute_capability is None or not runtime_policy.is_packed_wigner_enabled( + compute_capability + ): + return None + panel = _wignerd_panel_op(edge_quat) + return panel, panel + + +def run_cute_wignerd( + edge_quat: torch.Tensor, + wigner_calc: Any, + *, + packed_wigner: bool = False, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Run K4 through its registered custom-op boundary.""" + return _run_cute_wignerd_impl( + edge_quat, + wigner_calc, + packed_wigner=packed_wigner, + ) diff --git a/deepmd/kernels/cute/neo/message_grid_gaunt_sm90.py b/deepmd/kernels/cute/neo/message_grid_gaunt_sm90.py new file mode 100644 index 0000000000..169279f209 --- /dev/null +++ b/deepmd/kernels/cute/neo/message_grid_gaunt_sm90.py @@ -0,0 +1,323 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Exact normalized-Gaunt message-grid product for the Neo SM90 path.""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, +) + +import cutlass +import cutlass.cute as cute +import cutlass.utils as cute_utils +import torch +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from .compile_cache import ( + device_aware_lru_cache, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN202, ANN204, TC002 + +COEFF_DIM = 48 +CHANNELS = 64 +THREADS = 256 +GROUPS = 4 +THREADS_PER_GROUP = CHANNELS +VALUES_PER_NODE = COEFF_DIM * CHANNELS +EXPECTED_COMPACT_PATHS = 1968 +EXPECTED_ORDERED_PATHS = 3833 +SUPPORT_GAP_THRESHOLD = 1.0e-4 +SM90_CAPABILITY = (9, 0) +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + +# (left row, right row, C[p,i,j], d[i]G[p,i,j], d[j]G[p,i,j]). +# The last value is zero on a diagonal path. +GauntTerm = tuple[int, int, float, float, float] + + +@dataclass(frozen=True, eq=False) +class Sm90GauntSchedule: + """Normalized-Gaunt rows captured as CuTe compile-time constants.""" + + rows: tuple[tuple[GauntTerm, ...], ...] + output_groups: tuple[tuple[int, ...], ...] + + def __post_init__(self) -> None: + if len(self.rows) != COEFF_DIM: + raise ValueError("Gaunt schedule must contain 48 rows") + if len(self.output_groups) != GROUPS: + raise ValueError("Gaunt schedule must contain four groups") + assigned = tuple(row for group in self.output_groups for row in group) + if tuple(sorted(assigned)) != tuple(range(COEFF_DIM)): + raise ValueError("each coefficient row must occur in exactly one group") + if sum(len(row) for row in self.rows) != EXPECTED_COMPACT_PATHS: + raise ValueError("Gaunt schedule must contain 1,968 paths") + for output_row, row in enumerate(self.rows): + for left_row, right_row, forward, adj_left, adj_right in row: + if not (0 <= left_row <= right_row < COEFF_DIM): + raise ValueError( + f"invalid compact path in row {output_row}: " + f"({left_row}, {right_row})" + ) + if forward == 0.0 or adj_left == 0.0: + raise ValueError("Gaunt paths must not contain zero weights") + if (left_row == right_row) != (adj_right == 0.0): + raise ValueError( + "only diagonal Gaunt paths may omit the mirrored adjoint" + ) + + +def build_sm90_gaunt_schedule( + to_grid: torch.Tensor, + from_grid: torch.Tensor, +) -> Sm90GauntSchedule: + """Recover and certify the fixed Neo Gaunt support from its projectors.""" + if ( + tuple(to_grid.shape) != (152, COEFF_DIM) + or tuple(from_grid.shape) != (COEFF_DIM, 152) + or to_grid.dtype != torch.float32 + or from_grid.dtype != torch.float32 + ): + raise ValueError("Neo Gaunt requires FP32 (152,48)/(48,152) projectors") + + to_cpu = to_grid.detach().to(device="cpu", dtype=torch.float64) + from_cpu = from_grid.detach().to(device="cpu", dtype=torch.float64) + tensor = torch.einsum("pg,gi,gj->pij", from_cpu, to_cpu, to_cpu) + support = tensor.abs() > SUPPORT_GAP_THRESHOLD + ordered_paths = int(support.sum()) + if ordered_paths != EXPECTED_ORDERED_PATHS: + raise ValueError( + "Neo projector support changed: expected " + f"{EXPECTED_ORDERED_PATHS} paths, got {ordered_paths}" + ) + minimum_signal = float(tensor.abs()[support].min()) + maximum_residual = float(tensor.abs()[~support].max()) + if minimum_signal <= 10.0 * SUPPORT_GAP_THRESHOLD: + raise ValueError("Neo Gaunt support no longer has a certified magnitude gap") + if maximum_residual >= SUPPORT_GAP_THRESHOLD: + raise ValueError("Neo Gaunt structural-zero residual exceeds its certificate") + + degree_weight = tuple( + 2 * degree + 1 + for degree in range(4) + for _order in range(-degree, degree + 1) + for _frame in range(3) + ) + weights = torch.tensor(degree_weight, device="cpu", dtype=torch.float64) + normalized = tensor / weights[:, None, None] + permutation_error = max( + float((normalized - normalized.permute(permutation)).abs().max()) + for permutation in ( + (0, 1, 2), + (0, 2, 1), + (1, 0, 2), + (1, 2, 0), + (2, 0, 1), + (2, 1, 0), + ) + ) + if permutation_error > 2.0e-6: + raise ValueError( + f"Neo normalized-Gaunt symmetry changed: max error {permutation_error:.3e}" + ) + + rows: list[tuple[GauntTerm, ...]] = [] + for output_row in range(COEFF_DIM): + terms: list[GauntTerm] = [] + for left_row in range(COEFF_DIM): + for right_row in range(left_row, COEFF_DIM): + if not bool(support[output_row, left_row, right_row]): + continue + forward = float(tensor[output_row, left_row, right_row].float()) + symmetric = normalized[output_row, left_row, right_row] + adj_left = float((symmetric * degree_weight[left_row]).float()) + adj_right = ( + 0.0 + if left_row == right_row + else float((symmetric * degree_weight[right_row]).float()) + ) + terms.append((left_row, right_row, forward, adj_left, adj_right)) + rows.append(tuple(terms)) + + group_rows: list[list[int]] = [[] for _ in range(GROUPS)] + group_loads = [0] * GROUPS + for output_row in sorted( + range(COEFF_DIM), + key=lambda row: (-len(rows[row]), row), + ): + group = min(range(GROUPS), key=lambda item: (group_loads[item], item)) + group_rows[group].append(output_row) + group_loads[group] += len(rows[output_row]) + if max(group_loads) - min(group_loads) > 32: + raise ValueError(f"Neo Gaunt static groups are imbalanced: {group_loads}") + return Sm90GauntSchedule( + tuple(rows), + tuple(tuple(sorted(group)) for group in group_rows), + ) + + +class _Sm90GauntForward: + def __init__(self, schedule: Sm90GauntSchedule) -> None: + self.schedule = schedule + + @cute.jit + def __call__( + self, + left: cute.Tensor, + right: cute.Tensor, + output: cute.Tensor, + stream: CUstream, + ): + operand_layout = cute.make_layout( + (2, COEFF_DIM, CHANNELS), + stride=(VALUES_PER_NODE, CHANNELS, 1), + ) + self.kernel(left, right, output, operand_layout).launch( + grid=(left.shape[0], 1, 1), + block=[THREADS, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + left: cute.Tensor, + right: cute.Tensor, + output: cute.Tensor, + operand_layout: cute.Layout, + ): + tidx, _, _ = cute.arch.thread_idx() + node, _, _ = cute.arch.block_idx() + group = tidx >> 6 + channel = tidx & (CHANNELS - 1) + + smem = cute_utils.SmemAllocator() + operands = smem.allocate_tensor(cutlass.Float32, operand_layout, 16) + for side in cutlass.range_constexpr(2): + for slot in cutlass.range_constexpr(VALUES_PER_NODE // THREADS): + linear = tidx + slot * THREADS + row = linear >> 6 + scalar_channel = linear & (CHANNELS - 1) + value = left[node, row, scalar_channel].to(cutlass.Float32) + if cutlass.const_expr(side == 1): + value = right[node, row, scalar_channel].to(cutlass.Float32) + operands[side, row, scalar_channel] = value + cute.arch.sync_threads() + + if group == 0: + self._accumulate_group(operands, output, node, channel, 0) + elif group == 1: + self._accumulate_group(operands, output, node, channel, 1) + elif group == 2: + self._accumulate_group(operands, output, node, channel, 2) + else: + self._accumulate_group(operands, output, node, channel, 3) + + @cute.jit + def _accumulate_group( + self, + operands: cute.Tensor, + output: cute.Tensor, + node: cutlass.Int32, + channel: cutlass.Int32, + group: cutlass.Constexpr, + ): + rows = self.schedule.output_groups[group] + for row_slot in cutlass.range_constexpr(len(rows)): + output_row = rows[row_slot] + accumulator = cutlass.Float32(0.0) + terms = self.schedule.rows[output_row] + for path in cutlass.range_constexpr(len(terms)): + left_row, right_row, coefficient_value, _, _ = terms[path] + product = ( + operands[0, left_row, channel] * operands[1, right_row, channel] + ) + if cutlass.const_expr(left_row != right_row): + product = product + ( + operands[0, right_row, channel] * operands[1, left_row, channel] + ) + accumulator = accumulator + cutlass.Float32(coefficient_value) * product + output[node, output_row, channel] = accumulator + + +def _fake_dense() -> cute.Tensor: + return make_fake_compact_tensor( + cutlass.Float32, + (cute.sym_int64(), COEFF_DIM, CHANNELS), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + + +@device_aware_lru_cache(maxsize=8) +def _compiled_sm90_gaunt_forward(schedule: Sm90GauntSchedule) -> Callable: + dense = _fake_dense() + return cute.compile( + _Sm90GauntForward(schedule), + dense, + dense, + dense, + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +def _validate_dense(*tensors: torch.Tensor) -> None: + reference = tensors[0] + if ( + tuple(reference.shape[1:]) != (COEFF_DIM, CHANNELS) + or reference.shape[0] <= 0 + or not reference.is_cuda + or reference.dtype != torch.float32 + or tuple(torch.cuda.get_device_capability(reference.device)) != SM90_CAPABILITY + ): + raise ValueError("Gaunt operands must be SM90 FP32 (N,48,64)") + if any( + tensor.shape != reference.shape + or tensor.device != reference.device + or tensor.dtype != torch.float32 + or not tensor.is_contiguous() + or tensor.data_ptr() % 16 != 0 + for tensor in tensors + ): + raise ValueError("Gaunt operands must match and be 16-byte aligned") + + +def run_sm90_gaunt_forward( + left: torch.Tensor, + right: torch.Tensor, + schedule: Sm90GauntSchedule, +) -> torch.Tensor: + """Evaluate the exact normalized-Gaunt product on SM90.""" + _validate_dense(left, right) + output = torch.empty_like(left) + with torch.cuda.device(left.device): + _compiled_sm90_gaunt_forward(schedule)(left, right, output) + return output + + +__all__ = [ + "Sm90GauntSchedule", + "build_sm90_gaunt_schedule", + "run_sm90_gaunt_forward", +] diff --git a/deepmd/kernels/cute/neo/message_grid_readout_sm90.py b/deepmd/kernels/cute/neo/message_grid_readout_sm90.py new file mode 100644 index 0000000000..1ed8b0d680 --- /dev/null +++ b/deepmd/kernels/cute/neo/message_grid_readout_sm90.py @@ -0,0 +1,849 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Fused strict-FP32 Neo message-grid input adjoint for SM90.""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, + Any, +) + +import cutlass +import cutlass.cute as cute +import cutlass.utils as cute_utils +import torch +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from .compile_cache import ( + device_aware_lru_cache, +) +from .message_grid_gaunt_sm90 import ( + Sm90GauntSchedule, + build_sm90_gaunt_schedule, +) +from .runtime_policy import ( + SM90_CAPABILITY, + uses_strict_fp32_matmul, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN001, ANN202, ANN204, TC002 + +DEGREE_COUNT = 16 +FRAME_COUNT = 3 +FOCUS_COUNT = 2 +CHANNELS = 32 +FOLDED_CHANNELS = FOCUS_COUNT * CHANNELS +PACKED_COEFF_DIM = DEGREE_COUNT * FRAME_COUNT + +FRAME_CONTRACT_THREADS = 128 +FRAME_CONTRACT_M_TILE = 128 +FRAME_CONTRACT_N_TILE = 64 +FRAME_CONTRACT_K_TILE = 16 +FRAME_CONTRACT_N_TILES = 2 +FRAME_CONTRACT_WIDTH = FRAME_COUNT * CHANNELS +FRAME_CONTRACT_SMEM_PADDING = 4 + +READOUT_THREADS = 256 +READOUT_GROUPS = 4 +VALUES_PER_PANEL = PACKED_COEFF_DIM * FOLDED_CHANNELS +ROWS_PER_GROUP = PACKED_COEFF_DIM // READOUT_GROUPS +DEGREE_CHANNEL_VALUES_PER_NODE = DEGREE_COUNT * CHANNELS +DEGREES_PER_THREAD = DEGREE_CHANNEL_VALUES_PER_NODE // READOUT_THREADS +LEFT_PANEL = 0 +RIGHT_PANEL = 1 +GATED_COEFFICIENT_PANEL = 2 +GRAD_LEFT_PANEL = 3 +GRAD_RIGHT_PANEL = 4 +WORKSPACE_PANELS = 5 + +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + + +@dataclass(frozen=True) +class Sm90MessageGridState: + """Certified Gaunt schedule and packed immutable readout weights.""" + + schedule: Sm90GauntSchedule + frame_contract: torch.Tensor + residual_scale: torch.Tensor + frame_expand_t: torch.Tensor + + +def _tensor_cache_key(tensor: torch.Tensor | None) -> tuple[Any, ...] | None: + if tensor is None: + return None + return ( + tensor.data_ptr(), + tensor._version, + tensor.dtype, + tensor.device, + tuple(tensor.shape), + tuple(tensor.stride()), + ) + + +def prepare_sm90_message_grid_state(net: Any) -> Sm90MessageGridState: + """Prepare and cache the fixed Neo readout contract outside the hot path.""" + projector = net.projector + key = ( + _tensor_cache_key(projector.to_grid_mat), + _tensor_cache_key(projector.from_grid_mat), + _tensor_cache_key(net.frame_contract.weight), + _tensor_cache_key(net.frame_contract.degree_index), + _tensor_cache_key(net.frame_expand.weight), + _tensor_cache_key(net.frame_expand.degree_index), + _tensor_cache_key(net.residual_scale), + ) + cached = getattr(net, "_deepmd_cute_sm90_message_grid_state", None) + if cached is not None and cached[0] == key: + return cached[1] + + schedule = build_sm90_gaunt_schedule( + projector.to_grid_mat, + projector.from_grid_mat, + ) + frame_contract = net.frame_contract.weight.index_select( + 0, + net.frame_contract.degree_index, + ).view(DEGREE_COUNT, FRAME_COUNT, CHANNELS, CHANNELS) + frame_expand = net.frame_expand.weight.index_select( + 0, + net.frame_expand.degree_index, + ).view(DEGREE_COUNT, CHANNELS, FRAME_COUNT, CHANNELS) + residual_scale = net.residual_scale + if residual_scale is None: + residual_scale = torch.ones( + (FOCUS_COUNT, CHANNELS), + device=frame_contract.device, + dtype=torch.float32, + ) + state = Sm90MessageGridState( + schedule=schedule, + frame_contract=frame_contract.contiguous(), + residual_scale=residual_scale.view(FOCUS_COUNT, CHANNELS).contiguous(), + frame_expand_t=frame_expand.permute(0, 2, 3, 1).contiguous(), + ) + net._deepmd_cute_sm90_message_grid_state = (key, state) + return state + + +@cute.jit +def _make_frame_contract_mma(): + atoms_layout = cute.make_layout( + (FRAME_CONTRACT_THREADS // 16, 16, 1), + stride=(16, 1, 0), + ) + permutation_m = cute.make_layout( + (atoms_layout.shape[0], 4), + stride=(4, 1), + ) + permutation_n = cute.make_layout( + (atoms_layout.shape[1], 4), + stride=(4, 1), + ) + return cute.make_tiled_mma( + cute.nvgpu.MmaUniversalOp(cutlass.Float32), + atoms_layout, + permutation_mnk=(permutation_m, permutation_n, None), + ) + + +class _FrameContractAdjoint: + """Batched ``N x 32 @ 32 x 96`` strict-FP32 FrameContract adjoint.""" + + @cute.jit + def __call__( + self, + grad_out: cute.Tensor, + frame_contract: cute.Tensor, + residual_scale: cute.Tensor, + coefficient_slab: cute.Tensor, + stream: CUstream, + ): + s_a_layout = cute.make_layout( + (FRAME_CONTRACT_M_TILE, FRAME_CONTRACT_K_TILE), + stride=(1, FRAME_CONTRACT_M_TILE + FRAME_CONTRACT_SMEM_PADDING), + ) + s_b_layout = cute.make_layout( + (FRAME_CONTRACT_N_TILE, FRAME_CONTRACT_K_TILE), + stride=(1, FRAME_CONTRACT_N_TILE + FRAME_CONTRACT_SMEM_PADDING), + ) + output_reference_layout = cute.make_layout( + (FRAME_CONTRACT_M_TILE, FRAME_CONTRACT_N_TILE), + stride=(FRAME_CONTRACT_N_TILE, 1), + ) + self.kernel( + grad_out, + frame_contract, + residual_scale, + coefficient_slab, + s_a_layout, + s_b_layout, + output_reference_layout, + _make_frame_contract_mma(), + ).launch( + grid=( + cute.ceil_div(grad_out.shape[0], FRAME_CONTRACT_M_TILE), + FRAME_CONTRACT_N_TILES, + DEGREE_COUNT * FOCUS_COUNT, + ), + block=[FRAME_CONTRACT_THREADS, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + grad_out: cute.Tensor, + frame_contract: cute.Tensor, + residual_scale: cute.Tensor, + coefficient_slab: cute.Tensor, + s_a_layout: cute.Layout, + s_b_layout: cute.Layout, + output_reference_layout: cute.Layout, + tiled_mma: cute.TiledMma, + ): + tidx, _, _ = cute.arch.thread_idx() + m_tile, n_tile, batch = cute.arch.block_idx() + m_base = m_tile * FRAME_CONTRACT_M_TILE + n_base = n_tile * FRAME_CONTRACT_N_TILE + degree = batch // FOCUS_COUNT + focus = batch - degree * FOCUS_COUNT + + smem = cute_utils.SmemAllocator() + s_a = smem.allocate_tensor(cutlass.Float32, s_a_layout, 16) + s_b = smem.allocate_tensor(cutlass.Float32, s_b_layout, 16) + thr_mma = tiled_mma.get_slice(tidx) + t_s_a = thr_mma.partition_A(s_a) + t_s_b = thr_mma.partition_B(s_b) + r_a = tiled_mma.make_fragment_A(t_s_a) + r_b = tiled_mma.make_fragment_B(t_s_b) + + # The reference supplies the C-fragment partition. The epilogue maps + # logical columns directly into the packed coefficient slab. + output_reference = cute.make_tensor( + coefficient_slab.iterator, + output_reference_layout, + ) + t_c_reference = thr_mma.partition_C(output_reference) + accumulator = tiled_mma.make_fragment_C(t_c_reference) + accumulator.fill(0.0) + k_blocks = cute.size(r_a, mode=[2]) + + for k_tile in cutlass.range_constexpr(CHANNELS // FRAME_CONTRACT_K_TILE): + self._stage_operands( + grad_out, + frame_contract, + residual_scale, + s_a, + s_b, + tidx, + m_base, + n_base, + degree, + focus, + k_tile * FRAME_CONTRACT_K_TILE, + ) + for k_block in cutlass.range(k_blocks, unroll_full=True): + cute.autovec_copy( + t_s_a[None, None, k_block], + r_a[None, None, k_block], + ) + cute.autovec_copy( + t_s_b[None, None, k_block], + r_b[None, None, k_block], + ) + cute.gemm( + tiled_mma, + accumulator, + r_a[None, None, k_block], + r_b[None, None, k_block], + accumulator, + ) + cute.arch.sync_threads() + + self._store_compact_epilogue( + coefficient_slab, + accumulator, + thr_mma, + output_reference, + m_base, + n_base, + degree, + focus, + ) + + @cute.jit + def _stage_operands( + self, + grad_out: cute.Tensor, + frame_contract: cute.Tensor, + residual_scale: cute.Tensor, + s_a: cute.Tensor, + s_b: cute.Tensor, + tidx: cutlass.Int32, + m_base: cutlass.Int32, + n_base: cutlass.Int32, + degree: cutlass.Int32, + focus: cutlass.Int32, + k_base: cutlass.Constexpr[int], + ): + a_slots = ( + FRAME_CONTRACT_M_TILE * FRAME_CONTRACT_K_TILE + FRAME_CONTRACT_THREADS - 1 + ) // FRAME_CONTRACT_THREADS + for slot in cutlass.range_constexpr(a_slots): + linear = tidx + slot * FRAME_CONTRACT_THREADS + row = linear // FRAME_CONTRACT_K_TILE + k = linear - row * FRAME_CONTRACT_K_TILE + node = m_base + row + value = cutlass.Float32(0.0) + if node < grad_out.shape[0]: + output_channel = k_base + k + value = grad_out[node, degree, focus, output_channel].to( + cutlass.Float32 + ) * residual_scale[focus, output_channel].to(cutlass.Float32) + s_a[row, k] = value + + b_slots = ( + FRAME_CONTRACT_N_TILE * FRAME_CONTRACT_K_TILE + FRAME_CONTRACT_THREADS - 1 + ) // FRAME_CONTRACT_THREADS + for slot in cutlass.range_constexpr(b_slots): + linear = tidx + slot * FRAME_CONTRACT_THREADS + output_column = linear // FRAME_CONTRACT_K_TILE + k = linear - output_column * FRAME_CONTRACT_K_TILE + logical_column = n_base + output_column + value = cutlass.Float32(0.0) + if logical_column < FRAME_CONTRACT_WIDTH: + frame = logical_column // CHANNELS + channel = logical_column - frame * CHANNELS + value = frame_contract[ + degree, + frame, + channel, + k_base + k, + ].to(cutlass.Float32) + s_b[output_column, k] = value + cute.arch.sync_threads() + + @cute.jit + def _store_compact_epilogue( + self, + coefficient_slab: cute.Tensor, + accumulator: cute.Tensor, + thr_mma, + output_reference: cute.Tensor, + m_base: cutlass.Int32, + n_base: cutlass.Int32, + degree: cutlass.Int32, + focus: cutlass.Int32, + ): + accumulator.store(accumulator.load()) + identity = cute.make_identity_tensor(output_reference.shape) + coordinates = thr_mma.partition_C(identity) + for value_idx in range(cute.size(accumulator.shape)): + coordinate = coordinates[value_idx] + node = m_base + coordinate[0] + logical_column = n_base + coordinate[1] + if ( + node < coefficient_slab.shape[0] + and logical_column < FRAME_CONTRACT_WIDTH + ): + frame = logical_column // CHANNELS + channel = logical_column - frame * CHANNELS + packed = degree * FRAME_COUNT + frame + hidden = focus * CHANNELS + channel + coefficient_slab[node, packed, hidden] = accumulator[value_idx].to( + cutlass.Float32 + ) + + +class _FusedReadoutAdjoint: + """Fuse gate, normalized-Gaunt adjoint, and FrameExpand transpose.""" + + def __init__(self, schedule: Sm90GauntSchedule) -> None: + self.schedule = schedule + + @cute.jit + def __call__( + self, + coefficient_slab: cute.Tensor, + scalar_gate: cute.Tensor, + product: cute.Tensor, + left: cute.Tensor, + right: cute.Tensor, + frame_expand_t: cute.Tensor, + grad_query: cute.Tensor, + grad_context: cute.Tensor, + grad_scalar_gate: cute.Tensor, + grad_scalar_out: cute.Tensor, + stream: CUstream, + ): + workspace_layout = cute.make_layout( + (WORKSPACE_PANELS, PACKED_COEFF_DIM, FOLDED_CHANNELS), + stride=(VALUES_PER_PANEL, FOLDED_CHANNELS, 1), + ) + statistic_layout = cute.make_layout( + (READOUT_GROUPS, FOLDED_CHANNELS), + stride=(FOLDED_CHANNELS, 1), + ) + self.kernel( + coefficient_slab, + scalar_gate, + product, + left, + right, + frame_expand_t, + grad_query, + grad_context, + grad_scalar_gate, + grad_scalar_out, + workspace_layout, + statistic_layout, + ).launch( + grid=(left.shape[0], 1, 1), + block=[READOUT_THREADS, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + coefficient_slab: cute.Tensor, + scalar_gate: cute.Tensor, + product: cute.Tensor, + left: cute.Tensor, + right: cute.Tensor, + frame_expand_t: cute.Tensor, + grad_query: cute.Tensor, + grad_context: cute.Tensor, + grad_scalar_gate: cute.Tensor, + grad_scalar_out: cute.Tensor, + workspace_layout: cute.Layout, + statistic_layout: cute.Layout, + ): + tidx, _, _ = cute.arch.thread_idx() + node, _, _ = cute.arch.block_idx() + group = tidx >> 6 + folded_channel = tidx & (FOLDED_CHANNELS - 1) + gate = scalar_gate[node, folded_channel].to(cutlass.Float32) + + smem = cute_utils.SmemAllocator() + workspace = smem.allocate_tensor(cutlass.Float32, workspace_layout, 16) + statistic_partials = smem.allocate_tensor( + cutlass.Float32, + statistic_layout, + 16, + ) + statistic = cutlass.Float32(0.0) + scalar_out = cutlass.Float32(0.0) + + for operand in cutlass.range_constexpr(3): + for slot in cutlass.range_constexpr(ROWS_PER_GROUP): + row = group * ROWS_PER_GROUP + slot + value = left[node, row, folded_channel].to(cutlass.Float32) + if cutlass.const_expr(operand == RIGHT_PANEL): + value = right[node, row, folded_channel].to(cutlass.Float32) + if cutlass.const_expr(operand == GATED_COEFFICIENT_PANEL): + value = coefficient_slab[node, row, folded_channel].to( + cutlass.Float32 + ) + statistic = statistic + value * product[ + node, + row, + folded_channel, + ].to(cutlass.Float32) + if row == 0: + scalar_out = value + value = value * gate + workspace[operand, row, folded_channel] = value + + statistic_partials[group, folded_channel] = statistic + cute.arch.sync_threads() + + if group == 0: + total = statistic_partials[0, folded_channel] + for statistic_group in cutlass.range_constexpr( + 1, + READOUT_GROUPS, + 1, + ): + total = ( + total + + statistic_partials[ + statistic_group, + folded_channel, + ] + ) + grad_scalar_gate[node, folded_channel] = total + grad_scalar_out[node, folded_channel] = scalar_out + + if group == 0: + self._accumulate_group(workspace, folded_channel, 0) + elif group == 1: + self._accumulate_group(workspace, folded_channel, 1) + elif group == 2: + self._accumulate_group(workspace, folded_channel, 2) + else: + self._accumulate_group(workspace, folded_channel, 3) + cute.arch.sync_threads() + + for slot in cutlass.range_constexpr(DEGREES_PER_THREAD): + linear = tidx + slot * READOUT_THREADS + input_channel = linear & (CHANNELS - 1) + degree = linear >> 5 + accumulator_left_0 = cutlass.Float32(0.0) + accumulator_left_1 = cutlass.Float32(0.0) + accumulator_right_0 = cutlass.Float32(0.0) + accumulator_right_1 = cutlass.Float32(0.0) + for frame in cutlass.range_constexpr(FRAME_COUNT): + packed_row = degree * FRAME_COUNT + frame + for output_channel in cutlass.range_constexpr(CHANNELS): + weight = frame_expand_t[ + degree, + frame, + output_channel, + input_channel, + ].to(cutlass.Float32) + accumulator_left_0 = accumulator_left_0 + ( + workspace[GRAD_LEFT_PANEL, packed_row, output_channel] * weight + ) + accumulator_right_0 = accumulator_right_0 + ( + workspace[GRAD_RIGHT_PANEL, packed_row, output_channel] * weight + ) + accumulator_left_1 = accumulator_left_1 + ( + workspace[ + GRAD_LEFT_PANEL, + packed_row, + CHANNELS + output_channel, + ] + * weight + ) + accumulator_right_1 = accumulator_right_1 + ( + workspace[ + GRAD_RIGHT_PANEL, + packed_row, + CHANNELS + output_channel, + ] + * weight + ) + grad_query[node, degree, 0, input_channel] = accumulator_left_0 + grad_query[node, degree, 1, input_channel] = accumulator_left_1 + grad_context[node, degree, 0, input_channel] = accumulator_right_0 + grad_context[node, degree, 1, input_channel] = accumulator_right_1 + + @cute.jit + def _accumulate_group( + self, + workspace: cute.Tensor, + folded_channel: cutlass.Int32, + group: cutlass.Constexpr, + ): + rows = self.schedule.output_groups[group] + for row_slot in cutlass.range_constexpr(len(rows)): + input_row = rows[row_slot] + accumulator_left = cutlass.Float32(0.0) + accumulator_right = cutlass.Float32(0.0) + terms = self.schedule.rows[input_row] + for path in cutlass.range_constexpr(len(terms)): + ( + left_row, + right_row, + _, + left_weight_value, + right_weight_value, + ) = terms[path] + common = ( + cutlass.Float32(left_weight_value) + * workspace[GATED_COEFFICIENT_PANEL, left_row, folded_channel] + ) + accumulator_left = accumulator_left + ( + common * workspace[RIGHT_PANEL, right_row, folded_channel] + ) + accumulator_right = accumulator_right + ( + common * workspace[LEFT_PANEL, right_row, folded_channel] + ) + if cutlass.const_expr(left_row != right_row): + mirrored_common = ( + cutlass.Float32(right_weight_value) + * workspace[ + GATED_COEFFICIENT_PANEL, + right_row, + folded_channel, + ] + ) + accumulator_left = accumulator_left + ( + mirrored_common + * workspace[RIGHT_PANEL, left_row, folded_channel] + ) + accumulator_right = accumulator_right + ( + mirrored_common + * workspace[LEFT_PANEL, left_row, folded_channel] + ) + workspace[GRAD_LEFT_PANEL, input_row, folded_channel] = accumulator_left + workspace[GRAD_RIGHT_PANEL, input_row, folded_channel] = accumulator_right + + +def _fake_state() -> cute.Tensor: + return make_fake_compact_tensor( + cutlass.Float32, + (cute.sym_int64(), DEGREE_COUNT, FOCUS_COUNT, CHANNELS), + stride_order=(3, 2, 1, 0), + **FAKE_TENSOR_KW, + ) + + +def _fake_frame_weight() -> cute.Tensor: + return make_fake_compact_tensor( + cutlass.Float32, + (DEGREE_COUNT, FRAME_COUNT, CHANNELS, CHANNELS), + stride_order=(3, 2, 1, 0), + **FAKE_TENSOR_KW, + ) + + +def _fake_focus_channel() -> cute.Tensor: + return make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, CHANNELS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + + +def _fake_packed() -> cute.Tensor: + return make_fake_compact_tensor( + cutlass.Float32, + (cute.sym_int64(), PACKED_COEFF_DIM, FOLDED_CHANNELS), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + + +def _fake_node_channel() -> cute.Tensor: + return make_fake_compact_tensor( + cutlass.Float32, + (cute.sym_int64(), FOLDED_CHANNELS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + + +@device_aware_lru_cache(maxsize=4) +def _compiled_frame_contract_adjoint() -> Callable: + return cute.compile( + _FrameContractAdjoint(), + _fake_state(), + _fake_frame_weight(), + _fake_focus_channel(), + _fake_packed(), + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +@device_aware_lru_cache(maxsize=8) +def _compiled_fused_readout_adjoint(schedule: Sm90GauntSchedule) -> Callable: + packed = _fake_packed() + node_channel = _fake_node_channel() + state = _fake_state() + return cute.compile( + _FusedReadoutAdjoint(schedule), + packed, + node_channel, + packed, + packed, + packed, + _fake_frame_weight(), + state, + state, + node_channel, + node_channel, + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +def _validate_tensor( + name: str, + tensor: torch.Tensor, + shape: tuple[int, ...], + device: torch.device, +) -> None: + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}") + if tensor.dtype != torch.float32: + raise TypeError(f"{name} must be FP32") + if tensor.device != device or not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous on {device}") + if tensor.data_ptr() % 16: + raise ValueError(f"{name} must be at least 16-byte aligned") + + +def _validate_runtime( + grad_out: torch.Tensor, + scalar_gate: torch.Tensor, + product: torch.Tensor, + left: torch.Tensor, + right: torch.Tensor, + state: Sm90MessageGridState, +) -> int: + device = left.device + if ( + device.type != "cuda" + or tuple(torch.cuda.get_device_capability(device)) != SM90_CAPABILITY + ): + raise RuntimeError("the fused message-grid readout requires SM90") + if not uses_strict_fp32_matmul(): + raise RuntimeError("the fused message-grid readout requires strict FP32") + nodes = int(left.shape[0]) + if nodes <= 0: + raise ValueError("the fused message-grid readout requires N > 0") + packed_shape = (nodes, PACKED_COEFF_DIM, FOLDED_CHANNELS) + state_shape = (nodes, DEGREE_COUNT, FOCUS_COUNT, CHANNELS) + node_channel_shape = (nodes, FOLDED_CHANNELS) + for name, tensor, shape in ( + ("grad_out", grad_out, state_shape), + ("scalar_gate", scalar_gate, node_channel_shape), + ("product", product, packed_shape), + ("left", left, packed_shape), + ("right", right, packed_shape), + ( + "frame_contract", + state.frame_contract, + (DEGREE_COUNT, FRAME_COUNT, CHANNELS, CHANNELS), + ), + ("residual_scale", state.residual_scale, (FOCUS_COUNT, CHANNELS)), + ( + "frame_expand_t", + state.frame_expand_t, + (DEGREE_COUNT, FRAME_COUNT, CHANNELS, CHANNELS), + ), + ): + _validate_tensor(name, tensor, shape, device) + return nodes + + +def run_sm90_message_grid_backward( + net: Any, + query_flat: torch.Tensor, + context_flat: torch.Tensor, + grad_out_flat: torch.Tensor, + product_flat: torch.Tensor, + state: Sm90MessageGridState, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return query/context adjoints without expanded global adjoint slabs.""" + from .k1_message_grid_packed import ( + _as_flat_ndfc, + _focus_linear_backward_input, + _frame_expand_packed, + _swiglu_backward_input, + _validate_contract, + ) + + query, context = _validate_contract(net, query_flat, context_flat) + nodes = int(query_flat.shape[0]) + scalar_pair = torch.cat([query[:, 0], context[:, 0]], dim=-1).to(net.dtype) + scalar_gate = ( + torch.sigmoid(net.scalar_gate(scalar_pair)) + .reshape( + nodes, + FOLDED_CHANNELS, + ) + .contiguous() + ) + left = _frame_expand_packed(net.frame_expand, query).view( + nodes, + PACKED_COEFF_DIM, + FOLDED_CHANNELS, + ) + right = _frame_expand_packed(net.frame_expand, context).view( + nodes, + PACKED_COEFF_DIM, + FOLDED_CHANNELS, + ) + grad_out = ( + _as_flat_ndfc( + net, + "grad_out", + grad_out_flat, + like=query_flat, + ) + .to(net.dtype) + .contiguous() + ) + _validate_runtime(grad_out, scalar_gate, product_flat, left, right, state) + + coefficient_slab = torch.empty_like(left) + grad_query = torch.empty( + query.shape, + device=query.device, + dtype=query.dtype, + ) + grad_context = torch.empty( + context.shape, + device=context.device, + dtype=context.dtype, + ) + grad_scalar_gate = torch.empty_like(scalar_gate) + grad_scalar_out = torch.empty_like(scalar_gate) + with torch.cuda.device(left.device): + _compiled_frame_contract_adjoint()( + grad_out, + state.frame_contract, + state.residual_scale, + coefficient_slab, + ) + _compiled_fused_readout_adjoint(state.schedule)( + coefficient_slab, + scalar_gate, + product_flat, + left, + right, + state.frame_expand_t, + grad_query, + grad_context, + grad_scalar_gate, + grad_scalar_out, + ) + + gate = scalar_gate.view(nodes, FOCUS_COUNT, CHANNELS) + grad_scalar_logits = grad_scalar_gate.view_as(gate) * gate * (1.0 - gate) + grad_scalar_pair = _focus_linear_backward_input( + net.scalar_gate, + grad_scalar_logits, + ) + _swiglu_backward_input( + scalar_pair, + grad_scalar_out.view(nodes, FOCUS_COUNT, CHANNELS), + ) + grad_query[:, 0].add_(grad_scalar_pair[:, :, :CHANNELS]) + grad_context[:, 0].add_(grad_scalar_pair[:, :, CHANNELS:]) + return ( + grad_query.reshape_as(query_flat).to(dtype=query_flat.dtype), + grad_context.reshape_as(context_flat).to(dtype=context_flat.dtype), + ) + + +__all__ = [ + "Sm90MessageGridState", + "prepare_sm90_message_grid_state", + "run_sm90_message_grid_backward", +] diff --git a/deepmd/kernels/cute/neo/output_grid_kernels/__init__.py b/deepmd/kernels/cute/neo/output_grid_kernels/__init__.py new file mode 100644 index 0000000000..3412625810 --- /dev/null +++ b/deepmd/kernels/cute/neo/output_grid_kernels/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""CuTe kernels for guarded non-K1 DeePMD operations.""" diff --git a/deepmd/kernels/cute/neo/output_grid_kernels/cute_readout_l0.py b/deepmd/kernels/cute/neo/output_grid_kernels/cute_readout_l0.py new file mode 100644 index 0000000000..793367a4e1 --- /dev/null +++ b/deepmd/kernels/cute/neo/output_grid_kernels/cute_readout_l0.py @@ -0,0 +1,834 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tiled strict-FP32 Neo output readout for degree zero only.""" + +from __future__ import ( + annotations, +) + +from collections.abc import ( + Callable, +) +from functools import ( + lru_cache, +) + +import cutlass +import cutlass.cute as cute +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from .cute_tiled_grid_product import ( + FAKE_TENSOR_KW, + GRID_SIZE, + PACKED_COEFF_DIM, + STAGES, + THREADS, + TILE_K, + TILE_M, + TILE_N, + TiledOutputGridProductBackward, +) + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN001, ANN201, ANN202, ANN204, TC002, TC003 + + +HIDDEN_CHANNELS = 192 +GRAM_FORWARD_THREADS = HIDDEN_CHANNELS +GRAM_ELEMENTS = PACKED_COEFF_DIM * PACKED_COEFF_DIM + + +class TiledReadoutL0GramForward: + """Evaluate the channelwise Gram bilinear with one CTA per node.""" + + @cute.jit + def __call__( + self, + left: cute.Tensor, + right: cute.Tensor, + gram: cute.Tensor, + out: cute.Tensor, + stream: CUstream, + ): + gram_layout = cute.make_layout( + (PACKED_COEFF_DIM, PACKED_COEFF_DIM), + stride=(PACKED_COEFF_DIM, 1), + ) + right_layout = cute.make_layout( + (PACKED_COEFF_DIM, HIDDEN_CHANNELS), + stride=(HIDDEN_CHANNELS, 1), + ) + self.kernel( + left, + right, + gram, + out, + gram_layout, + right_layout, + ).launch( + grid=(left.shape[0], 1, 1), + block=[GRAM_FORWARD_THREADS, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + left: cute.Tensor, + right: cute.Tensor, + gram: cute.Tensor, + out: cute.Tensor, + gram_layout: cute.Layout, + right_layout: cute.Layout, + ): + channel, _, _ = cute.arch.thread_idx() + node, _, _ = cute.arch.block_idx() + + smem = cutlass.utils.SmemAllocator() + s_gram = smem.allocate_tensor(cutlass.Float32, gram_layout, 16) + s_right = smem.allocate_tensor(cutlass.Float32, right_layout, 16) + + for linear in cutlass.range( + channel, + GRAM_ELEMENTS, + GRAM_FORWARD_THREADS, + unroll=1, + ): + row = linear // PACKED_COEFF_DIM + col = linear - row * PACKED_COEFF_DIM + s_gram[row, col] = gram[row, col].to(cutlass.Float32) + for coeff in cutlass.range(0, PACKED_COEFF_DIM, 1, unroll=1): + s_right[coeff, channel] = right[node, coeff, channel].to(cutlass.Float32) + cute.arch.sync_threads() + + value = cutlass.Float32(0.0) + for row in cutlass.range(0, PACKED_COEFF_DIM, 1, unroll=1): + transformed_right = cutlass.Float32(0.0) + for col in cutlass.range(0, PACKED_COEFF_DIM, 1, unroll=1): + transformed_right += s_gram[row, col] * s_right[col, channel] + value += left[node, row, channel].to(cutlass.Float32) * transformed_right + out[node, channel] = value.to(out.element_type) + + +class TiledReadoutL0GramBackward(TiledOutputGridProductBackward): + """Apply the frozen 48x48 Gram matrix to one 64-channel tile.""" + + def __init__(self) -> None: + super().__init__(HIDDEN_CHANNELS) + + @cute.jit + def __call__( + self, + dq0: cute.Tensor, + left: cute.Tensor, + right: cute.Tensor, + gram: cute.Tensor, + grad_left: cute.Tensor, + grad_right: cute.Tensor, + stream: CUstream, + ): + sA_layout = cute.make_layout( + (TILE_M, TILE_K, STAGES), + stride=(1, TILE_M + 4, TILE_K * (TILE_M + 4)), + ) + sB_layout = cute.make_layout( + (TILE_N, TILE_K, STAGES), + stride=(1, TILE_N, TILE_K * TILE_N), + ) + dq0_layout = cute.make_layout((TILE_N,), stride=(1,)) + + copy_a_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + left.element_type, + num_bits_per_copy=left.element_type.width, + ) + copy_a_layout = cute.make_layout( + (THREADS // TILE_K, TILE_K), + stride=(TILE_K, 1), + ) + tiled_copy_A = cute.make_tiled_copy_tv( + copy_a_atom, + copy_a_layout, + cute.make_layout((1, 1)), + ) + + vector = 4 + copy_b_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + left.element_type, + num_bits_per_copy=left.element_type.width * vector, + ) + copy_b_major = TILE_N // vector + copy_b_layout = cute.make_layout( + (copy_b_major, THREADS // copy_b_major), + stride=(1, copy_b_major), + ) + tiled_copy_B = cute.make_tiled_copy_tv( + copy_b_atom, + copy_b_layout, + cute.make_layout((vector, 1)), + ) + + atoms_layout = cute.make_layout( + (THREADS // 16, 16, 1), + stride=(16, 1, 0), + ) + permutation_m = cute.make_layout( + (atoms_layout.shape[0], 4), + stride=(4, 1), + ) + permutation_n = cute.make_layout( + (atoms_layout.shape[1], 4), + stride=(4, 1), + ) + tiled_mma = cute.make_tiled_mma( + cute.nvgpu.MmaUniversalOp(cutlass.Float32), + atoms_layout, + permutation_mnk=(permutation_m, permutation_n, None), + ) + + self.kernel( + dq0, + left, + right, + gram, + grad_left, + grad_right, + sA_layout, + sB_layout, + dq0_layout, + tiled_copy_A, + tiled_copy_B, + tiled_mma, + ).launch( + grid=(left.shape[0], self.channel_tiles, 1), + block=[THREADS, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + dq0: cute.Tensor, + left: cute.Tensor, + right: cute.Tensor, + gram: cute.Tensor, + grad_left: cute.Tensor, + grad_right: cute.Tensor, + sA_layout: cute.Layout, + sB_layout: cute.Layout, + dq0_layout: cute.Layout, + tiled_copy_A: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_mma: cute.TiledMma, + ): + tidx, _, _ = cute.arch.thread_idx() + node, channel_tile, _ = cute.arch.block_idx() + + matrix_b_layout = cute.make_layout( + (HIDDEN_CHANNELS, PACKED_COEFF_DIM), + stride=(1, HIDDEN_CHANNELS), + ) + left_b = cute.make_tensor( + left[node, None, None].iterator, + matrix_b_layout, + ) + right_b = cute.make_tensor( + right[node, None, None].iterator, + matrix_b_layout, + ) + gram_t = cute.make_tensor( + gram.iterator, + cute.make_layout( + (PACKED_COEFF_DIM, PACKED_COEFF_DIM), + stride=(1, PACKED_COEFF_DIM), + ), + ) + + smem = cutlass.utils.SmemAllocator() + sA_left = smem.allocate_tensor(cutlass.Float32, sA_layout, 16) + sA_right = smem.allocate_tensor(cutlass.Float32, sA_layout, 16) + sB_left = smem.allocate_tensor(cutlass.Float32, sB_layout, 16) + sB_right = smem.allocate_tensor(cutlass.Float32, sB_layout, 16) + sDq0 = smem.allocate_tensor(cutlass.Float32, dq0_layout, 16) + + for local_channel in cutlass.range( + tidx, + TILE_N, + THREADS, + unroll=1, + ): + channel = channel_tile * TILE_N + local_channel + sDq0[local_channel] = dq0[node, channel].to(cutlass.Float32) + cute.arch.sync_threads() + + self._dual_gram_adjoint( + gram, + gram_t, + right_b, + left_b, + grad_left[node, None, None], + grad_right[node, None, None], + sA_left, + sA_right, + sB_left, + sB_right, + sDq0, + tiled_copy_A, + tiled_copy_B, + tiled_mma, + tidx, + channel_tile, + ) + + @cute.jit + def _dual_gram_adjoint( + self, + mA_left: cute.Tensor, + mA_right: cute.Tensor, + mB_left: cute.Tensor, + mB_right: cute.Tensor, + mOut_left: cute.Tensor, + mOut_right: cute.Tensor, + sA_left: cute.Tensor, + sA_right: cute.Tensor, + sB_left: cute.Tensor, + sB_right: cute.Tensor, + sDq0: cute.Tensor, + tiled_copy_A: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + channel_tile: cutlass.Int32, + ): + thr_mma = tiled_mma.get_slice(tidx) + gA_left = cute.local_tile( + mA_left, + tiler=self.cta_tiler, + coord=(0, 0, None), + proj=(1, None, 1), + ) + gA_right = cute.local_tile( + mA_right, + tiler=self.cta_tiler, + coord=(0, 0, None), + proj=(1, None, 1), + ) + gB_left = cute.local_tile( + mB_left, + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(None, 1, 1), + ) + gB_right = cute.local_tile( + mB_right, + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(None, 1, 1), + ) + gOut_left = cute.local_tile( + mOut_left, + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(1, 1, None), + ) + gOut_right = cute.local_tile( + mOut_right, + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(1, 1, None), + ) + + thr_copy_A = tiled_copy_A.get_slice(tidx) + thr_copy_B = tiled_copy_B.get_slice(tidx) + tAgA_left = thr_copy_A.partition_S(gA_left) + tAgA_right = thr_copy_A.partition_S(gA_right) + tAsA_left = thr_copy_A.partition_D(sA_left) + tAsA_right = thr_copy_A.partition_D(sA_right) + tBgB_left = thr_copy_B.partition_S(gB_left) + tBgB_right = thr_copy_B.partition_S(gB_right) + tBsB_left = thr_copy_B.partition_D(sB_left) + tBsB_right = thr_copy_B.partition_D(sB_right) + + cA = cute.local_tile( + cute.make_identity_tensor(mA_left.shape), + tiler=self.cta_tiler, + coord=(0, 0, None), + proj=(1, None, 1), + ) + tAcA = thr_copy_A.partition_S(cA) + tApA = cute.make_rmem_tensor( + cute.make_layout( + ( + tAsA_left.shape[0][1], + cute.size(tAsA_left, mode=[1]), + cute.size(tAsA_left, mode=[2]), + ), + stride=(cute.size(tAsA_left, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in range(tApA.shape[0]): + for row in range(tApA.shape[1]): + tApA[rest_v, row, 0] = cute.elem_less( + tAcA[(0, rest_v), row, 0, 0][0], + PACKED_COEFF_DIM, + ) + + k_tile_count = cute.size(tAgA_left, mode=[3]) + gmem_pipe_read = cutlass.Int32(0) + cute.copy( + tiled_copy_A, + tAgA_left[None, None, None, gmem_pipe_read], + tAsA_left[None, None, None, 0], + pred=tApA, + ) + cute.copy( + tiled_copy_A, + tAgA_right[None, None, None, gmem_pipe_read], + tAsA_right[None, None, None, 0], + pred=tApA, + ) + cute.copy( + tiled_copy_B, + tBgB_left[None, None, None, gmem_pipe_read], + tBsB_left[None, None, None, 0], + ) + cute.copy( + tiled_copy_B, + tBgB_right[None, None, None, gmem_pipe_read], + tBsB_right[None, None, None, 0], + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + for stage in range(1, STAGES - 1): + cute.copy( + tiled_copy_A, + tAgA_left[None, None, None, gmem_pipe_read], + tAsA_left[None, None, None, stage], + pred=tApA, + ) + cute.copy( + tiled_copy_A, + tAgA_right[None, None, None, gmem_pipe_read], + tAsA_right[None, None, None, stage], + pred=tApA, + ) + cute.copy( + tiled_copy_B, + tBgB_left[None, None, None, gmem_pipe_read], + tBsB_left[None, None, None, stage], + ) + cute.copy( + tiled_copy_B, + tBgB_right[None, None, None, gmem_pipe_read], + tBsB_right[None, None, None, stage], + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + + tCsA_left = thr_mma.partition_A(sA_left) + tCsA_right = thr_mma.partition_A(sA_right) + tCsB_left = thr_mma.partition_B(sB_left) + tCsB_right = thr_mma.partition_B(sB_right) + tCgOut_left = thr_mma.partition_C(gOut_left) + tCgOut_right = thr_mma.partition_C(gOut_right) + tCrA_left = tiled_mma.make_fragment_A(tCsA_left[None, None, None, 0]) + tCrA_right = tiled_mma.make_fragment_A(tCsA_right[None, None, None, 0]) + tCrB_left = tiled_mma.make_fragment_B(tCsB_left[None, None, None, 0]) + tCrB_right = tiled_mma.make_fragment_B(tCsB_right[None, None, None, 0]) + tCrOut_left = tiled_mma.make_fragment_C(tCgOut_left) + tCrOut_right = tiled_mma.make_fragment_C(tCgOut_right) + tCrOut_left.fill(0.0) + tCrOut_right.fill(0.0) + + smem_pipe_read = cutlass.Int32(0) + smem_pipe_write = cutlass.Int32(STAGES - 1) + tiles_issued = cutlass.Int32(STAGES - 1) + tCsA_left_p = tCsA_left[None, None, None, smem_pipe_read] + tCsA_right_p = tCsA_right[None, None, None, smem_pipe_read] + tCsB_left_p = tCsB_left[None, None, None, smem_pipe_read] + tCsB_right_p = tCsB_right[None, None, None, smem_pipe_read] + k_block_max = cute.size(tCrA_left, mode=[2]) + if k_block_max > 1: + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + cute.autovec_copy( + tCsA_left_p[None, None, 0], + tCrA_left[None, None, 0], + ) + cute.autovec_copy( + tCsA_right_p[None, None, 0], + tCrA_right[None, None, 0], + ) + cute.autovec_copy( + tCsB_left_p[None, None, 0], + tCrB_left[None, None, 0], + ) + cute.autovec_copy( + tCsB_right_p[None, None, 0], + tCrB_right[None, None, 0], + ) + + for _ in range(k_tile_count): + for k_block in range(k_block_max, unroll_full=True): + if k_block == k_block_max - 1: + tCsA_left_p = tCsA_left[None, None, None, smem_pipe_read] + tCsA_right_p = tCsA_right[None, None, None, smem_pipe_read] + tCsB_left_p = tCsB_left[None, None, None, smem_pipe_read] + tCsB_right_p = tCsB_right[None, None, None, smem_pipe_read] + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + k_block_next = (k_block + 1) % k_block_max + cute.autovec_copy( + tCsA_left_p[None, None, k_block_next], + tCrA_left[None, None, k_block_next], + ) + cute.autovec_copy( + tCsA_right_p[None, None, k_block_next], + tCrA_right[None, None, k_block_next], + ) + cute.autovec_copy( + tCsB_left_p[None, None, k_block_next], + tCrB_left[None, None, k_block_next], + ) + cute.autovec_copy( + tCsB_right_p[None, None, k_block_next], + tCrB_right[None, None, k_block_next], + ) + if k_block == 0 and tiles_issued < k_tile_count: + cute.copy( + tiled_copy_A, + tAgA_left[None, None, None, gmem_pipe_read], + tAsA_left[None, None, None, smem_pipe_write], + pred=tApA, + ) + cute.copy( + tiled_copy_A, + tAgA_right[None, None, None, gmem_pipe_read], + tAsA_right[None, None, None, smem_pipe_write], + pred=tApA, + ) + cute.copy( + tiled_copy_B, + tBgB_left[None, None, None, gmem_pipe_read], + tBsB_left[None, None, None, smem_pipe_write], + ) + cute.copy( + tiled_copy_B, + tBgB_right[None, None, None, gmem_pipe_read], + tBsB_right[None, None, None, smem_pipe_write], + ) + cute.gemm( + tiled_mma, + tCrOut_left, + tCrA_left[None, None, k_block], + tCrB_left[None, None, k_block], + tCrOut_left, + ) + cute.gemm( + tiled_mma, + tCrOut_right, + tCrA_right[None, None, k_block], + tCrB_right[None, None, k_block], + tCrOut_right, + ) + if k_block == 0: + cute.arch.cp_async_commit_group() + tiles_issued = tiles_issued + 1 + smem_pipe_write = smem_pipe_read + smem_pipe_read = smem_pipe_read + 1 + if smem_pipe_read == STAGES: + smem_pipe_read = cutlass.Int32(0) + gmem_pipe_read = ( + gmem_pipe_read + 1 + if gmem_pipe_read + 1 < k_tile_count + else cutlass.Int32(0) + ) + + cute.arch.cp_async_wait_group(0) + self.cta_sync_barrier.arrive_and_wait() + cOut = cute.make_identity_tensor(gOut_left.shape) + tCpOut = thr_mma.partition_C(cOut) + pred = cute.make_rmem_tensor(tCrOut_left.layout, cutlass.Boolean) + for idx in range(cute.size(tCrOut_left.shape)): + pred[idx] = cute.elem_less( + tCpOut[idx], + (PACKED_COEFF_DIM, TILE_N), + ) + if pred[idx]: + local_channel = tCpOut[idx][1] + scale = sDq0[local_channel].to(cutlass.Float32) + tCrOut_left[idx] = tCrOut_left[idx].to(cutlass.Float32) * scale + tCrOut_right[idx] = tCrOut_right[idx].to(cutlass.Float32) * scale + atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mOut_left.element_type, + ) + cute.copy(atom, tCrOut_left, tCgOut_left, pred=pred) + cute.copy(atom, tCrOut_right, tCgOut_right, pred=pred) + + +def compile_readout_l0_gram_forward( + device_index: int | None = None, + compute_capability: tuple[int, int] | None = None, +) -> Callable: + """Compile the dense-Gram forward with symbolic runtime node count.""" + import torch + + if device_index is None: + device_index = torch.cuda.current_device() + actual_capability = tuple(torch.cuda.get_device_capability(device_index)) + if ( + compute_capability is not None + and tuple(compute_capability) != actual_capability + ): + raise ValueError("compile target does not match the selected CUDA device") + with torch.cuda.device(device_index): + nodes = cute.sym_int64() + fake_coeff = make_fake_compact_tensor( + cutlass.Float32, + (nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + fake_gram = make_fake_compact_tensor( + cutlass.Float32, + (PACKED_COEFF_DIM, PACKED_COEFF_DIM), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_out = make_fake_compact_tensor( + cutlass.Float32, + (nodes, HIDDEN_CHANNELS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + TiledReadoutL0GramForward(), + fake_coeff, + fake_coeff, + fake_gram, + fake_out, + fake_stream, + options="--enable-tvm-ffi", + ) + + +def compile_readout_l0_gram_backward( + device_index: int | None = None, + compute_capability: tuple[int, int] | None = None, +) -> Callable: + """Compile the dense-Gram first-backward artifact.""" + import torch + + if device_index is None: + device_index = torch.cuda.current_device() + actual_capability = tuple(torch.cuda.get_device_capability(device_index)) + if ( + compute_capability is not None + and tuple(compute_capability) != actual_capability + ): + raise ValueError("compile target does not match the selected CUDA device") + with torch.cuda.device(device_index): + nodes = cute.sym_int64() + fake_coeff = make_fake_compact_tensor( + cutlass.Float32, + (nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + fake_q0 = make_fake_compact_tensor( + cutlass.Float32, + (nodes, HIDDEN_CHANNELS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_gram = make_fake_compact_tensor( + cutlass.Float32, + (PACKED_COEFF_DIM, PACKED_COEFF_DIM), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + TiledReadoutL0GramBackward(), + fake_q0, + fake_coeff, + fake_coeff, + fake_gram, + fake_coeff, + fake_coeff, + fake_stream, + options="--enable-tvm-ffi", + ) + + +@lru_cache(maxsize=16) +def _compiled_readout_l0_gram_forward( + device_index: int, + compute_capability: tuple[int, int], +) -> Callable: + return compile_readout_l0_gram_forward(device_index, compute_capability) + + +@lru_cache(maxsize=16) +def _compiled_readout_l0_gram_backward( + device_index: int, + compute_capability: tuple[int, int], +) -> Callable: + return compile_readout_l0_gram_backward(device_index, compute_capability) + + +def _compile_key(tensor) -> tuple[int, tuple[int, int]]: + import torch + + device_index = tensor.device.index + if device_index is None: + device_index = torch.cuda.current_device() + compute_capability = tuple(torch.cuda.get_device_capability(device_index)) + return int(device_index), compute_capability + + +def _validate_tensors(left, right, to_grid, from_grid, out=None) -> None: + import torch + + tensors = (left, right, to_grid, from_grid) + if ( + any(not tensor.is_cuda for tensor in tensors) + or any(tensor.dtype != torch.float32 for tensor in tensors) + or any(tensor.device != left.device for tensor in tensors) + or left.ndim != 3 + or left.shape[0] <= 0 + or tuple(left.shape[1:]) != (PACKED_COEFF_DIM, HIDDEN_CHANNELS) + or right.shape != left.shape + or tuple(to_grid.shape) != (GRID_SIZE, PACKED_COEFF_DIM) + or tuple(from_grid.shape) != (PACKED_COEFF_DIM, GRID_SIZE) + or any(not tensor.is_contiguous() for tensor in tensors) + or any(tensor.data_ptr() % 16 != 0 for tensor in tensors) + or torch.cuda.get_device_capability(left.device)[0] < 8 + ): + raise ValueError( + "readout l=0 requires contiguous CUDA FP32 left/right=(N,48,192), " + "to_grid=(152,48), and from_grid=(48,152) tensors on compute " + "capability 8.0+" + ) + if out is not None and ( + tuple(out.shape) != (left.shape[0], HIDDEN_CHANNELS) + or out.device != left.device + or out.dtype != left.dtype + or not out.is_contiguous() + or out.data_ptr() % 16 != 0 + ): + raise ValueError("readout l=0 output must be contiguous with shape (N,192)") + + +def _validate_gram_tensors(left, right, gram, out=None) -> None: + import torch + + tensors = (left, right, gram) + if ( + any(not tensor.is_cuda for tensor in tensors) + or any(tensor.dtype != torch.float32 for tensor in tensors) + or any(tensor.device != left.device for tensor in tensors) + or left.ndim != 3 + or left.shape[0] <= 0 + or tuple(left.shape[1:]) != (PACKED_COEFF_DIM, HIDDEN_CHANNELS) + or right.shape != left.shape + or tuple(gram.shape) != (PACKED_COEFF_DIM, PACKED_COEFF_DIM) + or any(not tensor.is_contiguous() for tensor in tensors) + or any(tensor.data_ptr() % 16 != 0 for tensor in tensors) + or torch.cuda.get_device_capability(left.device)[0] < 8 + ): + raise ValueError( + "Gram readout l=0 requires contiguous CUDA FP32 left/right=" + "(N,48,192) and gram=(48,48) tensors on compute capability 8.0+" + ) + if out is not None and ( + tuple(out.shape) != (left.shape[0], HIDDEN_CHANNELS) + or out.device != left.device + or out.dtype != left.dtype + or not out.is_contiguous() + or out.data_ptr() % 16 != 0 + ): + raise ValueError("readout l=0 output must be contiguous with shape (N,192)") + + +def run_readout_l0_gram(left, right, gram): + """Run ``left[:, :, h]^T G right[:, :, h]`` in strict FP32.""" + import torch + + _validate_gram_tensors(left, right, gram) + out = torch.empty( + (left.shape[0], HIDDEN_CHANNELS), + dtype=left.dtype, + device=left.device, + ) + _compiled_readout_l0_gram_forward(*_compile_key(left))( + left, + right, + gram, + out, + ) + return out + + +def run_readout_l0(left, right, to_grid, from_grid): + """Run the strict-FP32 row-zero readout forward.""" + _validate_tensors(left, right, to_grid, from_grid) + from ..readout_l0 import ( + build_readout_l0_gram, + ) + + return run_readout_l0_gram( + left, + right, + build_readout_l0_gram(to_grid, from_grid), + ) + + +def run_readout_l0_backward(dq0, left, right, to_grid, from_grid): + """Run first backward and return full `(N,48,192)` input adjoints.""" + import torch + + _validate_tensors(left, right, to_grid, from_grid) + if ( + tuple(dq0.shape) != (left.shape[0], HIDDEN_CHANNELS) + or dq0.device != left.device + or dq0.dtype != torch.float32 + or not dq0.is_contiguous() + ): + raise ValueError("dq0 must be contiguous CUDA FP32 with shape (N,192)") + grad_left = torch.empty_like(left) + grad_right = torch.empty_like(right) + compile_key = _compile_key(left) + from ..readout_l0 import ( + build_readout_l0_gram, + ) + + gram = build_readout_l0_gram(to_grid, from_grid) + _compiled_readout_l0_gram_backward(*compile_key)( + dq0, + left, + right, + gram, + grad_left, + grad_right, + ) + return grad_left, grad_right + + +__all__ = [ + "TiledReadoutL0GramBackward", + "TiledReadoutL0GramForward", + "run_readout_l0", + "run_readout_l0_backward", + "run_readout_l0_gram", +] diff --git a/deepmd/kernels/cute/neo/output_grid_kernels/cute_tiled_grid_product.py b/deepmd/kernels/cute/neo/output_grid_kernels/cute_tiled_grid_product.py new file mode 100644 index 0000000000..582fc8ef22 --- /dev/null +++ b/deepmd/kernels/cute/neo/output_grid_kernels/cute_tiled_grid_product.py @@ -0,0 +1,2915 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tiled strict-FP32 Neo output-grid product forward and first backward.""" + +from __future__ import ( + annotations, +) + +from collections.abc import ( + Callable, +) +from functools import ( + lru_cache, +) + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from .. import ( + runtime_policy, +) + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN001, ANN201, ANN202, ANN204, TC002, TC003 + + +PACKED_COEFF_DIM = 48 +GRID_SIZE = 152 +SUPPORTED_HIDDEN_CHANNELS = (96, 192) +TILE_M = 64 +TILE_N = 64 +TILE_K = 8 +C96_TAIL_TILE_N = 32 +C96_TAIL_CHANNEL_TILE = 2 +SM80_C96_TILE_N = 48 +SM80_C96_THREADS = 128 +MMA_ATOMS_N = 16 +THREADS = 128 +STAGES = 3 +GRID_TILES = 3 +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + + +def _validate_hidden_channels(hidden_channels: int) -> int: + hidden_channels = int(hidden_channels) + if hidden_channels not in SUPPORTED_HIDDEN_CHANNELS: + raise ValueError( + "tiled output-grid channel width must be one of " + f"{SUPPORTED_HIDDEN_CHANNELS}, got {hidden_channels}" + ) + return hidden_channels + + +class TiledOutputGridProductForward: + """Dual tiled projections, shared product, and tiled backprojection.""" + + def __init__( + self, + hidden_channels: int = 192, + *, + tile_n: int = TILE_N, + channel_tile_start: int = 0, + channel_tile_count: int | None = None, + ) -> None: + self.hidden_channels = _validate_hidden_channels(hidden_channels) + tile_n = int(tile_n) + if tile_n not in (C96_TAIL_TILE_N, SM80_C96_TILE_N, TILE_N): + raise ValueError("output-grid forward tile_n must be 32, 48, or 64") + if tile_n == SM80_C96_TILE_N and self.hidden_channels != 96: + raise ValueError("output-grid forward N=48 specializes C=96") + channel_tile_start = int(channel_tile_start) + total_channel_tiles = (self.hidden_channels + tile_n - 1) // tile_n + if channel_tile_count is None: + channel_tile_count = total_channel_tiles - channel_tile_start + channel_tile_count = int(channel_tile_count) + if ( + channel_tile_start < 0 + or channel_tile_count <= 0 + or channel_tile_start + channel_tile_count > total_channel_tiles + ): + raise ValueError("invalid output-grid forward channel-tile range") + if tile_n == C96_TAIL_TILE_N and ( + self.hidden_channels != 96 + or channel_tile_start != C96_TAIL_CHANNEL_TILE + or channel_tile_count != 1 + ): + raise ValueError("output-grid forward N=32 specializes the C96 tail panel") + self.cta_tiler = (TILE_M, tile_n, TILE_K) + self.channel_tile_start = channel_tile_start + self.channel_tiles = channel_tile_count + self.has_channel_residue = self.hidden_channels % tile_n != 0 + self.cta_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=THREADS, + ) + + @cute.jit + def __call__( + self, + left: cute.Tensor, + right: cute.Tensor, + to_grid: cute.Tensor, + from_grid: cute.Tensor, + out: cute.Tensor, + stream: CUstream, + ): + tile_n = self.cta_tiler[1] + sA_layout = cute.make_layout( + (TILE_M, TILE_K, STAGES), + stride=(1, TILE_M + 4, TILE_K * (TILE_M + 4)), + ) + sB_layout = cute.make_layout( + (tile_n, TILE_K, STAGES), + stride=(1, tile_n, TILE_K * tile_n), + ) + product_layout = cute.make_layout( + (GRID_SIZE, tile_n), + stride=(tile_n, 1), + ) + copy_a_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + left.element_type, + num_bits_per_copy=left.element_type.width, + ) + copy_a_layout = cute.make_layout( + (THREADS // TILE_K, TILE_K), + stride=(TILE_K, 1), + ) + tiled_copy_A = cute.make_tiled_copy_tv( + copy_a_atom, + copy_a_layout, + cute.make_layout((1, 1)), + ) + vector = 2 if cutlass.const_expr(tile_n == C96_TAIL_TILE_N) else 4 + if cutlass.const_expr(tile_n == SM80_C96_TILE_N): + copy_b_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + left.element_type, + num_bits_per_copy=left.element_type.width, + ) + copy_b_major = THREADS // TILE_K + copy_b_layout = cute.make_layout( + (copy_b_major, TILE_K), + stride=(1, copy_b_major), + ) + copy_b_value_layout = cute.make_layout( + (tile_n // copy_b_major, 1), + ) + else: + copy_b_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + left.element_type, + num_bits_per_copy=left.element_type.width * vector, + ) + copy_b_major = tile_n // vector + copy_b_layout = cute.make_layout( + (copy_b_major, THREADS // copy_b_major), + stride=(1, copy_b_major), + ) + copy_b_value_layout = cute.make_layout((vector, 1)) + tiled_copy_B = cute.make_tiled_copy_tv( + copy_b_atom, + copy_b_layout, + copy_b_value_layout, + ) + atoms_layout = cute.make_layout( + (THREADS // 16, 16, 1), + stride=(16, 1, 0), + ) + permutation_m = cute.make_layout( + (atoms_layout.shape[0], 4), + stride=(4, 1), + ) + values_n = tile_n // MMA_ATOMS_N + permutation_n = cute.make_layout( + (atoms_layout.shape[1], values_n), + stride=(values_n, 1), + ) + tiled_mma = cute.make_tiled_mma( + cute.nvgpu.MmaUniversalOp(cutlass.Float32), + atoms_layout, + permutation_mnk=(permutation_m, permutation_n, None), + ) + self.kernel( + left, + right, + to_grid, + from_grid, + out, + sA_layout, + sB_layout, + product_layout, + tiled_copy_A, + tiled_copy_B, + tiled_mma, + ).launch( + grid=(left.shape[0], self.channel_tiles, 1), + block=[THREADS, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + left: cute.Tensor, + right: cute.Tensor, + to_grid: cute.Tensor, + from_grid: cute.Tensor, + out: cute.Tensor, + sA_layout: cute.Layout, + sB_layout: cute.Layout, + product_layout: cute.Layout, + tiled_copy_A: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_mma: cute.TiledMma, + ): + tidx, _, _ = cute.arch.thread_idx() + node, channel_tile, _ = cute.arch.block_idx() + channel_tile = channel_tile + self.channel_tile_start + + left_node = left[node, None, None] + right_node = right[node, None, None] + out_node = out[node, None, None] + matrix_b_layout = cute.make_layout( + (self.hidden_channels, PACKED_COEFF_DIM), + stride=(1, self.hidden_channels), + ) + left_b = cute.make_tensor(left_node.iterator, matrix_b_layout) + right_b = cute.make_tensor(right_node.iterator, matrix_b_layout) + + smem = cutlass.utils.SmemAllocator() + sA = smem.allocate_tensor(cutlass.Float32, sA_layout, 16) + sB_left = smem.allocate_tensor(cutlass.Float32, sB_layout, 16) + sB_right = smem.allocate_tensor(cutlass.Float32, sB_layout, 16) + product = smem.allocate_tensor(cutlass.Float32, product_layout, 16) + + for grid_tile in cutlass.range_constexpr(GRID_TILES): + self._dual_projection_product( + to_grid, + left_b, + right_b, + product, + sA, + sB_left, + sB_right, + tiled_copy_A, + tiled_copy_B, + tiled_mma, + tidx, + grid_tile, + channel_tile, + ) + + product_b = cute.make_tensor( + product.iterator, + cute.make_layout( + (self.cta_tiler[1], GRID_SIZE), + stride=(1, self.cta_tiler[1]), + ), + ) + self._backproject( + from_grid, + product_b, + out_node, + sA, + tiled_copy_A, + tiled_mma, + tidx, + channel_tile, + ) + + @cute.jit + def _dual_projection_product( + self, + mA: cute.Tensor, + mB_left: cute.Tensor, + mB_right: cute.Tensor, + mProduct: cute.Tensor, + sA: cute.Tensor, + sB_left: cute.Tensor, + sB_right: cute.Tensor, + tiled_copy_A: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + grid_tile: cutlass.Constexpr, + channel_tile: cutlass.Int32, + ): + thr_mma = tiled_mma.get_slice(tidx) + gA = cute.local_tile( + mA, + tiler=self.cta_tiler, + coord=(grid_tile, 0, None), + proj=(1, None, 1), + ) + gB_left = cute.local_tile( + mB_left, + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(None, 1, 1), + ) + gB_right = cute.local_tile( + mB_right, + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(None, 1, 1), + ) + gProduct = cute.local_tile( + mProduct, + tiler=self.cta_tiler, + coord=(grid_tile, 0, None), + proj=(1, 1, None), + ) + + thr_copy_A = tiled_copy_A.get_slice(tidx) + thr_copy_B = tiled_copy_B.get_slice(tidx) + tAgA = thr_copy_A.partition_S(gA) + tAsA = thr_copy_A.partition_D(sA) + tBgB_left = thr_copy_B.partition_S(gB_left) + tBgB_right = thr_copy_B.partition_S(gB_right) + tBsB_left = thr_copy_B.partition_D(sB_left) + tBsB_right = thr_copy_B.partition_D(sB_right) + + if cutlass.const_expr(self.has_channel_residue): + cB = cute.local_tile( + cute.make_identity_tensor(mB_left.shape), + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(None, 1, 1), + ) + tBcB = thr_copy_B.partition_S(cB) + tBpB = cute.make_rmem_tensor( + cute.make_layout( + ( + tBsB_left.shape[0][1], + cute.size(tBsB_left, mode=[1]), + cute.size(tBsB_left, mode=[2]), + ), + stride=(cute.size(tBsB_left, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in range(tBpB.shape[0]): + for channel in range(tBpB.shape[1]): + tBpB[rest_v, channel, 0] = cute.elem_less( + tBcB[(0, rest_v), channel, 0, 0][0], + mB_left.shape[0], + ) + + cA = cute.local_tile( + cute.make_identity_tensor(mA.shape), + tiler=self.cta_tiler, + coord=(grid_tile, 0, None), + proj=(1, None, 1), + ) + tAcA = thr_copy_A.partition_S(cA) + tApA = cute.make_rmem_tensor( + cute.make_layout( + ( + tAsA.shape[0][1], + cute.size(tAsA, mode=[1]), + cute.size(tAsA, mode=[2]), + ), + stride=(cute.size(tAsA, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in range(tApA.shape[0]): + for row in range(tApA.shape[1]): + tApA[rest_v, row, 0] = cute.elem_less( + tAcA[(0, rest_v), row, 0, 0][0], + mA.shape[0], + ) + + k_pipe_max = cute.size(tAsA, mode=[3]) + k_tile_count = cute.size(tAgA, mode=[3]) + gmem_pipe_read = cutlass.Int32(0) + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, 0], + pred=tApA, + ) + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB_left[None, None, None, gmem_pipe_read], + tBsB_left[None, None, None, 0], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB_left[None, None, None, gmem_pipe_read], + tBsB_left[None, None, None, 0], + ) + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB_right[None, None, None, gmem_pipe_read], + tBsB_right[None, None, None, 0], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB_right[None, None, None, gmem_pipe_read], + tBsB_right[None, None, None, 0], + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + for stage in range(1, STAGES - 1): + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, stage], + pred=tApA, + ) + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB_left[None, None, None, gmem_pipe_read], + tBsB_left[None, None, None, stage], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB_left[None, None, None, gmem_pipe_read], + tBsB_left[None, None, None, stage], + ) + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB_right[None, None, None, gmem_pipe_read], + tBsB_right[None, None, None, stage], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB_right[None, None, None, gmem_pipe_read], + tBsB_right[None, None, None, stage], + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + + tCsA = thr_mma.partition_A(sA) + tCsB_left = thr_mma.partition_B(sB_left) + tCsB_right = thr_mma.partition_B(sB_right) + tCgProduct = thr_mma.partition_C(gProduct) + tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) + tCrB_left = tiled_mma.make_fragment_B(tCsB_left[None, None, None, 0]) + tCrB_right = tiled_mma.make_fragment_B(tCsB_right[None, None, None, 0]) + tCrLeft = tiled_mma.make_fragment_C(tCgProduct) + tCrRight = tiled_mma.make_fragment_C(tCgProduct) + tCrLeft.fill(0.0) + tCrRight.fill(0.0) + + smem_pipe_read = cutlass.Int32(0) + smem_pipe_write = cutlass.Int32(STAGES - 1) + tiles_issued = cutlass.Int32(STAGES - 1) + tCsA_p = tCsA[None, None, None, smem_pipe_read] + tCsB_left_p = tCsB_left[None, None, None, smem_pipe_read] + tCsB_right_p = tCsB_right[None, None, None, smem_pipe_read] + k_block_max = cute.size(tCrA, mode=[2]) + if k_block_max > 1: + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0]) + cute.autovec_copy( + tCsB_left_p[None, None, 0], + tCrB_left[None, None, 0], + ) + cute.autovec_copy( + tCsB_right_p[None, None, 0], + tCrB_right[None, None, 0], + ) + + for _ in range(k_tile_count): + for k_block in range(k_block_max, unroll_full=True): + if k_block == k_block_max - 1: + tCsA_p = tCsA[None, None, None, smem_pipe_read] + tCsB_left_p = tCsB_left[None, None, None, smem_pipe_read] + tCsB_right_p = tCsB_right[None, None, None, smem_pipe_read] + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + + k_block_next = (k_block + 1) % k_block_max + cute.autovec_copy( + tCsA_p[None, None, k_block_next], + tCrA[None, None, k_block_next], + ) + cute.autovec_copy( + tCsB_left_p[None, None, k_block_next], + tCrB_left[None, None, k_block_next], + ) + cute.autovec_copy( + tCsB_right_p[None, None, k_block_next], + tCrB_right[None, None, k_block_next], + ) + if k_block == 0 and tiles_issued < k_tile_count: + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, smem_pipe_write], + pred=tApA, + ) + + cute.gemm( + tiled_mma, + tCrLeft, + tCrA[None, None, k_block], + tCrB_left[None, None, k_block], + tCrLeft, + ) + cute.gemm( + tiled_mma, + tCrRight, + tCrA[None, None, k_block], + tCrB_right[None, None, k_block], + tCrRight, + ) + + if k_block == 0: + if tiles_issued < k_tile_count: + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB_left[None, None, None, gmem_pipe_read], + tBsB_left[None, None, None, smem_pipe_write], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB_left[None, None, None, gmem_pipe_read], + tBsB_left[None, None, None, smem_pipe_write], + ) + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB_right[None, None, None, gmem_pipe_read], + tBsB_right[None, None, None, smem_pipe_write], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB_right[None, None, None, gmem_pipe_read], + tBsB_right[None, None, None, smem_pipe_write], + ) + cute.arch.cp_async_commit_group() + tiles_issued = tiles_issued + 1 + smem_pipe_write = smem_pipe_read + smem_pipe_read = smem_pipe_read + 1 + if smem_pipe_read == STAGES: + smem_pipe_read = cutlass.Int32(0) + gmem_pipe_read = ( + gmem_pipe_read + 1 + if gmem_pipe_read + 1 < k_tile_count + else cutlass.Int32(0) + ) + + cute.arch.cp_async_wait_group(0) + self.cta_sync_barrier.arrive_and_wait() + cProduct = cute.make_identity_tensor(gProduct.shape) + tCpProduct = thr_mma.partition_C(cProduct) + pred = cute.make_rmem_tensor(tCrLeft.layout, cutlass.Boolean) + residue_m = GRID_SIZE - TILE_M * grid_tile + for idx in range(cute.size(tCrLeft.shape)): + pred[idx] = cute.elem_less( + tCpProduct[idx], + (residue_m, self.cta_tiler[1]), + ) + if pred[idx]: + tCrLeft[idx] = tCrLeft[idx].to(cutlass.Float32) * tCrRight[idx].to( + cutlass.Float32 + ) + atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mProduct.element_type, + ) + cute.copy(atom, tCrLeft, tCgProduct, pred=pred) + cute.arch.sync_threads() + + @cute.jit + def _backproject( + self, + mA: cute.Tensor, + mB_shared: cute.Tensor, + mOut: cute.Tensor, + sA: cute.Tensor, + tiled_copy_A: cute.TiledCopy, + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + channel_tile: cutlass.Int32, + ): + thr_mma = tiled_mma.get_slice(tidx) + gA = cute.local_tile( + mA, + tiler=self.cta_tiler, + coord=(0, 0, None), + proj=(1, None, 1), + ) + sB = cute.local_tile( + mB_shared, + tiler=self.cta_tiler, + coord=(0, 0, None), + proj=(None, 1, 1), + ) + gOut = cute.local_tile( + mOut, + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(1, 1, None), + ) + thr_copy_A = tiled_copy_A.get_slice(tidx) + tAgA = thr_copy_A.partition_S(gA) + tAsA = thr_copy_A.partition_D(sA) + + cA = cute.local_tile( + cute.make_identity_tensor(mA.shape), + tiler=self.cta_tiler, + coord=(0, 0, None), + proj=(1, None, 1), + ) + tAcA = thr_copy_A.partition_S(cA) + tApA = cute.make_rmem_tensor( + cute.make_layout( + ( + tAsA.shape[0][1], + cute.size(tAsA, mode=[1]), + cute.size(tAsA, mode=[2]), + ), + stride=(cute.size(tAsA, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in range(tApA.shape[0]): + for row in range(tApA.shape[1]): + tApA[rest_v, row, 0] = cute.elem_less( + tAcA[(0, rest_v), row, 0, 0][0], + PACKED_COEFF_DIM, + ) + + k_pipe_max = cute.size(tAsA, mode=[3]) + k_tile_count = cute.size(tAgA, mode=[3]) + gmem_pipe_read = cutlass.Int32(0) + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, 0], + pred=tApA, + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + for stage in range(1, STAGES - 1): + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, stage], + pred=tApA, + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + + tCsA = thr_mma.partition_A(sA) + tSsB = thr_mma.partition_B(sB) + tCgOut = thr_mma.partition_C(gOut) + tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) + tCrB = tiled_mma.make_fragment_B(tSsB[None, None, None, 0]) + tCrOut = tiled_mma.make_fragment_C(tCgOut) + tCrOut.fill(0.0) + + smem_pipe_read = cutlass.Int32(0) + smem_pipe_write = cutlass.Int32(STAGES - 1) + tiles_issued = cutlass.Int32(STAGES - 1) + logical_k_tile = cutlass.Int32(0) + tCsA_p = tCsA[None, None, None, smem_pipe_read] + k_block_max = cute.size(tCrA, mode=[2]) + if k_block_max > 1: + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0]) + cute.autovec_copy( + tSsB[None, None, 0, logical_k_tile], + tCrB[None, None, 0], + ) + + for _ in range(k_tile_count): + for k_block in range(k_block_max, unroll_full=True): + if k_block == k_block_max - 1: + tCsA_p = tCsA[None, None, None, smem_pipe_read] + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + k_block_next = (k_block + 1) % k_block_max + fragment_k_tile = logical_k_tile + if k_block_max > 1: + if k_block == k_block_max - 1: + fragment_k_tile = ( + logical_k_tile + 1 + if logical_k_tile + 1 < k_tile_count + else logical_k_tile + ) + cute.autovec_copy( + tCsA_p[None, None, k_block_next], + tCrA[None, None, k_block_next], + ) + cute.autovec_copy( + tSsB[None, None, k_block_next, fragment_k_tile], + tCrB[None, None, k_block_next], + ) + if k_block == 0 and tiles_issued < k_tile_count: + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, smem_pipe_write], + pred=tApA, + ) + cute.gemm( + tiled_mma, + tCrOut, + tCrA[None, None, k_block], + tCrB[None, None, k_block], + tCrOut, + ) + if k_block == 0: + cute.arch.cp_async_commit_group() + tiles_issued = tiles_issued + 1 + smem_pipe_write = smem_pipe_read + smem_pipe_read = smem_pipe_read + 1 + if smem_pipe_read == STAGES: + smem_pipe_read = cutlass.Int32(0) + gmem_pipe_read = ( + gmem_pipe_read + 1 + if gmem_pipe_read + 1 < k_tile_count + else cutlass.Int32(0) + ) + logical_k_tile = logical_k_tile + 1 + + cute.arch.cp_async_wait_group(0) + self.cta_sync_barrier.arrive_and_wait() + pred = cute.make_rmem_tensor(tCrOut.layout, cutlass.Boolean) + if cutlass.const_expr(self.has_channel_residue): + cOut = cute.local_tile( + cute.make_identity_tensor(mOut.shape), + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(1, 1, None), + ) + tCpOut = thr_mma.partition_C(cOut) + for idx in range(cute.size(tCrOut.shape)): + pred[idx] = cute.elem_less(tCpOut[idx], mOut.shape) + else: + cOut = cute.make_identity_tensor(gOut.shape) + tCpOut = thr_mma.partition_C(cOut) + for idx in range(cute.size(tCrOut.shape)): + pred[idx] = cute.elem_less( + tCpOut[idx], + (PACKED_COEFF_DIM, self.cta_tiler[1]), + ) + atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mOut.element_type, + ) + cute.copy(atom, tCrOut, tCgOut, pred=pred) + + +class TiledOutputGridProductBackward: + """Tiled dP, projection recomputation, and dual coefficient adjoints.""" + + def __init__( + self, + hidden_channels: int = 192, + *, + tile_n: int = TILE_N, + channel_tile_start: int = 0, + channel_tile_count: int | None = None, + ) -> None: + self.hidden_channels = _validate_hidden_channels(hidden_channels) + self.tile_k = TILE_K + tile_n = int(tile_n) + if tile_n not in (C96_TAIL_TILE_N, SM80_C96_TILE_N, TILE_N): + raise ValueError("output-grid backward tile_n must be 32, 48, or 64") + if tile_n == SM80_C96_TILE_N and self.hidden_channels != 96: + raise ValueError("output-grid N=48 specializes C=96") + channel_tile_start = int(channel_tile_start) + total_channel_tiles = (self.hidden_channels + tile_n - 1) // tile_n + if channel_tile_count is None: + channel_tile_count = total_channel_tiles - channel_tile_start + channel_tile_count = int(channel_tile_count) + if ( + channel_tile_start < 0 + or channel_tile_count <= 0 + or channel_tile_start + channel_tile_count > total_channel_tiles + ): + raise ValueError("invalid output-grid backward channel-tile range") + if tile_n == C96_TAIL_TILE_N and ( + self.hidden_channels != 96 + or channel_tile_start != C96_TAIL_CHANNEL_TILE + or channel_tile_count != 1 + ): + raise ValueError( + "output-grid backward N=32 specializes the C96 K=8 tail panel" + ) + threads = SM80_C96_THREADS if tile_n == SM80_C96_TILE_N else THREADS + self.sm80_c96_n48_panel = tile_n == SM80_C96_TILE_N + self.cta_tiler = (TILE_M, tile_n, self.tile_k) + self.channel_tile_start = channel_tile_start + self.channel_tiles = channel_tile_count + self.has_channel_residue = self.hidden_channels % tile_n != 0 + self.cta_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=threads, + ) + + @cute.jit + def __call__( + self, + grad_out: cute.Tensor, + left: cute.Tensor, + right: cute.Tensor, + to_grid: cute.Tensor, + from_grid: cute.Tensor, + grad_left: cute.Tensor, + grad_right: cute.Tensor, + stream: CUstream, + ): + tile_n = self.cta_tiler[1] + threads = SM80_C96_THREADS if tile_n == SM80_C96_TILE_N else THREADS + sA_row_layout = cute.make_layout( + (TILE_M, self.tile_k, STAGES), + stride=(1, TILE_M + 4, self.tile_k * (TILE_M + 4)), + ) + sA_col_layout = cute.make_layout( + (TILE_M, self.tile_k, STAGES), + stride=(1, TILE_M, self.tile_k * TILE_M), + ) + sB_layout = cute.make_layout( + (tile_n, self.tile_k, STAGES), + stride=(1, tile_n, self.tile_k * tile_n), + ) + grid_layout = cute.make_layout( + (GRID_SIZE, tile_n), + stride=(tile_n, 1), + ) + panel_layout = cute.make_layout( + (TILE_M, tile_n), + stride=(tile_n, 1), + ) + + copy_a_row_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + left.element_type, + num_bits_per_copy=left.element_type.width, + ) + copy_a_row_layout = cute.make_layout( + (threads // self.tile_k, self.tile_k), + stride=(self.tile_k, 1), + ) + tiled_copy_A_row = cute.make_tiled_copy_tv( + copy_a_row_atom, + copy_a_row_layout, + cute.make_layout((1, 1)), + ) + + vector = 4 + copy_a_col_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + left.element_type, + num_bits_per_copy=left.element_type.width * vector, + ) + copy_a_col_major = TILE_M // vector + copy_a_col_layout = cute.make_layout( + (copy_a_col_major, threads // copy_a_col_major), + stride=(1, copy_a_col_major), + ) + tiled_copy_A_col = cute.make_tiled_copy_tv( + copy_a_col_atom, + copy_a_col_layout, + cute.make_layout((vector, 1)), + ) + + if cutlass.const_expr(tile_n == SM80_C96_TILE_N): + # Keep all 128 threads in the copy/MMA contract. Each thread + # issues three scalar N copies, exactly covering 48x8 without a + # partial thread layout or an out-of-bounds 64-column staging tile. + copy_b_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + left.element_type, + num_bits_per_copy=left.element_type.width, + ) + copy_b_major = threads // self.tile_k + copy_b_layout = cute.make_layout( + (copy_b_major, self.tile_k), + stride=(1, copy_b_major), + ) + copy_b_value_layout = cute.make_layout( + (tile_n // copy_b_major, 1), + ) + else: + vector_b = 2 if cutlass.const_expr(tile_n == C96_TAIL_TILE_N) else vector + copy_b_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + left.element_type, + num_bits_per_copy=left.element_type.width * vector_b, + ) + copy_b_major = tile_n // vector_b + copy_b_layout = cute.make_layout( + (copy_b_major, threads // copy_b_major), + stride=(1, copy_b_major), + ) + copy_b_value_layout = cute.make_layout((vector_b, 1)) + tiled_copy_B = cute.make_tiled_copy_tv( + copy_b_atom, + copy_b_layout, + copy_b_value_layout, + ) + + # Follow the Ampere SGEMM thread topology: the universal-FMA atom is + # always tiled over 16 threads in N. N=48 assigns three consecutive N + # values to each thread. With 128 threads this produces a 32x48 MMA + # tile, which divides the 64x48 CTA exactly; a 96-thread 24x48 tile + # would create an unpredicated shared-memory fragment for rows 64..71. + atoms_n = MMA_ATOMS_N + atoms_m = threads // atoms_n + values_n = tile_n // atoms_n + atoms_layout = cute.make_layout( + (atoms_m, atoms_n, 1), + stride=(atoms_n, 1, 0), + ) + permutation_m = cute.make_layout( + (atoms_layout.shape[0], 4), + stride=(4, 1), + ) + permutation_n = cute.make_layout( + (atoms_layout.shape[1], values_n), + stride=(values_n, 1), + ) + tiled_mma = cute.make_tiled_mma( + cute.nvgpu.MmaUniversalOp(cutlass.Float32), + atoms_layout, + permutation_mnk=(permutation_m, permutation_n, None), + ) + + self.kernel( + grad_out, + left, + right, + to_grid, + from_grid, + grad_left, + grad_right, + sA_row_layout, + sA_col_layout, + sB_layout, + grid_layout, + panel_layout, + tiled_copy_A_row, + tiled_copy_A_col, + tiled_copy_B, + tiled_mma, + ).launch( + grid=(left.shape[0], self.channel_tiles, 1), + block=[threads, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + grad_out: cute.Tensor, + left: cute.Tensor, + right: cute.Tensor, + to_grid: cute.Tensor, + from_grid: cute.Tensor, + grad_left: cute.Tensor, + grad_right: cute.Tensor, + sA_row_layout: cute.Layout, + sA_col_layout: cute.Layout, + sB_layout: cute.Layout, + grid_layout: cute.Layout, + panel_layout: cute.Layout, + tiled_copy_A_row: cute.TiledCopy, + tiled_copy_A_col: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_mma: cute.TiledMma, + ): + tidx, _, _ = cute.arch.thread_idx() + node, channel_tile, _ = cute.arch.block_idx() + channel_tile = channel_tile + self.channel_tile_start + + matrix_b_layout = cute.make_layout( + (self.hidden_channels, PACKED_COEFF_DIM), + stride=(1, self.hidden_channels), + ) + grad_out_b = cute.make_tensor( + grad_out[node, None, None].iterator, + matrix_b_layout, + ) + left_b = cute.make_tensor( + left[node, None, None].iterator, + matrix_b_layout, + ) + right_b = cute.make_tensor( + right[node, None, None].iterator, + matrix_b_layout, + ) + grad_left_node = grad_left[node, None, None] + grad_right_node = grad_right[node, None, None] + from_grid_t = cute.make_tensor( + from_grid.iterator, + cute.make_layout( + (GRID_SIZE, PACKED_COEFF_DIM), + stride=(1, GRID_SIZE), + ), + ) + to_grid_t = cute.make_tensor( + to_grid.iterator, + cute.make_layout( + (PACKED_COEFF_DIM, GRID_SIZE), + stride=(1, PACKED_COEFF_DIM), + ), + ) + + smem = cutlass.utils.SmemAllocator() + sA_storage = smem.allocate_tensor( + cutlass.Float32, + sA_row_layout, + 16, + ) + sA_row = sA_storage + sA_col = cute.make_tensor(sA_storage.iterator, sA_col_layout) + if cutlass.const_expr(self.sm80_c96_n48_panel): + sB = smem.allocate_tensor(cutlass.Float32, sB_layout, 16) + adjoint_panel = smem.allocate_tensor( + cutlass.Float32, + panel_layout, + 16, + ) + self._panel_adjoint_backward( + grad_out_b, + left_b, + right_b, + to_grid, + from_grid_t, + to_grid_t, + grad_left_node, + grad_right_node, + adjoint_panel, + sA_row, + sA_col, + sB, + tiled_copy_A_row, + tiled_copy_A_col, + tiled_copy_B, + tiled_mma, + tidx, + channel_tile, + ) + else: + sB_left = smem.allocate_tensor(cutlass.Float32, sB_layout, 16) + sB_right = smem.allocate_tensor(cutlass.Float32, sB_layout, 16) + grad_left_grid = smem.allocate_tensor(cutlass.Float32, grid_layout, 16) + grad_right_grid = smem.allocate_tensor(cutlass.Float32, grid_layout, 16) + + for grid_tile in cutlass.range_constexpr(GRID_TILES): + self._single_projection_to_shared( + from_grid_t, + grad_out_b, + grad_left_grid, + sA_col, + sB_left, + tiled_copy_A_col, + tiled_copy_B, + tiled_mma, + tidx, + grid_tile, + channel_tile, + ) + + for grid_tile in cutlass.range_constexpr(GRID_TILES): + self._dual_projection_adjoint( + to_grid, + left_b, + right_b, + grad_left_grid, + grad_right_grid, + sA_row, + sB_left, + sB_right, + tiled_copy_A_row, + tiled_copy_B, + tiled_mma, + tidx, + grid_tile, + channel_tile, + ) + + grad_left_grid_b = cute.make_tensor( + grad_left_grid.iterator, + cute.make_layout( + (self.cta_tiler[1], GRID_SIZE), + stride=(1, self.cta_tiler[1]), + ), + ) + grad_right_grid_b = cute.make_tensor( + grad_right_grid.iterator, + cute.make_layout( + (self.cta_tiler[1], GRID_SIZE), + stride=(1, self.cta_tiler[1]), + ), + ) + self._dual_backproject( + to_grid_t, + grad_left_grid_b, + grad_right_grid_b, + grad_left_node, + grad_right_node, + sA_col, + tiled_copy_A_col, + tiled_mma, + tidx, + channel_tile, + ) + + @cute.jit + def _panel_adjoint_backward( + self, + grad_out_b: cute.Tensor, + left_b: cute.Tensor, + right_b: cute.Tensor, + to_grid: cute.Tensor, + from_grid_t: cute.Tensor, + to_grid_t: cute.Tensor, + grad_left_node: cute.Tensor, + grad_right_node: cute.Tensor, + adjoint_panel: cute.Tensor, + sA_row: cute.Tensor, + sA_col: cute.Tensor, + sB: cute.Tensor, + tiled_copy_A_row: cute.TiledCopy, + tiled_copy_A_col: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + channel_tile: cutlass.Int32, + ): + """Keep dP in registers and reuse one shared adjoint panel.""" + thr_mma = tiled_mma.get_slice(tidx) + gOut_left = cute.local_tile( + grad_left_node, + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(1, 1, None), + ) + gOut_right = cute.local_tile( + grad_right_node, + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(1, 1, None), + ) + tCgOut_left = thr_mma.partition_C(gOut_left) + tCgOut_right = thr_mma.partition_C(gOut_right) + tCrOut_left = tiled_mma.make_fragment_C(tCgOut_left) + tCrOut_right = tiled_mma.make_fragment_C(tCgOut_right) + tCrOut_left.fill(0.0) + tCrOut_right.fill(0.0) + + panel_b = cute.make_tensor( + adjoint_panel.iterator, + cute.make_layout( + (self.cta_tiler[1], TILE_M), + stride=(1, self.cta_tiler[1]), + ), + ) + for grid_tile in cutlass.range_constexpr(GRID_TILES): + tCrDP = self._projection_fragment( + from_grid_t, + grad_out_b, + adjoint_panel, + sA_col, + sB, + tiled_copy_A_col, + tiled_copy_B, + tiled_mma, + tidx, + grid_tile, + channel_tile, + ) + + tCrRight_grid = self._projection_fragment( + to_grid, + right_b, + adjoint_panel, + sA_row, + sB, + tiled_copy_A_row, + tiled_copy_B, + tiled_mma, + tidx, + grid_tile, + channel_tile, + ) + self._store_adjoint_panel( + tCrDP, + tCrRight_grid, + adjoint_panel, + tiled_mma, + tidx, + grid_tile, + channel_tile, + ) + self._backproject_panel_accumulate( + to_grid_t, + panel_b, + sA_col, + tiled_copy_A_col, + tiled_mma, + tidx, + grid_tile, + tCrOut_left, + ) + + tCrLeft_grid = self._projection_fragment( + to_grid, + left_b, + adjoint_panel, + sA_row, + sB, + tiled_copy_A_row, + tiled_copy_B, + tiled_mma, + tidx, + grid_tile, + channel_tile, + ) + self._store_adjoint_panel( + tCrDP, + tCrLeft_grid, + adjoint_panel, + tiled_mma, + tidx, + grid_tile, + channel_tile, + ) + self._backproject_panel_accumulate( + to_grid_t, + panel_b, + sA_col, + tiled_copy_A_col, + tiled_mma, + tidx, + grid_tile, + tCrOut_right, + ) + + pred = cute.make_rmem_tensor(tCrOut_left.layout, cutlass.Boolean) + if cutlass.const_expr(self.has_channel_residue): + cOut = cute.local_tile( + cute.make_identity_tensor(grad_left_node.shape), + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(1, 1, None), + ) + tCpOut = thr_mma.partition_C(cOut) + for idx in range(cute.size(tCrOut_left.shape)): + pred[idx] = cute.elem_less(tCpOut[idx], grad_left_node.shape) + else: + cOut = cute.make_identity_tensor(gOut_left.shape) + tCpOut = thr_mma.partition_C(cOut) + for idx in range(cute.size(tCrOut_left.shape)): + pred[idx] = cute.elem_less( + tCpOut[idx], + (PACKED_COEFF_DIM, self.cta_tiler[1]), + ) + atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + grad_left_node.element_type, + ) + cute.copy(atom, tCrOut_left, tCgOut_left, pred=pred) + cute.copy(atom, tCrOut_right, tCgOut_right, pred=pred) + + @cute.jit + def _projection_fragment( + self, + mA: cute.Tensor, + mB: cute.Tensor, + mC_layout: cute.Tensor, + sA: cute.Tensor, + sB: cute.Tensor, + tiled_copy_A: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + grid_tile: cutlass.Constexpr, + channel_tile: cutlass.Int32, + ): + """Project one 64-row grid panel and return its register fragment.""" + thr_mma = tiled_mma.get_slice(tidx) + gA = cute.local_tile( + mA, + tiler=self.cta_tiler, + coord=(grid_tile, 0, None), + proj=(1, None, 1), + ) + gB = cute.local_tile( + mB, + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(None, 1, 1), + ) + thr_copy_A = tiled_copy_A.get_slice(tidx) + thr_copy_B = tiled_copy_B.get_slice(tidx) + tAgA = thr_copy_A.partition_S(gA) + tAsA = thr_copy_A.partition_D(sA) + tBgB = thr_copy_B.partition_S(gB) + tBsB = thr_copy_B.partition_D(sB) + + if cutlass.const_expr(self.has_channel_residue): + cB = cute.local_tile( + cute.make_identity_tensor(mB.shape), + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(None, 1, 1), + ) + tBcB = thr_copy_B.partition_S(cB) + tBpB = cute.make_rmem_tensor( + cute.make_layout( + ( + tBsB.shape[0][1], + cute.size(tBsB, mode=[1]), + cute.size(tBsB, mode=[2]), + ), + stride=(cute.size(tBsB, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in range(tBpB.shape[0]): + for channel in range(tBpB.shape[1]): + tBpB[rest_v, channel, 0] = cute.elem_less( + tBcB[(0, rest_v), channel, 0, 0][0], + mB.shape[0], + ) + + cA = cute.local_tile( + cute.make_identity_tensor(mA.shape), + tiler=self.cta_tiler, + coord=(grid_tile, 0, None), + proj=(1, None, 1), + ) + tAcA = thr_copy_A.partition_S(cA) + tApA = cute.make_rmem_tensor( + cute.make_layout( + ( + tAsA.shape[0][1], + cute.size(tAsA, mode=[1]), + cute.size(tAsA, mode=[2]), + ), + stride=(cute.size(tAsA, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in range(tApA.shape[0]): + for row in range(tApA.shape[1]): + tApA[rest_v, row, 0] = cute.elem_less( + tAcA[(0, rest_v), row, 0, 0][0], + mA.shape[0], + ) + + k_tile_count = cute.size(tAgA, mode=[3]) + gmem_pipe_read = cutlass.Int32(0) + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, 0], + pred=tApA, + ) + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB[None, None, None, gmem_pipe_read], + tBsB[None, None, None, 0], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB[None, None, None, gmem_pipe_read], + tBsB[None, None, None, 0], + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + for stage in range(1, STAGES - 1): + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, stage], + pred=tApA, + ) + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB[None, None, None, gmem_pipe_read], + tBsB[None, None, None, stage], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB[None, None, None, gmem_pipe_read], + tBsB[None, None, None, stage], + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + + tCsA = thr_mma.partition_A(sA) + tCsB = thr_mma.partition_B(sB) + tCgC = thr_mma.partition_C(mC_layout) + tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) + tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0]) + tCrC = tiled_mma.make_fragment_C(tCgC) + tCrC.fill(0.0) + + smem_pipe_read = cutlass.Int32(0) + smem_pipe_write = cutlass.Int32(STAGES - 1) + tiles_issued = cutlass.Int32(STAGES - 1) + tCsA_p = tCsA[None, None, None, smem_pipe_read] + tCsB_p = tCsB[None, None, None, smem_pipe_read] + k_block_max = cute.size(tCrA, mode=[2]) + if k_block_max > 1: + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0]) + cute.autovec_copy(tCsB_p[None, None, 0], tCrB[None, None, 0]) + + for _ in range(k_tile_count): + for k_block in range(k_block_max, unroll_full=True): + if k_block == k_block_max - 1: + tCsA_p = tCsA[None, None, None, smem_pipe_read] + tCsB_p = tCsB[None, None, None, smem_pipe_read] + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + k_block_next = (k_block + 1) % k_block_max + cute.autovec_copy( + tCsA_p[None, None, k_block_next], + tCrA[None, None, k_block_next], + ) + cute.autovec_copy( + tCsB_p[None, None, k_block_next], + tCrB[None, None, k_block_next], + ) + if k_block == 0 and tiles_issued < k_tile_count: + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, smem_pipe_write], + pred=tApA, + ) + cute.gemm( + tiled_mma, + tCrC, + tCrA[None, None, k_block], + tCrB[None, None, k_block], + tCrC, + ) + if k_block == 0: + if tiles_issued < k_tile_count: + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB[None, None, None, gmem_pipe_read], + tBsB[None, None, None, smem_pipe_write], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB[None, None, None, gmem_pipe_read], + tBsB[None, None, None, smem_pipe_write], + ) + cute.arch.cp_async_commit_group() + tiles_issued = tiles_issued + 1 + smem_pipe_write = smem_pipe_read + smem_pipe_read = smem_pipe_read + 1 + if smem_pipe_read == STAGES: + smem_pipe_read = cutlass.Int32(0) + gmem_pipe_read = ( + gmem_pipe_read + 1 + if gmem_pipe_read + 1 < k_tile_count + else cutlass.Int32(0) + ) + + cute.arch.cp_async_wait_group(0) + self.cta_sync_barrier.arrive_and_wait() + return tCrC + + @cute.jit + def _store_adjoint_panel( + self, + tCrDP: cute.Tensor, + tCrBranch: cute.Tensor, + panel: cute.Tensor, + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + grid_tile: cutlass.Constexpr, + channel_tile: cutlass.Int32, + ): + thr_mma = tiled_mma.get_slice(tidx) + tCgPanel = thr_mma.partition_C(panel) + cPanel = cute.make_identity_tensor(panel.shape) + tCpPanel = thr_mma.partition_C(cPanel) + pred = cute.make_rmem_tensor(tCrBranch.layout, cutlass.Boolean) + residue_m = GRID_SIZE - TILE_M * grid_tile + residue_n = self.hidden_channels - self.cta_tiler[1] * channel_tile + for idx in range(cute.size(tCrBranch.shape)): + pred[idx] = cute.elem_less( + tCpPanel[idx], + (residue_m, residue_n), + ) + if pred[idx]: + tCrBranch[idx] = tCrDP[idx].to(cutlass.Float32) * tCrBranch[idx].to( + cutlass.Float32 + ) + atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + panel.element_type, + ) + cute.copy(atom, tCrBranch, tCgPanel, pred=pred) + cute.arch.sync_threads() + + @cute.jit + def _backproject_panel_accumulate( + self, + mA: cute.Tensor, + panel_b: cute.Tensor, + sA: cute.Tensor, + tiled_copy_A: cute.TiledCopy, + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + grid_tile: cutlass.Constexpr, + tCrOut: cute.Tensor, + ): + """Accumulate one 64-row adjoint panel into coefficient registers.""" + thr_mma = tiled_mma.get_slice(tidx) + gA = cute.local_tile( + mA, + tiler=self.cta_tiler, + coord=(0, 0, None), + proj=(1, None, 1), + ) + tAgA = tiled_copy_A.get_slice(tidx).partition_S(gA) + tAsA = tiled_copy_A.get_slice(tidx).partition_D(sA) + # Expose the K-tile mode required by the MMA B-fragment contract. + sB = cute.local_tile( + panel_b, + tiler=self.cta_tiler, + coord=(0, 0, None), + proj=(None, 1, 1), + ) + tSsB = thr_mma.partition_B(sB) + + cA = cute.local_tile( + cute.make_identity_tensor(mA.shape), + tiler=self.cta_tiler, + coord=(0, 0, None), + proj=(1, None, 1), + ) + tAcA = tiled_copy_A.get_slice(tidx).partition_S(cA) + tApA = cute.make_rmem_tensor( + cute.make_layout( + ( + tAsA.shape[0][1], + cute.size(tAsA, mode=[1]), + cute.size(tAsA, mode=[2]), + ), + stride=(cute.size(tAsA, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in range(tApA.shape[0]): + for row in range(tApA.shape[1]): + tApA[rest_v, row, 0] = cute.elem_less( + tAcA[(0, rest_v), row, 0, 0][0], + PACKED_COEFF_DIM, + ) + + panel_k_tiles = TILE_M // self.tile_k + if cutlass.const_expr(grid_tile == GRID_TILES - 1): + panel_k_tiles = (GRID_SIZE - TILE_M * grid_tile) // self.tile_k + panel_k_start = grid_tile * (TILE_M // self.tile_k) + gmem_pipe_read = cutlass.Int32(panel_k_start) + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, 0], + pred=tApA, + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + for stage in range(1, STAGES - 1): + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, stage], + pred=tApA, + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + + tCsA = thr_mma.partition_A(sA) + tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) + tCrB = tiled_mma.make_fragment_B(tSsB[None, None, None, 0]) + smem_pipe_read = cutlass.Int32(0) + smem_pipe_write = cutlass.Int32(STAGES - 1) + tiles_issued = cutlass.Int32(STAGES - 1) + logical_k_tile = cutlass.Int32(0) + tCsA_p = tCsA[None, None, None, smem_pipe_read] + k_block_max = cute.size(tCrA, mode=[2]) + if k_block_max > 1: + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0]) + cute.autovec_copy( + tSsB[None, None, 0, logical_k_tile], + tCrB[None, None, 0], + ) + + for _ in range(panel_k_tiles): + for k_block in range(k_block_max, unroll_full=True): + if k_block == k_block_max - 1: + tCsA_p = tCsA[None, None, None, smem_pipe_read] + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + k_block_next = (k_block + 1) % k_block_max + fragment_k_tile = logical_k_tile + if k_block_max > 1 and k_block == k_block_max - 1: + fragment_k_tile = ( + logical_k_tile + 1 + if logical_k_tile + 1 < panel_k_tiles + else logical_k_tile + ) + cute.autovec_copy( + tCsA_p[None, None, k_block_next], + tCrA[None, None, k_block_next], + ) + cute.autovec_copy( + tSsB[ + None, + None, + k_block_next, + fragment_k_tile, + ], + tCrB[None, None, k_block_next], + ) + if k_block == 0 and tiles_issued < panel_k_tiles: + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, smem_pipe_write], + pred=tApA, + ) + cute.gemm( + tiled_mma, + tCrOut, + tCrA[None, None, k_block], + tCrB[None, None, k_block], + tCrOut, + ) + if k_block == 0: + cute.arch.cp_async_commit_group() + tiles_issued = tiles_issued + 1 + smem_pipe_write = smem_pipe_read + smem_pipe_read = smem_pipe_read + 1 + if smem_pipe_read == STAGES: + smem_pipe_read = cutlass.Int32(0) + gmem_pipe_read = ( + gmem_pipe_read + 1 + if gmem_pipe_read + 1 < panel_k_start + panel_k_tiles + else cutlass.Int32(panel_k_start) + ) + logical_k_tile = logical_k_tile + 1 + + cute.arch.cp_async_wait_group(0) + self.cta_sync_barrier.arrive_and_wait() + + @cute.jit + def _single_projection_to_shared( + self, + mA: cute.Tensor, + mB: cute.Tensor, + mGrid: cute.Tensor, + sA: cute.Tensor, + sB: cute.Tensor, + tiled_copy_A: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + grid_tile: cutlass.Constexpr, + channel_tile: cutlass.Int32, + ): + thr_mma = tiled_mma.get_slice(tidx) + gA = cute.local_tile( + mA, + tiler=self.cta_tiler, + coord=(grid_tile, 0, None), + proj=(1, None, 1), + ) + gB = cute.local_tile( + mB, + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(None, 1, 1), + ) + gGrid = cute.local_tile( + mGrid, + tiler=self.cta_tiler, + coord=(grid_tile, 0, None), + proj=(1, 1, None), + ) + thr_copy_A = tiled_copy_A.get_slice(tidx) + thr_copy_B = tiled_copy_B.get_slice(tidx) + tAgA = thr_copy_A.partition_S(gA) + tAsA = thr_copy_A.partition_D(sA) + tBgB = thr_copy_B.partition_S(gB) + tBsB = thr_copy_B.partition_D(sB) + + if cutlass.const_expr(self.has_channel_residue): + cB = cute.local_tile( + cute.make_identity_tensor(mB.shape), + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(None, 1, 1), + ) + tBcB = thr_copy_B.partition_S(cB) + tBpB = cute.make_rmem_tensor( + cute.make_layout( + ( + tBsB.shape[0][1], + cute.size(tBsB, mode=[1]), + cute.size(tBsB, mode=[2]), + ), + stride=(cute.size(tBsB, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in range(tBpB.shape[0]): + for channel in range(tBpB.shape[1]): + tBpB[rest_v, channel, 0] = cute.elem_less( + tBcB[(0, rest_v), channel, 0, 0][0], + mB.shape[0], + ) + + cA = cute.local_tile( + cute.make_identity_tensor(mA.shape), + tiler=self.cta_tiler, + coord=(grid_tile, 0, None), + proj=(1, None, 1), + ) + tAcA = thr_copy_A.partition_S(cA) + tApA = cute.make_rmem_tensor( + cute.make_layout( + ( + tAsA.shape[0][1], + cute.size(tAsA, mode=[1]), + cute.size(tAsA, mode=[2]), + ), + stride=(cute.size(tAsA, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in range(tApA.shape[0]): + for row in range(tApA.shape[1]): + tApA[rest_v, row, 0] = cute.elem_less( + tAcA[(0, rest_v), row, 0, 0][0], + mA.shape[0], + ) + + k_tile_count = cute.size(tAgA, mode=[3]) + gmem_pipe_read = cutlass.Int32(0) + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, 0], + pred=tApA, + ) + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB[None, None, None, gmem_pipe_read], + tBsB[None, None, None, 0], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB[None, None, None, gmem_pipe_read], + tBsB[None, None, None, 0], + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + for stage in range(1, STAGES - 1): + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, stage], + pred=tApA, + ) + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB[None, None, None, gmem_pipe_read], + tBsB[None, None, None, stage], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB[None, None, None, gmem_pipe_read], + tBsB[None, None, None, stage], + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + + tCsA = thr_mma.partition_A(sA) + tCsB = thr_mma.partition_B(sB) + tCgGrid = thr_mma.partition_C(gGrid) + tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) + tCrB = tiled_mma.make_fragment_B(tCsB[None, None, None, 0]) + tCrGrid = tiled_mma.make_fragment_C(tCgGrid) + tCrGrid.fill(0.0) + + smem_pipe_read = cutlass.Int32(0) + smem_pipe_write = cutlass.Int32(STAGES - 1) + tiles_issued = cutlass.Int32(STAGES - 1) + tCsA_p = tCsA[None, None, None, smem_pipe_read] + tCsB_p = tCsB[None, None, None, smem_pipe_read] + k_block_max = cute.size(tCrA, mode=[2]) + if k_block_max > 1: + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0]) + cute.autovec_copy(tCsB_p[None, None, 0], tCrB[None, None, 0]) + + for _ in range(k_tile_count): + for k_block in range(k_block_max, unroll_full=True): + if k_block == k_block_max - 1: + tCsA_p = tCsA[None, None, None, smem_pipe_read] + tCsB_p = tCsB[None, None, None, smem_pipe_read] + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + k_block_next = (k_block + 1) % k_block_max + cute.autovec_copy( + tCsA_p[None, None, k_block_next], + tCrA[None, None, k_block_next], + ) + cute.autovec_copy( + tCsB_p[None, None, k_block_next], + tCrB[None, None, k_block_next], + ) + if k_block == 0 and tiles_issued < k_tile_count: + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, smem_pipe_write], + pred=tApA, + ) + cute.gemm( + tiled_mma, + tCrGrid, + tCrA[None, None, k_block], + tCrB[None, None, k_block], + tCrGrid, + ) + if k_block == 0: + if tiles_issued < k_tile_count: + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB[None, None, None, gmem_pipe_read], + tBsB[None, None, None, smem_pipe_write], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB[None, None, None, gmem_pipe_read], + tBsB[None, None, None, smem_pipe_write], + ) + cute.arch.cp_async_commit_group() + tiles_issued = tiles_issued + 1 + smem_pipe_write = smem_pipe_read + smem_pipe_read = smem_pipe_read + 1 + if smem_pipe_read == STAGES: + smem_pipe_read = cutlass.Int32(0) + gmem_pipe_read = ( + gmem_pipe_read + 1 + if gmem_pipe_read + 1 < k_tile_count + else cutlass.Int32(0) + ) + + cute.arch.cp_async_wait_group(0) + self.cta_sync_barrier.arrive_and_wait() + cGrid = cute.make_identity_tensor(gGrid.shape) + tCpGrid = thr_mma.partition_C(cGrid) + pred = cute.make_rmem_tensor(tCrGrid.layout, cutlass.Boolean) + residue_m = GRID_SIZE - TILE_M * grid_tile + for idx in range(cute.size(tCrGrid.shape)): + pred[idx] = cute.elem_less( + tCpGrid[idx], + (residue_m, self.cta_tiler[1]), + ) + atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mGrid.element_type, + ) + cute.copy(atom, tCrGrid, tCgGrid, pred=pred) + cute.arch.sync_threads() + + @cute.jit + def _dual_projection_adjoint( + self, + mA: cute.Tensor, + mB_left: cute.Tensor, + mB_right: cute.Tensor, + mGrad_left_grid: cute.Tensor, + mGrad_right_grid: cute.Tensor, + sA: cute.Tensor, + sB_left: cute.Tensor, + sB_right: cute.Tensor, + tiled_copy_A: cute.TiledCopy, + tiled_copy_B: cute.TiledCopy, + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + grid_tile: cutlass.Constexpr, + channel_tile: cutlass.Int32, + ): + thr_mma = tiled_mma.get_slice(tidx) + gA = cute.local_tile( + mA, + tiler=self.cta_tiler, + coord=(grid_tile, 0, None), + proj=(1, None, 1), + ) + gB_left = cute.local_tile( + mB_left, + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(None, 1, 1), + ) + gB_right = cute.local_tile( + mB_right, + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(None, 1, 1), + ) + gGrad_left = cute.local_tile( + mGrad_left_grid, + tiler=self.cta_tiler, + coord=(grid_tile, 0, None), + proj=(1, 1, None), + ) + gGrad_right = cute.local_tile( + mGrad_right_grid, + tiler=self.cta_tiler, + coord=(grid_tile, 0, None), + proj=(1, 1, None), + ) + + thr_copy_A = tiled_copy_A.get_slice(tidx) + thr_copy_B = tiled_copy_B.get_slice(tidx) + tAgA = thr_copy_A.partition_S(gA) + tAsA = thr_copy_A.partition_D(sA) + tBgB_left = thr_copy_B.partition_S(gB_left) + tBgB_right = thr_copy_B.partition_S(gB_right) + tBsB_left = thr_copy_B.partition_D(sB_left) + tBsB_right = thr_copy_B.partition_D(sB_right) + + if cutlass.const_expr(self.has_channel_residue): + cB = cute.local_tile( + cute.make_identity_tensor(mB_left.shape), + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(None, 1, 1), + ) + tBcB = thr_copy_B.partition_S(cB) + tBpB = cute.make_rmem_tensor( + cute.make_layout( + ( + tBsB_left.shape[0][1], + cute.size(tBsB_left, mode=[1]), + cute.size(tBsB_left, mode=[2]), + ), + stride=(cute.size(tBsB_left, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in range(tBpB.shape[0]): + for channel in range(tBpB.shape[1]): + tBpB[rest_v, channel, 0] = cute.elem_less( + tBcB[(0, rest_v), channel, 0, 0][0], + mB_left.shape[0], + ) + + cA = cute.local_tile( + cute.make_identity_tensor(mA.shape), + tiler=self.cta_tiler, + coord=(grid_tile, 0, None), + proj=(1, None, 1), + ) + tAcA = thr_copy_A.partition_S(cA) + tApA = cute.make_rmem_tensor( + cute.make_layout( + ( + tAsA.shape[0][1], + cute.size(tAsA, mode=[1]), + cute.size(tAsA, mode=[2]), + ), + stride=(cute.size(tAsA, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in range(tApA.shape[0]): + for row in range(tApA.shape[1]): + tApA[rest_v, row, 0] = cute.elem_less( + tAcA[(0, rest_v), row, 0, 0][0], + mA.shape[0], + ) + + k_tile_count = cute.size(tAgA, mode=[3]) + gmem_pipe_read = cutlass.Int32(0) + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, 0], + pred=tApA, + ) + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB_left[None, None, None, gmem_pipe_read], + tBsB_left[None, None, None, 0], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB_left[None, None, None, gmem_pipe_read], + tBsB_left[None, None, None, 0], + ) + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB_right[None, None, None, gmem_pipe_read], + tBsB_right[None, None, None, 0], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB_right[None, None, None, gmem_pipe_read], + tBsB_right[None, None, None, 0], + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + for stage in range(1, STAGES - 1): + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, stage], + pred=tApA, + ) + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB_left[None, None, None, gmem_pipe_read], + tBsB_left[None, None, None, stage], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB_left[None, None, None, gmem_pipe_read], + tBsB_left[None, None, None, stage], + ) + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB_right[None, None, None, gmem_pipe_read], + tBsB_right[None, None, None, stage], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB_right[None, None, None, gmem_pipe_read], + tBsB_right[None, None, None, stage], + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + + tCsA = thr_mma.partition_A(sA) + tCsB_left = thr_mma.partition_B(sB_left) + tCsB_right = thr_mma.partition_B(sB_right) + tCgGrad_left = thr_mma.partition_C(gGrad_left) + tCgGrad_right = thr_mma.partition_C(gGrad_right) + tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) + tCrB_left = tiled_mma.make_fragment_B(tCsB_left[None, None, None, 0]) + tCrB_right = tiled_mma.make_fragment_B(tCsB_right[None, None, None, 0]) + tCrLeft = tiled_mma.make_fragment_C(tCgGrad_left) + tCrRight = tiled_mma.make_fragment_C(tCgGrad_right) + tCrLeft.fill(0.0) + tCrRight.fill(0.0) + + smem_pipe_read = cutlass.Int32(0) + smem_pipe_write = cutlass.Int32(STAGES - 1) + tiles_issued = cutlass.Int32(STAGES - 1) + tCsA_p = tCsA[None, None, None, smem_pipe_read] + tCsB_left_p = tCsB_left[None, None, None, smem_pipe_read] + tCsB_right_p = tCsB_right[None, None, None, smem_pipe_read] + k_block_max = cute.size(tCrA, mode=[2]) + if k_block_max > 1: + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0]) + cute.autovec_copy( + tCsB_left_p[None, None, 0], + tCrB_left[None, None, 0], + ) + cute.autovec_copy( + tCsB_right_p[None, None, 0], + tCrB_right[None, None, 0], + ) + + for _ in range(k_tile_count): + for k_block in range(k_block_max, unroll_full=True): + if k_block == k_block_max - 1: + tCsA_p = tCsA[None, None, None, smem_pipe_read] + tCsB_left_p = tCsB_left[None, None, None, smem_pipe_read] + tCsB_right_p = tCsB_right[None, None, None, smem_pipe_read] + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + k_block_next = (k_block + 1) % k_block_max + cute.autovec_copy( + tCsA_p[None, None, k_block_next], + tCrA[None, None, k_block_next], + ) + cute.autovec_copy( + tCsB_left_p[None, None, k_block_next], + tCrB_left[None, None, k_block_next], + ) + cute.autovec_copy( + tCsB_right_p[None, None, k_block_next], + tCrB_right[None, None, k_block_next], + ) + if k_block == 0 and tiles_issued < k_tile_count: + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, smem_pipe_write], + pred=tApA, + ) + cute.gemm( + tiled_mma, + tCrLeft, + tCrA[None, None, k_block], + tCrB_left[None, None, k_block], + tCrLeft, + ) + cute.gemm( + tiled_mma, + tCrRight, + tCrA[None, None, k_block], + tCrB_right[None, None, k_block], + tCrRight, + ) + if k_block == 0: + if tiles_issued < k_tile_count: + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB_left[None, None, None, gmem_pipe_read], + tBsB_left[None, None, None, smem_pipe_write], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB_left[None, None, None, gmem_pipe_read], + tBsB_left[None, None, None, smem_pipe_write], + ) + if cutlass.const_expr(self.has_channel_residue): + cute.copy( + tiled_copy_B, + tBgB_right[None, None, None, gmem_pipe_read], + tBsB_right[None, None, None, smem_pipe_write], + pred=tBpB, + ) + else: + cute.copy( + tiled_copy_B, + tBgB_right[None, None, None, gmem_pipe_read], + tBsB_right[None, None, None, smem_pipe_write], + ) + cute.arch.cp_async_commit_group() + tiles_issued = tiles_issued + 1 + smem_pipe_write = smem_pipe_read + smem_pipe_read = smem_pipe_read + 1 + if smem_pipe_read == STAGES: + smem_pipe_read = cutlass.Int32(0) + gmem_pipe_read = ( + gmem_pipe_read + 1 + if gmem_pipe_read + 1 < k_tile_count + else cutlass.Int32(0) + ) + + cute.arch.cp_async_wait_group(0) + self.cta_sync_barrier.arrive_and_wait() + cGrid = cute.make_identity_tensor(gGrad_left.shape) + tCpGrid = thr_mma.partition_C(cGrid) + pred = cute.make_rmem_tensor(tCrLeft.layout, cutlass.Boolean) + tCrDP = tiled_mma.make_fragment_C(tCgGrad_left) + tCrDP.fill(0.0) + residue_m = GRID_SIZE - TILE_M * grid_tile + for idx in range(cute.size(tCrLeft.shape)): + pred[idx] = cute.elem_less( + tCpGrid[idx], + (residue_m, self.cta_tiler[1]), + ) + atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mGrad_left_grid.element_type, + ) + cute.copy(atom, tCgGrad_left, tCrDP, pred=pred) + for idx in range(cute.size(tCrLeft.shape)): + if pred[idx]: + left_value = tCrLeft[idx].to(cutlass.Float32) + right_value = tCrRight[idx].to(cutlass.Float32) + grad_product = tCrDP[idx].to(cutlass.Float32) + tCrLeft[idx] = grad_product * right_value + tCrRight[idx] = grad_product * left_value + cute.copy(atom, tCrLeft, tCgGrad_left, pred=pred) + cute.copy(atom, tCrRight, tCgGrad_right, pred=pred) + cute.arch.sync_threads() + + @cute.jit + def _dual_backproject( + self, + mA: cute.Tensor, + mB_left_shared: cute.Tensor, + mB_right_shared: cute.Tensor, + mOut_left: cute.Tensor, + mOut_right: cute.Tensor, + sA: cute.Tensor, + tiled_copy_A: cute.TiledCopy, + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + channel_tile: cutlass.Int32, + ): + thr_mma = tiled_mma.get_slice(tidx) + gA = cute.local_tile( + mA, + tiler=self.cta_tiler, + coord=(0, 0, None), + proj=(1, None, 1), + ) + sB_left = cute.local_tile( + mB_left_shared, + tiler=self.cta_tiler, + coord=(0, 0, None), + proj=(None, 1, 1), + ) + sB_right = cute.local_tile( + mB_right_shared, + tiler=self.cta_tiler, + coord=(0, 0, None), + proj=(None, 1, 1), + ) + gOut_left = cute.local_tile( + mOut_left, + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(1, 1, None), + ) + gOut_right = cute.local_tile( + mOut_right, + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(1, 1, None), + ) + thr_copy_A = tiled_copy_A.get_slice(tidx) + tAgA = thr_copy_A.partition_S(gA) + tAsA = thr_copy_A.partition_D(sA) + + cA = cute.local_tile( + cute.make_identity_tensor(mA.shape), + tiler=self.cta_tiler, + coord=(0, 0, None), + proj=(1, None, 1), + ) + tAcA = thr_copy_A.partition_S(cA) + tApA = cute.make_rmem_tensor( + cute.make_layout( + ( + tAsA.shape[0][1], + cute.size(tAsA, mode=[1]), + cute.size(tAsA, mode=[2]), + ), + stride=(cute.size(tAsA, mode=[1]), 1, 0), + ), + cutlass.Boolean, + ) + for rest_v in range(tApA.shape[0]): + for row in range(tApA.shape[1]): + tApA[rest_v, row, 0] = cute.elem_less( + tAcA[(0, rest_v), row, 0, 0][0], + PACKED_COEFF_DIM, + ) + + k_tile_count = cute.size(tAgA, mode=[3]) + gmem_pipe_read = cutlass.Int32(0) + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, 0], + pred=tApA, + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + for stage in range(1, STAGES - 1): + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, stage], + pred=tApA, + ) + cute.arch.cp_async_commit_group() + gmem_pipe_read = gmem_pipe_read + 1 + + tCsA = thr_mma.partition_A(sA) + tSsB_left = thr_mma.partition_B(sB_left) + tSsB_right = thr_mma.partition_B(sB_right) + tCgOut_left = thr_mma.partition_C(gOut_left) + tCgOut_right = thr_mma.partition_C(gOut_right) + tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) + tCrB_left = tiled_mma.make_fragment_B(tSsB_left[None, None, None, 0]) + tCrB_right = tiled_mma.make_fragment_B(tSsB_right[None, None, None, 0]) + tCrOut_left = tiled_mma.make_fragment_C(tCgOut_left) + tCrOut_right = tiled_mma.make_fragment_C(tCgOut_right) + tCrOut_left.fill(0.0) + tCrOut_right.fill(0.0) + + smem_pipe_read = cutlass.Int32(0) + smem_pipe_write = cutlass.Int32(STAGES - 1) + tiles_issued = cutlass.Int32(STAGES - 1) + logical_k_tile = cutlass.Int32(0) + tCsA_p = tCsA[None, None, None, smem_pipe_read] + k_block_max = cute.size(tCrA, mode=[2]) + if k_block_max > 1: + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0]) + cute.autovec_copy( + tSsB_left[None, None, 0, logical_k_tile], + tCrB_left[None, None, 0], + ) + cute.autovec_copy( + tSsB_right[None, None, 0, logical_k_tile], + tCrB_right[None, None, 0], + ) + + for _ in range(k_tile_count): + for k_block in range(k_block_max, unroll_full=True): + if k_block == k_block_max - 1: + tCsA_p = tCsA[None, None, None, smem_pipe_read] + cute.arch.cp_async_wait_group(STAGES - 2) + self.cta_sync_barrier.arrive_and_wait() + k_block_next = (k_block + 1) % k_block_max + fragment_k_tile = logical_k_tile + if k_block_max > 1: + if k_block == k_block_max - 1: + fragment_k_tile = ( + logical_k_tile + 1 + if logical_k_tile + 1 < k_tile_count + else logical_k_tile + ) + cute.autovec_copy( + tCsA_p[None, None, k_block_next], + tCrA[None, None, k_block_next], + ) + cute.autovec_copy( + tSsB_left[ + None, + None, + k_block_next, + fragment_k_tile, + ], + tCrB_left[None, None, k_block_next], + ) + cute.autovec_copy( + tSsB_right[ + None, + None, + k_block_next, + fragment_k_tile, + ], + tCrB_right[None, None, k_block_next], + ) + if k_block == 0 and tiles_issued < k_tile_count: + cute.copy( + tiled_copy_A, + tAgA[None, None, None, gmem_pipe_read], + tAsA[None, None, None, smem_pipe_write], + pred=tApA, + ) + cute.gemm( + tiled_mma, + tCrOut_left, + tCrA[None, None, k_block], + tCrB_left[None, None, k_block], + tCrOut_left, + ) + cute.gemm( + tiled_mma, + tCrOut_right, + tCrA[None, None, k_block], + tCrB_right[None, None, k_block], + tCrOut_right, + ) + if k_block == 0: + cute.arch.cp_async_commit_group() + tiles_issued = tiles_issued + 1 + smem_pipe_write = smem_pipe_read + smem_pipe_read = smem_pipe_read + 1 + if smem_pipe_read == STAGES: + smem_pipe_read = cutlass.Int32(0) + gmem_pipe_read = ( + gmem_pipe_read + 1 + if gmem_pipe_read + 1 < k_tile_count + else cutlass.Int32(0) + ) + logical_k_tile = logical_k_tile + 1 + + cute.arch.cp_async_wait_group(0) + self.cta_sync_barrier.arrive_and_wait() + pred = cute.make_rmem_tensor(tCrOut_left.layout, cutlass.Boolean) + if cutlass.const_expr(self.has_channel_residue): + cOut = cute.local_tile( + cute.make_identity_tensor(mOut_left.shape), + tiler=self.cta_tiler, + coord=(0, channel_tile, None), + proj=(1, 1, None), + ) + tCpOut = thr_mma.partition_C(cOut) + for idx in range(cute.size(tCrOut_left.shape)): + pred[idx] = cute.elem_less(tCpOut[idx], mOut_left.shape) + else: + cOut = cute.make_identity_tensor(gOut_left.shape) + tCpOut = thr_mma.partition_C(cOut) + for idx in range(cute.size(tCrOut_left.shape)): + pred[idx] = cute.elem_less( + tCpOut[idx], + (PACKED_COEFF_DIM, self.cta_tiler[1]), + ) + atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + mOut_left.element_type, + ) + cute.copy(atom, tCrOut_left, tCgOut_left, pred=pred) + cute.copy(atom, tCrOut_right, tCgOut_right, pred=pred) + + +def compile_tiled_output_grid_product( + device_index: int | None = None, + compute_capability: tuple[int, int] | None = None, + hidden_channels: int = 192, + tile_n: int = TILE_N, + channel_tile_start: int = 0, + channel_tile_count: int | None = None, +) -> Callable: + """Compile one forward artifact with symbolic runtime node count.""" + import torch + + hidden_channels = _validate_hidden_channels(hidden_channels) + tile_n = int(tile_n) + if tile_n not in (C96_TAIL_TILE_N, SM80_C96_TILE_N, TILE_N): + raise ValueError("output-grid forward tile_n must be 32, 48, or 64") + if device_index is None: + device_index = torch.cuda.current_device() + actual_capability = tuple(torch.cuda.get_device_capability(device_index)) + if ( + compute_capability is not None + and tuple(compute_capability) != actual_capability + ): + raise ValueError("compile target does not match the selected CUDA device") + if tile_n == SM80_C96_TILE_N and ( + actual_capability not in runtime_policy.SM80_PROFILE_CAPABILITIES + or hidden_channels != 96 + ): + raise ValueError("output-grid forward N=48 requires SM80-family and C=96") + if tile_n == C96_TAIL_TILE_N and ( + actual_capability != (9, 0) + or hidden_channels != 96 + or int(channel_tile_start) != C96_TAIL_CHANNEL_TILE + or int(channel_tile_count or 0) != 1 + ): + raise ValueError("output-grid forward N=32 tail requires sm90 and C=96") + if ( + tile_n == TILE_N + and channel_tile_count is not None + and ( + actual_capability != (9, 0) + or hidden_channels != 96 + or int(channel_tile_start) != 0 + or int(channel_tile_count) != 1 + ) + ): + raise ValueError( + "partial output-grid forward N=64 launch requires sm90 C=96 base panel" + ) + with torch.cuda.device(device_index): + nodes = cute.sym_int64() + fake_coeff = make_fake_compact_tensor( + cutlass.Float32, + (nodes, PACKED_COEFF_DIM, hidden_channels), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + fake_to_grid = make_fake_compact_tensor( + cutlass.Float32, + (GRID_SIZE, PACKED_COEFF_DIM), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_from_grid = make_fake_compact_tensor( + cutlass.Float32, + (PACKED_COEFF_DIM, GRID_SIZE), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + TiledOutputGridProductForward( + hidden_channels, + tile_n=tile_n, + channel_tile_start=channel_tile_start, + channel_tile_count=channel_tile_count, + ), + fake_coeff, + fake_coeff, + fake_to_grid, + fake_from_grid, + fake_coeff, + fake_stream, + options="--enable-tvm-ffi", + ) + + +def compile_tiled_output_grid_product_backward( + device_index: int | None = None, + compute_capability: tuple[int, int] | None = None, + hidden_channels: int = 192, + tile_n: int = TILE_N, + channel_tile_start: int = 0, + channel_tile_count: int | None = None, +) -> Callable: + """Compile one first-backward artifact with symbolic node count.""" + import torch + + hidden_channels = _validate_hidden_channels(hidden_channels) + tile_n = int(tile_n) + if tile_n not in (C96_TAIL_TILE_N, SM80_C96_TILE_N, TILE_N): + raise ValueError("output-grid backward tile_n must be 32, 48, or 64") + if device_index is None: + device_index = torch.cuda.current_device() + actual_capability = tuple(torch.cuda.get_device_capability(device_index)) + if ( + compute_capability is not None + and tuple(compute_capability) != actual_capability + ): + raise ValueError("compile target does not match the selected CUDA device") + if tile_n == SM80_C96_TILE_N and ( + actual_capability not in runtime_policy.SM80_PROFILE_CAPABILITIES + or hidden_channels != 96 + ): + raise ValueError( + "output-grid N=48 panel adjoint requires SM80-family, C=96, and K=8" + ) + if tile_n == C96_TAIL_TILE_N and ( + actual_capability != (9, 0) + or hidden_channels != 96 + or int(channel_tile_start) != C96_TAIL_CHANNEL_TILE + or int(channel_tile_count or 0) != 1 + ): + raise ValueError("output-grid backward N=32 tail requires sm90, C=96, and K=8") + if ( + tile_n == TILE_N + and channel_tile_count is not None + and ( + actual_capability != (9, 0) + or hidden_channels != 96 + or int(channel_tile_start) != 0 + or int(channel_tile_count) != 1 + ) + ): + raise ValueError( + "partial output-grid backward N=64 launch requires sm90 C=96 K=8 base panel" + ) + with torch.cuda.device(device_index): + nodes = cute.sym_int64() + fake_coeff = make_fake_compact_tensor( + cutlass.Float32, + (nodes, PACKED_COEFF_DIM, hidden_channels), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + fake_to_grid = make_fake_compact_tensor( + cutlass.Float32, + (GRID_SIZE, PACKED_COEFF_DIM), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_from_grid = make_fake_compact_tensor( + cutlass.Float32, + (PACKED_COEFF_DIM, GRID_SIZE), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + TiledOutputGridProductBackward( + hidden_channels, + tile_n=tile_n, + channel_tile_start=channel_tile_start, + channel_tile_count=channel_tile_count, + ), + fake_coeff, + fake_coeff, + fake_coeff, + fake_to_grid, + fake_from_grid, + fake_coeff, + fake_coeff, + fake_stream, + options="--enable-tvm-ffi", + ) + + +@lru_cache(maxsize=16) +def _compiled_tiled_forward( + device_index: int, + compute_capability: tuple[int, int], + hidden_channels: int, + tile_n: int, + channel_tile_start: int, + channel_tile_count: int | None, +) -> Callable: + return compile_tiled_output_grid_product( + device_index, + compute_capability, + hidden_channels, + tile_n, + channel_tile_start, + channel_tile_count, + ) + + +@lru_cache(maxsize=16) +def _compiled_tiled_backward( + device_index: int, + compute_capability: tuple[int, int], + hidden_channels: int, + tile_n: int, + channel_tile_start: int, + channel_tile_count: int | None, +) -> Callable: + return compile_tiled_output_grid_product_backward( + device_index, + compute_capability, + hidden_channels, + tile_n, + channel_tile_start, + channel_tile_count, + ) + + +def _compile_key(tensor) -> tuple[int, tuple[int, int], int]: + import torch + + device_index = tensor.device.index + if device_index is None: + device_index = torch.cuda.current_device() + compute_capability = tuple(torch.cuda.get_device_capability(device_index)) + return int(device_index), compute_capability, int(tensor.shape[2]) + + +def _validate_tensors(left, right, to_grid, from_grid, out=None) -> None: + import torch + + tensors = (left, right, to_grid, from_grid) + if ( + any(not tensor.is_cuda for tensor in tensors) + or any(tensor.dtype != torch.float32 for tensor in tensors) + or any(tensor.device != left.device for tensor in tensors) + or left.ndim != 3 + or left.shape[0] <= 0 + or int(left.shape[1]) != PACKED_COEFF_DIM + or int(left.shape[2]) not in SUPPORTED_HIDDEN_CHANNELS + or right.shape != left.shape + or tuple(to_grid.shape) != (GRID_SIZE, PACKED_COEFF_DIM) + or tuple(from_grid.shape) != (PACKED_COEFF_DIM, GRID_SIZE) + or any(not tensor.is_contiguous() for tensor in tensors) + or torch.cuda.get_device_capability(left.device)[0] < 8 + ): + raise ValueError( + "tiled output-grid product requires contiguous CUDA FP32 " + "left/right=(N,48,C) with C in {96,192}, to_grid=(152,48), and " + "from_grid=(48,152) tensors on compute capability 8.0+" + ) + if out is not None and ( + out.shape != left.shape + or out.device != left.device + or out.dtype != left.dtype + or not out.is_contiguous() + ): + raise ValueError("tiled output-grid output must match left") + + +def run_tiled_output_grid_product( + left, + right, + to_grid, + from_grid, + *, + use_sm80_c96_n48: bool = False, + use_sm90_c96_asymmetric_panels: bool = False, +): + """Run the fused strict-FP32 tiled output-grid forward.""" + import torch + + _validate_tensors(left, right, to_grid, from_grid) + out = torch.empty_like(left) + return run_tiled_output_grid_product_out( + left, + right, + to_grid, + from_grid, + out, + use_sm80_c96_n48=use_sm80_c96_n48, + use_sm90_c96_asymmetric_panels=use_sm90_c96_asymmetric_panels, + ) + + +def run_tiled_output_grid_product_out( + left, + right, + to_grid, + from_grid, + out, + *, + use_sm80_c96_n48: bool = False, + use_sm90_c96_asymmetric_panels: bool = False, +): + """Run the fused forward into a caller-provided output tensor.""" + _validate_tensors(left, right, to_grid, from_grid, out) + compile_key = _compile_key(left) + if use_sm80_c96_n48 and use_sm90_c96_asymmetric_panels: + raise ValueError("output-grid panel specializations are mutually exclusive") + if use_sm80_c96_n48 and ( + compile_key[1] not in runtime_policy.SM80_PROFILE_CAPABILITIES + or compile_key[2] != 96 + ): + raise ValueError("output-grid forward N=48 requires SM80-family and C=96") + if use_sm90_c96_asymmetric_panels and ( + compile_key[1] != (9, 0) or compile_key[2] != 96 + ): + raise ValueError("output-grid asymmetric forward requires sm90 and C=96") + if use_sm90_c96_asymmetric_panels: + _compiled_tiled_forward( + *compile_key, + TILE_N, + 0, + 1, + )( + left, + right, + to_grid, + from_grid, + out, + ) + _compiled_tiled_forward( + *compile_key, + C96_TAIL_TILE_N, + C96_TAIL_CHANNEL_TILE, + 1, + )( + left, + right, + to_grid, + from_grid, + out, + ) + return out + tile_n = SM80_C96_TILE_N if use_sm80_c96_n48 else TILE_N + _compiled_tiled_forward(*compile_key, tile_n, 0, None)( + left, + right, + to_grid, + from_grid, + out, + ) + return out + + +def run_tiled_output_grid_product_backward( + grad_out, + left, + right, + to_grid, + from_grid, + *, + use_sm80_c96_n48_panel: bool = False, + use_sm90_c96_asymmetric_panels: bool = False, +): + """Run the complete fused first backward for left and right inputs.""" + import torch + + _validate_tensors(left, right, to_grid, from_grid) + if ( + grad_out.shape != left.shape + or grad_out.device != left.device + or grad_out.dtype != left.dtype + or not grad_out.is_contiguous() + ): + raise ValueError("grad_out must be contiguous and match left") + compile_key = _compile_key(left) + if use_sm80_c96_n48_panel and use_sm90_c96_asymmetric_panels: + raise ValueError("output-grid panel specializations are mutually exclusive") + if use_sm80_c96_n48_panel and ( + compile_key[1] not in runtime_policy.SM80_PROFILE_CAPABILITIES + or compile_key[2] != 96 + ): + raise ValueError( + "output-grid N=48 panel adjoint requires SM80-family, C=96, and K=8" + ) + if use_sm90_c96_asymmetric_panels and ( + compile_key[1] != (9, 0) or compile_key[2] != 96 + ): + raise ValueError("output-grid asymmetric backward requires sm90, C=96, and K=8") + grad_left = torch.empty_like(left) + grad_right = torch.empty_like(right) + if use_sm90_c96_asymmetric_panels: + _compiled_tiled_backward( + *compile_key, + TILE_N, + 0, + 1, + )( + grad_out, + left, + right, + to_grid, + from_grid, + grad_left, + grad_right, + ) + _compiled_tiled_backward( + *compile_key, + C96_TAIL_TILE_N, + C96_TAIL_CHANNEL_TILE, + 1, + )( + grad_out, + left, + right, + to_grid, + from_grid, + grad_left, + grad_right, + ) + return grad_left, grad_right + tile_n = SM80_C96_TILE_N if use_sm80_c96_n48_panel else TILE_N + _compiled_tiled_backward( + *compile_key, + tile_n, + 0, + None, + )( + grad_out, + left, + right, + to_grid, + from_grid, + grad_left, + grad_right, + ) + return grad_left, grad_right + + +__all__ = [ + "TiledOutputGridProductBackward", + "TiledOutputGridProductForward", + "run_tiled_output_grid_product", + "run_tiled_output_grid_product_backward", + "run_tiled_output_grid_product_out", +] diff --git a/deepmd/kernels/cute/neo/output_grid_product.py b/deepmd/kernels/cute/neo/output_grid_product.py new file mode 100644 index 0000000000..4d106f0a4e --- /dev/null +++ b/deepmd/kernels/cute/neo/output_grid_product.py @@ -0,0 +1,329 @@ +# 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 supported Neo grid MLPs.""" + +from __future__ import ( + annotations, +) + +from typing import ( + Any, +) + +import torch + +from . import ( + runtime_policy, +) +from .runtime_policy import ( + PORTABLE_TILED_BACKEND, + SUPPORTED_HIDDEN_CHANNELS, + select_output_grid_backend, +) + +COEFF_DIM = 16 +N_FRAMES = 3 +PACKED_COEFF_DIM = COEFF_DIM * N_FRAMES +GRID_SIZE = 152 + + +def _exact_hidden_channels( + left: torch.Tensor, + n_frames: int, +) -> int | None: + if left.ndim != 4 or left.shape[0] <= 0 or int(n_frames) != N_FRAMES: + return None + for hidden_channels in SUPPORTED_HIDDEN_CHANNELS: + if tuple(left.shape[1:]) == ( + COEFF_DIM, + 1, + N_FRAMES * hidden_channels, + ): + return hidden_channels + return None + + +def _has_exact_contract( + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, + n_frames: int, +) -> bool: + hidden_channels = _exact_hidden_channels(left, n_frames) + if ( + hidden_channels is None + or not left.is_cuda + or left.dtype != torch.float32 + or right.device != left.device + or right.dtype != left.dtype + or to_grid.device != left.device + or from_grid.device != left.device + or to_grid.dtype != left.dtype + or from_grid.dtype != left.dtype + or right.shape != left.shape + or tuple(to_grid.shape) != (GRID_SIZE, PACKED_COEFF_DIM) + or tuple(from_grid.shape) != (PACKED_COEFF_DIM, GRID_SIZE) + or not left.is_contiguous() + or not right.is_contiguous() + or not to_grid.is_contiguous() + or not from_grid.is_contiguous() + or to_grid.requires_grad + or from_grid.requires_grad + or torch.is_autocast_enabled("cuda") + or not runtime_policy.uses_strict_fp32_matmul() + ): + return False + max_intermediate_values = left.shape[0] * GRID_SIZE * hidden_channels + if max_intermediate_values > runtime_policy.INT32_MAX: + return False + compute_capability = tuple(torch.cuda.get_device_capability(left.device)) + return ( + select_output_grid_backend(compute_capability, hidden_channels) + == PORTABLE_TILED_BACKEND + ) + + +def _validate_exact_contract( + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, + n_frames: int, +) -> int: + hidden_channels = _exact_hidden_channels(left, n_frames) + if hidden_channels is None or not _has_exact_contract( + left, + right, + to_grid, + from_grid, + n_frames, + ): + raise ValueError( + "the fused output GridMLP kernel requires contiguous CUDA FP32 " + "tensors with Neo's (D=16, F=1, frames=3, G=152, " + "C in {96, 192}) contract" + ) + return hidden_channels + + +def _output_grid_product_impl( + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, + n_frames: int, +) -> torch.Tensor: + hidden_channels = _validate_exact_contract( + left, + right, + to_grid, + from_grid, + n_frames, + ) + from .output_grid_kernels.cute_tiled_grid_product import ( + run_tiled_output_grid_product, + ) + + nodes = left.shape[0] + left_flat = left.detach().view(nodes, PACKED_COEFF_DIM, hidden_channels) + right_flat = right.detach().view(nodes, PACKED_COEFF_DIM, hidden_channels) + compute_capability = tuple(torch.cuda.get_device_capability(left.device)) + use_sm80_c96_n48 = ( + hidden_channels == 96 + and compute_capability in runtime_policy.SM80_PROFILE_CAPABILITIES + and runtime_policy.is_output_grid_fwd_sm80_c96_n48_enabled(compute_capability) + ) + use_sm90_c96_asymmetric_panels = ( + hidden_channels == 96 + and compute_capability == (9, 0) + and runtime_policy.is_output_grid_sm90_c96_asymmetric_panels_enabled( + compute_capability + ) + ) + out = run_tiled_output_grid_product( + left_flat, + right_flat, + to_grid.detach(), + from_grid.detach(), + use_sm80_c96_n48=use_sm80_c96_n48, + use_sm90_c96_asymmetric_panels=use_sm90_c96_asymmetric_panels, + ) + return out.view_as(left) + + +def _output_grid_product_bwd_impl( + grad_out: torch.Tensor, + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, + n_frames: int, +) -> tuple[torch.Tensor, torch.Tensor]: + hidden_channels = _validate_exact_contract( + left, + right, + to_grid, + from_grid, + n_frames, + ) + if ( + grad_out.shape != left.shape + or grad_out.dtype != left.dtype + or grad_out.device != left.device + ): + raise ValueError("grad_out must match the fused output GridMLP output") + from .output_grid_kernels.cute_tiled_grid_product import ( + run_tiled_output_grid_product_backward, + ) + + nodes = left.shape[0] + compute_capability = tuple(torch.cuda.get_device_capability(left.device)) + use_sm80_c96_n48_panel = ( + hidden_channels == 96 + and compute_capability in runtime_policy.SM80_PROFILE_CAPABILITIES + and runtime_policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled( + compute_capability + ) + ) + use_sm90_c96_asymmetric_panels = ( + hidden_channels == 96 + and compute_capability == (9, 0) + and runtime_policy.is_output_grid_sm90_c96_asymmetric_panels_enabled( + compute_capability + ) + ) + grad_left, grad_right = run_tiled_output_grid_product_backward( + grad_out.detach() + .contiguous() + .view( + nodes, + PACKED_COEFF_DIM, + hidden_channels, + ), + left.detach().view(nodes, PACKED_COEFF_DIM, hidden_channels), + right.detach().view(nodes, PACKED_COEFF_DIM, hidden_channels), + to_grid.detach(), + from_grid.detach(), + use_sm80_c96_n48_panel=use_sm80_c96_n48_panel, + use_sm90_c96_asymmetric_panels=use_sm90_c96_asymmetric_panels, + ) + return grad_left.view_as(left), grad_right.view_as(right) + + +_output_grid_product_op = torch.library.custom_op( + "sezm_cute::output_grid_product", + mutates_args=(), +)(_output_grid_product_impl) +_output_grid_product_bwd_op = torch.library.custom_op( + "sezm_cute::output_grid_product_bwd", + mutates_args=(), +)(_output_grid_product_bwd_impl) + + +@_output_grid_product_op.register_fake +def _output_grid_product_fake( + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, + n_frames: int, +) -> torch.Tensor: + del right, to_grid, from_grid, n_frames + return torch.empty(left.shape, dtype=left.dtype, device=left.device) + + +@_output_grid_product_bwd_op.register_fake +def _output_grid_product_bwd_fake( + grad_out: torch.Tensor, + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, + n_frames: int, +) -> tuple[torch.Tensor, torch.Tensor]: + del grad_out, to_grid, from_grid, n_frames + return ( + torch.empty(left.shape, dtype=left.dtype, device=left.device), + torch.empty(right.shape, dtype=right.dtype, device=right.device), + ) + + +def _setup_context( + ctx: Any, + inputs: tuple, + output: torch.Tensor, +) -> None: + del output + left, right, to_grid, from_grid, n_frames = inputs + ctx.save_for_backward(left, right, to_grid, from_grid) + ctx.n_frames = int(n_frames) + + +def _backward(ctx: Any, grad_out: torch.Tensor) -> tuple: + left, right, to_grid, from_grid = ctx.saved_tensors + grad_left, grad_right = _output_grid_product_bwd_op( + grad_out, + left, + right, + to_grid, + from_grid, + ctx.n_frames, + ) + return grad_left, grad_right, None, None, None + + +_output_grid_product_op.register_autograd( + _backward, + setup_context=_setup_context, +) + + +def output_grid_product_cute( + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, + *, + n_frames: int, +) -> torch.Tensor: + """Run one supported exact-shape fused grid contraction.""" + return _output_grid_product_op( + left, + right, + to_grid, + from_grid, + int(n_frames), + ) + + +def maybe_run_cute_output_grid_product( + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, + *, + n_frames: int, +) -> torch.Tensor | None: + """Return ``None`` unless the master gate and exact Neo contract match.""" + if not runtime_policy.is_cute_infer_enabled() or not _has_exact_contract( + left, + right, + to_grid, + from_grid, + n_frames, + ): + return None + return output_grid_product_cute( + left, + right, + to_grid, + from_grid, + n_frames=n_frames, + ) + + +__all__ = [ + "maybe_run_cute_output_grid_product", + "output_grid_product_cute", +] diff --git a/deepmd/kernels/cute/neo/readout_l0.py b/deepmd/kernels/cute/neo/readout_l0.py new file mode 100644 index 0000000000..db18740aca --- /dev/null +++ b/deepmd/kernels/cute/neo/readout_l0.py @@ -0,0 +1,774 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Strict-FP32 degree-zero readout for the exact Neo output GridMLP.""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, + Any, +) + +import torch + +from . import ( + runtime_policy, +) +from .runtime_policy import ( + PORTABLE_TILED_BACKEND, + select_output_grid_backend, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + +COEFF_DIM = 16 +N_FRAMES = 3 +PACKED_COEFF_DIM = 48 +GRID_SIZE = 152 +HIDDEN_CHANNELS = 192 +PACKED_WIDTH = N_FRAMES * HIDDEN_CHANNELS +_READOUT_INPUT_FOLD_CACHE = "_neo_sm80_readout_input_fold_cache" +_READOUT_INPUT_FOLD_HOOK = "_neo_sm80_readout_input_fold_hook" +_READOUT_INPUT_FOLD_LEFT = "_neo_sm80_readout_input_fold_left" +_READOUT_INPUT_FOLD_RIGHT = "_neo_sm80_readout_input_fold_right" +_READOUT_INPUT_FOLD_SCALAR = "_neo_sm80_readout_input_fold_scalar" +_READOUT_INPUT_FOLD_BUFFERS = ( + _READOUT_INPUT_FOLD_LEFT, + _READOUT_INPUT_FOLD_RIGHT, + _READOUT_INPUT_FOLD_SCALAR, +) +_READOUT_INPUT_FOLD_PREPARE_ERROR = ( + "the SM80 readout input fold is stale or missing; call " + "prepare_sm80_readout_input_fold(output_ffn) after loading, replacing, " + "mutating, or moving model state and before torch.compile" +) + + +def _has_exact_product_shape(left: torch.Tensor) -> bool: + return ( + left.ndim == 4 + and left.shape[0] > 0 + and tuple(left.shape[1:]) == (COEFF_DIM, 1, PACKED_WIDTH) + ) + + +def _has_exact_product_contract( + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, +) -> bool: + tensors = (left, right, to_grid, from_grid) + if ( + not _has_exact_product_shape(left) + or right.shape != left.shape + or tuple(to_grid.shape) != (GRID_SIZE, PACKED_COEFF_DIM) + or tuple(from_grid.shape) != (PACKED_COEFF_DIM, GRID_SIZE) + or any(not tensor.is_cuda for tensor in tensors) + or any(tensor.dtype != torch.float32 for tensor in tensors) + or any(tensor.device != left.device for tensor in tensors) + or any(not tensor.is_contiguous() for tensor in tensors) + or to_grid.requires_grad + or from_grid.requires_grad + ): + return False + compute_capability = tuple(torch.cuda.get_device_capability(left.device)) + return ( + select_output_grid_backend(compute_capability, HIDDEN_CHANNELS) + == PORTABLE_TILED_BACKEND + ) + + +def _validate_product_contract( + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, +) -> None: + if not _has_exact_product_contract(left, right, to_grid, from_grid): + raise ValueError( + "the readout l=0 kernel requires contiguous CUDA FP32 tensors " + "with Neo's left/right=(N,16,1,576), to_grid=(152,48), and " + "from_grid=(48,152) contract" + ) + + +def build_readout_l0_gram( + to_grid: torch.Tensor, + from_grid: torch.Tensor, +) -> torch.Tensor: + """Collapse the frozen row-zero projector into a dense FP32 Gram matrix.""" + if ( + tuple(to_grid.shape) != (GRID_SIZE, PACKED_COEFF_DIM) + or tuple(from_grid.shape) != (PACKED_COEFF_DIM, GRID_SIZE) + or to_grid.dtype != torch.float32 + or from_grid.dtype != torch.float32 + or to_grid.device != from_grid.device + or not to_grid.is_contiguous() + or not from_grid.is_contiguous() + or to_grid.requires_grad + or from_grid.requires_grad + ): + raise ValueError( + "readout l=0 Gram construction requires frozen contiguous FP32 " + "to_grid=(152,48) and from_grid=(48,152) tensors" + ) + with torch.no_grad(): + return torch.matmul( + to_grid.T, + from_grid[0, :, None] * to_grid, + ).contiguous() + + +def _has_exact_neo_readout_structure(output_ffn: Any) -> bool: + from deepmd.pt.model.descriptor.sezm_nn.ffn import ( + EquivariantFFN, + ) + from deepmd.pt.model.descriptor.sezm_nn.grid_net import ( + GridMLP, + SO3GridNet, + ) + + if type(output_ffn) is not EquivariantFFN: + return False + grid_net = output_ffn.act + if type(grid_net) is not SO3GridNet or type(grid_net.grid_op) is not GridMLP: + return False + grid_op = grid_net.grid_op + projector = grid_net.projector + return ( + output_ffn.lmax == 3 + and output_ffn.channels == 32 + and output_ffn.hidden_channels == 96 + and output_ffn.kmax == 1 + and output_ffn.grid_n_frames == N_FRAMES + and output_ffn.use_grid_net + and output_ffn.use_grid_mlp + and not output_ffn.use_grid_branch + and output_ffn.ffn_so3_grid + and not output_ffn.s2_activation + and not output_ffn.mlp_bias + and grid_net.lmax == 3 + and grid_net.channels == 96 + and grid_net.n_focus == 1 + and grid_net.n_frames == N_FRAMES + and grid_net.mode == "self" + and grid_net.op_type == "mlp" + and grid_net.layout == "ndfc" + and grid_net.frame_zero_index == 0 + and grid_net.frames == [0, -1, 1] + and grid_net.frame_expand is None + and grid_net.frame_contract is None + and grid_net.residual_scale is None + and grid_op.mode == "self" + and grid_op.channels == 96 + and grid_op.hidden_channels == HIDDEN_CHANNELS + and grid_op.n_frames == N_FRAMES + and tuple(output_ffn.so3_linear_1.weight.shape) == (4, 32, PACKED_WIDTH) + and tuple(grid_op.left_proj.weight.shape) == (HIDDEN_CHANNELS, HIDDEN_CHANNELS) + and tuple(grid_op.right_proj.weight.shape) == (HIDDEN_CHANNELS, HIDDEN_CHANNELS) + and tuple(grid_op.out_proj.weight.shape) == (HIDDEN_CHANNELS, 96) + and tuple(grid_net.scalar_gate.weight.shape) == (HIDDEN_CHANNELS, 96) + and tuple(output_ffn.so3_linear_2.weight.shape) == (4, 288, 32) + and output_ffn.so3_linear_1.bias is None + and grid_op.left_proj.bias is None + and grid_op.right_proj.bias is None + and grid_op.out_proj.bias is None + and grid_net.scalar_gate.bias is None + and output_ffn.so3_linear_2.bias is None + and tuple(projector.to_grid_mat.shape) == (GRID_SIZE, PACKED_COEFF_DIM) + and tuple(projector.from_grid_mat.shape) == (PACKED_COEFF_DIM, GRID_SIZE) + ) + + +def _state_uses_strict_fp32(output_ffn: Any, device: torch.device) -> bool: + for tensor in (*output_ffn.parameters(), *output_ffn.buffers()): + if tensor.is_floating_point() and ( + tensor.dtype != torch.float32 + or tensor.device != device + or not tensor.is_contiguous() + ): + return False + return True + + +def _inference_mode_is_frozen(output_ffn: Any) -> bool: + return not output_ffn.training and not any( + parameter.requires_grad for parameter in output_ffn.parameters() + ) + + +@torch.compiler.assume_constant_result +def _uses_strict_fp32_matmul() -> bool: + """Preserve the tested private helper while sharing the runtime policy.""" + return runtime_policy.uses_strict_fp32_matmul() + + +def _has_exact_neo_readout_contract( + output_ffn: Any, + ffn_in: torch.Tensor, +) -> bool: + if ( + not _has_exact_neo_readout_structure(output_ffn) + or not _inference_mode_is_frozen(output_ffn) + or ffn_in.ndim != 4 + or ffn_in.shape[0] <= 0 + or tuple(ffn_in.shape[1:]) != (COEFF_DIM, 1, 32) + or not ffn_in.is_cuda + or ffn_in.dtype != torch.float32 + or not ffn_in.is_contiguous() + or not _state_uses_strict_fp32(output_ffn, ffn_in.device) + or torch.is_autocast_enabled("cuda") + or not _uses_strict_fp32_matmul() + ): + return False + compute_capability = tuple(torch.cuda.get_device_capability(ffn_in.device)) + return ( + select_output_grid_backend(compute_capability, HIDDEN_CHANNELS) + == PORTABLE_TILED_BACKEND + ) + + +def _can_use_sm80_readout_input_fold( + output_ffn: Any, + ffn_in: torch.Tensor, +) -> bool: + """Fail closed unless the exact frozen strict-FP32 SM80 path is active.""" + if not _has_exact_neo_readout_contract(output_ffn, ffn_in): + return False + compute_capability = tuple(torch.cuda.get_device_capability(ffn_in.device)) + return runtime_policy.is_readout_input_fold_enabled(compute_capability) + + +def _readout_input_fold_sources(output_ffn: Any) -> tuple[torch.Tensor, ...]: + grid_net = output_ffn.act + grid_op = grid_net.grid_op + return ( + output_ffn.so3_linear_1.weight, + grid_op.left_proj.weight, + grid_op.right_proj.weight, + grid_net.scalar_gate.weight, + ) + + +def _readout_input_fold_cache_key( + sources: tuple[torch.Tensor, ...], +) -> tuple[tuple[Any, ...], ...]: + return tuple( + ( + tensor.data_ptr(), + tensor._version, + tensor.dtype, + tensor.device, + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.storage_offset(), + ) + for tensor in sources + ) + + +def _readout_input_fold_cache_matches( + output_ffn: Any, + sources: tuple[torch.Tensor, ...], + cache_key: tuple[tuple[Any, ...], ...], +) -> bool: + cache = getattr(output_ffn, _READOUT_INPUT_FOLD_CACHE, None) + return ( + cache is not None + and len(cache) == 2 + and len(cache[0]) == len(sources) + and all(cached is current for cached, current in zip(cache[0], sources)) + and cache[1] == cache_key + and all( + isinstance(getattr(output_ffn, name, None), torch.Tensor) + for name in _READOUT_INPUT_FOLD_BUFFERS + ) + ) + + +def _invalidate_sm80_readout_input_fold(output_ffn: Any) -> None: + setattr(output_ffn, _READOUT_INPUT_FOLD_CACHE, None) + for name in _READOUT_INPUT_FOLD_BUFFERS: + if name in output_ffn._buffers: + setattr(output_ffn, name, None) + + +def invalidate_neo_readout_input_fold(output_ffn: Any) -> None: + """Invalidate frozen readout weights after parameter topology changes.""" + _invalidate_sm80_readout_input_fold(output_ffn) + + +def _invalidate_sm80_readout_input_fold_after_load( + output_ffn: Any, + incompatible_keys: Any, +) -> None: + del incompatible_keys + _invalidate_sm80_readout_input_fold(output_ffn) + + +def _ensure_sm80_readout_input_fold_load_hook(output_ffn: Any) -> None: + if getattr(output_ffn, _READOUT_INPUT_FOLD_HOOK, False): + return + output_ffn.register_load_state_dict_post_hook( + _invalidate_sm80_readout_input_fold_after_load + ) + setattr(output_ffn, _READOUT_INPUT_FOLD_HOOK, True) + + +def _set_nonpersistent_buffer( + module: Any, + name: str, + tensor: torch.Tensor, +) -> None: + if name in module._buffers: + setattr(module, name, tensor) + else: + module.register_buffer(name, tensor, persistent=False) + + +def _synchronize_sm80_readout_input_fold_build( + folded_weights: tuple[torch.Tensor, torch.Tensor, torch.Tensor], +) -> None: + """Make a newly built CUDA cache safe for every later consumer stream.""" + device = folded_weights[0].device + if device.type != "cuda": + return + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "prepare the SM80 readout input fold before CUDA graph capture" + ) + ready = torch.cuda.Event() + ready.record(torch.cuda.current_stream(device)) + ready.synchronize() + + +def _build_sm80_readout_input_fold( + output_ffn: Any, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compose the frozen maps without crossing either readout nonlinearity.""" + input_weight, left_weight, right_weight, scalar_gate_weight = ( + _readout_input_fold_sources(output_ffn) + ) + grid_net = output_ffn.act + with torch.no_grad(): + split_weight = input_weight.reshape( + output_ffn.lmax + 1, + output_ffn.channels, + 2, + N_FRAMES, + 96, + ) + left_input = split_weight[:, :, 0] + right_input = split_weight[:, :, 1] + per_frame_input = torch.cat((left_input, right_input), dim=-1) + left_fold = torch.matmul(per_frame_input, left_weight) + right_fold = torch.matmul(per_frame_input, right_weight) + projected_left_weight = left_fold.reshape( + output_ffn.lmax + 1, + output_ffn.channels, + -1, + ).contiguous() + projected_right_weight = right_fold.reshape( + output_ffn.lmax + 1, + output_ffn.channels, + -1, + ).contiguous() + + frame_zero = grid_net.frame_zero_index + scalar_pair_weight = torch.cat( + ( + left_input[0, :, frame_zero], + right_input[0, :, frame_zero], + ), + dim=-1, + ) + scalar_gate_fold = torch.matmul(scalar_pair_weight, scalar_gate_weight) + scalar_aux_weight = torch.cat( + (scalar_pair_weight, scalar_gate_fold), + dim=-1, + ).contiguous() + return ( + projected_left_weight.detach().clone(), + projected_right_weight.detach().clone(), + scalar_aux_weight.detach().clone(), + ) + + +@torch.compiler.disable +def prepare_sm80_readout_input_fold( + output_ffn: Any, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Prepare immutable folded weights before compiling the frozen readout.""" + if torch.compiler.is_compiling(): + raise RuntimeError(_READOUT_INPUT_FOLD_PREPARE_ERROR) + if not _has_exact_neo_readout_structure( + output_ffn + ) or not _inference_mode_is_frozen(output_ffn): + raise ValueError("readout input folding requires the exact frozen Neo FFN") + + sources = _readout_input_fold_sources(output_ffn) + device = sources[0].device + if any( + tensor.dtype != torch.float32 + or tensor.device != device + or not tensor.is_contiguous() + for tensor in sources + ): + raise ValueError( + "readout input folding requires contiguous FP32 source weights " + "on one device" + ) + cache_key = _readout_input_fold_cache_key(sources) + if _readout_input_fold_cache_matches(output_ffn, sources, cache_key): + return ( + getattr(output_ffn, _READOUT_INPUT_FOLD_LEFT), + getattr(output_ffn, _READOUT_INPUT_FOLD_RIGHT), + getattr(output_ffn, _READOUT_INPUT_FOLD_SCALAR), + ) + + _invalidate_sm80_readout_input_fold(output_ffn) + folded_weights = _build_sm80_readout_input_fold(output_ffn) + _synchronize_sm80_readout_input_fold_build(folded_weights) + _ensure_sm80_readout_input_fold_load_hook(output_ffn) + for name, tensor in zip(_READOUT_INPUT_FOLD_BUFFERS, folded_weights): + _set_nonpersistent_buffer(output_ffn, name, tensor) + setattr( + output_ffn, + _READOUT_INPUT_FOLD_CACHE, + (sources, cache_key), + ) + return folded_weights + + +@torch.compiler.disable +def maybe_prepare_sm80_readout_input_fold( + output_ffn: Any, + compute_capability: tuple[int, int] | None = None, +) -> bool: + """Prepare only when a supported frozen Neo readout contract matches.""" + if not _has_exact_neo_readout_structure( + output_ffn + ) or not _inference_mode_is_frozen(output_ffn): + return False + sources = _readout_input_fold_sources(output_ffn) + device = sources[0].device + if device.type != "cuda": + return False + if compute_capability is None: + compute_capability = tuple(torch.cuda.get_device_capability(device)) + if not runtime_policy.is_readout_input_fold_enabled(compute_capability): + return False + if any( + tensor.dtype != torch.float32 + or tensor.device != device + or not tensor.is_contiguous() + for tensor in sources + ): + return False + prepare_sm80_readout_input_fold(output_ffn) + return True + + +def _get_prepared_sm80_readout_input_fold( + output_ffn: Any, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + cache = getattr(output_ffn, _READOUT_INPUT_FOLD_CACHE, None) + folded_weights = ( + getattr(output_ffn, _READOUT_INPUT_FOLD_LEFT, None), + getattr(output_ffn, _READOUT_INPUT_FOLD_RIGHT, None), + getattr(output_ffn, _READOUT_INPUT_FOLD_SCALAR, None), + ) + if ( + cache is None + or len(cache) != 2 + or len(cache[0]) != 4 + or any(not isinstance(weight, torch.Tensor) for weight in folded_weights) + ): + raise RuntimeError(_READOUT_INPUT_FOLD_PREPARE_ERROR) + + sources = _readout_input_fold_sources(output_ffn) + cached_sources, cache_key = cache + if ( + any(cached is not current for cached, current in zip(cached_sources, sources)) + or any( + source.dtype != cached[2] + or source.device != cached[3] + or tuple(source.shape) != cached[4] + for source, cached in zip(sources, cache_key) + ) + or any( + weight.dtype != torch.float32 + or weight.device != sources[0].device + or not weight.is_contiguous() + for weight in folded_weights + ) + ): + raise RuntimeError(_READOUT_INPUT_FOLD_PREPARE_ERROR) + + # Do not emit source ``_version`` counters into the graph. AOTAutograd + # represents them as unbacked symbolic integers and cannot lower the + # resulting assertion. Eager preparation validates versions before trace; + # SeZM's load-state hook invalidates both local and shared compiled graphs. + return folded_weights + + +def _get_sm80_readout_input_fold( + output_ffn: Any, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return prepared degree weights for operands and scalar auxiliaries.""" + if torch.compiler.is_compiling(): + return _get_prepared_sm80_readout_input_fold(output_ffn) + return prepare_sm80_readout_input_fold(output_ffn) + + +def _maybe_prepare_sm80_readout_input_fold( + output_ffn: Any, + ffn_in: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None: + """Project directly from the 32-channel input to both product operands.""" + if not _can_use_sm80_readout_input_fold(output_ffn, ffn_in): + return None + + left_weight, right_weight, scalar_weight = _get_sm80_readout_input_fold(output_ffn) + expanded_left_weight = left_weight.index_select( + 0, + output_ffn.so3_linear_1.expand_index, + ) + expanded_right_weight = right_weight.index_select( + 0, + output_ffn.so3_linear_1.expand_index, + ) + left = torch.einsum("ndfi,dio->ndfo", ffn_in, expanded_left_weight).contiguous() + right = torch.einsum("ndfi,dio->ndfo", ffn_in, expanded_right_weight).contiguous() + scalar_aux = torch.einsum("nfi,io->nfo", ffn_in[:, 0], scalar_weight) + scalar_pair, scalar_gate_logits = torch.split(scalar_aux, (192, 96), dim=-1) + return left, right, scalar_pair, scalar_gate_logits + + +def _run_neo_readout_l0( + output_ffn: Any, + ffn_in: torch.Tensor, + grid_product: Callable[ + [torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], + torch.Tensor, + ], +) -> torch.Tensor: + """Complete the exact Neo readout around one row-zero grid product.""" + if not _has_exact_neo_readout_structure(output_ffn): + raise ValueError("readout l=0 completion requires the exact Neo output FFN") + from deepmd.pt.model.descriptor.sezm_nn.grid_net import ( + _project_frames, + ) + + grid_net = output_ffn.act + grid_op = grid_net.grid_op + + prepared = _maybe_prepare_sm80_readout_input_fold(output_ffn, ffn_in) + if prepared is None: + projected = output_ffn.so3_linear_1(ffn_in) + left, right, scalar_pair = grid_net._prepare_self_pair(projected) + shape = (*left.shape[:-1], N_FRAMES, -1) + fused = torch.cat( + [left.reshape(shape), right.reshape(shape)], + dim=-1, + ).reshape(*left.shape[:-1], -1) + left = _project_frames(fused, grid_op.left_proj, N_FRAMES) + right = _project_frames(fused, grid_op.right_proj, N_FRAMES) + scalar_gate_logits = None + else: + left, right, scalar_pair, scalar_gate_logits = prepared + + q0 = grid_product( + left, + right, + grid_net.projector.to_grid_mat, + grid_net.projector.from_grid_mat, + ) + if tuple(q0.shape) != (ffn_in.shape[0], HIDDEN_CHANNELS): + raise ValueError("readout l=0 grid product must return shape (N,192)") + + q0 = torch.matmul(q0, grid_op.out_proj.weight) + scalar_out = grid_net.scalar_act(scalar_pair)[:, 0, :] + if scalar_gate_logits is None: + scalar_gate_logits = grid_net.scalar_gate(scalar_pair) + scalar_gate = torch.sigmoid(scalar_gate_logits)[:, 0, :] + scalar_coeff = q0 * scalar_gate + scalar_out + output_weight = output_ffn.so3_linear_2.weight[0, :96, :] + return ffn_in[:, 0, 0, :] + torch.matmul(scalar_coeff, output_weight) + + +def maybe_run_neo_readout_l0( + output_ffn: Any, + ffn_in: torch.Tensor, +) -> torch.Tensor | None: + """Return the optimized final `[N,32]` readout or ``None`` for fallback.""" + if not runtime_policy.is_cute_infer_enabled(): + return None + if not _inference_mode_is_frozen(output_ffn): + return None + if not _has_exact_neo_readout_contract(output_ffn, ffn_in): + return None + return _run_neo_readout_l0(output_ffn, ffn_in, readout_l0_product_cute) + + +def run_neo_output_readout( + output_ffn: Any, + ffn_in: torch.Tensor, + *, + parameters_frozen: bool = True, +) -> torch.Tensor: + """Return the residual-inclusive `[N,32]` output with generic fallback.""" + if parameters_frozen: + candidate = maybe_run_neo_readout_l0(output_ffn, ffn_in) + if candidate is not None: + return candidate + return (ffn_in + output_ffn(ffn_in))[:, 0:1, :, :].reshape( + ffn_in.shape[0], output_ffn.channels + ) + + +def _readout_l0_impl( + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, +) -> torch.Tensor: + _validate_product_contract(left, right, to_grid, from_grid) + from .output_grid_kernels.cute_readout_l0 import ( + run_readout_l0, + ) + + nodes = left.shape[0] + q0 = run_readout_l0( + left.detach().view(nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS), + right.detach().view(nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS), + to_grid.detach(), + from_grid.detach(), + ) + return q0 + + +def _readout_l0_bwd_impl( + dq0: torch.Tensor, + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + _validate_product_contract(left, right, to_grid, from_grid) + if ( + tuple(dq0.shape) != (left.shape[0], HIDDEN_CHANNELS) + or dq0.dtype != left.dtype + or dq0.device != left.device + or not dq0.is_contiguous() + ): + raise ValueError("dq0 must be contiguous and have shape (N,192)") + from .output_grid_kernels.cute_readout_l0 import ( + run_readout_l0_backward, + ) + + nodes = left.shape[0] + grad_left, grad_right = run_readout_l0_backward( + dq0.detach(), + left.detach().view(nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS), + right.detach().view(nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS), + to_grid.detach(), + from_grid.detach(), + ) + return grad_left.view_as(left), grad_right.view_as(right) + + +_readout_l0_op = torch.library.custom_op( + "sezm_cute::readout_l0", + mutates_args=(), +)(_readout_l0_impl) +_readout_l0_bwd_op = torch.library.custom_op( + "sezm_cute::readout_l0_bwd", + mutates_args=(), +)(_readout_l0_bwd_impl) + + +@_readout_l0_op.register_fake +def _readout_l0_fake( + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, +) -> torch.Tensor: + del right, to_grid, from_grid + return torch.empty( + (left.shape[0], HIDDEN_CHANNELS), + dtype=left.dtype, + device=left.device, + ) + + +@_readout_l0_bwd_op.register_fake +def _readout_l0_bwd_fake( + dq0: torch.Tensor, + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + del dq0, to_grid, from_grid + return ( + torch.empty(left.shape, dtype=left.dtype, device=left.device), + torch.empty(right.shape, dtype=right.dtype, device=right.device), + ) + + +def _setup_context( + ctx: Any, + inputs: tuple, + output: torch.Tensor, +) -> None: + del output + left, right, to_grid, from_grid = inputs + ctx.save_for_backward(left, right, to_grid, from_grid) + + +def _backward(ctx: Any, dq0: torch.Tensor) -> tuple: + left, right, to_grid, from_grid = ctx.saved_tensors + grad_left, grad_right = _readout_l0_bwd_op( + dq0.contiguous(), + left, + right, + to_grid, + from_grid, + ) + return grad_left, grad_right, None, None + + +_readout_l0_op.register_autograd( + _backward, + setup_context=_setup_context, +) + + +def readout_l0_product_cute( + left: torch.Tensor, + right: torch.Tensor, + to_grid: torch.Tensor, + from_grid: torch.Tensor, +) -> torch.Tensor: + """Run the exact-shape C=192 degree-zero grid contraction.""" + return _readout_l0_op(left, right, to_grid, from_grid) + + +__all__ = [ + "build_readout_l0_gram", + "invalidate_neo_readout_input_fold", + "maybe_prepare_sm80_readout_input_fold", + "maybe_run_neo_readout_l0", + "prepare_sm80_readout_input_fold", + "readout_l0_product_cute", + "run_neo_output_readout", +] diff --git a/deepmd/kernels/cute/neo/runtime_policy.py b/deepmd/kernels/cute/neo/runtime_policy.py new file mode 100644 index 0000000000..1598680c87 --- /dev/null +++ b/deepmd/kernels/cute/neo/runtime_policy.py @@ -0,0 +1,290 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Shared runtime policy for the opt-in Neo CuTe inference path.""" + +from __future__ import ( + annotations, +) + +import os + +import torch + +_TRUE_VALUES = frozenset({"1", "true", "yes", "on"}) +_FALSE_VALUES = frozenset({"0", "false", "no", "off"}) +NEO_CUTE_INFER_ENV = "DP_NEO_CUTE_INFER" +SM80_PROFILE_CAPABILITIES = frozenset({(8, 0), (8, 6)}) +SM90_CAPABILITY = (9, 0) +FUSED_SO2_GATE_CAPABILITIES = frozenset({(8, 9), (12, 0)}) +OUTPUT_GRID_SM90_C96_ASYMMETRIC_PANELS_ENV = ( + "DP_CUTE_OUTPUT_GRID_SM90_C96_ASYMMETRIC_PANELS" +) +READOUT_INPUT_FOLD_SM90_ENV = "DP_CUTE_READOUT_INPUT_FOLD_SM90" +SUPPORTED_K1_CAPABILITIES = SM80_PROFILE_CAPABILITIES | frozenset( + {(8, 9), (9, 0), (10, 0), (12, 0)} +) +_GIE_DEFAULT_CAPABILITIES = SM80_PROFILE_CAPABILITIES +INT32_MAX = (1 << 31) - 1 +K1_VALUES_PER_EDGE = 10 * 64 +K1_VALUES_PER_NODE = 16 * 64 +PORTABLE_TILED_BACKEND = "portable_tiled" +PYTORCH_BACKEND = "pytorch" +SUPPORTED_HIDDEN_CHANNELS = (96, 192) +_OUTPUT_GRID_ARCH_BACKENDS = { + "sm80": { + 96: PORTABLE_TILED_BACKEND, + 192: PORTABLE_TILED_BACKEND, + }, + "sm90": { + 96: PORTABLE_TILED_BACKEND, + 192: PORTABLE_TILED_BACKEND, + }, +} + + +def _env_override(name: str) -> bool | None: + value = os.environ.get(name) + if value is None or not value.strip(): + return None + normalized = value.strip().lower() + if normalized in _TRUE_VALUES: + return True + if normalized in _FALSE_VALUES: + return False + return False + + +def is_cute_infer_enabled() -> bool: + """Return whether the process opted into the full Neo CuTe K1 path. + + ``DP_NEO_CUTE_INFER`` is deliberately separate from the + ``DP_CUTE_INFER`` inner SO2 value-path selector. The full K1 replacement + may therefore coexist with ``DP_TRITON_INFER=2``. + """ + return _env_override(NEO_CUTE_INFER_ENV) is True + + +def _current_compute_capability() -> tuple[int, int] | None: + if not torch.cuda.is_available(): + return None + try: + return tuple(torch.cuda.get_device_capability()) + except RuntimeError: + return None + + +def output_grid_arch_key(compute_capability: tuple[int, int]) -> str: + """Return the architecture-family key used for output-grid dispatch.""" + if tuple(compute_capability) in SM80_PROFILE_CAPABILITIES: + return "sm80" + major, minor = compute_capability + return f"sm{int(major)}{int(minor)}" + + +def select_output_grid_backend( + compute_capability: tuple[int, int], + hidden_channels: int, +) -> str: + """Select a width-specific CuTe or PyTorch output-grid backend.""" + hidden_channels = int(hidden_channels) + if hidden_channels not in SUPPORTED_HIDDEN_CHANNELS: + return PYTORCH_BACKEND + architecture_backends = _OUTPUT_GRID_ARCH_BACKENDS.get( + output_grid_arch_key(compute_capability) + ) + if architecture_backends is None: + return PYTORCH_BACKEND + return architecture_backends.get(hidden_channels, PYTORCH_BACKEND) + + +def is_sm80_profile_enabled( + compute_capability: tuple[int, int] | None = None, +) -> bool: + """Return whether the shared SM80/SM86 profile is selected.""" + if compute_capability is None: + compute_capability = _current_compute_capability() + return ( + is_cute_infer_enabled() + and compute_capability is not None + and tuple(compute_capability) in SM80_PROFILE_CAPABILITIES + ) + + +def _sm80_profile_feature( + name: str, + compute_capability: tuple[int, int] | None = None, +) -> bool: + """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 + + +def _profile_or_explicit_feature( + name: str, + compute_capability: tuple[int, int] | None = None, +) -> bool: + """Default on for the SM80 profile; otherwise require explicit opt-in.""" + if is_sm80_profile_enabled(compute_capability): + return _env_override(name) is not False + return is_cute_infer_enabled() and _env_override(name) is True + + +@torch.compiler.assume_constant_result +def is_k1_thin_wrapper_enabled( + compute_capability: tuple[int, int] | None = None, +) -> bool: + """Select compile-visible K1 dispatch for the SM80 profile.""" + return _profile_or_explicit_feature( + "DP_CUTE_K1_THIN_WRAPPER", + compute_capability, + ) + + +@torch.compiler.assume_constant_result +def is_cute_strict_enabled() -> bool: + """Return whether expensive CuTe contract assertions are requested.""" + return _env_override("DP_CUTE_STRICT") is True + + +def is_output_grid_bwd_sm80_c96_n48_panel_enabled( + compute_capability: tuple[int, int] | None = None, +) -> bool: + """Select the C=96, N=48 SM80 panel adjoint.""" + return _sm80_profile_feature( + "DP_CUTE_OUTPUT_GRID_BWD_SM80_C96_N48_PANEL", + compute_capability, + ) + + +def is_output_grid_fwd_sm80_c96_n48_enabled( + compute_capability: tuple[int, int] | None = None, +) -> bool: + """Select the C=96, N=48 SM80 forward.""" + return _sm80_profile_feature( + "DP_CUTE_OUTPUT_GRID_FWD_SM80_C96_N48", + compute_capability, + ) + + +@torch.compiler.assume_constant_result +def is_output_grid_sm90_c96_asymmetric_panels_enabled( + compute_capability: tuple[int, int] | None = None, +) -> bool: + """Select the C96 N64+N32 panels only on exact SM90.""" + if compute_capability is None: + compute_capability = _current_compute_capability() + return _master_gated_feature( + OUTPUT_GRID_SM90_C96_ASYMMETRIC_PANELS_ENV, + default=compute_capability is not None + and tuple(compute_capability) == SM90_CAPABILITY, + ) + + +@torch.compiler.assume_constant_result +def is_readout_input_fold_sm80_enabled( + compute_capability: tuple[int, int] | None = None, +) -> bool: + """Select the frozen C=192 Neo readout fold on SM80.""" + return _sm80_profile_feature( + "DP_CUTE_READOUT_INPUT_FOLD_SM80", + compute_capability, + ) + + +@torch.compiler.assume_constant_result +def is_readout_input_fold_sm90_enabled( + compute_capability: tuple[int, int] | None = None, +) -> bool: + """Select the frozen C=192 Neo readout fold only on exact SM90.""" + if compute_capability is None: + compute_capability = _current_compute_capability() + return _master_gated_feature( + READOUT_INPUT_FOLD_SM90_ENV, + default=compute_capability is not None + and tuple(compute_capability) == SM90_CAPABILITY, + ) + + +def is_readout_input_fold_enabled( + compute_capability: tuple[int, int] | None = None, +) -> bool: + """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 supports this compute capability.""" + return tuple(compute_capability) in SUPPORTED_K1_CAPABILITIES + + +def k1_int32_indexing_is_safe( + *, + edge_count: int, + node_count: int, +) -> bool: + """Check every flattened K1 offset represented with signed Int32.""" + if edge_count < 0 or node_count < 0: + return False + return ( + edge_count <= INT32_MAX // K1_VALUES_PER_EDGE + and node_count <= INT32_MAX // K1_VALUES_PER_NODE + ) + + +@torch.compiler.assume_constant_result +def uses_strict_fp32_matmul() -> bool: + """Read CUDA matmul precision outside Dynamo and fail closed on TF32.""" + matmul = torch.backends.cuda.matmul + try: + precision = matmul.fp32_precision + except AttributeError: + precision = None + except RuntimeError: + return False + if precision is not None and precision != "none": + return precision == "ieee" + try: + return not matmul.allow_tf32 + except RuntimeError: + return False + + +def _master_gated_feature(name: str, *, default: bool) -> bool: + if not is_cute_infer_enabled() or not default: + return False + override = _env_override(name) + return override is not False + + +def is_gie_enabled(compute_capability: tuple[int, int]) -> bool: + """Select the optimized geometric-initial-embedding path.""" + return _master_gated_feature( + "DP_CUTE_GIE", + default=compute_capability in _GIE_DEFAULT_CAPABILITIES, + ) + + +def is_packed_wigner_enabled(compute_capability: tuple[int, int]) -> bool: + """Select packed Wigner storage required by the optimized K1 profiles.""" + return _master_gated_feature( + "DP_CUTE_K1_PACKED_WIGNER", + default=compute_capability in SUPPORTED_K1_CAPABILITIES, + ) + + +@torch.compiler.assume_constant_result +def is_k1_eager_island_enabled( + compute_capability: tuple[int, int] | None = None, +) -> bool: + """Select the SM80 neighbor-list eager island for the Neo K1 path.""" + if not is_cute_infer_enabled(): + return False + override = _env_override("DP_CUTE_K1_EAGER_ISLANDS") + if override is not None: + return override + if compute_capability is None: + compute_capability = _current_compute_capability() + return compute_capability is not None and tuple(compute_capability) == (8, 0) diff --git a/deepmd/kernels/cute/neo/sm90_k1/__init__.py b/deepmd/kernels/cute/neo/sm90_k1/__init__.py new file mode 100644 index 0000000000..353bc4a488 --- /dev/null +++ b/deepmd/kernels/cute/neo/sm90_k1/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""SM90 split-complex Neo K1 implementation.""" diff --git a/deepmd/kernels/cute/neo/sm90_k1/final_phase_c.py b/deepmd/kernels/cute/neo/sm90_k1/final_phase_c.py new file mode 100644 index 0000000000..6475566c30 --- /dev/null +++ b/deepmd/kernels/cute/neo/sm90_k1/final_phase_c.py @@ -0,0 +1,483 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Direct destination statistics for the SM90 final SO2Linear/Phase-C boundary. + +One CTA owns one ``(node, focus)`` segment. Its feature-owning threads keep +all output-degree statistics in registers while walking the destination's +edges in 64-edge chunks. The CTA writes the final node-scale ``a0`` and +``a1`` statistics directly, eliminating the global chunk partials and their +second reduction launch. + +The strict-FP32 node GEMMs are applied only after edge values have been +reduced to node-scale sufficient statistics. +""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, +) + +import cutlass +import cutlass.cute as cute +import torch +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from ..compile_cache import ( + device_aware_lru_cache, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN001, ANN202, ANN204, TC002 + +EDGE_CHUNK = 64 +FOCUS_COUNT = 2 +CHANNELS = 32 +DEGREE_COUNT = 16 +M0_WIDTH = 128 +M1_WIDTH = 96 +PACKED_WIGNER_VALUES = 46 +M0_THREADS = M0_WIDTH +M1_THREADS = M1_WIDTH * 2 +THREADS = M0_THREADS + M1_THREADS +M1_SCALE_VALUES = EDGE_CHUNK * (DEGREE_COUNT - 1) * 2 +M0_SCALE_VALUES = EDGE_CHUNK * DEGREE_COUNT +TOTAL_SCALE_VALUES = M0_SCALE_VALUES + M1_SCALE_VALUES +LOADS_PER_THREAD = (TOTAL_SCALE_VALUES + THREADS - 1) // THREADS +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + + +@dataclass(frozen=True) +class ExpandedFinalWeights: + """Dense final SO2Linear blocks selected for each output degree.""" + + w0: torch.Tensor + wc: torch.Tensor + + +@dataclass(frozen=True) +class ExpandedComplexWorkspace: + """Node-scale real and complex sufficient statistics.""" + + m0: torch.Tensor + m1: torch.Tensor + + @property + def storage_bytes(self) -> int: + return sum( + tensor.numel() * tensor.element_size() for tensor in (self.m0, self.m1) + ) + + +def prepare_expanded_final_weights( + w0: torch.Tensor, + wc: torch.Tensor, +) -> ExpandedFinalWeights: + """Select the dense input block needed by each full output degree.""" + if tuple(w0.shape) != (FOCUS_COUNT, M0_WIDTH, M0_WIDTH): + raise ValueError("w0 must have shape (2,128,128)") + if tuple(wc.shape) != (FOCUS_COUNT, M1_WIDTH, M1_WIDTH): + raise ValueError("wc must have shape (2,96,96)") + if w0.dtype != torch.float32 or wc.dtype != torch.complex64: + raise TypeError("w0/wc must be float32/complex64") + degree_by_q = (0, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3) + blocks0 = torch.stack( + [ + w0[:, :, degree * CHANNELS : (degree + 1) * CHANNELS] + for degree in degree_by_q + ], + dim=1, + ).contiguous() + blocks1 = torch.stack( + [ + wc[:, :, (degree - 1) * CHANNELS : degree * CHANNELS] + for degree in degree_by_q[1:] + ], + dim=1, + ).contiguous() + return ExpandedFinalWeights(w0=blocks0, wc=blocks1) + + +def _require_sm90_strict_fp32(device: torch.device) -> None: + if device.type != "cuda": + raise ValueError("SM90 final Phase C requires CUDA") + if tuple(torch.cuda.get_device_capability(device)) != (9, 0): + raise RuntimeError("SM90 final Phase C requires compute capability 9.0") + 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'") + + +@dataclass(frozen=True) +class DirectStatisticsForwardResult: + """Forward output and the final node-scale sufficient statistics.""" + + output: torch.Tensor + workspace: ExpandedComplexWorkspace + + @property + def statistics_storage_bytes(self) -> int: + return self.workspace.storage_bytes + + +@cute.jit +def _packed_index_runtime(q, row_slot): + base = cutlass.Int32(0) + width = cutlass.Int32(1) + local = cutlass.Int32(0) + if q >= 9: + base = cutlass.Int32(25) + width = cutlass.Int32(7) + local = q - 9 + elif q >= 4: + base = cutlass.Int32(10) + width = cutlass.Int32(5) + local = q - 4 + elif q >= 1: + base = cutlass.Int32(1) + width = cutlass.Int32(3) + local = q - 1 + return base + row_slot * width + local + + +class CuteDirectNodeStatistics: + """Accumulate all real and complex statistics in one destination CTA.""" + + @cute.jit + def __call__( + self, + m0: cute.Tensor, + m1_ri: cute.Tensor, + dt_packed: cute.Tensor, + beta: cute.Tensor, + dst_ptr: cute.Tensor, + a0: cute.Tensor, + a1_ri: cute.Tensor, + stream: CUstream, + ): + self.kernel(m0, m1_ri, dt_packed, beta, dst_ptr, a0, a1_ri).launch( + grid=[a0.shape[2], FOCUS_COUNT, 1], + block=[THREADS, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + m0: cute.Tensor, + m1_ri: cute.Tensor, + dt_packed: cute.Tensor, + beta: cute.Tensor, + dst_ptr: cute.Tensor, + a0: cute.Tensor, + a1_ri: cute.Tensor, + ): + tidx, _, _ = cute.arch.thread_idx() + node, focus, _ = cute.arch.block_idx() + node_lo = dst_ptr[node] + node_hi = dst_ptr[node + 1] + edge_count = node_hi - node_lo + chunk_count = (edge_count + EDGE_CHUNK - 1) // EDGE_CHUNK + + smem = cutlass.utils.SmemAllocator() + m0_scale_storage = smem.allocate_tensor(cutlass.Float32, M0_SCALE_VALUES) + m0_scales = cute.make_tensor( + m0_scale_storage.iterator, + cute.make_layout( + (EDGE_CHUNK, DEGREE_COUNT), + stride=(DEGREE_COUNT, 1), + ), + ) + m1_scale_storage = smem.allocate_tensor(cutlass.Float32, M1_SCALE_VALUES) + m1_scales = cute.make_tensor( + m1_scale_storage.iterator, + cute.make_layout( + (EDGE_CHUNK, DEGREE_COUNT - 1, 2), + stride=((DEGREE_COUNT - 1) * 2, 2, 1), + ), + ) + + # A single 16-value register fragment serves either a real feature or + # one component of a complex feature. Complex threads use entries + # [0, 15), avoiding two live accumulator arrays in generated code. + accumulators = cute.make_rmem_tensor( + cute.make_layout((DEGREE_COUNT,), stride=(1,)), + cutlass.Float32, + ) + accumulators.fill(0.0) + + for chunk_slot in cutlass.range(chunk_count, unroll=1): + lo = node_lo + chunk_slot * EDGE_CHUNK + hi = lo + EDGE_CHUNK + if node_hi < hi: + hi = node_hi + + for load_slot in cutlass.range_constexpr(LOADS_PER_THREAD): + linear = tidx + load_slot * THREADS + if linear < M0_SCALE_VALUES: + edge_slot = linear // DEGREE_COUNT + q = linear - edge_slot * DEGREE_COUNT + edge = lo + edge_slot + value = cutlass.Float32(0.0) + if edge < hi: + panel = _packed_index_runtime(q, cutlass.Int32(0)) + value = beta[edge, focus].to(cutlass.Float32) * dt_packed[ + edge, panel + ].to(cutlass.Float32) + m0_scales[edge_slot, q] = value + elif linear < TOTAL_SCALE_VALUES: + item = linear - M0_SCALE_VALUES + edge_slot = item // ((DEGREE_COUNT - 1) * 2) + remainder = item - edge_slot * (DEGREE_COUNT - 1) * 2 + q1 = remainder // 2 + component = remainder - q1 * 2 + edge = lo + edge_slot + value = cutlass.Float32(0.0) + if edge < hi: + panel = _packed_index_runtime(q1 + 1, component + 1) + value = beta[edge, focus].to(cutlass.Float32) * dt_packed[ + edge, panel + ].to(cutlass.Float32) + m1_scales[edge_slot, q1, component] = value + cute.arch.sync_threads() + + if tidx < M0_THREADS: + feature = tidx + for edge_slot in cutlass.range_constexpr(EDGE_CHUNK): + edge = lo + edge_slot + if edge < hi: + x = m0[focus, edge, feature].to(cutlass.Float32) + for q in cutlass.range_constexpr(DEGREE_COUNT): + value = accumulators[q].to(cutlass.Float32) + value += m0_scales[edge_slot, q] * x + accumulators[q] = value + elif tidx < THREADS: + complex_thread = tidx - M0_THREADS + feature = complex_thread // 2 + component = complex_thread - feature * 2 + for edge_slot in cutlass.range_constexpr(EDGE_CHUNK): + edge = lo + edge_slot + if edge < hi: + xr = m1_ri[focus, edge, feature, 0].to(cutlass.Float32) + xi = m1_ri[focus, edge, feature, 1].to(cutlass.Float32) + for q1 in cutlass.range_constexpr(DEGREE_COUNT - 1): + dr = m1_scales[edge_slot, q1, 0] + di = m1_scales[edge_slot, q1, 1] + value = accumulators[q1].to(cutlass.Float32) + if component == 0: + value += dr * xr + di * xi + else: + value += dr * xi - di * xr + accumulators[q1] = value + + # Every thread must finish reading shared scales before the next + # 64-edge chunk overwrites them. + cute.arch.sync_threads() + + if tidx < M0_THREADS: + feature = tidx + for q in cutlass.range_constexpr(DEGREE_COUNT): + a0[focus, q, node, feature] = accumulators[q] + elif tidx < THREADS: + complex_thread = tidx - M0_THREADS + feature = complex_thread // 2 + component = complex_thread - feature * 2 + for q1 in cutlass.range_constexpr(DEGREE_COUNT - 1): + a1_ri[focus, q1, node, feature, component] = accumulators[q1] + + +def _fake_m0_edges(): + return make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, cute.sym_int64(), M0_WIDTH), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + + +def _fake_m1_edges_ri(): + return make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, cute.sym_int64(), M1_WIDTH, 2), + stride_order=(3, 2, 1, 0), + **FAKE_TENSOR_KW, + ) + + +def _fake_dt(): + return make_fake_compact_tensor( + cutlass.Float32, + (cute.sym_int64(), PACKED_WIGNER_VALUES), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + + +def _fake_beta(): + return make_fake_compact_tensor( + cutlass.Float32, + (cute.sym_int64(), FOCUS_COUNT), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + + +def _fake_index(): + return make_fake_compact_tensor( + cutlass.Int32, + (cute.sym_int64(),), + stride_order=(0,), + **FAKE_TENSOR_KW, + ) + + +def _fake_a0(): + return make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, DEGREE_COUNT, cute.sym_int64(), M0_WIDTH), + stride_order=(3, 2, 1, 0), + **FAKE_TENSOR_KW, + ) + + +def _fake_a1_ri(): + return make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, DEGREE_COUNT - 1, cute.sym_int64(), M1_WIDTH, 2), + stride_order=(4, 3, 2, 1, 0), + **FAKE_TENSOR_KW, + ) + + +@device_aware_lru_cache(maxsize=2) +def _compiled_direct_statistics() -> Callable: + return cute.compile( + CuteDirectNodeStatistics(), + _fake_m0_edges(), + _fake_m1_edges_ri(), + _fake_dt(), + _fake_beta(), + _fake_index(), + _fake_a0(), + _fake_a1_ri(), + stream=make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +def _validate_forward_inputs( + m0: torch.Tensor, + m1: torch.Tensor, + dt_packed: torch.Tensor, + beta: torch.Tensor, + dst_ptr: torch.Tensor, +) -> tuple[torch.device, int]: + device = m0.device + _require_sm90_strict_fp32(device) + edge_count = int(m0.shape[1]) + node_count = int(dst_ptr.numel() - 1) + expected = ( + ("m0", m0, (FOCUS_COUNT, edge_count, M0_WIDTH), torch.float32), + ("m1", m1, (FOCUS_COUNT, edge_count, M1_WIDTH), torch.complex64), + ( + "dt_packed", + dt_packed, + (edge_count, PACKED_WIGNER_VALUES), + torch.float32, + ), + ("beta", beta, (edge_count, FOCUS_COUNT), torch.float32), + ("dst_ptr", dst_ptr, (node_count + 1,), torch.int32), + ) + for name, tensor, shape, dtype in expected: + if tuple(tensor.shape) != shape or tensor.dtype != dtype: + raise ValueError(f"{name} must have shape {shape} and dtype {dtype}") + if tensor.device != device or not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous on {device}") + if tensor.data_ptr() % 16: + raise ValueError(f"{name} must be at least 16-byte aligned") + return device, node_count + + +def run_direct_statistics_forward( + *, + m0: torch.Tensor, + m1: torch.Tensor, + dt_packed: torch.Tensor, + beta: torch.Tensor, + dst_ptr: torch.Tensor, + weights: ExpandedFinalWeights, + edge_chunk: int = EDGE_CHUNK, + chunk_slots: int | None = None, +) -> DirectStatisticsForwardResult: + """Build final node statistics directly and apply node-scale SO2Linear. + + ``edge_chunk`` and ``chunk_slots`` retain the chunked-forward argument + contract used by the SM90 K1 runner. + Only the 64-edge internal schedule is supported; ``chunk_slots`` is not an + allocation dimension in this implementation. + """ + if edge_chunk != EDGE_CHUNK: + raise ValueError(f"direct statistics requires edge_chunk={EDGE_CHUNK}") + if chunk_slots is not None and chunk_slots <= 0: + raise ValueError("chunk_slots must be positive when provided") + device, node_count = _validate_forward_inputs(m0, m1, dt_packed, beta, dst_ptr) + a0 = torch.empty( + (FOCUS_COUNT, DEGREE_COUNT, node_count, M0_WIDTH), + device=device, + dtype=torch.float32, + ) + a1 = torch.empty( + (FOCUS_COUNT, DEGREE_COUNT - 1, node_count, M1_WIDTH), + device=device, + dtype=torch.complex64, + ) + with torch.cuda.device(device): + _compiled_direct_statistics()( + m0, + torch.view_as_real(m1), + dt_packed, + beta, + dst_ptr, + a0, + torch.view_as_real(a1), + ) + out0 = torch.bmm(a0.flatten(0, 1), weights.w0.flatten(0, 1)).view( + FOCUS_COUNT, DEGREE_COUNT, node_count, CHANNELS + ) + out1 = torch.bmm(a1.flatten(0, 1), weights.wc.flatten(0, 1)).view( + FOCUS_COUNT, DEGREE_COUNT - 1, node_count, CHANNELS + ) + output = out0.permute(2, 0, 1, 3).contiguous() + output[:, :, 1:] += out1.real.permute(2, 0, 1, 3) + return DirectStatisticsForwardResult( + output=output, + workspace=ExpandedComplexWorkspace(m0=a0, m1=a1), + ) + + +__all__ = [ + "EDGE_CHUNK", + "THREADS", + "DirectStatisticsForwardResult", + "ExpandedFinalWeights", + "prepare_expanded_final_weights", + "run_direct_statistics_forward", +] diff --git a/deepmd/kernels/cute/neo/sm90_k1/output_gate.py b/deepmd/kernels/cute/neo/sm90_k1/output_gate.py new file mode 100644 index 0000000000..480e8217e2 --- /dev/null +++ b/deepmd/kernels/cute/neo/sm90_k1/output_gate.py @@ -0,0 +1,204 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""SM90 CuTe epilogue for chunked final-SO2Linear sufficient statistics.""" + +from __future__ import ( + annotations, +) + +import math +from typing import ( + TYPE_CHECKING, +) + +import cutlass +import cutlass.cute as cute +import torch +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from ..compile_cache import ( + device_aware_lru_cache, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN001, ANN202, TC002 + +FOCUS_COUNT = 2 +DEGREE_COUNT = 16 +CHANNELS = 32 +HIDDEN = FOCUS_COUNT * CHANNELS +THREADS = HIDDEN +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + + +@cute.jit +def _sigmoid(value): + one = cutlass.Float32(1.0) + return one / (one + cute.exp(-value)) + + +@cute.jit +def _chunked_final_output_gate_jit( + raw: cute.Tensor, + x_wide: cute.Tensor, + norm_scale: cute.Tensor, + gate_weight: cute.Tensor, + rotate_inv_rescale: cute.Tensor, + out: cute.Tensor, + eps: cutlass.Constexpr[float], + stream: CUstream, +): + _chunked_final_output_gate_kernel( + raw, + x_wide, + norm_scale, + gate_weight, + rotate_inv_rescale, + out, + eps, + ).launch( + grid=[raw.shape[0], 1, 1], + block=[THREADS, 1, 1], + stream=stream, + ) + + +@cute.kernel +def _chunked_final_output_gate_kernel( + raw: cute.Tensor, + x_wide: cute.Tensor, + norm_scale: cute.Tensor, + gate_weight: cute.Tensor, + rotate_inv_rescale: cute.Tensor, + out: cute.Tensor, + eps: cutlass.Constexpr[float], +): + tidx, _, _ = cute.arch.thread_idx() + node, _, _ = cute.arch.block_idx() + focus = tidx // CHANNELS + channel = tidx - focus * CHANNELS + + x = x_wide[node, 0, tidx].to(cutlass.Float32) + square_sum = cute.arch.warp_reduction_sum(x * x) + inv_rms = cute.rsqrt(square_sum / cutlass.Float32(CHANNELS) + cutlass.Float32(eps)) + logit_part = ( + x + * inv_rms + * norm_scale[focus, channel].to(cutlass.Float32) + * gate_weight[channel, focus, 0].to(cutlass.Float32) + ) + gate = _sigmoid(cute.arch.warp_reduction_sum(logit_part)) + for degree in cutlass.range_constexpr(DEGREE_COUNT): + value = raw[node, focus, degree, channel].to(cutlass.Float32) + value *= rotate_inv_rescale[degree].to(cutlass.Float32) + out[node, degree, tidx] = value * gate + + +def _fake_float(shape: tuple[object, ...], stride_order: tuple[int, ...]): + return make_fake_compact_tensor( + cutlass.Float32, + shape, + stride_order=stride_order, + **FAKE_TENSOR_KW, + ) + + +@device_aware_lru_cache(maxsize=8) +def _compiled_chunked_final_output_gate(eps: float) -> Callable: + nodes = cute.sym_int64() + return cute.compile( + _chunked_final_output_gate_jit, + _fake_float( + (nodes, FOCUS_COUNT, DEGREE_COUNT, CHANNELS), + (3, 2, 1, 0), + ), + _fake_float((nodes, DEGREE_COUNT, HIDDEN), (2, 1, 0)), + _fake_float((FOCUS_COUNT, CHANNELS), (1, 0)), + _fake_float((CHANNELS, FOCUS_COUNT, 1), (2, 1, 0)), + _fake_float((DEGREE_COUNT,), (0,)), + _fake_float((nodes, DEGREE_COUNT, HIDDEN), (2, 1, 0)), + eps, + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +def _require_tensor( + name: str, + tensor: torch.Tensor, + shape: tuple[int, ...], + device: torch.device, +) -> None: + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}") + if tensor.dtype != torch.float32 or tensor.device != device: + raise ValueError(f"{name} must be contiguous FP32 on {device}") + if not tensor.is_contiguous() or tensor.data_ptr() % 16: + raise ValueError(f"{name} must be contiguous and 16-byte aligned") + + +def run_chunked_final_output_gate( + *, + raw: torch.Tensor, + x_wide: torch.Tensor, + norm_scale: torch.Tensor, + gate_weight: torch.Tensor, + rotate_inv_rescale: torch.Tensor, + eps: float, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Apply the real Neo rotate-rescale and output gate to node statistics.""" + if not math.isfinite(eps) or eps <= 0.0: + raise ValueError("output-gate epsilon must be finite and positive") + device = raw.device + if device.type != "cuda" or tuple(torch.cuda.get_device_capability(device)) != ( + 9, + 0, + ): + raise RuntimeError("final output gate requires SM90") + if torch.backends.cuda.matmul.allow_tf32: + raise RuntimeError("strict FP32 requires allow_tf32=False") + node_count = int(raw.shape[0]) + _require_tensor( + "raw", + raw, + (node_count, FOCUS_COUNT, DEGREE_COUNT, CHANNELS), + device, + ) + _require_tensor("x_wide", x_wide, (node_count, DEGREE_COUNT, HIDDEN), device) + _require_tensor("norm_scale", norm_scale, (FOCUS_COUNT, CHANNELS), device) + _require_tensor("gate_weight", gate_weight, (CHANNELS, FOCUS_COUNT, 1), device) + _require_tensor("rotate_inv_rescale", rotate_inv_rescale, (DEGREE_COUNT,), device) + if out is None: + out = torch.empty( + (node_count, DEGREE_COUNT, HIDDEN), + device=device, + dtype=torch.float32, + ) + else: + _require_tensor("out", out, (node_count, DEGREE_COUNT, HIDDEN), device) + with torch.cuda.device(device): + _compiled_chunked_final_output_gate(float(eps))( + raw, + x_wide, + norm_scale, + gate_weight, + rotate_inv_rescale, + out, + ) + return out + + +__all__ = ["run_chunked_final_output_gate"] diff --git a/deepmd/kernels/cute/neo/sm90_k1/persistent.py b/deepmd/kernels/cute/neo/sm90_k1/persistent.py new file mode 100644 index 0000000000..3c1a00641c --- /dev/null +++ b/deepmd/kernels/cute/neo/sm90_k1/persistent.py @@ -0,0 +1,593 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Persistent-complex strict-FP32 Neo SO2 stack for SM90. + +The fixed Neo ``m=1`` block is a complex 96-wide representation of the dense +real block ``[[U,V],[-V,U]]``. The first two frozen SO2Linear layers and gated +residuals remain in this representation without intermediate packing. +""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, +) + +import cutlass +import cutlass.cute as cute +import torch +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) +from torch import ( + Tensor, +) + +from ..compile_cache import ( + device_aware_lru_cache, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN001, ANN202, TC002 + +FOCUS_COUNT = 2 +CHANNELS = 32 +M0_ROWS = 4 +M1_ROWS = 3 +M0_WIDTH = M0_ROWS * CHANNELS +M1_WIDTH = M1_ROWS * CHANNELS +PAIR_WIDTH = 2 * M1_WIDTH +GATE_GROUPS = 3 +GATE_WIDTH = GATE_GROUPS * CHANNELS +GATED_LAYERS = 2 +STACK_LAYERS = 3 + +ROWS_PER_BLOCK = 8 +THREADS = ROWS_PER_BLOCK * CHANNELS +WEIGHT_SMEM_STRIDE = GATE_WIDTH + 1 +WEIGHT_VALUES = CHANNELS * GATE_WIDTH +WEIGHT_LOADS_PER_THREAD = WEIGHT_VALUES // THREADS +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + +__all__ = [ + "NeoPersistentComplexSaved", + "NeoPersistentComplexState", + "NeoPersistentComplexWeights", + "prepare_neo_persistent_complex_weights", + "validate_neo_persistent_complex_state", +] + + +@dataclass(frozen=True) +class NeoPersistentComplexState: + """Focus-major state consumed directly by the persistent stack. + + ``m0`` has shape ``(2,E,128)`` and dtype ``float32``. ``m1`` has shape + ``(2,E,96)`` and dtype ``complex64``. Both tensors are contiguous; no + packing or transposition is needed by the real or complex batched GEMMs. + """ + + m0: Tensor + m1: Tensor + + @property + def edge_count(self) -> int: + return int(self.m0.shape[1]) + + @property + def storage_bytes(self) -> int: + return sum( + tensor.numel() * tensor.element_size() for tensor in (self.m0, self.m1) + ) + + +@dataclass(frozen=True) +class NeoPersistentComplexWeights: + """Frozen strict-FP32 operands in forward and input-adjoint orientation.""" + + w0: Tensor + wc: Tensor + w0_h: Tensor + wc_h: Tensor + gate: Tensor + + +@dataclass(frozen=True) +class NeoPersistentComplexSaved: + """Minimal exact gate state for the two nonlinear layers.""" + + z0: tuple[Tensor, Tensor] + z1: tuple[Tensor, Tensor] + + @property + def storage_bytes(self) -> int: + return sum( + tensor.numel() * tensor.element_size() for tensor in (*self.z0, *self.z1) + ) + + +def _require_shape(name: str, tensor: Tensor, shape: tuple[int, ...]) -> None: + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}") + + +def _require_frozen_cuda_tensor( + name: str, + tensor: Tensor, + *, + dtype: torch.dtype, + device: torch.device, +) -> None: + if tensor.dtype != dtype: + raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}, got {tensor.device}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + if tensor.requires_grad: + raise ValueError(f"{name} must be frozen; parameter gradients are out of scope") + + +def validate_neo_persistent_complex_state( + state: NeoPersistentComplexState, + *, + name: str = "state", +) -> None: + """Validate the direct Phase-A/Phase-C split interface.""" + if ( + state.m0.ndim != 3 + or state.m0.shape[0] != FOCUS_COUNT + or state.m0.shape[2] != M0_WIDTH + ): + raise ValueError( + f"{name}.m0 must have shape (2,E,128), got {tuple(state.m0.shape)}" + ) + _require_shape( + f"{name}.m1", + state.m1, + (FOCUS_COUNT, state.edge_count, M1_WIDTH), + ) + if state.edge_count <= 0: + raise ValueError(f"{name} requires E > 0") + if state.m0.dtype != torch.float32 or state.m1.dtype != torch.complex64: + raise TypeError(f"{name} requires float32 m0 and complex64 m1") + if state.m0.device != state.m1.device: + raise ValueError(f"{name}.m0 and {name}.m1 must share a device") + if not state.m0.is_contiguous() or not state.m1.is_contiguous(): + raise ValueError(f"{name} tensors must be focus-major contiguous") + + +def prepare_neo_persistent_complex_weights( + w0: Tensor, + wp: Tensor, + gate: Tensor, +) -> NeoPersistentComplexWeights: + """Convert exact block-real frozen weights to persistent complex weights. + + ``w0`` and ``wp`` use the live ``(input, output)`` orientation consumed by + ``torch.bmm``. The pair block must be exactly ``[[U,V],[-V,U]]``. + """ + _require_shape("w0", w0, (STACK_LAYERS, FOCUS_COUNT, M0_WIDTH, M0_WIDTH)) + _require_shape( + "wp", + wp, + (STACK_LAYERS, FOCUS_COUNT, PAIR_WIDTH, PAIR_WIDTH), + ) + _require_shape( + "gate", + gate, + (GATED_LAYERS, FOCUS_COUNT, CHANNELS, GATE_WIDTH), + ) + if not w0.is_cuda: + raise ValueError("weights must be CUDA tensors") + device = w0.device + for name, tensor in (("w0", w0), ("wp", wp), ("gate", gate)): + _require_frozen_cuda_tensor( + name, + tensor, + dtype=torch.float32, + device=device, + ) + + u = wp[:, :, :M1_WIDTH, :M1_WIDTH] + v = wp[:, :, :M1_WIDTH, M1_WIDTH:] + if not torch.equal(wp[:, :, M1_WIDTH:, :M1_WIDTH], -v): + raise ValueError("wp lower-left block must equal -V exactly") + if not torch.equal(wp[:, :, M1_WIDTH:, M1_WIDTH:], u): + raise ValueError("wp lower-right block must equal U exactly") + + w0_live = w0.detach().contiguous() + wc_live = torch.complex(u, v).contiguous() + w0_h = w0_live.transpose(-2, -1).contiguous() + wc_h = wc_live.conj().transpose(-2, -1).contiguous() + + return NeoPersistentComplexWeights( + w0=w0_live, + wc=wc_live, + w0_h=w0_h, + wc_h=wc_h, + gate=gate.detach().contiguous(), + ) + + +@cute.jit +def _sigmoid(value): + return cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cute.exp(-value)) + + +@cute.jit +def _stage_gate_weight(gate_weight, shared_weight, focus, tidx): + for load_slot in cutlass.range_constexpr(WEIGHT_LOADS_PER_THREAD): + linear = tidx + load_slot * THREADS + source_channel = linear // GATE_WIDTH + gate_channel = linear - source_channel * GATE_WIDTH + shared_weight[source_channel * WEIGHT_SMEM_STRIDE + gate_channel] = gate_weight[ + focus, source_channel, gate_channel + ].to(cutlass.Float32) + + +@cute.jit +def _load_gate_values(shared_weight, scalar_rows, row_slot, channel): + gate0_logit = cutlass.Float32(0.0) + gate1_logit = cutlass.Float32(0.0) + gate2_logit = cutlass.Float32(0.0) + scalar_base = row_slot * CHANNELS + for source_channel in cutlass.range_constexpr(CHANNELS): + source = scalar_rows[scalar_base + source_channel] + weight_base = source_channel * WEIGHT_SMEM_STRIDE + channel + gate0_logit += source * shared_weight[weight_base] + gate1_logit += source * shared_weight[weight_base + CHANNELS] + gate2_logit += source * shared_weight[weight_base + 2 * CHANNELS] + return _sigmoid(gate0_logit), _sigmoid(gate1_logit), _sigmoid(gate2_logit) + + +@cute.jit +def _select_gate(gate0, gate1, gate2, group): + gate = gate0 + if cutlass.const_expr(group == 1): + gate = gate1 + if cutlass.const_expr(group == 2): + gate = gate2 + return gate + + +@cute.jit +def _persistent_gate_forward_jit( + residual0: cute.Tensor, + residual1_ri: cute.Tensor, + z0: cute.Tensor, + z1_ri: cute.Tensor, + gate_weight: cute.Tensor, + out0: cute.Tensor, + out1_ri: cute.Tensor, + stream: CUstream, +): + edge_count = z0.shape[1] + _persistent_gate_forward_kernel( + residual0, + residual1_ri, + z0, + z1_ri, + gate_weight, + out0, + out1_ri, + ).launch( + grid=[cute.ceil_div(edge_count, ROWS_PER_BLOCK), FOCUS_COUNT, 1], + block=[THREADS, 1, 1], + stream=stream, + ) + + +@cute.kernel +def _persistent_gate_forward_kernel( + residual0: cute.Tensor, + residual1_ri: cute.Tensor, + z0: cute.Tensor, + z1_ri: cute.Tensor, + gate_weight: cute.Tensor, + out0: cute.Tensor, + out1_ri: cute.Tensor, +): + tidx, _, _ = cute.arch.thread_idx() + edge_block, focus, _ = cute.arch.block_idx() + row_slot = tidx // CHANNELS + channel = tidx - row_slot * CHANNELS + edge = edge_block * ROWS_PER_BLOCK + row_slot + edge_count = z0.shape[1] + + smem = cutlass.utils.SmemAllocator() + shared_weight = smem.allocate_tensor( + cutlass.Float32, + CHANNELS * WEIGHT_SMEM_STRIDE, + ) + scalar_rows = smem.allocate_tensor( + cutlass.Float32, + ROWS_PER_BLOCK * CHANNELS, + ) + _stage_gate_weight(gate_weight, shared_weight, focus, tidx) + scalar = cutlass.Float32(0.0) + if edge < edge_count: + scalar = z0[focus, edge, channel].to(cutlass.Float32) + scalar_rows[row_slot * CHANNELS + channel] = scalar + cute.arch.sync_threads() + + if edge < edge_count: + gate0, gate1, gate2 = _load_gate_values( + shared_weight, + scalar_rows, + row_slot, + channel, + ) + out0[focus, edge, channel] = ( + residual0[focus, edge, channel].to(cutlass.Float32) + + scalar * _sigmoid(scalar) + ).to(out0.element_type) + for group in cutlass.range_constexpr(GATE_GROUPS): + gate = _select_gate(gate0, gate1, gate2, group) + m0_col = (group + 1) * CHANNELS + channel + out0[focus, edge, m0_col] = ( + residual0[focus, edge, m0_col].to(cutlass.Float32) + + z0[focus, edge, m0_col].to(cutlass.Float32) * gate + ).to(out0.element_type) + m1_col = group * CHANNELS + channel + for component in cutlass.range_constexpr(2): + out1_ri[focus, edge, m1_col, component] = ( + residual1_ri[focus, edge, m1_col, component].to(cutlass.Float32) + + z1_ri[focus, edge, m1_col, component].to(cutlass.Float32) * gate + ).to(out1_ri.element_type) + + +@cute.jit +def _persistent_gate_adjoint_jit( + grad0: cute.Tensor, + grad1_ri: cute.Tensor, + z0: cute.Tensor, + z1_ri: cute.Tensor, + gate_weight: cute.Tensor, + grad_z0: cute.Tensor, + grad_z1_ri: cute.Tensor, + stream: CUstream, +): + edge_count = z0.shape[1] + _persistent_gate_adjoint_kernel( + grad0, + grad1_ri, + z0, + z1_ri, + gate_weight, + grad_z0, + grad_z1_ri, + ).launch( + grid=[cute.ceil_div(edge_count, ROWS_PER_BLOCK), FOCUS_COUNT, 1], + block=[THREADS, 1, 1], + stream=stream, + ) + + +@cute.kernel +def _persistent_gate_adjoint_kernel( + grad0: cute.Tensor, + grad1_ri: cute.Tensor, + z0: cute.Tensor, + z1_ri: cute.Tensor, + gate_weight: cute.Tensor, + grad_z0: cute.Tensor, + grad_z1_ri: cute.Tensor, +): + tidx, _, _ = cute.arch.thread_idx() + edge_block, focus, _ = cute.arch.block_idx() + row_slot = tidx // CHANNELS + channel = tidx - row_slot * CHANNELS + edge = edge_block * ROWS_PER_BLOCK + row_slot + edge_count = z0.shape[1] + + smem = cutlass.utils.SmemAllocator() + shared_weight = smem.allocate_tensor( + cutlass.Float32, + CHANNELS * WEIGHT_SMEM_STRIDE, + ) + scalar_rows = smem.allocate_tensor( + cutlass.Float32, + ROWS_PER_BLOCK * CHANNELS, + ) + grad_logits = smem.allocate_tensor( + cutlass.Float32, + ROWS_PER_BLOCK * GATE_WIDTH, + ) + _stage_gate_weight(gate_weight, shared_weight, focus, tidx) + scalar = cutlass.Float32(0.0) + if edge < edge_count: + scalar = z0[focus, edge, channel].to(cutlass.Float32) + scalar_rows[row_slot * CHANNELS + channel] = scalar + cute.arch.sync_threads() + + grad_scalar = cutlass.Float32(0.0) + grad_logit0 = cutlass.Float32(0.0) + grad_logit1 = cutlass.Float32(0.0) + grad_logit2 = cutlass.Float32(0.0) + if edge < edge_count: + gate0, gate1, gate2 = _load_gate_values( + shared_weight, + scalar_rows, + row_slot, + channel, + ) + scalar_sigmoid = _sigmoid(scalar) + grad_scalar = ( + grad0[focus, edge, channel].to(cutlass.Float32) + * scalar_sigmoid + * (cutlass.Float32(1.0) + scalar * (cutlass.Float32(1.0) - scalar_sigmoid)) + ) + + for group in cutlass.range_constexpr(GATE_GROUPS): + gate = _select_gate(gate0, gate1, gate2, group) + m0_col = (group + 1) * CHANNELS + channel + upstream0 = grad0[focus, edge, m0_col].to(cutlass.Float32) + value0 = z0[focus, edge, m0_col].to(cutlass.Float32) + grad_z0[focus, edge, m0_col] = (upstream0 * gate).to(grad_z0.element_type) + contribution = upstream0 * value0 + + m1_col = group * CHANNELS + channel + for component in cutlass.range_constexpr(2): + upstream1 = grad1_ri[focus, edge, m1_col, component].to(cutlass.Float32) + value1 = z1_ri[focus, edge, m1_col, component].to(cutlass.Float32) + grad_z1_ri[focus, edge, m1_col, component] = (upstream1 * gate).to( + grad_z1_ri.element_type + ) + contribution += upstream1 * value1 + + gate_derivative = gate * (cutlass.Float32(1.0) - gate) + if cutlass.const_expr(group == 0): + grad_logit0 = contribution * gate_derivative + if cutlass.const_expr(group == 1): + grad_logit1 = contribution * gate_derivative + if cutlass.const_expr(group == 2): + grad_logit2 = contribution * gate_derivative + + grad_base = row_slot * GATE_WIDTH + channel + grad_logits[grad_base] = grad_logit0 + grad_logits[grad_base + CHANNELS] = grad_logit1 + grad_logits[grad_base + 2 * CHANNELS] = grad_logit2 + cute.arch.sync_threads() + + if edge < edge_count: + for gate_channel in cutlass.range_constexpr(GATE_WIDTH): + grad_scalar += ( + grad_logits[row_slot * GATE_WIDTH + gate_channel] + * shared_weight[channel * WEIGHT_SMEM_STRIDE + gate_channel] + ) + grad_z0[focus, edge, channel] = grad_scalar.to(grad_z0.element_type) + + +def _fake_m0(): + edge_count = cute.sym_int64() + return make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, edge_count, M0_WIDTH), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + + +def _fake_m1_ri(): + edge_count = cute.sym_int64() + return make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, edge_count, M1_WIDTH, 2), + stride_order=(3, 2, 1, 0), + **FAKE_TENSOR_KW, + ) + + +def _fake_gate_weight(): + return make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, CHANNELS, GATE_WIDTH), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + + +def _compile_forward() -> Callable: + return cute.compile( + _persistent_gate_forward_jit, + _fake_m0(), + _fake_m1_ri(), + _fake_m0(), + _fake_m1_ri(), + _fake_gate_weight(), + _fake_m0(), + _fake_m1_ri(), + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +def _compile_adjoint() -> Callable: + return cute.compile( + _persistent_gate_adjoint_jit, + _fake_m0(), + _fake_m1_ri(), + _fake_m0(), + _fake_m1_ri(), + _fake_gate_weight(), + _fake_m0(), + _fake_m1_ri(), + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +@device_aware_lru_cache(maxsize=2) +def _compiled_forward() -> Callable: + return _compile_forward() + + +@device_aware_lru_cache(maxsize=2) +def _compiled_adjoint() -> Callable: + return _compile_adjoint() + + +def _m1_real_view(m1: Tensor) -> Tensor: + view = torch.view_as_real(m1) + if not view.is_contiguous(): + raise ValueError("complex state must expose a contiguous real/imag view") + return view + + +def _run_gate_forward( + residual: NeoPersistentComplexState, + z: NeoPersistentComplexState, + gate_weight: Tensor, + out: NeoPersistentComplexState, +) -> None: + with torch.cuda.device(z.m0.device): + _compiled_forward()( + residual.m0, + _m1_real_view(residual.m1), + z.m0, + _m1_real_view(z.m1), + gate_weight, + out.m0, + _m1_real_view(out.m1), + ) + + +def _run_gate_adjoint( + grad: NeoPersistentComplexState, + z: NeoPersistentComplexState, + gate_weight: Tensor, + out: NeoPersistentComplexState, +) -> None: + with torch.cuda.device(z.m0.device): + _compiled_adjoint()( + grad.m0, + _m1_real_view(grad.m1), + z.m0, + _m1_real_view(z.m1), + gate_weight, + out.m0, + _m1_real_view(out.m1), + ) + + +def _empty_state_like(state: NeoPersistentComplexState) -> NeoPersistentComplexState: + return NeoPersistentComplexState( + m0=torch.empty_like(state.m0), + m1=torch.empty_like(state.m1), + ) diff --git a/deepmd/kernels/cute/neo/sm90_k1/phase_a.py b/deepmd/kernels/cute/neo/sm90_k1/phase_a.py new file mode 100644 index 0000000000..476e86cdd2 --- /dev/null +++ b/deepmd/kernels/cute/neo/sm90_k1/phase_a.py @@ -0,0 +1,554 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Direct strict-FP32 Phase-A producer for the persistent-complex SO2 stack. + +The generic boundary materializes Neo's reduced SO2 state as block-real +``(E,2,10,32)`` and then launches a second kernel to transpose/split it into +focus-major ``m=0`` real and interleaved ``m=1`` complex panels. This producer +applies the same packed-Wigner rotation, compact radial +maps, and rank-1 channel basis, but writes the persistent representation +directly: + +* reduced rows 0..3 -> ``m0[focus, edge, 4 * channel]``; +* reduced rows 4..6 -> the real component of ``m1``; +* reduced rows 7..9 -> the imaginary component of ``m1``. + +No full-edge block-real slab exists on this path. ``N`` and ``E`` +remain runtime dimensions; only the Neo representation contract is static. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, +) + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as cute_utils +import torch +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from ..compile_cache import ( + device_aware_lru_cache, +) +from ..k1_wigner_layout import ( + PACKED_VALUE_COUNT, +) +from .persistent import ( + NeoPersistentComplexState, + validate_neo_persistent_complex_state, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN001, ANN202, ANN204 + +EDGE_TILE = 32 +THREADS = 256 +FOCUS_COUNT = 2 +FOCUS_DIM = 32 +FULL_CHANNELS = FOCUS_COUNT * FOCUS_DIM +M0_ROWS = 4 +M1_ROWS = 3 +M0_WIDTH = M0_ROWS * FOCUS_DIM +M1_WIDTH = M1_ROWS * FOCUS_DIM +RADIAL_COMPACT = 25 + +D_CACHE_BYTES = EDGE_TILE * PACKED_VALUE_COUNT * 4 +RADIAL_CACHE_BYTES = EDGE_TILE * RADIAL_COMPACT * 4 +SRC_CACHE_BYTES = EDGE_TILE * 4 +BASIS_CACHE_BYTES = FULL_CHANNELS * 4 +CTA_SHARED_BYTES = ( + D_CACHE_BYTES + RADIAL_CACHE_BYTES + SRC_CACHE_BYTES + BASIS_CACHE_BYTES +) + +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + +__all__ = [ + "CTA_SHARED_BYTES", + "run_neo_phase_a_persistent_complex_fp32", +] + + +def _fake_x_wide(): + return make_fake_compact_tensor( + cutlass.Float32, + (cute.sym_int64(), 16 * FULL_CHANNELS), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + + +def _fake_src(): + return make_fake_compact_tensor( + cutlass.Int32, + (cute.sym_int64(),), + stride_order=(0,), + **FAKE_TENSOR_KW, + ) + + +def _fake_edge_matrix(columns: int): + return make_fake_compact_tensor( + cutlass.Float32, + (cute.sym_int64(), columns), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + + +def _fake_channel_basis(): + return make_fake_compact_tensor( + cutlass.Float32, + (FULL_CHANNELS,), + stride_order=(0,), + **FAKE_TENSOR_KW, + ) + + +def _fake_m0(): + return make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, cute.sym_int64(), M0_WIDTH), + stride_order=(2, 1, 0), + **FAKE_TENSOR_KW, + ) + + +def _fake_m1_ri(): + return make_fake_compact_tensor( + cutlass.Float32, + (FOCUS_COUNT, cute.sym_int64(), M1_WIDTH, 2), + stride_order=(3, 2, 1, 0), + **FAKE_TENSOR_KW, + ) + + +@cute.jit +def _rotation_row_cached( + x_wide, + d_cache, + edge_row, + src_node, + channel, + panel_start: cutlass.Constexpr[int], + full_start: cutlass.Constexpr[int], + width: cutlass.Constexpr[int], +): + """Evaluate one packed-Wigner row with the Phase-A reduction order.""" + value = cutlass.Float32(0.0) + for local_col in cutlass.range_constexpr(width): + value += d_cache[edge_row, panel_start + local_col].to( + cutlass.Float32 + ) * x_wide[ + src_node, + (full_start + local_col) * FULL_CHANNELS + channel, + ].to(cutlass.Float32) + return value + + +class CuteNeoPhaseAPersistentComplexFP32: + """Produce native persistent-complex panels without a dense boundary.""" + + @cute.jit + def __call__( + self, + x_wide, + src, + d_full, + radial_compact, + channel_basis, + m0, + m1_ri, + stream: cuda.CUstream = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT), + ): + d_layout = cute.make_layout( + (EDGE_TILE, PACKED_VALUE_COUNT), + stride=(PACKED_VALUE_COUNT, 1), + ) + radial_layout = cute.make_layout( + (EDGE_TILE, RADIAL_COMPACT), + stride=(RADIAL_COMPACT, 1), + ) + edge_layout = cute.make_layout((EDGE_TILE,), stride=(1,)) + basis_layout = cute.make_layout((FULL_CHANNELS,), stride=(1,)) + self.kernel( + x_wide, + src, + d_full, + radial_compact, + channel_basis, + m0, + m1_ri, + d_layout, + radial_layout, + edge_layout, + basis_layout, + ).launch( + grid=(cute.ceil_div(m0.shape[1], EDGE_TILE), 1, 1), + block=[THREADS, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + x_wide, + src, + d_full, + radial_compact, + channel_basis, + m0, + m1_ri, + d_layout, + radial_layout, + edge_layout, + basis_layout, + ): + tidx, _, _ = cute.arch.thread_idx() + edge_tile, _, _ = cute.arch.block_idx() + edge_count = m0.shape[1] + + smem = cute_utils.SmemAllocator() + d_cache = smem.allocate_tensor(cutlass.Float32, d_layout, 16) + radial_cache = smem.allocate_tensor(cutlass.Float32, radial_layout, 16) + src_cache = smem.allocate_tensor(cutlass.Int32, edge_layout, 16) + basis_cache = smem.allocate_tensor(cutlass.Float32, basis_layout, 16) + + self._load_edge_state( + src, + d_full, + radial_compact, + channel_basis, + d_cache, + radial_cache, + src_cache, + basis_cache, + edge_count, + tidx, + edge_tile, + ) + self._produce_split_panels( + x_wide, + d_cache, + radial_cache, + src_cache, + basis_cache, + m0, + m1_ri, + edge_count, + tidx, + edge_tile, + ) + + @cute.jit + def _load_edge_state( + self, + src, + d_full, + radial_compact, + channel_basis, + d_cache, + radial_cache, + src_cache, + basis_cache, + edge_count, + tidx, + edge_tile, + ): + d_slots = (EDGE_TILE * PACKED_VALUE_COUNT + THREADS - 1) // THREADS + for slot in cutlass.range_constexpr(d_slots): + linear = tidx + slot * THREADS + if linear < EDGE_TILE * PACKED_VALUE_COUNT: + edge_row = linear // PACKED_VALUE_COUNT + column = linear - edge_row * PACKED_VALUE_COUNT + edge = edge_tile * EDGE_TILE + edge_row + value = cutlass.Float32(0.0) + if edge < edge_count: + value = d_full[edge, column].to(cutlass.Float32) + d_cache[edge_row, column] = value + + radial_slots = (EDGE_TILE * RADIAL_COMPACT + THREADS - 1) // THREADS + for slot in cutlass.range_constexpr(radial_slots): + linear = tidx + slot * THREADS + if linear < EDGE_TILE * RADIAL_COMPACT: + edge_row = linear // RADIAL_COMPACT + column = linear - edge_row * RADIAL_COMPACT + edge = edge_tile * EDGE_TILE + edge_row + value = cutlass.Float32(0.0) + if edge < edge_count: + value = radial_compact[edge, column].to(cutlass.Float32) + radial_cache[edge_row, column] = value + + if tidx < EDGE_TILE: + edge = edge_tile * EDGE_TILE + tidx + src_node = cutlass.Int32(0) + if edge < edge_count: + src_node = src[edge] + src_cache[tidx] = src_node + if tidx < FULL_CHANNELS: + basis_cache[tidx] = channel_basis[tidx].to(cutlass.Float32) + cute.arch.sync_threads() + + @cute.jit + def _produce_split_panels( + self, + x_wide, + d_cache, + radial_cache, + src_cache, + basis_cache, + m0, + m1_ri, + edge_count, + tidx, + edge_tile, + ): + tasks = (EDGE_TILE * FULL_CHANNELS) // THREADS + for task in cutlass.range_constexpr(tasks): + linear = tidx + task * THREADS + edge_row = linear // FULL_CHANNELS + channel = linear - edge_row * FULL_CHANNELS + edge = edge_tile * EDGE_TILE + edge_row + + if edge < edge_count: + focus = channel // FOCUS_DIM + focus_channel = channel - focus * FOCUS_DIM + src_node = src_cache[edge_row] + x0 = _rotation_row_cached( + x_wide, d_cache, edge_row, src_node, channel, 0, 0, 1 + ) + x1 = _rotation_row_cached( + x_wide, d_cache, edge_row, src_node, channel, 1, 1, 3 + ) + x2 = _rotation_row_cached( + x_wide, d_cache, edge_row, src_node, channel, 10, 4, 5 + ) + x3 = _rotation_row_cached( + x_wide, d_cache, edge_row, src_node, channel, 25, 9, 7 + ) + x4 = _rotation_row_cached( + x_wide, d_cache, edge_row, src_node, channel, 4, 1, 3 + ) + x5 = _rotation_row_cached( + x_wide, d_cache, edge_row, src_node, channel, 15, 4, 5 + ) + x6 = _rotation_row_cached( + x_wide, d_cache, edge_row, src_node, channel, 32, 9, 7 + ) + x7 = _rotation_row_cached( + x_wide, d_cache, edge_row, src_node, channel, 7, 1, 3 + ) + x8 = _rotation_row_cached( + x_wide, d_cache, edge_row, src_node, channel, 20, 4, 5 + ) + x9 = _rotation_row_cached( + x_wide, d_cache, edge_row, src_node, channel, 39, 9, 7 + ) + + basis = basis_cache[channel].to(cutlass.Float32) + y0 = ( + radial_cache[edge_row, 0] * x0 + + radial_cache[edge_row, 4] * x1 + + radial_cache[edge_row, 8] * x2 + + radial_cache[edge_row, 12] * x3 + ) * basis + y1 = ( + radial_cache[edge_row, 1] * x0 + + radial_cache[edge_row, 5] * x1 + + radial_cache[edge_row, 9] * x2 + + radial_cache[edge_row, 13] * x3 + ) * basis + y2 = ( + radial_cache[edge_row, 2] * x0 + + radial_cache[edge_row, 6] * x1 + + radial_cache[edge_row, 10] * x2 + + radial_cache[edge_row, 14] * x3 + ) * basis + y3 = ( + radial_cache[edge_row, 3] * x0 + + radial_cache[edge_row, 7] * x1 + + radial_cache[edge_row, 11] * x2 + + radial_cache[edge_row, 15] * x3 + ) * basis + y4 = ( + radial_cache[edge_row, 16] * x4 + + radial_cache[edge_row, 19] * x5 + + radial_cache[edge_row, 22] * x6 + ) * basis + y5 = ( + radial_cache[edge_row, 17] * x4 + + radial_cache[edge_row, 20] * x5 + + radial_cache[edge_row, 23] * x6 + ) * basis + y6 = ( + radial_cache[edge_row, 18] * x4 + + radial_cache[edge_row, 21] * x5 + + radial_cache[edge_row, 24] * x6 + ) * basis + y7 = ( + radial_cache[edge_row, 16] * x7 + + radial_cache[edge_row, 19] * x8 + + radial_cache[edge_row, 22] * x9 + ) * basis + y8 = ( + radial_cache[edge_row, 17] * x7 + + radial_cache[edge_row, 20] * x8 + + radial_cache[edge_row, 23] * x9 + ) * basis + y9 = ( + radial_cache[edge_row, 18] * x7 + + radial_cache[edge_row, 21] * x8 + + radial_cache[edge_row, 24] * x9 + ) * basis + + m0[focus, edge, focus_channel] = y0 + m0[focus, edge, FOCUS_DIM + focus_channel] = y1 + m0[focus, edge, 2 * FOCUS_DIM + focus_channel] = y2 + m0[focus, edge, 3 * FOCUS_DIM + focus_channel] = y3 + m1_ri[focus, edge, focus_channel, 0] = y4 + m1_ri[focus, edge, focus_channel, 1] = y7 + m1_ri[focus, edge, FOCUS_DIM + focus_channel, 0] = y5 + m1_ri[focus, edge, FOCUS_DIM + focus_channel, 1] = y8 + m1_ri[focus, edge, 2 * FOCUS_DIM + focus_channel, 0] = y6 + m1_ri[focus, edge, 2 * FOCUS_DIM + focus_channel, 1] = y9 + + +@device_aware_lru_cache(maxsize=8) +def _compiled_producer( + device_index: int, + compute_capability: tuple[int, int], +) -> Callable: + if compute_capability != (9, 0): + raise RuntimeError("direct split Phase A requires SM90") + with torch.cuda.device(device_index): + return cute.compile( + CuteNeoPhaseAPersistentComplexFP32(), + _fake_x_wide(), + _fake_src(), + _fake_edge_matrix(PACKED_VALUE_COUNT), + _fake_edge_matrix(RADIAL_COMPACT), + _fake_channel_basis(), + _fake_m0(), + _fake_m1_ri(), + stream=make_fake_stream(use_tvm_ffi_env_stream=False), + options="--enable-tvm-ffi", + ) + + +def _validate_inputs( + x_wide: torch.Tensor, + src: torch.Tensor, + d_full: torch.Tensor, + radial_compact: torch.Tensor, + channel_basis: torch.Tensor, +) -> tuple[int, torch.Tensor]: + if x_wide.ndim != 3 or tuple(x_wide.shape[1:]) != (16, FULL_CHANNELS): + raise ValueError(f"x_wide must have shape (N,16,64), got {x_wide.shape}") + if x_wide.shape[0] <= 0: + raise ValueError("direct split Phase A requires N > 0") + if src.ndim != 1 or src.dtype not in (torch.int32, torch.int64): + raise TypeError("src must be a one-dimensional int32 or int64 tensor") + edge_count = src.numel() + if edge_count <= 0: + raise ValueError("direct split Phase A requires E > 0") + if tuple(d_full.shape) != (edge_count, PACKED_VALUE_COUNT): + raise ValueError(f"d_full must have shape {(edge_count, PACKED_VALUE_COUNT)}") + if tuple(radial_compact.shape) != (edge_count, RADIAL_COMPACT): + raise ValueError( + f"radial_compact must have shape {(edge_count, RADIAL_COMPACT)}" + ) + if tuple(channel_basis.shape) != (FULL_CHANNELS,): + raise ValueError("channel_basis must have shape (64,)") + + float_tensors = (x_wide, d_full, radial_compact, channel_basis) + if any(t.dtype != torch.float32 or not t.is_cuda for t in float_tensors): + raise TypeError("all Phase-A floating-point operands must be CUDA float32") + if not src.is_cuda: + raise TypeError("src must be a CUDA tensor") + if any(t.device != x_wide.device for t in (*float_tensors[1:], src)): + raise ValueError("all Phase-A operands must share x_wide.device") + if any(not t.is_contiguous() for t in float_tensors): + raise ValueError("all Phase-A floating-point operands must be contiguous") + src_i32 = ( + src + if src.dtype == torch.int32 and src.is_contiguous() + else src.to(dtype=torch.int32).contiguous() + ) + return edge_count, src_i32 + + +def _allocate_state(edge_count: int, device: torch.device) -> NeoPersistentComplexState: + return NeoPersistentComplexState( + m0=torch.empty( + (FOCUS_COUNT, edge_count, M0_WIDTH), + dtype=torch.float32, + device=device, + ), + m1=torch.empty( + (FOCUS_COUNT, edge_count, M1_WIDTH), + dtype=torch.complex64, + device=device, + ), + ) + + +def run_neo_phase_a_persistent_complex_fp32( + *, + x_wide: torch.Tensor, + src: torch.Tensor, + d_full: torch.Tensor, + radial_compact: torch.Tensor, + channel_basis: torch.Tensor, + out: NeoPersistentComplexState | None = None, +) -> NeoPersistentComplexState: + """Write Phase A directly into the persistent stack's native split state.""" + edge_count, src_i32 = _validate_inputs( + x_wide, + src, + d_full, + radial_compact, + channel_basis, + ) + if out is None: + out = _allocate_state(edge_count, x_wide.device) + validate_neo_persistent_complex_state(out, name="out") + if out.edge_count != edge_count or out.m0.device != x_wide.device: + raise ValueError("out must have matching E and share x_wide.device") + + 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'") + device_index = x_wide.device.index + if device_index is None: + raise RuntimeError("direct split Phase A requires CUDA") + compute_capability = tuple(torch.cuda.get_device_capability(device_index)) + compiled = _compiled_producer(device_index, compute_capability) + m1_ri = torch.view_as_real(out.m1) + if not m1_ri.is_contiguous(): + raise ValueError("out.m1 must expose a contiguous interleaved real/imag view") + stream = cuda.CUstream(torch.cuda.current_stream(x_wide.device).cuda_stream) + compiled( + x_wide.view(x_wide.shape[0], 16 * FULL_CHANNELS), + src_i32, + d_full, + radial_compact, + channel_basis, + out.m0, + m1_ri, + stream=stream, + ) + return out diff --git a/deepmd/kernels/cute/neo/sm90_k1/phase_a_backward.py b/deepmd/kernels/cute/neo/sm90_k1/phase_a_backward.py new file mode 100644 index 0000000000..85a980799d --- /dev/null +++ b/deepmd/kernels/cute/neo/sm90_k1/phase_a_backward.py @@ -0,0 +1,918 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Exact split-complex input adjoint for the SM90 Neo Phase-A boundary. + +The incoming adjoint remains in the persistent stack's focus-major contract: +``m0`` is float32 ``(2,E,128)`` and ``m1`` is complex64 ``(2,E,96)``. One +CTA owns one source node and reads those panels directly while recomputing the +packed-Wigner rotation. It emits the exact input adjoints for node features, +packed Wigner values, compact radial maps, and the rank-1 channel basis without +ever reconstructing an ``(E,2,10,32)`` block-real gradient slab. + +Source CSR is a caller-owned edge-cache property. It preserves physical edge +order while giving each node exclusive ownership of its feature adjoint, so no +atomics or ``(E,16,64)`` edge-local reduction tensor is required. +""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, +) + +import cutlass +import cutlass.cute as cute +import cutlass.utils as cute_utils +import torch +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) +from torch import ( + Tensor, +) + +from ..compile_cache import ( + device_aware_lru_cache, +) +from ..k1_kernels.cute_neo_radial_phase_a_backward_node import ( + COMPACT_WIDTH, + DEGREE_COUNT, + FOCUS_COUNT, + FOCUS_HIDDEN, + GROUPS_PER_CTA, + GROUPS_PER_WARP, + HIDDEN, + PACKED_WIGNER_VALUES, + REDUCED_COUNT, + SHARED_ROW_PITCH, + WARP_REDUCTION_GROUP, + _grad_x_value, + _recompute_local_value, + _warp_owned_grad_compact, + _warp_owned_grad_d, +) +from .persistent import ( + M0_WIDTH, + M1_WIDTH, + NeoPersistentComplexState, + validate_neo_persistent_complex_state, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN001, ANN202, TC002 + +THREADS = HIDDEN +REDUCE_THREADS = 32 +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + +__all__ = [ + "NeoPhaseAPersistentComplexAdjoints", + "NeoPhaseAPersistentComplexBackwardWorkspace", + "allocate_neo_phase_a_persistent_complex_backward", + "run_neo_phase_a_persistent_complex_backward_fp32", +] + + +@dataclass(frozen=True) +class NeoPhaseAPersistentComplexAdjoints: + """Differentiable input adjoints for the strict-FP32 Phase-A producer.""" + + grad_x_wide: Tensor + grad_d_full: Tensor + grad_radial_compact: Tensor + grad_channel_basis: Tensor + + +@dataclass(frozen=True) +class NeoPhaseAPersistentComplexBackwardWorkspace: + """Small deterministic channel-basis reduction workspace.""" + + grad_basis_by_node: Tensor + + @property + def storage_bytes(self) -> int: + return self.grad_basis_by_node.numel() * self.grad_basis_by_node.element_size() + + +@dataclass(frozen=True) +class _BackwardParams: + grad_m0: cute.Tensor + grad_m1_ri: cute.Tensor + radial_compact: cute.Tensor + channel_basis: cute.Tensor + x_wide: cute.Tensor + source_order: cute.Tensor + source_ptr: cute.Tensor + d_full: cute.Tensor + grad_x_wide: cute.Tensor + grad_d_full: cute.Tensor + grad_radial_compact: cute.Tensor + grad_basis_by_node: cute.Tensor + + +@cute.jit +def _split_grad_value( + grad_m0, + grad_m1_ri, + edge, + reduced: cutlass.Constexpr[int], + channel, +): + focus = channel // FOCUS_HIDDEN + focus_channel = channel - focus * FOCUS_HIDDEN + if cutlass.const_expr(reduced < 4): + return grad_m0[ + focus, + edge, + reduced * FOCUS_HIDDEN + focus_channel, + ].to(cutlass.Float32) + if cutlass.const_expr(reduced < 7): + return grad_m1_ri[ + focus, + edge, + (reduced - 4) * FOCUS_HIDDEN + focus_channel, + 0, + ].to(cutlass.Float32) + return grad_m1_ri[ + focus, + edge, + (reduced - 7) * FOCUS_HIDDEN + focus_channel, + 1, + ].to(cutlass.Float32) + + +@cute.jit +def _radial_output_value( + local_values, + compact, + reduced: cutlass.Constexpr[int], + channel, +): + """Recompute the pre-basis radial output for one reduced row/channel.""" + value = cutlass.Float32(0.0) + if cutlass.const_expr(reduced < 4): + for input_row in cutlass.range_constexpr(4): + value += ( + compact[input_row * 4 + reduced] + * local_values[input_row * SHARED_ROW_PITCH + channel] + ) + elif cutlass.const_expr(reduced < 7): + output_row = reduced - 4 + for input_row in cutlass.range_constexpr(3): + value += ( + compact[16 + input_row * 3 + output_row] + * local_values[(4 + input_row) * SHARED_ROW_PITCH + channel] + ) + else: + output_row = reduced - 7 + for input_row in cutlass.range_constexpr(3): + value += ( + compact[16 + input_row * 3 + output_row] + * local_values[(7 + input_row) * SHARED_ROW_PITCH + channel] + ) + return value + + +@cute.jit +def _phase_a_split_backward_jit( + grad_m0: cute.Tensor, + grad_m1_ri: cute.Tensor, + radial_compact: cute.Tensor, + channel_basis: cute.Tensor, + x_wide: cute.Tensor, + source_order: cute.Tensor, + source_ptr: cute.Tensor, + d_full: cute.Tensor, + grad_x_wide: cute.Tensor, + grad_d_full: cute.Tensor, + grad_radial_compact: cute.Tensor, + grad_basis_by_node: cute.Tensor, + grad_channel_basis: cute.Tensor, + stream: CUstream, +): + params = _BackwardParams( + grad_m0=grad_m0, + grad_m1_ri=grad_m1_ri, + radial_compact=radial_compact, + channel_basis=channel_basis, + x_wide=x_wide, + source_order=source_order, + source_ptr=source_ptr, + d_full=d_full, + grad_x_wide=grad_x_wide, + grad_d_full=grad_d_full, + grad_radial_compact=grad_radial_compact, + grad_basis_by_node=grad_basis_by_node, + ) + _phase_a_split_backward_kernel(params).launch( + grid=[x_wide.shape[0], 1, 1], + block=[THREADS, 1, 1], + stream=stream, + ) + _reduce_channel_basis_kernel(grad_basis_by_node, grad_channel_basis).launch( + grid=[HIDDEN, 1, 1], + block=[REDUCE_THREADS, 1, 1], + stream=stream, + ) + + +@cute.kernel +def _phase_a_split_backward_kernel(params: _BackwardParams): + channel, _, _ = cute.arch.thread_idx() + node, _, _ = cute.arch.block_idx() + x_row_pitch = SHARED_ROW_PITCH + + smem = cute_utils.SmemAllocator() + x_values = smem.allocate_tensor( + cutlass.Float32, + DEGREE_COUNT * x_row_pitch, + ) + focus_grad = smem.allocate_tensor( + cutlass.Float32, + REDUCED_COUNT * SHARED_ROW_PITCH, + ) + # Primal local rows are overwritten by their adjoints after grad-radial and + # grad-basis have consumed them. + local_values = smem.allocate_tensor( + cutlass.Float32, + REDUCED_COUNT * SHARED_ROW_PITCH, + ) + d_values = smem.allocate_tensor(cutlass.Float32, PACKED_WIGNER_VALUES) + compact = smem.allocate_tensor(cutlass.Float32, COMPACT_WIDTH) + + for full_row in cutlass.range_constexpr(DEGREE_COUNT): + x_values[full_row * x_row_pitch + channel] = params.x_wide[ + node, + full_row * HIDDEN + channel, + ].to(cutlass.Float32) + cute.arch.sync_threads() + + grad_x_0 = cutlass.Float32(0.0) + grad_x_1 = cutlass.Float32(0.0) + grad_x_2 = cutlass.Float32(0.0) + grad_x_3 = cutlass.Float32(0.0) + grad_x_4 = cutlass.Float32(0.0) + grad_x_5 = cutlass.Float32(0.0) + grad_x_6 = cutlass.Float32(0.0) + grad_x_7 = cutlass.Float32(0.0) + grad_x_8 = cutlass.Float32(0.0) + grad_x_9 = cutlass.Float32(0.0) + grad_x_10 = cutlass.Float32(0.0) + grad_x_11 = cutlass.Float32(0.0) + grad_x_12 = cutlass.Float32(0.0) + grad_x_13 = cutlass.Float32(0.0) + grad_x_14 = cutlass.Float32(0.0) + grad_x_15 = cutlass.Float32(0.0) + grad_basis = cutlass.Float32(0.0) + + lo = params.source_ptr[node] + hi = params.source_ptr[node + 1] + for slot in cutlass.range(lo, hi, 1, unroll=1): + edge = params.source_order[slot] + for reduced in cutlass.range_constexpr(REDUCED_COUNT): + focus_grad[reduced * SHARED_ROW_PITCH + channel] = _split_grad_value( + params.grad_m0, + params.grad_m1_ri, + edge, + reduced, + channel, + ) + if channel < COMPACT_WIDTH: + compact[channel] = params.radial_compact[edge, channel].to(cutlass.Float32) + if channel < PACKED_WIGNER_VALUES: + d_values[channel] = params.d_full[edge, channel].to(cutlass.Float32) + cute.arch.sync_threads() + + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 0, + 0, + 0, + 1, + SHARED_ROW_PITCH, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 1, + 1, + 1, + 3, + SHARED_ROW_PITCH, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 2, + 10, + 4, + 5, + SHARED_ROW_PITCH, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 3, + 25, + 9, + 7, + SHARED_ROW_PITCH, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 4, + 4, + 1, + 3, + SHARED_ROW_PITCH, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 5, + 15, + 4, + 5, + SHARED_ROW_PITCH, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 6, + 32, + 9, + 7, + SHARED_ROW_PITCH, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 7, + 7, + 1, + 3, + SHARED_ROW_PITCH, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 8, + 20, + 4, + 5, + SHARED_ROW_PITCH, + x_row_pitch, + ) + _recompute_local_value( + local_values, + x_values, + d_values, + channel, + 9, + 39, + 9, + 7, + SHARED_ROW_PITCH, + x_row_pitch, + ) + cute.arch.sync_threads() + + for reduced in cutlass.range_constexpr(REDUCED_COUNT): + grad_basis += focus_grad[ + reduced * SHARED_ROW_PITCH + channel + ] * _radial_output_value( + local_values, + compact, + reduced, + channel, + ) + + lane = channel % 32 + warp = channel // 32 + subgroup = lane // WARP_REDUCTION_GROUP + subgroup_lane = lane % WARP_REDUCTION_GROUP + group = warp * GROUPS_PER_WARP + subgroup + for batch in cutlass.range_constexpr( + (COMPACT_WIDTH + GROUPS_PER_CTA - 1) // GROUPS_PER_CTA + ): + compact_idx = batch * GROUPS_PER_CTA + group + safe_compact_idx = compact_idx + if compact_idx >= COMPACT_WIDTH: + safe_compact_idx = cutlass.Int32(0) + grad_value = _warp_owned_grad_compact( + focus_grad, + local_values, + params.channel_basis, + safe_compact_idx, + subgroup_lane, + SHARED_ROW_PITCH, + ) + if compact_idx < COMPACT_WIDTH: + if subgroup_lane == 0: + params.grad_radial_compact[edge, compact_idx] = grad_value.to( + params.grad_radial_compact.element_type + ) + # Every warp reads the full shared primal panel while reducing its + # assigned compact columns. Do not let an early warp reuse that panel + # for input adjoints until all compact reductions have finished. + cute.arch.sync_threads() + + basis = params.channel_basis[channel].to(cutlass.Float32) + for reduced in cutlass.range_constexpr(REDUCED_COUNT): + grad_value = cutlass.Float32(0.0) + if reduced < 4: + for output_row in cutlass.range_constexpr(4): + grad_value += ( + focus_grad[output_row * SHARED_ROW_PITCH + channel] + * compact[reduced * 4 + output_row] + ) + elif reduced < 7: + input_row = reduced - 4 + for output_row in cutlass.range_constexpr(3): + grad_value += ( + focus_grad[(4 + output_row) * SHARED_ROW_PITCH + channel] + * compact[16 + input_row * 3 + output_row] + ) + else: + input_row = reduced - 7 + for output_row in cutlass.range_constexpr(3): + grad_value += ( + focus_grad[(7 + output_row) * SHARED_ROW_PITCH + channel] + * compact[16 + input_row * 3 + output_row] + ) + local_values[reduced * SHARED_ROW_PITCH + channel] = grad_value * basis + cute.arch.sync_threads() + + for batch in cutlass.range_constexpr( + (PACKED_WIGNER_VALUES + GROUPS_PER_CTA - 1) // GROUPS_PER_CTA + ): + panel_idx = batch * GROUPS_PER_CTA + group + safe_panel_idx = panel_idx + if panel_idx >= PACKED_WIGNER_VALUES: + safe_panel_idx = cutlass.Int32(0) + grad_d_value = _warp_owned_grad_d( + local_values, + x_values, + safe_panel_idx, + subgroup_lane, + SHARED_ROW_PITCH, + x_row_pitch, + ) + if panel_idx < PACKED_WIGNER_VALUES: + if subgroup_lane == 0: + params.grad_d_full[edge, panel_idx] = grad_d_value.to( + params.grad_d_full.element_type + ) + + grad_x_0 += _grad_x_value( + local_values, + d_values, + channel, + 0, + 0, + 0, + 1, + 1, + SHARED_ROW_PITCH, + ) + grad_x_1 += _grad_x_value( + local_values, + d_values, + channel, + 1, + 0, + 1, + 3, + 3, + SHARED_ROW_PITCH, + ) + grad_x_2 += _grad_x_value( + local_values, + d_values, + channel, + 1, + 1, + 1, + 3, + 3, + SHARED_ROW_PITCH, + ) + grad_x_3 += _grad_x_value( + local_values, + d_values, + channel, + 1, + 2, + 1, + 3, + 3, + SHARED_ROW_PITCH, + ) + grad_x_4 += _grad_x_value( + local_values, + d_values, + channel, + 2, + 0, + 10, + 5, + 3, + SHARED_ROW_PITCH, + ) + grad_x_5 += _grad_x_value( + local_values, + d_values, + channel, + 2, + 1, + 10, + 5, + 3, + SHARED_ROW_PITCH, + ) + grad_x_6 += _grad_x_value( + local_values, + d_values, + channel, + 2, + 2, + 10, + 5, + 3, + SHARED_ROW_PITCH, + ) + grad_x_7 += _grad_x_value( + local_values, + d_values, + channel, + 2, + 3, + 10, + 5, + 3, + SHARED_ROW_PITCH, + ) + grad_x_8 += _grad_x_value( + local_values, + d_values, + channel, + 2, + 4, + 10, + 5, + 3, + SHARED_ROW_PITCH, + ) + grad_x_9 += _grad_x_value( + local_values, + d_values, + channel, + 3, + 0, + 25, + 7, + 3, + SHARED_ROW_PITCH, + ) + grad_x_10 += _grad_x_value( + local_values, + d_values, + channel, + 3, + 1, + 25, + 7, + 3, + SHARED_ROW_PITCH, + ) + grad_x_11 += _grad_x_value( + local_values, + d_values, + channel, + 3, + 2, + 25, + 7, + 3, + SHARED_ROW_PITCH, + ) + grad_x_12 += _grad_x_value( + local_values, + d_values, + channel, + 3, + 3, + 25, + 7, + 3, + SHARED_ROW_PITCH, + ) + grad_x_13 += _grad_x_value( + local_values, + d_values, + channel, + 3, + 4, + 25, + 7, + 3, + SHARED_ROW_PITCH, + ) + grad_x_14 += _grad_x_value( + local_values, + d_values, + channel, + 3, + 5, + 25, + 7, + 3, + SHARED_ROW_PITCH, + ) + grad_x_15 += _grad_x_value( + local_values, + d_values, + channel, + 3, + 6, + 25, + 7, + 3, + SHARED_ROW_PITCH, + ) + cute.arch.sync_threads() + + params.grad_x_wide[node, 0 * HIDDEN + channel] = grad_x_0 + params.grad_x_wide[node, 1 * HIDDEN + channel] = grad_x_1 + params.grad_x_wide[node, 2 * HIDDEN + channel] = grad_x_2 + params.grad_x_wide[node, 3 * HIDDEN + channel] = grad_x_3 + params.grad_x_wide[node, 4 * HIDDEN + channel] = grad_x_4 + params.grad_x_wide[node, 5 * HIDDEN + channel] = grad_x_5 + params.grad_x_wide[node, 6 * HIDDEN + channel] = grad_x_6 + params.grad_x_wide[node, 7 * HIDDEN + channel] = grad_x_7 + params.grad_x_wide[node, 8 * HIDDEN + channel] = grad_x_8 + params.grad_x_wide[node, 9 * HIDDEN + channel] = grad_x_9 + params.grad_x_wide[node, 10 * HIDDEN + channel] = grad_x_10 + params.grad_x_wide[node, 11 * HIDDEN + channel] = grad_x_11 + params.grad_x_wide[node, 12 * HIDDEN + channel] = grad_x_12 + params.grad_x_wide[node, 13 * HIDDEN + channel] = grad_x_13 + params.grad_x_wide[node, 14 * HIDDEN + channel] = grad_x_14 + params.grad_x_wide[node, 15 * HIDDEN + channel] = grad_x_15 + params.grad_basis_by_node[node, channel] = grad_basis + + +@cute.kernel +def _reduce_channel_basis_kernel( + grad_basis_by_node: cute.Tensor, + grad_channel_basis: cute.Tensor, +): + tidx, _, _ = cute.arch.thread_idx() + channel, _, _ = cute.arch.block_idx() + value = cutlass.Float32(0.0) + for node in cutlass.range( + tidx, + grad_basis_by_node.shape[0], + REDUCE_THREADS, + unroll=1, + ): + value += grad_basis_by_node[node, channel].to(cutlass.Float32) + value = cute.arch.warp_reduction_sum(value) + if tidx == 0: + grad_channel_basis[channel] = value + + +def _fake(dtype, shape: tuple, stride_order: tuple[int, ...]): + return make_fake_compact_tensor( + dtype, + shape, + stride_order=stride_order, + **FAKE_TENSOR_KW, + ) + + +@device_aware_lru_cache(maxsize=8) +def _compiled_backward( + device_index: int, + compute_capability: tuple[int, int], +) -> Callable: + if compute_capability != (9, 0): + raise RuntimeError("the direct split Phase-A adjoint is sm_90-only") + edge_count = cute.sym_int64() + node_count = cute.sym_int64() + source_ptr_count = cute.sym_int64() + with torch.cuda.device(device_index): + return cute.compile( + _phase_a_split_backward_jit, + _fake(cutlass.Float32, (FOCUS_COUNT, edge_count, M0_WIDTH), (2, 1, 0)), + _fake( + cutlass.Float32, + (FOCUS_COUNT, edge_count, M1_WIDTH, 2), + (3, 2, 1, 0), + ), + _fake(cutlass.Float32, (edge_count, COMPACT_WIDTH), (1, 0)), + _fake(cutlass.Float32, (HIDDEN,), (0,)), + _fake(cutlass.Float32, (node_count, DEGREE_COUNT * HIDDEN), (1, 0)), + _fake(cutlass.Int32, (edge_count,), (0,)), + _fake(cutlass.Int32, (source_ptr_count,), (0,)), + _fake(cutlass.Float32, (edge_count, PACKED_WIGNER_VALUES), (1, 0)), + _fake(cutlass.Float32, (node_count, DEGREE_COUNT * HIDDEN), (1, 0)), + _fake(cutlass.Float32, (edge_count, PACKED_WIGNER_VALUES), (1, 0)), + _fake(cutlass.Float32, (edge_count, COMPACT_WIDTH), (1, 0)), + _fake(cutlass.Float32, (node_count, HIDDEN), (1, 0)), + _fake(cutlass.Float32, (HIDDEN,), (0,)), + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +def allocate_neo_phase_a_persistent_complex_backward( + *, + edge_count: int, + node_count: int, + device: torch.device, +) -> tuple[ + NeoPhaseAPersistentComplexAdjoints, + NeoPhaseAPersistentComplexBackwardWorkspace, +]: + """Allocate caller-owned outputs and the bounded basis-reduction workspace.""" + if edge_count <= 0 or node_count <= 0: + raise ValueError("the Phase-A adjoint requires N > 0 and E > 0") + opts = {"device": device, "dtype": torch.float32} + outputs = NeoPhaseAPersistentComplexAdjoints( + grad_x_wide=torch.empty( + (node_count, DEGREE_COUNT, HIDDEN), + device=opts["device"], + dtype=opts["dtype"], + ), + grad_d_full=torch.empty( + (edge_count, PACKED_WIGNER_VALUES), + device=opts["device"], + dtype=opts["dtype"], + ), + grad_radial_compact=torch.empty( + (edge_count, COMPACT_WIDTH), + device=opts["device"], + dtype=opts["dtype"], + ), + grad_channel_basis=torch.empty( + (HIDDEN,), + device=opts["device"], + dtype=opts["dtype"], + ), + ) + workspace = NeoPhaseAPersistentComplexBackwardWorkspace( + grad_basis_by_node=torch.empty( + (node_count, HIDDEN), + device=opts["device"], + dtype=opts["dtype"], + ) + ) + return outputs, workspace + + +def _expect( + name: str, + tensor: Tensor, + shape: tuple[int, ...], + device: torch.device, + dtype: torch.dtype, +) -> None: + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}") + if tensor.device != device or tensor.dtype != dtype or not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous {dtype} on {device}") + + +def run_neo_phase_a_persistent_complex_backward_fp32( + *, + grad_state: NeoPersistentComplexState, + radial_compact: Tensor, + channel_basis: Tensor, + x_wide: Tensor, + source_order: Tensor, + source_ptr: Tensor, + d_full: Tensor, + outputs: NeoPhaseAPersistentComplexAdjoints | None = None, + workspace: NeoPhaseAPersistentComplexBackwardWorkspace | None = None, +) -> NeoPhaseAPersistentComplexAdjoints: + """Run the direct split-gradient Phase-A input adjoint.""" + if torch.is_grad_enabled(): + raise RuntimeError("the explicit Phase-A adjoint must run under no_grad") + validate_neo_persistent_complex_state(grad_state, name="grad_state") + device = grad_state.m0.device + edge_count = grad_state.edge_count + if x_wide.ndim != 3 or tuple(x_wide.shape[1:]) != (DEGREE_COUNT, HIDDEN): + raise ValueError("x_wide must have shape (N,16,64)") + node_count = x_wide.shape[0] + if node_count <= 0: + raise ValueError("the Phase-A adjoint requires N > 0") + + specs = ( + ("radial_compact", radial_compact, (edge_count, COMPACT_WIDTH), torch.float32), + ("channel_basis", channel_basis, (HIDDEN,), torch.float32), + ("x_wide", x_wide, (node_count, DEGREE_COUNT, HIDDEN), torch.float32), + ("source_order", source_order, (edge_count,), torch.int32), + ("source_ptr", source_ptr, (node_count + 1,), torch.int32), + ("d_full", d_full, (edge_count, PACKED_WIGNER_VALUES), torch.float32), + ) + for name, tensor, shape, dtype in specs: + _expect(name, tensor, shape, device, dtype) + + if outputs is None or workspace is None: + if outputs is not None or workspace is not None: + raise ValueError("outputs and workspace must be supplied together") + outputs, workspace = allocate_neo_phase_a_persistent_complex_backward( + edge_count=edge_count, + node_count=node_count, + device=device, + ) + output_specs = ( + ("grad_x_wide", outputs.grad_x_wide, (node_count, DEGREE_COUNT, HIDDEN)), + ("grad_d_full", outputs.grad_d_full, (edge_count, PACKED_WIGNER_VALUES)), + ( + "grad_radial_compact", + outputs.grad_radial_compact, + (edge_count, COMPACT_WIDTH), + ), + ("grad_channel_basis", outputs.grad_channel_basis, (HIDDEN,)), + ( + "grad_basis_by_node", + workspace.grad_basis_by_node, + (node_count, HIDDEN), + ), + ) + for name, tensor, shape in output_specs: + _expect(name, tensor, shape, device, torch.float32) + + 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'") + device_index = device.index + if device_index is None: + raise RuntimeError("the direct split Phase-A adjoint requires CUDA") + compute_capability = tuple(torch.cuda.get_device_capability(device_index)) + kernel = _compiled_backward(device_index, compute_capability) + grad_m1_ri = torch.view_as_real(grad_state.m1) + if not grad_m1_ri.is_contiguous(): + raise ValueError("grad_state.m1 must expose interleaved real/imag storage") + with torch.cuda.device(device): + kernel( + grad_state.m0, + grad_m1_ri, + radial_compact, + channel_basis, + x_wide.view(node_count, DEGREE_COUNT * HIDDEN), + source_order, + source_ptr, + d_full, + outputs.grad_x_wide.view(node_count, DEGREE_COUNT * HIDDEN), + outputs.grad_d_full, + outputs.grad_radial_compact, + workspace.grad_basis_by_node, + outputs.grad_channel_basis, + ) + return outputs diff --git a/deepmd/kernels/cute/neo/sm90_k1/phase_c_attention_backward.py b/deepmd/kernels/cute/neo/sm90_k1/phase_c_attention_backward.py new file mode 100644 index 0000000000..b51f0f40af --- /dev/null +++ b/deepmd/kernels/cute/neo/sm90_k1/phase_c_attention_backward.py @@ -0,0 +1,838 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Grouped SM90 final Phase-C/attention adjoint with split-state gather. + +Four independent 64-thread groups process four CSR edges +per round while sharing one destination node's reverse-linear ``b0/b1`` +panels. Each ``(edge, focus)`` warp also retains the expanded +``grad_m0`` and complex ``grad_m1`` accumulators in registers. The same panel +loads therefore serve both the Phase-C scale contraction and split-state +adjoint, replacing the separate ``_expanded_adjoint_gather`` launch. + +The focus competition, segmented attention softmax, envelope, and Q/K +adjoints remain fused. No edge-sized temporary is added. +Source-K is atomically accumulated, so callers must clear ``grad_k_node`` +before every invocation. +""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, +) + +import cutlass +import cutlass.cute as cute +import torch +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from ..compile_cache import ( + device_aware_lru_cache, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN001, ANN202, TC002 + +FOCUS_COUNT = 2 +CHANNELS = 32 +DEGREE_COUNT = 16 +M0_WIDTH = 128 +M1_WIDTH = 96 +PACKED_WIGNER_VALUES = 46 +MAX_EDGES_PER_NODE = 128 +QK_SCALE = CHANNELS**-0.5 +B0_NODE_VALUES = FOCUS_COUNT * DEGREE_COUNT * M0_WIDTH +B1_NODE_VALUES = FOCUS_COUNT * (DEGREE_COUNT - 1) * M1_WIDTH * 2 +PACKED_M0 = (0, 1, 2, 3, 10, 11, 12, 13, 14, 25, 26, 27, 28, 29, 30, 31) +PACKED_RE = (4, 5, 6, 15, 16, 17, 18, 19, 32, 33, 34, 35, 36, 37, 38) +PACKED_IM = (7, 8, 9, 20, 21, 22, 23, 24, 39, 40, 41, 42, 43, 44, 45) +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + +EDGE_GROUPS = 4 +GROUP_WIDTH = FOCUS_COUNT * CHANNELS +THREADS = EDGE_GROUPS * GROUP_WIDTH +M0_ROWS = M0_WIDTH // CHANNELS +M1_ROWS = M1_WIDTH // CHANNELS + + +@cute.jit +def _warp_sum(value): + return cute.arch.warp_reduction_sum(value) + + +def _fake(dtype, shape: tuple[object, ...], stride_order: tuple[int, ...]): + return make_fake_compact_tensor( + dtype, + shape, + stride_order=stride_order, + **FAKE_TENSOR_KW, + ) + + +def _require_tensor( + name: str, + tensor: torch.Tensor, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> None: + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}") + if ( + tensor.dtype != dtype + or tensor.device != device + or not tensor.is_cuda + or not tensor.is_contiguous() + or tensor.data_ptr() % 16 + ): + raise ValueError( + f"{name} must be contiguous, 16-byte-aligned {dtype} on {device}" + ) + + +__all__ = [ + "GroupedExpandedFinalPhaseCAttentionAdjointOutputs", + "allocate_grouped_expanded_final_phase_c_attention_adjoint_outputs", + "compile_grouped_expanded_final_phase_c_attention_adjoint", + "run_grouped_expanded_final_phase_c_attention_adjoint", +] + + +@dataclass(frozen=True) +class GroupedExpandedFinalPhaseCAttentionAdjointOutputs: + """True adjoints emitted by the one-launch boundary.""" + + grad_m0: torch.Tensor + grad_m1: torch.Tensor + grad_dt: torch.Tensor + grad_logits: torch.Tensor + grad_edge: torch.Tensor + grad_focus_src: torch.Tensor + grad_q_node: torch.Tensor + grad_k_node: torch.Tensor + + +@cute.jit +def _grouped_expanded_adjoint_jit( + b0: cute.Tensor, + b1_ri: cute.Tensor, + m0: cute.Tensor, + m1_ri: cute.Tensor, + dt_packed: cute.Tensor, + beta: cute.Tensor, + alpha: cute.Tensor, + focus_alpha: cute.Tensor, + focus_src: cute.Tensor, + focus_weight: cute.Tensor, + focus_scale: cute.Tensor, + q_node: cute.Tensor, + k_node: cute.Tensor, + edge_gate: cute.Tensor, + src: cute.Tensor, + dst_ptr: cute.Tensor, + grad_m0: cute.Tensor, + grad_m1_ri: cute.Tensor, + grad_dt: cute.Tensor, + grad_logits: cute.Tensor, + grad_edge: cute.Tensor, + grad_focus_src: cute.Tensor, + grad_q_node: cute.Tensor, + grad_k_node: cute.Tensor, + focus_eps: cutlass.Constexpr[float], + focus_tau: cutlass.Constexpr[float], + label_smoothing: cutlass.Constexpr[float], + qk_scale: cutlass.Constexpr[float], + stream: CUstream, +): + _grouped_expanded_adjoint_kernel( + b0, + b1_ri, + m0, + m1_ri, + dt_packed, + beta, + alpha, + focus_alpha, + focus_src, + focus_weight, + focus_scale, + q_node, + k_node, + edge_gate, + src, + dst_ptr, + grad_m0, + grad_m1_ri, + grad_dt, + grad_logits, + grad_edge, + grad_focus_src, + grad_q_node, + grad_k_node, + focus_eps, + focus_tau, + label_smoothing, + qk_scale, + ).launch( + grid=[dst_ptr.shape[0] - 1, 1, 1], + block=[THREADS, 1, 1], + stream=stream, + ) + + +@cute.kernel +def _grouped_expanded_adjoint_kernel( + b0: cute.Tensor, + b1_ri: cute.Tensor, + m0: cute.Tensor, + m1_ri: cute.Tensor, + dt_packed: cute.Tensor, + beta: cute.Tensor, + alpha: cute.Tensor, + focus_alpha: cute.Tensor, + focus_src: cute.Tensor, + focus_weight: cute.Tensor, + focus_scale: cute.Tensor, + q_node: cute.Tensor, + k_node: cute.Tensor, + edge_gate: cute.Tensor, + src: cute.Tensor, + dst_ptr: cute.Tensor, + grad_m0: cute.Tensor, + grad_m1_ri: cute.Tensor, + grad_dt: cute.Tensor, + grad_logits: cute.Tensor, + grad_edge: cute.Tensor, + grad_focus_src: cute.Tensor, + grad_q_node: cute.Tensor, + grad_k_node: cute.Tensor, + focus_eps: cutlass.Constexpr[float], + focus_tau: cutlass.Constexpr[float], + label_smoothing: cutlass.Constexpr[float], + qk_scale: cutlass.Constexpr[float], +): + tidx, _, _ = cute.arch.thread_idx() + node, _, _ = cute.arch.block_idx() + group = tidx // GROUP_WIDTH + local = tidx - group * GROUP_WIDTH + focus = local // CHANNELS + lane = local - focus * CHANNELS + node_lo = dst_ptr[node] + node_hi = dst_ptr[node + 1] + degree = node_hi - node_lo + rounds = (degree + EDGE_GROUPS - 1) // EDGE_GROUPS + + smem = cutlass.utils.SmemAllocator() + b0_node = smem.allocate_tensor(cutlass.Float32, B0_NODE_VALUES) + b1_node = smem.allocate_tensor(cutlass.Float32, B1_NODE_VALUES) + beta_adjoint = smem.allocate_tensor( + cutlass.Float32, + MAX_EDGES_PER_NODE * FOCUS_COUNT, + ) + dt_partial = smem.allocate_tensor( + cutlass.Float32, + EDGE_GROUPS * FOCUS_COUNT * PACKED_WIGNER_VALUES, + ) + softmax_dot = smem.allocate_tensor(cutlass.Float32, FOCUS_COUNT) + q_partial = smem.allocate_tensor( + cutlass.Float32, + EDGE_GROUPS * FOCUS_COUNT * CHANNELS, + ) + + for linear in cutlass.range(tidx, B0_NODE_VALUES, THREADS, unroll=1): + quotient = linear // M0_WIDTH + feature = linear - quotient * M0_WIDTH + panel_focus = quotient // DEGREE_COUNT + q = quotient - panel_focus * DEGREE_COUNT + b0_node[linear] = b0[panel_focus, q, node, feature].to(cutlass.Float32) + for linear in cutlass.range(tidx, B1_NODE_VALUES, THREADS, unroll=1): + quotient = linear // 2 + component = linear - quotient * 2 + feature_quotient = quotient // M1_WIDTH + feature = quotient - feature_quotient * M1_WIDTH + panel_focus = feature_quotient // (DEGREE_COUNT - 1) + q1 = feature_quotient - panel_focus * (DEGREE_COUNT - 1) + b1_node[linear] = b1_ri[ + panel_focus, + q1, + node, + feature, + component, + ].to(cutlass.Float32) + cute.arch.sync_threads() + + m0_fragment = cute.make_rmem_tensor( + cute.make_layout((M0_ROWS,), stride=(1,)), + cutlass.Float32, + ) + m1_real_fragment = cute.make_rmem_tensor( + cute.make_layout((M1_ROWS,), stride=(1,)), + cutlass.Float32, + ) + m1_imag_fragment = cute.make_rmem_tensor( + cute.make_layout((M1_ROWS,), stride=(1,)), + cutlass.Float32, + ) + grad_m0_fragment = cute.make_rmem_tensor( + cute.make_layout((M0_ROWS,), stride=(1,)), + cutlass.Float32, + ) + grad_m1_real_fragment = cute.make_rmem_tensor( + cute.make_layout((M1_ROWS,), stride=(1,)), + cutlass.Float32, + ) + grad_m1_imag_fragment = cute.make_rmem_tensor( + cute.make_layout((M1_ROWS,), stride=(1,)), + cutlass.Float32, + ) + + # All four groups execute the same round count, making CTA barriers safe. + for edge_round in cutlass.range(rounds, unroll=1): + edge_slot = edge_round * EDGE_GROUPS + group + edge = node_lo + edge_slot + edge_active = edge_slot < degree + + for row in cutlass.range_constexpr(M0_ROWS): + feature = row * CHANNELS + lane + value = cutlass.Float32(0.0) + if edge_active: + value = m0[focus, edge, feature].to(cutlass.Float32) + m0_fragment[row] = value + grad_m0_fragment[row] = cutlass.Float32(0.0) + for row in cutlass.range_constexpr(M1_ROWS): + feature = row * CHANNELS + lane + real = cutlass.Float32(0.0) + imag = cutlass.Float32(0.0) + if edge_active: + real = m1_ri[focus, edge, feature, 0].to(cutlass.Float32) + imag = m1_ri[focus, edge, feature, 1].to(cutlass.Float32) + m1_real_fragment[row] = real + m1_imag_fragment[row] = imag + grad_m1_real_fragment[row] = cutlass.Float32(0.0) + grad_m1_imag_fragment[row] = cutlass.Float32(0.0) + + beta_value = cutlass.Float32(0.0) + if edge_active: + beta_value = beta[edge, focus].to(cutlass.Float32) + grad_beta_value = cutlass.Float32(0.0) + dt_base = (group * FOCUS_COUNT + focus) * PACKED_WIGNER_VALUES + + # Ascending q order matches the existing expanded gather exactly. + for q in cutlass.range_constexpr(DEGREE_COUNT): + panel0 = PACKED_M0[q] + dt0 = cutlass.Float32(0.0) + if edge_active: + dt0 = dt_packed[edge, panel0].to(cutlass.Float32) + scalar0_lane = cutlass.Float32(0.0) + for row in cutlass.range_constexpr(M0_ROWS): + feature = row * CHANNELS + lane + b0_offset = (focus * DEGREE_COUNT + q) * M0_WIDTH + feature + b_value = b0_node[b0_offset] + grad_m0_fragment[row] += beta_value * dt0 * b_value + scalar0_lane += b_value * m0_fragment[row] + scalar0 = _warp_sum(scalar0_lane) + if lane == 0: + dt_partial[dt_base + panel0] = beta_value * scalar0 + grad_beta_value += dt0 * scalar0 + + if cutlass.const_expr(q > 0): + panel_re = PACKED_RE[q - 1] + panel_im = PACKED_IM[q - 1] + dt_re = cutlass.Float32(0.0) + dt_im = cutlass.Float32(0.0) + if edge_active: + dt_re = dt_packed[edge, panel_re].to(cutlass.Float32) + dt_im = dt_packed[edge, panel_im].to(cutlass.Float32) + scalar1_re_lane = cutlass.Float32(0.0) + scalar1_im_lane = cutlass.Float32(0.0) + for row in cutlass.range_constexpr(M1_ROWS): + feature = row * CHANNELS + lane + b1_offset = ( + (focus * (DEGREE_COUNT - 1) + q - 1) * M1_WIDTH + feature + ) * 2 + br = b1_node[b1_offset] + bi = b1_node[b1_offset + 1] + grad_m1_real_fragment[row] += beta_value * (dt_re * br - dt_im * bi) + grad_m1_imag_fragment[row] += beta_value * (dt_re * bi + dt_im * br) + xr = m1_real_fragment[row] + xi = m1_imag_fragment[row] + scalar1_re_lane += br * xr + bi * xi + scalar1_im_lane += br * xi - bi * xr + scalar1_re = _warp_sum(scalar1_re_lane) + scalar1_im = _warp_sum(scalar1_im_lane) + if lane == 0: + dt_partial[dt_base + panel_re] = beta_value * scalar1_re + dt_partial[dt_base + panel_im] = beta_value * scalar1_im + grad_beta_value += dt_re * scalar1_re + dt_im * scalar1_im + + if edge_active: + for row in cutlass.range_constexpr(M0_ROWS): + feature = row * CHANNELS + lane + grad_m0[focus, edge, feature] = grad_m0_fragment[row] + for row in cutlass.range_constexpr(M1_ROWS): + feature = row * CHANNELS + lane + grad_m1_ri[focus, edge, feature, 0] = grad_m1_real_fragment[row] + grad_m1_ri[focus, edge, feature, 1] = grad_m1_imag_fragment[row] + if lane == 0: + beta_adjoint[edge_slot * FOCUS_COUNT + focus] = grad_beta_value + cute.arch.sync_threads() + + if tidx < EDGE_GROUPS * PACKED_WIGNER_VALUES: + output_group = tidx // PACKED_WIGNER_VALUES + panel = tidx - output_group * PACKED_WIGNER_VALUES + output_slot = edge_round * EDGE_GROUPS + output_group + if output_slot < degree: + output_edge = node_lo + output_slot + focus0 = output_group * FOCUS_COUNT * PACKED_WIGNER_VALUES + grad_dt[output_edge, panel] = ( + dt_partial[focus0 + panel] + + dt_partial[focus0 + PACKED_WIGNER_VALUES + panel] + ) + cute.arch.sync_threads() + + # Match the forward path's ascending-edge segmented-softmax reduction order. + if group == 0 and lane == 0: + value = cutlass.Float32(0.0) + for edge_slot in cutlass.range(degree, unroll=1): + edge = node_lo + edge_slot + grad_attention = beta_adjoint[ + edge_slot * FOCUS_COUNT + focus + ] * focus_alpha[edge, focus].to(cutlass.Float32) + value += alpha[edge, focus].to(cutlass.Float32) * grad_attention + softmax_dot[focus] = value + cute.arch.sync_threads() + + keep = cutlass.Float32(1.0 - label_smoothing) + smooth = cutlass.Float32(label_smoothing / FOCUS_COUNT) + inv_tau = cutlass.Float32(1.0 / focus_tau) + grad_q = cutlass.Float32(0.0) + + for edge_slot in cutlass.range(group, degree, EDGE_GROUPS, unroll=1): + edge = node_lo + edge_slot + source = src[edge] + grad_beta_value = beta_adjoint[edge_slot * FOCUS_COUNT + focus] + alpha_value = alpha[edge, focus].to(cutlass.Float32) + focus_value = focus_alpha[edge, focus].to(cutlass.Float32) + grad_attention = grad_beta_value * focus_value + grad_logit = alpha_value * (grad_attention - softmax_dot[focus]) + if lane == 0: + grad_logits[edge, focus] = grad_logit + + scaled_grad_logit = grad_logit * cutlass.Float32(qk_scale) + grad_q += scaled_grad_logit * k_node[source, focus, lane].to(cutlass.Float32) + grad_k = scaled_grad_logit * q_node[node, focus, lane].to(cutlass.Float32) + k_offset = (source * FOCUS_COUNT + focus) * CHANNELS + lane + k_ptr = grad_k_node.iterator + k_offset + cute.arch.atomic_add( + k_ptr.llvm_ptr, + grad_k, + sem="relaxed", + scope="gpu", + ) + + probability0 = (focus_alpha[edge, 0].to(cutlass.Float32) - smooth) / keep + probability1 = (focus_alpha[edge, 1].to(cutlass.Float32) - smooth) / keep + focus_grad0 = ( + beta_adjoint[edge_slot * FOCUS_COUNT] + * alpha[edge, 0].to(cutlass.Float32) + * keep + ) + focus_grad1 = ( + beta_adjoint[edge_slot * FOCUS_COUNT + 1] + * alpha[edge, 1].to(cutlass.Float32) + * keep + ) + focus_dot = focus_grad0 * probability0 + focus_grad1 * probability1 + probability = probability0 + focus_grad_probability = focus_grad0 + if focus == 1: + probability = probability1 + focus_grad_probability = focus_grad1 + focus_grad_logit = probability * (focus_grad_probability - focus_dot) * inv_tau + + focus_value_raw = focus_src[edge, focus, lane].to(cutlass.Float32) + scale = focus_scale[focus, lane].to(cutlass.Float32) + weight = focus_weight[lane, focus].to(cutlass.Float32) + inv_rms = cute.rsqrt( + _warp_sum(focus_value_raw * focus_value_raw) / cutlass.Float32(CHANNELS) + + cutlass.Float32(focus_eps) + ) + grad_scaled = focus_grad_logit * weight * scale + coeff = _warp_sum(grad_scaled * focus_value_raw) / cutlass.Float32(CHANNELS) + grad_focus_src[focus, edge, lane] = ( + grad_scaled * inv_rms + - focus_value_raw * inv_rms * inv_rms * inv_rms * coeff + ) + + if local == 0: + grad_attention0 = beta_adjoint[edge_slot * FOCUS_COUNT] * focus_alpha[ + edge, 0 + ].to(cutlass.Float32) + grad_attention1 = beta_adjoint[edge_slot * FOCUS_COUNT + 1] * focus_alpha[ + edge, 1 + ].to(cutlass.Float32) + grad_logit0 = alpha[edge, 0].to(cutlass.Float32) * ( + grad_attention0 - softmax_dot[0] + ) + grad_logit1 = alpha[edge, 1].to(cutlass.Float32) * ( + grad_attention1 - softmax_dot[1] + ) + gate = edge_gate[edge].to(cutlass.Float32) + value = cutlass.Float32(0.0) + if gate > cutlass.Float32(0.0): + value = cutlass.Float32(2.0) * (grad_logit0 + grad_logit1) / gate + grad_edge[edge] = value + + q_partial[group * GROUP_WIDTH + local] = grad_q + cute.arch.sync_threads() + if group == 0: + total = cutlass.Float32(0.0) + for partial_group in cutlass.range_constexpr(EDGE_GROUPS): + total += q_partial[partial_group * GROUP_WIDTH + local] + grad_q_node[node, focus, lane] = total + + +@device_aware_lru_cache(maxsize=8) +def compile_grouped_expanded_final_phase_c_attention_adjoint( + focus_eps: float, + focus_tau: float, + label_smoothing: float, + qk_scale: float = QK_SCALE, +) -> Callable: + """Compile the fixed four-group SM90 schedule.""" + edges = cute.sym_int64() + nodes = cute.sym_int64() + b0 = _fake( + cutlass.Float32, + (FOCUS_COUNT, DEGREE_COUNT, nodes, M0_WIDTH), + (3, 2, 1, 0), + ) + b1 = _fake( + cutlass.Float32, + (FOCUS_COUNT, DEGREE_COUNT - 1, nodes, M1_WIDTH, 2), + (4, 3, 2, 1, 0), + ) + m0 = _fake(cutlass.Float32, (FOCUS_COUNT, edges, M0_WIDTH), (2, 1, 0)) + m1 = _fake( + cutlass.Float32, + (FOCUS_COUNT, edges, M1_WIDTH, 2), + (3, 2, 1, 0), + ) + dt = _fake(cutlass.Float32, (edges, PACKED_WIGNER_VALUES), (1, 0)) + edge_focus = _fake(cutlass.Float32, (edges, FOCUS_COUNT), (1, 0)) + focus_src = _fake( + cutlass.Float32, + (edges, FOCUS_COUNT, CHANNELS), + (2, 1, 0), + ) + focus_weight = _fake(cutlass.Float32, (CHANNELS, FOCUS_COUNT), (1, 0)) + focus_scale = _fake(cutlass.Float32, (FOCUS_COUNT, CHANNELS), (1, 0)) + node_focus = _fake( + cutlass.Float32, + (nodes, FOCUS_COUNT, CHANNELS), + (2, 1, 0), + ) + edge_scalar = _fake(cutlass.Float32, (edges,), (0,)) + edge_index = _fake(cutlass.Int32, (edges,), (0,)) + dst_ptr = _fake(cutlass.Int32, (cute.sym_int64(),), (0,)) + grad_focus_src = _fake( + cutlass.Float32, + (FOCUS_COUNT, edges, CHANNELS), + (2, 1, 0), + ) + return cute.compile( + _grouped_expanded_adjoint_jit, + b0, + b1, + m0, + m1, + dt, + edge_focus, + edge_focus, + edge_focus, + focus_src, + focus_weight, + focus_scale, + node_focus, + node_focus, + edge_scalar, + edge_index, + dst_ptr, + _fake(cutlass.Float32, (FOCUS_COUNT, edges, M0_WIDTH), (2, 1, 0)), + _fake( + cutlass.Float32, + (FOCUS_COUNT, edges, M1_WIDTH, 2), + (3, 2, 1, 0), + ), + _fake(cutlass.Float32, (edges, PACKED_WIGNER_VALUES), (1, 0)), + _fake(cutlass.Float32, (edges, FOCUS_COUNT), (1, 0)), + _fake(cutlass.Float32, (edges,), (0,)), + grad_focus_src, + _fake(cutlass.Float32, (nodes, FOCUS_COUNT, CHANNELS), (2, 1, 0)), + _fake(cutlass.Float32, (nodes, FOCUS_COUNT, CHANNELS), (2, 1, 0)), + float(focus_eps), + float(focus_tau), + float(label_smoothing), + float(qk_scale), + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +def allocate_grouped_expanded_final_phase_c_attention_adjoint_outputs( + *, + edge_count: int, + node_count: int, + device: torch.device, + grad_m0: torch.Tensor | None = None, + grad_m1: torch.Tensor | None = None, +) -> GroupedExpandedFinalPhaseCAttentionAdjointOutputs: + """Allocate outputs, optionally reusing dead split-state storage.""" + if (grad_m0 is None) != (grad_m1 is None): + raise ValueError("grad_m0 and grad_m1 must be supplied together") + opts = {"device": device, "dtype": torch.float32} + return GroupedExpandedFinalPhaseCAttentionAdjointOutputs( + grad_m0=( + grad_m0 + if grad_m0 is not None + else torch.empty( + (FOCUS_COUNT, edge_count, M0_WIDTH), + device=opts["device"], + dtype=opts["dtype"], + ) + ), + grad_m1=( + grad_m1 + if grad_m1 is not None + else torch.empty( + (FOCUS_COUNT, edge_count, M1_WIDTH), + device=device, + dtype=torch.complex64, + ) + ), + grad_dt=torch.empty( + (edge_count, PACKED_WIGNER_VALUES), + device=opts["device"], + dtype=opts["dtype"], + ), + grad_logits=torch.empty( + (edge_count, FOCUS_COUNT), + device=opts["device"], + dtype=opts["dtype"], + ), + grad_edge=torch.empty( + (edge_count,), + device=opts["device"], + dtype=opts["dtype"], + ), + grad_focus_src=torch.empty( + (FOCUS_COUNT, edge_count, CHANNELS), + device=opts["device"], + dtype=opts["dtype"], + ), + grad_q_node=torch.empty( + (node_count, FOCUS_COUNT, CHANNELS), + device=opts["device"], + dtype=opts["dtype"], + ), + grad_k_node=torch.zeros( + (node_count, FOCUS_COUNT, CHANNELS), + device=opts["device"], + dtype=opts["dtype"], + ), + ) + + +def run_grouped_expanded_final_phase_c_attention_adjoint( + *, + b0: torch.Tensor, + b1: torch.Tensor, + m0: torch.Tensor, + m1: torch.Tensor, + dt_packed: torch.Tensor, + beta: torch.Tensor, + alpha: torch.Tensor, + focus_alpha: torch.Tensor, + focus_src: torch.Tensor, + focus_weight: torch.Tensor, + focus_scale: torch.Tensor, + q_node: torch.Tensor, + k_node: torch.Tensor, + edge_gate: torch.Tensor, + src: torch.Tensor, + dst_ptr: torch.Tensor, + focus_eps: float, + focus_tau: float, + label_smoothing: float, + max_degree: int, + qk_scale: float = QK_SCALE, + outputs: GroupedExpandedFinalPhaseCAttentionAdjointOutputs | None = None, +) -> GroupedExpandedFinalPhaseCAttentionAdjointOutputs: + """Run the adjoint; reusable outputs require clearing ``grad_k_node``.""" + device = b0.device + if device.type != "cuda" or tuple(torch.cuda.get_device_capability(device)) != ( + 9, + 0, + ): + raise RuntimeError("grouped expanded Phase-C adjoint requires SM90") + 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'") + if max_degree < 0 or max_degree > MAX_EDGES_PER_NODE: + raise ValueError( + f"max_degree must be in [0,{MAX_EDGES_PER_NODE}], got {max_degree}" + ) + + node_count = int(dst_ptr.numel() - 1) + edge_count = int(src.numel()) + expected_inputs = ( + ("b0", b0, (FOCUS_COUNT, DEGREE_COUNT, node_count, M0_WIDTH), torch.float32), + ( + "b1", + b1, + (FOCUS_COUNT, DEGREE_COUNT - 1, node_count, M1_WIDTH), + torch.complex64, + ), + ("m0", m0, (FOCUS_COUNT, edge_count, M0_WIDTH), torch.float32), + ("m1", m1, (FOCUS_COUNT, edge_count, M1_WIDTH), torch.complex64), + ( + "dt_packed", + dt_packed, + (edge_count, PACKED_WIGNER_VALUES), + torch.float32, + ), + ("beta", beta, (edge_count, FOCUS_COUNT), torch.float32), + ("alpha", alpha, (edge_count, FOCUS_COUNT), torch.float32), + ("focus_alpha", focus_alpha, (edge_count, FOCUS_COUNT), torch.float32), + ( + "focus_src", + focus_src, + (edge_count, FOCUS_COUNT, CHANNELS), + torch.float32, + ), + ( + "focus_weight", + focus_weight, + (CHANNELS, FOCUS_COUNT), + torch.float32, + ), + ( + "focus_scale", + focus_scale, + (FOCUS_COUNT, CHANNELS), + torch.float32, + ), + ( + "q_node", + q_node, + (node_count, FOCUS_COUNT, CHANNELS), + torch.float32, + ), + ( + "k_node", + k_node, + (node_count, FOCUS_COUNT, CHANNELS), + torch.float32, + ), + ("edge_gate", edge_gate, (edge_count,), torch.float32), + ("src", src, (edge_count,), torch.int32), + ("dst_ptr", dst_ptr, (node_count + 1,), torch.int32), + ) + for name, tensor, shape, dtype in expected_inputs: + _require_tensor(name, tensor, shape, dtype, device) + + if outputs is None: + outputs = allocate_grouped_expanded_final_phase_c_attention_adjoint_outputs( + edge_count=edge_count, + node_count=node_count, + device=device, + ) + expected_outputs = ( + ( + "grad_m0", + outputs.grad_m0, + (FOCUS_COUNT, edge_count, M0_WIDTH), + torch.float32, + ), + ( + "grad_m1", + outputs.grad_m1, + (FOCUS_COUNT, edge_count, M1_WIDTH), + torch.complex64, + ), + ("grad_dt", outputs.grad_dt, (edge_count, PACKED_WIGNER_VALUES), torch.float32), + ("grad_logits", outputs.grad_logits, (edge_count, FOCUS_COUNT), torch.float32), + ("grad_edge", outputs.grad_edge, (edge_count,), torch.float32), + ( + "grad_focus_src", + outputs.grad_focus_src, + (FOCUS_COUNT, edge_count, CHANNELS), + torch.float32, + ), + ( + "grad_q_node", + outputs.grad_q_node, + (node_count, FOCUS_COUNT, CHANNELS), + torch.float32, + ), + ( + "grad_k_node", + outputs.grad_k_node, + (node_count, FOCUS_COUNT, CHANNELS), + torch.float32, + ), + ) + for name, tensor, shape, dtype in expected_outputs: + _require_tensor(name, tensor, shape, dtype, device) + + with torch.cuda.device(device): + compile_grouped_expanded_final_phase_c_attention_adjoint( + float(focus_eps), + float(focus_tau), + float(label_smoothing), + float(qk_scale), + )( + b0, + torch.view_as_real(b1), + m0, + torch.view_as_real(m1), + dt_packed, + beta, + alpha, + focus_alpha, + focus_src, + focus_weight, + focus_scale, + q_node, + k_node, + edge_gate, + src, + dst_ptr, + outputs.grad_m0, + torch.view_as_real(outputs.grad_m1), + outputs.grad_dt, + outputs.grad_logits, + outputs.grad_edge, + outputs.grad_focus_src, + outputs.grad_q_node, + outputs.grad_k_node, + ) + return outputs diff --git a/deepmd/kernels/cute/neo/sm90_k1/prefix.py b/deepmd/kernels/cute/neo/sm90_k1/prefix.py new file mode 100644 index 0000000000..c678069083 --- /dev/null +++ b/deepmd/kernels/cute/neo/sm90_k1/prefix.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""In-place forward and reverse recurrences for the SM90 persistent prefix. + +The CuTe gate is elementwise alias-safe for ``out=residual``: every output +element reads only the corresponding residual element, while all gate values +are derived from the separate ``z`` tensors. The input state may therefore +become the running state without changing the recurrence. + +The reverse recurrence overwrites dead saved preactivations and aliases the +running residual through ``torch.baddbmm`` to minimize edge-sized storage. +""" + +from __future__ import ( + annotations, +) + +import torch + +from .persistent import ( + NeoPersistentComplexSaved, + NeoPersistentComplexState, + NeoPersistentComplexWeights, + _empty_state_like, + _run_gate_adjoint, + _run_gate_forward, + validate_neo_persistent_complex_state, +) + +GATED_LAYERS = 2 + +__all__ = [ + "run_persistent_prefix_forward_inplace", + "run_persistent_prefix_input_adjoint_destructive_saved", +] + + +def _saved(z0: list[torch.Tensor], z1: list[torch.Tensor]) -> NeoPersistentComplexSaved: + return NeoPersistentComplexSaved( + z0=(z0[0], z0[1]), + z1=(z1[0], z1[1]), + ) + + +def run_persistent_prefix_forward_inplace( + state: NeoPersistentComplexState, + weights: NeoPersistentComplexWeights, +) -> tuple[NeoPersistentComplexState, NeoPersistentComplexSaved]: + """Overwrite the caller-owned running state after each separate GEMM.""" + validate_neo_persistent_complex_state(state) + saved_m0: list[torch.Tensor] = [] + saved_m1: list[torch.Tensor] = [] + for layer in range(GATED_LAYERS): + z = _empty_state_like(state) + torch.bmm(state.m0, weights.w0[layer], out=z.m0) + torch.bmm(state.m1, weights.wc[layer], out=z.m1) + saved_m0.append(z.m0) + saved_m1.append(z.m1) + _run_gate_forward(state, z, weights.gate[layer], state) + return state, _saved(saved_m0, saved_m1) + + +def run_persistent_prefix_input_adjoint_destructive_saved( + grad_out: NeoPersistentComplexState, + saved: NeoPersistentComplexSaved, + weights: NeoPersistentComplexWeights, +) -> NeoPersistentComplexState: + """Overwrite each dead saved preactivation with its exact gate adjoint. + + The gate kernel stages the scalar row before writing any output. Every + remaining preactivation element is read and then replaced by its own + adjoint, so input/output aliasing is safe. Backward visits layer 1 before + layer 0; each saved state is therefore dead when it becomes ``grad_z``. + """ + validate_neo_persistent_complex_state(grad_out, name="grad_out") + running = grad_out + for layer in range(GATED_LAYERS - 1, -1, -1): + grad_z = NeoPersistentComplexState(saved.z0[layer], saved.z1[layer]) + _run_gate_adjoint(running, grad_z, weights.gate[layer], grad_z) + torch.baddbmm( + running.m0, + grad_z.m0, + weights.w0_h[layer], + out=running.m0, + ) + torch.baddbmm( + running.m1, + grad_z.m1, + weights.wc_h[layer], + out=running.m1, + ) + return running diff --git a/deepmd/kernels/cute/neo/sm90_k1/radial.py b/deepmd/kernels/cute/neo/sm90_k1/radial.py new file mode 100644 index 0000000000..7e71024db7 --- /dev/null +++ b/deepmd/kernels/cute/neo/sm90_k1/radial.py @@ -0,0 +1,306 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Compact radial state for the SM90 persistent-complex Neo K1 path. + +The fused Phase-A kernel computes three logically separate values in one +edge CTA: the packed-Wigner rotation, the 25-value compact radial map, and the +64-value degree-zero radial attention feature. A persistent-complex Phase A +must not call that kernel merely to recover the latter two values because doing +so would also allocate and write the discarded ``(E,2,10,32)`` real stack. + +This module retains the same FP32 reduction order for the two +radial projections while omitting the dense SO2 output. The resulting compact +state is consumed directly by the split-complex Phase A and its adjoint. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, +) + +import cutlass +import cutlass.cute as cute +import torch +from cuda.bindings.driver import ( + CUstream, +) +from cutlass.cute.runtime import ( + make_fake_compact_tensor, + make_fake_stream, +) + +from ..compile_cache import ( + device_aware_lru_cache, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + ) + + +# CuTe JIT functions use DSL-inferred argument and return types. +# ruff: noqa: ANN001, ANN202, TC002 + +RADIAL_WIDTH = 4 * 32 +COMPACT_WIDTH = 25 +ATTENTION_WIDTH = 64 +THREADS = 64 +FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} + +__all__ = [ + "project_neo_radial_input_adjoint_fp32", + "run_neo_radial_state_forward_fp32", +] + + +@cute.jit +def _radial_state_forward_jit( + radial, + combined_weight, + hidden_weight, + compact_out, + attention_out, + stream: CUstream, +): + edge_count, _ = radial.shape + _radial_state_forward_kernel( + radial, + combined_weight, + hidden_weight, + compact_out, + attention_out, + ).launch( + grid=[edge_count, 1, 1], + block=[THREADS, 1, 1], + stream=stream, + ) + + +@cute.kernel +def _radial_state_forward_kernel( + radial, + combined_weight, + hidden_weight, + compact_out, + attention_out, +): + channel, _, _ = cute.arch.thread_idx() + edge, _, _ = cute.arch.block_idx() + + if channel < COMPACT_WIDTH: + compact = cutlass.Float32(0.0) + for radial_channel in cutlass.range_constexpr(RADIAL_WIDTH): + compact += radial[edge, radial_channel].to( + cutlass.Float32 + ) * combined_weight[radial_channel, channel].to(cutlass.Float32) + compact_out[edge, channel] = compact.to(compact_out.element_type) + + attention = cutlass.Float32(0.0) + for radial_channel in cutlass.range_constexpr(32): + attention += radial[edge, radial_channel].to(cutlass.Float32) * hidden_weight[ + radial_channel, channel + ].to(cutlass.Float32) + attention_out[edge, channel] = attention.to(attention_out.element_type) + + +@device_aware_lru_cache(maxsize=4) +def _compiled_radial_state_forward() -> Callable: + edge_count = cute.sym_int64() + fake_radial = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, RADIAL_WIDTH), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_combined = make_fake_compact_tensor( + cutlass.Float32, + (RADIAL_WIDTH, COMPACT_WIDTH), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_hidden = make_fake_compact_tensor( + cutlass.Float32, + (32, ATTENTION_WIDTH), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_compact = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, COMPACT_WIDTH), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + fake_attention = make_fake_compact_tensor( + cutlass.Float32, + (edge_count, ATTENTION_WIDTH), + stride_order=(1, 0), + **FAKE_TENSOR_KW, + ) + return cute.compile( + _radial_state_forward_jit, + fake_radial, + fake_combined, + fake_hidden, + fake_compact, + fake_attention, + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +def _expect_fp32_cuda( + name: str, + tensor: torch.Tensor, + shape: tuple[int, ...], + device: torch.device, +) -> None: + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}") + if ( + tensor.dtype != torch.float32 + or tensor.device != device + or not tensor.is_cuda + or not tensor.is_contiguous() + ): + raise ValueError(f"{name} must be contiguous CUDA float32 on {device}") + + +def run_neo_radial_state_forward_fp32( + *, + radial_feat: torch.Tensor, + combined_weight: torch.Tensor, + hidden_weight: torch.Tensor, + compact_out: torch.Tensor | None = None, + attention_out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Produce the 25-value Phase-A map and 64-value attention feature.""" + if radial_feat.ndim != 3 or tuple(radial_feat.shape[1:]) != (4, 32): + raise ValueError( + f"radial_feat must have shape (E,4,32), got {tuple(radial_feat.shape)}" + ) + if not radial_feat.is_cuda: + raise ValueError("radial_feat must be a CUDA tensor") + device = radial_feat.device + edge_count = radial_feat.shape[0] + _expect_fp32_cuda( + "radial_feat", + radial_feat, + (edge_count, 4, 32), + device, + ) + _expect_fp32_cuda( + "combined_weight", + combined_weight, + (RADIAL_WIDTH, COMPACT_WIDTH), + device, + ) + _expect_fp32_cuda( + "hidden_weight", + hidden_weight, + (32, ATTENTION_WIDTH), + device, + ) + if edge_count <= 0: + raise ValueError("persistent-complex K1 requires E > 0") + if tuple(torch.cuda.get_device_capability(device)) != (9, 0): + raise RuntimeError("persistent-complex K1 requires SM90") + 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'") + + if compact_out is None: + compact_out = torch.empty( + (edge_count, COMPACT_WIDTH), + dtype=torch.float32, + device=device, + ) + if attention_out is None: + attention_out = torch.empty( + (edge_count, ATTENTION_WIDTH), + dtype=torch.float32, + device=device, + ) + _expect_fp32_cuda( + "compact_out", + compact_out, + (edge_count, COMPACT_WIDTH), + device, + ) + _expect_fp32_cuda( + "attention_out", + attention_out, + (edge_count, ATTENTION_WIDTH), + device, + ) + + with torch.cuda.device(device): + _compiled_radial_state_forward()( + radial_feat.view(edge_count, RADIAL_WIDTH), + combined_weight, + hidden_weight, + compact_out, + attention_out, + ) + return compact_out, attention_out + + +def project_neo_radial_input_adjoint_fp32( + *, + grad_compact: torch.Tensor, + grad_logits: torch.Tensor, + combined_weight: torch.Tensor, + combined_attention_weight: torch.Tensor, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Project compact Phase-A and attention adjoints to ``(E,4,32)``. + + This intentionally follows the strict-FP32 cuBLAS route. No approximation + is introduced: the first matrix product is the compact radial adjoint and + the degree-zero slice receives the independent attention-logit adjoint. + """ + edge_count = grad_compact.shape[0] + device = grad_compact.device + _expect_fp32_cuda( + "grad_compact", + grad_compact, + (edge_count, COMPACT_WIDTH), + device, + ) + _expect_fp32_cuda( + "grad_logits", + grad_logits, + (edge_count, 2), + device, + ) + _expect_fp32_cuda( + "combined_weight", + combined_weight, + (RADIAL_WIDTH, COMPACT_WIDTH), + device, + ) + _expect_fp32_cuda( + "combined_attention_weight", + combined_attention_weight, + (32, 2), + device, + ) + if out is None: + out = torch.empty( + (edge_count, 4, 32), + dtype=torch.float32, + device=device, + ) + _expect_fp32_cuda("out", out, (edge_count, 4, 32), device) + + out_flat = out.view(edge_count, RADIAL_WIDTH) + torch.mm(grad_compact, combined_weight.transpose(0, 1), out=out_flat) + out_flat[:, :32].addmm_( + grad_logits, + combined_attention_weight.transpose(0, 1), + ) + return out diff --git a/deepmd/kernels/cute/neo/sm90_k1/runner.py b/deepmd/kernels/cute/neo/sm90_k1/runner.py new file mode 100644 index 0000000000..ce4decb20e --- /dev/null +++ b/deepmd/kernels/cute/neo/sm90_k1/runner.py @@ -0,0 +1,661 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Split-complex strict-FP32 SM90 implementation of the complete Neo K1. + +The value path uses one split-complex edge representation throughout: + +* compact radial projection and direct split-complex Phase A; +* two persistent split-complex gated residual layers; +* the third residual SO2Linear commuted through Phase C and evaluated from + 64-edge destination-CSR sufficient statistics; +* the Neo output gate followed by the message-grid/readout. + +No ``(E,2,10,32)`` block-real SO2 slab is constructed, and no split/block-real +pack or unpack kernel is used. +""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) +from typing import ( + Any, +) + +import torch +from torch import ( + Tensor, +) + +from ..k1 import ( + _compile_output_gate_backward, + _equivariant_rmsnorm_backward, + _runner_compile_identity, + _so3_linear_backward_input, + _x_wide_manual_backward, +) +from ..k1_runner import ( + NeoFullCuteBackward, +) +from ..k1_so2linear import ( + cached_neo_so2_linear_weights, +) +from ..message_grid_readout_sm90 import ( + prepare_sm90_message_grid_state, + run_sm90_message_grid_backward, +) +from .final_phase_c import ( + ExpandedFinalWeights, + prepare_expanded_final_weights, + run_direct_statistics_forward, +) +from .output_gate import ( + run_chunked_final_output_gate, +) +from .persistent import ( + NeoPersistentComplexSaved, + NeoPersistentComplexState, + NeoPersistentComplexWeights, + prepare_neo_persistent_complex_weights, +) +from .phase_a import ( + run_neo_phase_a_persistent_complex_fp32, +) +from .phase_a_backward import ( + run_neo_phase_a_persistent_complex_backward_fp32, +) +from .phase_c_attention_backward import ( + allocate_grouped_expanded_final_phase_c_attention_adjoint_outputs, + run_grouped_expanded_final_phase_c_attention_adjoint, +) +from .prefix import ( + run_persistent_prefix_forward_inplace, + run_persistent_prefix_input_adjoint_destructive_saved, +) +from .radial import ( + project_neo_radial_input_adjoint_fp32, + run_neo_radial_state_forward_fp32, +) + +FOCUS_COUNT = 2 +CHANNELS = 32 +HIDDEN = FOCUS_COUNT * CHANNELS +DEGREE_COUNT = 16 +M0_WIDTH = 128 +M1_WIDTH = 96 +GATED_LAYERS = 2 +_PERSISTENT_WEIGHT_CACHE = "_deepmd_cute_neo_sm90_persistent_weights" +_FINAL_WEIGHT_CACHE = "_deepmd_cute_neo_sm90_final_weights" + + +def _tensor_version_key(tensor: Tensor) -> tuple[Any, ...]: + return ( + tensor.data_ptr(), + tensor._version, + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + tensor.device, + ) + + +def _prepare_persistent_weights(so2: Any) -> NeoPersistentComplexWeights: + """Pack the three SO2Linear blocks and two scalar gates once.""" + linears = tuple(so2.so2_linears) + nonlinearities = tuple(so2.non_linearities) + if len(linears) != 3 or len(nonlinearities) != 3: + raise NotImplementedError("SM90 K1 requires three SO2 layers") + if any(type(norm).__name__ != "Identity" for norm in so2.so2_inter_norms): + raise NotImplementedError("SM90 K1 requires disabled inter-layer norms") + for nonlinearity in nonlinearities[:GATED_LAYERS]: + if getattr(nonlinearity, "layout", None) != "fndc": + raise NotImplementedError("SM90 K1 requires fndc gates") + activation = getattr( + getattr(nonlinearity, "scalar_act", None), + "activation", + None, + ) + if str(activation).lower() != "silu": + raise NotImplementedError("SM90 K1 requires SiLU gates") + if getattr(nonlinearity.gate_linear, "bias", None) is not None: + raise NotImplementedError("SM90 K1 does not support gate bias") + + sources = tuple( + tensor + for linear in linears + for tensor in (linear.weight_m0, linear.weight_m[0]) + ) + tuple( + nonlinearity.gate_linear.weight + for nonlinearity in nonlinearities[:GATED_LAYERS] + ) + cache_key = tuple(_tensor_version_key(tensor) for tensor in sources) + cached = getattr(so2, _PERSISTENT_WEIGHT_CACHE, None) + if isinstance(cached, tuple) and len(cached) == 3 and cached[0] == cache_key: + return cached[1] + + w0_layers: list[Tensor] = [] + wp_layers: list[Tensor] = [] + for linear in linears: + w0, wp = cached_neo_so2_linear_weights(linear) + w0_layers.append(w0) + wp_layers.append(wp) + gate_layers = [ + nonlinearity.gate_linear.weight.detach() + .view(CHANNELS, FOCUS_COUNT, 3 * CHANNELS) + .permute(1, 0, 2) + .contiguous() + for nonlinearity in nonlinearities[:GATED_LAYERS] + ] + packed = prepare_neo_persistent_complex_weights( + torch.stack(w0_layers, dim=0).contiguous(), + torch.stack(wp_layers, dim=0).contiguous(), + torch.stack(gate_layers, dim=0).contiguous(), + ) + setattr(so2, _PERSISTENT_WEIGHT_CACHE, (cache_key, packed, sources)) + return packed + + +def _prepare_final_weights(so2: Any) -> ExpandedFinalWeights: + """Cache the residual-folded final SO2Linear in node-GEMM form.""" + linear = tuple(so2.so2_linears)[-1] + sources = (linear.weight_m0, linear.weight_m[0]) + key = tuple(_tensor_version_key(tensor) for tensor in sources) + cached = getattr(so2, _FINAL_WEIGHT_CACHE, None) + if isinstance(cached, tuple) and len(cached) == 3 and cached[0] == key: + return cached[1] + + w0, wp = cached_neo_so2_linear_weights(linear) + if tuple(w0.shape) != (FOCUS_COUNT, M0_WIDTH, M0_WIDTH): + raise ValueError("SM90 K1 requires final m=0 weights (2,128,128)") + if tuple(wp.shape) != (FOCUS_COUNT, 2 * M1_WIDTH, 2 * M1_WIDTH): + raise ValueError("SM90 K1 requires final pair weights (2,192,192)") + u = wp[:, :M1_WIDTH, :M1_WIDTH] + v = wp[:, :M1_WIDTH, M1_WIDTH:] + if not torch.equal(wp[:, M1_WIDTH:, :M1_WIDTH], -v): + raise ValueError("final pair weight lower-left block must be -V") + if not torch.equal(wp[:, M1_WIDTH:, M1_WIDTH:], u): + raise ValueError("final pair weight lower-right block must be U") + weights = prepare_expanded_final_weights( + (w0 + torch.eye(M0_WIDTH, device=w0.device, dtype=torch.float32)).contiguous(), + ( + torch.complex(u, v) + + torch.eye(M1_WIDTH, device=w0.device, dtype=torch.complex64) + ).contiguous(), + ) + setattr(so2, _FINAL_WEIGHT_CACHE, (key, weights, sources)) + return weights + + +__all__ = ["NeoSm90K1Runner"] + + +@dataclass(frozen=True) +class _PersistentPrefixForward: + """Pre-final split state and exact gate preactivations for its adjoint.""" + + state: NeoPersistentComplexState + saved: NeoPersistentComplexSaved + + +def _run_persistent_prefix_forward( + state: NeoPersistentComplexState, + weights: NeoPersistentComplexWeights, +) -> _PersistentPrefixForward: + """Run two gated layers while reusing the direct Phase-A state buffer.""" + current, saved = run_persistent_prefix_forward_inplace(state, weights) + return _PersistentPrefixForward( + state=current, + saved=saved, + ) + + +def _run_persistent_prefix_input_adjoint( + grad_out: NeoPersistentComplexState, + saved: NeoPersistentComplexSaved, + weights: NeoPersistentComplexWeights, +) -> NeoPersistentComplexState: + """Overwrite dead gate checkpoints and the running residual in-place.""" + return run_persistent_prefix_input_adjoint_destructive_saved( + grad_out, + saved, + weights, + ) + + +def _run_final_reverse_panels( + *, + grad_out: Tensor, + weights: Any, +) -> tuple[Tensor, Tensor]: + """Form the two strict-FP32 reverse node panels once.""" + node_count = int(grad_out.shape[0]) + grad_node = grad_out.permute(1, 2, 0, 3).contiguous() + b0 = torch.bmm( + grad_node.flatten(0, 1), + weights.w0.flatten(0, 1).transpose(-1, -2), + ).view(FOCUS_COUNT, DEGREE_COUNT, node_count, M0_WIDTH) + grad_node1 = grad_node[:, 1:].to(torch.complex64).contiguous() + b1 = torch.bmm( + grad_node1.flatten(0, 1), + weights.wc.flatten(0, 1).conj().transpose(-1, -2), + ).view(FOCUS_COUNT, DEGREE_COUNT - 1, node_count, M1_WIDTH) + return b0, b1 + + +def _qk_node_input_adjoint( + runner: NeoSm90K1Runner, + grad_q_node: Tensor, + grad_k_node: Tensor, +) -> Tensor: + """Map fused Q/K node adjoints into the wide SO2 input.""" + so2 = runner.so2 + x_wide = runner.x_wide.detach() + x_l0 = x_wide[:, 0, :].reshape(runner.node_count, FOCUS_COUNT, CHANNELS) + grad_x_wide = torch.empty_like(x_wide, memory_format=torch.contiguous_format) + runner.qk_node_input_adjoint( + x_l0.contiguous(), + grad_q_node, + grad_k_node, + so2.attn_q_proj.weight.detach() + .float() + .view(CHANNELS, FOCUS_COUNT, CHANNELS) + .contiguous(), + so2.attn_k_proj.weight.detach() + .float() + .view(CHANNELS, FOCUS_COUNT, CHANNELS) + .contiguous(), + so2.attn_qk_norm.adam_scale.detach().float().contiguous(), + grad_x_wide.view(runner.node_count, DEGREE_COUNT * HIDDEN), + ) + return grad_x_wide + + +class NeoSm90K1Runner(NeoFullCuteBackward): + """Complete Neo K1 runner with one native split representation.""" + + uses_native_sm90_path = True + + def _build_forward_graph(self) -> None: + torch_module = self.torch + so2 = self.so2 + block = self.block + node_count = self.node_count + edge_count = self.edge_count + if self.compute_capability != (9, 0): + raise RuntimeError("split-complex K1 requires SM90") + + self.structural_scratch = None + self.use_full_node = block.node_lmax == block.lmax + x_so2 = self.x if self.use_full_node else self.x[:, : block.mp_ebed_dim] + x_pre = block.pre_so2_norm(x_so2) + self.x_wide = ( + so2.pre_focus_mix( + x_pre.reshape( + node_count, + x_so2.shape[1], + block.channels, + ).unsqueeze(2) + ) + .squeeze(2) + .contiguous() + ) + + self.radial_compact, radial_l0 = run_neo_radial_state_forward_fp32( + radial_feat=self.radial.detach().contiguous(), + combined_weight=self.combined_radial, + hidden_weight=so2.radial_hidden_proj.weight.detach().contiguous(), + ) + phase_a_state = run_neo_phase_a_persistent_complex_fp32( + x_wide=self.x_wide.detach(), + src=self.src_i32, + d_full=self.d.detach(), + radial_compact=self.radial_compact, + channel_basis=so2.radial_degree_mixer.channel_basis.detach() + .view(HIDDEN) + .contiguous(), + ) + self.focus_gate_src = ( + phase_a_state.m0[:, :, :CHANNELS].permute(1, 0, 2).contiguous() + ) + + self.persistent_weights = _prepare_persistent_weights(so2) + prefix = _run_persistent_prefix_forward( + phase_a_state, + self.persistent_weights, + ) + self.phase_c_state = prefix.state + self.persistent_saved = prefix.saved + del phase_a_state, prefix + + x_l0_node = self.x_wide[:, 0, :].reshape( + node_count, + FOCUS_COUNT, + CHANNELS, + ) + self.focus_alpha = torch_module.empty( + edge_count, + FOCUS_COUNT, + device=self.x.device, + dtype=torch_module.float32, + ) + self.q_node = torch_module.empty_like( + x_l0_node, + memory_format=torch_module.contiguous_format, + ) + self.k_node = torch_module.empty_like( + x_l0_node, + memory_format=torch_module.contiguous_format, + ) + self.attention_prelude_forward( + self.focus_gate_src.view(edge_count, HIDDEN), + x_l0_node.contiguous(), + so2.adamw_focus_compete_w.detach().float().contiguous(), + so2.focus_compete_norm.adam_scale.detach().float().contiguous(), + so2.attn_q_proj.weight.detach() + .float() + .view(CHANNELS, FOCUS_COUNT, CHANNELS) + .contiguous(), + so2.attn_k_proj.weight.detach() + .float() + .view(CHANNELS, FOCUS_COUNT, CHANNELS) + .contiguous(), + so2.attn_qk_norm.adam_scale.detach().float().contiguous(), + self.focus_alpha, + self.q_node, + self.k_node, + ) + + self.attn_logits = torch_module.empty( + edge_count, + FOCUS_COUNT, + device=self.x.device, + dtype=torch_module.float32, + ) + self.qk_edge_forward( + self.q_node, + self.k_node, + radial_l0.view(edge_count, FOCUS_COUNT, CHANNELS), + so2.adamw_attn_logit_w.detach().contiguous(), + self.src_i32, + self.dst_i32, + self.attn_logits, + ) + del radial_l0 + self.softmax_fwd( + self.attn_logits, + self.edge_gate, + self.dst_ptr_i32, + so2.adamw_attn_z_bias_raw.detach() + .reshape(FOCUS_COUNT) + .float() + .contiguous(), + self.alpha, + self.group_max, + self.denom, + ) + # First input adjoints need alpha, Q, K, and edge metadata, but not the + # materialized logits or the null-mass parameter-gradient statistics. + self.attn_logits = None + self.group_max = None + self.denom = None + + self.beta = (self.alpha * self.focus_alpha).contiguous() + self.final_weights = _prepare_final_weights(so2) + raw_phase_c = run_direct_statistics_forward( + m0=self.phase_c_state.m0, + m1=self.phase_c_state.m1, + dt_packed=self.dt.detach(), + beta=self.beta, + dst_ptr=self.dst_ptr_i32, + weights=self.final_weights, + ).output + self.phase_c_out = run_chunked_final_output_gate( + raw=raw_phase_c, + x_wide=self.x_wide.detach(), + norm_scale=so2.attn_output_gate_norm.adam_scale.detach() + .float() + .reshape(FOCUS_COUNT, CHANNELS) + .contiguous(), + gate_weight=so2.adamw_attn_gate_w.detach() + .float() + .reshape(CHANNELS, FOCUS_COUNT, 1) + .contiguous(), + rotate_inv_rescale=so2.rotate_inv_rescale_full.detach().contiguous(), + eps=float(so2.attn_output_gate_norm.eps), + ).to(dtype=so2.compute_dtype) + # The partial/node statistics and ungated output have no backward role: + # the exact adjoint recomputes from grad_out and the pre-final state. + del raw_phase_c + + out = self.phase_c_out.detach().to(dtype=so2.dtype) + self.out_gate_flat = out.detach() + self.message_grid_product = None + self.message_grid_sm90_state = None + if so2.message_node_grid_product is not None: + if self.packed_message_grid: + from ..k1_message_grid_packed import ( + run_packed_message_grid_forward, + ) + + self.message_grid_sm90_state = prepare_sm90_message_grid_state( + so2.message_node_grid_product + ) + grid_out, product = run_packed_message_grid_forward( + so2.message_node_grid_product, + out, + self.x_wide, + return_product=True, + sm90_state=self.message_grid_sm90_state, + ) + self.message_grid_product = product.detach() + else: + grid_out = so2.message_node_grid_product(out, self.x_wide) + out = out + grid_out + + self.post_mix_input = out.detach() + out = so2.post_focus_mix(out.unsqueeze(2)).squeeze(2) + self.post_norm_input = out.unsqueeze(2).detach() + so2_out = block.post_so2_norm(self.post_norm_input) + if self.use_full_node: + self.final = so2_out + else: + final = self.x.new_zeros(self.x.shape) + final[:, : block.mp_ebed_dim] = so2_out + self.final = final + + def input_adjoint(self, grad_out: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """Return K1 input adjoints through the split-complex reverse path.""" + return _runner_backward(self, grad_out) + + +def _final_manual_backward_sm90( + runner: NeoSm90K1Runner, + grad_out: Tensor, +) -> tuple[Tensor, Tensor]: + """Use the one-slab tiled message-grid adjoint in the final K1 boundary.""" + if not ( + runner.packed_message_grid + and runner.so2.message_node_grid_product is not None + and runner.message_grid_sm90_state is not None + ): + raise RuntimeError("SM90 K1 requires the packed message-grid path") + + so2 = runner.so2 + block = runner.block + if runner.use_full_node: + grad_so2_out = grad_out + else: + grad_so2_out = grad_out[:, : block.mp_ebed_dim, :, :] + + phase = runner.phase_c_out.detach() + x_wide = runner.x_wide.detach() + post_in = runner.post_mix_input.unsqueeze(2) + grad_post_norm_in = _equivariant_rmsnorm_backward( + block.post_so2_norm, + runner.post_norm_input, + grad_so2_out, + ) + grad_post_mix = _so3_linear_backward_input( + so2.post_focus_mix, + post_in, + grad_post_norm_in.squeeze(2).unsqueeze(2), + ).squeeze(2) + + message_grid_product = runner.message_grid_product + runner.message_grid_product = None + grad_out_gate_flat, grad_grid_context = run_sm90_message_grid_backward( + so2.message_node_grid_product, + runner.out_gate_flat, + x_wide, + grad_post_mix, + message_grid_product, + runner.message_grid_sm90_state, + ) + del message_grid_product + runner.message_grid_sm90_state = None + + grad_out_gate_flat.add_(grad_post_mix) + # FrameExpand's input adjoint preserves its degree-major einsum stride. + # The output-gate kernel updates one flat node panel in place, so establish + # that writable contract once at this consumer boundary. + grad_x_wide_down = grad_grid_context.contiguous() + grad_phase = grad_out_gate_flat.contiguous() + output_gate_backward = _compile_output_gate_backward( + _runner_compile_identity(runner), + float(so2.attn_output_gate_norm.eps), + ) + output_gate_backward( + grad_phase.view(runner.node_count, DEGREE_COUNT * HIDDEN), + phase.contiguous().view(runner.node_count, DEGREE_COUNT * HIDDEN), + x_wide.contiguous().view(runner.node_count, DEGREE_COUNT * HIDDEN), + so2.attn_output_gate_norm.adam_scale.detach() + .float() + .reshape(FOCUS_COUNT, CHANNELS) + .contiguous(), + so2.adamw_attn_gate_w.detach() + .float() + .reshape(CHANNELS, FOCUS_COUNT, 1) + .contiguous(), + grad_phase.view(runner.node_count, DEGREE_COUNT * HIDDEN), + grad_x_wide_down.view(runner.node_count, DEGREE_COUNT * HIDDEN), + ) + return grad_phase.reshape_as(phase), grad_x_wide_down + + +def _runner_backward( + runner: NeoSm90K1Runner, + grad_out: Tensor, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: + so2 = runner.so2 + grad_phase, grad_x_wide_down = _final_manual_backward_sm90(runner, grad_out) + runner.phase_c_out = None + runner.out_gate_flat = None + runner.post_mix_input = None + runner.post_norm_input = None + grad_node = ( + grad_phase.view( + runner.node_count, + DEGREE_COUNT, + FOCUS_COUNT, + CHANNELS, + ) + .permute(0, 2, 1, 3) + .contiguous() + ) + grad_node.mul_(runner.rotate.view(1, 1, DEGREE_COUNT, 1)) + b0, b1 = _run_final_reverse_panels( + grad_out=grad_node, + weights=runner.final_weights, + ) + fused_outputs = allocate_grouped_expanded_final_phase_c_attention_adjoint_outputs( + edge_count=runner.edge_count, + node_count=runner.node_count, + device=grad_node.device, + grad_m0=runner.phase_c_state.m0, + grad_m1=runner.phase_c_state.m1, + ) + fused_adjoint = run_grouped_expanded_final_phase_c_attention_adjoint( + b0=b0, + b1=b1, + m0=runner.phase_c_state.m0, + m1=runner.phase_c_state.m1, + dt_packed=runner.dt.detach(), + beta=runner.beta, + alpha=runner.alpha, + focus_alpha=runner.focus_alpha, + focus_src=runner.focus_gate_src, + focus_weight=so2.adamw_focus_compete_w.detach().float().contiguous(), + focus_scale=so2.focus_compete_norm.adam_scale.detach() + .float() + .reshape(FOCUS_COUNT, CHANNELS) + .contiguous(), + q_node=runner.q_node, + k_node=runner.k_node, + edge_gate=runner.edge_gate, + src=runner.src_i32, + dst_ptr=runner.dst_ptr_i32, + focus_eps=float(so2.focus_compete_norm.eps), + focus_tau=float(so2.focus_softmax_tau), + label_smoothing=float(so2.focus_label_smoothing), + max_degree=128, + qk_scale=CHANNELS**-0.5, + outputs=fused_outputs, + ) + grad_stack = NeoPersistentComplexState( + fused_adjoint.grad_m0, + fused_adjoint.grad_m1, + ) + grad_dt = fused_adjoint.grad_dt + grad_logits = fused_adjoint.grad_logits + grad_edge = fused_adjoint.grad_edge + grad_focus_src = fused_adjoint.grad_focus_src + grad_q_node = fused_adjoint.grad_q_node + grad_k_node = fused_adjoint.grad_k_node + del fused_adjoint, b0, b1, grad_node + runner.phase_c_state = None + runner.beta = None + runner.final_weights = None + runner.focus_gate_src = None + runner.alpha = None + runner.focus_alpha = None + + grad_phase_a = _run_persistent_prefix_input_adjoint( + grad_stack, + runner.persistent_saved, + runner.persistent_weights, + ) + runner.persistent_saved = None + # focus_gate_src aliases the first m=0 row of the direct Phase-A result. + grad_phase_a.m0[:, :, :CHANNELS].add_(grad_focus_src) + del grad_stack, grad_focus_src + phase_a = run_neo_phase_a_persistent_complex_backward_fp32( + grad_state=grad_phase_a, + radial_compact=runner.radial_compact, + channel_basis=so2.radial_degree_mixer.channel_basis.detach() + .view(HIDDEN) + .contiguous(), + x_wide=runner.x_wide.detach(), + source_order=runner.source_order_i32, + source_ptr=runner.source_ptr_i32, + d_full=runner.d.detach(), + ) + runner.radial_compact = None + runner.persistent_weights = None + grad_radial = project_neo_radial_input_adjoint_fp32( + grad_compact=phase_a.grad_radial_compact, + grad_logits=grad_logits, + combined_weight=runner.combined_radial, + combined_attention_weight=runner.combined_attention_radial, + ) + grad_x_wide = phase_a.grad_x_wide.view_as(runner.x_wide) + grad_x_wide.add_(_qk_node_input_adjoint(runner, grad_q_node, grad_k_node)) + runner.q_node = None + runner.k_node = None + grad_x_wide.add_(grad_x_wide_down) + grad_x = _x_wide_manual_backward(runner, grad_x_wide) + runner.grad_edge = grad_edge + + return grad_x, phase_a.grad_d_full, grad_dt, grad_radial diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index f468379565..24457cd300 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -782,18 +782,19 @@ def _eval_lower_strategy( list(inner.get_sel()), return_mode="edges", ) - predict = inner.forward_common_lower( - edge_schema.coord, - edge_schema.atype, - edge_schema.edge_index, - edge_schema.edge_vec, - edge_schema.edge_scatter_index, - edge_schema.edge_mask, - fparam=fparam, - aparam=aparam, - charge_spin=charge_spin, - input_prec=coord.dtype, - ) + with self.dp._frozen_parameter_context(): + predict = inner.forward_common_lower( + edge_schema.coord, + edge_schema.atype, + edge_schema.edge_index, + edge_schema.edge_vec, + edge_schema.edge_scatter_index, + edge_schema.edge_mask, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + input_prec=coord.dtype, + ) else: ext_coord, ext_atype, nlist, mapping = self._nlist_builder.build( coord, @@ -802,16 +803,17 @@ def _eval_lower_strategy( self.rcut, list(inner.get_sel()), ) - model_lower = inner.forward_common_lower( - ext_coord, - ext_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - charge_spin=charge_spin, - ) + with self.dp._frozen_parameter_context(): + model_lower = inner.forward_common_lower( + ext_coord, + ext_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + charge_spin=charge_spin, + ) predict = communicate_extended_output( model_lower, self.output_def, diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index 6e05ae884f..66c578b4b6 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -52,6 +52,7 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) +from deepmd.kernels.cute.neo import runtime_policy as cute_runtime_policy from deepmd.kernels.utils import ( use_amp_infer, ) @@ -106,6 +107,9 @@ safe_norm, safe_numpy_to_tensor, ) +from .sezm_nn.edge_cache import ( + build_sorted_edge_index_metadata, +) if TYPE_CHECKING: from collections.abc import ( @@ -1129,6 +1133,7 @@ def forward( force_embedding: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, spin: torch.Tensor | None = None, + edge_index_sorted_by_dst: bool = False, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -1170,6 +1175,10 @@ def forward( initial SO(3) backbone state before the interaction blocks. charge_spin Frame-level charge and spin conditions with shape (nf, 2). + spin + Optional per-atom spin vectors. + edge_index_sorted_by_dst + Host-side provenance that ``edge_index[1]`` is nondecreasing. Returns ------- @@ -1205,6 +1214,7 @@ def forward( edge_index=edge_index, edge_vec=edge_vec, edge_mask=edge_mask, + edge_index_sorted_by_dst=edge_index_sorted_by_dst, force_embedding=force_embedding, charge_spin=charge_spin, spin=spin, @@ -1282,6 +1292,11 @@ def forward( # the model is roll-equivariant, so inference fixes gamma. random_gamma=self.random_gamma and self.training, wigner_calc=self.wigner_calc, + packed_wigner_candidate=self._packed_wigner_candidate( + type_ebed.device, + type_ebed.dtype, + extended_coord.dtype, + ), build_wigner=self._need_full_wigner, ) @@ -1362,29 +1377,46 @@ def forward( # === Step 10. Fuse edge type features into radial features (fp32+) === with nvtx_range("radial_fuse"): + block_dtype = ( + torch.float32 if edge_cache.D_packed is not None else self.dtype + ) if radial_feat is not None: radial_feat = radial_feat + rearrange( edge_cache.edge_type_feat, "E C -> E 1 C" ) - radial_feat = radial_feat.to(dtype=self.dtype) + radial_feat = radial_feat.to(dtype=block_dtype) rad_feat_per_block = [ radial_feat[:, :rad_len, :] for rad_len in self.rad_sizes_per_block ] # list of (E, lmax+1, C) else: rad_feat_per_block = [] - # === Step 11. Convert to self.dtype and run blocks === + # === Step 11. Convert to the block runtime dtype and run blocks === # The block stage is skipped entirely when there are no interaction # blocks (zero-block descriptor) or no valid edges, sparing the working # edge-cache dtype cast that only the blocks consume. with nvtx_range("blocks"): - x = x.to(dtype=self.dtype) # (N, D, 1, C) + x = x.to(dtype=block_dtype) # (N, D, 1, C) if force_embedding is not None: - x = x + force_embedding.to(dtype=self.dtype) + x = x + force_embedding.to(dtype=block_dtype) if self.blocks and edge_cache.src.numel() > 0: - edge_cache = edge_cache_to_dtype(edge_cache, self.dtype) + edge_cache = edge_cache_to_dtype(edge_cache, block_dtype) + k1_dst_ptr, k1_source_order, k1_source_ptr = ( + self._prepare_cute_k1_sorted_metadata( + edge_cache, + n_nodes, + ) + ) with self._compute_mode_ctx(extended_coord.device): - x = self._forward_blocks(x, edge_cache, rad_feat_per_block) + x = self._forward_blocks( + x, + edge_cache, + rad_feat_per_block, + comm_dict=comm_dict, + k1_dst_ptr=k1_dst_ptr, + k1_source_order=k1_source_order, + k1_source_ptr=k1_source_ptr, + ) # === Step 12. Final l=0 output mixing === with nvtx_range("output_ffn"): @@ -1410,6 +1442,7 @@ def forward_with_edges( edge_index: torch.Tensor, edge_vec: torch.Tensor, edge_mask: torch.Tensor, + edge_index_sorted_by_dst: bool = False, force_embedding: torch.Tensor | None = None, charge_spin: torch.Tensor | None = None, spin: torch.Tensor | None = None, @@ -1443,6 +1476,8 @@ def forward_with_edges( Edge vectors with shape (E, 3) in Ã…. edge_mask Edge mask with shape (E,). + edge_index_sorted_by_dst + Host-side provenance that ``edge_index[1]`` is nondecreasing. force_embedding Optional precomputed equivariant force embedding with shape ``(nf * nloc, D, 1, channels)``, where @@ -1541,6 +1576,12 @@ def forward_with_edges( # the model is roll-equivariant, so inference fixes gamma. random_gamma=self.random_gamma and self.training, wigner_calc=self.wigner_calc, + packed_wigner_candidate=self._packed_wigner_candidate( + type_ebed.device, + type_ebed.dtype, + extended_coord.dtype, + ), + destinations_sorted=edge_index_sorted_by_dst, build_wigner=self._need_full_wigner, node_partial_exchange=node_partial_exchange, ) @@ -1615,26 +1656,41 @@ def forward_with_edges( # === Step 9. Fuse edge type features into radial features (fp32+) === with nvtx_range("radial_fuse"): - radial_feat = radial_feat.to(dtype=self.dtype) + block_dtype = ( + torch.float32 if edge_cache.D_packed is not None else self.dtype + ) + radial_feat = radial_feat.to(dtype=block_dtype) radial_feat = radial_feat + rearrange( - edge_cache.edge_type_feat.to(dtype=self.dtype), "E C -> E 1 C" + edge_cache.edge_type_feat.to(dtype=block_dtype), "E C -> E 1 C" ) rad_feat_per_block = [ radial_feat[:, :rad_len, :] for rad_len in self.rad_sizes_per_block ] - # === Step 10. Convert to self.dtype and run blocks === + # === Step 10. Convert to the block runtime dtype and run blocks === # The block stage is skipped entirely for the zero-block descriptor, # sparing the working edge-cache dtype cast that only the blocks consume. with nvtx_range("blocks"): - x = x.to(dtype=self.dtype) # (N, D, 1, C) + x = x.to(dtype=block_dtype) # (N, D, 1, C) if force_embedding is not None: - x = x + force_embedding.to(dtype=self.dtype) + x = x + force_embedding.to(dtype=block_dtype) if self.blocks: - edge_cache = edge_cache_to_dtype(edge_cache, self.dtype) + edge_cache = edge_cache_to_dtype(edge_cache, block_dtype) + k1_dst_ptr, k1_source_order, k1_source_ptr = ( + self._prepare_cute_k1_sorted_metadata( + edge_cache, + n_nodes, + ) + ) with self._compute_mode_ctx(extended_coord.device): x = self._forward_blocks( - x, edge_cache, rad_feat_per_block, comm_dict=comm_dict + x, + edge_cache, + rad_feat_per_block, + comm_dict=comm_dict, + k1_dst_ptr=k1_dst_ptr, + k1_source_order=k1_source_order, + k1_source_ptr=k1_source_ptr, ) # === Step 11. Keep the owned-atom rows for the read-out === @@ -1654,12 +1710,60 @@ def forward_with_edges( descriptor = x_scalar.reshape(nf, out_nloc, self.channels) # (nf, nloc, C) return descriptor.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), x.contiguous() + @torch.jit.unused + def _run_output_readout(self, ffn_in: torch.Tensor) -> torch.Tensor: + """Return the residual-inclusive scalar output with guarded CuTe routing.""" + from deepmd.kernels.cute.neo.readout_l0 import ( + run_neo_output_readout, + ) + + return run_neo_output_readout( + self.output_ffn, + ffn_in, + parameters_frozen=self._readout_parameters_are_frozen(), + ) + + def _readout_parameters_are_frozen(self) -> bool: + """Return whether descriptor inference can exclude every weight gradient.""" + return not self.training and not any( + parameter.requires_grad for parameter in self.parameters() + ) + + def _prepare_cute_k1_sorted_metadata( + self, + edge_cache: EdgeFeatureCache, + n_nodes: int, + ) -> tuple[ + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: + """Build one per-forward CSR tensor bundle for eligible K1 blocks.""" + if ( + self.training + or not edge_cache.destinations_sorted + or edge_cache.D_packed is None + or edge_cache.edge_src_gate is not None + ): + return None, None, None + if not cute_runtime_policy.is_cute_infer_enabled(): + return None, None, None + return build_sorted_edge_index_metadata( + edge_cache.src, + edge_cache.dst, + n_nodes, + validate_sorted=cute_runtime_policy.is_cute_strict_enabled(), + ) + def _forward_blocks( self, x: torch.Tensor, edge_cache: EdgeFeatureCache, radial_feat_per_block: list[torch.Tensor], comm_dict: dict[str, torch.Tensor] | None = None, + k1_dst_ptr: torch.Tensor | None = None, + k1_source_order: torch.Tensor | None = None, + k1_source_ptr: torch.Tensor | None = None, ) -> torch.Tensor: """ Run the interaction blocks with optional depth attention. @@ -1695,6 +1799,9 @@ def _forward_blocks( edge_cache, blk_radial, comm_dict=self._block_comm(i, comm_dict), + k1_dst_ptr=k1_dst_ptr, + k1_source_order=k1_source_order, + k1_source_ptr=k1_source_ptr, ) return x @@ -1723,6 +1830,9 @@ def node_l0_extractor(v: torch.Tensor) -> torch.Tensor: blk_radial, unit_history=truncated_unit_history, comm_dict=self._block_comm(i, comm_dict), + k1_dst_ptr=k1_dst_ptr, + k1_source_order=k1_source_order, + k1_source_ptr=k1_source_ptr, ) unit_history.append(so2_unit_output) unit_history.extend(ffn_unit_outputs) @@ -1756,6 +1866,9 @@ def node_l0_extractor(v: torch.Tensor) -> torch.Tensor: blk_radial, unit_history=truncated_block_history, comm_dict=self._block_comm(i, comm_dict), + k1_dst_ptr=k1_dst_ptr, + k1_source_order=k1_source_order, + k1_source_ptr=k1_source_ptr, ) block_history.append(block_summary) x = block_output @@ -1770,6 +1883,34 @@ def node_l0_extractor(v: torch.Tensor) -> torch.Tensor: ).to(dtype=self.dtype) return x + def _packed_wigner_candidate( + self, + device: torch.device, + dtype: torch.dtype, + geometry_dtype: torch.dtype, + ) -> bool: + """Return whether all blocks satisfy packed K1's stable contract.""" + from deepmd.kernels.cute.neo.k1 import ( + is_packed_wigner_candidate, + ) + + return is_packed_wigner_candidate( + blocks=self.blocks, + training=self.training, + device=device, + dtype=dtype, + producer_modules=( + self.radial_basis, + self.radial_embedding, + self.edge_envelope, + self.wigner_calc, + *((self.inner_clamp,) if self.inner_clamp is not None else ()), + *((self.bridging_switch,) if self.bridging_switch is not None else ()), + ), + producer_dtypes=(dtype, geometry_dtype), + has_edge_src_gate=self.bridging_switch is not None, + ) + def _apply_readout(self, x: torch.Tensor, n_rows: int) -> torch.Tensor: """Fold the node tensor into the scalar (``l=0``) descriptor. @@ -1808,6 +1949,15 @@ def _apply_readout(self, x: torch.Tensor, n_rows: int) -> torch.Tensor: x_ro = x[:, : self.node_readout_dim, :, :].to(dtype=self.compute_dtype) for layer in self.readout_pre_layers: x_ro = x_ro + layer(x_ro) + if not self.readout_pre_layers: + if ( + not torch.jit.is_scripting() + and cute_runtime_policy.is_cute_infer_enabled() + ): + return self._run_output_readout(x_ro).reshape( + n_rows, 1, 1, self.channels + ) + return (x_ro + self.output_ffn(x_ro))[:, 0:1, :, :] return (x_ro + self.output_ffn(x_ro))[:, 0:1, :, :] def _edge_quaternion(self, edge_cache: EdgeFeatureCache) -> torch.Tensor: @@ -1851,6 +2001,18 @@ def _build_gie_zonal_coupling( the blocks are skipped (all-Cartesian model) the full coupling is reconstructed from the edge quaternion via the m=0-only path. """ + if edge_cache.D_packed is not None: + if self.gie_zonal_wigner_calc is None: + return None + mp_coupling = edge_cache.D_packed.index_select( + 1, + self.gie.packed_zonal_offsets, + ) + extra_coupling = self.gie_zonal_wigner_calc.forward_zonal( + self._edge_quaternion(edge_cache), + lmin=self.lmax + 1, + ) + return torch.cat([mp_coupling, extra_coupling], dim=1) if edge_cache.Dt_full is None: calc = self.gie_zonal_wigner_calc or self.wigner_calc return calc.forward_zonal(self._edge_quaternion(edge_cache), lmin=1) diff --git a/deepmd/pt/model/descriptor/sezm_nn/block.py b/deepmd/pt/model/descriptor/sezm_nn/block.py index 6b170a8935..7f94eb3a72 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/block.py +++ b/deepmd/pt/model/descriptor/sezm_nn/block.py @@ -26,6 +26,9 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) +from deepmd.kernels.cute.neo.runtime_policy import ( + is_cute_infer_enabled, +) from deepmd.pt.utils import ( env, ) @@ -673,6 +676,9 @@ def forward( radial_feat: torch.Tensor, unit_history: list[torch.Tensor] | None = None, comm_dict: dict[str, torch.Tensor] | None = None, + k1_dst_ptr: torch.Tensor | None = None, + k1_source_order: torch.Tensor | None = None, + k1_source_ptr: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, @@ -712,7 +718,16 @@ def forward( - full AttnRes path returns `(block_output, None, so2_unit_output, ffn_unit_outputs)` - block AttnRes path returns `(block_output, block_summary, None, None)` """ - return self._forward_impl(x, edge_cache, radial_feat, unit_history, comm_dict) + return self._forward_impl( + x, + edge_cache, + radial_feat, + unit_history, + comm_dict, + k1_dst_ptr, + k1_source_order, + k1_source_ptr, + ) def _extract_l0_from_canonical(self, value: torch.Tensor) -> torch.Tensor: """ @@ -751,6 +766,9 @@ def _run_so2_unit( edge_cache: EdgeFeatureCache, radial_feat: torch.Tensor, comm_dict: dict[str, torch.Tensor] | None = None, + k1_dst_ptr: torch.Tensor | None = None, + k1_source_order: torch.Tensor | None = None, + k1_source_ptr: torch.Tensor | None = None, ) -> torch.Tensor: """ Run the SO(2) unit without an outer block-level residual shortcut. @@ -786,21 +804,55 @@ def _run_so2_unit( x_, edge_cache_no_proj, radial_feat_, + k1_dst_ptr, + k1_source_order, + k1_source_ptr, ), x, radial_feat, use_reentrant=False, preserve_rng_state=True, ) - return self._run_so2_unit_impl(x, edge_cache, radial_feat) + return self._run_so2_unit_impl( + x, + edge_cache, + radial_feat, + k1_dst_ptr, + k1_source_order, + k1_source_ptr, + ) def _run_so2_unit_impl( self, x: torch.Tensor, edge_cache: EdgeFeatureCache, radial_feat: torch.Tensor, + k1_dst_ptr: torch.Tensor | None = None, + k1_source_order: torch.Tensor | None = None, + k1_source_ptr: torch.Tensor | None = None, ) -> torch.Tensor: """Run the SO(2) unit implementation.""" + if not self.training and is_cute_infer_enabled(): + from deepmd.kernels.cute.neo.k1 import ( + maybe_run_cute_k1, + ) + + cute_out = maybe_run_cute_k1( + self, + x, + edge_cache, + radial_feat, + dst_ptr=k1_dst_ptr, + source_order=k1_source_order, + source_ptr=k1_source_ptr, + ) + if cute_out is not None: + return cute_out + if edge_cache.D_packed is not None: + raise RuntimeError( + "packed Wigner cache reached an ineligible Neo K1 dispatch" + ) + n_node = x.shape[0] channels = self.channels use_full_node = self.node_lmax == self.lmax @@ -861,6 +913,9 @@ def _forward_with_residual_shortcuts( radial_feat: torch.Tensor, unit_history: list[torch.Tensor] | None = None, comm_dict: dict[str, torch.Tensor] | None = None, + k1_dst_ptr: torch.Tensor | None = None, + k1_source_order: torch.Tensor | None = None, + k1_source_ptr: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, @@ -891,7 +946,15 @@ def _forward_with_residual_shortcuts( Tuple `(block_output, None, None, None)`. """ with nvtx_range("so2_conv"): - so2_unit_output = self._run_so2_unit(x, edge_cache, radial_feat, comm_dict) + so2_unit_output = self._run_so2_unit( + x, + edge_cache, + radial_feat, + comm_dict, + k1_dst_ptr, + k1_source_order, + k1_source_ptr, + ) so2_state = x + so2_unit_output with nvtx_range("ffn"): @@ -910,6 +973,9 @@ def _forward_with_full_attn_res( radial_feat: torch.Tensor, unit_history: list[torch.Tensor] | None = None, comm_dict: dict[str, torch.Tensor] | None = None, + k1_dst_ptr: torch.Tensor | None = None, + k1_source_order: torch.Tensor | None = None, + k1_source_ptr: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, @@ -949,7 +1015,13 @@ def _forward_with_full_attn_res( current_x=x, ) so2_unit_output = self._run_so2_unit( - so2_input, edge_cache, radial_feat, comm_dict + so2_input, + edge_cache, + radial_feat, + comm_dict, + k1_dst_ptr, + k1_source_order, + k1_source_ptr, ) with nvtx_range("ffn"): @@ -978,6 +1050,9 @@ def _forward_with_block_attn_res( radial_feat: torch.Tensor, unit_history: list[torch.Tensor] | None = None, comm_dict: dict[str, torch.Tensor] | None = None, + k1_dst_ptr: torch.Tensor | None = None, + k1_source_order: torch.Tensor | None = None, + k1_source_ptr: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, @@ -1017,7 +1092,13 @@ def _forward_with_block_attn_res( current_x=x, ) so2_unit_output = self._run_so2_unit( - so2_input, edge_cache, radial_feat, comm_dict + so2_input, + edge_cache, + radial_feat, + comm_dict, + k1_dst_ptr, + k1_source_order, + k1_source_ptr, ) with nvtx_range("ffn"): diff --git a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py index ff1d497d8b..84f443ef1f 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py +++ b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py @@ -25,6 +25,10 @@ rearrange, ) +from deepmd.kernels.cute.neo.runtime_policy import ( + is_cute_infer_enabled, +) + from .utils import ( get_promoted_dtype, nvtx_range, @@ -68,10 +72,11 @@ class EdgeFeatureCache(NamedTuple): inv_sqrt_deg Inverse square root smooth degree normalization with shape (N, 1, 1). D_full - Block-diagonal Wigner-D matrix with shape (E, D, D) where D=(lmax+1)^2. - Used for efficient batched rotation. None if not available. + Block-diagonal Wigner-D matrix with shape (E, D, D). None when dense + Wigner storage is skipped. Dt_full - Transpose of D_full with shape (E, D, D). None if not available. + Transpose of D_full with shape (E, D, D). None when dense Wigner + storage is skipped. edge_quat Per-edge global-to-local quaternion actually used to build ``D_full`` and ``Dt_full`` with shape (E, 4). Includes the optional random local-Z roll. @@ -92,6 +97,12 @@ class EdgeFeatureCache(NamedTuple): by this gate to forbid any node whose local neighborhood enters the frozen zone from propagating information along its outgoing edges. + destinations_sorted + Host-side provenance indicating that ``dst`` is nondecreasing. CuTe K1 + may only consume caches carrying this guarantee. + D_packed + Opt-in Neo packed Wigner panel with shape (E, 46). This is kept + separate so the public dense fields have a stable rank. """ src: torch.Tensor @@ -108,6 +119,78 @@ class EdgeFeatureCache(NamedTuple): Dt_from_m_cache: dict[str, torch.Tensor] | None = None edge_src_gate: torch.Tensor | None = None edge_quat: torch.Tensor | None = None + destinations_sorted: bool = False + D_packed: torch.Tensor | None = None + + +def _separate_packed_wigner( + D_full: torch.Tensor | None, + Dt_full: torch.Tensor | None, +) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """Move the opt-in 46-value panel out of the dense Wigner fields.""" + if D_full is None or D_full.dim() != 2: + return D_full, Dt_full, None + if Dt_full is not D_full: + raise RuntimeError("packed Wigner forward and transpose must share storage") + return None, None, D_full + + +def build_sorted_edge_index_metadata( + src: torch.Tensor, + dst: torch.Tensor, + n_nodes: int, + *, + validate_sorted: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build destination and source CSR metadata for one sorted edge list.""" + if src.dim() != 1 or dst.dim() != 1: + raise ValueError("src and dst must be one-dimensional") + if src.shape != dst.shape: + raise ValueError("src and dst must have the same shape") + if src.device != dst.device: + raise ValueError("src and dst must be on the same device") + if src.dtype not in (torch.int32, torch.int64): + raise TypeError("src must have dtype int32 or int64") + if dst.dtype not in (torch.int32, torch.int64): + raise TypeError("dst must have dtype int32 or int64") + if n_nodes < 0: + raise ValueError("n_nodes must be non-negative") + if src.numel() > 2**31 - 1: + raise ValueError("sorted edge metadata requires E <= 2**31 - 1") + + src = src.contiguous() + dst = dst.contiguous() + if validate_sorted and dst.numel() > 1: + torch._assert_async( + torch.all(dst[1:] >= dst[:-1]), + "Neo K1 destinations_sorted=True requires monotonically " + "nondecreasing destination indices", + ) + dst_boundaries = torch.arange( + n_nodes + 1, + device=dst.device, + dtype=dst.dtype, + ) + dst_ptr = torch.searchsorted( + dst, + dst_boundaries, + out_int32=True, + ).contiguous() + + source_order_i64 = torch.argsort(src, stable=True) + sorted_src = src.index_select(0, source_order_i64) + src_boundaries = torch.arange( + n_nodes + 1, + device=src.device, + dtype=src.dtype, + ) + source_ptr = torch.searchsorted( + sorted_src, + src_boundaries, + out_int32=True, + ).contiguous() + source_order = source_order_i64.to(dtype=torch.int32).contiguous() + return dst_ptr, source_order, source_ptr def compute_edge_src_gate( @@ -248,6 +331,7 @@ def build_edge_cache( n_radial: int, random_gamma: bool, wigner_calc: WignerCalculatorFn, + packed_wigner_candidate: bool = False, build_wigner: bool = True, ) -> EdgeFeatureCache: """ @@ -310,6 +394,11 @@ def build_edge_cache( wigner_calc Callable that converts edge-aligned quaternions into packed Wigner-D blocks. + packed_wigner_candidate + Whether descriptor-level strict-FP32 K1 checks passed. Concrete edge + count and destination ordering are checked before panel generation. + build_wigner + Whether to materialize Wigner-D blocks for the SO(2) path. Returns ------- @@ -374,7 +463,18 @@ def build_edge_cache( random_gamma=random_gamma, wigner_calc=wigner_calc, build_full=build_wigner, + packed_wigner=( + build_wigner + and _packed_wigner_edges_eligible( + packed_wigner_candidate, + edge_count=dst.numel(), + node_count=n_nodes, + destinations_sorted=True, + runtime_dtypes=(edge_vec.dtype, edge_env.dtype, edge_rbf.dtype), + ) + ), ) # (E, D, D), (E, D, D), (E, 4) + D_full, Dt_full, D_packed = _separate_packed_wigner(D_full, Dt_full) edge_type_feat = build_edge_type_feat(type_ebed, src, dst) # (E, C) @@ -388,8 +488,10 @@ def build_edge_cache( edge_env=edge_env, D_full=D_full, Dt_full=Dt_full, + D_packed=D_packed, edge_quat=edge_quat, deg_norm_floor=deg_norm_floor, + destinations_sorted=True, ) @@ -412,6 +514,8 @@ def build_edge_cache_from_edges( edge_type_keep_mask: EdgeTypeKeepMaskFn, random_gamma: bool, wigner_calc: WignerCalculatorFn, + packed_wigner_candidate: bool = False, + destinations_sorted: bool = False, build_wigner: bool = True, node_partial_exchange: Callable[[torch.Tensor], torch.Tensor] | None = None, ) -> EdgeFeatureCache: @@ -460,6 +564,13 @@ def build_edge_cache_from_edges( wigner_calc Callable that converts edge-aligned quaternions into packed Wigner-D blocks. + packed_wigner_candidate + Whether descriptor-level strict-FP32 K1 checks passed. Concrete edge + count and destination ordering are checked before panel generation. + destinations_sorted + Host-side provenance that ``edge_index[1]`` is nondecreasing. + build_wigner + Whether to materialize Wigner-D blocks for the SO(2) path. Returns ------- @@ -469,6 +580,11 @@ def build_edge_cache_from_edges( n_nodes = type_ebed.shape[0] src = edge_index[0].to(dtype=torch.long) dst = edge_index[1].to(dtype=torch.long) + if is_cute_infer_enabled() and destinations_sorted and dst.numel() > 1: + torch._assert_async( + torch.all(dst[1:] >= dst[:-1]), + "destinations_sorted=True requires nondecreasing destination indices", + ) # === Step 1. Normalize mask and apply type exclusions === edge_keep = edge_mask.to(dtype=torch.bool) @@ -503,8 +619,20 @@ def build_edge_cache_from_edges( eps=eps, random_gamma=random_gamma, wigner_calc=wigner_calc, + packed_wigner=( + build_wigner + and bridging_switch is None + and _packed_wigner_edges_eligible( + packed_wigner_candidate, + edge_count=dst.numel(), + node_count=n_nodes, + destinations_sorted=destinations_sorted, + runtime_dtypes=(edge_vec.dtype, edge_env.dtype, edge_rbf.dtype), + ) + ), build_full=build_wigner, ) # (E, D, D), (E, D, D), (E, 4) + D_full, Dt_full, D_packed = _separate_packed_wigner(D_full, Dt_full) # === Step 5. Edge type features === edge_type_feat = build_edge_type_feat(type_ebed, src, dst) @@ -537,9 +665,11 @@ def build_edge_cache_from_edges( edge_env=edge_env, D_full=D_full, Dt_full=Dt_full, + D_packed=D_packed, edge_quat=edge_quat, deg_norm_floor=deg_norm_floor, edge_src_gate=edge_src_gate, + destinations_sorted=destinations_sorted, ) @@ -551,6 +681,7 @@ def _build_edge_wigner( random_gamma: bool, wigner_calc: WignerCalculatorFn, build_full: bool = True, + packed_wigner: bool = False, ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]: """ Build packed Wigner-D blocks from edge vectors. @@ -573,13 +704,14 @@ def _build_edge_wigner( False (all message-passing blocks take the Cartesian path), only the quaternion is returned and the blocks are ``None``; the geometric initial embedding reconstructs the zonal coupling from the quaternion. + packed_wigner + Whether the exact packed K1 eligibility contract has passed. Returns ------- tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor] - Packed Wigner-D matrices ``(D_full, Dt_full)`` with shape ``(E, D, D)`` - (or ``None`` when ``build_full`` is False) and the quaternion used to - build them with shape ``(E, 4)``. + Wigner data with dense shape ``(E, D, D)``, packed shape ``(E, 46)``, + or ``None`` when ``build_full`` is false, plus the edge quaternion. """ # === Step 1. Build edge-aligned quaternions === edge_quat = build_edge_quaternion( @@ -600,10 +732,44 @@ def _build_edge_wigner( # === Step 3. Convert quaternions to packed Wigner-D blocks === if not build_full: return None, None, edge_quat + if packed_wigner: + from deepmd.kernels.cute.neo.k4_wignerd import ( + run_cute_wignerd, + ) + + cute_wigner = run_cute_wignerd( + edge_quat, + wigner_calc, + packed_wigner=packed_wigner, + ) + if cute_wigner is not None: + return cute_wigner[0], cute_wigner[1], edge_quat D_full, Dt_full = wigner_calc(edge_quat) return D_full, Dt_full, edge_quat +def _packed_wigner_edges_eligible( + candidate: bool, + *, + edge_count: int, + node_count: int, + destinations_sorted: bool, + runtime_dtypes: tuple[torch.dtype, ...] = (), +) -> bool: + """Finish packed eligibility from scalar shape and provenance metadata.""" + from deepmd.kernels.cute.neo.k1 import ( + packed_wigner_edges_eligible, + ) + + return packed_wigner_edges_eligible( + candidate=candidate, + edge_count=edge_count, + node_count=node_count, + destinations_sorted=destinations_sorted, + runtime_dtypes=runtime_dtypes, + ) + + def _finalize_edge_cache( *, n_nodes: int, @@ -615,9 +781,11 @@ def _finalize_edge_cache( edge_env: torch.Tensor, D_full: torch.Tensor | None, Dt_full: torch.Tensor | None, + D_packed: torch.Tensor | None, edge_quat: torch.Tensor, deg_norm_floor: float, edge_src_gate: torch.Tensor | None = None, + destinations_sorted: bool = False, ) -> EdgeFeatureCache: """ Assemble the shared `EdgeFeatureCache` layout. @@ -639,11 +807,14 @@ def _finalize_edge_cache( edge_env Smooth edge envelope weights with shape (E, 1). D_full - Packed Wigner-D matrices with shape (E, D, D), or None when the + Dense Wigner-D matrices with shape (E, D, D), or None when the full Wigner-D construction is skipped (all-Cartesian model). Dt_full - Transposed packed Wigner-D matrices with shape (E, D, D), or None + Transposed dense Wigner-D matrices with shape (E, D, D), or None when the full Wigner-D construction is skipped. + D_packed + Optional Neo packed Wigner panel with shape (E, 46). Dense consumers + continue to observe stable-rank ``D_full`` and ``Dt_full`` fields. edge_quat Global-to-local quaternions used to build the Wigner-D matrices with shape (E, 4). @@ -655,12 +826,17 @@ def _finalize_edge_cache( edge_src_gate Optional per-edge SFPG weight with shape (E, 1). ``None`` in non-bridging mode. + destinations_sorted + Host-side provenance that ``dst`` is nondecreasing. Returns ------- EdgeFeatureCache Finalized per-edge cache shared by eager and compile paths. """ + if deg_norm_floor <= 0.0: + raise ValueError("deg_norm_floor must be positive") + # === Step 1. Build smooth destination degrees === with nvtx_range("degree"): deg = torch.zeros(n_nodes, dtype=edge_vec.dtype, device=edge_vec.device) # (N,) @@ -680,10 +856,12 @@ def _finalize_edge_cache( inv_sqrt_deg=inv_sqrt_deg, D_full=D_full, Dt_full=Dt_full, + D_packed=D_packed, D_to_m_cache={}, Dt_from_m_cache={}, edge_src_gate=edge_src_gate, edge_quat=edge_quat, + destinations_sorted=destinations_sorted, ) @@ -734,10 +912,12 @@ def _get_empty_edge_cache( inv_sqrt_deg=inv_sqrt_deg, D_full=None, Dt_full=None, + D_packed=None, D_to_m_cache={}, Dt_from_m_cache={}, edge_src_gate=None, edge_quat=empty_quat, + destinations_sorted=True, ) @@ -893,16 +1073,20 @@ def edge_cache_to_dtype( # Use local variables with explicit None check and assignment. _D_full = cache.D_full _Dt_full = cache.Dt_full + _D_packed = cache.D_packed _edge_src_gate = cache.edge_src_gate _edge_quat = cache.edge_quat D_full: torch.Tensor | None = None Dt_full: torch.Tensor | None = None + D_packed: torch.Tensor | None = None edge_src_gate: torch.Tensor | None = None edge_quat: torch.Tensor | None = None if _D_full is not None: D_full = _D_full.to(dtype=dtype) if _Dt_full is not None: Dt_full = _Dt_full.to(dtype=dtype) + if _D_packed is not None: + D_packed = _D_packed.to(dtype=dtype) if _edge_src_gate is not None: edge_src_gate = _edge_src_gate.to(dtype=dtype) if _edge_quat is not None: @@ -919,8 +1103,10 @@ def edge_cache_to_dtype( inv_sqrt_deg=cache.inv_sqrt_deg.to(dtype=dtype), D_full=D_full, Dt_full=Dt_full, + D_packed=D_packed, D_to_m_cache=None if cache.D_to_m_cache is None else {}, Dt_from_m_cache=None if cache.Dt_from_m_cache is None else {}, edge_src_gate=edge_src_gate, edge_quat=edge_quat, + destinations_sorted=cache.destinations_sorted, ) diff --git a/deepmd/pt/model/descriptor/sezm_nn/embedding.py b/deepmd/pt/model/descriptor/sezm_nn/embedding.py index b70357a943..dbdb63492d 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/embedding.py +++ b/deepmd/pt/model/descriptor/sezm_nn/embedding.py @@ -22,6 +22,10 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) +from deepmd.kernels.cute.neo.k1_wigner_layout import ( + PACKED_VALUE_COUNT, + ZONAL_PANEL_OFFSETS, +) from deepmd.pt.model.network.mlp import ( MLPLayer, ) @@ -196,6 +200,16 @@ def __init__( node_radial_l_index, persistent=True, ) + packed_zonal_offsets = torch.tensor( + ZONAL_PANEL_OFFSETS if self.lmax == 3 else (), + device=self.device, + dtype=torch.long, + ) + self.register_buffer( + "packed_zonal_offsets", + packed_zonal_offsets, + persistent=False, + ) # The l=1 coefficients (packed rows 1..3) are the first three entries of # the non-scalar sequence ``node_row_index = [1, 2, ..., D-1]``, so the # native neighbor-spin l=1 message folds in at these local positions. @@ -238,27 +252,55 @@ def forward( torch.Tensor Initial features to add with shape (N, D, C). l=0 is guaranteed zero. """ - # === Step 1. Initialize output === + # === Step 1. Validate the non-scalar contract === device = edge_cache.edge_vec.device dtype = edge_cache.edge_vec.dtype - out = torch.zeros( - n_nodes, self.ebed_dim, self.channels, device=device, dtype=dtype - ) # (N, D, C) if self.lmax == 0: - return out + return torch.zeros( + n_nodes, self.ebed_dim, self.channels, device=device, dtype=dtype + ) # === Step 2. Gather all m=0 columns (l >= 1) in one shot === # Advanced indexing pairs one packed non-scalar row with the zonal m=0 column # from the same degree block in Dt_full. if zonal_coupling is None: - Dt_full = edge_cache.Dt_full # (E, D, D) - zonal_coupling = Dt_full[ - :, - self.non_scalar_row_index, - self.zonal_m0_col_index_for_row, - ] # (E, D-1) - - # === Step 3. Broadcast radial features per row === + D_packed = edge_cache.D_packed + if D_packed is not None: + if self.lmax != 3 or D_packed.shape[1] != PACKED_VALUE_COUNT: + raise ValueError("packed Wigner zonal coupling requires Neo lmax=3") + zonal_coupling = D_packed.index_select( + 1, + self.packed_zonal_offsets, + ) + else: + Dt_full = edge_cache.Dt_full # (E, D, D) + if Dt_full is None: + raise RuntimeError("GIE requires dense or packed Wigner storage") + zonal_coupling = Dt_full[ + :, + self.non_scalar_row_index, + self.zonal_m0_col_index_for_row, + ] # (E, D-1) + + # === Step 3. Optional fused message construction and reduction === + if not self.training: + from deepmd.kernels.cute.neo.gie import ( + is_cute_gie_enabled, + maybe_run_cute_gie, + ) + + if is_cute_gie_enabled(device) and spin_l1_message is None: + cute_out = maybe_run_cute_gie( + self, + n_nodes=n_nodes, + edge_cache=edge_cache, + radial_feat=radial_feat, + zonal_coupling=zonal_coupling, + ) + if cute_out is not None: + return cute_out + + # === Step 4. Eager fallback: broadcast radial features per row === # Each non-scalar packed row reuses the radial feature of its degree l. radial_value_for_row = radial_feat.index_select( 1, self.radial_slot_index_for_row @@ -276,7 +318,7 @@ def forward( 1, self.l1_local_index, spin_l1_message ) - # === Step 4. Source Freeze Propagation Gate (optional) === + # === Step 5. Source Freeze Propagation Gate (optional) === # Mute messages emitted by nodes whose local neighborhood enters # the frozen zone. ``edge_src_gate`` is ``None`` outside bridging # mode so this is a no-op in normal training. @@ -286,7 +328,10 @@ def forward( dtype=non_scalar_message.dtype ).unsqueeze(-1) - # === Step 5. Scatter to nodes and normalize === + # === Step 6. Scatter to nodes and normalize === + out = torch.zeros( + n_nodes, self.ebed_dim, self.channels, device=device, dtype=dtype + ) # (N, D, C) # Avoid advanced-index writeback (out[:, non_scalar_row_index, :]) which produces a copy. non_scalar_out = out.new_zeros( n_nodes, self.non_scalar_row_index.numel(), self.channels diff --git a/deepmd/pt/model/descriptor/sezm_nn/grid_net.py b/deepmd/pt/model/descriptor/sezm_nn/grid_net.py index a2e5efdd4f..60cd916f33 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/grid_net.py +++ b/deepmd/pt/model/descriptor/sezm_nn/grid_net.py @@ -30,6 +30,7 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) +from deepmd.kernels.cute.neo import runtime_policy as cute_runtime_policy from deepmd.pt.utils import ( env, ) @@ -116,6 +117,13 @@ def _project_frames( return projected.reshape(n_batch, coeff_dim, n_focus, -1) +def _inference_mode_is_frozen(module: nn.Module) -> bool: + """Return whether first-order inference-only shortcuts are safe.""" + return not module.training and not any( + parameter.requires_grad for parameter in module.parameters() + ) + + class GridProduct(nn.Module): """Parameter-free quadratic grid product ``u(g) * v(g)``.""" @@ -204,6 +212,8 @@ def forward( *, to_grid: Callable[[torch.Tensor], torch.Tensor], from_grid: Callable[[torch.Tensor], torch.Tensor], + grid_product: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] + | None = None, ) -> torch.Tensor: """ Apply the polynomial point-wise MLP on coefficient operands. @@ -221,6 +231,8 @@ def forward( Invariant routing signal; unused on this path. to_grid, from_grid : Callable Coefficient/grid projectors supplied by the owning grid net. + grid_product : Callable, optional + Fused replacement for the middle grid projection/product path. Returns ------- @@ -240,7 +252,10 @@ def forward( right = _project_frames(right, self.right_proj, self.n_frames) # === Step 2. Quadratic product on the grid, projected back === - coeff = from_grid(to_grid(left) * to_grid(right)) + if grid_product is None: + coeff = from_grid(to_grid(left) * to_grid(right)) + else: + coeff = grid_product(left, right) return _project_frames(coeff, self.out_proj, self.n_frames) @@ -310,6 +325,8 @@ def forward( *, to_grid: Callable[[torch.Tensor], torch.Tensor], from_grid: Callable[[torch.Tensor], torch.Tensor], + grid_product: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] + | None = None, ) -> torch.Tensor: """ Apply scalar-routed grid branch mixing on coefficient operands. @@ -322,6 +339,8 @@ def forward( Invariant router source with shape ``(N, F, 2*C)``. to_grid, from_grid : Callable Coefficient/grid projectors supplied by the owning grid net. + grid_product : Callable, optional + Fused replacement for the single-branch middle grid product. Returns ------- @@ -332,6 +351,17 @@ def forward( left = _project_frames(left, self.left_proj, self.n_frames) right = _project_frames(right, self.right_proj, self.n_frames) + if ( + self.n_branches == 1 + and cute_runtime_policy.is_cute_infer_enabled() + and _inference_mode_is_frozen(self) + ): + if grid_product is None: + coeff = from_grid(to_grid(left) * to_grid(right)) + else: + coeff = grid_product(left, right) + return _project_frames(coeff, self.out_proj, self.n_frames) + # === Step 2. Quadratic branches on the grid, routed by scalars === value = to_grid(left) * to_grid(right) # (N, G, F, N_branches * C) n_batch, n_grid, n_focus, _ = value.shape @@ -575,13 +605,26 @@ def forward( input_dtype = query.dtype query_ndfc, shape_info = self._to_ndfc(query) left, right, scalar_pair = self._prepare_pair(query_ndfc, context) - coeff_out = self.grid_op( - left.to(dtype=self.dtype), - right.to(dtype=self.dtype), - scalar_pair, - to_grid=self._to_grid, - from_grid=self._from_grid, - ) + if isinstance(self.grid_op, (GridMLP, GridBranch)): + grid_product = ( + self._grid_product if _inference_mode_is_frozen(self) else None + ) + coeff_out = self.grid_op( + left.to(dtype=self.dtype), + right.to(dtype=self.dtype), + scalar_pair, + to_grid=self._to_grid, + from_grid=self._from_grid, + grid_product=grid_product, + ) + else: + coeff_out = self.grid_op( + left.to(dtype=self.dtype), + right.to(dtype=self.dtype), + scalar_pair, + to_grid=self._to_grid, + from_grid=self._from_grid, + ) coeff_out = self._apply_scalar_path(coeff_out, scalar_pair) coeff_out = self._contract_frames(coeff_out) coeff_out = self._apply_residual_scale(coeff_out) @@ -717,6 +760,26 @@ def _from_grid(self, grid: torch.Tensor) -> torch.Tensor: coeff = torch.einsum("dkg,ngfc->ndfkc", from_grid, grid) return coeff.reshape(n_batch, coeff_dim, n_focus, -1) + def _grid_product( + self, + left: torch.Tensor, + right: torch.Tensor, + ) -> torch.Tensor: + from deepmd.kernels.cute.neo.output_grid_product import ( + maybe_run_cute_output_grid_product, + ) + + candidate = maybe_run_cute_output_grid_product( + left, + right, + self.projector.to_grid_mat, + self.projector.from_grid_mat, + n_frames=self.n_frames, + ) + if candidate is not None: + return candidate + return self._from_grid(self._to_grid(left) * self._to_grid(right)) + def _to_ndfc(self, value: torch.Tensor) -> tuple[torch.Tensor, tuple[int, ...]]: # All grid operations run in the canonical ``(N, D, F, C)`` layout; the # ``fndc`` re-orientation folds the focus-major SO(2) mixing layout into the diff --git a/deepmd/pt/model/descriptor/sezm_nn/lora.py b/deepmd/pt/model/descriptor/sezm_nn/lora.py index 3db8befae2..ae3413a302 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/lora.py +++ b/deepmd/pt/model/descriptor/sezm_nn/lora.py @@ -487,7 +487,28 @@ def _clear_sezm_compile_cache(model: nn.Module) -> None: crash or silently skip LoRA parameters. Mirrors the pattern used in :meth:`SeZMModel.reset_head_for_mode`. """ + from deepmd.kernels.cute.neo.k1 import ( + invalidate_cute_k1_state, + ) + from deepmd.pt.model.model.sezm_model import ( + _clear_shared_sezm_compile_cache, + ) + + readout_invalidator = None + for m in model.modules(): + invalidate_cute_k1_state(m) + if hasattr(m, "_neo_sm80_readout_input_fold_cache"): + if readout_invalidator is None: + from deepmd.kernels.cute.neo.readout_l0 import ( + invalidate_neo_readout_input_fold, + ) + + readout_invalidator = invalidate_neo_readout_input_fold + readout_invalidator(m) + for name in tuple(vars(m)): + if name.startswith("_deepmd_cute_"): + delattr(m, name) core_cache = getattr(m, "compiled_core_compute_cache", None) if isinstance(core_cache, dict): core_cache.clear() @@ -501,6 +522,7 @@ def _clear_sezm_compile_cache(model: nn.Module) -> None: m._dens_compiled = False if hasattr(m, "_dens_pending_compile_t0"): m._dens_pending_compile_t0 = None + _clear_shared_sezm_compile_cache() def _swap_submodule(parent: nn.Module, attr: str, new_module: nn.Module) -> None: diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index 5d9d182d79..b6a918cc3e 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -258,7 +258,8 @@ * ``triton.cudagraphs=False`` cudagraphs capture autograd metadata only once. Higher-order gradients need fresh metadata per call, so cudagraphs would feed - stale autograd state into the second backward. + stale autograd state into the second backward. Packed K1 also retains + Python-owned forward state, so direct graph capture remains disabled. * ``max_fusion_size=8`` Caps kernel fusion complexity so Inductor's scheduler does not time out on the large edge-level reductions inside the @@ -526,6 +527,131 @@ SeZMModel_ = make_model(SeZMAtomicModel) +def _neo_cute_infer_enabled() -> bool: + """Return whether the model builder must destination-sort CuTe edges.""" + from deepmd.kernels.cute.neo.runtime_policy import ( + is_cute_infer_enabled, + ) + + return is_cute_infer_enabled() + + +@torch.compiler.assume_constant_result +def _neo_cute_nlist_eager_island_enabled(device: torch.device) -> bool: + """Return whether Neo requests an eager Toolkit-Ops neighbor-list call.""" + if not _neo_cute_infer_enabled() or device.type != "cuda": + return False + try: + compute_capability = tuple(torch.cuda.get_device_capability(device)) + except RuntimeError: + return False + from deepmd.kernels.cute.neo.runtime_policy import ( + is_k1_eager_island_enabled, + ) + + return is_k1_eager_island_enabled(compute_capability) + + +@torch.compiler.disable +def _build_neo_neighbor_list_eager_island( + builder: NeighborList, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None, + rcut: float, + sel: list[int], + *, + return_mode: str, +) -> Any: + """Build Neo neighbors outside Dynamo when the runtime policy requests it.""" + return builder.build( + coord, + atype, + box, + rcut, + sel, + return_mode=return_mode, + ) + + +def _build_neo_neighbor_list( + builder: NeighborList, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None, + rcut: float, + sel: list[int], + *, + return_mode: str, +) -> Any: + """Apply Neo runtime policy around a general neighbor-list strategy.""" + if isinstance(builder, NvNeighborList) and _neo_cute_nlist_eager_island_enabled( + coord.device + ): + return _build_neo_neighbor_list_eager_island( + builder, + coord, + atype, + box, + rcut, + sel, + return_mode=return_mode, + ) + return builder.build( + coord, + atype, + box, + rcut, + sel, + return_mode=return_mode, + ) + + +def _neo_cute_k1_requires_sorted_edges( + descriptor: Any, + *, + training: bool, + device: torch.device, +) -> bool: + """Return whether this descriptor can consume packed destination-sorted K1. + + ``forward_with_edges`` promotes coordinates and edge vectors to the + descriptor compute dtype before constructing Wigner data. Mirror that + post-cast contract here; model inputs may still be FP64 at this boundary. + """ + compute_dtype = descriptor.compute_dtype + return bool( + not training + and _neo_cute_infer_enabled() + and descriptor._packed_wigner_candidate( + device, + compute_dtype, + compute_dtype, + ) + ) + + +def _sort_edge_tensors_by_destination( + edge_index: torch.Tensor, + edge_vec: torch.Tensor, + edge_mask: torch.Tensor, + edge_scatter_index: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Stably sort aligned edge tensors by destination, then source.""" + if edge_index.shape[1] == 0: + return edge_index, edge_vec, edge_mask, edge_scatter_index + src = edge_index[0].to(dtype=torch.long) + dst = edge_index[1].to(dtype=torch.long) + source_stride = src.max().clamp_min(0) + 1 + permutation = torch.argsort(dst * source_stride + src, stable=True) + return ( + edge_index.index_select(1, permutation).contiguous(), + edge_vec.index_select(0, permutation).contiguous(), + edge_mask.index_select(0, permutation).contiguous(), + edge_scatter_index.index_select(1, permutation).contiguous(), + ) + + def _select_neighbor_builder(nf: int, device: torch.device) -> NeighborList: """Select the O(N) neighbor builder for the given batch shape and device. @@ -584,6 +710,13 @@ def _select_neighbor_builder(nf: int, device: torch.device) -> NeighborList: # knows which buffers were promoted and in what order. _SEZM_TASK_BUF_ORDER: dict[tuple[Any, ...], tuple[str, ...]] = {} + +def _clear_shared_sezm_compile_cache() -> None: + """Drop shared graphs that may contain frozen parameter constants.""" + _SEZM_COMPILE_CACHE.clear() + _SEZM_TASK_BUF_ORDER.clear() + + _ENV_BOOL_CHOICES = { "1": True, "true": True, @@ -731,6 +864,9 @@ def __init__( # Maps cache_key -> task_buf_order for this instance so forward() # knows which buffers to pass and in what order. object.__setattr__(self, "_task_buf_order_cache", {}) + self.register_load_state_dict_post_hook( + self._invalidate_compiled_state_after_load + ) # Training follows `use_compile`. Evaluation/inference samples env # policy at init time so path and precision stay fixed per model. @@ -766,6 +902,27 @@ def __init__( else None ) + @staticmethod + def _invalidate_compiled_state_after_load( + module: SeZMModel, + incompatible_keys: Any, + ) -> None: + """Discard graphs that may have captured state from before a load.""" + del incompatible_keys + module.compiled_core_compute_cache.clear() + module._task_buf_order_cache.clear() + object.__setattr__(module, "compiled_embedding", None) + object.__setattr__(module, "_embedding_task_buf_order", None) + object.__setattr__(module, "compiled_dens_compute", None) + module._dens_compiled = False + module._core_compute_pending_compile_t0 = None + module._core_compute_pending_compile_key = None + module._dens_pending_compile_t0 = None + # Shared callables can contain make_fx get_attr constants, including + # prepared CuTe readout folds. A checkpoint load therefore invalidates + # the process-level cache as well as this instance's local slots. + _clear_shared_sezm_compile_cache() + # ========================================================================= # Forward Methods # ========================================================================= @@ -1489,6 +1646,21 @@ def core_compute( descriptor_model = self.atomic_model.descriptor # === Step 1. Establish the force-autograd endpoint === + sort_edges_by_dst = _neo_cute_k1_requires_sorted_edges( + descriptor_model, + training=self.training, + device=edge_vec.device, + ) + if sort_edges_by_dst: + edge_index, edge_vec, edge_mask, edge_scatter_index = ( + _sort_edge_tensors_by_destination( + edge_index, + edge_vec, + edge_mask, + edge_scatter_index, + ) + ) + # Neighbor-list construction and periodic-image resolution are explicit # caller responsibilities. Once the edge displacements are supplied, # SeZM differentiates only the pure map ``(edge_vec, theta) -> E``. @@ -1559,6 +1731,7 @@ def core_compute( edge_index=edge_index, edge_vec=edge_vec, edge_mask=edge_mask, + edge_index_sorted_by_dst=sort_edges_by_dst, charge_spin=charge_spin, spin=spin, comm_dict=comm_dict, @@ -1739,10 +1912,16 @@ def core_compute_dens( descriptor_model = self.atomic_model.descriptor # === Step 1. Build compact sparse edges === + sort_edges_by_dst = _neo_cute_k1_requires_sorted_edges( + descriptor_model, + training=self.training, + device=extended_coord.device, + ) edge_index, edge_vec, edge_mask, _ = self.build_edge_list_from_nlist( extended_coord=extended_coord, nlist=nlist, mapping=mapping, + sort_by_destination=sort_edges_by_dst, ) # === Step 2. Force embedding === @@ -1760,6 +1939,7 @@ def core_compute_dens( edge_index=edge_index, edge_vec=edge_vec, edge_mask=edge_mask, + edge_index_sorted_by_dst=sort_edges_by_dst, force_embedding=force_embedding, charge_spin=charge_spin, ) @@ -1957,6 +2137,51 @@ def trace_and_compile( ) return + # Register Python-owned K1 state before make_fx starts. The opt-in thin + # path can then keep adjacent linears in this graph while the CuTe work + # remains opaque behind its existing custom op. + from deepmd.kernels.cute.neo import runtime_policy as cute_runtime_policy + + compute_capability = ( + tuple(torch.cuda.get_device_capability(coord.device)) + if coord.device.type == "cuda" + else None + ) + if not self.training and cute_runtime_policy.is_cute_infer_enabled(): + from deepmd.kernels.cute.neo.readout_l0 import ( + maybe_prepare_sm80_readout_input_fold, + ) + + maybe_prepare_sm80_readout_input_fold( + self.atomic_model.descriptor.output_ffn, + compute_capability, + ) + prepared_cute_k1 = False + if ( + not self.training + and compute_capability is not None + and cute_runtime_policy.is_cute_infer_enabled() + and cute_runtime_policy.is_supported_k1_capability(compute_capability) + ): + from deepmd.kernels.cute.neo.k1 import ( + prepare_cute_k1_blocks, + ) + + prepared_cute_k1 = prepare_cute_k1_blocks( + self.atomic_model.descriptor.blocks, + training=self.training, + device=coord.device, + dtype=coord.dtype, + ) + if prepared_cute_k1 and cute_runtime_policy.is_k1_thin_wrapper_enabled( + compute_capability + ): + from ..network.mlp import ( + enable_neo_cute_compile_visible_linears, + ) + + enable_neo_cute_compile_visible_linears(self) + log.info( "SeZM: start tracing and compiling (mode=%s, coord_corr=%s)", mode, @@ -2768,8 +2993,10 @@ def build_neighbor_list( ) -> EdgeNeighborList: """Build the unified edge-vector schema for the ``forward`` entry.""" nf, nloc = atype.shape[:2] - return _select_neighbor_builder(nf, coord.device).build( - coord.view(nf, nloc, 3), + coord = coord.view(nf, nloc, 3) + return _build_neo_neighbor_list( + _select_neighbor_builder(nf, coord.device), + coord, atype, box, self.get_rcut(), @@ -2795,8 +3022,10 @@ def build_extended_neighbor_list( contract order ``(extended_coord, extended_atype, nlist, mapping)``. """ nf, nloc = atype.shape[:2] - return _select_neighbor_builder(nf, coord.device).build( - coord.view(nf, nloc, 3), + coord = coord.view(nf, nloc, 3) + return _build_neo_neighbor_list( + _select_neighbor_builder(nf, coord.device), + coord, atype, box, self.get_rcut(), @@ -2810,6 +3039,7 @@ def build_edge_list_from_nlist( extended_coord: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None, + sort_by_destination: bool | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """ Build a compact edge list from a DeePMD padded neighbor list. @@ -2834,6 +3064,9 @@ def build_edge_list_from_nlist( DeePMD padded neighbor list with shape (nf, nloc, nsel). mapping Extended-to-local mapping with shape (nf, nall), or ``None``. + sort_by_destination + Whether to sort edges by destination and source. ``None`` preserves + the environment-controlled CuTe inference behavior. Returns ------- @@ -2860,12 +3093,24 @@ def build_edge_list_from_nlist( nlist, mapping, ) - return ( + edge_tensors = ( edge_schema.edge_index, edge_schema.edge_vec, edge_schema.edge_mask, edge_schema.edge_scatter_index, ) + should_sort = ( + _neo_cute_k1_requires_sorted_edges( + self.atomic_model.descriptor, + training=self.training, + device=extended_coord.device, + ) + if sort_by_destination is None + else bool(sort_by_destination) + ) + if should_sort: + return _sort_edge_tensors_by_destination(*edge_tensors) + return edge_tensors # ========================================================================= # Input Canonicalization diff --git a/deepmd/pt/model/network/mlp.py b/deepmd/pt/model/network/mlp.py index 13ea438f4f..b342fe75b6 100644 --- a/deepmd/pt/model/network/mlp.py +++ b/deepmd/pt/model/network/mlp.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import os from typing import ( Any, ClassVar, @@ -46,6 +47,41 @@ def empty_t(shape: tuple[int, ...], precision: torch.dtype) -> torch.Tensor: return torch.empty(shape, dtype=precision, device=device) +@torch.compiler.assume_constant_result +def _use_k1_compile_visible_linear( + input_device: torch.device | None = None, +) -> bool: + """Keep the SM80 linear topology stable for one compiled graph.""" + truthy = {"1", "true", "yes", "on"} + falsy = {"0", "false", "no", "off"} + cute_enabled = os.environ.get("DP_NEO_CUTE_INFER", "").strip().lower() + if cute_enabled not in truthy: + return False + thin_enabled = os.environ.get("DP_CUTE_K1_THIN_WRAPPER", "").strip().lower() + if thin_enabled in falsy: + return False + if thin_enabled in truthy: + return True + if input_device is not None and input_device.type != "cuda": + return False + if not torch.cuda.is_available(): + return False + try: + return tuple(torch.cuda.get_device_capability(input_device)) == (8, 0) + except RuntimeError: + return False + + +def _matmul_bias( + value: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, +) -> torch.Tensor: + """Avoid eager addmm's expanded-bias copy and expose the add to Inductor.""" + output = torch.matmul(value, weight) + return output if bias is None else output + bias + + class Identity(nn.Module): def __init__(self) -> None: super().__init__() @@ -86,6 +122,7 @@ def __init__( ) -> None: super().__init__() self.trainable = trainable + self._deepmd_cute_compile_visible_linear = False # only use_timestep when skip connection is established. self.use_timestep = use_timestep and ( num_out == num_in or num_out == num_in * 2 @@ -204,7 +241,20 @@ def forward( ori_prec = xx.dtype if not env.DP_DTYPE_PROMOTION_STRICT: xx = xx.to(self.prec) - yy = F.linear(xx, self.matrix.t(), self.bias) + if torch.jit.is_scripting(): + yy = F.linear(xx, self.matrix.t(), self.bias) + elif ( + not self.training + and xx.dtype == torch.float32 + and self.matrix.dtype == torch.float32 + and (self.bias is None or self.bias.dtype == torch.float32) + and not torch.is_autocast_enabled(xx.device.type) + and self._deepmd_cute_compile_visible_linear + and _use_k1_compile_visible_linear(xx.device) + ): + yy = _matmul_bias(xx, self.matrix, self.bias) + else: + yy = F.linear(xx, self.matrix.t(), self.bias) yy = self.activate(yy) yy = yy * self.idt if self.idt is not None else yy if self.resnet: @@ -278,6 +328,13 @@ def check_load_param(ss: str) -> nn.Parameter | None: return obj +def enable_neo_cute_compile_visible_linears(module: nn.Module) -> None: + """Select the alternate eval linear topology only inside one Neo model.""" + for child in module.modules(): + if isinstance(child, MLPLayer): + child._deepmd_cute_compile_visible_linear = True + + MLP_ = make_multilayer_network(MLPLayer, nn.Module) diff --git a/deepmd/pt/utils/nv_nlist.py b/deepmd/pt/utils/nv_nlist.py index 3c68251a7d..5b2b13ef33 100644 --- a/deepmd/pt/utils/nv_nlist.py +++ b/deepmd/pt/utils/nv_nlist.py @@ -178,6 +178,27 @@ def build( ``return_mode='edges'`` does not support ``pair_excl``; a :class:`NotImplementedError` is raised in that combination. """ + return self._build_impl( + coord, + atype, + box, + rcut, + sel, + return_mode=return_mode, + pair_excl=pair_excl, + ) + + def _build_impl( + self, + coord: Any, + atype: Any, + box: Any, + rcut: float, + sel: list[int], + return_mode: str = "extended", + pair_excl: PairExcludeMask | None = None, + ) -> tuple[Any, Any, Any, Any] | EdgeNeighborList: + """Implement the complete Toolkit-Ops build behind the dispatch guard.""" if return_mode == "edges" and pair_excl is not None: raise NotImplementedError( "pair_excl is not supported with return_mode='edges'; " diff --git a/doc/install/easy-install.md b/doc/install/easy-install.md index 1eef898316..69b0866cfe 100644 --- a/doc/install/easy-install.md +++ b/doc/install/easy-install.md @@ -189,6 +189,20 @@ pip install deepmd-kit :::::: +Optional CuTe inference kernels are available on Linux for the PyTorch backend: + +```bash +pip install "deepmd-kit[torch,cute]" +``` + +The `cute` extra installs the CUTLASS CuTe DSL and NVIDIA Alchemi Toolkit-Ops +runtime used by the optional CuTe paths. Set `DP_NEO_CUTE_INFER=1` to enable the +DPA4-Neo K1 implementation; it may coexist with `DP_TRITON_INFER=2`. +`DP_NEO_CUTE_INFER` and `DP_CUTE_INFER` are independent: the latter controls the +inner SO(2) value-path implementation. Requesting a CuTe path without its +runtime dependencies raises an error rather than silently selecting another +implementation. + The supported platform includes Linux x86-64 and aarch64 with GNU C Library 2.28 or above, macOS x86-64 and arm64, and Windows x86-64. > [!WARNING] diff --git a/pyproject.toml b/pyproject.toml index 3b64d3ec1b..dc4e9228dd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -145,6 +145,10 @@ cu12 = [ "nvidia-cudnn-cu12", "nvidia-cuda-nvcc-cu12", ] +cute = [ + 'nvidia-cutlass-dsl[cu13]>=4.6.1,<5; platform_system == "Linux"', + 'nvalchemi-toolkit-ops>=0.3.1; python_version >= "3.11" and platform_system == "Linux"', +] jax = [ # below is a funny workaround for # https://github.com/astral-sh/uv/issues/8601 diff --git a/source/tests/pt/model/test_descriptor_sezm_block_ffn_grid_fusion.py b/source/tests/pt/model/test_descriptor_sezm_block_ffn_grid_fusion.py new file mode 100644 index 0000000000..16a273bb19 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_block_ffn_grid_fusion.py @@ -0,0 +1,377 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Focused contracts for Neo's C=96 block-FFN grid fusion.""" + +from __future__ import ( + annotations, +) + +import pytest +import torch + +from deepmd.pt.model.descriptor.sezm import ( + DescrptSeZM, +) +from deepmd.pt.model.descriptor.sezm_nn.grid_net import ( + GridBranch, + GridMLP, + S2GridNet, +) + + +def test_single_branch_bypasses_router_and_accepts_fused_middle( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + branch = ( + GridBranch( + channels=2, + n_branches=1, + n_frames=3, + dtype=torch.float32, + trainable=False, + seed=7, + ) + .to("cpu") + .eval() + ) + monkeypatch.setattr( + branch.router, + "forward", + lambda value: pytest.fail("single-branch router must not run"), + ) + calls = 0 + + def fused_middle( + left: torch.Tensor, + right: torch.Tensor, + ) -> torch.Tensor: + nonlocal calls + calls += 1 + return left * right + + left = torch.randn(2, 4, 1, 6, device="cpu") + out = branch( + left, + torch.randn_like(left), + torch.randn(2, 1, 4, device="cpu"), + to_grid=lambda value: value, + from_grid=lambda value: value, + grid_product=fused_middle, + ) + + assert calls == 1 + assert out.shape == left.shape + + +def test_single_branch_keeps_pytorch_middle_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + branch = ( + GridBranch( + channels=2, + n_branches=1, + n_frames=1, + dtype=torch.float32, + trainable=False, + seed=11, + ) + .to("cpu") + .eval() + ) + monkeypatch.setattr( + branch.router, + "forward", + lambda value: pytest.fail("single-branch router must not run"), + ) + calls = {"to_grid": 0, "from_grid": 0} + + def to_grid(value: torch.Tensor) -> torch.Tensor: + calls["to_grid"] += 1 + return value + + def from_grid(value: torch.Tensor) -> torch.Tensor: + calls["from_grid"] += 1 + return value + + left = torch.randn(2, 4, 1, 2, device="cpu") + out = branch( + left, + torch.randn_like(left), + torch.randn(2, 1, 4, device="cpu"), + to_grid=to_grid, + from_grid=from_grid, + ) + + assert calls == {"to_grid": 2, "from_grid": 1} + assert out.shape == left.shape + + +def test_single_branch_opt_in_matches_routed_result( + monkeypatch: pytest.MonkeyPatch, +) -> None: + branch = ( + GridBranch( + channels=2, + n_branches=1, + n_frames=1, + dtype=torch.float32, + trainable=False, + seed=17, + ) + .to("cpu") + .eval() + ) + left = torch.randn(2, 4, 1, 2, device="cpu") + right = torch.randn_like(left) + scalar_pair = torch.randn(2, 1, 4, device="cpu") + original_forward = branch.router.forward + router_calls = 0 + + def tracked_router(value: torch.Tensor) -> torch.Tensor: + nonlocal router_calls + router_calls += 1 + return original_forward(value) + + monkeypatch.setattr(branch.router, "forward", tracked_router) + monkeypatch.delenv("DP_NEO_CUTE_INFER", raising=False) + routed = branch( + left, + right, + scalar_pair, + to_grid=lambda value: value, + from_grid=lambda value: value, + ) + monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + shortcut = branch( + left, + right, + scalar_pair, + to_grid=lambda value: value, + from_grid=lambda value: value, + ) + + assert router_calls == 1 + torch.testing.assert_close(shortcut, routed) + + +def test_multi_branch_preserves_softmax_router( + monkeypatch: pytest.MonkeyPatch, +) -> None: + branch = GridBranch( + channels=2, + n_branches=2, + n_frames=1, + dtype=torch.float32, + trainable=False, + seed=13, + ).to("cpu") + original_softmax = torch.softmax + softmax_calls = 0 + + def tracked_softmax( + value: torch.Tensor, + dim: int, + ) -> torch.Tensor: + nonlocal softmax_calls + softmax_calls += 1 + return original_softmax(value, dim=dim) + + monkeypatch.setattr(torch, "softmax", tracked_softmax) + left = torch.randn(2, 4, 1, 2, device="cpu") + out = branch( + left, + torch.randn_like(left), + torch.randn(2, 1, 4, device="cpu"), + to_grid=lambda value: value, + from_grid=lambda value: value, + grid_product=lambda left, right: pytest.fail( + "multi-branch routing must keep the generic path" + ), + ) + + assert softmax_calls == 1 + assert out.shape == left.shape + assert torch.isfinite(out).all() + + +def test_single_branch_training_preserves_router_zero_gradient( + monkeypatch: pytest.MonkeyPatch, +) -> None: + branch = ( + GridBranch( + channels=2, + n_branches=1, + n_frames=1, + dtype=torch.float32, + trainable=True, + seed=17, + ) + .to("cpu") + .train() + ) + original_softmax = torch.softmax + softmax_calls = 0 + + def tracked_softmax( + value: torch.Tensor, + dim: int, + ) -> torch.Tensor: + nonlocal softmax_calls + softmax_calls += 1 + return original_softmax(value, dim=dim) + + monkeypatch.setattr(torch, "softmax", tracked_softmax) + left = torch.randn(2, 4, 1, 2, device="cpu", requires_grad=True) + right = torch.randn_like(left) + out = branch( + left, + right, + torch.randn(2, 1, 4, device="cpu"), + to_grid=lambda value: value, + from_grid=lambda value: value, + grid_product=lambda left, right: pytest.fail( + "training must keep the generic router path" + ), + ) + out.sum().backward() + + assert softmax_calls == 1 + assert branch.router.weight.grad is not None + assert torch.count_nonzero(branch.router.weight.grad) == 0 + + +def test_single_branch_eval_with_trainable_parameters_preserves_router_gradient( + monkeypatch: pytest.MonkeyPatch, +) -> None: + branch = ( + GridBranch( + channels=2, + n_branches=1, + n_frames=1, + dtype=torch.float32, + trainable=True, + seed=19, + ) + .to("cpu") + .eval() + ) + left = torch.randn(2, 4, 1, 2, device="cpu", requires_grad=True) + out = branch( + left, + torch.randn_like(left), + torch.randn(2, 1, 4, device="cpu"), + to_grid=lambda value: value, + from_grid=lambda value: value, + grid_product=lambda left, right: pytest.fail( + "trainable eval must preserve the generic router path" + ), + ) + out.sum().backward() + + assert branch.router.weight.grad is not None + assert torch.count_nonzero(branch.router.weight.grad) == 0 + + +@pytest.mark.parametrize("training", (False, True)) +def test_grid_net_with_trainable_parameters_does_not_offer_fused_product( + monkeypatch: pytest.MonkeyPatch, + training: bool, +) -> None: + net = S2GridNet( + lmax=1, + mmax=1, + channels=2, + n_focus=1, + mode="self", + op_type="mlp", + dtype=torch.float32, + layout="ndfc", + coefficient_layout="packed", + grid_method="e3nn", + trainable=True, + seed=23, + ).to("cpu") + net.train(training) + monkeypatch.setattr( + net, + "_grid_product", + lambda left, right: pytest.fail( + "trainable grid net must not offer its first-order-only fused product" + ), + ) + + query = torch.randn(2, 4, 1, 4, device="cpu", requires_grad=True) + net(query).sum().backward() + assert query.grad is not None + + +def test_frozen_eval_grid_net_offers_fused_product( + monkeypatch: pytest.MonkeyPatch, +) -> None: + net = ( + S2GridNet( + lmax=1, + mmax=1, + channels=2, + n_focus=1, + mode="self", + op_type="mlp", + dtype=torch.float32, + layout="ndfc", + coefficient_layout="packed", + grid_method="e3nn", + trainable=False, + seed=29, + ) + .to("cpu") + .eval() + ) + calls = 0 + + def tracked_product(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: + nonlocal calls + calls += 1 + return net._from_grid(net._to_grid(left) * net._to_grid(right)) + + monkeypatch.setattr(net, "_grid_product", tracked_product) + output = net(torch.randn(2, 4, 1, 4, device="cpu")) + + assert calls == 1 + assert output.shape == (2, 4, 1, 2) + + +def test_exact_neo_has_two_c96_block_products_and_c192_readout() -> None: + descriptor = DescrptSeZM( + ntypes=2, + sel=4, + channels=32, + lmax=3, + mmax=1, + n_blocks=2, + so2_layers=3, + n_focus=2, + message_node_so3=True, + ffn_neurons=0, + ffn_so3_grid=True, + grid_branch=[0, 0, 1], + ffn_blocks=1, + so3_readout="mlp", + use_amp=False, + precision="float32", + trainable=False, + seed=42, + ) + + assert len(descriptor.blocks) == 2 + for block in descriptor.blocks: + assert len(block.ffns) == 1 + grid_op = block.ffns[0].act.grid_op + assert isinstance(grid_op, GridBranch) + assert grid_op.n_branches == 1 + assert grid_op.channels == 96 + + readout_grid_op = descriptor.output_ffn.act.grid_op + assert isinstance(readout_grid_op, GridMLP) + assert readout_grid_op.hidden_channels == 192 diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_compile_cache.py b/source/tests/pt/model/test_descriptor_sezm_cute_compile_cache.py new file mode 100644 index 0000000000..6dde077517 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_compile_cache.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Behavioral tests for device-aware CuTe compile caching.""" + +from __future__ import ( + annotations, +) + +import importlib.util +import sys +from pathlib import ( + Path, +) +from types import ( + SimpleNamespace, +) +from unittest import ( + mock, +) + +REPO_ROOT = Path(__file__).resolve().parents[4] +CACHE_PATH = REPO_ROOT / "deepmd/kernels/cute/neo/compile_cache.py" + + +def _load_cache_module(): + assert CACHE_PATH.is_file(), f"CuTe compile cache is missing: {CACHE_PATH}" + name = "sezm_cute_compile_cache_test" + spec = importlib.util.spec_from_file_location(name, CACHE_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + try: + spec.loader.exec_module(module) + finally: + sys.modules.pop(name, None) + return module + + +def test_cache_separates_device_and_compute_capability() -> None: + cache_module = _load_cache_module() + identity = [(0, 8, 0)] + calls: list[tuple[str, tuple[int, int, int]]] = [] + + @cache_module.device_aware_lru_cache( + maxsize=8, + identity_getter=lambda: identity[0], + ) + def compile_kernel(mode: str): + calls.append((mode, identity[0])) + return object() + + sm80_first = compile_kernel("strict-fp32") + assert compile_kernel("strict-fp32") is sm80_first + + identity[0] = (1, 9, 0) + sm90 = compile_kernel("strict-fp32") + assert sm90 is not sm80_first + + identity[0] = (0, 8, 0) + assert compile_kernel("strict-fp32") is sm80_first + assert calls == [ + ("strict-fp32", (0, 8, 0)), + ("strict-fp32", (1, 9, 0)), + ] + + +def test_cache_exposes_standard_cache_controls() -> None: + cache_module = _load_cache_module() + + @cache_module.device_aware_lru_cache( + maxsize=2, + identity_getter=lambda: (0, 8, 6), + ) + def compile_kernel(rows: int): + return object() + + compile_kernel(4) + assert compile_kernel.cache_info().currsize == 1 + compile_kernel.cache_clear() + assert compile_kernel.cache_info().currsize == 0 + + +def test_cache_compiles_inside_the_keyed_cuda_device() -> None: + cache_module = _load_cache_module() + entered: list[tuple[str, int]] = [] + + class DeviceContext: + def __init__(self, index: int) -> None: + self.index = index + + def __enter__(self) -> None: + entered.append(("enter", self.index)) + + def __exit__(self, *_args) -> None: + entered.append(("exit", self.index)) + + fake_torch = SimpleNamespace( + cuda=SimpleNamespace( + is_available=lambda: True, + device=lambda index: DeviceContext(index), + ) + ) + + @cache_module.device_aware_lru_cache( + maxsize=2, + identity_getter=lambda: (3, 9, 0), + ) + def compile_kernel() -> object: + entered.append(("compile", 3)) + return object() + + with mock.patch.dict(sys.modules, {"torch": fake_torch}): + compile_kernel() + + assert entered == [("enter", 3), ("compile", 3), ("exit", 3)] diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_envelope_softmax.py b/source/tests/pt/model/test_descriptor_sezm_cute_envelope_softmax.py new file mode 100644 index 0000000000..06b1bdd3a9 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_envelope_softmax.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Numerical acceptance for the CuTe envelope-gated softmax.""" + +from __future__ import ( + annotations, +) + +import importlib + +import pytest +import torch + +from deepmd.pt.model.descriptor.sezm_nn.attention import ( + segment_envelope_gated_softmax, +) + + +def _runtime_skip_reason() -> str | None: + if not torch.cuda.is_available(): + return "CuTe envelope-softmax acceptance requires CUDA" + try: + importlib.import_module("cutlass.cute") + importlib.import_module("cuda.bindings.driver") + except Exception as exc: # pragma: no cover - runtime dependent + return f"CuTe DSL runtime unavailable: {exc}" + return None + + +_SKIP_REASON = _runtime_skip_reason() + + +@pytest.mark.skipif(_SKIP_REASON is not None, reason=_SKIP_REASON or "unavailable") +@pytest.mark.parametrize( + ("logit_scale", "edge_gate"), + [ + (1.0, 1.0), + (15.0, 1.0e-2), + (110.0, 1.0e-23), + (-120.0, 1.0), + ], +) +def test_cute_envelope_softmax_matches_eager_across_extreme_frames( + logit_scale: float, + edge_gate: float, +) -> None: + from deepmd.kernels.cute.neo.k1_kernels.cute_envelope_gated_softmax import ( + compile_envelope_softmax_forward, + ) + + logits = torch.tensor( + [ + [logit_scale, logit_scale - 0.5], + [logit_scale - 1.0, logit_scale - 1.5], + [logit_scale - 0.25, logit_scale - 0.75], + [logit_scale - 2.0, logit_scale - 2.5], + ], + device="cuda", + dtype=torch.float32, + ) + gate = torch.full((4,), edge_gate, device="cuda", dtype=torch.float32) + dst = torch.tensor([0, 0, 1, 1], device="cuda", dtype=torch.long) + dst_ptr = torch.tensor([0, 2, 4], device="cuda", dtype=torch.int32) + z_bias_raw = torch.tensor([0.1, -0.3], device="cuda", dtype=torch.float32) + eps = 1.0e-7 + + expected = segment_envelope_gated_softmax( + logits.view(4, 2, 1), + gate, + dst, + 2, + z_bias_raw.view(2, 1), + eps, + ).view(4, 2) + actual = torch.empty_like(logits) + group_max = torch.empty((2, 2), device="cuda", dtype=torch.float32) + denom = torch.empty_like(group_max) + run = compile_envelope_softmax_forward(128, eps) + run(logits, gate, dst_ptr, z_bias_raw, actual, group_max, denom) + torch.cuda.synchronize() + + torch.testing.assert_close(actual, expected, atol=5.0e-5, rtol=5.0e-5) + assert torch.isfinite(actual).all() diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_gie.py b/source/tests/pt/model/test_descriptor_sezm_cute_gie.py new file mode 100644 index 0000000000..4deb334291 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_gie.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Differential tests for the opt-in CuTe geometric initial embedding.""" + +from __future__ import ( + annotations, +) + +import importlib.util +import sys +import unittest +from functools import ( + lru_cache, +) +from pathlib import ( + Path, +) +from types import ( + SimpleNamespace, +) + +import torch + +REPO_ROOT = Path(__file__).resolve().parents[4] +GIE_PATH = REPO_ROOT / "deepmd/kernels/cute/neo/gie.py" + + +@lru_cache(maxsize=1) +def _load_gie_module(): + if not GIE_PATH.is_file(): + raise FileNotFoundError(GIE_PATH) + name = "sezm_cute_gie_contract_test" + spec = importlib.util.spec_from_file_location(name, GIE_PATH) + if spec is None or spec.loader is None: + raise RuntimeError(f"could not load {GIE_PATH}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + try: + spec.loader.exec_module(module) + finally: + sys.modules.pop(name, None) + return module + + +def _degree_slots(lmax: int, *, device: torch.device) -> torch.Tensor: + degrees = torch.arange(1, lmax + 1, device=device, dtype=torch.long) + return torch.repeat_interleave(degrees - 1, 2 * degrees + 1) + + +def _zonal_indices( + lmax: int, *, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor]: + rows = torch.arange(1, (lmax + 1) ** 2, device=device, dtype=torch.long) + degrees = torch.arange(1, lmax + 1, device=device, dtype=torch.long) + degree_for_row = torch.repeat_interleave(degrees, 2 * degrees + 1) + return rows, degree_for_row * (degree_for_row + 1) + + +def _materialized_reference( + radial: torch.Tensor, + zonal: torch.Tensor, + inv_sqrt_deg: torch.Tensor, + dst: torch.Tensor, + gate: torch.Tensor, + *, + n_nodes: int, + lmax: int, +) -> torch.Tensor: + slots = _degree_slots(lmax, device=radial.device) + message = zonal.unsqueeze(-1) * radial.index_select(1, slots) + if gate.numel() != 0: + message = message * gate.reshape(-1, 1, 1) + non_scalar = radial.new_zeros(n_nodes, zonal.shape[1], radial.shape[2]) + non_scalar.index_add_(0, dst, message) + out = radial.new_zeros(n_nodes, zonal.shape[1] + 1, radial.shape[2]) + out[:, 1:, :] = non_scalar + return out * inv_sqrt_deg + + +def _inputs( + *, + n_nodes: int, + dst_values: tuple[int, ...], + lmax: int, + channels: int, + device: torch.device, + with_gate: bool, +) -> tuple[torch.Tensor, ...]: + edge_count = len(dst_values) + generator = torch.Generator(device=device).manual_seed(20260703 + edge_count) + radial = torch.randn( + edge_count, + lmax, + channels, + generator=generator, + device=device, + dtype=torch.float32, + requires_grad=True, + ) + dense_dt = torch.randn( + edge_count, + (lmax + 1) ** 2, + (lmax + 1) ** 2, + generator=generator, + device=device, + dtype=torch.float32, + requires_grad=True, + ) + rows, cols = _zonal_indices(lmax, device=device) + zonal = dense_dt[:, rows, cols] + inv_sqrt_deg = ( + torch.rand( + n_nodes, + 1, + 1, + generator=generator, + device=device, + dtype=torch.float32, + ) + + 0.25 + ).requires_grad_(True) + dst = torch.tensor(dst_values, device=device, dtype=torch.long) + if with_gate: + gate = torch.rand( + edge_count, + 1, + generator=generator, + device=device, + dtype=torch.float32, + requires_grad=True, + ) + else: + gate = torch.empty(0, device=device, dtype=torch.float32) + return radial, dense_dt, zonal, inv_sqrt_deg, dst, gate + + +class TestSeZMCuTeGIEContract(unittest.TestCase): + def test_backward_compile_key_separates_destination_dtypes(self): + gie = _load_gie_module() + common = ((0, 8, 0), 3, 32, True, (128, 32, 1)) + + int32_key = gie._backward_compile_key(*common, torch.int32) + int64_key = gie._backward_compile_key(*common, torch.int64) + + self.assertNotEqual(int32_key, int64_key) + + def test_contract_requires_sorted_dynamic_strict_fp32_inputs(self): + gie = _load_gie_module() + radial, _dense_dt, zonal, inv_sqrt_deg, dst, gate = _inputs( + n_nodes=5, + dst_values=(0, 0, 1, 3, 3, 4), + lmax=3, + channels=4, + device=torch.device("cpu"), + with_gate=True, + ) + module = SimpleNamespace( + lmax=3, + channels=4, + training=False, + non_scalar_row_index=torch.arange(1, 16, device=torch.device("cpu")), + radial_slot_index_for_row=_degree_slots(3, device=torch.device("cpu")), + ) + cache = SimpleNamespace( + dst=dst, + inv_sqrt_deg=inv_sqrt_deg, + edge_src_gate=gate, + destinations_sorted=True, + ) + self.assertTrue(gie.validate_gie_contract(module, 5, cache, radial, zonal)) + + cache.destinations_sorted = False + self.assertFalse(gie.validate_gie_contract(module, 5, cache, radial, zonal)) + cache.destinations_sorted = True + self.assertFalse( + gie.validate_gie_contract(module, 5, cache, radial.double(), zonal) + ) + module.non_scalar_row_index = torch.arange(15, device=torch.device("cpu")) + self.assertFalse(gie.validate_gie_contract(module, 5, cache, radial, zonal)) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required") + def test_cuda_forward_and_wigner_radial_degree_gate_gradients(self): + gie = _load_gie_module() + if not gie.SEZM_CUTE_GIE_AVAILABLE: + self.skipTest("CuTe DSL is not available") + for n_nodes, dst_values, with_gate in ( + (4, (0, 0, 2, 2, 2, 3), False), + (7, (0, 1, 1, 1, 4, 6, 6, 6, 6), True), + ): + with self.subTest(n_nodes=n_nodes, edges=len(dst_values), gate=with_gate): + expected_inputs = _inputs( + n_nodes=n_nodes, + dst_values=dst_values, + lmax=3, + channels=32, + device=torch.device("cuda"), + with_gate=with_gate, + ) + radial, dense_dt, zonal, inv_sqrt_deg, dst, gate = expected_inputs + weight = torch.randn_like( + radial.new_empty(n_nodes, 16, radial.shape[2]) + ) + expected = _materialized_reference( + radial, + zonal, + inv_sqrt_deg, + dst, + gate, + n_nodes=n_nodes, + lmax=3, + ) + expected_grads = torch.autograd.grad( + (expected * weight).sum(), + (radial, dense_dt, inv_sqrt_deg, *([gate] if with_gate else [])), + ) + + actual_inputs = _inputs( + n_nodes=n_nodes, + dst_values=dst_values, + lmax=3, + channels=32, + device=torch.device("cuda"), + with_gate=with_gate, + ) + radial, dense_dt, zonal, inv_sqrt_deg, dst, gate = actual_inputs + actual = gie.gie_fused_cuda( + radial, + zonal, + inv_sqrt_deg, + dst, + gate, + n_nodes=n_nodes, + lmax=3, + ) + actual_grads = torch.autograd.grad( + (actual * weight).sum(), + (radial, dense_dt, inv_sqrt_deg, *([gate] if with_gate else [])), + ) + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + for actual_grad, expected_grad in zip( + actual_grads, expected_grads, strict=True + ): + torch.testing.assert_close( + actual_grad, expected_grad, rtol=2e-5, atol=2e-5 + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1.py new file mode 100644 index 0000000000..f5a4af7b35 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1.py @@ -0,0 +1,555 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Contracts for the opt-in CuTe Neo K1 inference path.""" + +from __future__ import ( + annotations, +) + +import importlib +from dataclasses import ( + fields, +) +from types import ( + SimpleNamespace, +) + +import pytest +import torch + +from deepmd.kernels.cute.neo import k1 as _K1 +from deepmd.kernels.cute.neo import k1_runner as _K1_RUNNER +from deepmd.kernels.cute.neo import ( + runtime_policy, +) +from deepmd.pt.model.descriptor.sezm_nn.edge_cache import ( + EdgeFeatureCache, + _build_edge_wigner, + _separate_packed_wigner, +) +from deepmd.pt.model.descriptor.sezm_nn.norm import ( + EquivariantRMSNorm, +) + +NeoK1BackwardWorkspace = _K1_RUNNER.NeoK1BackwardWorkspace +NeoK1RuntimeConfig = _K1.NeoK1RuntimeConfig +NeoFullCuteBackward = _K1_RUNNER.NeoFullCuteBackward +StackCache = _K1_RUNNER.StackCache +_validate_runtime_config = _K1_RUNNER._validate_runtime_config +_uses_packed_message_grid = _K1_RUNNER._uses_packed_message_grid + + +class _Identity: + pass + + +def _neo_like_block(**overrides): + frame_contract = { + "coefficient_layout": "packed", + "n_frames": 3, + "channels": 32, + } + message_node_grid_product = SimpleNamespace( + layout="flat", + mode="cross", + op_type="glu", + n_focus=2, + n_frames=3, + channels=32, + dtype=torch.float32, + frame_expand=SimpleNamespace(**frame_contract), + frame_contract=SimpleNamespace(**frame_contract), + ) + so2 = SimpleNamespace( + lmax=3, + mmax=1, + ebed_dim_full=16, + reduced_dim=10, + channels=32, + n_focus=2, + so2_focus_dim=32, + hidden_channels=64, + mixing_layers=3, + n_atten_head=1, + head_dim=32, + radial_so2_mode="degree_channel", + radial_so2_rank=1, + so2_norm=False, + focus_compete=True, + focus_norm=True, + edge_cartesian=False, + node_cartesian_tp=None, + message_node_grid_product=message_node_grid_product, + atten_f_mix=False, + attn_v_proj=None, + attn_o_proj=None, + mlp_bias=False, + layer_scale=False, + use_so2_attn_res=False, + ) + for key, value in overrides.items(): + setattr(so2, key, value) + block = torch.nn.Module() + block.so2_conv = so2 + block.lmax = 3 + block.node_lmax = 3 + block.pre_so2_norm = _Identity() + block.post_so2_norm = EquivariantRMSNorm( + 3, + 32, + dtype=torch.float32, + trainable=False, + ) + block.runtime_weight = torch.nn.Parameter( + torch.ones(1, dtype=torch.float32, device="cpu"), + requires_grad=False, + ) + return block + + +def test_exact_neo_contract_is_supported() -> None: + block = _neo_like_block() + + assert _K1.get_neo_k1_spec(block).is_current_neo_target + assert _K1.is_supported_neo_k1_block(block) + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("lmax", 4), + ("mmax", 2), + ("channels", 64), + ("n_focus", 1), + ("mixing_layers", 2), + ("focus_compete", False), + ("edge_cartesian", True), + ], +) +def test_non_neo_contracts_fall_back(field: str, value: object) -> None: + assert not _K1.is_supported_neo_k1_block(_neo_like_block(**{field: value})) + + +def test_unsupported_message_grid_contract_falls_back() -> None: + block = _neo_like_block() + block.so2_conv.message_node_grid_product.op_type = "mlp" + + assert not _K1.is_supported_neo_k1_block(block) + + +def test_unsupported_post_norm_falls_back() -> None: + block = _neo_like_block() + block.post_so2_norm = torch.nn.LayerNorm(32, device="cpu") + + assert not _K1.is_supported_neo_k1_block(block) + + +def test_gate_expand_contract_cache_tracks_buffer_versions() -> None: + block = _neo_like_block() + expand_index = torch.tensor( + [0, 1, 2, 0, 1, 2, 0, 1, 2], + dtype=torch.long, + device="cpu", + ) + block.so2_conv.non_linearities = [SimpleNamespace(expand_index=expand_index)] + + assert _K1._gate_expand_index_structure_is_supported(block) + assert _K1._gate_expand_index_is_supported(block) + expand_index[0] = 2 + assert _K1._gate_expand_index_structure_is_supported(block) + assert not _K1._gate_expand_index_is_supported(block) + + +@pytest.mark.parametrize( + "capability", + sorted(runtime_policy.SUPPORTED_K1_CAPABILITIES), +) +def test_every_supported_architecture_has_a_validated_config( + capability: tuple[int, int], +) -> None: + config = _K1._architecture_default_config(capability) + + assert _validate_runtime_config(config, compute_capability=capability) is None + + +@pytest.mark.parametrize("capability", [(8, 0), (8, 6)]) +def test_sm80_family_uses_per_focus_so2_forward( + capability: tuple[int, int], +) -> None: + config = _K1._architecture_default_config(capability) + + assert config.per_focus_so2_fwd_pair + assert not config.native_sm90_path + assert not config.combined_so2_gate + + +def test_sm90_uses_native_split_complex_path() -> None: + config = _K1._architecture_default_config((9, 0)) + + assert config.native_sm90_path + assert not config.per_focus_so2_fwd_pair + assert not config.combined_so2_gate + + +def test_sm100_uses_shared_default_profile() -> None: + config = _K1._architecture_default_config((10, 0)) + + assert not config.native_sm90_path + assert not config.per_focus_so2_fwd_pair + assert not config.combined_so2_gate + + +@pytest.mark.parametrize("capability", [(8, 9), (12, 0)]) +def test_sm89_and_sm120_use_combined_so2_gate( + capability: tuple[int, int], +) -> None: + config = _K1._architecture_default_config(capability) + + assert config.combined_so2_gate + assert not config.native_sm90_path + assert not config.per_focus_so2_fwd_pair + + +def test_combined_so2_gate_rejects_misaligned_contiguous_tensor() -> None: + pytest.importorskip("cutlass.cute") + pytest.importorskip("cuda.bindings.driver") + from deepmd.kernels.cute.neo.k1_kernels.cute_neo_so2_gate_combined_fwd import ( + _require_16_byte_alignment, + ) + + storage = torch.empty(9, dtype=torch.float32, device="cpu") + aligned = storage[:8] + misaligned = storage[1:9] + assert aligned.data_ptr() % 16 == 0 + assert misaligned.is_contiguous() + _require_16_byte_alignment((aligned,)) + with pytest.raises(ValueError, match="16-byte aligned"): + _require_16_byte_alignment((misaligned,)) + + +def test_runtime_config_contains_only_reached_selectors() -> None: + assert {field.name for field in fields(NeoK1RuntimeConfig)} == { + "native_sm90_path", + "per_focus_so2_fwd_pair", + "combined_so2_gate", + } + + +def test_runtime_config_rejects_incompatible_capability() -> None: + config = _K1._architecture_default_config((8, 0)) + + with pytest.raises(RuntimeError, match="supported compute capability"): + _validate_runtime_config(config, compute_capability=(7, 5)) + + +def test_runtime_config_rejects_wrong_architecture_profile() -> None: + sm80_config = _K1._architecture_default_config((8, 0)) + with pytest.raises(RuntimeError, match="per-focus SO2"): + _validate_runtime_config(sm80_config, compute_capability=(12, 0)) + + sm90_config = _K1._architecture_default_config((9, 0)) + with pytest.raises(RuntimeError, match="native SM90 K1"): + _validate_runtime_config(sm90_config, compute_capability=(8, 0)) + + sm120_config = _K1._architecture_default_config((12, 0)) + with pytest.raises(RuntimeError, match="combined SO2/gate"): + _validate_runtime_config(sm120_config, compute_capability=(10, 0)) + with pytest.raises(RuntimeError, match="combined SO2/gate"): + _validate_runtime_config(NeoK1RuntimeConfig(), compute_capability=(12, 0)) + + +@pytest.mark.parametrize( + ("capability", "expected"), + [ + ((8, 0), True), + ((8, 6), True), + ((8, 9), False), + ((9, 0), True), + ((10, 0), False), + ((12, 0), False), + ], +) +def test_packed_message_grid_architecture_contract( + capability: tuple[int, int], + expected: bool, +) -> None: + assert _uses_packed_message_grid(capability) is expected + + +def test_packed_edge_eligibility_requires_sorted_strict_fp32() -> None: + kwargs = { + "candidate": True, + "edge_count": 12, + "node_count": 4, + "destinations_sorted": True, + "runtime_dtypes": (torch.float32, torch.float32), + } + + assert _K1.packed_wigner_edges_eligible(**kwargs) + assert not _K1.packed_wigner_edges_eligible( + **{**kwargs, "destinations_sorted": False} + ) + assert not _K1.packed_wigner_edges_eligible( + **{**kwargs, "runtime_dtypes": (torch.float64,)} + ) + assert _K1.packed_wigner_edges_eligible( + **{**kwargs, "edge_count": 3, "node_count": 4} + ) + + +def test_packed_wigner_has_a_separate_backward_compatible_cache_field() -> None: + dense = torch.empty((2, 16, 16), device="cpu") + dense_t = torch.empty_like(dense) + actual_dense, actual_dense_t, packed = _separate_packed_wigner(dense, dense_t) + assert actual_dense is dense + assert actual_dense_t is dense_t + assert packed is None + + panel = torch.empty((2, 46), device="cpu") + actual_dense, actual_dense_t, packed = _separate_packed_wigner(panel, panel) + assert actual_dense is None + assert actual_dense_t is None + assert packed is panel + assert EdgeFeatureCache._fields[-2:] == ("destinations_sorted", "D_packed") + + +def test_ineligible_wigner_build_does_not_import_cute( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + edge_vec = torch.tensor([[1.0, 0.0, 0.0]], device="cpu") + edge_len = torch.linalg.vector_norm(edge_vec, dim=-1, keepdim=True) + + def eager_wigner(edge_quat: torch.Tensor): + dense = edge_quat.new_zeros((edge_quat.shape[0], 1, 1)) + return dense, dense.transpose(-1, -2) + + dense, dense_t, edge_quat = _build_edge_wigner( + edge_vec=edge_vec, + edge_len=edge_len, + eps=1.0e-8, + random_gamma=False, + wigner_calc=eager_wigner, + packed_wigner=False, + ) + + assert dense is not None and tuple(dense.shape) == (1, 1, 1) + assert dense_t is not None and tuple(dense_t.shape) == (1, 1, 1) + assert tuple(edge_quat.shape) == (1, 4) + + +@pytest.mark.parametrize( + ("training", "dtype", "edge_count", "node_count", "destinations_sorted"), + [ + (True, torch.float32, 12, 4, True), + (False, torch.float64, 12, 4, True), + (False, torch.float32, 0, 4, True), + (False, torch.float32, 12, 4, False), + ], +) +def test_ineligible_runtime_contracts_fall_back( + monkeypatch: pytest.MonkeyPatch, + training: bool, + dtype: torch.dtype, + edge_count: int, + node_count: int, + destinations_sorted: bool, +) -> None: + monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + block = _neo_like_block().eval() + + assert not _K1.is_neo_k1_runtime_eligible( + block, + training=training, + device=torch.device("cuda", 0), + dtype=dtype, + edge_count=edge_count, + node_count=node_count, + destinations_sorted=destinations_sorted, + ) + + +def test_runtime_contract_allows_fewer_edges_than_nodes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + block = _neo_like_block().eval() + + assert _K1.is_neo_k1_runtime_eligible( + block, + training=False, + device=torch.device("cuda", 0), + dtype=torch.float32, + edge_count=3, + node_count=4, + destinations_sorted=True, + ) + + +def _cute_cuda_runtime_available() -> bool: + if not torch.cuda.is_available(): + return False + try: + importlib.import_module("cutlass.cute") + importlib.import_module("cuda.bindings.driver") + except Exception: # pragma: no cover - runtime dependent + return False + return True + + +@pytest.mark.skipif( + not _cute_cuda_runtime_available(), + reason="CuTe attention-prelude regression requires CUDA and CuTe DSL", +) +def test_attention_prelude_initializes_all_qk_rows_when_edges_are_sparse() -> None: + from deepmd.kernels.cute.neo.k1_kernels.cute_neo_focus_src_backward import ( + compile_neo_attention_prelude_forward, + ) + + torch.manual_seed(20260810) + device = torch.device("cuda") + edge_count = 2 + node_count = 3 + focus = torch.randn(edge_count, 64, device=device) + x_l0 = torch.randn(node_count, 2, 32, device=device) + focus_weight = torch.randn(32, 2, device=device) + focus_scale = torch.randn(2, 32, device=device) + q_weight = torch.randn(32, 2, 32, device=device) + k_weight = torch.randn_like(q_weight) + qk_scale = torch.randn(2, 32, device=device) + focus_alpha = torch.full((edge_count, 2), torch.nan, device=device) + q_node = torch.full((node_count, 2, 32), torch.nan, device=device) + k_node = torch.full_like(q_node, torch.nan) + focus_eps = 1.0e-5 + qk_eps = 1.0e-5 + tau = 0.7 + label_smoothing = 0.1 + + focus_view = focus.view(edge_count, 2, 32) + focus_norm = ( + focus_view + * torch.rsqrt(focus_view.square().mean(dim=-1, keepdim=True) + focus_eps) + * focus_scale.unsqueeze(0) + ) + focus_logits = torch.stack( + [ + (focus_norm[:, focus_idx] * focus_weight[:, focus_idx]).sum(dim=-1) + for focus_idx in range(2) + ], + dim=-1, + ) + expected_alpha = torch.softmax(focus_logits / tau, dim=-1) + expected_alpha = expected_alpha * (1.0 - label_smoothing) + (label_smoothing / 2.0) + x_norm = ( + x_l0 + * torch.rsqrt(x_l0.square().mean(dim=-1, keepdim=True) + qk_eps) + * qk_scale.unsqueeze(0) + ) + expected_q = torch.einsum("nfi,ifo->nfo", x_norm, q_weight) + expected_k = torch.einsum("nfi,ifo->nfo", x_norm, k_weight) + + run = compile_neo_attention_prelude_forward( + focus_eps, + qk_eps, + tau, + label_smoothing, + ) + run( + focus, + x_l0, + focus_weight, + focus_scale, + q_weight, + k_weight, + qk_scale, + focus_alpha, + q_node, + k_node, + ) + torch.cuda.synchronize() + + assert torch.isfinite(q_node).all() + assert torch.isfinite(k_node).all() + torch.testing.assert_close(focus_alpha, expected_alpha, atol=5.0e-5, rtol=5.0e-5) + torch.testing.assert_close(q_node, expected_q, atol=5.0e-5, rtol=5.0e-5) + torch.testing.assert_close(k_node, expected_k, atol=5.0e-5, rtol=5.0e-5) + + +def test_exact_runtime_contract_is_eligible_without_autocast( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setattr(runtime_policy, "uses_strict_fp32_matmul", lambda: True) + block = _neo_like_block().eval() + kwargs = { + "training": False, + "device": torch.device("cuda", 0), + "dtype": torch.float32, + "edge_count": 12, + "node_count": 4, + "destinations_sorted": True, + } + + assert _K1.is_neo_k1_runtime_eligible(block, **kwargs) + monkeypatch.setattr(torch, "is_autocast_enabled", lambda device_type: True) + assert not _K1.is_neo_k1_runtime_eligible(block, **kwargs) + + +def test_runtime_contract_rejects_tf32_matmul_policy( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setattr(torch, "is_autocast_enabled", lambda device_type: False) + block = _neo_like_block().eval() + kwargs = { + "training": False, + "device": torch.device("cuda", 0), + "dtype": torch.float32, + "edge_count": 12, + "node_count": 4, + "destinations_sorted": True, + } + + monkeypatch.setattr(runtime_policy, "uses_strict_fp32_matmul", lambda: True) + assert _K1.is_neo_k1_runtime_eligible(block, **kwargs) + + monkeypatch.setattr(runtime_policy, "uses_strict_fp32_matmul", lambda: False) + assert not _K1.is_neo_k1_runtime_eligible(block, **kwargs) + + +def test_fake_native_gradient_layout_matches_runtime_contract() -> None: + x = torch.empty(5, 16, 1, 32, device="cpu") + + grad = _K1._fake_x_wide_grad_like(x, skip=True) + + assert grad.shape == x.shape + assert grad.stride() == (32, 5 * 32, 32, 1) + + +def test_custom_op_runtime_canonicalizes_misaligned_contiguous_view() -> None: + base = torch.arange(17, dtype=torch.float32, device="cpu") + offset_view = base[1:] + + assert offset_view.is_contiguous() + assert offset_view.storage_offset() == 1 + actual = _K1._aligned_contiguous(offset_view) + + torch.testing.assert_close(actual, offset_view) + assert actual.is_contiguous() + assert actual.storage_offset() == 0 + assert actual.data_ptr() % 16 == 0 + + +def test_custom_op_runtime_preserves_aligned_compact_tensor() -> None: + tensor = torch.arange(16, dtype=torch.float32, device="cpu") + + assert _K1._aligned_contiguous(tensor) is tensor + + +def test_stack_cache_retains_only_backward_state() -> None: + assert {field.name for field in fields(StackCache)} == { + "y", + "logits", + "non_linear", + "final", + } diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_dst_ptr.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_dst_ptr.py new file mode 100644 index 0000000000..ef12c702f2 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_dst_ptr.py @@ -0,0 +1,646 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Contracts for asynchronous K1 destination-pointer construction.""" + +from __future__ import ( + annotations, +) + +import ast +import importlib +import unittest +from pathlib import ( + Path, +) +from types import ( + SimpleNamespace, +) +from typing import ( + Any, +) +from unittest import ( + mock, +) + +try: + import torch +except ModuleNotFoundError: # pragma: no cover - lightweight source-test host + torch = None + + +REPO_ROOT = Path(__file__).resolve().parents[4] +K1_PATH = REPO_ROOT / "deepmd/kernels/cute/neo/k1.py" +EDGE_CACHE_PATH = REPO_ROOT / "deepmd/pt/model/descriptor/sezm_nn/edge_cache.py" + + +class _Tensor: + pass + + +def _function(tree: ast.AST, name: str) -> ast.FunctionDef: + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == name: + return node + raise AssertionError(f"function {name!r} is missing") + + +def _load_extracted_dst_ptr_function(*, strict: bool = False): + source = K1_PATH.read_text(encoding="utf-8") + tree = ast.parse(source) + function = _function(tree, "_dst_ptr_from_sorted") + module = ast.Module(body=[function], type_ignores=[]) + ast.fix_missing_locations(module) + namespace: dict[str, Any] = { + "Any": Any, + "Tensor": _Tensor, + "runtime_policy": SimpleNamespace( + is_cute_strict_enabled=lambda: strict, + ), + } + exec(compile(module, str(K1_PATH), "exec"), namespace) + return namespace["_dst_ptr_from_sorted"] + + +def _load_extracted_sorted_metadata_function(): + source = EDGE_CACHE_PATH.read_text(encoding="utf-8") + tree = ast.parse(source) + function = _function(tree, "build_sorted_edge_index_metadata") + module = ast.Module(body=[function], type_ignores=[]) + ast.fix_missing_locations(module) + namespace = {"torch": torch} + exec(compile(module, str(EDGE_CACHE_PATH), "exec"), namespace) + return namespace["build_sorted_edge_index_metadata"] + + +class _FakeDst: + device = "cuda:0" + dtype = "int64" + + def __init__(self): + self.contiguous_calls = 0 + + def contiguous(self): + self.contiguous_calls += 1 + return self + + +class _RecordingTorch: + int64 = "int64" + + def __init__(self): + self.calls: list[tuple[Any, ...]] = [] + self.boundaries = object() + self.result = object() + + def arange(self, stop, *, device, dtype): + self.calls.append(("arange", stop, device, dtype)) + return self.boundaries + + def searchsorted(self, sorted_sequence, values): + self.calls.append(("searchsorted", sorted_sequence, values)) + return self.result + + def bincount(self, *args, **kwargs): + raise AssertionError("CUDA bincount must not construct sorted dst_ptr") + + +class TestK1DstPtrExtracted(unittest.TestCase): + def test_sorted_path_operator_and_output_contract(self): + helper = _load_extracted_dst_ptr_function() + torch_module = _RecordingTorch() + dst = _FakeDst() + + result = helper( + torch_module, + dst, + 8, + destinations_sorted=True, + ) + + self.assertIs(result, torch_module.result) + self.assertEqual(dst.contiguous_calls, 1) + self.assertEqual( + torch_module.calls, + [ + ("arange", 9, dst.device, torch_module.int64), + ("searchsorted", dst, torch_module.boundaries), + ], + ) + + def test_unsorted_path_returns_before_tensor_work(self): + helper = _load_extracted_dst_ptr_function() + torch_module = _RecordingTorch() + + self.assertIsNone( + helper( + torch_module, + _FakeDst(), + 8, + destinations_sorted=False, + ) + ) + self.assertEqual(torch_module.calls, []) + + +@unittest.skipIf(torch is None, "destination-pointer differential requires PyTorch") +class TestK1DstPtrTorch(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.helper = staticmethod(_load_extracted_dst_ptr_function()) + + def test_matches_bincount_reference_and_preserves_layout(self): + assert torch is not None + devices = [torch.device("cpu")] + if torch.cuda.is_available(): + devices.append(torch.device("cuda")) + + cases = ( + (1, []), + (1, [0, 0, 0]), + (8, [0, 0, 2, 5, 5, 5, 7]), + (9, [1, 1, 1, 4, 8]), + (17, [0, 3, 3, 7, 7, 7, 15]), + ) + for device in devices: + for n_node, values in cases: + with self.subTest(device=device, n_node=n_node, values=values): + dst = torch.tensor(values, device=device, dtype=torch.int64) + counts = torch.bincount(dst, minlength=n_node) + expected = torch.empty( + n_node + 1, + device=device, + dtype=torch.int64, + ) + expected[0] = 0 + expected[1:] = counts.cumsum(0) + + actual = self.helper( + torch, + dst, + n_node, + destinations_sorted=True, + ) + + self.assertTrue(torch.equal(actual, expected)) + self.assertEqual(actual.dtype, torch.int64) + self.assertEqual(actual.device, dst.device) + self.assertTrue(actual.is_contiguous()) + self.assertEqual(actual.stride(), (1,)) + + def test_strict_mode_rejects_false_sorted_provenance(self): + assert torch is not None + helper = _load_extracted_dst_ptr_function(strict=True) + dst = torch.tensor([0, 2, 1, 3], dtype=torch.int64, device="cpu") + + with self.assertRaisesRegex(RuntimeError, "monotonically nondecreasing"): + helper( + torch, + dst, + 4, + destinations_sorted=True, + ) + + def test_strict_mode_accepts_duplicate_sorted_destinations(self): + assert torch is not None + helper = _load_extracted_dst_ptr_function(strict=True) + dst = torch.tensor([0, 0, 2, 3], dtype=torch.int64, device="cpu") + + actual = helper( + torch, + dst, + 4, + destinations_sorted=True, + ) + + self.assertTrue( + torch.equal( + actual, + torch.tensor([0, 2, 2, 3, 4], device="cpu"), + ) + ) + + def test_fullgraph_dynamic_compile_preserves_pointer_contract(self): + assert torch is not None + helper = self.helper + if torch.cuda.is_available(): + device = torch.device("cuda") + backend = "inductor" + else: + device = torch.device("cpu") + backend = "aot_eager" + + def build_ptr(dst): + return helper( + torch, + dst, + 8, + destinations_sorted=True, + ) + + compiled = torch.compile( + build_ptr, + backend=backend, + dynamic=True, + fullgraph=True, + ) + dst = torch.tensor( + [0, 0, 2, 5, 5, 5, 7], + device=device, + dtype=torch.int64, + ) + actual = compiled(dst) + expected = torch.tensor( + [0, 2, 2, 3, 3, 3, 6, 6, 7], + device=device, + dtype=torch.int64, + ) + + self.assertTrue(torch.equal(actual, expected)) + self.assertEqual(actual.stride(), (1,)) + + @unittest.skipUnless( + torch is not None and torch.cuda.is_available(), + "K1 custom-op registry lifetime regression requires CUDA", + ) + def test_compile_cold_registration_survives_into_packed_custom_op_runtime(self): + assert torch is not None + k1 = importlib.import_module("deepmd.kernels.cute.neo.k1") + prior_registry = dict(k1._REGISTRY) + prior_next_handle = k1._NEXT_HANDLE + k1._REGISTRY.clear() + k1._PACKED_RUNNER_CACHE.clear() + torch._dynamo.reset() + + def cleanup(): + torch._dynamo.reset() + k1._PACKED_RUNNER_CACHE.clear() + k1._REGISTRY.clear() + k1._REGISTRY.update(prior_registry) + k1._NEXT_HANDLE = prior_next_handle + + self.addCleanup(cleanup) + block = torch.nn.Module() + config = k1.NeoK1RuntimeConfig() + runtime_handles = [] + + class FakeRunner: + pass + + def fake_build(handle, x_arg, *args): + del args + entry = k1._REGISTRY[int(handle)] + self.assertIs(entry.block, block) + runtime_handles.append(int(handle)) + runner = FakeRunner() + runner.final = x_arg.detach().clone() + return runner + + def invoke( + x_arg, + d_arg, + dt_arg, + radial_arg, + edge_arg, + src_arg, + dst_arg, + ): + state = getattr(block, "_deepmd_cute_k1_state", None) + if state is None: + state = k1._register_cute_k1_state( + block, + x_arg.device.index, + config, + ) + dst_ptr = k1._dst_ptr_from_sorted( + torch, + dst_arg, + x_arg.shape[0], + destinations_sorted=True, + ) + edge_src_gate = edge_arg.new_empty((0,)) + return k1.cute_k1( + state.handle, + x_arg, + d_arg, + dt_arg, + radial_arg, + edge_arg, + src_arg, + dst_arg, + dst_ptr, + edge_src_gate, + ) + + disabled_invoke = torch.compiler.disable(invoke) + + def compiled_entry(*args): + return disabled_invoke(*args) + + device = torch.device("cuda:0") + x = torch.randn(2, 16, 1, 32, device=device) + d_full = torch.randn(3, 46, device=device) + dt_full = torch.randn_like(d_full) + radial = torch.randn(3, 4, 32, device=device) + edge_env = torch.ones(3, 1, device=device) + src = torch.tensor((0, 1, 0), dtype=torch.int64, device=device) + dst = torch.tensor((0, 0, 1), dtype=torch.int64, device=device) + + with mock.patch.object(k1, "_build_runner", new=fake_build): + compiled = torch.compile(compiled_entry, backend="eager", dynamic=True) + first = compiled(x, d_full, dt_full, radial, edge_env, src, dst) + second = compiled(x, d_full, dt_full, radial, edge_env, src, dst) + + state = block._deepmd_cute_k1_state + self.assertEqual(runtime_handles, [state.handle, state.handle]) + self.assertEqual(list(k1._REGISTRY), [state.handle]) + self.assertIs(k1._REGISTRY[state.handle].block, block) + torch.testing.assert_close(first, x, rtol=0.0, atol=0.0) + torch.testing.assert_close(second, x, rtol=0.0, atol=0.0) + + +@unittest.skipIf(torch is None, "sorted edge metadata requires PyTorch") +class TestSortedEdgeIndexMetadata(unittest.TestCase): + @staticmethod + def _cache(src, dst, node_count): + assert torch is not None + edge_cache = importlib.import_module( + "deepmd.pt.model.descriptor.sezm_nn.edge_cache" + ) + device = torch.device("cpu") + src = src.to(device=device) + dst = dst.to(device=device) + edge_count = src.numel() + return edge_cache.EdgeFeatureCache( + src=src, + dst=dst, + edge_type_feat=torch.empty(edge_count, 1, device=device), + edge_vec=torch.empty(edge_count, 3, device=device), + edge_rbf=torch.empty(edge_count, 1, device=device), + edge_env=torch.ones(edge_count, 1, device=device), + deg=torch.zeros(node_count, device=device), + inv_sqrt_deg=torch.ones(node_count, 1, 1, device=device), + destinations_sorted=True, + ) + + def test_builds_destination_and_indirect_source_csr(self): + assert torch is not None + edge_cache = importlib.import_module( + "deepmd.pt.model.descriptor.sezm_nn.edge_cache" + ) + src = torch.tensor( + [2, 0, 3, 0, 1, 2], + dtype=torch.int64, + device="cpu", + ) + dst = torch.tensor( + [0, 0, 1, 2, 2, 3], + dtype=torch.int64, + device="cpu", + ) + dst_ptr, source_order, source_ptr = edge_cache.build_sorted_edge_index_metadata( + src, dst, 4 + ) + + self.assertEqual(dst_ptr.dtype, torch.int32) + self.assertEqual(source_order.dtype, torch.int32) + self.assertEqual(source_ptr.dtype, torch.int32) + torch.testing.assert_close( + dst_ptr, + torch.tensor( + [0, 2, 3, 5, 6], + dtype=torch.int32, + device="cpu", + ), + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + src.index_select(0, source_order.to(torch.int64)), + torch.tensor( + [0, 0, 1, 2, 2, 3], + dtype=torch.int64, + device="cpu", + ), + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + source_ptr, + torch.tensor( + [0, 2, 3, 5, 6], + dtype=torch.int32, + device="cpu", + ), + rtol=0.0, + atol=0.0, + ) + + def test_empty_edges_build_valid_zero_csr(self): + assert torch is not None + edge_cache = importlib.import_module( + "deepmd.pt.model.descriptor.sezm_nn.edge_cache" + ) + dst_ptr, source_order, source_ptr = edge_cache.build_sorted_edge_index_metadata( + torch.empty(0, dtype=torch.int64, device="cpu"), + torch.empty(0, dtype=torch.int64, device="cpu"), + 4, + ) + + torch.testing.assert_close( + dst_ptr, + torch.zeros(5, dtype=torch.int32, device="cpu"), + rtol=0.0, + atol=0.0, + ) + self.assertEqual(source_order.numel(), 0) + torch.testing.assert_close( + source_ptr, + torch.zeros(5, dtype=torch.int32, device="cpu"), + rtol=0.0, + atol=0.0, + ) + + def test_strict_metadata_rejects_unsorted_destinations(self): + assert torch is not None + builder = _load_extracted_sorted_metadata_function() + src = torch.tensor([0, 1, 2, 3], dtype=torch.int64, device="cpu") + dst = torch.tensor([0, 2, 1, 3], dtype=torch.int64, device="cpu") + + with self.assertRaisesRegex(RuntimeError, "monotonically nondecreasing"): + builder( + src, + dst, + 4, + validate_sorted=True, + ) + + def test_strict_metadata_builder_is_fullgraph_compilable(self): + assert torch is not None + builder = _load_extracted_sorted_metadata_function() + dynamo = getattr(torch, "_dynamo", None) + if dynamo is None: + self.skipTest("torch._dynamo is unavailable") + + def build_ptr(src, dst): + dst_ptr, _, _ = builder( + src, + dst, + 4, + validate_sorted=True, + ) + return dst_ptr + + compiled = torch.compile( + build_ptr, + backend="eager", + dynamic=True, + fullgraph=True, + ) + src = torch.tensor([2, 0, 1], dtype=torch.int64, device="cpu") + dst = torch.tensor([0, 1, 2], dtype=torch.int64, device="cpu") + actual = compiled(src, dst) + + torch.testing.assert_close( + actual, + torch.tensor([0, 1, 2, 3, 3], dtype=torch.int32, device="cpu"), + rtol=0.0, + atol=0.0, + ) + + def test_dynamic_node_and_edge_counts_build_independent_metadata(self): + assert torch is not None + edge_cache = importlib.import_module( + "deepmd.pt.model.descriptor.sezm_nn.edge_cache" + ) + first = edge_cache.build_sorted_edge_index_metadata( + torch.tensor([0, 1, 2], dtype=torch.int64, device="cpu"), + torch.tensor([0, 1, 2], dtype=torch.int64, device="cpu"), + 4, + ) + second = edge_cache.build_sorted_edge_index_metadata( + torch.tensor( + [0, 2, 4, 1, 3], + dtype=torch.int64, + device="cpu", + ), + torch.tensor( + [0, 0, 2, 4, 5], + dtype=torch.int64, + device="cpu", + ), + 6, + ) + + torch.testing.assert_close( + first[0], + torch.tensor([0, 1, 2, 3, 3], dtype=torch.int32, device="cpu"), + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + second[0], + torch.tensor( + [0, 2, 2, 3, 3, 4, 5], + dtype=torch.int32, + device="cpu", + ), + rtol=0.0, + atol=0.0, + ) + self.assertNotEqual(first[0].data_ptr(), second[0].data_ptr()) + self.assertNotEqual(first[1].data_ptr(), second[1].data_ptr()) + self.assertNotEqual(first[2].data_ptr(), second[2].data_ptr()) + + def test_explicit_tensor_metadata_survives_disabled_k1_boundary(self): + assert torch is not None + edge_cache = importlib.import_module( + "deepmd.pt.model.descriptor.sezm_nn.edge_cache" + ) + dynamo = getattr(torch, "_dynamo", None) + if dynamo is None: + self.skipTest("torch._dynamo is unavailable") + + def k1_boundary( + x, + cache, + radial, + dst_ptr, + source_order, + source_ptr, + ): + del radial + assert dst_ptr is not None + assert source_order is not None + assert source_ptr is not None + return ( + x + + cache.edge_env.sum() + + dst_ptr.sum() + + source_order.sum() + + source_ptr.sum() + ) + + disabled_k1_boundary = torch.compiler.disable(k1_boundary) + + def consume(cache, x, radial, dst_ptr, source_order, source_ptr): + before_break = x + cache.edge_env.sum() + return disabled_k1_boundary( + before_break, + cache, + radial, + dst_ptr, + source_order, + source_ptr, + ) + + compiled = torch.compile(consume, backend="eager", dynamic=True) + cases = ( + ( + torch.tensor([0, 1, 2], dtype=torch.int64, device="cpu"), + torch.tensor([0, 1, 2], dtype=torch.int64, device="cpu"), + 4, + ), + ( + torch.tensor( + [0, 2, 4, 1, 3], + dtype=torch.int64, + device="cpu", + ), + torch.tensor( + [0, 0, 2, 4, 5], + dtype=torch.int64, + device="cpu", + ), + 6, + ), + ) + for src, dst, node_count in cases: + cache = self._cache(src, dst, node_count) + dst_ptr, source_order, source_ptr = ( + edge_cache.build_sorted_edge_index_metadata( + src, + dst, + node_count, + ) + ) + radial = torch.ones(src.numel(), 4, 1, device="cpu") + args = ( + cache, + torch.ones(1, device="cpu"), + radial, + dst_ptr, + source_order, + source_ptr, + ) + eager = consume(*args) + actual = compiled(*args) + torch.testing.assert_close( + actual, + eager, + rtol=0.0, + atol=0.0, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py new file mode 100644 index 0000000000..3394bf1356 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py @@ -0,0 +1,117 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Numerical differentials for the structural Neo K1 gate.""" + +from __future__ import ( + annotations, +) + +import importlib.util +import sys +import unittest +from pathlib import ( + Path, +) + +import torch + +REPO_ROOT = Path(__file__).resolve().parents[4] +HELPER_PATH = REPO_ROOT / "deepmd/kernels/cute/neo/k1_gate_structural.py" + + +def _load_module(path: Path, name: str): + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise RuntimeError(f"could not load {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + try: + spec.loader.exec_module(module) + finally: + sys.modules.pop(name, None) + return module + + +class TestStructuralGateHelpers(unittest.TestCase): + @classmethod + def setUpClass(cls): + if not HELPER_PATH.is_file(): + raise AssertionError(f"structural gate helper is missing: {HELPER_PATH}") + cls.helper = _load_module(HELPER_PATH, "structural_gate_helper_test") + + def setUp(self): + torch.manual_seed(20260703) + self.edge_count = 7 + self.gate_src = torch.randn( + self.edge_count, 2, 32, dtype=torch.float32, device="cpu" + ) + self.weight = torch.randn(32, 2 * 3 * 32, dtype=torch.float32, device="cpu") + + def test_focus_major_forward_matches_focus_linear(self): + actual = self.helper.focus_major_gate_linear_forward( + self.gate_src, + self.weight, + ) + expected = torch.einsum( + "efi,ifo->efo", + self.gate_src, + self.weight.view(32, 2, 3 * 32), + ) + self.assertEqual(actual.shape, (2, self.edge_count, 3 * 32)) + self.assertEqual(actual.stride(), (self.edge_count * 3 * 32, 3 * 32, 1)) + torch.testing.assert_close( + actual.permute(1, 0, 2), expected, atol=5e-5, rtol=5e-5 + ) + + def test_backward_addmm_accumulates_without_replacing_grad_y_storage(self): + grad_y = torch.randn( + self.edge_count, 2, 10, 32, dtype=torch.float32, device="cpu" + ) + grad_logits = torch.randn( + 2, self.edge_count, 3 * 32, dtype=torch.float32, device="cpu" + ) + expected = grad_y.clone() + weight = self.weight.view(32, 2, 3 * 32) + for focus in range(2): + expected[:, focus, 0, :] += grad_logits[focus] @ weight[:, focus, :].T + + pointer = grad_y.untyped_storage().data_ptr() + actual = self.helper.focus_major_gate_linear_backward_add_( + grad_y, + grad_logits, + self.weight, + ) + + self.assertIs(actual, grad_y) + self.assertEqual(actual.untyped_storage().data_ptr(), pointer) + torch.testing.assert_close(actual, expected, atol=5e-5, rtol=5e-5) + + def test_forward_wrapper_preserves_caller_owned_storage(self): + residual = torch.randn( + self.edge_count, 2, 10, 32, dtype=torch.float32, device="cpu" + ) + y = torch.randn_like(residual) + logits = torch.randn( + 2, self.edge_count, 3 * 32, dtype=torch.float32, device="cpu" + ) + pointer = residual.untyped_storage().data_ptr() + + def fake_kernel(residual_flat, y_flat, logits_arg, out_flat): + self.assertEqual(residual_flat.data_ptr(), out_flat.data_ptr()) + self.assertIs(logits_arg, logits) + out_flat.copy_(residual_flat + y_flat) + + actual = self.helper.run_structural_gate_forward( + fake_kernel, + residual, + y, + logits, + out=residual, + ) + + self.assertIs(actual, residual) + self.assertEqual(actual.untyped_storage().data_ptr(), pointer) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py new file mode 100644 index 0000000000..a977b4ce6f --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py @@ -0,0 +1,387 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Contracts for the in-place K1 SO2 residual/addmm helper.""" + +from __future__ import ( + annotations, +) + +import importlib.util +import sys +import unittest +from functools import ( + lru_cache, +) +from pathlib import ( + Path, +) + +import torch +from torch.utils._python_dispatch import ( + TorchDispatchMode, +) + +REPO_ROOT = Path(__file__).resolve().parents[4] +IMPLEMENTATION_PATH = REPO_ROOT / "deepmd/kernels/cute/neo/k1_so2linear.py" + + +@lru_cache(maxsize=1) +def _load_candidate(): + assert IMPLEMENTATION_PATH.is_file(), ( + f"in-place residual module is missing: {IMPLEMENTATION_PATH}" + ) + name = "sezm_cute_k1_inplace_residual_test" + spec = importlib.util.spec_from_file_location(name, IMPLEMENTATION_PATH) + if spec is None or spec.loader is None: + raise RuntimeError(f"could not load {IMPLEMENTATION_PATH}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + try: + spec.loader.exec_module(module) + finally: + sys.modules.pop(name, None) + return module + + +def _out_of_place_reference(grad_out, residual, w0_t, wpair_t): + edge_count = grad_out.shape[0] + grad_flat = grad_out.view(edge_count, 2, 10 * 32) + residual_flat = residual.view(edge_count, 2, 10 * 32) + out = torch.empty_like(residual) + out_flat = out.view(edge_count, 2, 10 * 32) + for focus in range(2): + torch.addmm( + residual_flat[:, focus, : 4 * 32], + grad_flat[:, focus, : 4 * 32], + w0_t[focus], + out=out_flat[:, focus, : 4 * 32], + ) + torch.addmm( + residual_flat[:, focus, 4 * 32 :], + grad_flat[:, focus, 4 * 32 :], + wpair_t[focus], + out=out_flat[:, focus, 4 * 32 :], + ) + return out + + +class _DispatchRecorder(TorchDispatchMode): + def __init__(self): + super().__init__() + self.calls = [] + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + self.calls.append((func, args, kwargs)) + return func(*args, **kwargs) + + +class TestK1InplaceResidual(unittest.TestCase): + def setUp(self): + torch.manual_seed(20260630) + self.edge_count = 7 + self.grad_out = torch.randn( + self.edge_count, + 2, + 10, + 32, + dtype=torch.float32, + device="cpu", + ) + self.residual = torch.randn_like(self.grad_out) + self.w0_t = torch.randn( + 2, + 4 * 32, + 4 * 32, + dtype=torch.float32, + device="cpu", + ) + self.wpair_t = torch.randn( + 2, + 6 * 32, + 6 * 32, + dtype=torch.float32, + device="cpu", + ) + + def test_inplace_result_matches_out_of_place_equations(self): + candidate = _load_candidate() + expected = _out_of_place_reference( + self.grad_out, + self.residual, + self.w0_t, + self.wpair_t, + ) + residual = self.residual.clone() + storage_ptr = residual.untyped_storage().data_ptr() + + actual = candidate.neo_so2_linear_backward_residual_inplace( + residual, + self.grad_out, + self.w0_t, + self.wpair_t, + ) + + self.assertIs(actual, residual) + self.assertEqual(actual.untyped_storage().data_ptr(), storage_ptr) + self.assertEqual(actual.stride(), (2 * 10 * 32, 10 * 32, 32, 1)) + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + + def test_only_four_inplace_addmm_ops_cross_the_dispatch_boundary(self): + candidate = _load_candidate() + residual = self.residual.clone() + recorder = _DispatchRecorder() + + with recorder: + candidate.neo_so2_linear_backward_residual_inplace( + residual, + self.grad_out, + self.w0_t, + self.wpair_t, + ) + + addmm_calls = [ + args + for func, args, _kwargs in recorder.calls + if func is torch.ops.aten.addmm_.default + ] + self.assertEqual(len(addmm_calls), 4) + forbidden = { + torch.ops.aten.clone.default, + torch.ops.aten.contiguous.default, + torch.ops.aten.copy_.default, + } + self.assertFalse( + any(func in forbidden for func, _args, _kwargs in recorder.calls) + ) + + for residual_block, grad_block, weight in addmm_calls: + width = residual_block.shape[1] + self.assertIn(width, (4 * 32, 6 * 32)) + self.assertEqual(residual_block.stride(), (2 * 10 * 32, 1)) + self.assertEqual(grad_block.stride(), (2 * 10 * 32, 1)) + self.assertEqual(weight.stride(), (width, 1)) + self.assertTrue(candidate._has_direct_cublas_layout(residual_block)) + self.assertTrue(candidate._has_direct_cublas_layout(grad_block)) + self.assertTrue(candidate._has_direct_cublas_layout(weight)) + + def test_layout_predicate_and_fp32_scope_are_explicit(self): + candidate = _load_candidate() + + self.assertFalse(hasattr(candidate, "is_cublas_compatible_matrix")) + self.assertIn("PyTorch 2.10", candidate._has_direct_cublas_layout.__doc__) + self.assertIn("dtype/device", candidate._has_direct_cublas_layout.__doc__) + self.assertIn("highest", candidate.__doc__) + self.assertIn("TF32", candidate.__doc__) + + def test_meta_execution_preserves_outer_shape_and_stride_contract(self): + candidate = _load_candidate() + residual = torch.empty(11, 2, 10, 32, dtype=torch.float32, device="meta") + grad_out = torch.empty_like(residual) + w0_t = torch.empty(2, 4 * 32, 4 * 32, dtype=torch.float32, device="meta") + wpair_t = torch.empty( + 2, + 6 * 32, + 6 * 32, + dtype=torch.float32, + device="meta", + ) + + result = candidate.neo_so2_linear_backward_residual_inplace( + residual, + grad_out, + w0_t, + wpair_t, + ) + + self.assertIs(result, residual) + self.assertEqual(result.shape, (11, 2, 10, 32)) + self.assertEqual(result.stride(), (2 * 10 * 32, 10 * 32, 32, 1)) + + def test_aliasing_grad_out_is_rejected_for_final_layer_safety(self): + candidate = _load_candidate() + shared = self.residual.clone() + + with self.assertRaisesRegex(ValueError, "must not alias"): + candidate.neo_so2_linear_backward_residual_inplace( + shared, + shared, + self.w0_t, + self.wpair_t, + ) + + def test_exact_cross_storage_grad_out_alias_is_rejected(self): + candidate = _load_candidate() + byte_count = self.residual.numel() * self.residual.element_size() + backing = bytearray(byte_count) + residual = torch.frombuffer( + backing, + dtype=torch.float32, + count=self.residual.numel(), + ).view_as(self.residual) + grad_out = torch.frombuffer( + backing, + dtype=torch.float32, + count=self.grad_out.numel(), + ).view_as(self.grad_out) + self.assertNotEqual( + residual.untyped_storage()._cdata, + grad_out.untyped_storage()._cdata, + ) + self.assertEqual(residual.data_ptr(), grad_out.data_ptr()) + self.assertFalse(torch._C._overlaps(residual, grad_out)) + + with self.assertRaisesRegex( + ValueError, + "residual and grad_out must not alias", + ): + candidate.neo_so2_linear_backward_residual_inplace( + residual, + grad_out, + self.w0_t, + self.wpair_t, + ) + + def test_residual_overlapping_w0_is_rejected(self): + candidate = _load_candidate() + storage = torch.randn( + 2 * 4 * 32 * 4 * 32, + dtype=torch.float32, + device="cpu", + ) + residual = storage[: self.residual.numel()].view_as(self.residual) + w0_t = storage.view(2, 4 * 32, 4 * 32) + self.assertTrue(residual.is_contiguous()) + self.assertTrue(w0_t.is_contiguous()) + self.assertTrue(torch._C._overlaps(residual, w0_t)) + + with self.assertRaisesRegex(ValueError, "residual and w0_t must not alias"): + candidate.neo_so2_linear_backward_residual_inplace( + residual, + self.grad_out, + w0_t, + self.wpair_t, + ) + + def test_partial_cross_storage_w0_alias_is_rejected(self): + candidate = _load_candidate() + overlap_elements = 17 + weight_elements = self.w0_t.numel() + weight_offset = self.residual.numel() - overlap_elements + backing = bytearray( + (weight_offset + weight_elements) * self.residual.element_size() + ) + residual = torch.frombuffer( + backing, + dtype=torch.float32, + count=self.residual.numel(), + ).view_as(self.residual) + w0_t = torch.frombuffer( + backing, + dtype=torch.float32, + count=weight_elements, + offset=weight_offset * self.residual.element_size(), + ).view_as(self.w0_t) + self.assertNotEqual( + residual.untyped_storage()._cdata, + w0_t.untyped_storage()._cdata, + ) + self.assertNotEqual(residual.data_ptr(), w0_t.data_ptr()) + self.assertFalse(torch._C._overlaps(residual, w0_t)) + + with self.assertRaisesRegex(ValueError, "residual and w0_t must not alias"): + candidate.neo_so2_linear_backward_residual_inplace( + residual, + self.grad_out, + w0_t, + self.wpair_t, + ) + + def test_residual_overlapping_wpair_is_rejected(self): + candidate = _load_candidate() + storage = torch.randn( + 2 * 6 * 32 * 6 * 32, + dtype=torch.float32, + device="cpu", + ) + residual = storage[: self.residual.numel()].view_as(self.residual) + wpair_t = storage.view(2, 6 * 32, 6 * 32) + self.assertTrue(residual.is_contiguous()) + self.assertTrue(wpair_t.is_contiguous()) + self.assertTrue(torch._C._overlaps(residual, wpair_t)) + + with self.assertRaisesRegex(ValueError, "residual and wpair_t must not alias"): + candidate.neo_so2_linear_backward_residual_inplace( + residual, + self.grad_out, + self.w0_t, + wpair_t, + ) + + def test_partial_cross_storage_wpair_alias_is_rejected(self): + candidate = _load_candidate() + overlap_elements = 17 + weight_elements = self.wpair_t.numel() + weight_offset = self.residual.numel() - overlap_elements + backing = bytearray( + (weight_offset + weight_elements) * self.residual.element_size() + ) + residual = torch.frombuffer( + backing, + dtype=torch.float32, + count=self.residual.numel(), + ).view_as(self.residual) + wpair_t = torch.frombuffer( + backing, + dtype=torch.float32, + count=weight_elements, + offset=weight_offset * self.residual.element_size(), + ).view_as(self.wpair_t) + self.assertNotEqual( + residual.untyped_storage()._cdata, + wpair_t.untyped_storage()._cdata, + ) + self.assertNotEqual(residual.data_ptr(), wpair_t.data_ptr()) + self.assertFalse(torch._C._overlaps(residual, wpair_t)) + + with self.assertRaisesRegex( + ValueError, + "residual and wpair_t must not alias", + ): + candidate.neo_so2_linear_backward_residual_inplace( + residual, + self.grad_out, + self.w0_t, + wpair_t, + ) + + def test_non_fp32_and_noncanonical_layouts_are_rejected(self): + candidate = _load_candidate() + with self.assertRaisesRegex(TypeError, "float32"): + candidate.neo_so2_linear_backward_residual_inplace( + self.residual.double(), + self.grad_out.double(), + self.w0_t.double(), + self.wpair_t.double(), + ) + + residual = torch.empty( + 2, + self.edge_count, + 10, + 32, + dtype=torch.float32, + device="cpu", + ).transpose(0, 1) + self.assertEqual(residual.shape, self.residual.shape) + with self.assertRaisesRegex(ValueError, "contiguous"): + candidate.neo_so2_linear_backward_residual_inplace( + residual, + self.grad_out, + self.w0_t, + self.wpair_t, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_manual_adjoints.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_manual_adjoints.py new file mode 100644 index 0000000000..85d96d195a --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_manual_adjoints.py @@ -0,0 +1,360 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Parity tests for the real-module adjoints used by Neo K1 backward.""" + +from __future__ import ( + annotations, +) + +import importlib +from types import ( + SimpleNamespace, +) + +import pytest +import torch + +from deepmd.kernels.cute.neo import ( + k1, +) +from deepmd.pt.model.descriptor.sezm_nn.activation import ( + SwiGLU, +) +from deepmd.pt.model.descriptor.sezm_nn.grid_net import ( + SO3GridNet, +) +from deepmd.pt.model.descriptor.sezm_nn.norm import ( + EquivariantRMSNorm, +) +from deepmd.pt.model.descriptor.sezm_nn.so3 import ( + FocusLinear, +) +from deepmd.pt.utils import ( + env, +) + + +@pytest.fixture(autouse=True) +def _construct_modules_on_cpu(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(env, "DEVICE", torch.device("cpu")) + + +def _randn(*shape: int, requires_grad: bool = False) -> torch.Tensor: + return torch.randn( + *shape, + dtype=torch.float64, + device="cpu", + requires_grad=requires_grad, + ) + + +def test_equivariant_rmsnorm_manual_adjoint_matches_real_module() -> None: + norm = EquivariantRMSNorm( + lmax=2, + channels=4, + n_focus=2, + eps=2.0e-6, + dtype=torch.float64, + trainable=False, + ) + with torch.no_grad(): + norm.adam_scale.copy_( + torch.linspace( + 0.5, + 1.5, + norm.adam_scale.numel(), + dtype=torch.float64, + device="cpu", + ).reshape_as(norm.adam_scale) + ) + norm.bias.copy_( + torch.linspace( + -0.2, + 0.2, + norm.bias.numel(), + dtype=torch.float64, + device="cpu", + ).reshape_as(norm.bias) + ) + x = _randn(3, 9, 2, 4, requires_grad=True) + grad_out = _randn(*x.shape) + + reference_grad = torch.autograd.grad(norm(x), x, grad_out)[0] + actual_grad = k1._equivariant_rmsnorm_backward( + norm, + x.detach(), + grad_out, + ) + + torch.testing.assert_close(actual_grad, reference_grad, atol=1.0e-12, rtol=1.0e-12) + + +def test_focus_linear_manual_adjoint_matches_real_module() -> None: + linear = FocusLinear( + in_channels=6, + out_channels=5, + n_focus=2, + dtype=torch.float64, + bias=True, + trainable=False, + seed=None, + ) + x = _randn(7, 2, 6, requires_grad=True) + grad_out = _randn(7, 2, 5) + + reference = linear(x) + reference_grad = torch.autograd.grad(reference, x, grad_out)[0] + + torch.testing.assert_close(k1._focus_linear_forward(linear, x), reference) + torch.testing.assert_close( + k1._focus_linear_backward_input(linear, grad_out), + reference_grad, + atol=1.0e-12, + rtol=1.0e-12, + ) + + +def test_swiglu_manual_adjoint_matches_real_module() -> None: + activation = SwiGLU() + x = _randn(6, 3, 10, requires_grad=True) + grad_out = _randn(6, 3, 5) + + reference = activation(x) + reference_grad = torch.autograd.grad(reference, x, grad_out)[0] + + torch.testing.assert_close(k1._swiglu_forward(x), reference) + torch.testing.assert_close( + k1._swiglu_backward_input(x.detach(), grad_out), + reference_grad, + atol=1.0e-12, + rtol=1.0e-12, + ) + + +def test_grid_cross_glu_flat_manual_adjoint_matches_real_module() -> None: + net = SO3GridNet( + lmax=3, + kmax=1, + channels=4, + n_focus=2, + mode="cross", + op_type="glu", + dtype=torch.float64, + layout="flat", + coefficient_layout="packed", + residual_scale_init=0.25, + trainable=False, + seed=None, + ) + coeff_dim = net.projector.coeff_dim // net.n_frames + query = _randn(3, coeff_dim, net.n_focus * net.channels, requires_grad=True) + context = _randn(3, coeff_dim, net.n_focus * net.channels, requires_grad=True) + grad_out = _randn(3, coeff_dim, net.n_focus * net.channels) + + reference_query, reference_context = torch.autograd.grad( + net(query, context), + (query, context), + grad_out, + ) + actual_query, actual_context = k1._so3_grid_cross_glu_flat_backward( + net, + query.detach(), + context.detach(), + grad_out, + ) + + torch.testing.assert_close( + actual_query, + reference_query, + atol=1.0e-11, + rtol=1.0e-11, + ) + torch.testing.assert_close( + actual_context, + reference_context, + atol=1.0e-11, + rtol=1.0e-11, + ) + + +def _cute_cuda_runtime_available() -> bool: + if not torch.cuda.is_available(): + return False + try: + importlib.import_module("cutlass.cute") + importlib.import_module("cuda.bindings.driver") + except Exception: # pragma: no cover - runtime dependent + return False + return True + + +@pytest.mark.skipif( + not _cute_cuda_runtime_available(), + reason="SM90 message-grid differential requires CUDA and CuTe DSL", +) +def test_sm90_message_grid_forward_and_adjoint_match_real_module() -> None: + if tuple(torch.cuda.get_device_capability()) != (9, 0): + pytest.skip("SM90 message-grid differential requires an SM90 GPU") + + from deepmd.kernels.cute.neo.k1_message_grid_packed import ( + run_packed_message_grid_forward, + ) + from deepmd.kernels.cute.neo.message_grid_readout_sm90 import ( + prepare_sm90_message_grid_state, + run_sm90_message_grid_backward, + ) + + prior_precision = torch.get_float32_matmul_precision() + prior_tf32 = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = False + torch.set_float32_matmul_precision("highest") + try: + net = SO3GridNet( + lmax=3, + kmax=1, + channels=32, + n_focus=2, + mode="cross", + op_type="glu", + dtype=torch.float32, + layout="flat", + coefficient_layout="packed", + residual_scale_init=0.25, + trainable=False, + seed=None, + ).to("cuda") + generator = torch.Generator(device="cuda").manual_seed(20260816) + query = ( + 0.1 + * torch.randn( + 5, + 16, + 64, + generator=generator, + device="cuda", + dtype=torch.float32, + ) + ).requires_grad_(True) + context = ( + 0.1 + * torch.randn( + 16, + 5, + 64, + generator=generator, + device="cuda", + dtype=torch.float32, + ).permute(1, 0, 2) + ).requires_grad_(True) + assert context.stride() == (64, 5 * 64, 1) + grad_out = torch.randn( + query.shape, + generator=generator, + device="cuda", + dtype=torch.float32, + ) + + reference = net(query, context) + reference_query, reference_context = torch.autograd.grad( + reference, + (query, context), + grad_out, + ) + state = prepare_sm90_message_grid_state(net) + actual, product = run_packed_message_grid_forward( + net, + query.detach(), + context.detach(), + return_product=True, + sm90_state=state, + ) + actual_query, actual_context = run_sm90_message_grid_backward( + net, + query.detach(), + context.detach(), + grad_out, + product, + state, + ) + torch.cuda.synchronize() + + torch.testing.assert_close(actual, reference, atol=5.0e-5, rtol=5.0e-5) + torch.testing.assert_close( + actual_query, + reference_query, + atol=5.0e-5, + rtol=5.0e-5, + ) + torch.testing.assert_close( + actual_context, + reference_context, + atol=5.0e-5, + rtol=5.0e-5, + ) + finally: + torch.backends.cuda.matmul.allow_tf32 = prior_tf32 + torch.set_float32_matmul_precision(prior_precision) + + +@pytest.mark.skipif( + not _cute_cuda_runtime_available(), + reason="Q/K CuTe adjoint regression requires CUDA and CuTe DSL", +) +def test_qk_manual_adjoint_matches_autograd_with_sparse_edges() -> None: + from deepmd.kernels.cute.neo.k1_kernels.cute_neo_qk_edge import ( + compile_neo_qk_edge_backward, + compile_neo_qk_node_input_adjoint, + ) + + torch.manual_seed(20260813) + device = torch.device("cuda", torch.cuda.current_device()) + node_count = 7 + src = torch.tensor([0, 1, 1], dtype=torch.int32, device=device) + dst = torch.tensor([2, 2, 3], dtype=torch.int32, device=device) + edge_count = src.numel() + eps = 1.0e-5 + scale = 32.0**-0.5 + x_wide = torch.randn( + node_count, + 16, + 64, + dtype=torch.float32, + device=device, + requires_grad=True, + ) + q_weight = torch.randn(32, 2, 32, dtype=torch.float32, device=device) + k_weight = torch.randn_like(q_weight) + norm_scale = torch.randn(2, 32, dtype=torch.float32, device=device) + grad_logits = torch.randn(edge_count, 2, dtype=torch.float32, device=device) + + x_l0 = x_wide[:, 0, :].reshape(node_count, 2, 32) + x_norm = ( + x_l0 + * torch.rsqrt(x_l0.square().mean(dim=-1, keepdim=True) + eps) + * norm_scale.unsqueeze(0) + ) + q_node = torch.einsum("nfi,ifo->nfo", x_norm, q_weight) + k_node = torch.einsum("nfi,ifo->nfo", x_norm, k_weight) + logits = (q_node[dst.long()] * k_node[src.long()]).sum(dim=-1) * scale + reference = torch.autograd.grad(logits, x_wide, grad_logits)[0] + + runner = SimpleNamespace( + so2=SimpleNamespace( + attn_q_proj=SimpleNamespace(weight=q_weight.reshape(32, 64)), + attn_k_proj=SimpleNamespace(weight=k_weight.reshape(32, 64)), + attn_qk_norm=SimpleNamespace(adam_scale=norm_scale, eps=eps), + ), + node_count=node_count, + edge_count=edge_count, + x_wide=x_wide.detach(), + q_node=q_node.detach().contiguous(), + k_node=k_node.detach().contiguous(), + src_i32=src, + dst_i32=dst, + qk_edge_backward=compile_neo_qk_edge_backward(scale), + qk_node_input_adjoint=compile_neo_qk_node_input_adjoint(eps), + ) + actual = k1._qk_manual_backward(runner, grad_logits) + torch.cuda.synchronize(device) + + torch.testing.assert_close(actual, reference, atol=5.0e-5, rtol=5.0e-5) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_packed_impl.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_packed_impl.py new file mode 100644 index 0000000000..acc4e94659 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_packed_impl.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Implementation differentials for the live packed Neo Wigner path.""" + +from __future__ import ( + annotations, +) + +import importlib + +import pytest +import torch + +TOL = 5.0e-5 + + +def _cute_runtime_skip_reason() -> str | None: + if not torch.cuda.is_available(): + return "packed implementation differentials require CUDA" + try: + importlib.import_module("cutlass.cute") + importlib.import_module("cuda.bindings.driver") + except Exception as exc: # pragma: no cover - runtime dependent + return f"packed implementation differentials require CuTe DSL: {exc}" + return None + + +def _randn( + *shape: int, + generator: torch.Generator, + scale: float = 0.1, +) -> torch.Tensor: + return scale * torch.randn( + *shape, + generator=generator, + device="cuda", + dtype=torch.float32, + ) + + +def _dense_wigner(edge_count: int, generator: torch.Generator) -> torch.Tensor: + from deepmd.kernels.cute.neo import k1_wigner_layout as layout + + dense = torch.zeros( + edge_count, + 16, + 16, + device="cuda", + dtype=torch.float32, + ) + for start, stop in zip( + layout.FULL_BLOCK_OFFSETS[:-1], + layout.FULL_BLOCK_OFFSETS[1:], + strict=True, + ): + dense[:, start:stop, start:stop] = _randn( + edge_count, + stop - start, + stop - start, + generator=generator, + ) + return dense + + +_CUTE_SKIP_REASON = _cute_runtime_skip_reason() + + +@pytest.mark.skipif( + _CUTE_SKIP_REASON is not None, + reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable", +) +def test_panel_native_quaternion_backward_matches_dense_torch(): + from deepmd.kernels.cute.neo import k1_wigner_layout as layout + from deepmd.kernels.cute.neo import ( + k4_wignerd, + ) + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + WignerDCalculator, + ) + + generator = torch.Generator(device="cuda").manual_seed(20260702) + q_value = _randn(3, 4, generator=generator, scale=1.0) + q_value = q_value / q_value.norm(dim=-1, keepdim=True) + grad_panel = _randn(3, 46, generator=generator) + + q_panel = q_value.detach().requires_grad_(True) + panel = k4_wignerd._wignerd_panel_op(q_panel) + grad_q_panel = torch.autograd.grad((panel * grad_panel).sum(), q_panel)[0] + + q_dense = q_value.detach().requires_grad_(True) + dense, dense_t = WignerDCalculator(lmax=3, dtype=q_dense.dtype).to("cuda")(q_dense) + entries = tuple(layout.iter_packed_entries()) + rows = [entry.full_row for entry in entries] + cols = [entry.full_col for entry in entries] + dense_weight = torch.zeros_like(dense) + dense_weight[:, rows, cols] = grad_panel + dense_loss = (dense * dense_weight).sum() + 0.0 * dense_t.sum() + grad_q_dense = torch.autograd.grad(dense_loss, q_dense)[0] + + torch.testing.assert_close( + panel, + dense[..., rows, cols], + rtol=TOL, + atol=TOL, + ) + torch.testing.assert_close(grad_q_panel, grad_q_dense, rtol=TOL, atol=TOL) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py new file mode 100644 index 0000000000..4197bf4163 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py @@ -0,0 +1,323 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Storage-contract tests for the exact in-place Phase-C adjoint.""" + +from __future__ import ( + annotations, +) + +import itertools +from dataclasses import ( + fields, + replace, +) + +import pytest +import torch + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or torch.cuda.get_device_capability() not in {(8, 0), (9, 0)}, + reason="in-place Phase-C tests require sm_80 or sm_90", +) + +OUTPUT_FIELDS = ( + "grad_stack", + "grad_wigner_dt", + "grad_logits", + "grad_edge", + "grad_z_partial", + "grad_z", + "grad_focus_src", +) +NONALIASED_OUTPUT_FIELDS = OUTPUT_FIELDS[1:] +INPUT_FIELDS = ( + "grad_out", + "stack", + "wigner_dt", + "alpha", + "focus_alpha", + "dst_ptr", + "rotate_inv_rescale", + "logits", + "edge_gate", + "z_bias_raw", + "group_max", + "denom", + "focus_src", + "focus_weight", + "focus_scale", +) +REDZONE_ELEMENTS = 4 +SENTINEL = 937.25 + + +def _phase_c_api(): + pytest.importorskip("cutlass") + from deepmd.kernels.cute.neo.k1_kernels.cute_neo_phase_c_backward_layout_runner import ( + CuteNeoPhaseCBackwardLayout, + NeoPhaseCBackwardLayoutOutputs, + ) + + return CuteNeoPhaseCBackwardLayout, NeoPhaseCBackwardLayoutOutputs + + +def _runner(): + Runner, _ = _phase_c_api() + return Runner( + focus_eps=1.0e-8, + focus_tau=1.0, + focus_label_smoothing=0.0, + ) + + +def _inputs(dst_ptr_values: list[int]): + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed( + 20260722 + dst_ptr_values[-1] + ) + node_count = len(dst_ptr_values) - 1 + edge_count = dst_ptr_values[-1] + rand = lambda *shape: torch.randn( # noqa: E731 + *shape, + device=device, + dtype=torch.float32, + generator=generator, + ) + return { + "grad_out": rand(node_count, 16, 64), + "stack": rand(edge_count, 2, 10, 32), + "wigner_dt": rand(edge_count, 46), + "alpha": torch.softmax(rand(edge_count, 2), dim=0), + "focus_alpha": torch.softmax(rand(edge_count, 2), dim=1), + "dst_ptr": torch.tensor(dst_ptr_values, device=device, dtype=torch.int32), + "rotate_inv_rescale": rand(16), + "logits": rand(edge_count, 2), + "edge_gate": rand(edge_count).abs(), + "z_bias_raw": rand(2), + "group_max": rand(node_count, 2), + "denom": rand(node_count, 2).abs().add_(0.5), + "focus_src": rand(edge_count, 2, 32), + "focus_weight": rand(32, 2), + "focus_scale": rand(2, 32), + } + + +def _clone_inputs(inputs): + return {name: value.clone() for name, value in inputs.items()} + + +def _allocate_outputs(inputs): + _, Outputs = _phase_c_api() + edge_count = inputs["stack"].shape[0] + node_count = inputs["grad_out"].shape[0] + options = {"device": inputs["stack"].device, "dtype": torch.float32} + return Outputs( + grad_stack=inputs["stack"], + grad_wigner_dt=torch.empty(edge_count, 46, **options), + grad_logits=torch.empty(edge_count, 2, **options), + grad_edge=torch.empty(edge_count, **options), + grad_z_partial=torch.empty(node_count, 2, **options), + grad_z=torch.empty(2, **options), + grad_focus_src=torch.empty(2, edge_count, 32, **options), + ) + + +def _call(runner, inputs, outputs): + return runner( + inputs["grad_out"], + inputs["stack"], + inputs["wigner_dt"], + inputs["alpha"], + inputs["focus_alpha"], + inputs["dst_ptr"], + inputs["rotate_inv_rescale"], + inputs["logits"], + inputs["edge_gate"], + inputs["z_bias_raw"], + inputs["group_max"], + inputs["denom"], + inputs["focus_src"], + inputs["focus_weight"], + inputs["focus_scale"], + outputs, + ) + + +def _redzoned_like(tensor: torch.Tensor): + flat = torch.full( + (tensor.numel() + 2 * REDZONE_ELEMENTS,), + SENTINEL, + device=tensor.device, + dtype=tensor.dtype, + ) + value = flat[REDZONE_ELEMENTS : REDZONE_ELEMENTS + tensor.numel()].view_as(tensor) + return value, (flat[:REDZONE_ELEMENTS], flat[-REDZONE_ELEMENTS:]) + + +def _misaligned_like(tensor: torch.Tensor): + flat = torch.empty(tensor.numel() + 1, device=tensor.device, dtype=tensor.dtype) + value = flat[1:].view_as(tensor) + value.copy_(tensor) + assert value.data_ptr() % 16 != 0 + return value + + +def _assert_outputs_close(actual, expected): + for field in fields(type(expected)): + torch.testing.assert_close( + getattr(actual, field.name), + getattr(expected, field.name), + atol=1.0e-6, + rtol=1.0e-6, + ) + + +@pytest.mark.parametrize( + "dst_ptr_values", + ( + [0, 0, 2, 2, 7], + [0, 1, 1, 4, 12, 12], + [0, 0, 37, 37], + ), +) +def test_exact_inplace_matches_independent_run_and_preserves_redzones( + dst_ptr_values, +): + runner = _runner() + reference_inputs = _inputs(dst_ptr_values) + reference_outputs = _allocate_outputs(reference_inputs) + _call(runner, reference_inputs, reference_outputs) + + inputs = _inputs(dst_ptr_values) + stack, stack_redzones = _redzoned_like(inputs["stack"]) + stack.copy_(inputs["stack"]) + inputs["stack"] = stack + outputs = _allocate_outputs(inputs) + redzones = {"grad_stack": stack_redzones} + for field_name in NONALIASED_OUTPUT_FIELDS: + value, field_redzones = _redzoned_like(getattr(outputs, field_name)) + outputs = replace(outputs, **{field_name: value}) + redzones[field_name] = field_redzones + + _call(runner, inputs, outputs) + torch.cuda.synchronize() + + assert outputs.grad_stack.data_ptr() == inputs["stack"].data_ptr() + _assert_outputs_close(outputs, reference_outputs) + for field_name, (prefix, suffix) in redzones.items(): + assert torch.equal(prefix, torch.full_like(prefix, SENTINEL)), field_name + assert torch.equal(suffix, torch.full_like(suffix, SENTINEL)), field_name + + +def test_grad_stack_must_be_the_exact_input_view(): + inputs = _inputs([0, 3, 7]) + outputs = replace(_allocate_outputs(inputs), grad_stack=inputs["stack"].clone()) + + with pytest.raises(ValueError, match="exact in-place stack view"): + _call(_runner(), inputs, outputs) + + +def test_partial_or_shifted_stack_alias_is_rejected(): + inputs = _inputs([0, 3, 7]) + edge_count = inputs["stack"].shape[0] + base = torch.empty( + edge_count + 1, + 2, + 10, + 32, + device="cuda", + dtype=torch.float32, + ) + base[1:].copy_(inputs["stack"]) + inputs["stack"] = base[1:] + outputs = replace(_allocate_outputs(inputs), grad_stack=base[:-1]) + + with pytest.raises(ValueError, match="exact in-place stack view"): + _call(_runner(), inputs, outputs) + + +@pytest.mark.parametrize( + "tensor_name", + INPUT_FIELDS + tuple(f"outputs.{name}" for name in NONALIASED_OUTPUT_FIELDS), +) +def test_compiled_tensors_require_16_byte_alignment(tensor_name): + inputs = _inputs([0, 3, 7]) + outputs = _allocate_outputs(inputs) + + if tensor_name.startswith("outputs."): + field_name = tensor_name.removeprefix("outputs.") + outputs = replace( + outputs, + **{field_name: _misaligned_like(getattr(outputs, field_name))}, + ) + else: + inputs[tensor_name] = _misaligned_like(inputs[tensor_name]) + if tensor_name == "stack": + outputs = replace(outputs, grad_stack=inputs["stack"]) + + with pytest.raises(ValueError, match="must be 16-byte aligned"): + _call(_runner(), inputs, outputs) + + +@pytest.mark.parametrize("field_name", NONALIASED_OUTPUT_FIELDS) +def test_nonaliased_outputs_reject_input_overlap(field_name): + inputs = _inputs([0, 3, 7]) + outputs = _allocate_outputs(inputs) + target = getattr(outputs, field_name) + overlapping = inputs["stack"].view(-1)[: target.numel()].view_as(target) + outputs = replace(outputs, **{field_name: overlapping}) + + with pytest.raises(ValueError, match="must not overlap"): + _call(_runner(), inputs, outputs) + + +@pytest.mark.parametrize( + ("first_name", "second_name"), + tuple(itertools.combinations(NONALIASED_OUTPUT_FIELDS, 2)), +) +def test_nonaliased_output_pairs_must_not_overlap(first_name, second_name): + inputs = _inputs([0, 3, 7]) + outputs = _allocate_outputs(inputs) + first = getattr(outputs, first_name) + second = getattr(outputs, second_name) + slab = torch.empty( + max(first.numel(), second.numel()), + device="cuda", + dtype=torch.float32, + ) + outputs = replace( + outputs, + **{ + first_name: slab[: first.numel()].view_as(first), + second_name: slab[: second.numel()].view_as(second), + }, + ) + + with pytest.raises(ValueError, match="must not overlap output"): + _call(_runner(), inputs, outputs) + + +def test_retained_runner_supports_dynamic_edge_counts(): + runner = _runner() + for dst_ptr_values in ([0, 2, 7], [0, 0, 1, 9, 9]): + inputs = _inputs(dst_ptr_values) + original = _clone_inputs(inputs) + outputs = _allocate_outputs(inputs) + _call(runner, inputs, outputs) + first = { + field.name: getattr(outputs, field.name).clone() + for field in fields(outputs) + } + + inputs = original + outputs = _allocate_outputs(inputs) + _call(runner, inputs, outputs) + torch.cuda.synchronize() + for field_name, expected in first.items(): + torch.testing.assert_close( + getattr(outputs, field_name), + expected, + atol=1.0e-6, + rtol=1.0e-6, + ) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_focus_fusion.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_focus_fusion.py new file mode 100644 index 0000000000..afa6e94f7b --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_focus_fusion.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""SM80/SM90 differential coverage for Phase-C focus-source load fusion.""" + +from __future__ import ( + annotations, +) + +import unittest + +import torch + + +def _has_supported_gpu() -> bool: + return torch.cuda.is_available() and torch.cuda.get_device_capability() in { + (8, 0), + (9, 0), + } + + +def test_source_csr_preserves_equal_source_order() -> None: + from deepmd.kernels.cute.neo.k1_radial_phase_a_node import ( + build_source_csr, + ) + + src = torch.tensor([2, 0, 2, 1, 0, 2], dtype=torch.int64, device="cpu") + source_csr = build_source_csr(src, node_count=3) + + torch.testing.assert_close( + source_csr.source_order, + torch.tensor([1, 4, 3, 0, 2, 5], dtype=torch.int32, device="cpu"), + ) + + +@unittest.skipUnless(_has_supported_gpu(), "requires an SM80 or SM90 CUDA device") +class TestRadialFocusSourceFusion(unittest.TestCase): + def test_fused_load_matches_standalone_scalar_lane_add(self) -> None: + from deepmd.kernels.cute.neo.k1_radial_phase_a_node import ( + build_source_csr, + prepare_batched_radial_projection_weight, + run_neo_radial_phase_a_backward_node_tiled, + ) + + torch.manual_seed(18072026) + device = torch.device("cuda") + edge_count = 7 + node_count = 3 + + src = torch.tensor([0, 2, 1, 0, 2, 1, 2], device=device, dtype=torch.int64) + source_csr = build_source_csr(src, node_count) + grad_stack = torch.randn(edge_count, 2, 10, 32, device=device) + grad_focus_src = torch.randn(2, edge_count, 32, device=device) + grad_logits = torch.randn(edge_count, 2, device=device) + radial_compact = torch.randn(edge_count, 25, device=device) + combined_weight = torch.randn(128, 25, device=device) + attention_weight = torch.randn(32, 2, device=device) + channel_basis = torch.randn(64, device=device) + x_wide = torch.randn(node_count, 16 * 64, device=device) + d_full = torch.randn(edge_count, 46, device=device) + projection_weight = prepare_batched_radial_projection_weight( + combined_weight, + attention_weight, + ) + + standalone_grad_stack = grad_stack.clone() + standalone_grad_stack[:, :, 0, :].add_(grad_focus_src.permute(1, 0, 2)) + standalone = run_neo_radial_phase_a_backward_node_tiled( + standalone_grad_stack.view(edge_count, 2 * 10 * 32), + grad_logits, + radial_compact, + channel_basis, + x_wide, + source_csr.source_order, + source_csr.source_ptr, + d_full, + grad_focus_src_focus=torch.zeros_like(grad_focus_src), + batched_radial_projection_weight=projection_weight, + ) + fused = run_neo_radial_phase_a_backward_node_tiled( + grad_stack.view(edge_count, 2 * 10 * 32), + grad_logits, + radial_compact, + channel_basis, + x_wide, + source_csr.source_order, + source_csr.source_ptr, + d_full, + grad_focus_src_focus=grad_focus_src, + batched_radial_projection_weight=projection_weight, + ) + + torch.testing.assert_close( + fused.grad_x_wide, + standalone.grad_x_wide, + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + fused.grad_d_full, + standalone.grad_d_full, + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + fused.grad_radial_m0, + standalone.grad_radial_m0, + rtol=0.0, + atol=0.0, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_pitch68_safe.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_pitch68_safe.py new file mode 100644 index 0000000000..70771144f1 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_pitch68_safe.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Safety coverage for the fixed pitch-68 radial backward path.""" + +from __future__ import ( + annotations, +) + +import ast +import unittest +from pathlib import ( + Path, +) + +REPO_ROOT = Path(__file__).resolve().parents[4] +CUTE_PATH = REPO_ROOT / "deepmd/kernels/cute/neo" +WRAPPER_PATH = CUTE_PATH / "k1_radial_phase_a_node.py" +KERNEL_PATH = CUTE_PATH / "k1_kernels/cute_neo_radial_phase_a_backward_node.py" + +WARP_HELPERS = {"_warp_owned_grad_compact", "_warp_owned_grad_d"} + + +def _call_name(node: ast.Call) -> str | None: + if isinstance(node.func, ast.Name): + return node.func.id + if isinstance(node.func, ast.Attribute): + return node.func.attr + return None + + +def _is_constexpr_if(node: ast.If) -> bool: + test = node.test + return isinstance(test, ast.Call) and _call_name(test) == "const_expr" + + +def _function(tree: ast.AST, name: str) -> ast.FunctionDef: + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == name: + return node + raise AssertionError(f"missing function {name}") + + +def _dynamic_if_ancestors( + node: ast.AST, + parents: dict[ast.AST, ast.AST], + boundary: ast.FunctionDef, +) -> list[ast.If]: + result: list[ast.If] = [] + current = parents.get(node) + while current is not None and current is not boundary: + if isinstance(current, ast.If) and not _is_constexpr_if(current): + result.append(current) + current = parents.get(current) + return result + + +class TestRadialPitch68SafeStatic(unittest.TestCase): + def test_pitch68_specialization_contract(self) -> None: + wrapper_source = WRAPPER_PATH.read_text(encoding="utf-8") + kernel_source = KERNEL_PATH.read_text(encoding="utf-8") + + self.assertIn("SHARED_ROW_PITCH = HIDDEN + 4", kernel_source) + self.assertNotIn("pitch68_safe", wrapper_source) + self.assertNotIn("pitch68_safe", kernel_source) + + def test_warp_collectives_are_not_runtime_predicated(self) -> None: + source = KERNEL_PATH.read_text(encoding="utf-8") + tree = ast.parse(source) + parents = { + child: parent + for parent in ast.walk(tree) + for child in ast.iter_child_nodes(parent) + } + kernel = _function(tree, "neo_radial_phase_a_backward_node_kernel") + + helper_calls = [ + node + for node in ast.walk(kernel) + if isinstance(node, ast.Call) and _call_name(node) in WARP_HELPERS + ] + self.assertEqual( + sorted(_call_name(node) for node in helper_calls), + sorted(WARP_HELPERS), + ) + for call in helper_calls: + self.assertEqual( + _dynamic_if_ancestors(call, parents, kernel), + [], + f"{_call_name(call)} must be reached by every warp lane", + ) + + compact_call = next( + node + for node in helper_calls + if _call_name(node) == "_warp_owned_grad_compact" + ) + panel_call = next( + node for node in helper_calls if _call_name(node) == "_warp_owned_grad_d" + ) + self.assertEqual(ast.unparse(compact_call.args[3]), "safe_compact_idx") + self.assertEqual(ast.unparse(panel_call.args[2]), "safe_panel_idx") + + for helper_name in WARP_HELPERS: + helper = _function(tree, helper_name) + reductions = [ + node + for node in ast.walk(helper) + if isinstance(node, ast.Call) + and _call_name(node) == "warp_reduction_sum" + ] + self.assertEqual(len(reductions), 1) + self.assertEqual( + _dynamic_if_ancestors(reductions[0], parents, helper), + [], + f"warp collective in {helper_name} must be unconditional", + ) + + assignments = [ + node for node in ast.walk(kernel) if isinstance(node, ast.Assign) + ] + compact_write = next( + node + for node in assignments + if ast.unparse(node.targets[0]) == "grad_compact[compact_idx]" + ) + panel_write = next( + node + for node in assignments + if ast.unparse(node.targets[0]) == "params.grad_d_full[edge, panel_idx]" + ) + self.assertEqual( + { + ast.unparse(guard.test) + for guard in _dynamic_if_ancestors(compact_write, parents, kernel) + }, + {"compact_idx < COMPACT_WIDTH", "subgroup_lane == 0"}, + ) + self.assertEqual( + { + ast.unparse(guard.test) + for guard in _dynamic_if_ancestors(panel_write, parents, kernel) + }, + {"panel_idx < PACKED_WIGNER_VALUES", "subgroup_lane == 0"}, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py new file mode 100644 index 0000000000..3c2cad637a --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py @@ -0,0 +1,234 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""CPU contracts for the batched Neo radial adjoint projection.""" + +from __future__ import ( + annotations, +) + +import importlib.util +import sys +import unittest +from pathlib import ( + Path, +) +from types import ( + ModuleType, + SimpleNamespace, +) +from unittest import ( + mock, +) + +import torch + +REPO_ROOT = Path(__file__).resolve().parents[4] +CUTE_PATH = REPO_ROOT / "deepmd/kernels/cute/neo" +_TEST_PACKAGE = "sezm_cute_k1_radial_projection_contract_test" + + +def _load_cute_source_module(name: str) -> ModuleType: + qualified_name = f"{_TEST_PACKAGE}.{name}" + spec = importlib.util.spec_from_file_location( + qualified_name, + CUTE_PATH / f"{name}.py", + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"could not load {name}") + module = importlib.util.module_from_spec(spec) + sys.modules[qualified_name] = module + spec.loader.exec_module(module) + return module + + +_package = ModuleType(_TEST_PACKAGE) +_package.__path__ = [str(CUTE_PATH)] +sys.modules[_TEST_PACKAGE] = _package +_load_cute_source_module("compile_cache") +_RADIAL_PROJECTION = _load_cute_source_module("k1_radial_phase_a_node") +_load_cute_source_module("runtime_policy") +_K1_RUNNER = _load_cute_source_module("k1_runner") + +COMPACT_WIDTH = _RADIAL_PROJECTION.COMPACT_WIDTH +FOCUS_COUNT = _RADIAL_PROJECTION.FOCUS_COUNT +FOCUS_HIDDEN = _RADIAL_PROJECTION.FOCUS_HIDDEN +PROJECTION_INPUT_WIDTH = _RADIAL_PROJECTION.PROJECTION_INPUT_WIDTH +RADIAL_WIDTH = _RADIAL_PROJECTION.RADIAL_WIDTH +_project_batched_radial_adjoint = _RADIAL_PROJECTION._project_batched_radial_adjoint +prepare_batched_radial_projection_weight = ( + _RADIAL_PROJECTION.prepare_batched_radial_projection_weight +) +_batched_radial_projection_weight = _K1_RUNNER._batched_radial_projection_weight + + +class TestBatchedRadialProjection(unittest.TestCase): + def test_precombined_single_gemm_matches_two_projection_algebra(self): + torch.manual_seed(20260718) + edge_count = 7 + grad_compact = torch.randn( + edge_count, COMPACT_WIDTH, dtype=torch.float32, device="cpu" + ) + grad_logits = torch.randn( + edge_count, FOCUS_COUNT, dtype=torch.float32, device="cpu" + ) + combined_weight = torch.randn( + RADIAL_WIDTH, + COMPACT_WIDTH, + dtype=torch.float32, + device="cpu", + ) + attention_weight = torch.randn( + FOCUS_HIDDEN, + FOCUS_COUNT, + dtype=torch.float32, + device="cpu", + ) + projection_weight = prepare_batched_radial_projection_weight( + combined_weight, + attention_weight, + ) + + expected = grad_compact @ combined_weight.transpose(0, 1) + expected[:, :FOCUS_HIDDEN].add_(grad_logits @ attention_weight.transpose(0, 1)) + workspace = torch.full( + (edge_count, FOCUS_COUNT * 10 * FOCUS_HIDDEN), + torch.nan, + dtype=torch.float32, + device="cpu", + ) + actual = torch.empty( + edge_count, + RADIAL_WIDTH, + dtype=torch.float32, + device="cpu", + ) + workspace[:, :COMPACT_WIDTH].copy_(grad_compact) + workspace[:, COMPACT_WIDTH:PROJECTION_INPUT_WIDTH].copy_(grad_logits) + + with ( + mock.patch.object(torch, "cat", wraps=torch.cat) as cat, + mock.patch.object(torch, "mm", wraps=torch.mm) as mm, + ): + _project_batched_radial_adjoint( + projection_weight, + actual, + workspace, + ) + + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=2e-5) + self.assertEqual(mm.call_count, 1) + self.assertEqual(cat.call_count, 0) + gemm_input, gemm_weight = mm.call_args.args + self.assertEqual(tuple(gemm_input.shape), (edge_count, PROJECTION_INPUT_WIDTH)) + self.assertEqual( + tuple(gemm_weight.shape), (PROJECTION_INPUT_WIDTH, RADIAL_WIDTH) + ) + self.assertEqual(gemm_input.stride(), (workspace.shape[1], 1)) + self.assertEqual( + gemm_input.untyped_storage().data_ptr(), + workspace.untyped_storage().data_ptr(), + ) + self.assertIs(mm.call_args.kwargs["out"], actual) + torch.testing.assert_close( + workspace[:, :COMPACT_WIDTH], + grad_compact, + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + workspace[:, COMPACT_WIDTH:PROJECTION_INPUT_WIDTH], + grad_logits, + rtol=0.0, + atol=0.0, + ) + self.assertTrue(torch.isnan(workspace[:, PROJECTION_INPUT_WIDTH:]).all()) + + def test_precombined_weight_layout_and_cache_contract(self): + combined_weight = torch.randn( + RADIAL_WIDTH, + COMPACT_WIDTH, + dtype=torch.float32, + device="cpu", + ) + attention_weight = torch.randn( + FOCUS_HIDDEN, + FOCUS_COUNT, + dtype=torch.float32, + device="cpu", + ) + owner = SimpleNamespace() + + first = _batched_radial_projection_weight( + owner, + combined_weight, + attention_weight, + ) + second = _batched_radial_projection_weight( + owner, + combined_weight, + attention_weight, + ) + + self.assertIs(first, second) + self.assertEqual(tuple(first.shape), (PROJECTION_INPUT_WIDTH, RADIAL_WIDTH)) + self.assertTrue(first.is_contiguous()) + torch.testing.assert_close( + first[:COMPACT_WIDTH], + combined_weight.transpose(0, 1), + rtol=0.0, + atol=0.0, + ) + attention_panel = torch.zeros( + FOCUS_COUNT, + 2 * FOCUS_HIDDEN, + device="cpu", + ) + attention_panel[:, :FOCUS_HIDDEN].copy_(attention_weight.transpose(0, 1)) + torch.testing.assert_close( + first[COMPACT_WIDTH:, : 2 * FOCUS_HIDDEN], + attention_panel, + rtol=0.0, + atol=0.0, + ) + self.assertTrue(torch.count_nonzero(first[COMPACT_WIDTH:, 64:]) == 0) + + attention_weight.add_(1.0) + updated = _batched_radial_projection_weight( + owner, + combined_weight, + attention_weight, + ) + self.assertIsNot(first, updated) + + def test_projection_requires_strict_fp32_and_full_consumed_workspace(self): + combined_weight = torch.randn( + RADIAL_WIDTH, + COMPACT_WIDTH, + device="cpu", + ) + attention_weight = torch.randn( + FOCUS_HIDDEN, + FOCUS_COUNT, + device="cpu", + ) + with self.assertRaisesRegex(TypeError, "torch.float32"): + prepare_batched_radial_projection_weight( + combined_weight.double(), + attention_weight.double(), + ) + + projection_weight = prepare_batched_radial_projection_weight( + combined_weight, + attention_weight, + ) + edge_count = 2 + with self.assertRaisesRegex(ValueError, "consumed_workspace must have shape"): + _project_batched_radial_adjoint( + projection_weight, + torch.empty(edge_count, RADIAL_WIDTH, device="cpu"), + torch.empty(edge_count, PROJECTION_INPUT_WIDTH, device="cpu"), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py new file mode 100644 index 0000000000..8bed9d3942 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py @@ -0,0 +1,516 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Contracts for the shared structural-gate vec4 path.""" + +from __future__ import ( + annotations, +) + +import ast +import importlib.util +import math +import random +import struct +import unittest +from pathlib import ( + Path, +) +from typing import ( + TYPE_CHECKING, +) + +if TYPE_CHECKING: + from types import ( + ModuleType, + ) + + +REPO_ROOT = Path(__file__).resolve().parents[4] +CUTE_ROOT = REPO_ROOT / "deepmd/kernels/cute/neo" +KERNEL_PATH = CUTE_ROOT / "k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py" +WRAPPER_PATH = CUTE_ROOT / "k1_gate_structural.py" +DYNAMIC_EDGE_COUNTS = (0, 1, 7, 8, 9, 31, 32, 33) + + +def _f32(value: float) -> float: + return struct.unpack("f", struct.pack("f", value))[0] + + +def _sigmoid_f32(value: float) -> float: + value = _f32(value) + return _f32(_f32(1.0) / _f32(_f32(1.0) + _f32(math.exp(_f32(-value))))) + + +def _mul_add_f32(lhs: float, rhs: float, addend: float) -> float: + return _f32(_f32(lhs * rhs) + addend) + + +def _reference_forward( + residual: list[float], + y: list[float], + logits: list[float], + edge_count: int, +) -> list[float]: + out = residual.copy() + for edge in range(edge_count): + for focus in range(2): + row = edge * 2 + focus + row_base = row * 10 * 32 + logits_base = (focus * edge_count + edge) * 3 * 32 + gates = [ + [ + _sigmoid_f32(logits[logits_base + gate * 32 + channel]) + for channel in range(32) + ] + for gate in range(3) + ] + for channel in range(32): + index = row_base + channel + y0 = y[index] + out[index] = _mul_add_f32(y0, _sigmoid_f32(y0), out[index]) + for degree in range(1, 10): + gate = gates[(degree - 1) % 3] + for channel in range(32): + index = row_base + degree * 32 + channel + out[index] = _mul_add_f32( + y[index], + gate[channel], + out[index], + ) + return out + + +def _vec4_forward_inplace_model( + out: list[float], + y: list[float], + logits: list[float], + edge_count: int, +) -> None: + rows = edge_count * 2 + for block_row in range((rows + 15) // 16): + for row_slot in range(16): + row = block_row * 16 + row_slot + if row >= rows: + continue + edge, focus = divmod(row, 2) + row_base = row * 10 * 32 + logits_base = (focus * edge_count + edge) * 3 * 32 + for channel_group in range(8): + channel_base = channel_group * 4 + for lane in range(4): + index = row_base + channel_base + lane + y0 = y[index] + out[index] = _mul_add_f32(y0, _sigmoid_f32(y0), out[index]) + for gate_index in range(3): + gates = [ + _sigmoid_f32( + logits[logits_base + gate_index * 32 + channel_base + lane] + ) + for lane in range(4) + ] + for repeat in range(3): + degree = 1 + gate_index + repeat * 3 + for lane in range(4): + index = row_base + degree * 32 + channel_base + lane + out[index] = _mul_add_f32( + y[index], + gates[lane], + out[index], + ) + + +def _backward_reference( + grad_out: list[float], + y: list[float], + logits: list[float], + edge_count: int, +) -> tuple[list[float], list[float]]: + grad_y = [0.0] * len(y) + grad_logits = [0.0] * len(logits) + for edge in range(edge_count): + for focus in range(2): + row = edge * 2 + focus + row_base = row * 10 * 32 + logits_base = (focus * edge_count + edge) * 3 * 32 + for channel in range(32): + y0 = _f32(y[row_base + channel]) + sig0 = _sigmoid_f32(y0) + grad0 = _f32(grad_out[row_base + channel]) + inner = _f32(_f32(1.0) + _f32(y0 * _f32(_f32(1.0) - sig0))) + grad_y[row_base + channel] = _f32(_f32(grad0 * sig0) * inner) + for gate_index in range(3): + gate_offset = logits_base + gate_index * 32 + channel + gate = _sigmoid_f32(logits[gate_offset]) + grad_logit = _f32(0.0) + for repeat in range(3): + degree = 1 + gate_index + repeat * 3 + index = row_base + degree * 32 + channel + gout = _f32(grad_out[index]) + grad_y[index] = _f32(gout * gate) + term = _f32(gout * _f32(y[index])) + term = _f32(term * gate) + term = _f32(term * _f32(_f32(1.0) - gate)) + grad_logit = _f32(grad_logit + term) + grad_logits[gate_offset] = grad_logit + return grad_y, grad_logits + + +def _module_constants(tree: ast.Module) -> dict[str, object]: + def constant_value(node: ast.expr, namespace: dict[str, object]) -> object: + if isinstance(node, ast.Constant): + return node.value + if isinstance(node, ast.Name): + return namespace[node.id] + if isinstance(node, ast.BinOp): + lhs = constant_value(node.left, namespace) + rhs = constant_value(node.right, namespace) + if isinstance(node.op, ast.FloorDiv): + return lhs // rhs + if isinstance(node.op, ast.Mult): + return lhs * rhs + raise ValueError("not a supported module constant") + + constants: dict[str, object] = {} + namespace: dict[str, object] = {} + for node in tree.body: + if not isinstance(node, ast.Assign) or len(node.targets) != 1: + continue + target = node.targets[0] + if not isinstance(target, ast.Name): + continue + try: + value = constant_value(node.value, namespace) + except (KeyError, TypeError, ValueError, ZeroDivisionError): + continue + constants[target.id] = value + namespace[target.id] = value + return constants + + +def _load_wrapper_module() -> ModuleType: + spec = importlib.util.spec_from_file_location( + "_deepmd_test_k1_gate_structural", + WRAPPER_PATH, + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load {WRAPPER_PATH}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class TestSM80StructuralGateVec4Contract(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.kernel_source = KERNEL_PATH.read_text() + cls.constants = _module_constants(ast.parse(cls.kernel_source)) + + def test_vector_layout_and_launch_are_fixed(self): + self.assertEqual(self.constants["FOCUS_COUNT"], 2) + self.assertEqual(self.constants["REDUCED_COUNT"], 10) + self.assertEqual(self.constants["CHANNELS"], 32) + self.assertEqual(self.constants["VECTOR_WIDTH"], 4) + self.assertEqual(self.constants["CHANNEL_GROUPS"], 8) + self.assertEqual(self.constants["ROWS_PER_BLOCK"], 16) + self.assertEqual(self.constants["THREADS"], 128) + self.assertIn("residual.element_type.width * VECTOR_WIDTH", self.kernel_source) + self.assertIn("cute.make_tiled_copy_tv", self.kernel_source) + self.assertIn("(1, CHANNEL_GROUPS)", self.kernel_source) + self.assertIn("(1, VECTOR_WIDTH)", self.kernel_source) + + def test_kernel_is_strict_fp32_edge_major_only(self): + for forbidden in ( + "Float16", + "BFloat16", + "TensorFloat32", + ): + self.assertNotIn(forbidden, self.kernel_source) + self.assertIn("cutlass.Float32", self.kernel_source) + self.assertIn('"assumed_align": 16', self.kernel_source) + self.assertIn("_guard_vec4_dispatch", self.kernel_source) + self.assertIn("runtime_policy.SUPPORTED_K1_CAPABILITIES", self.kernel_source) + + +class TestSM80StructuralGateVec4Arithmetic(unittest.TestCase): + def test_forward_dynamic_edges_and_block_tails(self): + for edge_count in DYNAMIC_EDGE_COUNTS: + with self.subTest(edge_count=edge_count): + rng = random.Random(20260721 + edge_count) + values = edge_count * 2 * 10 * 32 + residual = [_f32(rng.uniform(-2.0, 2.0)) for _ in range(values)] + y = [_f32(rng.uniform(-2.0, 2.0)) for _ in range(values)] + logits = [ + _f32(rng.uniform(-4.0, 4.0)) for _ in range(2 * edge_count * 3 * 32) + ] + expected = _reference_forward(residual, y, logits, edge_count) + actual = residual.copy() + _vec4_forward_inplace_model(actual, y, logits, edge_count) + self.assertEqual(actual, expected) + + def test_backward_dynamic_edges_and_block_tails(self): + for edge_count in DYNAMIC_EDGE_COUNTS: + with self.subTest(edge_count=edge_count): + rng = random.Random(20260722 + edge_count) + values = edge_count * 2 * 10 * 32 + grad_out = [_f32(rng.uniform(-2.0, 2.0)) for _ in range(values)] + y = [_f32(rng.uniform(-2.0, 2.0)) for _ in range(values)] + logits = [ + _f32(rng.uniform(-4.0, 4.0)) for _ in range(2 * edge_count * 3 * 32) + ] + grad_y, grad_logits = _backward_reference( + grad_out, + y, + logits, + edge_count, + ) + self.assertEqual(len(grad_y), values) + self.assertEqual(len(grad_logits), len(logits)) + + +try: + import torch +except ImportError: + torch = None + + +@unittest.skipUnless(torch is not None, "PyTorch is required") +class TestSM80StructuralGateVec4Dispatch(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.dispatch = staticmethod( + _load_wrapper_module()._dispatch_aligned_vec4_kernel + ) + + def test_aligned_fp32_tensor_dispatches(self): + calls = [] + tensor = torch.empty(16, dtype=torch.float32, device="cpu") + + def kernel(value): + calls.append(value) + return "dispatched" + + result = self.dispatch(kernel, ("value",), tensor) + self.assertEqual(result, "dispatched") + self.assertEqual(calls, [tensor]) + + def test_misaligned_storage_offset_is_rejected_before_dispatch(self): + calls = [] + tensor = torch.empty(17, dtype=torch.float32, device="cpu")[1:] + self.assertTrue(tensor.is_contiguous()) + self.assertEqual(tensor.storage_offset(), 1) + + with self.assertRaisesRegex(ValueError, "storage offsets divisible by 4"): + self.dispatch(lambda value: calls.append(value), ("value",), tensor) + self.assertEqual(calls, []) + + def test_misaligned_pointer_is_rejected_before_dispatch(self): + class MisalignedTensor: + dtype = torch.float32 + shape = (4,) + + @staticmethod + def numel(): + return 4 + + @staticmethod + def is_contiguous(): + return True + + @staticmethod + def stride(): + return (1,) + + @staticmethod + def storage_offset(): + return 0 + + @staticmethod + def data_ptr(): + return 4 + + calls = [] + with self.assertRaisesRegex(ValueError, "16-byte-aligned"): + self.dispatch( + lambda value: calls.append(value), + ("value",), + MisalignedTensor(), + ) + self.assertEqual(calls, []) + + def test_noncompact_and_non_fp32_tensors_are_rejected(self): + noncompact = torch.empty(4, 8, dtype=torch.float32, device="cpu").T + with self.assertRaisesRegex(ValueError, "compact tensors"): + self.dispatch(lambda value: value, ("value",), noncompact) + + with self.assertRaisesRegex(TypeError, "requires float32"): + self.dispatch( + lambda value: value, + ("value",), + torch.empty(4, dtype=torch.float64, device="cpu"), + ) + + def test_empty_tensor_does_not_launch(self): + calls = [] + result = self.dispatch( + lambda value: calls.append(value), + ("value",), + torch.empty(0, dtype=torch.float32, device="cpu"), + ) + self.assertIsNone(result) + self.assertEqual(calls, []) + + +@unittest.skipUnless( + torch is not None + and torch.cuda.is_available() + and tuple(torch.cuda.get_device_capability()) + in {(8, 0), (8, 6), (8, 9), (9, 0), (10, 0), (12, 0)}, + "A supported Neo K1 CUDA runtime is required", +) +class TestSM80StructuralGateVec4CudaDifferential(unittest.TestCase): + @classmethod + def setUpClass(cls): + from deepmd.kernels.cute.neo.k1_kernels.cute_neo_gate_split_structural_vec4_sm80 import ( + compile_neo_gate_split_structural_vec4_sm80_backward, + compile_neo_gate_split_structural_vec4_sm80_forward, + ) + + capability = tuple(torch.cuda.get_device_capability()) + compile_identity = (torch.cuda.current_device(), *capability) + cls.vec4_forward = staticmethod( + compile_neo_gate_split_structural_vec4_sm80_forward(compile_identity) + ) + cls.vec4_backward = staticmethod( + compile_neo_gate_split_structural_vec4_sm80_backward(compile_identity) + ) + + def test_forward_dynamic_edges_tails_and_empty(self): + for edge_count in DYNAMIC_EDGE_COUNTS: + with self.subTest(edge_count=edge_count): + torch.manual_seed(20260721 + edge_count) + residual = torch.randn( + edge_count, + 2, + 10, + 32, + device="cuda", + dtype=torch.float32, + ) + y = torch.randn_like(residual) + logits = torch.randn( + 2, + edge_count, + 3 * 32, + device="cuda", + dtype=torch.float32, + ) + gates = torch.sigmoid( + logits.permute(1, 0, 2).reshape(edge_count, 2, 3, 32) + ) + expected = residual.clone() + expected[:, :, 0, :].add_(torch.nn.functional.silu(y[:, :, 0, :])) + for degree in range(1, 10): + expected[:, :, degree, :].add_( + y[:, :, degree, :] * gates[:, :, (degree - 1) % 3, :] + ) + + vec4 = residual.clone() + self.vec4_forward( + vec4.view(edge_count * 2, 10 * 32), + y.view(edge_count * 2, 10 * 32), + logits, + vec4.view(edge_count * 2, 10 * 32), + ) + torch.testing.assert_close(vec4, expected, atol=5e-5, rtol=5e-5) + + def test_backward_dynamic_edges_tails_and_empty(self): + gate_indices = torch.tensor( + [0, 1, 2, 0, 1, 2, 0, 1, 2], + device="cuda", + ) + for edge_count in DYNAMIC_EDGE_COUNTS: + with self.subTest(edge_count=edge_count): + torch.manual_seed(20260722 + edge_count) + grad_out = torch.randn( + edge_count, + 2, + 10, + 32, + device="cuda", + dtype=torch.float32, + ) + y = torch.randn_like(grad_out) + logits = torch.randn( + 2, + edge_count, + 3 * 32, + device="cuda", + dtype=torch.float32, + ) + vec4_grad_y = torch.empty_like(y) + vec4_grad_logits = torch.empty_like(logits) + self.vec4_backward( + grad_out.view(edge_count * 2, 10 * 32), + y.view(edge_count * 2, 10 * 32), + logits, + vec4_grad_y.view(edge_count * 2, 10 * 32), + vec4_grad_logits, + ) + + y_reference = y.detach().clone().requires_grad_(True) + logits_reference = logits.detach().clone().requires_grad_(True) + gates = torch.sigmoid( + logits_reference.permute(1, 0, 2).reshape( + edge_count, + 2, + 3, + 32, + ) + ) + output = torch.cat( + ( + torch.nn.functional.silu(y_reference[:, :, :1, :]), + y_reference[:, :, 1:, :] * gates.index_select(2, gate_indices), + ), + dim=2, + ) + expected_grad_y, expected_grad_logits = torch.autograd.grad( + output, + (y_reference, logits_reference), + grad_out, + ) + torch.testing.assert_close( + vec4_grad_y, + expected_grad_y, + atol=5e-5, + rtol=5e-5, + ) + torch.testing.assert_close( + vec4_grad_logits, + expected_grad_logits, + atol=5e-5, + rtol=5e-5, + ) + + def test_misaligned_cuda_view_is_rejected(self): + edge_count = 1 + storage = torch.empty( + edge_count * 2 * 10 * 32 + 1, + device="cuda", + dtype=torch.float32, + ) + residual = storage[1:].view(edge_count * 2, 10 * 32) + y = torch.empty_like(residual) + logits = torch.empty( + 2, + edge_count, + 3 * 32, + device="cuda", + dtype=torch.float32, + ) + with self.assertRaisesRegex(ValueError, "storage offsets divisible by 4"): + self.vec4_forward(residual, y, logits, residual) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_memory_reuse.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_memory_reuse.py new file mode 100644 index 0000000000..76b4ae5b55 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_memory_reuse.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Behavioral contracts for K1 structural storage reuse.""" + +from __future__ import ( + annotations, +) + +import importlib.util +import sys +from pathlib import ( + Path, +) +from types import ( + SimpleNamespace, +) + +from .test_descriptor_sezm_cute_k1 import ( + _K1, + NeoFullCuteBackward, + NeoK1BackwardWorkspace, + NeoK1RuntimeConfig, + StackCache, + _validate_runtime_config, +) + +REPO_ROOT = Path(__file__).resolve().parents[4] +STRUCTURAL_HELPER = REPO_ROOT / "deepmd/kernels/cute/neo/k1_gate_structural.py" +SO2_HELPER = REPO_ROOT / "deepmd/kernels/cute/neo/k1_so2linear.py" + + +def _load_module(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + if spec is None or spec.loader is None: + raise RuntimeError(f"could not load {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + try: + spec.loader.exec_module(module) + finally: + sys.modules.pop(name, None) + return module + + +def _valid_structural_config() -> NeoK1RuntimeConfig: + return NeoK1RuntimeConfig(per_focus_so2_fwd_pair=True) + + +def _storage_ptr(tensor) -> int: + return tensor.untyped_storage().data_ptr() + + +def _workspace_inputs(edge_count: int): + torch = _K1.torch + return { + "edge_count": edge_count, + "node_count": 3, + "d_full": torch.empty(edge_count, 46, dtype=torch.float32, device="cpu"), + "dt_full": torch.empty(edge_count, 46, dtype=torch.float32, device="cpu"), + "radial": torch.empty(edge_count, 4, 32, dtype=torch.float32, device="cpu"), + } + + +def _so2_linear(torch): + return SimpleNamespace( + lmax=3, + mmax=1, + in_channels=32, + out_channels=32, + n_focus=2, + mlp_bias=False, + weight_m0=torch.randn(4 * 32, 2 * 4 * 32, device="cpu"), + weight_m=(torch.randn(3 * 32, 2 * 2 * 3 * 32, device="cpu"),), + ) + + +def test_structural_memory_reuse_profile_is_valid() -> None: + config = _valid_structural_config() + assert _validate_runtime_config(config, compute_capability=(8, 0)) is None + + +def test_final_so2_single_input_fold_matches_residual_plus_linear(): + torch = _K1.torch + helper = _load_module("sezm_cute_k1_so2_test", SO2_HELPER) + torch.manual_seed(20260705) + x_local = torch.randn(3, 2, 10, 32, dtype=torch.float32, device="cpu") + linear = _so2_linear(torch) + + linear_only = helper.run_neo_so2_linear_manual(linear, x_local) + folded = helper.run_neo_so2_linear_manual( + linear, + x_local, + add_residual=True, + ) + + torch.testing.assert_close(folded, linear_only + x_local, atol=5e-5, rtol=5e-5) + + +def test_per_focus_pair_forward_matches_batched_forward(): + torch = _K1.torch + helper = _load_module("sezm_cute_k1_so2_per_focus_test", SO2_HELPER) + torch.manual_seed(20260718) + x_local = torch.randn(17, 2, 10, 32, dtype=torch.float32, device="cpu") + linear = _so2_linear(torch) + + batched = helper.run_neo_so2_linear_manual(linear, x_local) + per_focus = helper.run_neo_so2_linear_manual( + linear, + x_local, + per_focus_pair=True, + ) + + torch.testing.assert_close(per_focus, batched, atol=5e-5, rtol=5e-5) + + +def test_structural_workspace_reuses_saved_slabs_and_gate_panel(): + torch = _K1.torch + edge_count = 5 + phase_c_stack = torch.empty( + edge_count, 2, 10, 32, dtype=torch.float32, device="cpu" + ) + phase_c_y = torch.empty_like(phase_c_stack) + radial_scratch = torch.empty( + 2, edge_count, 3 * 32, dtype=torch.float32, device="cpu" + ) + + workspace = NeoK1BackwardWorkspace( + torch, + **_workspace_inputs(edge_count), + phase_c_stack=phase_c_stack, + phase_c_y=phase_c_y, + radial_scratch=radial_scratch, + structural_memory_reuse=True, + ) + + stack_storage = _storage_ptr(phase_c_stack) + assert _storage_ptr(workspace.grad_stack_focus) == stack_storage + assert _storage_ptr(workspace.grad_y) == stack_storage + assert _storage_ptr(workspace.grad_x_rot) == stack_storage + assert _storage_ptr(workspace.grad_d) == stack_storage + assert _storage_ptr(workspace.grad_radial_flat) == _storage_ptr(radial_scratch) + assert _storage_ptr(workspace.grad_mixed_slab) == _storage_ptr(phase_c_y) + assert workspace.grad_gate_logits is None + + +def test_single_input_workspace_omits_phase_c_y_and_keeps_mixed_output_distinct(): + torch = _K1.torch + edge_count = 5 + phase_c_stack = torch.empty( + edge_count, 2, 10, 32, dtype=torch.float32, device="cpu" + ) + radial_scratch = torch.empty(edge_count, 4, 32, dtype=torch.float32, device="cpu") + + workspace = NeoK1BackwardWorkspace( + torch, + **_workspace_inputs(edge_count), + phase_c_stack=phase_c_stack, + phase_c_y=None, + radial_scratch=radial_scratch, + structural_memory_reuse=True, + phase_c_single_input_reuse=True, + ) + + assert _storage_ptr(workspace.grad_stack_focus) == _storage_ptr(phase_c_stack) + assert _storage_ptr(workspace.grad_mixed_slab) != _storage_ptr(phase_c_stack) + assert _storage_ptr(workspace.grad_radial_flat) == _storage_ptr(radial_scratch) + assert workspace.grad_mixed_slab.shape == phase_c_stack.shape + + +def test_runner_routes_saved_structural_buffers_into_lazy_workspace(): + torch = _K1.torch + edge_count = 5 + runner = object.__new__(NeoFullCuteBackward) + runner.torch = torch + runner.config = _valid_structural_config() + runner.edge_count = edge_count + runner.node_count = 3 + runner.d = torch.empty(edge_count, 46, dtype=torch.float32, device="cpu") + runner.dt = torch.empty_like(runner.d) + runner.radial = torch.empty(edge_count, 4, 32, dtype=torch.float32, device="cpu") + runner.phase_c_stack = torch.empty( + edge_count, 2, 10, 32, dtype=torch.float32, device="cpu" + ) + runner.phase_c_y = None + first_logits = torch.empty(2, edge_count, 3 * 32, dtype=torch.float32, device="cpu") + runner.stack_caches = [ + StackCache( + torch.empty_like(runner.phase_c_stack), first_logits, object(), False + ), + StackCache( + torch.empty_like(runner.phase_c_stack), + torch.empty_like(first_logits), + object(), + False, + ), + StackCache(torch.empty_like(runner.phase_c_stack), None, object(), True), + ] + runner._backward_workspace = None + + workspace = runner.ensure_backward_workspace() + + assert runner.ensure_backward_workspace() is workspace + assert _storage_ptr(workspace.grad_stack_focus) == _storage_ptr( + runner.phase_c_stack + ) + assert _storage_ptr(workspace.grad_radial_flat) == _storage_ptr(first_logits) + assert _storage_ptr(workspace.grad_mixed_slab) != _storage_ptr(runner.phase_c_stack) + + +def test_structural_backward_can_overwrite_consumed_logits(): + torch = _K1.torch + helper = _load_module("sezm_cute_k1_structural_test", STRUCTURAL_HELPER) + edge_count = 3 + grad_out = torch.randn(edge_count, 2, 10, 32, dtype=torch.float32, device="cpu") + y = torch.randn_like(grad_out) + logits = torch.randn(2, edge_count, 3 * 32, dtype=torch.float32, device="cpu") + grad_y = torch.empty_like(y) + + def fake_kernel(grad_out_flat, y_flat, logits_in, grad_y_flat, grad_logits): + assert logits_in is logits + assert grad_logits is logits + grad_y_flat.copy_(grad_out_flat + y_flat) + grad_logits.fill_(7.0) + + result = helper.run_structural_gate_backward( + fake_kernel, + grad_out, + y, + logits, + grad_y, + grad_logits=None, + overwrite_logits=True, + ) + + assert result is logits + torch.testing.assert_close(grad_y, grad_out + y) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_output_grid_product.py b/source/tests/pt/model/test_descriptor_sezm_cute_output_grid_product.py new file mode 100644 index 0000000000..00335a9258 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_output_grid_product.py @@ -0,0 +1,539 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Behavioral differentials for the fused Neo output-grid product.""" + +from __future__ import ( + annotations, +) + +import importlib +import sys +import types + +import pytest +import torch + +TOL = 5.0e-5 +N_FRAMES = 3 +COEFF_DIM = 16 +GRID_SIZE = 152 +SUPPORTED_HIDDEN_CHANNELS = (96, 192) + + +def _cute_runtime_skip_reason() -> str | None: + if not torch.cuda.is_available(): + return "output-grid differentials require CUDA" + if tuple(torch.cuda.get_device_capability()) not in {(8, 0), (8, 6), (9, 0)}: + return "output-grid CuTe dispatch supports only sm80, sm86, and sm90" + try: + importlib.import_module("cutlass.cute") + importlib.import_module("cuda.bindings.driver") + except Exception as exc: # pragma: no cover - runtime dependent + return f"output-grid differentials require the CuTe DSL runtime: {exc}" + return None + + +_CUTE_SKIP_REASON = _cute_runtime_skip_reason() + + +def _sm80_skip_reason() -> str | None: + if _CUTE_SKIP_REASON is not None: + return _CUTE_SKIP_REASON + if tuple(torch.cuda.get_device_capability()) not in {(8, 0), (8, 6)}: + return "specialized output-grid differentials require sm80 or sm86" + return None + + +_SM80_SKIP_REASON = _sm80_skip_reason() + + +def _reference(left, right, to_grid, from_grid): + nodes = left.shape[0] + hidden_channels = left.shape[-1] // N_FRAMES + left_flat = left.reshape(nodes, COEFF_DIM * N_FRAMES, hidden_channels) + right_flat = right.reshape(nodes, COEFF_DIM * N_FRAMES, hidden_channels) + left_grid = torch.einsum("gj,njc->ngc", to_grid, left_flat) + right_grid = torch.einsum("gj,njc->ngc", to_grid, right_flat) + out = torch.einsum("jg,ngc->njc", from_grid, left_grid * right_grid) + return out.reshape_as(left) + + +def test_grid_mlp_accepts_fused_middle_callback(): + from deepmd.pt.model.descriptor.sezm_nn.grid_net import ( + GridMLP, + ) + + module = GridMLP( + channels=2, + mode="self", + n_frames=3, + dtype=torch.float32, + trainable=False, + seed=7, + ).to("cpu") + left = torch.randn(2, 4, 1, 6, device="cpu") + right = torch.randn_like(left) + scalar_pair = torch.empty(2, 1, 4, device="cpu") + calls = 0 + + def fused_middle(projected_left, projected_right): + nonlocal calls + calls += 1 + return projected_left * projected_right + + out = module( + left, + right, + scalar_pair, + to_grid=lambda value: value, + from_grid=lambda value: value, + grid_product=fused_middle, + ) + + assert calls == 1 + assert out.shape == left.shape + + +@pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) +def test_dispatch_falls_back_without_exact_cuda_contract( + monkeypatch: pytest.MonkeyPatch, + hidden_channels: int, +): + from deepmd.kernels.cute.neo import ( + output_grid_product, + ) + + monkeypatch.setattr( + output_grid_product.runtime_policy, + "is_cute_infer_enabled", + lambda: True, + ) + left = torch.randn(2, COEFF_DIM, 1, N_FRAMES * hidden_channels, device="cpu") + right = torch.randn_like(left) + to_grid = torch.randn(GRID_SIZE, COEFF_DIM * N_FRAMES, device="cpu") + from_grid = torch.randn(COEFF_DIM * N_FRAMES, GRID_SIZE, device="cpu") + + assert ( + output_grid_product.maybe_run_cute_output_grid_product( + left, + right, + to_grid, + from_grid, + n_frames=N_FRAMES, + ) + is None + ) + + +@pytest.mark.parametrize( + ("hidden_channels", "expected"), + [(96, 96), (192, 192), (128, None)], +) +def test_exact_shape_guard_accepts_only_validated_widths( + hidden_channels: int, + expected: int | None, +) -> None: + from deepmd.kernels.cute.neo.output_grid_product import ( + _exact_hidden_channels, + ) + + left = torch.empty(2, COEFF_DIM, 1, N_FRAMES * hidden_channels, device="cpu") + assert _exact_hidden_channels(left, N_FRAMES) == expected + + +@pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) +def test_fake_registrations_allocate_canonical_strides(hidden_channels: int) -> None: + from deepmd.kernels.cute.neo import ( + output_grid_product, + ) + + width = N_FRAMES * hidden_channels + shape = (2, COEFF_DIM, 1, width) + canonical_stride = (COEFF_DIM * width, width, width, 1) + left = torch.empty_strided( + shape, + (canonical_stride[0], canonical_stride[1], 7, 1), + device="cpu", + ) + right = torch.empty_strided( + shape, + (canonical_stride[0], canonical_stride[1], 11, 1), + device="cpu", + ) + grad_out = torch.empty_strided( + shape, + (canonical_stride[0], canonical_stride[1], 13, 1), + device="cpu", + ) + to_grid = torch.empty(GRID_SIZE, COEFF_DIM * N_FRAMES, device="cpu") + from_grid = torch.empty(COEFF_DIM * N_FRAMES, GRID_SIZE, device="cpu") + + out = output_grid_product._output_grid_product_fake( + left, + right, + to_grid, + from_grid, + N_FRAMES, + ) + grad_left, grad_right = output_grid_product._output_grid_product_bwd_fake( + grad_out, + left, + right, + to_grid, + from_grid, + N_FRAMES, + ) + + assert out.stride() == canonical_stride + assert grad_left.stride() == canonical_stride + assert grad_right.stride() == canonical_stride + + +@pytest.mark.parametrize( + ("compute_capability", "hidden_channels", "policy_enabled", "expected"), + [ + ((9, 0), 96, True, True), + ((9, 0), 96, False, False), + ((9, 0), 192, True, False), + ((9, 1), 96, True, False), + ], +) +def test_sm90_c96_asymmetric_panel_dispatch_is_explicit( + monkeypatch: pytest.MonkeyPatch, + compute_capability: tuple[int, int], + hidden_channels: int, + policy_enabled: bool, + expected: bool, +) -> None: + from deepmd.kernels.cute.neo import ( + output_grid_product, + ) + + kernel_module_name = ( + "deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product" + ) + kernel_module = types.ModuleType(kernel_module_name) + calls: dict[str, dict[str, bool]] = {} + + def run_forward(left, right, to_grid, from_grid, **kwargs): + del right, to_grid, from_grid + calls["forward"] = kwargs + return torch.empty_like(left) + + def run_backward(grad_out, left, right, to_grid, from_grid, **kwargs): + del grad_out, to_grid, from_grid + calls["backward"] = kwargs + return torch.empty_like(left), torch.empty_like(right) + + kernel_module.run_tiled_output_grid_product = run_forward + kernel_module.run_tiled_output_grid_product_backward = run_backward + monkeypatch.setitem(sys.modules, kernel_module_name, kernel_module) + monkeypatch.setattr( + output_grid_product, + "_validate_exact_contract", + lambda *_args: hidden_channels, + ) + monkeypatch.setattr( + torch.cuda, + "get_device_capability", + lambda _device=None: compute_capability, + ) + monkeypatch.setattr( + output_grid_product.runtime_policy, + "is_output_grid_fwd_sm80_c96_n48_enabled", + lambda _compute_capability: False, + ) + monkeypatch.setattr( + output_grid_product.runtime_policy, + "is_output_grid_bwd_sm80_c96_n48_panel_enabled", + lambda _compute_capability: False, + ) + monkeypatch.setattr( + output_grid_product.runtime_policy, + "is_output_grid_sm90_c96_asymmetric_panels_enabled", + lambda _compute_capability: policy_enabled, + ) + + width = N_FRAMES * hidden_channels + left = torch.empty(2, COEFF_DIM, 1, width, device="cpu") + right = torch.empty_like(left) + grad_out = torch.empty_like(left) + to_grid = torch.empty( + GRID_SIZE, + COEFF_DIM * N_FRAMES, + device="cpu", + ) + from_grid = torch.empty( + COEFF_DIM * N_FRAMES, + GRID_SIZE, + device="cpu", + ) + + output_grid_product._output_grid_product_impl( + left, + right, + to_grid, + from_grid, + N_FRAMES, + ) + output_grid_product._output_grid_product_bwd_impl( + grad_out, + left, + right, + to_grid, + from_grid, + N_FRAMES, + ) + + assert calls["forward"] == { + "use_sm80_c96_n48": False, + "use_sm90_c96_asymmetric_panels": expected, + } + assert calls["backward"] == { + "use_sm80_c96_n48_panel": False, + "use_sm90_c96_asymmetric_panels": expected, + } + + +@pytest.mark.skipif( + _CUTE_SKIP_REASON is not None, + reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable", +) +class TestOutputGridProductCuda: + @staticmethod + def _inputs(nodes: int, hidden_channels: int): + generator = torch.Generator(device="cuda").manual_seed( + 20260703 + nodes + hidden_channels + ) + left = 0.1 * torch.randn( + nodes, + COEFF_DIM, + 1, + N_FRAMES * hidden_channels, + device="cuda", + generator=generator, + ) + right = 0.1 * torch.randn( + left.shape, + device="cuda", + generator=generator, + ) + to_grid = 0.1 * torch.randn( + GRID_SIZE, + COEFF_DIM * N_FRAMES, + device="cuda", + generator=generator, + ) + from_grid = 0.1 * torch.randn( + COEFF_DIM * N_FRAMES, + GRID_SIZE, + device="cuda", + generator=generator, + ) + return left, right, to_grid, from_grid + + @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) + @pytest.mark.parametrize("nodes", [1, 7, 65]) + def test_forward_and_first_backward_match_strict_fp32( + self, + nodes: int, + hidden_channels: int, + ): + from deepmd.kernels.cute.neo.output_grid_product import ( + output_grid_product_cute, + ) + + left, right, to_grid, from_grid = self._inputs(nodes, hidden_channels) + left_ref = left.detach().clone().requires_grad_(True) + right_ref = right.detach().clone().requires_grad_(True) + left_actual = left.detach().clone().requires_grad_(True) + right_actual = right.detach().clone().requires_grad_(True) + grad = torch.randn_like(left) + + expected = _reference(left_ref, right_ref, to_grid, from_grid) + expected_grads = torch.autograd.grad( + expected, + (left_ref, right_ref), + grad, + ) + actual = output_grid_product_cute( + left_actual, + right_actual, + to_grid, + from_grid, + n_frames=N_FRAMES, + ) + actual_grads = torch.autograd.grad( + actual, + (left_actual, right_actual), + grad, + ) + + torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL) + torch.testing.assert_close( + actual_grads[0], expected_grads[0], atol=TOL, rtol=TOL + ) + torch.testing.assert_close( + actual_grads[1], expected_grads[1], atol=TOL, rtol=TOL + ) + + @pytest.mark.skipif( + _SM80_SKIP_REASON is not None, + reason=_SM80_SKIP_REASON or "sm80-family GPU is unavailable", + ) + @pytest.mark.parametrize("nodes", [1, 7, 65]) + def test_sm80_c96_n48_forward_custom_op_matches_strict_fp32( + self, + monkeypatch: pytest.MonkeyPatch, + nodes: int, + ): + from deepmd.kernels.cute.neo import ( + output_grid_product, + ) + + monkeypatch.setattr( + output_grid_product.runtime_policy, + "is_output_grid_fwd_sm80_c96_n48_enabled", + lambda _compute_capability: True, + ) + left, right, to_grid, from_grid = self._inputs(nodes, 96) + expected = _reference(left, right, to_grid, from_grid) + actual = output_grid_product.output_grid_product_cute( + left, + right, + to_grid, + from_grid, + n_frames=N_FRAMES, + ) + torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL) + + @pytest.mark.skipif( + _SM80_SKIP_REASON is not None, + reason=_SM80_SKIP_REASON or "sm80-family GPU is unavailable", + ) + def test_sm80_c96_n48_panel_custom_op_matches_strict_fp32( + self, + monkeypatch: pytest.MonkeyPatch, + ): + from deepmd.kernels.cute.neo import ( + output_grid_product, + ) + + monkeypatch.setattr( + output_grid_product.runtime_policy, + "is_output_grid_bwd_sm80_c96_n48_panel_enabled", + lambda _compute_capability: True, + ) + left, right, to_grid, from_grid = self._inputs(7, 96) + left_ref = left.detach().clone().requires_grad_(True) + right_ref = right.detach().clone().requires_grad_(True) + left_actual = left.detach().clone().requires_grad_(True) + right_actual = right.detach().clone().requires_grad_(True) + grad = torch.randn_like(left) + + expected = _reference(left_ref, right_ref, to_grid, from_grid) + expected_grads = torch.autograd.grad( + expected, + (left_ref, right_ref), + grad, + ) + actual = output_grid_product.output_grid_product_cute( + left_actual, + right_actual, + to_grid, + from_grid, + n_frames=N_FRAMES, + ) + actual_grads = torch.autograd.grad( + actual, + (left_actual, right_actual), + grad, + ) + + torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL) + torch.testing.assert_close( + actual_grads[0], expected_grads[0], atol=TOL, rtol=TOL + ) + torch.testing.assert_close( + actual_grads[1], expected_grads[1], atol=TOL, rtol=TOL + ) + + @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) + def test_dispatch_uses_only_the_master_gate( + self, + monkeypatch: pytest.MonkeyPatch, + hidden_channels: int, + ): + from deepmd.kernels.cute.neo import ( + output_grid_product, + ) + + left, right, to_grid, from_grid = self._inputs(2, hidden_channels) + monkeypatch.setattr( + output_grid_product.runtime_policy, + "is_cute_infer_enabled", + lambda: False, + ) + assert ( + output_grid_product.maybe_run_cute_output_grid_product( + left, + right, + to_grid, + from_grid, + n_frames=N_FRAMES, + ) + is None + ) + + monkeypatch.setattr( + output_grid_product.runtime_policy, + "is_cute_infer_enabled", + lambda: True, + ) + actual = output_grid_product.maybe_run_cute_output_grid_product( + left, + right, + to_grid, + from_grid, + n_frames=N_FRAMES, + ) + assert actual is not None + torch.testing.assert_close( + actual, + _reference(left, right, to_grid, from_grid), + atol=TOL, + rtol=TOL, + ) + + @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) + def test_dispatch_declines_non_strict_matmul_state( + self, + monkeypatch: pytest.MonkeyPatch, + hidden_channels: int, + ): + from deepmd.kernels.cute.neo import ( + output_grid_product, + ) + + left, right, to_grid, from_grid = self._inputs(2, hidden_channels) + monkeypatch.setattr( + output_grid_product.runtime_policy, + "is_cute_infer_enabled", + lambda: True, + ) + monkeypatch.setattr( + output_grid_product.runtime_policy, + "uses_strict_fp32_matmul", + lambda: False, + ) + + assert ( + output_grid_product.maybe_run_cute_output_grid_product( + left, + right, + to_grid, + from_grid, + n_frames=N_FRAMES, + ) + is None + ) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_readout_l0.py b/source/tests/pt/model/test_descriptor_sezm_cute_readout_l0.py new file mode 100644 index 0000000000..93fd80108e --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_readout_l0.py @@ -0,0 +1,860 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Focused algebra, dispatch, custom-op, and fullgraph readout checks.""" + +from __future__ import ( + annotations, +) + +import importlib + +import pytest +import torch + +COEFF_DIM = 16 +N_FRAMES = 3 +PACKED_COEFF_DIM = COEFF_DIM * N_FRAMES +GRID_SIZE = 152 +HIDDEN_CHANNELS = 192 +PACKED_WIDTH = N_FRAMES * HIDDEN_CHANNELS +TOL = 5.0e-5 + + +def _cute_runtime_skip_reason() -> str | None: + if not torch.cuda.is_available(): + return "Neo readout l=0 differentials require CUDA" + if tuple(torch.cuda.get_device_capability()) not in {(8, 0), (9, 0)}: + return "Neo readout l=0 differentials require sm80 or sm90" + try: + importlib.import_module("cutlass.cute") + importlib.import_module("cuda.bindings.driver") + except Exception as exc: # pragma: no cover - runtime dependent + return f"Neo readout l=0 differentials require CuTe DSL: {exc}" + return None + + +_CUTE_SKIP_REASON = _cute_runtime_skip_reason() + + +def _cpu_inputs(nodes: int = 2, hidden_channels: int = HIDDEN_CHANNELS): + width = N_FRAMES * hidden_channels + left = torch.randn(nodes, COEFF_DIM, 1, width, device="cpu") + right = torch.randn_like(left) + to_grid = torch.randn(GRID_SIZE, PACKED_COEFF_DIM, device="cpu") + from_grid = torch.randn(PACKED_COEFF_DIM, GRID_SIZE, device="cpu") + return left, right, to_grid, from_grid + + +def _output_ffn( + *, + hidden_channels: int = 96, + trainable: bool = False, + device: str = "cpu", +): + from deepmd.pt.model.descriptor.sezm_nn.ffn import ( + EquivariantFFN, + ) + + return ( + EquivariantFFN( + lmax=3, + channels=32, + hidden_channels=hidden_channels, + kmax=1, + grid_mlp=True, + grid_branch=0, + dtype=torch.float32, + s2_activation=False, + ffn_so3_grid=True, + activation_function="silu", + glu_activation=True, + mlp_bias=False, + trainable=trainable, + seed=29, + ) + .to(device) + .eval() + ) + + +def _neo_descriptor(): + from deepmd.pt.model.descriptor.sezm import ( + DescrptSeZM, + ) + + return DescrptSeZM( + ntypes=2, + sel=4, + channels=32, + lmax=3, + mmax=1, + n_blocks=2, + so2_layers=3, + n_focus=2, + message_node_so3=True, + ffn_neurons=0, + ffn_so3_grid=True, + grid_branch=[0, 0, 1], + ffn_blocks=1, + so3_readout="mlp", + use_amp=False, + precision="float32", + trainable=False, + seed=42, + ).eval() + + +@pytest.fixture +def cpu_neo_descriptor(monkeypatch): + from deepmd.pt.model.network import ( + mlp, + ) + from deepmd.pt.utils import ( + env, + ) + from deepmd.pt.utils import utils as pt_utils + + cpu = torch.device("cpu") + monkeypatch.setattr(env, "DEVICE", cpu) + monkeypatch.setattr(pt_utils, "DEVICE", cpu) + monkeypatch.setattr(mlp, "device", cpu) + return _neo_descriptor().to("cpu") + + +def _capture_matmul_precision_state() -> tuple[str | None, str | None, str]: + matmul = torch.backends.cuda.matmul + try: + global_precision = torch.backends.fp32_precision + except AttributeError: + return None, None, torch.get_float32_matmul_precision() + + torch.backends.fp32_precision = "none" + backend_precision = matmul.fp32_precision + matmul.fp32_precision = "none" + try: + legacy_precision = torch.get_float32_matmul_precision() + finally: + torch.backends.fp32_precision = global_precision + matmul.fp32_precision = backend_precision + return backend_precision, global_precision, legacy_precision + + +def _restore_matmul_precision_state( + state: tuple[str | None, str | None, str], +) -> None: + backend_precision, global_precision, legacy_precision = state + matmul = torch.backends.cuda.matmul + if backend_precision is None: + torch.set_float32_matmul_precision(legacy_precision) + return + matmul.fp32_precision = "none" + if global_precision is not None: + torch.backends.fp32_precision = "none" + torch.set_float32_matmul_precision(legacy_precision) + if global_precision is not None: + torch.backends.fp32_precision = global_precision + matmul.fp32_precision = backend_precision + + +@pytest.fixture +def matmul_precision_state(): + state = _capture_matmul_precision_state() + try: + yield torch.backends.cuda.matmul + finally: + _restore_matmul_precision_state(state) + + +def _set_new_matmul_precision(matmul, precision: str) -> None: + try: + matmul.fp32_precision = precision + except AttributeError: + pytest.skip("new CUDA matmul precision API is unavailable") + + +def _reference_product(left, right, to_grid, from_grid): + nodes = left.shape[0] + left_flat = left.reshape(nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS) + right_flat = right.reshape(nodes, PACKED_COEFF_DIM, HIDDEN_CHANNELS) + left_grid = torch.einsum("gj,njh->ngh", to_grid, left_flat) + right_grid = torch.einsum("gj,njh->ngh", to_grid, right_flat) + return torch.einsum( + "g,ngh->nh", + from_grid[0], + left_grid * right_grid, + ) + + +def test_sm80_readout_input_fold_selector_uses_only_full_neo_gate(monkeypatch): + from deepmd.kernels.cute.neo import ( + runtime_policy, + ) + + monkeypatch.delenv("DP_NEO_CUTE_INFER", raising=False) + monkeypatch.setenv("DP_CUTE_INFER", "1") + monkeypatch.delenv("DP_CUTE_READOUT_INPUT_FOLD_SM80", raising=False) + assert not runtime_policy.is_cute_infer_enabled() + assert not runtime_policy.is_readout_input_fold_sm80_enabled((8, 0)) + + monkeypatch.setenv("DP_CUTE_INFER", "0") + monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + assert runtime_policy.is_cute_infer_enabled() + assert runtime_policy.is_readout_input_fold_sm80_enabled((8, 0)) + monkeypatch.setenv("DP_CUTE_READOUT_INPUT_FOLD_SM80", "0") + assert not runtime_policy.is_readout_input_fold_sm80_enabled((8, 0)) + monkeypatch.setenv("DP_CUTE_READOUT_INPUT_FOLD_SM80", "1") + assert runtime_policy.is_readout_input_fold_sm80_enabled((8, 0)) + assert not runtime_policy.is_readout_input_fold_sm80_enabled((9, 0)) + + +def test_sm90_readout_input_fold_selector_uses_only_full_neo_gate(monkeypatch): + from deepmd.kernels.cute.neo import ( + runtime_policy, + ) + + monkeypatch.delenv("DP_NEO_CUTE_INFER", raising=False) + monkeypatch.setenv("DP_CUTE_INFER", "1") + monkeypatch.delenv("DP_CUTE_READOUT_INPUT_FOLD_SM90", raising=False) + assert not runtime_policy.is_readout_input_fold_sm90_enabled((9, 0)) + + monkeypatch.setenv("DP_CUTE_INFER", "0") + monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + assert runtime_policy.is_readout_input_fold_sm90_enabled((9, 0)) + assert runtime_policy.is_readout_input_fold_enabled((9, 0)) + monkeypatch.setenv("DP_CUTE_READOUT_INPUT_FOLD_SM90", "0") + assert not runtime_policy.is_readout_input_fold_sm90_enabled((9, 0)) + monkeypatch.setenv("DP_CUTE_READOUT_INPUT_FOLD_SM90", "1") + assert runtime_policy.is_readout_input_fold_sm90_enabled((9, 0)) + assert not runtime_policy.is_readout_input_fold_sm90_enabled((8, 0)) + + +def test_readout_input_fold_guard_accepts_validated_architectures(monkeypatch): + from deepmd.kernels.cute.neo import ( + readout_l0, + runtime_policy, + ) + + module = _output_ffn() + value = torch.randn(2, COEFF_DIM, 1, 32, device="cpu") + capabilities = [] + + def select_readout(capability): + capabilities.append(capability) + return capability in {(8, 0), (9, 0)} + + monkeypatch.setattr( + runtime_policy, + "is_readout_input_fold_enabled", + select_readout, + ) + monkeypatch.setattr(readout_l0, "_has_exact_neo_readout_contract", lambda *_: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_: (8, 0)) + assert readout_l0._can_use_sm80_readout_input_fold(module, value) + + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *_: (9, 0)) + assert readout_l0._can_use_sm80_readout_input_fold(module, value) + assert capabilities == [(8, 0), (9, 0)] + + +def test_sm80_readout_input_fold_matches_forward_and_input_vjp(monkeypatch): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + module = _output_ffn() + with torch.no_grad(): + module.so3_linear_2.weight.normal_(std=0.03) + value = torch.randn(3, COEFF_DIM, 1, 32, device="cpu") + cotangent = torch.randn(3, 32, device="cpu") + + monkeypatch.setattr( + readout_l0, + "_can_use_sm80_readout_input_fold", + lambda *_: False, + ) + value_ref = value.detach().clone().requires_grad_(True) + expected = readout_l0._run_neo_readout_l0( + module, + value_ref, + _reference_product, + ) + expected_grad = torch.autograd.grad(expected, value_ref, cotangent)[0] + + monkeypatch.setattr( + readout_l0, + "_can_use_sm80_readout_input_fold", + lambda *_: True, + ) + + def forbid_staged_projection(*args, **kwargs): + del args, kwargs + pytest.fail("folded readout must not execute a staged input projection") + + monkeypatch.setattr(module.so3_linear_1, "forward", forbid_staged_projection) + monkeypatch.setattr( + module.act.grid_op.left_proj, + "forward", + forbid_staged_projection, + ) + monkeypatch.setattr( + module.act.grid_op.right_proj, + "forward", + forbid_staged_projection, + ) + monkeypatch.setattr(module.act.scalar_gate, "forward", forbid_staged_projection) + + def compact_reference(left, right, to_grid, from_grid): + assert left.is_contiguous() + assert right.is_contiguous() + return _reference_product(left, right, to_grid, from_grid) + + value_actual = value.detach().clone().requires_grad_(True) + actual = readout_l0._run_neo_readout_l0( + module, + value_actual, + compact_reference, + ) + actual_grad = torch.autograd.grad(actual, value_actual, cotangent)[0] + + torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL) + torch.testing.assert_close(actual_grad, expected_grad, atol=TOL, rtol=TOL) + + +def test_sm80_readout_input_fold_is_fullgraph_safe(monkeypatch): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + module = _output_ffn() + with torch.no_grad(): + module.so3_linear_2.weight.normal_(std=0.03) + monkeypatch.setattr( + readout_l0, + "_can_use_sm80_readout_input_fold", + lambda *_: True, + ) + + def folded(value): + return readout_l0._run_neo_readout_l0( + module, + value, + _reference_product, + ) + + value = torch.randn(2, COEFF_DIM, 1, 32, device="cpu") + readout_l0.prepare_sm80_readout_input_fold(module) + expected = folded(value) + compiled = torch.compile(folded, backend="eager", fullgraph=True) + actual = compiled(value) + torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL) + + +def test_sm80_readout_input_fold_cache_refreshes_after_inplace_change(monkeypatch): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + module = _output_ffn() + ready_calls = [] + monkeypatch.setattr( + readout_l0, + "_synchronize_sm80_readout_input_fold_build", + lambda weights: ready_calls.append(weights), + ) + first = readout_l0.prepare_sm80_readout_input_fold(module) + first_cache = getattr(module, readout_l0._READOUT_INPUT_FOLD_CACHE) + assert getattr(module, readout_l0._READOUT_INPUT_FOLD_LEFT) is first[0] + assert not any("readout_input_fold" in key for key in module.state_dict()) + + cached = readout_l0.prepare_sm80_readout_input_fold(module) + assert all( + actual is expected for actual, expected in zip(cached, first, strict=True) + ) + assert len(ready_calls) == 1 + with torch.no_grad(): + module.act.grid_op.left_proj.weight.add_(0.01) + second = readout_l0.prepare_sm80_readout_input_fold(module) + second_cache = getattr(module, readout_l0._READOUT_INPUT_FOLD_CACHE) + + assert second_cache is not first_cache + assert second[0] is not first[0] + assert not torch.equal(second[0], first[0]) + assert len(ready_calls) == 2 + + +def test_sm80_readout_input_fold_cache_refreshes_after_parameter_replacement(): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + module = _output_ffn() + first = readout_l0.prepare_sm80_readout_input_fold(module) + old_weight = module.act.grid_op.left_proj.weight + module.act.grid_op.left_proj.weight = torch.nn.Parameter( + old_weight.detach().clone().add_(0.02), + requires_grad=False, + ) + second = readout_l0.prepare_sm80_readout_input_fold(module) + + assert second[0] is not first[0] + assert not torch.equal(second[0], first[0]) + + +@pytest.mark.parametrize("assign", [False, True]) +def test_sm80_readout_input_fold_load_state_dict_invalidates_cache(assign): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + module = _output_ffn() + first = readout_l0.prepare_sm80_readout_input_fold(module) + state = {key: value.clone() for key, value in module.state_dict().items()} + left_key = next(key for key in state if key.endswith("grid_op.left_proj.weight")) + state[left_key].add_(0.03) + + module.load_state_dict(state, assign=assign) + + assert getattr(module, readout_l0._READOUT_INPUT_FOLD_CACHE) is None + assert all( + getattr(module, name) is None for name in readout_l0._READOUT_INPUT_FOLD_BUFFERS + ) + second = readout_l0.prepare_sm80_readout_input_fold(module) + assert second[0] is not first[0] + assert not torch.equal(second[0], first[0]) + assert not any("readout_input_fold" in key for key in module.state_dict()) + + +def test_sm80_readout_input_fold_buffers_follow_device_moves(): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + module = _output_ffn() + readout_l0.prepare_sm80_readout_input_fold(module) + + module.to("meta") + + assert all( + getattr(module, name).device.type == "meta" + for name in readout_l0._READOUT_INPUT_FOLD_BUFFERS + ) + refreshed = readout_l0.prepare_sm80_readout_input_fold(module) + assert all(weight.device.type == "meta" for weight in refreshed) + + +def test_sm80_readout_input_fold_compile_requires_explicit_preparation( + monkeypatch, +): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + module = _output_ffn() + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + with pytest.raises(RuntimeError, match="prepare_sm80_readout_input_fold"): + readout_l0._get_sm80_readout_input_fold(module) + + +def test_sm80_readout_input_fold_compiled_lookup_returns_prepared_buffers( + monkeypatch, +): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + module = _output_ffn() + expected = readout_l0.prepare_sm80_readout_input_fold(module) + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + cached = readout_l0._get_sm80_readout_input_fold(module) + assert all( + actual is reference for actual, reference in zip(cached, expected, strict=True) + ) + + +@pytest.mark.parametrize( + ("hidden_channels", "expected"), + [(HIDDEN_CHANNELS, True), (96, False), (128, False)], +) +def test_shape_guard_accepts_only_exact_c192_readout( + hidden_channels: int, + expected: bool, +): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + left, _, _, _ = _cpu_inputs(hidden_channels=hidden_channels) + assert readout_l0._has_exact_product_shape(left) is expected + + +def test_fake_registrations_return_canonical_strides(): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + shape = (2, COEFF_DIM, 1, PACKED_WIDTH) + canonical = (COEFF_DIM * PACKED_WIDTH, PACKED_WIDTH, PACKED_WIDTH, 1) + left = torch.empty_strided(shape, (canonical[0], canonical[1], 7, 1), device="cpu") + right = torch.empty_strided( + shape, (canonical[0], canonical[1], 11, 1), device="cpu" + ) + dq0 = torch.empty(2, HIDDEN_CHANNELS, device="cpu") + to_grid = torch.empty(GRID_SIZE, PACKED_COEFF_DIM, device="cpu") + from_grid = torch.empty(PACKED_COEFF_DIM, GRID_SIZE, device="cpu") + + q0 = readout_l0._readout_l0_fake(left, right, to_grid, from_grid) + grad_left, grad_right = readout_l0._readout_l0_bwd_fake( + dq0, left, right, to_grid, from_grid + ) + + assert q0.shape == (2, HIDDEN_CHANNELS) + assert q0.stride() == (HIDDEN_CHANNELS, 1) + assert grad_left.stride() == canonical + assert grad_right.stride() == canonical + + +def test_custom_ops_use_canonical_fake_metadata(): + from torch._subclasses.fake_tensor import ( + FakeTensorMode, + ) + + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + shape = (2, COEFF_DIM, 1, PACKED_WIDTH) + canonical = (COEFF_DIM * PACKED_WIDTH, PACKED_WIDTH, PACKED_WIDTH, 1) + with FakeTensorMode(): + left = torch.empty_strided( + shape, (canonical[0], canonical[1], 7, 1), device="cuda" + ) + right = torch.empty_strided( + shape, (canonical[0], canonical[1], 11, 1), device="cuda" + ) + dq0 = torch.empty(2, HIDDEN_CHANNELS, device="cuda") + to_grid = torch.empty(GRID_SIZE, PACKED_COEFF_DIM, device="cuda") + from_grid = torch.empty(PACKED_COEFF_DIM, GRID_SIZE, device="cuda") + q0 = readout_l0._readout_l0_op(left, right, to_grid, from_grid) + grad_left, grad_right = readout_l0._readout_l0_bwd_op( + dq0, left, right, to_grid, from_grid + ) + + assert q0.stride() == (HIDDEN_CHANNELS, 1) + assert grad_left.stride() == canonical + assert grad_right.stride() == canonical + + +def test_reference_completion_matches_module_output_and_input_vjp(): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + module = _output_ffn() + with torch.no_grad(): + module.so3_linear_2.weight.normal_(std=0.05) + value = 0.1 * torch.randn(3, COEFF_DIM, 1, 32, device="cpu") + seed = torch.randn(3, 32, device="cpu") + value_ref = value.detach().clone().requires_grad_(True) + value_actual = value.detach().clone().requires_grad_(True) + + expected = (value_ref + module(value_ref))[:, 0, 0, :] + expected_grad = torch.autograd.grad(expected, value_ref, seed)[0] + actual = readout_l0._run_neo_readout_l0(module, value_actual, _reference_product) + actual_grad = torch.autograd.grad(actual, value_actual, seed)[0] + + torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL) + torch.testing.assert_close(actual_grad, expected_grad, atol=TOL, rtol=TOL) + + +def test_exact_structure_and_frozen_inference_guards(): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + module = _output_ffn() + assert readout_l0._has_exact_neo_readout_structure(module) + assert readout_l0._inference_mode_is_frozen(module) + assert not readout_l0._has_exact_neo_readout_structure( + _output_ffn(hidden_channels=64) + ) + module.train() + assert not readout_l0._inference_mode_is_frozen(module) + module.eval() + next(module.parameters()).requires_grad_(True) + assert not readout_l0._inference_mode_is_frozen(module) + + +def test_module_boundary_falls_back_when_device_is_unsupported(monkeypatch): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + module = _output_ffn() + value = torch.randn(2, COEFF_DIM, 1, 32, device="cpu") + expected = (value + module(value))[:, 0, 0, :] + monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + + assert readout_l0.maybe_run_neo_readout_l0(module, value) is None + torch.testing.assert_close( + readout_l0.run_neo_output_readout(module, value), expected + ) + + +def test_descriptor_freeze_guard_and_readout_wiring(cpu_neo_descriptor, monkeypatch): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + descriptor = cpu_neo_descriptor + value = torch.randn(2, COEFF_DIM, 1, 32, device="cpu") + expected = torch.randn(2, 32, device="cpu") + calls = [] + + def record_readout(output_ffn, ffn_in, *, parameters_frozen): + calls.append((output_ffn, ffn_in, parameters_frozen)) + return expected + + monkeypatch.setattr(readout_l0, "run_neo_output_readout", record_readout) + assert descriptor._readout_parameters_are_frozen() + assert descriptor._run_output_readout(value) is expected + assert calls[-1] == (descriptor.output_ffn, value, True) + + next(descriptor.parameters()).requires_grad_(True) + descriptor._run_output_readout(value) + assert calls[-1] == (descriptor.output_ffn, value, False) + + +def test_trainable_descriptor_state_bypasses_candidate(monkeypatch): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + module = _output_ffn() + value = torch.randn(2, COEFF_DIM, 1, 32, device="cpu") + expected = (value + module(value))[:, 0, 0, :] + monkeypatch.setattr( + readout_l0, + "maybe_run_neo_readout_l0", + lambda *args, **kwargs: pytest.fail("candidate must be bypassed"), + ) + + actual = readout_l0.run_neo_output_readout(module, value, parameters_frozen=False) + torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize( + ("precision", "expected"), + [("highest", True), ("high", False)], +) +def test_legacy_matmul_precision_guard_is_eager_and_fullgraph_safe( + matmul_precision_state, + precision: str, + expected: bool, +): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + try: + matmul_precision_state.fp32_precision = "none" + except AttributeError: + pass + torch.set_float32_matmul_precision(precision) + value = torch.ones(1, device="cpu") + + def guarded(tensor): + return tensor if readout_l0._uses_strict_fp32_matmul() else -tensor + + eager = guarded(value) + fullgraph = torch.compile(guarded, backend="eager", fullgraph=True)(value) + assert readout_l0._uses_strict_fp32_matmul() is expected + assert torch.equal(fullgraph, eager) + + +@pytest.mark.parametrize( + ("precision", "expected"), + [("ieee", True), ("tf32", False)], +) +def test_modern_matmul_precision_guard_rejects_tf32( + matmul_precision_state, + precision: str, + expected: bool, +): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + _set_new_matmul_precision(matmul_precision_state, precision) + assert readout_l0._uses_strict_fp32_matmul() is expected + + +@pytest.mark.skipif( + _CUTE_SKIP_REASON is not None, + reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable", +) +class TestReadoutL0Cuda: + @staticmethod + def _inputs(nodes: int): + generator = torch.Generator(device="cuda").manual_seed(20260704 + nodes) + left = 0.1 * torch.randn( + nodes, + COEFF_DIM, + 1, + PACKED_WIDTH, + device="cuda", + generator=generator, + ) + right = 0.1 * torch.randn(left.shape, device="cuda", generator=generator) + to_grid = 0.1 * torch.randn( + GRID_SIZE, + PACKED_COEFF_DIM, + device="cuda", + generator=generator, + ) + from_grid = 0.1 * torch.randn( + PACKED_COEFF_DIM, + GRID_SIZE, + device="cuda", + generator=generator, + ) + return left, right, to_grid, from_grid + + @pytest.mark.parametrize("nodes", [1, 7, 65]) + def test_forward_and_input_vjp_match_strict_fp32( + self, + nodes: int, + ): + from deepmd.kernels.cute.neo.readout_l0 import ( + readout_l0_product_cute, + ) + + left, right, to_grid, from_grid = self._inputs(nodes) + left_ref = left.detach().clone().requires_grad_(True) + right_ref = right.detach().clone().requires_grad_(True) + left_actual = left.detach().clone().requires_grad_(True) + right_actual = right.detach().clone().requires_grad_(True) + dq0 = torch.randn(nodes, HIDDEN_CHANNELS, device="cuda") + + expected = _reference_product(left_ref, right_ref, to_grid, from_grid) + expected_grads = torch.autograd.grad(expected, (left_ref, right_ref), dq0) + actual = readout_l0_product_cute(left_actual, right_actual, to_grid, from_grid) + actual_grads = torch.autograd.grad(actual, (left_actual, right_actual), dq0) + + torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL) + torch.testing.assert_close( + actual_grads[0], expected_grads[0], atol=TOL, rtol=TOL + ) + torch.testing.assert_close( + actual_grads[1], expected_grads[1], atol=TOL, rtol=TOL + ) + + def test_opcheck_and_fullgraph_preserve_metadata_and_vjp(self): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + left, right, to_grid, from_grid = self._inputs(3) + dq0 = torch.randn(3, HIDDEN_CHANNELS, device="cuda") + torch.library.opcheck( + readout_l0._readout_l0_op, + (left, right, to_grid, from_grid), + test_utils=("test_schema", "test_faketensor"), + ) + compiled = torch.compile( + readout_l0.readout_l0_product_cute, + dynamic=True, + fullgraph=True, + ) + left.requires_grad_(True) + right.requires_grad_(True) + actual = compiled(left, right, to_grid, from_grid) + grads = torch.autograd.grad(actual, (left, right), dq0) + assert actual.shape == (3, HIDDEN_CHANNELS) + assert grads[0].shape == left.shape + assert grads[1].shape == right.shape + + def test_exact_module_fullgraph_uses_full_neo_gate(self, monkeypatch): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + module = _output_ffn(device="cuda") + with torch.no_grad(): + module.so3_linear_2.weight.normal_(std=0.05) + value = torch.randn(3, COEFF_DIM, 1, 32, device="cuda") + seed = torch.randn(3, 32, device="cuda") + value_ref = value.detach().clone().requires_grad_(True) + value_actual = value.detach().clone().requires_grad_(True) + monkeypatch.delenv("DP_NEO_CUTE_INFER", raising=False) + expected = (value_ref + module(value_ref))[:, 0, 0, :] + expected_grad = torch.autograd.grad(expected, value_ref, seed)[0] + + monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + readout_l0.prepare_sm80_readout_input_fold(module) + compiled = torch.compile( + lambda tensor: readout_l0.run_neo_output_readout(module, tensor), + dynamic=True, + fullgraph=True, + ) + actual = compiled(value_actual) + actual_grad = torch.autograd.grad(actual, value_actual, seed)[0] + + torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL) + torch.testing.assert_close(actual_grad, expected_grad, atol=TOL, rtol=TOL) + + def test_sm80_input_fold_preparation_is_cross_stream_ready(self): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + if tuple(torch.cuda.get_device_capability()) not in {(8, 0), (8, 6)}: + pytest.skip("readout input folding requires SM80 or SM86") + + module = _output_ffn(device="cuda") + expected = readout_l0._build_sm80_readout_input_fold(module) + torch.cuda.synchronize() + + producer = torch.cuda.Stream() + with torch.cuda.stream(producer): + prepared = readout_l0.prepare_sm80_readout_input_fold(module) + assert producer.query() + + consumer = torch.cuda.Stream() + with torch.cuda.stream(consumer): + observed = tuple(weight.clone() for weight in prepared) + consumer.synchronize() + for actual, reference in zip(observed, expected, strict=True): + torch.testing.assert_close(actual, reference, atol=0.0, rtol=0.0) + + def test_sm80_input_fold_fullgraph_matches_module_vjp(self, monkeypatch): + from deepmd.kernels.cute.neo import ( + readout_l0, + ) + + if tuple(torch.cuda.get_device_capability()) not in {(8, 0), (8, 6)}: + pytest.skip("readout input folding requires SM80 or SM86") + + module = _output_ffn(device="cuda") + with torch.no_grad(): + module.so3_linear_2.weight.normal_(std=0.05) + value = torch.randn(7, COEFF_DIM, 1, 32, device="cuda") + seed = torch.randn(7, 32, device="cuda") + value_ref = value.detach().clone().requires_grad_(True) + value_actual = value.detach().clone().requires_grad_(True) + expected = (value_ref + module(value_ref))[:, 0, 0, :] + expected_grad = torch.autograd.grad(expected, value_ref, seed)[0] + + monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setenv("DP_CUTE_READOUT_INPUT_FOLD_SM80", "1") + assert readout_l0._can_use_sm80_readout_input_fold(module, value_actual) + readout_l0.prepare_sm80_readout_input_fold(module) + compiled = torch.compile( + lambda tensor: readout_l0.run_neo_output_readout(module, tensor), + dynamic=True, + fullgraph=True, + ) + actual = compiled(value_actual) + actual_grad = torch.autograd.grad(actual, value_actual, seed)[0] + + torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL) + torch.testing.assert_close(actual_grad, expected_grad, atol=TOL, rtol=TOL) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_runtime_policy.py b/source/tests/pt/model/test_descriptor_sezm_cute_runtime_policy.py new file mode 100644 index 0000000000..89a7be4c92 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_runtime_policy.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Behavioral tests for CuTe inference feature policy.""" + +from __future__ import ( + annotations, +) + +import importlib.util +import os +import sys +from pathlib import ( + Path, +) +from unittest import ( + mock, +) + +REPO_ROOT = Path(__file__).resolve().parents[4] +POLICY_PATH = REPO_ROOT / "deepmd/kernels/cute/neo/runtime_policy.py" + + +def _load_policy(): + assert POLICY_PATH.is_file(), f"CuTe runtime policy is missing: {POLICY_PATH}" + name = "sezm_cute_runtime_policy_test" + spec = importlib.util.spec_from_file_location(name, POLICY_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + try: + spec.loader.exec_module(module) + finally: + sys.modules.pop(name, None) + return module + + +def test_full_neo_master_is_independent_from_inner_cute_selector() -> None: + policy = _load_policy() + cases = ( + ({}, False), + ({"DP_CUTE_INFER": "1"}, False), + ({"DP_NEO_CUTE_INFER": "1"}, True), + ({"DP_CUTE_INFER": "0", "DP_NEO_CUTE_INFER": "1"}, True), + ({"DP_CUTE_INFER": "1", "DP_NEO_CUTE_INFER": "0"}, False), + ({"DP_NEO_CUTE_INFER": "1", "DP_TRITON_INFER": "2"}, True), + ) + for environment, expected in cases: + with mock.patch.dict(os.environ, environment, clear=True): + assert policy.is_cute_infer_enabled() is expected + + +def test_master_switch_controls_every_sm80_subfeature() -> None: + policy = _load_policy() + with mock.patch.dict( + os.environ, + { + "DP_CUTE_GIE": "1", + "DP_CUTE_K1_PACKED_WIGNER": "1", + "DP_CUTE_K1_THIN_WRAPPER": "1", + "DP_CUTE_OUTPUT_GRID_BWD_SM80_C96_N48_PANEL": "1", + "DP_CUTE_OUTPUT_GRID_FWD_SM80_C96_N48": "1", + "DP_CUTE_READOUT_INPUT_FOLD_SM80": "1", + }, + clear=True, + ): + assert not policy.is_cute_infer_enabled() + for capability in policy.SM80_PROFILE_CAPABILITIES: + assert not policy.is_sm80_profile_enabled(capability) + assert not policy.is_gie_enabled(capability) + assert not policy.is_packed_wigner_enabled(capability) + assert not policy.is_k1_thin_wrapper_enabled(capability) + assert not policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled(capability) + assert not policy.is_output_grid_fwd_sm80_c96_n48_enabled(capability) + assert not policy.is_readout_input_fold_sm80_enabled(capability) + + +def test_sm80_family_profile_defaults_from_master_only() -> None: + policy = _load_policy() + with mock.patch.dict(os.environ, {"DP_NEO_CUTE_INFER": "1"}, clear=True): + assert policy.SM80_PROFILE_CAPABILITIES == frozenset({(8, 0), (8, 6)}) + for capability in policy.SM80_PROFILE_CAPABILITIES: + assert policy.is_sm80_profile_enabled(capability) + assert policy.is_gie_enabled(capability) + assert policy.is_packed_wigner_enabled(capability) + assert policy.is_k1_thin_wrapper_enabled(capability) + assert policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled(capability) + assert policy.is_output_grid_fwd_sm80_c96_n48_enabled(capability) + assert policy.is_readout_input_fold_sm80_enabled(capability) + + +def test_every_sm80_profile_feature_can_be_disabled() -> None: + policy = _load_policy() + checks = { + "DP_CUTE_GIE": lambda: policy.is_gie_enabled((8, 0)), + "DP_CUTE_K1_PACKED_WIGNER": lambda: policy.is_packed_wigner_enabled((8, 0)), + "DP_CUTE_K1_THIN_WRAPPER": lambda: policy.is_k1_thin_wrapper_enabled((8, 0)), + "DP_CUTE_OUTPUT_GRID_BWD_SM80_C96_N48_PANEL": lambda: ( + policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled((8, 0)) + ), + "DP_CUTE_OUTPUT_GRID_FWD_SM80_C96_N48": lambda: ( + policy.is_output_grid_fwd_sm80_c96_n48_enabled((8, 0)) + ), + "DP_CUTE_READOUT_INPUT_FOLD_SM80": lambda: ( + policy.is_readout_input_fold_sm80_enabled((8, 0)) + ), + } + for name, checker in checks.items(): + with mock.patch.dict( + os.environ, + {"DP_NEO_CUTE_INFER": "1", name: "0"}, + clear=True, + ): + assert not checker() + + +def test_non_sm80_family_uses_architecture_defaults_without_sm80_features() -> None: + policy = _load_policy() + with mock.patch.dict(os.environ, {"DP_NEO_CUTE_INFER": "1"}, clear=True): + capability = (8, 9) + assert not policy.is_sm80_profile_enabled(capability) + assert not policy.is_gie_enabled(capability) + assert policy.is_packed_wigner_enabled(capability) + assert not policy.is_k1_thin_wrapper_enabled(capability) + assert not policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled(capability) + assert not policy.is_output_grid_fwd_sm80_c96_n48_enabled(capability) + assert not policy.is_readout_input_fold_sm80_enabled(capability) + + +def test_sm90_c96_asymmetric_panels_default_is_exact_arch_disable_only() -> None: + policy = _load_policy() + switch = policy.OUTPUT_GRID_SM90_C96_ASYMMETRIC_PANELS_ENV + + with mock.patch.dict(os.environ, {"DP_NEO_CUTE_INFER": "1"}, clear=True): + assert policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((9, 0)) + assert not policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((8, 9)) + assert not policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((9, 1)) + + with mock.patch.dict( + os.environ, + {"DP_NEO_CUTE_INFER": "1", switch: "0"}, + clear=True, + ): + assert not policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((9, 0)) + + with mock.patch.dict(os.environ, {switch: "1"}, clear=True): + assert not policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((9, 0)) + + with mock.patch.dict( + os.environ, + {"DP_NEO_CUTE_INFER": "1", switch: "1"}, + clear=True, + ): + assert policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((9, 0)) + assert not policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((8, 0)) + + +def test_overrides_remain_architecture_safe() -> None: + policy = _load_policy() + with mock.patch.dict( + os.environ, + { + "DP_NEO_CUTE_INFER": "1", + "DP_CUTE_K1_THIN_WRAPPER": "1", + "DP_CUTE_OUTPUT_GRID_BWD_SM80_C96_N48_PANEL": "1", + "DP_CUTE_OUTPUT_GRID_FWD_SM80_C96_N48": "1", + "DP_CUTE_READOUT_INPUT_FOLD_SM80": "1", + }, + clear=True, + ): + assert policy.is_k1_thin_wrapper_enabled((9, 0)) + assert not policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled((9, 0)) + assert not policy.is_output_grid_fwd_sm80_c96_n48_enabled((9, 0)) + assert not policy.is_readout_input_fold_sm80_enabled((9, 0)) + + with mock.patch.dict( + os.environ, + { + "DP_NEO_CUTE_INFER": "1", + "DP_CUTE_GIE": "1", + "DP_CUTE_K1_PACKED_WIGNER": "1", + }, + clear=True, + ): + assert not policy.is_gie_enabled((9, 0)) + assert not policy.is_packed_wigner_enabled((10, 1)) + + +def test_int32_k1_capacity_checks_every_flattened_axis() -> None: + policy = _load_policy() + max_edges = policy.INT32_MAX // policy.K1_VALUES_PER_EDGE + max_nodes = policy.INT32_MAX // policy.K1_VALUES_PER_NODE + + assert policy.k1_int32_indexing_is_safe( + edge_count=max_edges, + node_count=max_nodes, + ) + assert not policy.k1_int32_indexing_is_safe( + edge_count=max_edges + 1, + node_count=1, + ) + assert not policy.k1_int32_indexing_is_safe( + edge_count=1, + node_count=max_nodes + 1, + ) + assert not policy.k1_int32_indexing_is_safe( + edge_count=-1, + node_count=1, + ) + assert not policy.k1_int32_indexing_is_safe( + edge_count=1, + node_count=-1, + ) + + +def test_strict_mode_is_explicitly_opt_in() -> None: + policy = _load_policy() + with mock.patch.dict(os.environ, {}, clear=True): + assert not policy.is_cute_strict_enabled() + with mock.patch.dict(os.environ, {"DP_CUTE_STRICT": "1"}, clear=True): + assert policy.is_cute_strict_enabled() + + +def test_neighbor_list_eager_island_policy_is_owned_by_neo_runtime() -> None: + policy = _load_policy() + cases = ( + ({}, (8, 0), False), + ({"DP_NEO_CUTE_INFER": "1"}, (8, 0), True), + ({"DP_NEO_CUTE_INFER": "1"}, (8, 6), False), + ({"DP_NEO_CUTE_INFER": "1"}, (9, 0), False), + ( + { + "DP_NEO_CUTE_INFER": "1", + "DP_CUTE_K1_EAGER_ISLANDS": "0", + }, + (8, 0), + False, + ), + ( + { + "DP_NEO_CUTE_INFER": "1", + "DP_CUTE_K1_EAGER_ISLANDS": "1", + }, + (9, 0), + True, + ), + ) + for environment, capability, expected in cases: + with mock.patch.dict(os.environ, environment, clear=True): + assert policy.is_k1_eager_island_enabled(capability) is expected diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_tiled_grid_product.py b/source/tests/pt/model/test_descriptor_sezm_cute_tiled_grid_product.py new file mode 100644 index 0000000000..f8ef09c246 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_tiled_grid_product.py @@ -0,0 +1,387 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Strict-FP32 differentials for the tiled Neo output-grid product.""" + +from __future__ import ( + annotations, +) + +import ast +import importlib +from pathlib import ( + Path, +) + +import pytest +import torch + +TOL = 5.0e-5 +PACKED_COEFF_DIM = 48 +GRID_SIZE = 152 +SUPPORTED_HIDDEN_CHANNELS = (96, 192) +REPO_ROOT = Path(__file__).resolve().parents[4] +TILED_KERNEL_PATH = ( + REPO_ROOT + / "deepmd/kernels/cute/neo/output_grid_kernels" + / "cute_tiled_grid_product.py" +) +MESSAGE_GRID_PATH = ( + REPO_ROOT + / "deepmd/kernels/cute/neo/k1_kernels" + / "cute_neo_message_grid_product.py" +) + + +def _cute_runtime_skip_reason() -> str | None: + if not torch.cuda.is_available(): + return "tiled output-grid differentials require CUDA" + if torch.cuda.get_device_capability()[0] < 8: + return "tiled output-grid differentials require compute capability 8.0+" + try: + importlib.import_module("cutlass.cute") + importlib.import_module("cuda.bindings.driver") + except Exception as exc: # pragma: no cover - runtime dependent + return f"tiled output-grid differentials require CuTe DSL: {exc}" + return None + + +_CUTE_SKIP_REASON = _cute_runtime_skip_reason() + + +def _sm80_skip_reason() -> str | None: + if _CUTE_SKIP_REASON is not None: + return _CUTE_SKIP_REASON + if tuple(torch.cuda.get_device_capability()) not in {(8, 0), (8, 6)}: + return "specialized output-grid differentials require sm80 or sm86" + return None + + +_SM80_SKIP_REASON = _sm80_skip_reason() + + +def _method_node(tree: ast.AST, name: str) -> ast.FunctionDef: + matches = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == name + ] + assert len(matches) == 1 + return matches[0] + + +def test_sm80_c96_n48_panel_tiles_shared_input_before_partition_b() -> None: + """Guard the CuTe B-fragment layout contract without requiring CUDA.""" + tree = ast.parse(TILED_KERNEL_PATH.read_text(encoding="utf-8")) + method = _method_node(tree, "_backproject_panel_accumulate") + + assignments = { + node.targets[0].id: node.value + for node in ast.walk(method) + if isinstance(node, ast.Assign) + and len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + } + shared_b = assignments["sB"] + assert isinstance(shared_b, ast.Call) + assert ast.unparse(shared_b.func) == "cute.local_tile" + assert ast.unparse(shared_b.args[0]) == "panel_b" + keywords = { + keyword.arg: ast.unparse(keyword.value) for keyword in shared_b.keywords + } + assert keywords == { + "tiler": "self.cta_tiler", + "coord": "(0, 0, None)", + "proj": "(None, 1, 1)", + } + + partition_b = assignments["tSsB"] + assert isinstance(partition_b, ast.Call) + assert ast.unparse(partition_b.func) == "thr_mma.partition_B" + assert ast.unparse(partition_b.args[0]) == "sB" + assert "thr_mma.partition_B(panel_b)" not in ast.unparse(method) + + +def test_sm80_c96_n48_panel_retains_one_ordered_accumulator() -> None: + """Guard the strict-FP32 K-tile order of the panel decomposition.""" + tree = ast.parse(TILED_KERNEL_PATH.read_text(encoding="utf-8")) + parent = _method_node(tree, "_panel_adjoint_backward") + helper = _method_node(tree, "_backproject_panel_accumulate") + parent_source = ast.unparse(parent) + helper_source = ast.unparse(helper) + + assert "for grid_tile in cutlass.range_constexpr(GRID_TILES)" in parent_source + assert parent_source.count("self._backproject_panel_accumulate(") == 2 + assert "tCrOut_left.fill(0.0)" in parent_source + assert "tCrOut_right.fill(0.0)" in parent_source + assert ".fill(0.0)" not in helper_source + assert "panel_k_start = grid_tile * (TILE_M // self.tile_k)" in helper_source + assert "logical_k_tile = cutlass.Int32(0)" in helper_source + assert "logical_k_tile = logical_k_tile + 1" in helper_source + + +def test_packed_message_grid_initializes_the_panel_adjoint_selector() -> None: + """Keep the manually constructed backward operation compile-complete.""" + tree = ast.parse(MESSAGE_GRID_PATH.read_text(encoding="utf-8")) + factory = _method_node(tree, "_make_grid_operation") + assigned_attributes = { + target.attr + for node in ast.walk(factory) + if isinstance(node, ast.Assign) + for target in node.targets + if isinstance(target, ast.Attribute) + and isinstance(target.value, ast.Name) + and target.value.id == "operation" + } + assert "sm80_c96_n48_panel" in assigned_attributes + assert "channel_tile_start" in assigned_attributes + assert "panel_adjoint" not in assigned_attributes + + +def _reference(left, right, to_grid, from_grid): + left_grid = torch.einsum("gj,njc->ngc", to_grid, left) + right_grid = torch.einsum("gj,njc->ngc", to_grid, right) + return torch.einsum("jg,ngc->njc", from_grid, left_grid * right_grid) + + +def _reference_backward(grad_out, left, right, to_grid, from_grid): + grad_product = torch.einsum("jg,njc->ngc", from_grid, grad_out) + left_grid = torch.einsum("gj,njc->ngc", to_grid, left) + right_grid = torch.einsum("gj,njc->ngc", to_grid, right) + grad_left = torch.einsum("gj,ngc->njc", to_grid, grad_product * right_grid) + grad_right = torch.einsum("gj,ngc->njc", to_grid, grad_product * left_grid) + return grad_left, grad_right + + +def _inputs(nodes: int, hidden_channels: int): + generator = torch.Generator(device="cuda").manual_seed( + 20260703 + nodes + hidden_channels + ) + left = 0.1 * torch.randn( + nodes, + PACKED_COEFF_DIM, + hidden_channels, + device="cuda", + generator=generator, + ) + right = 0.1 * torch.randn( + left.shape, + device="cuda", + generator=generator, + ) + to_grid = 0.1 * torch.randn( + GRID_SIZE, + PACKED_COEFF_DIM, + device="cuda", + generator=generator, + ) + from_grid = 0.1 * torch.randn( + PACKED_COEFF_DIM, + GRID_SIZE, + device="cuda", + generator=generator, + ) + return left, right, to_grid, from_grid + + +def test_architecture_policy_shares_sm80_backend_with_sm86(): + from deepmd.kernels.cute.neo.runtime_policy import ( + PORTABLE_TILED_BACKEND, + PYTORCH_BACKEND, + output_grid_arch_key, + select_output_grid_backend, + ) + + for hidden_channels in SUPPORTED_HIDDEN_CHANNELS: + assert select_output_grid_backend((8, 0), hidden_channels) == ( + PORTABLE_TILED_BACKEND + ) + assert select_output_grid_backend((8, 6), hidden_channels) == ( + PORTABLE_TILED_BACKEND + ) + assert select_output_grid_backend((9, 0), hidden_channels) == ( + PORTABLE_TILED_BACKEND + ) + for capability in ((7, 5), (8, 9), (10, 0), (12, 0)): + assert ( + select_output_grid_backend(capability, hidden_channels) + == PYTORCH_BACKEND + ) + assert select_output_grid_backend((8, 0), 128) == PYTORCH_BACKEND + assert select_output_grid_backend((8, 6), 128) == PYTORCH_BACKEND + assert select_output_grid_backend((9, 0), 128) == PYTORCH_BACKEND + assert output_grid_arch_key((8, 0)) == "sm80" + assert output_grid_arch_key((8, 6)) == "sm80" + assert output_grid_arch_key((9, 0)) == "sm90" + + +@pytest.mark.skipif( + _CUTE_SKIP_REASON is not None, + reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable", +) +class TestTiledOutputGridProductForwardCuda: + @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) + @pytest.mark.parametrize("nodes", [1, 7, 65]) + def test_matches_strict_fp32_reference( + self, + nodes: int, + hidden_channels: int, + ): + from deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product import ( + run_tiled_output_grid_product, + ) + + left, right, to_grid, from_grid = _inputs(nodes, hidden_channels) + expected = _reference(left, right, to_grid, from_grid) + actual = run_tiled_output_grid_product(left, right, to_grid, from_grid) + torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL) + + @pytest.mark.skipif( + _SM80_SKIP_REASON is not None, + reason=_SM80_SKIP_REASON or "sm80-family GPU is unavailable", + ) + @pytest.mark.parametrize("nodes", [1, 7, 65]) + def test_sm80_c96_n48_matches_strict_fp32_reference(self, nodes: int): + from deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product import ( + run_tiled_output_grid_product, + ) + + left, right, to_grid, from_grid = _inputs(nodes, 96) + expected = _reference(left, right, to_grid, from_grid) + actual = run_tiled_output_grid_product( + left, + right, + to_grid, + from_grid, + use_sm80_c96_n48=True, + ) + torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL) + + @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) + def test_one_compile_accepts_symbolic_node_counts(self, hidden_channels: int): + from deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product import ( + _compiled_tiled_forward, + run_tiled_output_grid_product, + ) + + before = _compiled_tiled_forward.cache_info() + for nodes in (3, 19): + left, right, to_grid, from_grid = _inputs(nodes, hidden_channels) + actual = run_tiled_output_grid_product( + left, + right, + to_grid, + from_grid, + ) + torch.testing.assert_close( + actual, + _reference(left, right, to_grid, from_grid), + atol=TOL, + rtol=TOL, + ) + after = _compiled_tiled_forward.cache_info() + assert after.misses - before.misses <= 1 + assert after.hits > before.hits + + @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) + def test_launches_on_current_non_default_stream(self, hidden_channels: int): + from deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product import ( + run_tiled_output_grid_product, + ) + + left, right, to_grid, from_grid = _inputs(11, hidden_channels) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + actual = run_tiled_output_grid_product( + left, + right, + to_grid, + from_grid, + ) + expected = _reference(left, right, to_grid, from_grid) + torch.cuda.current_stream().wait_stream(stream) + torch.testing.assert_close(actual, expected, atol=TOL, rtol=TOL) + + +@pytest.mark.skipif( + _CUTE_SKIP_REASON is not None, + reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable", +) +class TestTiledOutputGridProductBackwardCuda: + @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) + @pytest.mark.parametrize("nodes", [1, 7, 65]) + def test_first_backward_matches_strict_fp32_reference( + self, + nodes: int, + hidden_channels: int, + ): + from deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product import ( + run_tiled_output_grid_product_backward, + ) + + left, right, to_grid, from_grid = _inputs(nodes, hidden_channels) + grad_out = torch.randn_like(left) + expected = _reference_backward(grad_out, left, right, to_grid, from_grid) + actual = run_tiled_output_grid_product_backward( + grad_out, + left, + right, + to_grid, + from_grid, + ) + torch.testing.assert_close(actual[0], expected[0], atol=TOL, rtol=TOL) + torch.testing.assert_close(actual[1], expected[1], atol=TOL, rtol=TOL) + + @pytest.mark.skipif( + _SM80_SKIP_REASON is not None, + reason=_SM80_SKIP_REASON or "sm80-family GPU is unavailable", + ) + @pytest.mark.parametrize("nodes", [1, 7, 65]) + def test_sm80_c96_n48_panel_matches_strict_fp32_reference( + self, + nodes: int, + ): + from deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product import ( + run_tiled_output_grid_product_backward, + ) + + left, right, to_grid, from_grid = _inputs(nodes, 96) + grad_out = torch.randn_like(left) + expected = _reference_backward(grad_out, left, right, to_grid, from_grid) + actual = run_tiled_output_grid_product_backward( + grad_out, + left, + right, + to_grid, + from_grid, + use_sm80_c96_n48_panel=True, + ) + torch.testing.assert_close(actual[0], expected[0], atol=TOL, rtol=TOL) + torch.testing.assert_close(actual[1], expected[1], atol=TOL, rtol=TOL) + + @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) + def test_one_backward_compile_accepts_symbolic_node_counts( + self, + hidden_channels: int, + ): + from deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product import ( + _compiled_tiled_backward, + run_tiled_output_grid_product_backward, + ) + + before = _compiled_tiled_backward.cache_info() + for nodes in (3, 19): + left, right, to_grid, from_grid = _inputs(nodes, hidden_channels) + grad_out = torch.randn_like(left) + expected = _reference_backward(grad_out, left, right, to_grid, from_grid) + actual = run_tiled_output_grid_product_backward( + grad_out, + left, + right, + to_grid, + from_grid, + ) + torch.testing.assert_close(actual[0], expected[0], atol=TOL, rtol=TOL) + torch.testing.assert_close(actual[1], expected[1], atol=TOL, rtol=TOL) + after = _compiled_tiled_backward.cache_info() + assert after.misses - before.misses <= 1 + assert after.hits > before.hits diff --git a/source/tests/pt/model/test_descriptor_sezm_readout_l0_algebra.py b/source/tests/pt/model/test_descriptor_sezm_readout_l0_algebra.py new file mode 100644 index 0000000000..e4b83402dd --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_readout_l0_algebra.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""CPU algebra checks for the Neo degree-zero readout boundary.""" + +from __future__ import ( + annotations, +) + +import torch + +COEFF_DIM = 16 +N_FRAMES = 3 +GRID_SIZE = 152 +HIDDEN_CHANNELS = 192 + + +def _readout_reference(left, right, to_grid, from_grid): + left_grid = torch.einsum("gj,njh->ngh", to_grid, left) + right_grid = torch.einsum("gj,njh->ngh", to_grid, right) + return torch.einsum("g,ngh->nh", from_grid[0], left_grid * right_grid) + + +def _gram_reference(left, right, gram): + transformed_right = torch.einsum("ij,njh->nih", gram, right) + return torch.sum(left * transformed_right, dim=1) + + +def _output_ffn(): + from deepmd.pt.model.descriptor.sezm_nn.ffn import ( + EquivariantFFN, + ) + + return ( + EquivariantFFN( + lmax=3, + channels=32, + hidden_channels=96, + kmax=1, + grid_mlp=True, + grid_branch=0, + dtype=torch.float32, + s2_activation=False, + ffn_so3_grid=True, + activation_function="silu", + glu_activation=True, + mlp_bias=False, + trainable=False, + seed=17, + ) + .to("cpu") + .eval() + ) + + +def test_final_slice_has_exactly_zero_l_greater_than_zero_output_cotangent(): + module = _output_ffn() + value = torch.randn( + 2, + COEFF_DIM, + 1, + 32, + device="cpu", + requires_grad=True, + ) + ffn_out = module(value) + seed = torch.randn(2, 32, device="cpu") + + cotangent = torch.autograd.grad( + (value + ffn_out)[:, 0, 0, :], + ffn_out, + seed, + )[0] + + assert torch.equal(cotangent[:, 0, 0, :], seed) + assert torch.count_nonzero(cotangent[:, 1:, :, :]) == 0 + + +def test_degree_zero_nonzero_frame_projector_rows_are_structural_zero(): + projector = _output_ffn().act.projector + from_grid = projector.from_grid_mat.reshape( + COEFF_DIM, + N_FRAMES, + GRID_SIZE, + ) + + assert projector.frame_set == [0, -1, 1] + assert torch.count_nonzero(from_grid[0, 0]) > 0 + assert torch.count_nonzero(from_grid[0, 1]) == 0 + assert torch.count_nonzero(from_grid[0, 2]) == 0 + + +def test_scalar_readout_matches_row_zero_of_generic_grid_product(): + generator = torch.Generator().manual_seed(20260704) + left = 0.1 * torch.randn( + 3, + COEFF_DIM * N_FRAMES, + HIDDEN_CHANNELS, + device="cpu", + generator=generator, + ) + right = 0.1 * torch.randn( + left.shape, + device="cpu", + generator=generator, + ) + projector = _output_ffn().act.projector + to_grid = projector.to_grid_mat + from_grid = projector.from_grid_mat + + left_grid = torch.einsum("gj,njh->ngh", to_grid, left) + right_grid = torch.einsum("gj,njh->ngh", to_grid, right) + generic = torch.einsum( + "jg,ngh->njh", + from_grid, + left_grid * right_grid, + ) + + actual = _readout_reference(left, right, to_grid, from_grid) + + assert torch.allclose(actual, generic[:, 0, :], atol=5.0e-5, rtol=5.0e-5) + + +def test_dense_gram_matches_both_projected_input_adjoints(): + from deepmd.kernels.cute.neo.readout_l0 import ( + build_readout_l0_gram, + ) + + generator = torch.Generator().manual_seed(20260718) + left = 0.1 * torch.randn( + 2, + COEFF_DIM * N_FRAMES, + HIDDEN_CHANNELS, + device="cpu", + generator=generator, + ) + right = 0.1 * torch.randn(left.shape, device="cpu", generator=generator) + dq0 = 0.1 * torch.randn( + left.shape[0], + HIDDEN_CHANNELS, + device="cpu", + generator=generator, + ) + projector = _output_ffn().act.projector + gram = build_readout_l0_gram( + projector.to_grid_mat, + projector.from_grid_mat, + ) + + left_ref = left.detach().clone().requires_grad_(True) + right_ref = right.detach().clone().requires_grad_(True) + q0 = _readout_reference( + left_ref, + right_ref, + projector.to_grid_mat, + projector.from_grid_mat, + ) + expected_left, expected_right = torch.autograd.grad( + q0, + (left_ref, right_ref), + dq0, + ) + actual_left = dq0[:, None, :] * torch.einsum( + "ij,njh->nih", + gram, + right, + ) + actual_right = dq0[:, None, :] * torch.einsum( + "ji,njh->nih", + gram, + left, + ) + + assert gram.shape == (COEFF_DIM * N_FRAMES, COEFF_DIM * N_FRAMES) + assert gram.dtype == torch.float32 + assert gram.is_contiguous() + assert not gram.requires_grad + torch.testing.assert_close(actual_left, expected_left, atol=5.0e-5, rtol=5.0e-5) + torch.testing.assert_close( + actual_right, + expected_right, + atol=5.0e-5, + rtol=5.0e-5, + ) + + +def test_dense_gram_forward_matches_row_zero_grid_projection(): + from deepmd.kernels.cute.neo.readout_l0 import ( + build_readout_l0_gram, + ) + + generator = torch.Generator().manual_seed(20260720) + left = 0.1 * torch.randn( + 3, + COEFF_DIM * N_FRAMES, + HIDDEN_CHANNELS, + device="cpu", + generator=generator, + ) + right = 0.1 * torch.randn(left.shape, device="cpu", generator=generator) + projector = _output_ffn().act.projector + gram = build_readout_l0_gram( + projector.to_grid_mat, + projector.from_grid_mat, + ) + + expected = _readout_reference( + left, + right, + projector.to_grid_mat, + projector.from_grid_mat, + ) + actual = _gram_reference(left, right, gram) + + torch.testing.assert_close(actual, expected, atol=5.0e-5, rtol=5.0e-5) + + +def test_dense_gram_forward_is_channelwise_without_cross_channel_mixing(): + generator = torch.Generator().manual_seed(20260721) + left = torch.randn( + 2, + COEFF_DIM * N_FRAMES, + HIDDEN_CHANNELS, + device="cpu", + generator=generator, + ) + right = torch.randn(left.shape, device="cpu", generator=generator) + gram = torch.randn( + COEFF_DIM * N_FRAMES, + COEFF_DIM * N_FRAMES, + device="cpu", + generator=generator, + ) + changed_channel = 37 + + baseline = _gram_reference(left, right, gram) + changed_right = right.clone() + changed_right[:, :, changed_channel].mul_(1.5) + changed = _gram_reference(left, changed_right, gram) + + unaffected = torch.ones(HIDDEN_CHANNELS, dtype=torch.bool, device="cpu") + unaffected[changed_channel] = False + torch.testing.assert_close(changed[:, unaffected], baseline[:, unaffected]) + torch.testing.assert_close( + changed[:, changed_channel], + 1.5 * baseline[:, changed_channel], + ) diff --git a/source/tests/pt/model/test_mlp.py b/source/tests/pt/model/test_mlp.py index 5f067cb0e6..1f141301a5 100644 --- a/source/tests/pt/model/test_mlp.py +++ b/source/tests/pt/model/test_mlp.py @@ -1,6 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import itertools +import os import unittest +from unittest import ( + mock, +) import numpy as np import torch @@ -11,6 +15,7 @@ NativeLayer, NativeNet, ) +from deepmd.pt.model.network import mlp as mlp_module from deepmd.pt.model.network.mlp import ( MLP, EmbeddingNet, @@ -84,6 +89,99 @@ def test_jit(self) -> None: ml1 = MLPLayer.deserialize(ml.serialize()) model = torch.jit.script(ml1) + def test_thin_k1_eval_linear_traces_without_addmm(self) -> None: + from torch.fx.experimental.proxy_tensor import ( + make_fx, + ) + + layer = MLPLayer( + 5, + 8, + bias=True, + activation_function="none", + precision="float32", + trainable=False, + ).to(env.DEVICE) + layer.eval() + mlp_module.enable_neo_cute_compile_visible_linears(layer) + value = torch.randn(7, 5, device=env.DEVICE, dtype=torch.float32) + environment = { + "DP_NEO_CUTE_INFER": "1", + "DP_CUTE_K1_THIN_WRAPPER": "1", + } + + with mock.patch.dict(os.environ, environment, clear=False): + graph = make_fx(layer)(value) + with mock.patch.object( + mlp_module.F, + "linear", + side_effect=AssertionError("thin eval path reached aten::linear"), + ): + actual = layer(value) + + expected = mlp_module.F.linear(value, layer.matrix.t(), layer.bias) + torch.testing.assert_close(actual, expected) + targets = { + node.target for node in graph.graph.nodes if node.op == "call_function" + } + self.assertIn(torch.ops.aten.mm.default, targets) + self.assertNotIn(torch.ops.aten.addmm.default, targets) + + def test_thin_k1_eval_linear_is_scoped_to_marked_model(self) -> None: + layer = MLPLayer( + 5, + 8, + bias=True, + activation_function="none", + precision="float32", + trainable=False, + ).to(env.DEVICE) + layer.eval() + value = torch.randn(7, 5, device=env.DEVICE, dtype=torch.float32) + environment = { + "DP_NEO_CUTE_INFER": "1", + "DP_CUTE_K1_THIN_WRAPPER": "1", + } + + with ( + mock.patch.dict(os.environ, environment, clear=False), + mock.patch.object( + mlp_module, + "_matmul_bias", + side_effect=AssertionError("unmarked MLP reached Neo-only topology"), + ), + ): + actual = layer(value) + + expected = mlp_module.F.linear(value, layer.matrix.t(), layer.bias) + torch.testing.assert_close(actual, expected) + + def test_thin_k1_linear_defaults_on_only_for_sm80(self) -> None: + with ( + mock.patch.dict( + os.environ, + {"DP_NEO_CUTE_INFER": "1"}, + clear=True, + ), + mock.patch.object(torch.cuda, "is_available", return_value=True), + mock.patch.object(torch.cuda, "get_device_capability", return_value=(8, 0)), + ): + self.assertTrue(mlp_module._use_k1_compile_visible_linear()) + + with ( + mock.patch.dict( + os.environ, + { + "DP_NEO_CUTE_INFER": "1", + "DP_CUTE_K1_THIN_WRAPPER": "0", + }, + clear=True, + ), + mock.patch.object(torch.cuda, "is_available", return_value=True), + mock.patch.object(torch.cuda, "get_device_capability", return_value=(8, 0)), + ): + self.assertFalse(mlp_module._use_k1_compile_visible_linear()) + class TestMLP(unittest.TestCase): def setUp(self) -> None: diff --git a/source/tests/pt/model/test_nlist_backend.py b/source/tests/pt/model/test_nlist_backend.py index d38262dacd..6d243a6b4e 100644 --- a/source/tests/pt/model/test_nlist_backend.py +++ b/source/tests/pt/model/test_nlist_backend.py @@ -10,6 +10,9 @@ """ import copy +from types import ( + SimpleNamespace, +) import numpy as np import pytest @@ -172,6 +175,55 @@ def test_self_built_model_forces_native(pt_files, monkeypatch) -> None: assert deep_eval._nlist_builder is None +def test_edge_strategy_freezes_parameters_only_during_inference_call() -> None: + from deepmd.pt.infer.deep_eval import DeepEval as PTDeepEval + + class ProbeModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.ones((), device="cpu")) + self.saw_frozen = False + + def get_sel(self) -> list[int]: + return [4] + + def forward_common_lower(self, *args, **kwargs) -> dict[str, torch.Tensor]: + self.saw_frozen = not self.weight.requires_grad + return {} + + class ProbeBuilder: + def build(self, coord, atype, *args, **kwargs) -> SimpleNamespace: + return SimpleNamespace( + coord=coord, + atype=atype, + edge_index=torch.empty((2, 0), dtype=torch.long, device="cpu"), + edge_vec=torch.empty((0, 3), dtype=coord.dtype, device="cpu"), + edge_scatter_index=torch.empty((2, 0), dtype=torch.long, device="cpu"), + edge_mask=torch.empty((0,), dtype=torch.bool, device="cpu"), + ) + + inner = ProbeModel() + deep_eval = object.__new__(PTDeepEval) + deep_eval.dp = ModelWrapper(inner) + deep_eval._uses_edge_schema = True + deep_eval._nlist_builder = ProbeBuilder() + deep_eval.rcut = 4.0 + + PTDeepEval._eval_lower_strategy( + deep_eval, + torch.zeros((1, 1, 3), device="cpu"), + torch.zeros((1, 1), dtype=torch.long, device="cpu"), + None, + None, + None, + None, + False, + ) + + assert inner.saw_frozen + assert inner.weight.requires_grad + + # --- equivalence with the native dense builder ------------------------------ diff --git a/source/tests/pt/model/test_sezm_cute_efs.py b/source/tests/pt/model/test_sezm_cute_efs.py new file mode 100644 index 0000000000..08e42e6ba4 --- /dev/null +++ b/source/tests/pt/model/test_sezm_cute_efs.py @@ -0,0 +1,452 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""End-to-end strict-FP32 acceptance tests for the Neo CuTe path.""" + +from __future__ import ( + annotations, +) + +import contextlib +import importlib +import os +import subprocess +import sys +from pathlib import ( + Path, +) +from typing import ( + TYPE_CHECKING, + Any, +) + +import pytest +import torch + +if TYPE_CHECKING: + from collections.abc import ( + Iterator, + ) + + +TOL = 5.0e-5 +OUTPUT_KEYS = ("energy", "force", "atom_energy", "virial") +SM80_CAPABILITIES = frozenset({(8, 0), (8, 6)}) +SUPPORTED_CAPABILITIES = SM80_CAPABILITIES | frozenset( + {(8, 9), (9, 0), (10, 0), (12, 0)} +) +REPO_ROOT = Path(__file__).resolve().parents[4] +_CHILD_MODE = "--neo-cute-efs-child" + + +def _cute_runtime_skip_reason() -> str | None: + if not torch.cuda.is_available(): + return "Neo CuTe E/F/S parity requires CUDA" + capability = tuple(torch.cuda.get_device_capability()) + if capability not in SUPPORTED_CAPABILITIES: + return f"Neo CuTe E/F/S parity does not support {capability}" + try: + importlib.import_module("cutlass.cute") + importlib.import_module("cuda.bindings.driver") + except Exception as exc: # pragma: no cover - runtime dependent + return f"Neo CuTe E/F/S acceptance requires the CuTe DSL runtime: {exc}" + return None + + +_CUTE_SKIP_REASON = _cute_runtime_skip_reason() + + +@contextlib.contextmanager +def _strict_fp32() -> Iterator[None]: + matmul = torch.backends.cuda.matmul + cudnn = torch.backends.cudnn + prior_matmul_tf32 = matmul.allow_tf32 + prior_cudnn_tf32 = cudnn.allow_tf32 + prior_precision = torch.get_float32_matmul_precision() + try: + matmul.allow_tf32 = False + cudnn.allow_tf32 = False + torch.set_float32_matmul_precision("highest") + yield + finally: + matmul.allow_tf32 = prior_matmul_tf32 + cudnn.allow_tf32 = prior_cudnn_tf32 + torch.set_float32_matmul_precision(prior_precision) + + +def _neo_model(*, use_compile: bool) -> torch.nn.Module: + from deepmd.pt.model.model import ( + get_sezm_model, + ) + + model = get_sezm_model( + { + "type": "SeZM", + "type_map": ["O", "H"], + "descriptor": { + "type": "SeZM", + "sel": 32, + "rcut": 3.0, + "channels": 32, + "n_radial": 16, + "use_env_seed": True, + "lmax": 3, + "mmax": 1, + "n_blocks": 2, + "so2_layers": 3, + "radial_so2_mode": "degree_channel", + "radial_so2_rank": 1, + "n_focus": 2, + "focus_dim": 0, + "n_atten_head": 1, + "message_node_so3": True, + "ffn_neurons": 0, + "ffn_so3_grid": True, + "grid_mlp": False, + "grid_branch": [0, 0, 1], + "ffn_blocks": 1, + "so3_readout": "mlp", + "use_amp": False, + "precision": "float32", + "seed": 42, + }, + "fitting_net": { + "neuron": [0], + "precision": "float32", + "seed": 42, + }, + "use_compile": use_compile, + "enable_tf32": False, + } + ).to(device="cuda", dtype=torch.float32) + model.eval() + for parameter in model.parameters(): + parameter.requires_grad_(False) + return model + + +def _water_frame() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + coord = torch.tensor( + [ + [ + [0.70, 0.80, 0.90], + [1.58, 0.80, 0.90], + [0.43, 1.63, 0.90], + [3.20, 3.00, 2.80], + [4.08, 3.00, 2.80], + [2.93, 3.83, 2.80], + ] + ], + device="cuda", + dtype=torch.float32, + ) + atype = torch.tensor( + [[0, 1, 1, 0, 1, 1]], + device="cuda", + dtype=torch.int32, + ) + box = torch.tensor( + [[5.4, 0.0, 0.0, 0.0, 5.2, 0.0, 0.0, 0.0, 5.0]], + device="cuda", + dtype=torch.float32, + ) + return coord, atype, box + + +def _triclinic_frame() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + box_matrix = torch.tensor( + [ + [5.2, 0.0, 0.0], + [0.7, 4.8, 0.0], + [0.3, 0.5, 5.0], + ], + device="cuda", + dtype=torch.float32, + ) + fractional = torch.tensor( + [ + [0.08, 0.11, 0.14], + [0.25, 0.14, 0.18], + [0.12, 0.31, 0.22], + [0.54, 0.57, 0.49], + [0.72, 0.59, 0.51], + [0.57, 0.76, 0.55], + [0.88, 0.16, 0.81], + [0.06, 0.22, 0.84], + ], + device="cuda", + dtype=torch.float32, + ) + coord = (fractional @ box_matrix).unsqueeze(0) + atype = torch.tensor( + [[0, 1, 1, 0, 1, 1, 0, 1]], + device="cuda", + dtype=torch.int32, + ) + return coord, atype, box_matrix.reshape(1, 9) + + +def _detached_outputs(outputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + return {name: outputs[name].detach().cpu().clone() for name in OUTPUT_KEYS} + + +_FRAME_FACTORIES = { + "water": _water_frame, + "triclinic": _triclinic_frame, +} + + +def _assert_runtime_profile() -> dict[str, bool]: + from deepmd.kernels.cute.neo import ( + k1, + runtime_policy, + ) + + capability = tuple(torch.cuda.get_device_capability()) + if capability not in SUPPORTED_CAPABILITIES: + raise AssertionError(f"unsupported K1 capability {capability}") + policy_checks = { + "master": runtime_policy.is_cute_infer_enabled(), + "triton_infer_2": os.environ.get("DP_TRITON_INFER") == "2", + "supported": runtime_policy.is_supported_k1_capability(capability), + "packed_wigner": runtime_policy.is_packed_wigner_enabled(capability), + } + config = k1._architecture_default_config(capability) + if capability in SM80_CAPABILITIES: + architecture_checks = { + "sm80_profile": runtime_policy.is_sm80_profile_enabled(capability), + "gie": runtime_policy.is_gie_enabled(capability), + "thin_wrapper": runtime_policy.is_k1_thin_wrapper_enabled(capability), + "output_grid_bwd_panel": ( + runtime_policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled(capability) + ), + "output_grid_fwd": ( + runtime_policy.is_output_grid_fwd_sm80_c96_n48_enabled(capability) + ), + "readout_fold": runtime_policy.is_readout_input_fold_sm80_enabled( + capability + ), + "per_focus_so2_forward": config.per_focus_so2_fwd_pair, + "no_native_sm90_path": not config.native_sm90_path, + } + elif capability == (9, 0): + architecture_checks = { + "native_sm90_path": config.native_sm90_path, + "shared_so2_forward": not config.per_focus_so2_fwd_pair, + "sm90_output_grid": ( + runtime_policy.is_output_grid_sm90_c96_asymmetric_panels_enabled( + capability + ) + ), + "sm90_readout_fold": runtime_policy.is_readout_input_fold_sm90_enabled( + capability + ), + } + elif capability in runtime_policy.FUSED_SO2_GATE_CAPABILITIES: + architecture_checks = { + "combined_so2_gate": config.combined_so2_gate, + "shared_so2_forward": not config.per_focus_so2_fwd_pair, + "no_native_sm90_path": not config.native_sm90_path, + } + else: + architecture_checks = { + "no_native_sm90_path": not config.native_sm90_path, + "shared_so2_forward": not config.per_focus_so2_fwd_pair, + } + selected = {**policy_checks, **architecture_checks} + missing = sorted(name for name, enabled in selected.items() if not enabled) + if missing: + raise AssertionError( + "DP_NEO_CUTE_INFER did not select the expected runtime profile: " + + ", ".join(missing) + ) + return selected + + +def _run_efs_child( + *, + frame_name: str, + use_compile: bool, + output_path: Path, +) -> None: + from deepmd.kernels.cute.neo import ( + k1, + ) + + torch.manual_seed(20260726) + torch.cuda.manual_seed_all(20260726) + runner_calls = 0 + original_build_runner = k1._build_runner + + def counted_build_runner(*args: Any, **kwargs: Any) -> Any: + nonlocal runner_calls + runner_calls += 1 + return original_build_runner(*args, **kwargs) + + if use_compile: + k1._build_runner = counted_build_runner + try: + model = _neo_model(use_compile=use_compile) + state_keys_before = tuple(model.state_dict()) + coord, atype, box = _FRAME_FACTORIES[frame_name]() + with _strict_fp32(): + profile = _assert_runtime_profile() if use_compile else {} + run_count = 2 if use_compile else 1 + runs = [ + _detached_outputs(model(coord, atype, box=box)) + for _ in range(run_count) + ] + torch.cuda.synchronize() + state_keys_after = tuple(model.state_dict()) + finally: + if use_compile: + k1._build_runner = original_build_runner + + if state_keys_after != state_keys_before: + raise AssertionError("CuTe warmup changed the model state_dict contract") + if use_compile and runner_calls == 0: + raise AssertionError("compiled CuTe run did not instantiate the K1 runner") + torch.save( + { + "outputs": runs, + "profile": profile, + "runner_calls": runner_calls, + "use_compile": use_compile, + }, + output_path, + ) + + +def _clean_child_environment() -> dict[str, str]: + environment = os.environ.copy() + for name in tuple(environment): + if name.startswith(("DP_CUTE_", "DP_NEO_CUTE_")): + environment.pop(name) + for name in ( + "DP_COMPILE_INFER", + "DP_INTERFACE_PREC", + "DP_TF32_INFER", + "DP_TRITON_INFER", + "NVIDIA_TF32_OVERRIDE", + ): + environment.pop(name, None) + pythonpath = environment.get("PYTHONPATH") + environment["PYTHONPATH"] = ( + str(REPO_ROOT) if not pythonpath else f"{REPO_ROOT}{os.pathsep}{pythonpath}" + ) + environment.update( + { + "DP_INTERFACE_PREC": "low", + "DP_TF32_INFER": "0", + "NVIDIA_TF32_OVERRIDE": "0", + "PYTHONDONTWRITEBYTECODE": "1", + } + ) + return environment + + +def _child_environment(*, use_compile: bool) -> dict[str, str]: + environment = _clean_child_environment() + if use_compile: + environment.update( + { + "DP_COMPILE_INFER": "1", + "DP_CUTE_STRICT": "1", + "DP_NEO_CUTE_INFER": "1", + "DP_TRITON_INFER": "2", + } + ) + else: + environment.update( + { + "DP_COMPILE_INFER": "0", + "DP_CUTE_INFER": "0", + "DP_NEO_CUTE_INFER": "0", + "DP_TRITON_INFER": "0", + } + ) + return environment + + +def _run_child_process( + *, + frame_name: str, + use_compile: bool, + output_path: Path, +) -> dict[str, Any]: + result = subprocess.run( + [ + sys.executable, + str(Path(__file__).resolve()), + _CHILD_MODE, + frame_name, + "compiled" if use_compile else "eager", + str(output_path), + ], + cwd=REPO_ROOT, + env=_child_environment(use_compile=use_compile), + capture_output=True, + text=True, + timeout=900, + check=False, + ) + if result.returncode: + pytest.fail( + "Neo CuTe E/F/S child failed.\n" + f"stdout:\n{result.stdout}\n" + f"stderr:\n{result.stderr}" + ) + return torch.load(output_path, map_location="cpu", weights_only=True) + + +@pytest.mark.skipif( + _CUTE_SKIP_REASON is not None, + reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable", +) +@pytest.mark.parametrize("frame_name", tuple(_FRAME_FACTORIES)) +def test_neo_cute_compiled_profile_matches_eager_at_shared_5e5( + tmp_path: Path, + frame_name: str, +) -> None: + eager = _run_child_process( + frame_name=frame_name, + use_compile=False, + output_path=tmp_path / f"{frame_name}-eager.pt", + ) + compiled = _run_child_process( + frame_name=frame_name, + use_compile=True, + output_path=tmp_path / f"{frame_name}-compiled.pt", + ) + + assert not eager["use_compile"] + assert compiled["use_compile"] + assert compiled["runner_calls"] > 0 + assert compiled["profile"] + assert all(compiled["profile"].values()) + expected = eager["outputs"][0] + for run_index, actual in enumerate(compiled["outputs"]): + for name in OUTPUT_KEYS: + torch.testing.assert_close( + actual[name], + expected[name], + atol=TOL, + rtol=TOL, + msg=( + f"compiled Neo CuTe {frame_name} run {run_index} {name} " + "differs from eager PyTorch" + ), + ) + + +def _main() -> None: + if len(sys.argv) != 5 or sys.argv[1] != _CHILD_MODE: + raise SystemExit("this test module is only executable in E/F/S child mode") + _run_efs_child( + frame_name=sys.argv[2], + use_compile=sys.argv[3] == "compiled", + output_path=Path(sys.argv[4]), + ) + + +if __name__ == "__main__": + _main() diff --git a/source/tests/pt/model/test_sezm_cute_multigpu_dispatch.py b/source/tests/pt/model/test_sezm_cute_multigpu_dispatch.py new file mode 100644 index 0000000000..3ec7830d1e --- /dev/null +++ b/source/tests/pt/model/test_sezm_cute_multigpu_dispatch.py @@ -0,0 +1,279 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""CPU/mock coverage for operand-device-aware Neo CuTe dispatch.""" + +from __future__ import ( + annotations, +) + +import importlib +import os +from types import ( + SimpleNamespace, +) +from unittest import ( + mock, +) + +import pytest +import torch + +from deepmd.kernels.cute.neo import ( + k1, + runtime_policy, +) +from deepmd.pt.model.network import ( + mlp, +) + + +def _cuda_operand(index: int = 1) -> SimpleNamespace: + return SimpleNamespace( + device=torch.device("cuda", index), + is_cuda=True, + ) + + +def test_k1_custom_ops_are_direct_and_thin_path_uses_operand_device() -> None: + operand = _cuda_operand() + sentinel = object() + + with mock.patch.object(k1, "_cute_k1_impl", return_value=sentinel) as direct: + actual = k1.cute_k1( + 1, + operand, + object(), + object(), + object(), + object(), + object(), + object(), + object(), + object(), + object(), + object(), + ) + + assert actual is sentinel + direct.assert_called_once() + + context = object() + with mock.patch.object( + k1, + "_k1_packed_direct_registered_backward_impl", + return_value=sentinel, + ) as backward: + actual = k1._k1_packed_direct_backward(context, operand, None) + + assert actual is sentinel + backward.assert_called_once_with(context, operand) + + with ( + mock.patch.object( + k1, + "_cuda_compute_capability", + return_value=(9, 0), + ) as capability, + mock.patch.object( + runtime_policy, + "is_k1_thin_wrapper_enabled", + return_value=False, + ) as thin_selector, + mock.patch.object( + k1, + "_maybe_run_cute_k1_fallback", + return_value=sentinel, + ) as fallback, + ): + actual = k1.maybe_run_cute_k1(object(), operand, object(), object()) + + assert actual is sentinel + capability.assert_called_once_with(1) + thin_selector.assert_called_once_with((9, 0)) + fallback.assert_called_once() + + +def test_k1_prepare_uses_requested_device_instead_of_current_device() -> None: + device = torch.device("cuda", 1) + with ( + mock.patch.object( + k1, + "_cuda_compute_capability", + return_value=(8, 6), + ) as capability, + mock.patch.object( + k1, + "is_supported_k1_compute_capability", + return_value=False, + ) as supported, + mock.patch.object( + torch.cuda, + "current_device", + side_effect=AssertionError("must not read the current CUDA device"), + ), + ): + assert not k1.prepare_cute_k1_blocks( + [object()], + training=False, + device=device, + dtype=torch.float32, + ) + + capability.assert_called_once_with(1) + supported.assert_called_once_with((8, 6)) + + +def test_k4_custom_ops_are_called_directly() -> None: + pytest.importorskip("cutlass.cute") + k4_wignerd = importlib.import_module("deepmd.kernels.cute.neo.k4_wignerd") + operand = _cuda_operand() + sentinel = object() + + with mock.patch.object( + k4_wignerd, + "_run_cute_wignerd_impl", + return_value=sentinel, + ) as forward: + actual = k4_wignerd.run_cute_wignerd(operand, object()) + + assert actual is sentinel + forward.assert_called_once_with(operand, mock.ANY, packed_wigner=False) + + context = SimpleNamespace(saved_tensors=(operand,)) + grad_panel = object() + with mock.patch.object( + k4_wignerd, + "_wignerd_panel_registered_backward_impl", + return_value=sentinel, + ) as backward: + actual = k4_wignerd._wignerd_panel_backward(context, grad_panel) + + assert actual is sentinel + backward.assert_called_once_with(context, grad_panel) + + +def test_mlp_selector_and_forward_use_input_device() -> None: + cuda1 = torch.device("cuda", 1) + + def capability(device: torch.device | None = None) -> tuple[int, int]: + return (8, 0) if device == cuda1 else (9, 0) + + with ( + mock.patch.dict( + os.environ, + {"DP_NEO_CUTE_INFER": "1"}, + clear=True, + ), + mock.patch.object(torch.cuda, "is_available", return_value=True), + mock.patch.object( + torch.cuda, + "get_device_capability", + side_effect=capability, + ) as get_capability, + ): + assert mlp._use_k1_compile_visible_linear(cuda1) + assert not mlp._use_k1_compile_visible_linear(torch.device("cuda", 0)) + assert not mlp._use_k1_compile_visible_linear(torch.device("cpu")) + + assert get_capability.call_args_list == [ + mock.call(cuda1), + mock.call(torch.device("cuda", 0)), + ] + + layer = mlp.MLPLayer( + 4, + 4, + activation_function="none", + precision="float32", + ).to("cpu") + layer.eval() + mlp.enable_neo_cute_compile_visible_linears(layer) + value = torch.randn(3, 4, device="cpu") + with mock.patch.object( + mlp, + "_use_k1_compile_visible_linear", + return_value=False, + ) as selector: + layer(value) + selector.assert_called_once_with(value.device) + + +def test_nv_neighbor_list_selector_and_build_use_coordinate_device() -> None: + sezm_model = importlib.import_module("deepmd.pt.model.model.sezm_model") + cuda1 = torch.device("cuda", 1) + + def capability(device: torch.device | None = None) -> tuple[int, int]: + return (8, 0) if device == cuda1 else (9, 0) + + with ( + mock.patch.dict( + os.environ, + {"DP_NEO_CUTE_INFER": "1"}, + clear=True, + ), + mock.patch.object(torch.cuda, "is_available", return_value=True), + mock.patch.object( + torch.cuda, + "get_device_capability", + side_effect=capability, + ) as get_capability, + ): + assert sezm_model._neo_cute_nlist_eager_island_enabled(cuda1) + assert not sezm_model._neo_cute_nlist_eager_island_enabled( + torch.device("cuda", 0) + ) + assert not sezm_model._neo_cute_nlist_eager_island_enabled(torch.device("cpu")) + + assert get_capability.call_args_list == [ + mock.call(cuda1), + mock.call(torch.device("cuda", 0)), + ] + + builder = sezm_model.NvNeighborList() + sentinel = object() + with ( + mock.patch.object( + sezm_model, + "_neo_cute_nlist_eager_island_enabled", + return_value=False, + ) as selector, + mock.patch.object(builder, "build", return_value=sentinel) as build, + ): + actual = sezm_model._build_neo_neighbor_list( + builder, + mock.Mock(device=cuda1), + object(), + None, + 4.0, + [32], + return_mode="edges", + ) + + assert actual is sentinel + selector.assert_called_once_with(cuda1) + build.assert_called_once() + + with ( + mock.patch.object( + sezm_model, + "_neo_cute_nlist_eager_island_enabled", + return_value=True, + ), + mock.patch.object( + sezm_model, + "_build_neo_neighbor_list_eager_island", + return_value=sentinel, + ) as eager, + ): + actual = sezm_model._build_neo_neighbor_list( + builder, + mock.Mock(device=cuda1), + object(), + None, + 4.0, + [32], + return_mode="edges", + ) + + assert actual is sentinel + eager.assert_called_once() diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index b9d07df06b..cd5f6a65fd 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -40,6 +40,7 @@ from deepmd.pt.model.model.sezm_model import ( InnerPotential, SeZMModel, + _sort_edge_tensors_by_destination, ) from deepmd.pt.model.model.sezm_native_spin_model import ( SeZMNativeSpinModel, @@ -230,6 +231,39 @@ def setUp(self) -> None: self.device = env.DEVICE torch.manual_seed(2024) + def test_destination_edge_sort_is_empty_safe_and_keeps_alignment(self) -> None: + empty_index = torch.empty((2, 0), dtype=torch.long, device="cpu") + empty_vec = torch.empty((0, 3), device="cpu") + empty_mask = torch.empty((0,), dtype=torch.bool, device="cpu") + empty_scatter = torch.empty((2, 0), dtype=torch.long, device="cpu") + + empty_result = _sort_edge_tensors_by_destination( + empty_index, + empty_vec, + empty_mask, + empty_scatter, + ) + self.assertEqual(tuple(empty_result[0].shape), (2, 0)) + + edge_index = torch.tensor([[2, 0, 1, 0], [1, 0, 1, 0]], device="cpu") + edge_vec = torch.arange(12, dtype=torch.float32, device="cpu").view(4, 3) + edge_mask = torch.tensor([True, False, True, True], device="cpu") + edge_scatter = edge_index + 10 + sorted_index, sorted_vec, sorted_mask, sorted_scatter = ( + _sort_edge_tensors_by_destination( + edge_index, + edge_vec, + edge_mask, + edge_scatter, + ) + ) + permutation = torch.tensor([1, 3, 2, 0], device="cpu") + + self.assertTrue(torch.equal(sorted_index, edge_index[:, permutation])) + self.assertTrue(torch.equal(sorted_vec, edge_vec[permutation])) + self.assertTrue(torch.equal(sorted_mask, edge_mask[permutation])) + self.assertTrue(torch.equal(sorted_scatter, edge_scatter[:, permutation])) + @staticmethod def _randomize_params(model: torch.nn.Module, seed: int = 1234) -> None: """Fill all parameters with small random values. @@ -429,6 +463,111 @@ def _make_frame_with_natoms( ) return coord, atype, box + def test_neo_cute_flag_is_noop_for_ineligible_model(self) -> None: + """An unsupported SeZM shape must remain bit-identical with the flag on.""" + from deepmd.kernels.cute.neo import k1 as cute_k1 + + cpu = torch.device("cpu") + with ( + mock.patch.object(env, "DEVICE", cpu), + mock.patch.object(self, "device", cpu), + torch.device(cpu), + ): + coord, atype, _, _, _, _ = self._make_tiny_frame() + params = self._build_model_params(use_compile=False) + params["descriptor"]["seed"] = None + params["descriptor"]["use_env_seed"] = False + params["fitting_net"]["seed"] = None + model = get_sezm_model(params).to(device=cpu) + self._randomize_params(model) + model.eval() + for parameter in model.parameters(): + parameter.requires_grad_(False) + + extended_coord, extended_atype, mapping, nlist = ( + extend_input_and_build_neighbor_list( + coord, + atype, + model.get_rcut(), + model.get_sel(), + mixed_types=model.mixed_types(), + box=None, + ) + ) + edge = edge_schema_from_extended( + extended_coord, + extended_atype, + nlist, + mapping, + ) + + def run_lower() -> dict[str, torch.Tensor]: + return model.forward_common_lower( + edge.coord, + edge.atype, + edge.edge_index, + edge.edge_vec, + edge.edge_scatter_index, + edge.edge_mask, + ) + + with mock.patch.dict( + os.environ, + {"DP_NEO_CUTE_INFER": "0"}, + clear=False, + ): + reference = run_lower() + with ( + mock.patch.dict( + os.environ, + {"DP_NEO_CUTE_INFER": "1"}, + clear=False, + ), + mock.patch.object( + cute_k1, + "maybe_run_cute_k1", + wraps=cute_k1.maybe_run_cute_k1, + ) as maybe_run, + ): + actual = run_lower() + + self.assertGreater(maybe_run.call_count, 0) + self.assertEqual(actual.keys(), reference.keys()) + for name, expected in reference.items(): + if not isinstance(expected, torch.Tensor): + continue + self.assertTrue( + torch.equal(actual[name], expected), + msg=f"DP_NEO_CUTE_INFER changed ineligible-model {name}", + ) + + def test_neo_cute_sort_guard_uses_post_cast_geometry_dtype(self) -> None: + """Raw FP64 coordinates do not hide an otherwise eligible FP32 K1.""" + from deepmd.pt.model.model import ( + sezm_model, + ) + + descriptor = mock.Mock() + descriptor.compute_dtype = torch.float32 + descriptor._packed_wigner_candidate.return_value = True + device = torch.device("cuda") + + with mock.patch.object( + sezm_model, "_neo_cute_infer_enabled", return_value=True + ): + actual = sezm_model._neo_cute_k1_requires_sorted_edges( + descriptor, + training=False, + device=device, + ) + + self.assertTrue(actual) + descriptor._packed_wigner_candidate.assert_called_once_with( + device, + torch.float32, + torch.float32, + ) + def test_trace_pad_dim_trim_returns_contiguous(self) -> None: """Trimmed trace inputs stay contiguous so strides mirror runtime layout. @@ -590,6 +729,103 @@ def test_compile_cache_slots_and_eval_shape_change(self) -> None: model_cmp.compiled_core_compute_cache[eval_key], callable_eval_first ) + def test_load_state_dict_invalidates_local_and_shared_compile_caches(self) -> None: + """Checkpoint loads must discard every callable that captured old state.""" + from deepmd.pt.model.model import sezm_model as sezm_model_module + + model = get_sezm_model(self._build_model_params(use_compile=True)) + self.assertTrue( + all( + getattr(hook, "__self__", None) is None + for hook in model._load_state_dict_post_hooks.values() + ) + ) + local_callable = object() + shared_callable = object() + model.compiled_core_compute_cache[(False, False)] = local_callable + model._task_buf_order_cache[(False, False)] = ("task",) + object.__setattr__(model, "compiled_embedding", local_callable) + object.__setattr__(model, "_embedding_task_buf_order", ("task",)) + object.__setattr__(model, "compiled_dens_compute", local_callable) + model._dens_compiled = True + + with ( + mock.patch.dict( + sezm_model_module._SEZM_COMPILE_CACHE, + {("shared",): shared_callable}, + clear=True, + ), + mock.patch.dict( + sezm_model_module._SEZM_TASK_BUF_ORDER, + {("shared",): ("task",)}, + clear=True, + ), + ): + model.load_state_dict(model.state_dict()) + self.assertEqual(model.compiled_core_compute_cache, {}) + self.assertEqual(model._task_buf_order_cache, {}) + self.assertIsNone(model.compiled_embedding) + self.assertIsNone(model._embedding_task_buf_order) + self.assertIsNone(model.compiled_dens_compute) + self.assertFalse(model._dens_compiled) + self.assertEqual(sezm_model_module._SEZM_COMPILE_CACHE, {}) + self.assertEqual(sezm_model_module._SEZM_TASK_BUF_ORDER, {}) + + @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) + def test_eval_compile_retraces_after_load_state_dict(self) -> None: + """The public eval path must not reuse a make_fx graph after reloading.""" + coord, atype, box, _, _, _ = self._make_tiny_frame() + initial = get_sezm_model(self._build_model_params(use_compile=False)) + replacement = get_sezm_model(self._build_model_params(use_compile=False)) + self._randomize_params(initial, seed=1234) + self._randomize_params(replacement, seed=5678) + with mock.patch.dict(os.environ, {"DP_COMPILE_INFER": "1"}, clear=False): + compiled = get_sezm_model(self._build_model_params(use_compile=True)) + compiled.load_state_dict(initial.state_dict()) + initial.eval() + replacement.eval() + compiled.eval() + + # Exercise the public make_fx/AOT cache lifecycle without depending on + # Inductor's dynamic-View lowering for this intentionally tiny frame. + with mock.patch( + "torch._inductor.compile_fx.compile_fx_inner", + side_effect=lambda graph, inputs: graph.forward, + ): + before = compiled(coord, atype, box=box) + old_callable = compiled.compiled_core_compute_cache[(False, False)] + + compiled.load_state_dict(replacement.state_dict()) + self.assertEqual(compiled.compiled_core_compute_cache, {}) + expected = replacement(coord, atype, box=box) + actual = compiled(coord, atype, box=box) + self.assertIsNot( + compiled.compiled_core_compute_cache[(False, False)], + old_callable, + ) + self.assertFalse(torch.equal(actual["energy"], before["energy"])) + _assert_close_with_strict_warning( + actual["energy"], + expected["energy"], + atol=1.0e-6, + rtol=1.0e-6, + msg="eval energy mismatch after checkpoint reload", + ) + _assert_close_with_strict_warning( + actual["force"], + expected["force"], + atol=1.0e-6, + rtol=1.0e-6, + msg="eval force mismatch after checkpoint reload", + ) + _assert_close_with_strict_warning( + actual["virial"], + expected["virial"], + atol=1.0e-5, + rtol=1.0e-5, + msg="eval virial mismatch after checkpoint reload", + ) + @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) def test_charge_spin_condition_matches_compile(self) -> None: """Charge/spin conditions should work through the compiled energy path.""" @@ -732,6 +968,41 @@ def test_fixed_edge_geometry_matches_standard_cache(self) -> None: torch.testing.assert_close(cache_std.D_full, cache_sparse.D_full[:n_real]) torch.testing.assert_close(cache_std.Dt_full, cache_sparse.Dt_full[:n_real]) + from deepmd.pt.model.descriptor.sezm_nn import edge_cache as edge_cache_module + + with ( + mock.patch.dict(os.environ, {"DP_CUTE_INFER": "1"}, clear=False), + mock.patch.object( + edge_cache_module, + "_build_edge_wigner", + wraps=edge_cache_module._build_edge_wigner, + ) as build_edge_wigner, + ): + cache_sfpg = build_edge_cache_from_edges( + type_ebed=type_ebed, + atype_flat=atype_loc.reshape(-1), + edge_index=edge_index, + edge_vec=edge_vec, + edge_mask=edge_mask, + compute_dtype=descriptor.compute_dtype, + eps=descriptor.eps, + deg_norm_floor=descriptor.deg_norm_floor, + inner_clamp=descriptor.inner_clamp, + bridging_switch=torch.ones_like, + edge_envelope=descriptor.edge_envelope, + radial_basis=descriptor.radial_basis, + has_exclude_types=False, + edge_type_keep_mask=descriptor._edge_type_keep_mask, + random_gamma=False, + wigner_calc=descriptor.wigner_calc, + packed_wigner_candidate=True, + destinations_sorted=True, + ) + + self.assertFalse(build_edge_wigner.call_args.kwargs["packed_wigner"]) + self.assertEqual(cache_sfpg.D_full.dim(), 3) + self.assertIsNotNone(cache_sfpg.edge_src_gate) + @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON) def test_eval_compile_policy(self) -> None: """Eval should stay eager by default and compile only with env override.""" From b5c48d7f992851266c662c775b705b0f0160ebad Mon Sep 17 00:00:00 2001 From: atulcthakur Date: Wed, 19 Aug 2026 16:50:25 -0700 Subject: [PATCH 02/11] fix(pt): address Neo CuTe review findings --- deepmd/kernels/cute/neo/k1.py | 74 +++++++++-------- .../k1_kernels/cute_envelope_gated_softmax.py | 1 + ...neo_gate_linear_residual_backward_fused.py | 1 + .../cute_neo_phase_a_radial_forward.py | 38 ++++++++- .../cute_neo_phase_c_backward_layout.py | 10 --- ...cute_neo_phase_c_backward_layout_runner.py | 4 - .../k1_kernels/cute_neo_phase_c_onepass.py | 4 + .../cute_neo_so2_gate_combined_fwd.py | 3 +- .../cute/neo/k1_radial_phase_a_node.py | 7 +- deepmd/kernels/cute/neo/k1_runner.py | 15 +++- deepmd/kernels/cute/neo/k4_wignerd.py | 9 ++- .../cute/neo/message_grid_gaunt_sm90.py | 5 +- deepmd/kernels/cute/neo/readout_l0.py | 17 +++- .../kernels/cute/neo/sm90_k1/output_gate.py | 2 +- deepmd/kernels/cute/neo/sm90_k1/phase_a.py | 3 +- .../neo/sm90_k1/phase_c_attention_backward.py | 14 ++-- deepmd/kernels/cute/neo/sm90_k1/radial.py | 4 + deepmd/kernels/cute/neo/sm90_k1/runner.py | 1 - deepmd/pt/model/model/sezm_model.py | 28 +++++-- doc/install/easy-install.md | 3 +- pyproject.toml | 2 +- .../pt/model/test_descriptor_sezm_cute_gie.py | 26 +----- .../test_descriptor_sezm_cute_k1_dst_ptr.py | 6 +- ...descriptor_sezm_cute_k1_gate_structural.py | 27 +------ ...escriptor_sezm_cute_k1_inplace_residual.py | 29 +------ ...scriptor_sezm_cute_k1_radial_projection.py | 34 +------- ..._sezm_cute_k1_structural_gate_vec4_sm80.py | 81 +++++++++++++------ source/tests/pt/model/test_sezm_model.py | 2 + 28 files changed, 240 insertions(+), 210 deletions(-) diff --git a/deepmd/kernels/cute/neo/k1.py b/deepmd/kernels/cute/neo/k1.py index 279fd62ece..7e227c332a 100644 --- a/deepmd/kernels/cute/neo/k1.py +++ b/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,19 +524,21 @@ 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 @@ -543,7 +546,8 @@ 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)) diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_envelope_gated_softmax.py b/deepmd/kernels/cute/neo/k1_kernels/cute_envelope_gated_softmax.py index 8a6512689b..45a996e816 100644 --- a/deepmd/kernels/cute/neo/k1_kernels/cute_envelope_gated_softmax.py +++ b/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, ) diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py index 6c8658b887..27d356caa0 100644 --- a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py +++ b/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, ) diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_a_radial_forward.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_a_radial_forward.py index 0855d3b426..488806d314 100644 --- a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_a_radial_forward.py +++ b/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, diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout.py index e4a55bfd2f..43e3f7ef14 100644 --- a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout.py +++ b/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, diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout_runner.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout_runner.py index 5f016886b5..f33059e9be 100644 --- a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout_runner.py +++ b/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, diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_onepass.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_onepass.py index d366e3d1ab..ee489af8bc 100644 --- a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_onepass.py +++ b/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 diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_so2_gate_combined_fwd.py b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_so2_gate_combined_fwd.py index 4dec70c29a..9a1de48779 100644 --- a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_so2_gate_combined_fwd.py +++ b/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), diff --git a/deepmd/kernels/cute/neo/k1_radial_phase_a_node.py b/deepmd/kernels/cute/neo/k1_radial_phase_a_node.py index f9384fbdfa..ac968839ba 100644 --- a/deepmd/kernels/cute/neo/k1_radial_phase_a_node.py +++ b/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 diff --git a/deepmd/kernels/cute/neo/k1_runner.py b/deepmd/kernels/cute/neo/k1_runner.py index c74256b972..af793d5a5d 100644 --- a/deepmd/kernels/cute/neo/k1_runner.py +++ b/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, diff --git a/deepmd/kernels/cute/neo/k4_wignerd.py b/deepmd/kernels/cute/neo/k4_wignerd.py index 49ee535200..b02e66fa52 100644 --- a/deepmd/kernels/cute/neo/k4_wignerd.py +++ b/deepmd/kernels/cute/neo/k4_wignerd.py @@ -8,6 +8,7 @@ annotations, ) +import threading from dataclasses import ( dataclass, ) @@ -796,6 +797,7 @@ def _cached_wignerd_panel_backward(threads: int) -> Callable: _TABLE_CACHE: dict[tuple[str, int], tuple[torch.Tensor, ...]] = {} +_TABLE_CACHE_LOCK = threading.Lock() def _device_cache_key(device: torch.device) -> tuple[str, int]: @@ -807,8 +809,11 @@ def _get_lmax3_tables(device: torch.device) -> tuple[torch.Tensor, ...]: key = _device_cache_key(device) tables = _TABLE_CACHE.get(key) if tables is None: - tables = _build_l2_l3_tables(torch.float32, device) - _TABLE_CACHE[key] = tables + with _TABLE_CACHE_LOCK: + tables = _TABLE_CACHE.get(key) + if tables is None: + tables = _build_l2_l3_tables(torch.float32, device) + _TABLE_CACHE[key] = tables return tables diff --git a/deepmd/kernels/cute/neo/message_grid_gaunt_sm90.py b/deepmd/kernels/cute/neo/message_grid_gaunt_sm90.py index 169279f209..2bbc087ac8 100644 --- a/deepmd/kernels/cute/neo/message_grid_gaunt_sm90.py +++ b/deepmd/kernels/cute/neo/message_grid_gaunt_sm90.py @@ -218,9 +218,10 @@ def kernel( linear = tidx + slot * THREADS row = linear >> 6 scalar_channel = linear & (CHANNELS - 1) - value = left[node, row, scalar_channel].to(cutlass.Float32) if cutlass.const_expr(side == 1): value = right[node, row, scalar_channel].to(cutlass.Float32) + else: + value = left[node, row, scalar_channel].to(cutlass.Float32) operands[side, row, scalar_channel] = value cute.arch.sync_threads() @@ -309,8 +310,8 @@ def run_sm90_gaunt_forward( schedule: Sm90GauntSchedule, ) -> torch.Tensor: """Evaluate the exact normalized-Gaunt product on SM90.""" - _validate_dense(left, right) output = torch.empty_like(left) + _validate_dense(left, right, output) with torch.cuda.device(left.device): _compiled_sm90_gaunt_forward(schedule)(left, right, output) return output diff --git a/deepmd/kernels/cute/neo/readout_l0.py b/deepmd/kernels/cute/neo/readout_l0.py index db18740aca..0d51e6c1fa 100644 --- a/deepmd/kernels/cute/neo/readout_l0.py +++ b/deepmd/kernels/cute/neo/readout_l0.py @@ -284,7 +284,9 @@ def _readout_input_fold_cache_matches( cache is not None and len(cache) == 2 and len(cache[0]) == len(sources) - and all(cached is current for cached, current in zip(cache[0], sources)) + and all( + cached is current for cached, current in zip(cache[0], sources, strict=True) + ) and cache[1] == cache_key and all( isinstance(getattr(output_ffn, name, None), torch.Tensor) @@ -437,7 +439,11 @@ def prepare_sm80_readout_input_fold( folded_weights = _build_sm80_readout_input_fold(output_ffn) _synchronize_sm80_readout_input_fold_build(folded_weights) _ensure_sm80_readout_input_fold_load_hook(output_ffn) - for name, tensor in zip(_READOUT_INPUT_FOLD_BUFFERS, folded_weights): + for name, tensor in zip( + _READOUT_INPUT_FOLD_BUFFERS, + folded_weights, + strict=True, + ): _set_nonpersistent_buffer(output_ffn, name, tensor) setattr( output_ffn, @@ -496,12 +502,15 @@ def _get_prepared_sm80_readout_input_fold( sources = _readout_input_fold_sources(output_ffn) cached_sources, cache_key = cache if ( - any(cached is not current for cached, current in zip(cached_sources, sources)) + any( + cached is not current + for cached, current in zip(cached_sources, sources, strict=True) + ) or any( source.dtype != cached[2] or source.device != cached[3] or tuple(source.shape) != cached[4] - for source, cached in zip(sources, cache_key) + for source, cached in zip(sources, cache_key, strict=True) ) or any( weight.dtype != torch.float32 diff --git a/deepmd/kernels/cute/neo/sm90_k1/output_gate.py b/deepmd/kernels/cute/neo/sm90_k1/output_gate.py index 480e8217e2..5ba25e3f41 100644 --- a/deepmd/kernels/cute/neo/sm90_k1/output_gate.py +++ b/deepmd/kernels/cute/neo/sm90_k1/output_gate.py @@ -144,7 +144,7 @@ def _require_tensor( if tuple(tensor.shape) != shape: raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}") if tensor.dtype != torch.float32 or tensor.device != device: - raise ValueError(f"{name} must be contiguous FP32 on {device}") + raise ValueError(f"{name} must be FP32 on {device}") if not tensor.is_contiguous() or tensor.data_ptr() % 16: raise ValueError(f"{name} must be contiguous and 16-byte aligned") diff --git a/deepmd/kernels/cute/neo/sm90_k1/phase_a.py b/deepmd/kernels/cute/neo/sm90_k1/phase_a.py index 476e86cdd2..e0f3b481a7 100644 --- a/deepmd/kernels/cute/neo/sm90_k1/phase_a.py +++ b/deepmd/kernels/cute/neo/sm90_k1/phase_a.py @@ -75,6 +75,7 @@ ) FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} +DEFAULT_STREAM = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT) __all__ = [ "CTA_SHARED_BYTES", @@ -172,7 +173,7 @@ def __call__( channel_basis, m0, m1_ri, - stream: cuda.CUstream = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT), + stream: cuda.CUstream = DEFAULT_STREAM, ): d_layout = cute.make_layout( (EDGE_TILE, PACKED_VALUE_COUNT), diff --git a/deepmd/kernels/cute/neo/sm90_k1/phase_c_attention_backward.py b/deepmd/kernels/cute/neo/sm90_k1/phase_c_attention_backward.py index b51f0f40af..f1f0430357 100644 --- a/deepmd/kernels/cute/neo/sm90_k1/phase_c_attention_backward.py +++ b/deepmd/kernels/cute/neo/sm90_k1/phase_c_attention_backward.py @@ -681,7 +681,6 @@ def run_grouped_expanded_final_phase_c_attention_adjoint( focus_eps: float, focus_tau: float, label_smoothing: float, - max_degree: int, qk_scale: float = QK_SCALE, outputs: GroupedExpandedFinalPhaseCAttentionAdjointOutputs | None = None, ) -> GroupedExpandedFinalPhaseCAttentionAdjointOutputs: @@ -696,10 +695,6 @@ def run_grouped_expanded_final_phase_c_attention_adjoint( raise RuntimeError("strict FP32 requires allow_tf32=False") if torch.get_float32_matmul_precision() != "highest": raise RuntimeError("strict FP32 requires float32 matmul precision 'highest'") - if max_degree < 0 or max_degree > MAX_EDGES_PER_NODE: - raise ValueError( - f"max_degree must be in [0,{MAX_EDGES_PER_NODE}], got {max_degree}" - ) node_count = int(dst_ptr.numel() - 1) edge_count = int(src.numel()) @@ -758,6 +753,15 @@ def run_grouped_expanded_final_phase_c_attention_adjoint( ) for name, tensor, shape, dtype in expected_inputs: _require_tensor(name, tensor, shape, dtype, device) + if dst_ptr.numel() > 1: + destination_degrees = dst_ptr[1:] - dst_ptr[:-1] + torch._assert_async( + torch.all( + (destination_degrees >= 0) & (destination_degrees <= MAX_EDGES_PER_NODE) + ), + "SM90 Phase-C backward requires destination degrees in " + f"[0, {MAX_EDGES_PER_NODE}]", + ) if outputs is None: outputs = allocate_grouped_expanded_final_phase_c_attention_adjoint_outputs( diff --git a/deepmd/kernels/cute/neo/sm90_k1/radial.py b/deepmd/kernels/cute/neo/sm90_k1/radial.py index 7e71024db7..cb7cb017f7 100644 --- a/deepmd/kernels/cute/neo/sm90_k1/radial.py +++ b/deepmd/kernels/cute/neo/sm90_k1/radial.py @@ -263,6 +263,10 @@ def project_neo_radial_input_adjoint_fp32( is introduced: the first matrix product is the compact radial adjoint and the degree-zero slice receives the independent attention-logit adjoint. """ + 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 = grad_compact.shape[0] device = grad_compact.device _expect_fp32_cuda( diff --git a/deepmd/kernels/cute/neo/sm90_k1/runner.py b/deepmd/kernels/cute/neo/sm90_k1/runner.py index ce4decb20e..1f8f37339e 100644 --- a/deepmd/kernels/cute/neo/sm90_k1/runner.py +++ b/deepmd/kernels/cute/neo/sm90_k1/runner.py @@ -600,7 +600,6 @@ def _runner_backward( focus_eps=float(so2.focus_compete_norm.eps), focus_tau=float(so2.focus_softmax_tau), label_smoothing=float(so2.focus_label_smoothing), - max_degree=128, qk_scale=CHANNELS**-0.5, outputs=fused_outputs, ) diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index b6a918cc3e..eccaae84b2 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -352,7 +352,8 @@ NOTE 10 -- Tail dummy edges --------------------------- -The edge-schema builders append two masked edges at the end of every batch. +The edge-schema builders append two masked edges to every batch. They remain +trailing unless a later destination sort permutes them with the real edges. Real edge compaction happens via ``torch.nonzero(valid_mask)``, whose output length is data-dependent and can be zero in sparse or single-atom systems (e.g. isolated-atom @@ -362,9 +363,9 @@ ``dynamic=True``. A pair of dummy slots also gives Inductor's batched matmul lowering a static ``E >= 2`` edge-axis bound, avoiding data-dependent layout guards on ``E == 1`` that would otherwise cause -an extra recompile when the first batch contains no real edges. Each -dummy's ``edge_mask`` is ``False`` so it contributes exactly zero to -every downstream sum or gather. +an extra recompile when the first batch contains no real edges. Each dummy's +``edge_mask`` is ``False`` wherever sorting places it, so it contributes +exactly zero to every downstream sum or gather. NOTE 11 -- Edge-vector leaf (gather outside the AD region) ---------------------------------------------------------- @@ -918,6 +919,20 @@ def _invalidate_compiled_state_after_load( module._core_compute_pending_compile_t0 = None module._core_compute_pending_compile_key = None module._dens_pending_compile_t0 = None + k1_invalidator = None + for child in module.modules(): + if not ( + hasattr(child, "_deepmd_cute_k1_state") + or hasattr(child, "_deepmd_cute_gate_expand_contract") + ): + continue + if k1_invalidator is None: + from deepmd.kernels.cute.neo.k1 import ( + invalidate_cute_k1_state, + ) + + k1_invalidator = invalidate_cute_k1_state + k1_invalidator(child) # Shared callables can contain make_fx get_attr constants, including # prepared CuTe readout folds. A checkpoint load therefore invalidates # the process-level cache as well as this instance's local slots. @@ -3076,7 +3091,10 @@ def build_edge_list_from_nlist( edge_vec Edge vectors with shape (E+2, 3). edge_mask - Boolean mask with shape (E+2,). The two trailing elements are ``False``. + Boolean mask with shape (E+2,). The two padded elements are + ``False``. They are trailing only when sorting is disabled; + destination sorting may move them into the interior, so consumers + must select valid edges through this mask rather than by position. edge_scatter_index Scatter-domain (src, dst) indices with shape (2, E+2), aligned 1:1 with ``edge_index`` and ``edge_vec``. diff --git a/doc/install/easy-install.md b/doc/install/easy-install.md index 69b0866cfe..a77d2b35d7 100644 --- a/doc/install/easy-install.md +++ b/doc/install/easy-install.md @@ -189,7 +189,8 @@ pip install deepmd-kit :::::: -Optional CuTe inference kernels are available on Linux for the PyTorch backend: +Optional CuTe inference kernels are available on Linux with Python 3.11 or +newer for the PyTorch backend: ```bash pip install "deepmd-kit[torch,cute]" diff --git a/pyproject.toml b/pyproject.toml index dc4e9228dd..ec79b7ca64 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -146,7 +146,7 @@ cu12 = [ "nvidia-cuda-nvcc-cu12", ] cute = [ - 'nvidia-cutlass-dsl[cu13]>=4.6.1,<5; platform_system == "Linux"', + 'nvidia-cutlass-dsl[cu13]>=4.6.1,<5; python_version >= "3.11" and platform_system == "Linux"', 'nvalchemi-toolkit-ops>=0.3.1; python_version >= "3.11" and platform_system == "Linux"', ] jax = [ diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_gie.py b/source/tests/pt/model/test_descriptor_sezm_cute_gie.py index 4deb334291..c03da7bcc9 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_gie.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_gie.py @@ -6,40 +6,18 @@ annotations, ) -import importlib.util -import sys import unittest -from functools import ( - lru_cache, -) -from pathlib import ( - Path, -) from types import ( SimpleNamespace, ) import torch -REPO_ROOT = Path(__file__).resolve().parents[4] -GIE_PATH = REPO_ROOT / "deepmd/kernels/cute/neo/gie.py" +from deepmd.kernels.cute.neo import gie as gie_module -@lru_cache(maxsize=1) def _load_gie_module(): - if not GIE_PATH.is_file(): - raise FileNotFoundError(GIE_PATH) - name = "sezm_cute_gie_contract_test" - spec = importlib.util.spec_from_file_location(name, GIE_PATH) - if spec is None or spec.loader is None: - raise RuntimeError(f"could not load {GIE_PATH}") - module = importlib.util.module_from_spec(spec) - sys.modules[name] = module - try: - spec.loader.exec_module(module) - finally: - sys.modules.pop(name, None) - return module + return gie_module def _degree_slots(lmax: int, *, device: torch.device) -> torch.Tensor: diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_dst_ptr.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_dst_ptr.py index ef12c702f2..662cb29ffa 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_dst_ptr.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_dst_ptr.py @@ -267,9 +267,6 @@ def test_compile_cold_registration_survives_into_packed_custom_op_runtime(self): k1 = importlib.import_module("deepmd.kernels.cute.neo.k1") prior_registry = dict(k1._REGISTRY) prior_next_handle = k1._NEXT_HANDLE - k1._REGISTRY.clear() - k1._PACKED_RUNNER_CACHE.clear() - torch._dynamo.reset() def cleanup(): torch._dynamo.reset() @@ -279,6 +276,9 @@ def cleanup(): k1._NEXT_HANDLE = prior_next_handle self.addCleanup(cleanup) + k1._REGISTRY.clear() + k1._PACKED_RUNNER_CACHE.clear() + torch._dynamo.reset() block = torch.nn.Module() config = k1.NeoK1RuntimeConfig() runtime_handles = [] diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py index 3394bf1356..de39840bd7 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py @@ -6,38 +6,19 @@ annotations, ) -import importlib.util -import sys import unittest -from pathlib import ( - Path, -) import torch -REPO_ROOT = Path(__file__).resolve().parents[4] -HELPER_PATH = REPO_ROOT / "deepmd/kernels/cute/neo/k1_gate_structural.py" - - -def _load_module(path: Path, name: str): - spec = importlib.util.spec_from_file_location(name, path) - if spec is None or spec.loader is None: - raise RuntimeError(f"could not load {path}") - module = importlib.util.module_from_spec(spec) - sys.modules[name] = module - try: - spec.loader.exec_module(module) - finally: - sys.modules.pop(name, None) - return module +from deepmd.kernels.cute.neo import ( + k1_gate_structural, +) class TestStructuralGateHelpers(unittest.TestCase): @classmethod def setUpClass(cls): - if not HELPER_PATH.is_file(): - raise AssertionError(f"structural gate helper is missing: {HELPER_PATH}") - cls.helper = _load_module(HELPER_PATH, "structural_gate_helper_test") + cls.helper = k1_gate_structural def setUp(self): torch.manual_seed(20260703) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py index a977b4ce6f..1e348b9018 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py @@ -6,41 +6,20 @@ annotations, ) -import importlib.util -import sys import unittest -from functools import ( - lru_cache, -) -from pathlib import ( - Path, -) import torch from torch.utils._python_dispatch import ( TorchDispatchMode, ) -REPO_ROOT = Path(__file__).resolve().parents[4] -IMPLEMENTATION_PATH = REPO_ROOT / "deepmd/kernels/cute/neo/k1_so2linear.py" +from deepmd.kernels.cute.neo import ( + k1_so2linear, +) -@lru_cache(maxsize=1) def _load_candidate(): - assert IMPLEMENTATION_PATH.is_file(), ( - f"in-place residual module is missing: {IMPLEMENTATION_PATH}" - ) - name = "sezm_cute_k1_inplace_residual_test" - spec = importlib.util.spec_from_file_location(name, IMPLEMENTATION_PATH) - if spec is None or spec.loader is None: - raise RuntimeError(f"could not load {IMPLEMENTATION_PATH}") - module = importlib.util.module_from_spec(spec) - sys.modules[name] = module - try: - spec.loader.exec_module(module) - finally: - sys.modules.pop(name, None) - return module + return k1_so2linear def _out_of_place_reference(grad_out, residual, w0_t, wpair_t): diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py index 3c2cad637a..6d382e030d 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py @@ -6,14 +6,8 @@ annotations, ) -import importlib.util -import sys import unittest -from pathlib import ( - Path, -) from types import ( - ModuleType, SimpleNamespace, ) from unittest import ( @@ -22,32 +16,8 @@ import torch -REPO_ROOT = Path(__file__).resolve().parents[4] -CUTE_PATH = REPO_ROOT / "deepmd/kernels/cute/neo" -_TEST_PACKAGE = "sezm_cute_k1_radial_projection_contract_test" - - -def _load_cute_source_module(name: str) -> ModuleType: - qualified_name = f"{_TEST_PACKAGE}.{name}" - spec = importlib.util.spec_from_file_location( - qualified_name, - CUTE_PATH / f"{name}.py", - ) - if spec is None or spec.loader is None: - raise RuntimeError(f"could not load {name}") - module = importlib.util.module_from_spec(spec) - sys.modules[qualified_name] = module - spec.loader.exec_module(module) - return module - - -_package = ModuleType(_TEST_PACKAGE) -_package.__path__ = [str(CUTE_PATH)] -sys.modules[_TEST_PACKAGE] = _package -_load_cute_source_module("compile_cache") -_RADIAL_PROJECTION = _load_cute_source_module("k1_radial_phase_a_node") -_load_cute_source_module("runtime_policy") -_K1_RUNNER = _load_cute_source_module("k1_runner") +from deepmd.kernels.cute.neo import k1_radial_phase_a_node as _RADIAL_PROJECTION +from deepmd.kernels.cute.neo import k1_runner as _K1_RUNNER COMPACT_WIDTH = _RADIAL_PROJECTION.COMPACT_WIDTH FOCUS_COUNT = _RADIAL_PROJECTION.FOCUS_COUNT diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py index 8bed9d3942..92d153b090 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py @@ -7,7 +7,7 @@ ) import ast -import importlib.util +import importlib import math import random import struct @@ -15,20 +15,10 @@ from pathlib import ( Path, ) -from typing import ( - TYPE_CHECKING, -) - -if TYPE_CHECKING: - from types import ( - ModuleType, - ) - REPO_ROOT = Path(__file__).resolve().parents[4] CUTE_ROOT = REPO_ROOT / "deepmd/kernels/cute/neo" KERNEL_PATH = CUTE_ROOT / "k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py" -WRAPPER_PATH = CUTE_ROOT / "k1_gate_structural.py" DYNAMIC_EDGE_COUNTS = (0, 1, 7, 8, 9, 31, 32, 33) @@ -155,6 +145,51 @@ def _backward_reference( return grad_y, grad_logits +def _vec4_backward_model( + grad_out: list[float], + y: list[float], + logits: list[float], + edge_count: int, +) -> tuple[list[float], list[float]]: + grad_y = [0.0] * len(y) + grad_logits = [0.0] * len(logits) + rows = edge_count * 2 + for block_row in range((rows + 15) // 16): + for row_slot in range(16): + row = block_row * 16 + row_slot + if row >= rows: + continue + edge, focus = divmod(row, 2) + row_base = row * 10 * 32 + logits_base = (focus * edge_count + edge) * 3 * 32 + for channel_group in range(8): + channel_base = channel_group * 4 + for lane in range(4): + channel = channel_base + lane + y0 = _f32(y[row_base + channel]) + sig0 = _sigmoid_f32(y0) + grad0 = _f32(grad_out[row_base + channel]) + inner = _f32(_f32(1.0) + _f32(y0 * _f32(_f32(1.0) - sig0))) + grad_y[row_base + channel] = _f32(_f32(grad0 * sig0) * inner) + for gate_index in range(3): + for lane in range(4): + channel = channel_base + lane + gate_offset = logits_base + gate_index * 32 + channel + gate = _sigmoid_f32(logits[gate_offset]) + grad_logit = _f32(0.0) + for repeat in range(3): + degree = 1 + gate_index + repeat * 3 + index = row_base + degree * 32 + channel + gout = _f32(grad_out[index]) + grad_y[index] = _f32(gout * gate) + term = _f32(gout * _f32(y[index])) + term = _f32(term * gate) + term = _f32(term * _f32(_f32(1.0) - gate)) + grad_logit = _f32(grad_logit + term) + grad_logits[gate_offset] = grad_logit + return grad_y, grad_logits + + def _module_constants(tree: ast.Module) -> dict[str, object]: def constant_value(node: ast.expr, namespace: dict[str, object]) -> object: if isinstance(node, ast.Constant): @@ -187,16 +222,8 @@ def constant_value(node: ast.expr, namespace: dict[str, object]) -> object: return constants -def _load_wrapper_module() -> ModuleType: - spec = importlib.util.spec_from_file_location( - "_deepmd_test_k1_gate_structural", - WRAPPER_PATH, - ) - if spec is None or spec.loader is None: - raise RuntimeError(f"cannot load {WRAPPER_PATH}") - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module +def _load_wrapper_module(): + return importlib.import_module("deepmd.kernels.cute.neo.k1_gate_structural") class TestSM80StructuralGateVec4Contract(unittest.TestCase): @@ -257,14 +284,20 @@ def test_backward_dynamic_edges_and_block_tails(self): logits = [ _f32(rng.uniform(-4.0, 4.0)) for _ in range(2 * edge_count * 3 * 32) ] - grad_y, grad_logits = _backward_reference( + expected_grad_y, expected_grad_logits = _backward_reference( + grad_out, + y, + logits, + edge_count, + ) + actual_grad_y, actual_grad_logits = _vec4_backward_model( grad_out, y, logits, edge_count, ) - self.assertEqual(len(grad_y), values) - self.assertEqual(len(grad_logits), len(logits)) + self.assertEqual(actual_grad_y, expected_grad_y) + self.assertEqual(actual_grad_logits, expected_grad_logits) try: diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index cd5f6a65fd..7f5cc5c7ad 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -748,6 +748,7 @@ def test_load_state_dict_invalidates_local_and_shared_compile_caches(self) -> No object.__setattr__(model, "_embedding_task_buf_order", ("task",)) object.__setattr__(model, "compiled_dens_compute", local_callable) model._dens_compiled = True + model._deepmd_cute_k1_state = False with ( mock.patch.dict( @@ -768,6 +769,7 @@ def test_load_state_dict_invalidates_local_and_shared_compile_caches(self) -> No self.assertIsNone(model._embedding_task_buf_order) self.assertIsNone(model.compiled_dens_compute) self.assertFalse(model._dens_compiled) + self.assertFalse(hasattr(model, "_deepmd_cute_k1_state")) self.assertEqual(sezm_model_module._SEZM_COMPILE_CACHE, {}) self.assertEqual(sezm_model_module._SEZM_TASK_BUF_ORDER, {}) From e93ba332df2a25a196a4c6f8d5b5c8e92e7a2ece Mon Sep 17 00:00:00 2001 From: OutisLi Date: Wed, 2 Sep 2026 23:01:13 +0800 Subject: [PATCH 03/11] perf(pt): replace SeZM CuTe inference kernels --- deepmd/dpmodel/descriptor/dpa4_nn/so2.py | 7 - deepmd/kernels/cute/neo/__init__.py | 3 - deepmd/pt/model/descriptor/sezm.py | 6 +- deepmd/pt/model/descriptor/sezm_nn/block.py | 8 +- .../pt/model/descriptor/sezm_nn/edge_cache.py | 6 +- .../pt/model/descriptor/sezm_nn/embedding.py | 10 +- .../pt/model/descriptor/sezm_nn/grid_net.py | 4 +- deepmd/pt/model/descriptor/sezm_nn/lora.py | 8 +- deepmd/pt/model/descriptor/sezm_nn/so2.py | 47 +- deepmd/pt/model/model/sezm_model.py | 14 +- deepmd/pt/model/network/mlp.py | 2 +- deepmd/pt_expt/descriptor/dpa4_nn/so2.py | 44 +- deepmd/pt_expt/kernels/cute/sezm/__init__.py | 55 +- deepmd/pt_expt/kernels/cute/sezm/backward.py | 692 ------------------ .../kernels/cute/sezm}/compile_cache.py | 0 deepmd/pt_expt/kernels/cute/sezm/forward.py | 444 ----------- .../neo => pt_expt/kernels/cute/sezm}/gie.py | 2 +- .../neo => pt_expt/kernels/cute/sezm}/k1.py | 0 .../kernels/cute/sezm}/k1_gate_structural.py | 0 .../kernels/cute/sezm}/k1_kernels/__init__.py | 0 .../k1_kernels/cute_envelope_gated_softmax.py | 0 .../k1_kernels/cute_neo_focus_src_backward.py | 0 ...neo_gate_linear_residual_backward_fused.py | 0 ...ute_neo_gate_split_structural_vec4_sm80.py | 0 .../cute_neo_message_grid_product.py | 0 .../cute_neo_output_gate_backward.py | 0 .../cute_neo_phase_a_radial_forward.py | 0 .../cute_neo_phase_c_backward_layout.py | 0 ...cute_neo_phase_c_backward_layout_runner.py | 0 .../k1_kernels/cute_neo_phase_c_onepass.py | 0 .../cute/sezm}/k1_kernels/cute_neo_qk_edge.py | 0 .../cute_neo_radial_phase_a_backward_node.py | 0 .../cute_neo_so2_gate_combined_fwd.py | 0 .../cute/sezm}/k1_message_grid_packed.py | 0 .../cute/sezm}/k1_radial_phase_a_node.py | 0 .../kernels/cute/sezm}/k1_runner.py | 0 .../kernels/cute/sezm}/k1_so2linear.py | 0 .../kernels/cute/sezm}/k1_wigner_layout.py | 0 .../kernels/cute/sezm}/k4_wignerd.py | 0 .../cute/sezm}/message_grid_gaunt_sm90.py | 0 .../cute/sezm}/message_grid_readout_sm90.py | 0 deepmd/pt_expt/kernels/cute/sezm/operator.py | 324 -------- .../sezm}/output_grid_kernels/__init__.py | 0 .../output_grid_kernels/cute_readout_l0.py | 0 .../cute_tiled_grid_product.py | 0 .../kernels/cute/sezm}/output_grid_product.py | 0 .../kernels/cute/sezm}/readout_l0.py | 2 +- .../kernels/cute/sezm}/runtime_policy.py | 14 +- .../kernels/cute/sezm}/sm90_k1/__init__.py | 0 .../cute/sezm}/sm90_k1/final_phase_c.py | 0 .../kernels/cute/sezm}/sm90_k1/output_gate.py | 0 .../kernels/cute/sezm}/sm90_k1/persistent.py | 0 .../kernels/cute/sezm}/sm90_k1/phase_a.py | 0 .../cute/sezm}/sm90_k1/phase_a_backward.py | 0 .../sm90_k1/phase_c_attention_backward.py | 0 .../kernels/cute/sezm}/sm90_k1/prefix.py | 0 .../kernels/cute/sezm}/sm90_k1/radial.py | 0 .../kernels/cute/sezm}/sm90_k1/runner.py | 0 deepmd/pt_expt/kernels/utils.py | 17 +- doc/install/easy-install.md | 12 +- doc/model/dpa4.md | 22 +- ...t_descriptor_sezm_block_ffn_grid_fusion.py | 8 +- ...test_descriptor_sezm_cute_compile_cache.py | 2 +- ...t_descriptor_sezm_cute_envelope_softmax.py | 2 +- .../pt/model/test_descriptor_sezm_cute_gie.py | 2 +- .../pt/model/test_descriptor_sezm_cute_k1.py | 24 +- .../test_descriptor_sezm_cute_k1_dst_ptr.py | 4 +- ...descriptor_sezm_cute_k1_gate_structural.py | 2 +- ...escriptor_sezm_cute_k1_inplace_residual.py | 2 +- ...descriptor_sezm_cute_k1_manual_adjoints.py | 12 +- ...est_descriptor_sezm_cute_k1_packed_impl.py | 10 +- ...or_sezm_cute_k1_phase_c_inplace_adjoint.py | 2 +- ...riptor_sezm_cute_k1_radial_focus_fusion.py | 4 +- ...riptor_sezm_cute_k1_radial_pitch68_safe.py | 2 +- ...scriptor_sezm_cute_k1_radial_projection.py | 6 +- ..._sezm_cute_k1_structural_gate_vec4_sm80.py | 8 +- ...or_sezm_cute_k1_structural_memory_reuse.py | 4 +- ...escriptor_sezm_cute_output_grid_product.py | 26 +- .../test_descriptor_sezm_cute_readout_l0.py | 76 +- ...est_descriptor_sezm_cute_runtime_policy.py | 40 +- ...descriptor_sezm_cute_tiled_grid_product.py | 20 +- ...test_descriptor_sezm_readout_l0_algebra.py | 4 +- source/tests/pt/model/test_mlp.py | 8 +- source/tests/pt/model/test_sezm_cute_efs.py | 11 +- .../model/test_sezm_cute_multigpu_dispatch.py | 14 +- source/tests/pt/model/test_sezm_model.py | 8 +- 86 files changed, 248 insertions(+), 1796 deletions(-) delete mode 100644 deepmd/kernels/cute/neo/__init__.py delete mode 100644 deepmd/pt_expt/kernels/cute/sezm/backward.py rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/compile_cache.py (100%) delete mode 100644 deepmd/pt_expt/kernels/cute/sezm/forward.py rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/gie.py (99%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_gate_structural.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_kernels/__init__.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_kernels/cute_envelope_gated_softmax.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_kernels/cute_neo_focus_src_backward.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_kernels/cute_neo_message_grid_product.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_kernels/cute_neo_output_gate_backward.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_kernels/cute_neo_phase_a_radial_forward.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_kernels/cute_neo_phase_c_backward_layout.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_kernels/cute_neo_phase_c_backward_layout_runner.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_kernels/cute_neo_phase_c_onepass.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_kernels/cute_neo_qk_edge.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_kernels/cute_neo_radial_phase_a_backward_node.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_kernels/cute_neo_so2_gate_combined_fwd.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_message_grid_packed.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_radial_phase_a_node.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_runner.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_so2linear.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k1_wigner_layout.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/k4_wignerd.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/message_grid_gaunt_sm90.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/message_grid_readout_sm90.py (100%) delete mode 100644 deepmd/pt_expt/kernels/cute/sezm/operator.py rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/output_grid_kernels/__init__.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/output_grid_kernels/cute_readout_l0.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/output_grid_kernels/cute_tiled_grid_product.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/output_grid_product.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/readout_l0.py (99%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/runtime_policy.py (96%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/sm90_k1/__init__.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/sm90_k1/final_phase_c.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/sm90_k1/output_gate.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/sm90_k1/persistent.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/sm90_k1/phase_a.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/sm90_k1/phase_a_backward.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/sm90_k1/phase_c_attention_backward.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/sm90_k1/prefix.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/sm90_k1/radial.py (100%) rename deepmd/{kernels/cute/neo => pt_expt/kernels/cute/sezm}/sm90_k1/runner.py (100%) diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py index 323fe3bbfe..a594ad9d09 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py @@ -1616,7 +1616,6 @@ def __init__( # value path: every hook stays ``None`` and ``so2_message`` takes the # dense branch. The ``pt_expt`` backend binds the selected implementation. self._triton_value_path = None - self._cute_value_path = None self._cutile_value_path = None # === Step 14. Optional fused training seams === @@ -2375,12 +2374,6 @@ def so2_message( if self._cached_edge_csr_fn is not None: self._cached_edge_csr_fn(edge_cache, "src", x.shape[0]) x_local, rad_feat = self._triton_value_path(x, edge_cache, radial_feat) - elif self._cute_value_path is not None and not training: - # === Steps 1-5 (fused CuTe operator). The operator folds - # rotate_to_local, radial degree mixing, the multi-layer gated SO(2) - # stack, and the focus competition into the bucketed kernels; the - # per-edge focus-major intermediates stay resident on chip. === - x_local, rad_feat = self._cute_value_path(x, edge_cache, radial_feat) else: # === Steps 1-3. Rotation, radial mixing and the focus-major cast === x_local, rad_feat = self._rotate_mix(x, edge_cache, radial_feat) diff --git a/deepmd/kernels/cute/neo/__init__.py b/deepmd/kernels/cute/neo/__init__.py deleted file mode 100644 index e98fa6b6db..0000000000 --- a/deepmd/kernels/cute/neo/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: LGPL-3.0-or-later -"""Neo-specialized CuTe kernels and PyTorch integration.""" diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index c8cb6694d4..70bf750827 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -52,7 +52,6 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) -from deepmd.kernels.cute.neo import runtime_policy as cute_runtime_policy from deepmd.pt.utils import ( env, ) @@ -66,6 +65,7 @@ from deepmd.pt.utils.update_sel import ( UpdateSel, ) +from deepmd.pt_expt.kernels.cute.sezm import runtime_policy as cute_runtime_policy from deepmd.pt_expt.kernels.utils import ( cuda_infer_level, use_amp_infer, @@ -1776,7 +1776,7 @@ def forward_with_edges( @torch.jit.unused def _run_output_readout(self, ffn_in: torch.Tensor) -> torch.Tensor: """Return the residual-inclusive scalar output with guarded CuTe routing.""" - from deepmd.kernels.cute.neo.readout_l0 import ( + from deepmd.pt_expt.kernels.cute.sezm.readout_l0 import ( run_neo_output_readout, ) @@ -1953,7 +1953,7 @@ def _packed_wigner_candidate( geometry_dtype: torch.dtype, ) -> bool: """Return whether all blocks satisfy packed K1's stable contract.""" - from deepmd.kernels.cute.neo.k1 import ( + from deepmd.pt_expt.kernels.cute.sezm.k1 import ( is_packed_wigner_candidate, ) diff --git a/deepmd/pt/model/descriptor/sezm_nn/block.py b/deepmd/pt/model/descriptor/sezm_nn/block.py index ca439c6fe0..7ef5a97f3d 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/block.py +++ b/deepmd/pt/model/descriptor/sezm_nn/block.py @@ -26,9 +26,6 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) -from deepmd.kernels.cute.neo.runtime_policy import ( - is_cute_infer_enabled, -) from deepmd.pt.utils import ( env, ) @@ -36,6 +33,9 @@ PRECISION_DICT, RESERVED_PRECISION_DICT, ) +from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import ( + is_cute_infer_enabled, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -834,7 +834,7 @@ def _run_so2_unit_impl( ) -> torch.Tensor: """Run the SO(2) unit implementation.""" if not self.training and is_cute_infer_enabled(): - from deepmd.kernels.cute.neo.k1 import ( + from deepmd.pt_expt.kernels.cute.sezm.k1 import ( maybe_run_cute_k1, ) diff --git a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py index 3d94a57fb9..2e21b1b8be 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py +++ b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py @@ -25,7 +25,7 @@ rearrange, ) -from deepmd.kernels.cute.neo.runtime_policy import ( +from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import ( is_cute_infer_enabled, ) @@ -795,7 +795,7 @@ def _build_edge_wigner( if not build_full: return None, None, edge_quat if packed_wigner: - from deepmd.kernels.cute.neo.k4_wignerd import ( + from deepmd.pt_expt.kernels.cute.sezm.k4_wignerd import ( run_cute_wignerd, ) @@ -819,7 +819,7 @@ def _packed_wigner_edges_eligible( runtime_dtypes: tuple[torch.dtype, ...] = (), ) -> bool: """Finish packed eligibility from scalar shape and provenance metadata.""" - from deepmd.kernels.cute.neo.k1 import ( + from deepmd.pt_expt.kernels.cute.sezm.k1 import ( packed_wigner_edges_eligible, ) diff --git a/deepmd/pt/model/descriptor/sezm_nn/embedding.py b/deepmd/pt/model/descriptor/sezm_nn/embedding.py index 71fa44906f..5bfe764495 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/embedding.py +++ b/deepmd/pt/model/descriptor/sezm_nn/embedding.py @@ -22,10 +22,6 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) -from deepmd.kernels.cute.neo.k1_wigner_layout import ( - PACKED_VALUE_COUNT, - ZONAL_PANEL_OFFSETS, -) from deepmd.pt.model.network.mlp import ( MLPLayer, ) @@ -39,6 +35,10 @@ from deepmd.pt.utils.utils import ( get_generator, ) +from deepmd.pt_expt.kernels.cute.sezm.k1_wigner_layout import ( + PACKED_VALUE_COUNT, + ZONAL_PANEL_OFFSETS, +) from deepmd.pt_expt.kernels.utils import ( cuda_infer_level, ) @@ -305,7 +305,7 @@ def forward( # === Step 3. Optional fused message construction and reduction === if not self.training: - from deepmd.kernels.cute.neo.gie import ( + from deepmd.pt_expt.kernels.cute.sezm.gie import ( is_cute_gie_enabled, maybe_run_cute_gie, ) diff --git a/deepmd/pt/model/descriptor/sezm_nn/grid_net.py b/deepmd/pt/model/descriptor/sezm_nn/grid_net.py index e599943d67..bc6b0aa1c6 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/grid_net.py +++ b/deepmd/pt/model/descriptor/sezm_nn/grid_net.py @@ -30,13 +30,13 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) -from deepmd.kernels.cute.neo import runtime_policy as cute_runtime_policy from deepmd.pt.utils import ( env, ) from deepmd.pt.utils.utils import ( get_generator, ) +from deepmd.pt_expt.kernels.cute.sezm import runtime_policy as cute_runtime_policy from deepmd.pt_expt.kernels.utils import ( cuda_infer_level, triton_train_level, @@ -1062,7 +1062,7 @@ def _pair_grid( and cute_runtime_policy.is_cute_infer_enabled() and _inference_mode_is_frozen(self) ): - from deepmd.kernels.cute.neo.output_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid_product import ( maybe_run_cute_output_grid_product, ) diff --git a/deepmd/pt/model/descriptor/sezm_nn/lora.py b/deepmd/pt/model/descriptor/sezm_nn/lora.py index ae3413a302..1b60dfa5de 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/lora.py +++ b/deepmd/pt/model/descriptor/sezm_nn/lora.py @@ -487,12 +487,12 @@ def _clear_sezm_compile_cache(model: nn.Module) -> None: crash or silently skip LoRA parameters. Mirrors the pattern used in :meth:`SeZMModel.reset_head_for_mode`. """ - from deepmd.kernels.cute.neo.k1 import ( - invalidate_cute_k1_state, - ) from deepmd.pt.model.model.sezm_model import ( _clear_shared_sezm_compile_cache, ) + from deepmd.pt_expt.kernels.cute.sezm.k1 import ( + invalidate_cute_k1_state, + ) readout_invalidator = None @@ -500,7 +500,7 @@ def _clear_sezm_compile_cache(model: nn.Module) -> None: invalidate_cute_k1_state(m) if hasattr(m, "_neo_sm80_readout_input_fold_cache"): if readout_invalidator is None: - from deepmd.kernels.cute.neo.readout_l0 import ( + from deepmd.pt_expt.kernels.cute.sezm.readout_l0 import ( invalidate_neo_readout_input_fold, ) diff --git a/deepmd/pt/model/descriptor/sezm_nn/so2.py b/deepmd/pt/model/descriptor/sezm_nn/so2.py index a366babdd2..1c4ed431e0 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/so2.py +++ b/deepmd/pt/model/descriptor/sezm_nn/so2.py @@ -37,7 +37,6 @@ cuda_train_enabled, triton_infer_level, triton_train_level, - use_cute_infer, use_cutile_infer, ) from deepmd.utils.version import ( @@ -1215,31 +1214,27 @@ def __init__( self.compute_dtype = get_promoted_dtype(self.dtype) # Opt-in fused fast paths, selected by ``DP_TRITON_INFER`` / # ``DP_TRITON_TRAIN`` (cumulative levels, see :func:`triton_infer_level` - # and :func:`triton_train_level`) and ``DP_CUTE_INFER``. Each is read + # and :func:`triton_train_level`). Each is read # once at construction so it becomes a compile-time constant in the # traced (``make_fx``) graph. Level 1 replaces the dense ``bmm`` # rotation with universal Triton kernels; level 2 additionally binds the # table-configured fused value path; inference level 3 routes the mixing # stack through the fp16x3 tensor-core operator on swept shapes. - # ``DP_CUTE_INFER`` selects the experimental CuTe value-path operator - # instead; both gates claim the same ``so2_message`` value path, so - # enabling them together has no coherent meaning and is rejected at - # construction. The CuTe, cuTile and hand-written CUDA paths remain - # inference-only. The fused value-path entries are bound at the end of - # construction (once every submodule exists) and stay ``None`` when the - # backend is unavailable or the block layout is unsupported. + # ``DP_CUTILE_INFER`` selects a complete alternative SO(2) path and is + # therefore mutually exclusive with Triton. CuTe dispatch is resolved at + # the enclosing SeZM block, where its exact-shape K1 kernel can coexist + # with these fallback paths. cuTile and the hand-written CUDA paths remain + # inference-only. Fused entries stay ``None`` when their backend is + # unavailable or the block layout is unsupported. self.triton_infer_level = triton_infer_level() self.triton_train_level = triton_train_level() self.use_triton_infer = self.triton_infer_level >= 1 - self.use_cute_infer = use_cute_infer() self.use_cutile_infer = use_cutile_infer() - if sum((self.use_triton_infer, self.use_cute_infer, self.use_cutile_infer)) > 1: + if self.use_triton_infer and self.use_cutile_infer: raise ValueError( - "DP_TRITON_INFER, DP_CUTE_INFER and DP_CUTILE_INFER are mutually " - "exclusive: each selects a complete accelerated inference path. " - "Enable exactly one of them." + "DP_TRITON_INFER and DP_CUTILE_INFER are mutually exclusive: " + "each selects a complete accelerated SO(2) inference path." ) - self._cute_value_path = None self._triton_value_path = None self._cutile_value_path = None @@ -1753,19 +1748,9 @@ def __init__( if SEGMENT_SOFTMAX_TRITON_AVAILABLE: self._segment_softmax_fn = segment_softmax - # === Step 17. Optional fused CuTe SO(2) value-path operator === - # Experimental alternative backend; mutually exclusive with the Triton - # flag (enforced above). - if self.use_cute_infer: - from deepmd.pt_expt.kernels.cute.sezm import ( - make_cute_value_path, - ) - - self._cute_value_path = make_cute_value_path(self) - - # === Step 18. Optional fused cuTile SO(2) value-path operators === - # Complete cuTile inference path, mutually exclusive with the two gates - # above. The factory validates the block layout and returns ``None`` + # === Step 17. Optional fused cuTile SO(2) value-path operators === + # Complete cuTile inference path, mutually exclusive with Triton. The + # factory validates the block layout and returns ``None`` # otherwise, leaving the dense reference path in charge. if self.use_cutile_infer: from deepmd.pt_expt.kernels.cutile.sezm.so2_value_path import ( @@ -2460,12 +2445,6 @@ def so2_message( # inter-layer activations off the traced graph. === cached_edge_csr(edge_cache, "src", x.shape[0]) x_local, rad_feat = self._triton_value_path(x, edge_cache, radial_feat) - elif self._cute_value_path is not None and not self.training: - # === Steps 1-5 (fused CuTe operator). The operator folds - # rotate_to_local, radial degree mixing, the multi-layer gated SO(2) - # stack, and the focus competition into the bucketed kernels; the - # per-edge focus-major intermediates stay resident on chip. === - x_local, rad_feat = self._cute_value_path(x, edge_cache, radial_feat) elif self._triton_rotate_mix is not None and active_triton_level(self) >= 1: # === Steps 1-3 (fused rotate-mix operator). One edge-parallel # kernel gathers the source features, applies the block-diagonal diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index cdc313e0d9..39aa8f13ed 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -532,7 +532,7 @@ def _neo_cute_infer_enabled() -> bool: """Return whether the model builder must destination-sort CuTe edges.""" - from deepmd.kernels.cute.neo.runtime_policy import ( + from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import ( is_cute_infer_enabled, ) @@ -548,7 +548,7 @@ def _neo_cute_nlist_eager_island_enabled(device: torch.device) -> bool: compute_capability = tuple(torch.cuda.get_device_capability(device)) except RuntimeError: return False - from deepmd.kernels.cute.neo.runtime_policy import ( + from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import ( is_k1_eager_island_enabled, ) @@ -929,7 +929,7 @@ def _invalidate_compiled_state_after_load( ): continue if k1_invalidator is None: - from deepmd.kernels.cute.neo.k1 import ( + from deepmd.pt_expt.kernels.cute.sezm.k1 import ( invalidate_cute_k1_state, ) @@ -2156,7 +2156,9 @@ def trace_and_compile( # Register Python-owned K1 state before make_fx starts. The opt-in thin # path can then keep adjacent linears in this graph while the CuTe work # remains opaque behind its existing custom op. - from deepmd.kernels.cute.neo import runtime_policy as cute_runtime_policy + from deepmd.pt_expt.kernels.cute.sezm import ( + runtime_policy as cute_runtime_policy, + ) compute_capability = ( tuple(torch.cuda.get_device_capability(coord.device)) @@ -2164,7 +2166,7 @@ def trace_and_compile( else None ) if not self.training and cute_runtime_policy.is_cute_infer_enabled(): - from deepmd.kernels.cute.neo.readout_l0 import ( + from deepmd.pt_expt.kernels.cute.sezm.readout_l0 import ( maybe_prepare_sm80_readout_input_fold, ) @@ -2179,7 +2181,7 @@ def trace_and_compile( and cute_runtime_policy.is_cute_infer_enabled() and cute_runtime_policy.is_supported_k1_capability(compute_capability) ): - from deepmd.kernels.cute.neo.k1 import ( + from deepmd.pt_expt.kernels.cute.sezm.k1 import ( prepare_cute_k1_blocks, ) diff --git a/deepmd/pt/model/network/mlp.py b/deepmd/pt/model/network/mlp.py index b342fe75b6..51bc476845 100644 --- a/deepmd/pt/model/network/mlp.py +++ b/deepmd/pt/model/network/mlp.py @@ -54,7 +54,7 @@ def _use_k1_compile_visible_linear( """Keep the SM80 linear topology stable for one compiled graph.""" truthy = {"1", "true", "yes", "on"} falsy = {"0", "false", "no", "off"} - cute_enabled = os.environ.get("DP_NEO_CUTE_INFER", "").strip().lower() + cute_enabled = os.environ.get("DP_CUTE_INFER", "").strip().lower() if cute_enabled not in truthy: return False thin_enabled = os.environ.get("DP_CUTE_K1_THIN_WRAPPER", "").strip().lower() diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/so2.py b/deepmd/pt_expt/descriptor/dpa4_nn/so2.py index b244f33e6a..595169e1e3 100644 --- a/deepmd/pt_expt/descriptor/dpa4_nn/so2.py +++ b/deepmd/pt_expt/descriptor/dpa4_nn/so2.py @@ -9,11 +9,11 @@ - the two rotation hot paths of :class:`SO2Convolution`, and - the low-rank branch of :class:`DynamicRadialDegreeMixer`. -Triton, CuTe, and cuTile are mutually exclusive complete SO(2) paths. The -hand-written CUDA operators form an independent cumulative layer and take -precedence where their factories bind. Every gate is resolved at construction -so export records a static dispatch choice; training and unsupported layouts -retain the dpmodel reference path. +Triton and cuTile are mutually exclusive complete SO(2) paths. The hand-written +CUDA operators form an independent cumulative layer and take precedence where +their factories bind. Every gate is resolved at construction so export records a +static dispatch choice; training and unsupported layouts retain the dpmodel +reference path. """ from __future__ import ( @@ -40,7 +40,6 @@ cuda_train_enabled, triton_infer_level, triton_train_level, - use_cute_infer, use_cutile_infer, ) @@ -182,22 +181,19 @@ class SO2Convolution(SO2ConvolutionDP): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) # The inference gates are read once at construction so they become - # compile-time constants in the traced (``make_fx``) graph. Triton, - # CuTe and cuTile claim the same SO(2) value path and are mutually - # exclusive; the hand-written CUDA operators form an independent, - # cumulative layer and take precedence where their factories bind. + # compile-time constants in the traced (``make_fx``) graph. Triton and + # cuTile claim the same SO(2) value path and are mutually exclusive; the + # hand-written CUDA operators form an independent, cumulative layer and + # take precedence where their factories bind. self.triton_infer_level = triton_infer_level() self.triton_train_level = triton_train_level() self.use_triton_infer = self.triton_infer_level >= 1 - self.use_cute_infer = use_cute_infer() self.use_cutile_infer = use_cutile_infer() - if sum((self.use_triton_infer, self.use_cute_infer, self.use_cutile_infer)) > 1: + if self.use_triton_infer and self.use_cutile_infer: raise ValueError( - "DP_TRITON_INFER, DP_CUTE_INFER and DP_CUTILE_INFER are mutually " - "exclusive: each selects a complete accelerated inference path. " - "Enable exactly one of them." + "DP_TRITON_INFER and DP_CUTILE_INFER are mutually exclusive: " + "each selects a complete accelerated SO(2) inference path." ) - self._cute_value_path = None self._triton_value_path = None self._cutile_value_path = None self._cached_edge_csr_fn = cached_edge_csr @@ -305,19 +301,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self._cuda_conv_fn = make_cuda_so2_conv(self) - # === Step 14. Optional fused CuTe SO(2) value-path operator === - # Experimental alternative backend; mutually exclusive with the Triton - # flag (enforced above). - if self.use_cute_infer: - from deepmd.pt_expt.kernels.cute.sezm import ( - make_cute_value_path, - ) - - self._cute_value_path = make_cute_value_path(self) - - # === Step 15. Optional fused cuTile SO(2) value-path operators === - # Complete cuTile inference path, mutually exclusive with the two gates - # above. The factory validates the block layout and returns ``None`` + # === Step 14. Optional fused cuTile SO(2) value-path operators === + # Complete cuTile inference path, mutually exclusive with Triton. The + # factory validates the block layout and returns ``None`` # otherwise, leaving the dense reference path in charge. if self.use_cutile_infer: from deepmd.pt_expt.kernels.cutile.sezm.so2_value_path import ( diff --git a/deepmd/pt_expt/kernels/cute/sezm/__init__.py b/deepmd/pt_expt/kernels/cute/sezm/__init__.py index 70f5cdca4b..56baf26081 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/__init__.py +++ b/deepmd/pt_expt/kernels/cute/sezm/__init__.py @@ -1,54 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -""" -CuTe-DSL fused SO(2) value-path operator for SeZM / DPA4. - -This package hosts a single bucketed CuTe operator that folds the entire per-edge -value path of :class:`~deepmd.pt_expt.descriptor.dpa4_nn.so2.SO2Convolution` -(``rotate_to_local`` -> radial degree mix -> the three-layer gated SO(2) mixing -stack -> focus competition) into a fused forward kernel and a matching -recompute backward kernel, keeping the per-edge intermediates on chip. It is an -opt-in inference path enabled by ``DP_CUTE_INFER``; the final local features are -handed to the committed flash-attention aggregation for rotate-back and scatter. -Kernel entry points are internal implementation details of the SeZM descriptor; -the package-level API only exposes availability and the value-path factory. - -Current limitations -------------------- -Performance - On H20 / fp32 the operator is about 2.8x slower than the compiled Triton + - flash-attention path (roughly 489 / 724 ms versus 174 / 262 ms per force step - at 2000 / 4000 atoms). Peak memory is at parity with, or marginally below, - the compiled path (about 0.5 / 0.8 GB lower) and roughly 1.68x below the - eager path. The bottleneck is the recompute backward, which dominates the - kernel time: its occupancy is capped by the block-diagonal weight held - resident in shared memory, and both the forward and backward GEMMs run at the - hand-written plateau of about 21% of fp32 peak (versus about 52% for cuBLAS). - -Deployment - This is a Python-inference-only path. The ``cutlass.cute`` kernels are - nvcc / NVRTC JIT-compiled at runtime and do not bake into the AOTInductor - ``.pt2`` artifact, so the operator is unavailable to the LAMMPS / GPUMD C++ - inference path. ``DP_CUTE_INFER`` and ``DP_TRITON_INFER`` both claim the - fused SO(2) value path and are mutually exclusive; enabling both is - rejected at construction. - -Correctness - The force is bit-exact against the eager reference (energy relative error - about 1e-9, force relative error about 5e-7 in fp32). -""" - -from __future__ import ( - annotations, -) - -from .forward import ( - SEZM_CUTE_AVAILABLE, -) -from .operator import ( - make_cute_value_path, -) - -__all__ = [ - "SEZM_CUTE_AVAILABLE", - "make_cute_value_path", -] +"""CuTe kernels and PyTorch integration for supported SeZM layouts.""" diff --git a/deepmd/pt_expt/kernels/cute/sezm/backward.py b/deepmd/pt_expt/kernels/cute/sezm/backward.py deleted file mode 100644 index cf1227e683..0000000000 --- a/deepmd/pt_expt/kernels/cute/sezm/backward.py +++ /dev/null @@ -1,692 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -# pyright: reportMissingImports=false -# ruff: noqa: ANN001, ANN201, ANN202, ANN204, ANN205 -""" -CuTe-DSL fused recompute-backward kernel for the SeZM SO(2) value path. - -Given the upstream gradients ``g_out`` (w.r.t. the pre-focus-compete local -features ``x_local``) and ``g_fgate`` (w.r.t. the pre-mixing ``l = 0`` scalar), -one bucketed kernel recomputes the forward value path from the small saved -inputs (``x``, ``D_to_m``, ``Kc``) entirely on chip and backpropagates it, -emitting the position-path gradients that carry the ``edge_vec`` -> force -dependence:: - - grad_x node-feature gradient, scattered to source nodes - grad_D_to_m Wigner-rotation gradient, atomically summed over focus - grad_Kc radial degree-kernel gradient, atomically summed over focus - -The weights are frozen on the inference force path. No ``E x D_m x C`` -intermediate reaches DRAM: the kernel recomputes the forward storing only the two -gated-layer pre-activations in shared memory, then backpropagates the residual -stack (gated-activation backward in registers/smem), the radial degree mix, and -the rotation. This kernel is specialized to the three-layer ``[gated, gated, -identity]`` mixing stack of the deployed configuration. - -Buffers per CTA (one focus of a bucket of ``B`` edges): four ``B x (D_m*Cf)`` -scratch tensors plus one ``max_block^2`` weight/gate scratch, inside the sm_90 -shared-memory limit at ``B = 16``. All accumulation is fp32 IEEE. -""" - -from __future__ import ( - annotations, -) - -import torch - -from .forward import ( - SEZM_CUTE_AVAILABLE, -) - -if SEZM_CUTE_AVAILABLE: - import cuda.bindings.driver as cuda - import cutlass - import cutlass.cute as cute - import cutlass.cute.math as cmath - from cutlass.cute.runtime import ( - from_dlpack, - ) - - from .forward import ( - ForwardRunner, - ) - - class BackwardProgram: - """Bucketed fused recompute-backward program for the SO(2) value path. - - Parameters - ---------- - lmax, mmax, cf, n_focus, n_layers, bucket, threads, rb, rn - Kernel configuration (see :class:`.forward.ForwardProgram`). - """ - - def __init__( - self, - *, - lmax: int, - mmax: int, - cf: int, - n_focus: int, - n_layers: int, - bucket: int, - threads: int, - rb: int, - rn: int, - ) -> None: - self.lmax, self.mmax, self.cf, self.nf, self.nl = ( - lmax, - mmax, - cf, - n_focus, - n_layers, - ) - self._B, self._T, self._RB, self._RN = bucket, threads, rb, rn - self.D = (lmax + 1) ** 2 - self.Dm = (lmax + 1) + sum(2 * (lmax - m + 1) for m in range(1, mmax + 1)) - self.Cw = n_focus * cf - self.gate_out = lmax * cf - self.ngroup = 1 + 2 * mmax - self.FLAT = self.Dm * cf - groups = [lmax + 1] + [2 * (lmax - m + 1) for m in range(1, mmax + 1)] - self._blocks: list[tuple[int, int]] = [] - off = 0 - for g in groups: - self._blocks.append((off, g * cf)) - off += g * cf - self._max_sb = max(sb for _, sb in self._blocks) - self._scr = max( - self._max_sb * self._max_sb, - cf * self.gate_out + 2 * bucket * self.gate_out, - ) - assert self._B % rb == 0 - for _, sb in self._blocks: - assert sb % rn == 0 - - @cute.jit - def __call__( - self, - mGout, - mGfg, - mX, - mSrc, - mDtoM, - mKc, - mCB, - mW, - mGW, - mExpand, - mGx, - mGD, - mGKc, - n_edge: cutlass.Int32, - n_bucket: cutlass.Int32, - stream: cuda.CUstream, - ): - self.kernel( - mGout, - mGfg, - mX, - mSrc, - mDtoM, - mKc, - mCB, - mW, - mGW, - mExpand, - mGx, - mGD, - mGKc, - n_edge, - ).launch(grid=[n_bucket, self.nf, 1], block=[self._T, 1, 1], stream=stream) - - @cute.kernel - def kernel( - self, - mGout, - mGfg, - mX, - mSrc, - mDtoM, - mKc, - mCB, - mW, - mGW, - mExpand, - mGx, - mGD, - mGKc, - n_edge: cutlass.Int32, - ): - D = cutlass.const_expr(self.D) - Dm = cutlass.const_expr(self.Dm) - CF = cutlass.const_expr(self.cf) - FLAT = cutlass.const_expr(self.FLAT) - B = cutlass.const_expr(self._B) - T = cutlass.const_expr(self._T) - - tidx, _, _ = cute.arch.thread_idx() - bucket, focus, _ = cute.arch.block_idx() - e0 = bucket * B - - smem = cutlass.utils.SmemAllocator() - b0 = smem.allocate_tensor( - cutlass.Float32, cute.make_layout((B, FLAT), stride=(FLAT, 1)), 16 - ) - b1 = smem.allocate_tensor( - cutlass.Float32, cute.make_layout((B, FLAT), stride=(FLAT, 1)), 16 - ) - b2 = smem.allocate_tensor( - cutlass.Float32, cute.make_layout((B, FLAT), stride=(FLAT, 1)), 16 - ) - b3 = smem.allocate_tensor( - cutlass.Float32, cute.make_layout((B, FLAT), stride=(FLAT, 1)), 16 - ) - s_w = smem.allocate_tensor( - cutlass.Float32, - cute.make_layout((cutlass.const_expr(self._scr),), stride=(1,)), - 16, - ) - - # === Step 1. Forward recompute: store z_0 -> b1, z_1 -> b2 === - # b0 carries the running layer input h_l; b3 is the activation temp. - i = tidx - while i < B * FLAT: - b = i // FLAT - rem = i % FLAT - dm = rem // CF - c = rem % CF - e = e0 + b - eq = e if e < n_edge else n_edge - 1 - src = mSrc[eq] - acc = cutlass.Float32(0.0) - for k in cutlass.range_constexpr(D): - acc += mDtoM[eq, dm, k] * mX[src, k, focus * CF + c] - b1[b, rem] = acc - i += T - cute.arch.sync_threads() - i = tidx - while i < B * FLAT: - b = i // FLAT - rem = i % FLAT - o = rem // CF - c = rem % CF - e = e0 + b - eq = e if e < n_edge else n_edge - 1 - acc = cutlass.Float32(0.0) - for ii in cutlass.range_constexpr(Dm): - acc += mKc[eq, o, ii] * b1[b, ii * CF + c] - b0[b, rem] = acc * mCB[focus * CF + c] - i += T - cute.arch.sync_threads() - for lyr in cutlass.range_constexpr(self.nl - 1): - zbuf = b1 if lyr == 0 else b2 - self._gemm_fwd(b0, zbuf, s_w, mW, lyr, focus, tidx) - self._gated_fwd(zbuf, b3, s_w, mGW, mExpand, lyr, focus, tidx) - i = tidx - while i < B * FLAT: - b = i // FLAT - rem = i % FLAT - b0[b, rem] = b0[b, rem] + b3[b, rem] - i += T - cute.arch.sync_threads() - - # === Step 2. Reverse the residual stack; b0 accumulates grad_h === - i = tidx - while i < B * FLAT: - b = i // FLAT - rem = i % FLAT - o = rem // CF - c = rem % CF - b0[b, rem] = mGout[e0 + b, focus, o, c] - i += T - cute.arch.sync_threads() - # layer 2 (identity): grad_z = grad_out; grad_h += W_2^T @ grad_z - self._gemm_bwd(b0, b3, s_w, mW, self.nl - 1, focus, tidx) - i = tidx - while i < B * FLAT: - b = i // FLAT - rem = i % FLAT - b0[b, rem] = b0[b, rem] + b3[b, rem] - i += T - cute.arch.sync_threads() - # layer 1 (gated): grad_z_1 -> b2 in place; grad_h += W_1^T @ grad_z_1 - self._gated_bwd(b2, b0, s_w, mGW, mExpand, 1, focus, tidx) - self._gemm_bwd(b2, b3, s_w, mW, 1, focus, tidx) - i = tidx - while i < B * FLAT: - b = i // FLAT - rem = i % FLAT - b0[b, rem] = b0[b, rem] + b3[b, rem] - i += T - cute.arch.sync_threads() - # layer 0 (gated): grad_z_0 -> b1 in place; grad_h += W_0^T @ grad_z_0 - self._gated_bwd(b1, b0, s_w, mGW, mExpand, 0, focus, tidx) - self._gemm_bwd(b1, b3, s_w, mW, 0, focus, tidx) - i = tidx - while i < B * FLAT: - b = i // FLAT - rem = i % FLAT - b0[b, rem] = b0[b, rem] + b3[b, rem] - i += T - cute.arch.sync_threads() - # add the focus-competition gradient into the l=0 row - i = tidx - while i < B * CF: - b = i // CF - c = i % CF - b0[b, c] = b0[b, c] + mGfg[e0 + b, focus, c] - i += T - cute.arch.sync_threads() - - # === Step 3. Radial + rotation backward; b0 holds grad_h0 === - # recompute x_rot -> b1 - i = tidx - while i < B * FLAT: - b = i // FLAT - rem = i % FLAT - dm = rem // CF - c = rem % CF - e = e0 + b - eq = e if e < n_edge else n_edge - 1 - src = mSrc[eq] - acc = cutlass.Float32(0.0) - for k in cutlass.range_constexpr(D): - acc += mDtoM[eq, dm, k] * mX[src, k, focus * CF + c] - b1[b, rem] = acc - i += T - cute.arch.sync_threads() - # grad_Kc[o, ii] = sum_c channel_basis[c] * grad_h0[o, c] * x_rot[ii, c] - i = tidx - while i < B * Dm * Dm: - b = i // (Dm * Dm) - rem = i % (Dm * Dm) - o = rem // Dm - ii = rem % Dm - e = e0 + b - acc = cutlass.Float32(0.0) - for c in cutlass.range_constexpr(CF): - acc += mCB[focus * CF + c] * b0[b, o * CF + c] * b1[b, ii * CF + c] - cute.arch.atomic_add(mGKc.iterator + mGKc.layout((e, o, ii)), acc) - i += T - # grad_x_rot[ii, c] = channel_basis[c] * sum_o Kc[o, ii] * grad_h0[o, c] -> b2 - i = tidx - while i < B * Dm * CF: - b = i // (Dm * CF) - rem = i % (Dm * CF) - ii = rem // CF - c = rem % CF - e = e0 + b - eq = e if e < n_edge else n_edge - 1 - acc = cutlass.Float32(0.0) - for o in cutlass.range_constexpr(Dm): - acc += mKc[eq, o, ii] * b0[b, o * CF + c] - b2[b, ii * CF + c] = acc * mCB[focus * CF + c] - i += T - cute.arch.sync_threads() - # grad_x[src, k, c] += sum_ii D_to_m[ii, k] * grad_x_rot[ii, c] - i = tidx - while i < B * D * CF: - b = i // (D * CF) - rem = i % (D * CF) - k = rem // CF - c = rem % CF - e = e0 + b - eq = e if e < n_edge else n_edge - 1 - src = mSrc[eq] - acc = cutlass.Float32(0.0) - for ii in cutlass.range_constexpr(Dm): - acc += mDtoM[eq, ii, k] * b2[b, ii * CF + c] - cute.arch.atomic_add( - mGx.iterator + mGx.layout((src, k, focus * CF + c)), acc - ) - i += T - # grad_D_to_m[ii, k] = sum_c grad_x_rot[ii, c] * x_src[k, c] - i = tidx - while i < B * Dm * D: - b = i // (Dm * D) - rem = i % (Dm * D) - ii = rem // D - k = rem % D - e = e0 + b - eq = e if e < n_edge else n_edge - 1 - src = mSrc[eq] - acc = cutlass.Float32(0.0) - for c in cutlass.range_constexpr(CF): - acc += b2[b, ii * CF + c] * mX[src, k, focus * CF + c] - cute.arch.atomic_add(mGD.iterator + mGD.layout((e, ii, k)), acc) - i += T - - @cute.jit - def _gemm_fwd(self, hbuf, zbuf, s_w, mW, lyr, focus, tidx): - """Recompute ``zbuf = hbuf @ W[lyr, focus]`` (block-diagonal).""" - T = cutlass.const_expr(self._T) - RB = cutlass.const_expr(self._RB) - RN = cutlass.const_expr(self._RN) - B = cutlass.const_expr(self._B) - for ob, sb in cutlass.const_expr(self._blocks): - j = tidx - while j < sb * sb: - s_w[(j // sb) * sb + (j % sb)] = mW[ - lyr, focus, ob + (j // sb), ob + (j % sb) - ] - j += T - cute.arch.sync_threads() - n_mt = cutlass.const_expr((B // RB) * (sb // RN)) - racc = cute.make_rmem_tensor( - cute.make_layout((RB, RN)), cutlass.Float32 - ) - mt = tidx - while mt < n_mt: - bi = (mt // (sb // RN)) * RB - nj = (mt % (sb // RN)) * RN - for r in cutlass.range_constexpr(RB): - for s in cutlass.range_constexpr(RN): - racc[r, s] = cutlass.Float32(0.0) - for k in cutlass.range(sb): - for r in cutlass.range_constexpr(RB): - a_rk = hbuf[bi + r, ob + k] - for s in cutlass.range_constexpr(RN): - racc[r, s] += a_rk * s_w[k * sb + nj + s] - for r in cutlass.range_constexpr(RB): - for s in cutlass.range_constexpr(RN): - zbuf[bi + r, ob + nj + s] = racc[r, s] - mt += T - cute.arch.sync_threads() - - @cute.jit - def _gemm_bwd(self, gzbuf, ghbuf, s_w, mW, lyr, focus, tidx): - """Compute ``ghbuf = W[lyr, focus]^T @ gzbuf`` (block-diagonal).""" - T = cutlass.const_expr(self._T) - RB = cutlass.const_expr(self._RB) - RN = cutlass.const_expr(self._RN) - B = cutlass.const_expr(self._B) - for ob, sb in cutlass.const_expr(self._blocks): - j = tidx - while j < sb * sb: - s_w[(j // sb) * sb + (j % sb)] = mW[ - lyr, focus, ob + (j // sb), ob + (j % sb) - ] - j += T - cute.arch.sync_threads() - n_mt = cutlass.const_expr((B // RB) * (sb // RN)) - racc = cute.make_rmem_tensor( - cute.make_layout((RB, RN)), cutlass.Float32 - ) - mt = tidx - while mt < n_mt: - bi = (mt // (sb // RN)) * RB - kj = (mt % (sb // RN)) * RN # W in-index (grad_h output column) - for r in cutlass.range_constexpr(RB): - for s in cutlass.range_constexpr(RN): - racc[r, s] = cutlass.Float32(0.0) - for n in cutlass.range(sb): # sum over the W out-index - for r in cutlass.range_constexpr(RB): - gz_rn = gzbuf[bi + r, ob + n] - for s in cutlass.range_constexpr(RN): - racc[r, s] += gz_rn * s_w[(kj + s) * sb + n] - for r in cutlass.range_constexpr(RB): - for s in cutlass.range_constexpr(RN): - ghbuf[bi + r, ob + kj + s] = racc[r, s] - mt += T - cute.arch.sync_threads() - - @cute.jit - def _gated_fwd(self, zbuf, abuf, s_w, mGW, mExpand, lyr, focus, tidx): - """Recompute ``abuf = GatedActivation(zbuf)`` (silu l=0, gate l>0).""" - CF = cutlass.const_expr(self.cf) - GO = cutlass.const_expr(self.gate_out) - Dm = cutlass.const_expr(self.Dm) - B = cutlass.const_expr(self._B) - T = cutlass.const_expr(self._T) - SIG_OFF = cutlass.const_expr(self.cf * self.gate_out) - j = tidx - while j < CF * GO: - s_w[(j // GO) * GO + (j % GO)] = mGW[lyr, focus, j // GO, j % GO] - j += T - cute.arch.sync_threads() - j = tidx - while j < B * GO: - b = j // GO - o = j % GO - acc = cutlass.Float32(0.0) - for ii in cutlass.range_constexpr(CF): - acc += zbuf[b, ii] * s_w[ii * GO + o] - s_w[SIG_OFF + b * GO + o] = cutlass.Float32(1.0) / ( - cutlass.Float32(1.0) + cmath.exp(-acc) - ) - j += T - cute.arch.sync_threads() - j = tidx - while j < B * CF: - b = j // CF - c = j % CF - z = zbuf[b, c] - abuf[b, c] = z / (cutlass.Float32(1.0) + cmath.exp(-z)) - j += T - j = tidx - while j < B * (Dm - 1) * CF: - b = j // ((Dm - 1) * CF) - rem = j % ((Dm - 1) * CF) - d1 = rem // CF - c = rem % CF - lidx = mExpand[d1] - abuf[b, (d1 + 1) * CF + c] = ( - zbuf[b, (d1 + 1) * CF + c] * s_w[SIG_OFF + b * GO + lidx * CF + c] - ) - j += T - cute.arch.sync_threads() - - @cute.jit - def _gated_bwd(self, zbuf, gabuf, s_w, mGW, mExpand, lyr, focus, tidx): - """Backprop the gated activation in place: ``zbuf`` (z_l) -> grad_z_l. - - With ``g_a = gabuf`` the incoming gradient and the recomputed gate - sigmoids ``sig``:: - - grad_z[dm, c] = g_a[dm, c] * sig[expand[dm-1], c] (dm > 0) - g_sig[L, c] = sum_{dm: expand[dm-1]=L} g_a[dm, c] * z[dm, c] - grad_z[0, i] = g_a[0, i] * silu'(z[0, i]) - + sum_{o'} Wg[i, o'] * g_sig * sig*(1-sig) - """ - CF = cutlass.const_expr(self.cf) - GO = cutlass.const_expr(self.gate_out) - Dm = cutlass.const_expr(self.Dm) - LMAX = cutlass.const_expr(self.lmax) - NG = cutlass.const_expr(self.ngroup) - B = cutlass.const_expr(self._B) - T = cutlass.const_expr(self._T) - SIG_OFF = cutlass.const_expr(self.cf * self.gate_out) - GGL_OFF = cutlass.const_expr(self.cf * self.gate_out + B * self.gate_out) - j = tidx - while j < CF * GO: - s_w[(j // GO) * GO + (j % GO)] = mGW[lyr, focus, j // GO, j % GO] - j += T - cute.arch.sync_threads() - j = tidx - while j < B * GO: - b = j // GO - o = j % GO - acc = cutlass.Float32(0.0) - for ii in cutlass.range_constexpr(CF): - acc += zbuf[b, ii] * s_w[ii * GO + o] - s_w[SIG_OFF + b * GO + o] = cutlass.Float32(1.0) / ( - cutlass.Float32(1.0) + cmath.exp(-acc) - ) - j += T - cute.arch.sync_threads() - # g_gl[L, c] = (sum over the (1 + 2*mmax) |m| groups) * sigmoid'(gate). - # For degree l = L + 1 the contributing coefficients are dm = 1 + L + k*lmax. - j = tidx - while j < B * GO: - b = j // GO - o = j % GO - gate_l = o // CF - c = o % CF - gsig = cutlass.Float32(0.0) - for kk in cutlass.range_constexpr(NG): - dm = 1 + gate_l + kk * LMAX - gsig += gabuf[b, dm * CF + c] * zbuf[b, dm * CF + c] - s = s_w[SIG_OFF + b * GO + o] - s_w[GGL_OFF + b * GO + o] = gsig * s * (cutlass.Float32(1.0) - s) - j += T - cute.arch.sync_threads() - # grad_z[dm>0]: reads sig, writes disjoint from the l=0 slice. - j = tidx - while j < B * (Dm - 1) * CF: - b = j // ((Dm - 1) * CF) - rem = j % ((Dm - 1) * CF) - d1 = rem // CF - c = rem % CF - lidx = mExpand[d1] - zbuf[b, (d1 + 1) * CF + c] = ( - gabuf[b, (d1 + 1) * CF + c] * s_w[SIG_OFF + b * GO + lidx * CF + c] - ) - j += T - # grad_z[0]: reads z[0, i] before overwriting it. - j = tidx - while j < B * CF: - b = j // CF - ii = j % CF - z0 = zbuf[b, ii] - sg = cutlass.Float32(1.0) / (cutlass.Float32(1.0) + cmath.exp(-z0)) - silup = sg + z0 * sg * (cutlass.Float32(1.0) - sg) - acc = gabuf[b, ii] * silup - for o in cutlass.range_constexpr(GO): - acc += s_w[ii * GO + o] * s_w[GGL_OFF + b * GO + o] - zbuf[b, ii] = acc - j += T - cute.arch.sync_threads() - - class BackwardRunner: - """Compile-once driver for :class:`BackwardProgram` (force-path gradients). - - Parameters - ---------- - weights - Packed weights (see :class:`.forward.ForwardRunner`). - lmax, mmax, cf, n_focus, n_layers, bucket, threads, rb, rn - Kernel configuration. - """ - - def __init__( - self, - weights, - *, - lmax: int, - mmax: int, - cf: int, - n_focus: int, - n_layers: int, - bucket: int, - threads: int, - rb: int, - rn: int, - ) -> None: - self.op = BackwardProgram( - lmax=lmax, - mmax=mmax, - cf=cf, - n_focus=n_focus, - n_layers=n_layers, - bucket=bucket, - threads=threads, - rb=rb, - rn=rn, - ) - self._B = bucket - self.nf, self.cf, self.Dm, self.D = n_focus, cf, self.op.Dm, self.op.D - self._compiled = None - fr = ForwardRunner( - weights, - lmax=lmax, - mmax=mmax, - cf=cf, - n_focus=n_focus, - n_layers=n_layers, - bucket=bucket, - threads=threads, - rb=rb, - rn=rn, - ) - self.m_w, self.m_gw, self.m_cb, self.m_expand = ( - fr.m_w, - fr.m_gw, - fr.m_cb, - fr.m_expand, - ) - - @staticmethod - def _dyn(t: torch.Tensor, leading: int): - return from_dlpack(t, assumed_align=16).mark_layout_dynamic( - leading_dim=leading - ) - - def __call__( - self, - x: torch.Tensor, - src: torch.Tensor, - d_to_m: torch.Tensor, - kc: torch.Tensor, - g_out: torch.Tensor, - g_fgate: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Emit the value-path gradients. - - Parameters - ---------- - x, src, d_to_m, kc - Forward inputs (see :meth:`.forward.ForwardRunner.__call__`). - g_out : torch.Tensor - Gradient w.r.t. ``x_local`` with shape (E, F, D_m, Cf). - g_fgate : torch.Tensor - Gradient w.r.t. ``focus_gate`` with shape (E, F, Cf). - - Returns - ------- - grad_x : torch.Tensor - Node-feature gradient with shape (N, D, C_wide). - grad_d_to_m : torch.Tensor - Wigner-rotation gradient with shape (E, D_m, D). - grad_kc : torch.Tensor - Radial degree-kernel gradient with shape (E, D_m, D_m). - """ - n_edge = src.shape[0] - n_bucket = (n_edge + self._B - 1) // self._B - n_pad = n_bucket * self._B - s32 = src.to(torch.int32) - g_out_p, g_fg_p = g_out, g_fgate - if n_pad > n_edge: - s32 = torch.cat([s32, s32.new_zeros(n_pad - n_edge)]) - g_out_p = torch.cat( - [g_out, g_out.new_zeros(n_pad - n_edge, *g_out.shape[1:])] - ) - g_fg_p = torch.cat( - [g_fgate, g_fgate.new_zeros(n_pad - n_edge, *g_fgate.shape[1:])] - ) - grad_x = torch.zeros_like(x) - grad_d = torch.zeros( - n_pad, self.Dm, self.D, device=x.device, dtype=torch.float32 - ) - grad_kc = torch.zeros( - n_pad, self.Dm, self.Dm, device=x.device, dtype=torch.float32 - ) - views = ( - self._dyn(g_out_p, 3), - self._dyn(g_fg_p, 2), - self._dyn(x, 2), - self._dyn(s32, 0), - self._dyn(d_to_m, 2), - self._dyn(kc, 2), - from_dlpack(self.m_cb, assumed_align=16), - from_dlpack(self.m_w, assumed_align=16), - from_dlpack(self.m_gw, assumed_align=16), - from_dlpack(self.m_expand, assumed_align=16), - self._dyn(grad_x, 2), - self._dyn(grad_d, 2), - self._dyn(grad_kc, 2), - ) - args = (*views, cutlass.Int32(n_edge), cutlass.Int32(n_bucket)) - stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream) - if self._compiled is None: - self._compiled = cute.compile(self.op, *args, stream=stream) - self._compiled(*args, stream=stream) - return grad_x, grad_d[:n_edge], grad_kc[:n_edge] diff --git a/deepmd/kernels/cute/neo/compile_cache.py b/deepmd/pt_expt/kernels/cute/sezm/compile_cache.py similarity index 100% rename from deepmd/kernels/cute/neo/compile_cache.py rename to deepmd/pt_expt/kernels/cute/sezm/compile_cache.py diff --git a/deepmd/pt_expt/kernels/cute/sezm/forward.py b/deepmd/pt_expt/kernels/cute/sezm/forward.py deleted file mode 100644 index 0d6ca5ec8a..0000000000 --- a/deepmd/pt_expt/kernels/cute/sezm/forward.py +++ /dev/null @@ -1,444 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -# pyright: reportMissingImports=false -# ruff: noqa: ANN001, ANN201, ANN204, ANN205 -""" -CuTe-DSL fused forward kernel for the SeZM SO(2) value path. - -One bucketed kernel folds the entire per-edge value path of ``SO2Convolution`` -into a single launch:: - - gather x[src] -> rotate_to_local (D_to_m) [prologue, in smem] - -> radial degree mix (Kc, channel_basis) [prologue, in smem] - -> 3x (block-diagonal SO2Linear + GatedActivation + residual) [in smem] - -> x_local (E, F, D_m, Cf) [+ pre-mixing l=0 scalar] - -The block-diagonal ``SO2Linear`` weight is staged in shared memory once per -bucket and reused across the bucket's ``B`` edges (register-blocked ``RB x RN`` -FMA micro-tile), so the ``E x D_m x C`` intermediates of all three mixing layers -stay resident on chip and never reach DRAM. The focus competition (a per-edge -softmax of the pre-mixing ``l = 0`` feature) is applied outside the kernel from -the returned ``focus_gate`` scalar. - -Grid ``(n_bucket, n_focus)``; one CTA owns ``B`` edges of one focus stream. -All accumulation is fp32 IEEE (no TF32) to keep the potential-energy surface -smooth. -""" - -from __future__ import ( - annotations, -) - -import torch - -try: - import cuda.bindings.driver as cuda - import cutlass - import cutlass.cute as cute - import cutlass.cute.math as cmath - from cutlass.cute.runtime import ( - from_dlpack, - ) - - SEZM_CUTE_AVAILABLE = True -except Exception: # pragma: no cover - import guard for non-CuTe environments - SEZM_CUTE_AVAILABLE = False - - -if SEZM_CUTE_AVAILABLE: - - class ForwardProgram: - """Bucketed fused forward program for the SO(2) value path. - - Parameters - ---------- - lmax : int - Maximum spherical harmonic degree. - mmax : int - Maximum SO(2) order retained in the reduced layout. - cf : int - Per-focus channel width ``Cf``. - n_focus : int - Number of focus streams ``F``. - n_layers : int - Number of SO(2) mixing layers. - bucket : int - Edges processed per CTA ``B``. - threads : int - Threads per CTA. - rb, rn : int - Register micro-tile dimensions (``RB`` bucket rows, ``RN`` output - columns per thread) of the block-diagonal GEMM. - """ - - def __init__( - self, - *, - lmax: int, - mmax: int, - cf: int, - n_focus: int, - n_layers: int, - bucket: int, - threads: int, - rb: int, - rn: int, - ) -> None: - self.lmax, self.mmax, self.cf, self.nf, self.nl = ( - lmax, - mmax, - cf, - n_focus, - n_layers, - ) - self._B, self._T, self._RB, self._RN = bucket, threads, rb, rn - self.D = (lmax + 1) ** 2 - self.Dm = (lmax + 1) + sum(2 * (lmax - m + 1) for m in range(1, mmax + 1)) - self.Cw = n_focus * cf - self.gate_out = lmax * cf - self.FLAT = self.Dm * cf - # Block-diagonal |m| block widths in the flattened coeff*channel axis: - # m = 0 spans (lmax + 1) coefficients, each |m| > 0 spans 2*(lmax-m+1). - groups = [lmax + 1] + [2 * (lmax - m + 1) for m in range(1, mmax + 1)] - self._blocks: list[tuple[int, int]] = [] - off = 0 - for g in groups: - self._blocks.append((off, g * cf)) - off += g * cf - self._max_sb = max(sb for _, sb in self._blocks) - # Shared scratch reused for the resident weight block and, during the - # gated activation, the gate weight plus per-edge sigmoid buffer. - self._scr = max( - self._max_sb * self._max_sb, cf * self.gate_out + bucket * self.gate_out - ) - assert self._B % rb == 0 - for _, sb in self._blocks: - assert sb % rn == 0 - - @cute.jit - def __call__( - self, - mX, - mSrc, - mDtoM, - mKc, - mCB, - mW, - mGW, - mExpand, - mOut, - mFocusGate, - n_edge: cutlass.Int32, - n_bucket: cutlass.Int32, - stream: cuda.CUstream, - ): - self.kernel( - mX, mSrc, mDtoM, mKc, mCB, mW, mGW, mExpand, mOut, mFocusGate, n_edge - ).launch(grid=[n_bucket, self.nf, 1], block=[self._T, 1, 1], stream=stream) - - @cute.kernel - def kernel( - self, - mX, - mSrc, - mDtoM, - mKc, - mCB, - mW, - mGW, - mExpand, - mOut, - mFocusGate, - n_edge: cutlass.Int32, - ): - D = cutlass.const_expr(self.D) - Dm = cutlass.const_expr(self.Dm) - CF = cutlass.const_expr(self.cf) - FLAT = cutlass.const_expr(self.FLAT) - GO = cutlass.const_expr(self.gate_out) - B = cutlass.const_expr(self._B) - T = cutlass.const_expr(self._T) - RB = cutlass.const_expr(self._RB) - RN = cutlass.const_expr(self._RN) - SCR = cutlass.const_expr(self._scr) - GATE_OFF = cutlass.const_expr(self.cf * self.gate_out) - NGATE = cutlass.const_expr(self.nl - 1) - - tidx, _, _ = cute.arch.thread_idx() - bucket, focus, _ = cute.arch.block_idx() - e0 = bucket * B - - smem = cutlass.utils.SmemAllocator() - s_cur = smem.allocate_tensor( - cutlass.Float32, cute.make_layout((B, FLAT), stride=(FLAT, 1)), 16 - ) - s_tmp = smem.allocate_tensor( - cutlass.Float32, cute.make_layout((B, FLAT), stride=(FLAT, 1)), 16 - ) - s_scr = smem.allocate_tensor( - cutlass.Float32, cute.make_layout((SCR,), stride=(1,)), 16 - ) - - # === Step 1. rotate_to_local: s_tmp[b, dm*Cf+c] = sum_k D_to_m x[src] === - # Padding edges (e >= n_edge) clamp their read index; their output rows - # are sliced off by the caller. - i = tidx - while i < B * FLAT: - b = i // FLAT - rem = i % FLAT - dm = rem // CF - c = rem % CF - e = e0 + b - eq = e if e < n_edge else n_edge - 1 - src = mSrc[eq] - acc = cutlass.Float32(0.0) - for k in cutlass.range_constexpr(D): - acc += mDtoM[eq, dm, k] * mX[src, k, focus * CF + c] - s_tmp[b, rem] = acc - i += T - cute.arch.sync_threads() - - # === Step 2. radial degree mix: s_cur = channel_basis * (Kc @ x_rot) === - i = tidx - while i < B * FLAT: - b = i // FLAT - rem = i % FLAT - o = rem // CF - c = rem % CF - e = e0 + b - eq = e if e < n_edge else n_edge - 1 - acc = cutlass.Float32(0.0) - for ii in cutlass.range_constexpr(Dm): - acc += mKc[eq, o, ii] * s_tmp[b, ii * CF + c] - s_cur[b, rem] = acc * mCB[focus * CF + c] - i += T - cute.arch.sync_threads() - - # === Step 3. Emit the pre-mixing l=0 scalar for the focus competition === - i = tidx - while i < B * CF: - b = i // CF - c = i % CF - mFocusGate[e0 + b, focus, c] = s_cur[b, c] - i += T - cute.arch.sync_threads() - - # === Step 4. Multi-layer gated SO(2) mixing (block-diagonal, residual) === - for lyr in cutlass.range_constexpr(self.nl): - # --- SO2Linear: s_tmp = s_cur @ W[lyr, focus] (per |m| block) --- - for ob, sb in cutlass.const_expr(self._blocks): - j = tidx - while j < sb * sb: - k = j // sb - n = j % sb - s_scr[k * sb + n] = mW[lyr, focus, ob + k, ob + n] - j += T - cute.arch.sync_threads() - n_mt = cutlass.const_expr((B // RB) * (sb // RN)) - racc = cute.make_rmem_tensor( - cute.make_layout((RB, RN)), cutlass.Float32 - ) - mt = tidx - while mt < n_mt: - bi = (mt // (sb // RN)) * RB - nj = (mt % (sb // RN)) * RN - for r in cutlass.range_constexpr(RB): - for s in cutlass.range_constexpr(RN): - racc[r, s] = cutlass.Float32(0.0) - for k in cutlass.range(sb): - for r in cutlass.range_constexpr(RB): - a_rk = s_cur[bi + r, ob + k] - for s in cutlass.range_constexpr(RN): - racc[r, s] += a_rk * s_scr[k * sb + nj + s] - for r in cutlass.range_constexpr(RB): - for s in cutlass.range_constexpr(RN): - s_tmp[bi + r, ob + nj + s] = racc[r, s] - mt += T - cute.arch.sync_threads() - - # --- GatedActivation (gated layers) or identity (last layer) --- - if lyr < NGATE: - # gate FocusLinear weight -> s_scr[0 : Cf*GO] - j = tidx - while j < CF * GO: - ii = j // GO - o = j % GO - s_scr[ii * GO + o] = mGW[lyr, focus, ii, o] - j += T - cute.arch.sync_threads() - # gate sigmoids from the l=0 scalar -> s_scr[GATE_OFF + b*GO + o] - j = tidx - while j < B * GO: - b = j // GO - o = j % GO - acc = cutlass.Float32(0.0) - for ii in cutlass.range_constexpr(CF): - acc += s_tmp[b, ii] * s_scr[ii * GO + o] - s_scr[GATE_OFF + b * GO + o] = cutlass.Float32(1.0) / ( - cutlass.Float32(1.0) + cmath.exp(-acc) - ) - j += T - cute.arch.sync_threads() - # l=0: silu(z) = z / (1 + exp(-z)) - j = tidx - while j < B * CF: - b = j // CF - c = j % CF - z = s_tmp[b, c] - s_tmp[b, c] = z / (cutlass.Float32(1.0) + cmath.exp(-z)) - j += T - # l>0: z * sigmoid(gate[expand[dm-1]]) - j = tidx - while j < B * (Dm - 1) * CF: - b = j // ((Dm - 1) * CF) - rem = j % ((Dm - 1) * CF) - d1 = rem // CF - c = rem % CF - lidx = mExpand[d1] - z = s_tmp[b, (d1 + 1) * CF + c] - s_tmp[b, (d1 + 1) * CF + c] = ( - z * s_scr[GATE_OFF + b * GO + lidx * CF + c] - ) - j += T - cute.arch.sync_threads() - - # --- residual add: s_cur += activation(s_tmp) --- - i = tidx - while i < B * FLAT: - b = i // FLAT - rem = i % FLAT - s_cur[b, rem] = s_cur[b, rem] + s_tmp[b, rem] - i += T - cute.arch.sync_threads() - - # === Step 5. Write the pre-focus-compete local features (E, F, D_m, Cf) === - i = tidx - while i < B * FLAT: - b = i // FLAT - rem = i % FLAT - o = rem // CF - c = rem % CF - mOut[e0 + b, focus, o, c] = s_cur[b, rem] - i += T - - class ForwardRunner: - """Compile-once driver for :class:`ForwardProgram`. - - Prepares the static packed weights on construction and dispatches the - compiled kernel over dynamic edge counts. - - Parameters - ---------- - weights - Packed weights exposing ``so2_w`` (L, F, D_m*Cf, D_m*Cf), - ``gate_w`` (L, Cf, F, lmax*Cf), ``has_gate`` (L,), and - ``channel_basis`` (C_wide,). - lmax, mmax, cf, n_focus, n_layers, bucket, threads, rb, rn - Kernel configuration (see :class:`ForwardProgram`). - """ - - def __init__( - self, - weights, - *, - lmax: int, - mmax: int, - cf: int, - n_focus: int, - n_layers: int, - bucket: int, - threads: int, - rb: int, - rn: int, - ) -> None: - self.op = ForwardProgram( - lmax=lmax, - mmax=mmax, - cf=cf, - n_focus=n_focus, - n_layers=n_layers, - bucket=bucket, - threads=threads, - rb=rb, - rn=rn, - ) - self._B = bucket - self.nf, self.cf, self.Dm = n_focus, cf, self.op.Dm - self._compiled = None - self._pack(weights, lmax, cf, n_focus, n_layers) - - def _pack(self, w, lmax: int, cf: int, nf: int, nl: int) -> None: - dev = w.so2_w.device - self.m_w = w.so2_w.detach().contiguous() - gate = torch.zeros(nl, nf, cf, lmax * cf, device=dev, dtype=torch.float32) - for layer in range(nl): - if bool(w.has_gate[layer]): - gate[layer] = w.gate_w[layer].detach().permute(1, 0, 2).contiguous() - self.m_gw = gate.contiguous() - self.m_cb = w.channel_basis.detach().contiguous() - # m-major degree index l(dm); the gate expand maps dm>0 -> (l-1). - l_index = list(range(lmax + 1)) - for m in range(1, self.op.mmax + 1): - l_index += list(range(m, lmax + 1)) * 2 - self.m_expand = torch.tensor( - [li - 1 for li in l_index[1:]], device=dev, dtype=torch.int32 - ).contiguous() - - @staticmethod - def _dyn(t: torch.Tensor, leading: int): - return from_dlpack(t, assumed_align=16).mark_layout_dynamic( - leading_dim=leading - ) - - def __call__( - self, - x: torch.Tensor, - src: torch.Tensor, - d_to_m: torch.Tensor, - kc: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Run the fused forward. - - Parameters - ---------- - x : torch.Tensor - Node features with shape (N, D, C_wide). - src : torch.Tensor - Per-edge source-node indices with shape (E,). - d_to_m : torch.Tensor - Row-projected Wigner-D with shape (E, D_m, D). - kc : torch.Tensor - Radial degree kernel with shape (E, D_m, D_m). - - Returns - ------- - x_local : torch.Tensor - Pre-focus-compete local features with shape (E, F, D_m, Cf). - focus_gate : torch.Tensor - Pre-mixing l=0 scalar with shape (E, F, Cf). - """ - n_edge = src.shape[0] - n_bucket = (n_edge + self._B - 1) // self._B - n_pad = n_bucket * self._B - s32 = src.to(torch.int32) - if n_pad > n_edge: - s32 = torch.cat([s32, s32.new_zeros(n_pad - n_edge)]) - out = x.new_empty(n_pad, self.nf, self.Dm, self.cf) - fgate = x.new_empty(n_pad, self.nf, self.cf) - views = ( - self._dyn(x, 2), - self._dyn(s32, 0), - self._dyn(d_to_m, 2), - self._dyn(kc, 2), - from_dlpack(self.m_cb, assumed_align=16), - from_dlpack(self.m_w, assumed_align=16), - from_dlpack(self.m_gw, assumed_align=16), - from_dlpack(self.m_expand, assumed_align=16), - self._dyn(out, 3), - self._dyn(fgate, 2), - ) - args = (*views, cutlass.Int32(n_edge), cutlass.Int32(n_bucket)) - stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream) - if self._compiled is None: - self._compiled = cute.compile(self.op, *args, stream=stream) - self._compiled(*args, stream=stream) - return out[:n_edge], fgate[:n_edge] diff --git a/deepmd/kernels/cute/neo/gie.py b/deepmd/pt_expt/kernels/cute/sezm/gie.py similarity index 99% rename from deepmd/kernels/cute/neo/gie.py rename to deepmd/pt_expt/kernels/cute/sezm/gie.py index 70e747b8c5..8208164483 100644 --- a/deepmd/kernels/cute/neo/gie.py +++ b/deepmd/pt_expt/kernels/cute/sezm/gie.py @@ -11,7 +11,7 @@ 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 enables this +``DP_CUTE_INFER`` is the master opt-in. The SM80/SM86 path enables this fusion by default; ``DP_CUTE_GIE=0`` disables it explicitly. """ diff --git a/deepmd/kernels/cute/neo/k1.py b/deepmd/pt_expt/kernels/cute/sezm/k1.py similarity index 100% rename from deepmd/kernels/cute/neo/k1.py rename to deepmd/pt_expt/kernels/cute/sezm/k1.py diff --git a/deepmd/kernels/cute/neo/k1_gate_structural.py b/deepmd/pt_expt/kernels/cute/sezm/k1_gate_structural.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_gate_structural.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_gate_structural.py diff --git a/deepmd/kernels/cute/neo/k1_kernels/__init__.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/__init__.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_kernels/__init__.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_kernels/__init__.py diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_envelope_gated_softmax.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_envelope_gated_softmax.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_kernels/cute_envelope_gated_softmax.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_envelope_gated_softmax.py diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_focus_src_backward.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_focus_src_backward.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_kernels/cute_neo_focus_src_backward.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_focus_src_backward.py diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_message_grid_product.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_message_grid_product.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_kernels/cute_neo_message_grid_product.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_message_grid_product.py diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_output_gate_backward.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_output_gate_backward.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_kernels/cute_neo_output_gate_backward.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_output_gate_backward.py diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_a_radial_forward.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_a_radial_forward.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_a_radial_forward.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_a_radial_forward.py diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_backward_layout.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_backward_layout.py diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout_runner.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_backward_layout_runner.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout_runner.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_backward_layout_runner.py diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_onepass.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_onepass.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_onepass.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_onepass.py diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_qk_edge.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_qk_edge.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_kernels/cute_neo_qk_edge.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_qk_edge.py diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_radial_phase_a_backward_node.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_radial_phase_a_backward_node.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_kernels/cute_neo_radial_phase_a_backward_node.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_radial_phase_a_backward_node.py diff --git a/deepmd/kernels/cute/neo/k1_kernels/cute_neo_so2_gate_combined_fwd.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_so2_gate_combined_fwd.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_kernels/cute_neo_so2_gate_combined_fwd.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_so2_gate_combined_fwd.py diff --git a/deepmd/kernels/cute/neo/k1_message_grid_packed.py b/deepmd/pt_expt/kernels/cute/sezm/k1_message_grid_packed.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_message_grid_packed.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_message_grid_packed.py diff --git a/deepmd/kernels/cute/neo/k1_radial_phase_a_node.py b/deepmd/pt_expt/kernels/cute/sezm/k1_radial_phase_a_node.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_radial_phase_a_node.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_radial_phase_a_node.py diff --git a/deepmd/kernels/cute/neo/k1_runner.py b/deepmd/pt_expt/kernels/cute/sezm/k1_runner.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_runner.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_runner.py diff --git a/deepmd/kernels/cute/neo/k1_so2linear.py b/deepmd/pt_expt/kernels/cute/sezm/k1_so2linear.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_so2linear.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_so2linear.py diff --git a/deepmd/kernels/cute/neo/k1_wigner_layout.py b/deepmd/pt_expt/kernels/cute/sezm/k1_wigner_layout.py similarity index 100% rename from deepmd/kernels/cute/neo/k1_wigner_layout.py rename to deepmd/pt_expt/kernels/cute/sezm/k1_wigner_layout.py diff --git a/deepmd/kernels/cute/neo/k4_wignerd.py b/deepmd/pt_expt/kernels/cute/sezm/k4_wignerd.py similarity index 100% rename from deepmd/kernels/cute/neo/k4_wignerd.py rename to deepmd/pt_expt/kernels/cute/sezm/k4_wignerd.py diff --git a/deepmd/kernels/cute/neo/message_grid_gaunt_sm90.py b/deepmd/pt_expt/kernels/cute/sezm/message_grid_gaunt_sm90.py similarity index 100% rename from deepmd/kernels/cute/neo/message_grid_gaunt_sm90.py rename to deepmd/pt_expt/kernels/cute/sezm/message_grid_gaunt_sm90.py diff --git a/deepmd/kernels/cute/neo/message_grid_readout_sm90.py b/deepmd/pt_expt/kernels/cute/sezm/message_grid_readout_sm90.py similarity index 100% rename from deepmd/kernels/cute/neo/message_grid_readout_sm90.py rename to deepmd/pt_expt/kernels/cute/sezm/message_grid_readout_sm90.py diff --git a/deepmd/pt_expt/kernels/cute/sezm/operator.py b/deepmd/pt_expt/kernels/cute/sezm/operator.py deleted file mode 100644 index a5a116f301..0000000000 --- a/deepmd/pt_expt/kernels/cute/sezm/operator.py +++ /dev/null @@ -1,324 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -""" -Autograd operator and entry point wiring the fused CuTe kernels into the SeZM -SO(2) convolution value path. - -:class:`_SO2ValuePathFunction` runs the fused forward kernel (saving only the -small inputs ``x``, ``D_to_m``, ``Kc``) and, on the force path, recomputes the -value path in the fused backward kernel -- so the per-edge ``E x D_m x C`` -intermediates stay off DRAM across the whole autograd graph. :func:`make_cute_value_path` -builds the per-convolution entry :class:`_CuteSO2ValuePath`, which computes the -radial-/scalar-only tensors (``D_to_m``, ``Kc``, ``rad_feat``) in ordinary -autograd, invokes the operator to produce the pre-focus-compete local features -and the pre-mixing ``l = 0`` scalar, and applies the focus competition -- exactly -as the reference path does. The packed weights are extracted lazily on the first -call so they reflect the loaded checkpoint. -""" - -from __future__ import ( - annotations, -) - -from dataclasses import ( - dataclass, -) -from typing import ( - TYPE_CHECKING, -) - -import torch - -from .forward import ( - SEZM_CUTE_AVAILABLE, -) - -if TYPE_CHECKING: - from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import ( - EdgeCache, - ) - from deepmd.dpmodel.descriptor.dpa4_nn.so2 import ( - SO2Convolution, - ) - -if SEZM_CUTE_AVAILABLE: - from .backward import ( - BackwardRunner, - ) - from .forward import ( - ForwardRunner, - ) - -# Validated configuration of the fused operator: the deployed three-layer -# ``[gated, gated, identity]`` mixing stack in the ``lmax = 3, mmax = 1`` layout. -_SUPPORTED_LMAX = 3 -_SUPPORTED_MMAX = 1 -_SUPPORTED_LAYERS = 3 - - -@dataclass -class _PackedWeights: - """Static weights of the SO(2) value path, packed for the fused kernels. - - Attributes - ---------- - so2_w : torch.Tensor - Assembled block-diagonal SO2Linear weight per layer with shape - (L, F, D_m*Cf, D_m*Cf), in ``(in, out)`` convention. - gate_w : torch.Tensor - GatedActivation ``FocusLinear`` weight per layer with shape - (L, Cf, F, lmax*Cf); zero for non-gated layers. - has_gate : torch.Tensor - Boolean per-layer flag with shape (L,). - channel_basis : torch.Tensor - Radial degree-mixer channel basis with shape (C_wide,). - """ - - so2_w: torch.Tensor - gate_w: torch.Tensor - has_gate: torch.Tensor - channel_basis: torch.Tensor - - -def _pack_weights(conv: SO2Convolution) -> _PackedWeights: - """Extract and pack the SO(2) value-path weights from a convolution block.""" - n_layers = conv.mixing_layers - so2_w = torch.stack( - [ - conv.so2_linears[layer]._build_so2_weight().permute(1, 0, 2).contiguous() - for layer in range(n_layers) - ] - ) - gate_w, has_gate = [], [] - for layer in range(n_layers): - non_linear = conv.non_linearities[layer] - if type(non_linear).__name__ == "GatedActivation" and non_linear.lmax > 0: - gate_w.append( - non_linear.gate_linear.weight.view( - conv.so2_focus_dim, conv.n_focus, conv.lmax * conv.so2_focus_dim - ).contiguous() - ) - has_gate.append(True) - else: - gate_w.append( - torch.zeros_like(gate_w[0]) - if gate_w - else torch.zeros( - conv.so2_focus_dim, - conv.n_focus, - conv.lmax * conv.so2_focus_dim, - device=so2_w.device, - dtype=so2_w.dtype, - ) - ) - has_gate.append(False) - return _PackedWeights( - so2_w=so2_w, - gate_w=torch.stack(gate_w), - has_gate=torch.tensor(has_gate, device=so2_w.device, dtype=torch.bool), - channel_basis=conv.radial_degree_mixer.channel_basis.reshape(-1).contiguous(), - ) - - -class _SO2ValuePathFunction(torch.autograd.Function): - """Fused CuTe forward with a recompute backward for the force path.""" - - @staticmethod - def forward(ctx, x, d_to_m, kc, src, fwd_runner, bwd_runner): # noqa: ANN001, ANN205 - with torch.no_grad(): - x_local, focus_gate = fwd_runner( - x.detach(), src, d_to_m.detach(), kc.detach() - ) - ctx.save_for_backward(x, d_to_m, kc) - ctx.src = src - ctx.bwd_runner = bwd_runner - return x_local, focus_gate - - @staticmethod - def backward(ctx, grad_local, grad_focus_gate): # noqa: ANN001, ANN205 - x, d_to_m, kc = ctx.saved_tensors - need = ctx.needs_input_grad - grad_x, grad_d_to_m, grad_kc = ctx.bwd_runner( - x.detach(), - ctx.src, - d_to_m.detach(), - kc.detach(), - grad_local.detach().contiguous(), - grad_focus_gate.detach().contiguous(), - ) - return ( - grad_x if need[0] else None, - grad_d_to_m if need[1] else None, - grad_kc if need[2] else None, - None, - None, - None, - ) - - -class _CuteSO2ValuePath: - """Per-convolution entry that runs the value path through the fused kernels. - - The convolution is held by reference so the packed weights are extracted - lazily on the first call (after the checkpoint is loaded) and the kernels are - compiled on first use. - - Parameters - ---------- - conv : SO2Convolution - The owning convolution block. - bucket_fwd, threads_fwd, bucket_bwd, threads_bwd, rb, rn : int - Fused-kernel launch configuration. - """ - - def __init__( - self, - conv: SO2Convolution, - *, - bucket_fwd: int = 32, - threads_fwd: int = 1024, - bucket_bwd: int = 16, - threads_bwd: int = 512, - rb: int = 4, - rn: int = 4, - ) -> None: - self._conv = conv - self._cfg = { - "lmax": conv.lmax, - "mmax": conv.mmax, - "cf": conv.so2_focus_dim, - "n_focus": conv.n_focus, - "n_layers": conv.mixing_layers, - "rb": rb, - "rn": rn, - } - self._launch = { - "bucket_fwd": bucket_fwd, - "threads_fwd": threads_fwd, - "bucket_bwd": bucket_bwd, - "threads_bwd": threads_bwd, - } - self._fwd_runner = None - self._bwd_runner = None - - def _build(self) -> None: - weights = _pack_weights(self._conv) - self._fwd_runner = ForwardRunner( - weights, - bucket=self._launch["bucket_fwd"], - threads=self._launch["threads_fwd"], - **self._cfg, - ) - self._bwd_runner = BackwardRunner( - weights, - bucket=self._launch["bucket_bwd"], - threads=self._launch["threads_bwd"], - **self._cfg, - ) - - def __call__( - self, - x: torch.Tensor, - edge_cache: EdgeCache, - radial_feat: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Compute the SO(2) local features and radial features via the fused op. - - Parameters - ---------- - x : torch.Tensor - Node features with shape (N, D, C_wide). - edge_cache : EdgeCache - Precomputed edge cache (provides ``src`` and the Wigner ``D_full``). - radial_feat : torch.Tensor - Per-edge radial features with shape (E, lmax+1, C). - - Returns - ------- - x_local : torch.Tensor - Post-focus-compete local features with shape (E, F, D_m, Cf). - rad_feat : torch.Tensor - Projected radial features with shape (E, D_m, C_wide); its ``l = 0`` - slice is consumed by the attention aggregation. - """ - if self._fwd_runner is None: - self._build() - conv = self._conv - src = edge_cache.src - - # === Step 1. Radial-/scalar-only tensors (kept in ordinary autograd) === - d_to_m = edge_cache.D_full[ - :, : conv.ebed_dim_full, : conv.ebed_dim_full - ].index_select( - 1, - conv.coeff_index_m, - ) - rad_feat = radial_feat[:, conv.degree_index_m, :] - rad_feat = conv.radial_hidden_proj(rad_feat) - mixer = conv.radial_degree_mixer - kernel_flat = mixer._project_radial(rad_feat) - compact = kernel_flat.view(src.shape[0], mixer.degree_kernel_size, mixer.rank) - kc = mixer._scatter_rank_kernel(compact).squeeze(-1) - - # === Step 2. Fused value path -> pre-focus-compete local + l=0 scalar === - x_local, focus_gate = _SO2ValuePathFunction.apply( - x, d_to_m, kc, src, self._fwd_runner, self._bwd_runner - ) - - # === Step 3. Cross-focus softmax competition (rotation-free scalars) === - if conv.focus_compete and conv.n_focus > 1: - alpha = conv._focus_alpha(focus_gate) - x_local = x_local * alpha.to(dtype=x_local.dtype).unsqueeze(-1).unsqueeze( - -1 - ) - return x_local, rad_feat - - -def _is_supported(conv: SO2Convolution) -> bool: - """Return whether ``conv`` matches the validated fused-operator configuration.""" - if ( - conv.lmax != _SUPPORTED_LMAX - or conv.mmax != _SUPPORTED_MMAX - or conv.mixing_layers != _SUPPORTED_LAYERS - or conv.node_wise_grid_product is not None - or conv.use_so2_attn_res - or conv.layer_scale - or conv.radial_degree_mixer is None - or conv.radial_hidden_proj is None - ): - return False - if any(type(norm).__name__ != "Identity" for norm in conv.so2_inter_norms): - return False - if any(linear.bias0 is not None for linear in conv.so2_linears): - return False - non_linears = conv.non_linearities - if any( - type(non_linears[layer]).__name__ != "GatedActivation" - or ( - getattr(non_linears[layer].scalar_act, "activation", None) - or getattr(non_linears[layer], "activation_function", None) - ) - != "silu" - for layer in range(_SUPPORTED_LAYERS - 1) - ): - return False - return type(non_linears[-1]).__name__ == "Identity" - - -def make_cute_value_path(conv: SO2Convolution) -> _CuteSO2ValuePath | None: - """Build the fused CuTe value-path entry for a convolution block. - - Parameters - ---------- - conv : SO2Convolution - The convolution block to accelerate. - - Returns - ------- - _CuteSO2ValuePath or None - The entry callable when the CuTe backend is available and ``conv`` matches - the validated configuration; otherwise ``None`` (the caller falls back to - the reference path). - """ - if not SEZM_CUTE_AVAILABLE or not _is_supported(conv): - return None - return _CuteSO2ValuePath(conv) diff --git a/deepmd/kernels/cute/neo/output_grid_kernels/__init__.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/__init__.py similarity index 100% rename from deepmd/kernels/cute/neo/output_grid_kernels/__init__.py rename to deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/__init__.py diff --git a/deepmd/kernels/cute/neo/output_grid_kernels/cute_readout_l0.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/cute_readout_l0.py similarity index 100% rename from deepmd/kernels/cute/neo/output_grid_kernels/cute_readout_l0.py rename to deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/cute_readout_l0.py diff --git a/deepmd/kernels/cute/neo/output_grid_kernels/cute_tiled_grid_product.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/cute_tiled_grid_product.py similarity index 100% rename from deepmd/kernels/cute/neo/output_grid_kernels/cute_tiled_grid_product.py rename to deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/cute_tiled_grid_product.py diff --git a/deepmd/kernels/cute/neo/output_grid_product.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid_product.py similarity index 100% rename from deepmd/kernels/cute/neo/output_grid_product.py rename to deepmd/pt_expt/kernels/cute/sezm/output_grid_product.py diff --git a/deepmd/kernels/cute/neo/readout_l0.py b/deepmd/pt_expt/kernels/cute/sezm/readout_l0.py similarity index 99% rename from deepmd/kernels/cute/neo/readout_l0.py rename to deepmd/pt_expt/kernels/cute/sezm/readout_l0.py index 0d51e6c1fa..3ba6e4403b 100644 --- a/deepmd/kernels/cute/neo/readout_l0.py +++ b/deepmd/pt_expt/kernels/cute/sezm/readout_l0.py @@ -638,7 +638,7 @@ def run_neo_output_readout( candidate = maybe_run_neo_readout_l0(output_ffn, ffn_in) if candidate is not None: return candidate - return (ffn_in + output_ffn(ffn_in))[:, 0:1, :, :].reshape( + return (ffn_in[:, 0:1, :, :] + output_ffn.forward_scalar(ffn_in)).reshape( ffn_in.shape[0], output_ffn.channels ) diff --git a/deepmd/kernels/cute/neo/runtime_policy.py b/deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py similarity index 96% rename from deepmd/kernels/cute/neo/runtime_policy.py rename to deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py index 1598680c87..c0bf24e754 100644 --- a/deepmd/kernels/cute/neo/runtime_policy.py +++ b/deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py @@ -10,9 +10,12 @@ import torch +from deepmd.pt_expt.kernels.utils import ( + use_cute_infer, +) + _TRUE_VALUES = frozenset({"1", "true", "yes", "on"}) _FALSE_VALUES = frozenset({"0", "false", "no", "off"}) -NEO_CUTE_INFER_ENV = "DP_NEO_CUTE_INFER" SM80_PROFILE_CAPABILITIES = frozenset({(8, 0), (8, 6)}) SM90_CAPABILITY = (9, 0) FUSED_SO2_GATE_CAPABILITIES = frozenset({(8, 9), (12, 0)}) @@ -55,13 +58,8 @@ def _env_override(name: str) -> bool | None: def is_cute_infer_enabled() -> bool: - """Return whether the process opted into the full Neo CuTe K1 path. - - ``DP_NEO_CUTE_INFER`` is deliberately separate from the - ``DP_CUTE_INFER`` inner SO2 value-path selector. The full K1 replacement - may therefore coexist with ``DP_TRITON_INFER=2``. - """ - return _env_override(NEO_CUTE_INFER_ENV) is True + """Return whether ``DP_CUTE_INFER`` enables the SeZM CuTe path.""" + return use_cute_infer() def _current_compute_capability() -> tuple[int, int] | None: diff --git a/deepmd/kernels/cute/neo/sm90_k1/__init__.py b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/__init__.py similarity index 100% rename from deepmd/kernels/cute/neo/sm90_k1/__init__.py rename to deepmd/pt_expt/kernels/cute/sezm/sm90_k1/__init__.py diff --git a/deepmd/kernels/cute/neo/sm90_k1/final_phase_c.py b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/final_phase_c.py similarity index 100% rename from deepmd/kernels/cute/neo/sm90_k1/final_phase_c.py rename to deepmd/pt_expt/kernels/cute/sezm/sm90_k1/final_phase_c.py diff --git a/deepmd/kernels/cute/neo/sm90_k1/output_gate.py b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/output_gate.py similarity index 100% rename from deepmd/kernels/cute/neo/sm90_k1/output_gate.py rename to deepmd/pt_expt/kernels/cute/sezm/sm90_k1/output_gate.py diff --git a/deepmd/kernels/cute/neo/sm90_k1/persistent.py b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/persistent.py similarity index 100% rename from deepmd/kernels/cute/neo/sm90_k1/persistent.py rename to deepmd/pt_expt/kernels/cute/sezm/sm90_k1/persistent.py diff --git a/deepmd/kernels/cute/neo/sm90_k1/phase_a.py b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_a.py similarity index 100% rename from deepmd/kernels/cute/neo/sm90_k1/phase_a.py rename to deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_a.py diff --git a/deepmd/kernels/cute/neo/sm90_k1/phase_a_backward.py b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_a_backward.py similarity index 100% rename from deepmd/kernels/cute/neo/sm90_k1/phase_a_backward.py rename to deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_a_backward.py diff --git a/deepmd/kernels/cute/neo/sm90_k1/phase_c_attention_backward.py b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_c_attention_backward.py similarity index 100% rename from deepmd/kernels/cute/neo/sm90_k1/phase_c_attention_backward.py rename to deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_c_attention_backward.py diff --git a/deepmd/kernels/cute/neo/sm90_k1/prefix.py b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/prefix.py similarity index 100% rename from deepmd/kernels/cute/neo/sm90_k1/prefix.py rename to deepmd/pt_expt/kernels/cute/sezm/sm90_k1/prefix.py diff --git a/deepmd/kernels/cute/neo/sm90_k1/radial.py b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/radial.py similarity index 100% rename from deepmd/kernels/cute/neo/sm90_k1/radial.py rename to deepmd/pt_expt/kernels/cute/sezm/sm90_k1/radial.py diff --git a/deepmd/kernels/cute/neo/sm90_k1/runner.py b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/runner.py similarity index 100% rename from deepmd/kernels/cute/neo/sm90_k1/runner.py rename to deepmd/pt_expt/kernels/cute/sezm/sm90_k1/runner.py diff --git a/deepmd/pt_expt/kernels/utils.py b/deepmd/pt_expt/kernels/utils.py index c6e54bd17f..2387b139dd 100644 --- a/deepmd/pt_expt/kernels/utils.py +++ b/deepmd/pt_expt/kernels/utils.py @@ -329,12 +329,12 @@ def operator_available(name: str) -> bool: def use_cute_infer() -> bool: - """Return whether the opt-in CuTe inference operator is enabled. + """Return whether the opt-in SeZM CuTe inference path is enabled. The flag is controlled by the ``DP_CUTE_INFER`` environment variable and is - read at module construction time. It selects the fused CuTe SO(2) value-path - operator (an independent path from ``DP_TRITON_INFER``) and only takes effect - during inference; training always uses the dense reference path. + read before accelerated state is prepared. CuTe replaces eligible SeZM spans + and leaves unsupported shapes to the selected Triton, CUDA, or reference + path. It only takes effect during inference. Returns ------- @@ -352,11 +352,10 @@ def use_cutile_infer() -> bool: written in the ``cuda.tile`` DSL and only takes effect during inference; training always uses the dense reference path. - The path is mutually exclusive with ``DP_TRITON_INFER`` and - ``DP_CUTE_INFER``: when it is enabled no Triton kernel executes, and a - convolution whose layout it does not support falls back to the dense - reference rather than to another accelerated backend. Enabling more than one - of the three is rejected at construction. + The path is mutually exclusive with ``DP_TRITON_INFER``: when it is enabled + no Triton kernel executes, and a convolution whose layout it does not support + falls back to the dense reference. ``DP_CUTE_INFER`` remains an independent + exact-shape overlay around the enclosing SeZM block. Returns ------- diff --git a/doc/install/easy-install.md b/doc/install/easy-install.md index 0021c2720d..e8a7f9cd77 100644 --- a/doc/install/easy-install.md +++ b/doc/install/easy-install.md @@ -199,12 +199,12 @@ pip install "deepmd-kit[torch,cute]" ``` The `cute` extra installs the CUTLASS CuTe DSL and NVIDIA Alchemi Toolkit-Ops -runtime used by the optional CuTe paths. Set `DP_NEO_CUTE_INFER=1` to enable the -DPA4-Neo K1 implementation; it may coexist with `DP_TRITON_INFER=2`. -`DP_NEO_CUTE_INFER` and `DP_CUTE_INFER` are independent: the latter controls the -inner SO(2) value-path implementation. Requesting a CuTe path without its -runtime dependencies raises an error rather than silently selecting another -implementation. +runtime used by the optional CuTe paths. Set `DP_CUTE_INFER=1` to enable the +DPA4-Neo implementation in PyTorch inference. It may coexist with +`DP_TRITON_INFER`: CuTe replaces exact supported spans, while unsupported model +shapes retain the selected Triton or reference path. CuTe is not captured in a +frozen `.pt2`; requesting it without the runtime dependencies raises an error +rather than silently selecting another implementation. The supported platform includes Linux x86-64 and aarch64 with GNU C Library 2.28 or above, macOS x86-64 and arm64, and Windows x86-64. diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 8733a576e3..ef87a26e6f 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -424,14 +424,15 @@ equivalent input-file option used during training validation: | `DP_TF32_INFER` | `validating.tf32_infer` | `0` (highest) | float32 matmul precision for inference: `0` highest, `1` high, `2` medium. Higher values improve throughput but make the potential energy surface less smooth. | | `DP_AMP_INFER` | `validating.amp_infer` | off | bf16 autocast inside the descriptor interaction blocks for inference, independently of `descriptor.use_amp`. Training AMP remains controlled by `descriptor.use_amp`. Usually keeps aggregate MAE similar but can make the potential energy surface less smooth. | | `DP_TRITON_INFER` | — | `0` | Triton inference kernel level `0`-`3` (CUDA eval only, compatible with `DP_COMPILE_INFER`). Levels `1` and `2` are exact float32; level `3` trades a small accuracy margin for a substantial speedup. Detailed below. | +| `DP_CUTE_INFER` | — | off | CuTe inference kernels for supported SeZM shapes (PyTorch CUDA eval only, compatible with `DP_TRITON_INFER` and `DP_COMPILE_INFER`). Unsupported shapes retain the selected fallback. Python inference only, and **not** captured in a frozen `.pt2`. Detailed below. | | `DP_CUTILE_INFER` | — | off | cuTile inference path (CUDA eval only, compatible with `DP_COMPILE_INFER`, mutually exclusive with `DP_TRITON_INFER`). Python inference only, and **not** captured in a frozen `.pt2`. Detailed below. | | `DP_CUDA_INFER` | — | `0` | Hand-written CUDA operator level `0`-`2` (CUDA eval only, stacks on top of `DP_TRITON_INFER`). Level `1` is faster on every GPU and checkpoint measured; level `2` additionally offers the fused convolution, which routes itself per checkpoint and falls back to the level-`1` behaviour where it would not pay. Detailed below. | Accepted boolean values for the other switches are `1`/`true`/`yes`/`on` and `0`/`false`/`no`/`off`; `DP_TRITON_INFER` accepts only the numeric levels. -`DP_TRITON_INFER`, `DP_CUTILE_INFER` and `DP_CUTE_INFER` each select a complete -accelerated inference path and are mutually exclusive; enabling more than one is -rejected when the model is constructed. +`DP_TRITON_INFER` and `DP_CUTILE_INFER` select mutually exclusive complete +SO(2) paths. `DP_CUTE_INFER` is an exact-shape overlay and may coexist with +either fallback. Shell exports take precedence over the input-file options and over values written in the input; they are read when the model is constructed and changing them afterward has no effect. @@ -449,6 +450,15 @@ shipping built in. On other GPUs the kernels fall back to conservative configurations, and `dp --pt freeze` tunes the missing entries on the local GPU before exporting, a one-off sweep of a few minutes baked into the `.pt2`. +`DP_CUTE_INFER` enables CuTe kernels for the current DPA4-Neo inference +contract. The path fuses the eligible K1 interaction block, packed Wigner-D +construction, geometric initial embedding, output grid product, and scalar +readout while retaining strict float32 arithmetic. Eligibility checks cover the +model structure, tensor layout, dtype, device, and GPU capability. Any +unsupported configuration falls through to the enabled Triton, CUDA, cuTile, or +reference implementation rather than entering a partially compatible kernel. +The kernels are JIT compiled at runtime and are not captured in a frozen `.pt2`. + `DP_CUTILE_INFER` replaces the whole SeZM edge pipeline — Wigner monomials, rotate-and-mix, the gated SO(2) mixing stack, the attention aggregation and the force / virial assembly — with kernels written in the `cuda.tile` DSL. On an @@ -552,9 +562,9 @@ of being unavailable to the frozen `.pt2` route. > chosen levels and whether each came from the environment or the default are > logged at export. A CPU-targeted archive disables GPU-only inference paths > and keeps the reference CPU implementation regardless of these settings. -> `DP_CUTILE_INFER` is the exception: -> its kernels are JIT compiled at runtime and do not bake into the artifact, so -> it applies to Python inference only and has no effect on a frozen model. A frozen `.pt2` runs a forward-only +> `DP_CUTILE_INFER` and `DP_CUTE_INFER` are the exceptions: their kernels are +> JIT compiled at runtime and do not bake into the artifact, so they apply to +> Python inference only and have no effect on a frozen model. A frozen `.pt2` runs a forward-only > package, so training-time memory-saving switches do not apply to it. ### Hardware selection diff --git a/source/tests/pt/model/test_descriptor_sezm_block_ffn_grid_fusion.py b/source/tests/pt/model/test_descriptor_sezm_block_ffn_grid_fusion.py index 16a273bb19..678d8a0763 100644 --- a/source/tests/pt/model/test_descriptor_sezm_block_ffn_grid_fusion.py +++ b/source/tests/pt/model/test_descriptor_sezm_block_ffn_grid_fusion.py @@ -22,7 +22,7 @@ def test_single_branch_bypasses_router_and_accepts_fused_middle( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setenv("DP_CUTE_INFER", "1") branch = ( GridBranch( channels=2, @@ -67,7 +67,7 @@ def fused_middle( def test_single_branch_keeps_pytorch_middle_fallback( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setenv("DP_CUTE_INFER", "1") branch = ( GridBranch( channels=2, @@ -135,7 +135,7 @@ def tracked_router(value: torch.Tensor) -> torch.Tensor: return original_forward(value) monkeypatch.setattr(branch.router, "forward", tracked_router) - monkeypatch.delenv("DP_NEO_CUTE_INFER", raising=False) + monkeypatch.delenv("DP_CUTE_INFER", raising=False) routed = branch( left, right, @@ -143,7 +143,7 @@ def tracked_router(value: torch.Tensor) -> torch.Tensor: to_grid=lambda value: value, from_grid=lambda value: value, ) - monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setenv("DP_CUTE_INFER", "1") shortcut = branch( left, right, diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_compile_cache.py b/source/tests/pt/model/test_descriptor_sezm_cute_compile_cache.py index 6dde077517..560c31ca2b 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_compile_cache.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_compile_cache.py @@ -19,7 +19,7 @@ ) REPO_ROOT = Path(__file__).resolve().parents[4] -CACHE_PATH = REPO_ROOT / "deepmd/kernels/cute/neo/compile_cache.py" +CACHE_PATH = REPO_ROOT / "deepmd/pt_expt/kernels/cute/sezm/compile_cache.py" def _load_cache_module(): diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_envelope_softmax.py b/source/tests/pt/model/test_descriptor_sezm_cute_envelope_softmax.py index 06b1bdd3a9..4ae651ba60 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_envelope_softmax.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_envelope_softmax.py @@ -44,7 +44,7 @@ def test_cute_envelope_softmax_matches_eager_across_extreme_frames( logit_scale: float, edge_gate: float, ) -> None: - from deepmd.kernels.cute.neo.k1_kernels.cute_envelope_gated_softmax import ( + from deepmd.pt_expt.kernels.cute.sezm.k1_kernels.cute_envelope_gated_softmax import ( compile_envelope_softmax_forward, ) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_gie.py b/source/tests/pt/model/test_descriptor_sezm_cute_gie.py index c03da7bcc9..9a132c997e 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_gie.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_gie.py @@ -13,7 +13,7 @@ import torch -from deepmd.kernels.cute.neo import gie as gie_module +from deepmd.pt_expt.kernels.cute.sezm import gie as gie_module def _load_gie_module(): diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1.py index f5a4af7b35..02c97b0cc5 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1.py @@ -17,11 +17,6 @@ import pytest import torch -from deepmd.kernels.cute.neo import k1 as _K1 -from deepmd.kernels.cute.neo import k1_runner as _K1_RUNNER -from deepmd.kernels.cute.neo import ( - runtime_policy, -) from deepmd.pt.model.descriptor.sezm_nn.edge_cache import ( EdgeFeatureCache, _build_edge_wigner, @@ -30,6 +25,11 @@ from deepmd.pt.model.descriptor.sezm_nn.norm import ( EquivariantRMSNorm, ) +from deepmd.pt_expt.kernels.cute.sezm import k1 as _K1 +from deepmd.pt_expt.kernels.cute.sezm import k1_runner as _K1_RUNNER +from deepmd.pt_expt.kernels.cute.sezm import ( + runtime_policy, +) NeoK1BackwardWorkspace = _K1_RUNNER.NeoK1BackwardWorkspace NeoK1RuntimeConfig = _K1.NeoK1RuntimeConfig @@ -213,7 +213,7 @@ def test_sm89_and_sm120_use_combined_so2_gate( def test_combined_so2_gate_rejects_misaligned_contiguous_tensor() -> None: pytest.importorskip("cutlass.cute") pytest.importorskip("cuda.bindings.driver") - from deepmd.kernels.cute.neo.k1_kernels.cute_neo_so2_gate_combined_fwd import ( + from deepmd.pt_expt.kernels.cute.sezm.k1_kernels.cute_neo_so2_gate_combined_fwd import ( _require_16_byte_alignment, ) @@ -316,7 +316,7 @@ def test_packed_wigner_has_a_separate_backward_compatible_cache_field() -> None: def test_ineligible_wigner_build_does_not_import_cute( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setenv("DP_CUTE_INFER", "1") edge_vec = torch.tensor([[1.0, 0.0, 0.0]], device="cpu") edge_len = torch.linalg.vector_norm(edge_vec, dim=-1, keepdim=True) @@ -355,7 +355,7 @@ def test_ineligible_runtime_contracts_fall_back( node_count: int, destinations_sorted: bool, ) -> None: - monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setenv("DP_CUTE_INFER", "1") block = _neo_like_block().eval() assert not _K1.is_neo_k1_runtime_eligible( @@ -372,7 +372,7 @@ def test_ineligible_runtime_contracts_fall_back( def test_runtime_contract_allows_fewer_edges_than_nodes( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setenv("DP_CUTE_INFER", "1") block = _neo_like_block().eval() assert _K1.is_neo_k1_runtime_eligible( @@ -402,7 +402,7 @@ def _cute_cuda_runtime_available() -> bool: reason="CuTe attention-prelude regression requires CUDA and CuTe DSL", ) def test_attention_prelude_initializes_all_qk_rows_when_edges_are_sparse() -> None: - from deepmd.kernels.cute.neo.k1_kernels.cute_neo_focus_src_backward import ( + from deepmd.pt_expt.kernels.cute.sezm.k1_kernels.cute_neo_focus_src_backward import ( compile_neo_attention_prelude_forward, ) @@ -478,7 +478,7 @@ def test_attention_prelude_initializes_all_qk_rows_when_edges_are_sparse() -> No def test_exact_runtime_contract_is_eligible_without_autocast( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setenv("DP_CUTE_INFER", "1") monkeypatch.setattr(runtime_policy, "uses_strict_fp32_matmul", lambda: True) block = _neo_like_block().eval() kwargs = { @@ -498,7 +498,7 @@ def test_exact_runtime_contract_is_eligible_without_autocast( def test_runtime_contract_rejects_tf32_matmul_policy( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setenv("DP_CUTE_INFER", "1") monkeypatch.setattr(torch, "is_autocast_enabled", lambda device_type: False) block = _neo_like_block().eval() kwargs = { diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_dst_ptr.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_dst_ptr.py index 662cb29ffa..ea4f573c89 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_dst_ptr.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_dst_ptr.py @@ -29,7 +29,7 @@ REPO_ROOT = Path(__file__).resolve().parents[4] -K1_PATH = REPO_ROOT / "deepmd/kernels/cute/neo/k1.py" +K1_PATH = REPO_ROOT / "deepmd/pt_expt/kernels/cute/sezm/k1.py" EDGE_CACHE_PATH = REPO_ROOT / "deepmd/pt/model/descriptor/sezm_nn/edge_cache.py" @@ -264,7 +264,7 @@ def build_ptr(dst): ) def test_compile_cold_registration_survives_into_packed_custom_op_runtime(self): assert torch is not None - k1 = importlib.import_module("deepmd.kernels.cute.neo.k1") + k1 = importlib.import_module("deepmd.pt_expt.kernels.cute.sezm.k1") prior_registry = dict(k1._REGISTRY) prior_next_handle = k1._NEXT_HANDLE diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py index de39840bd7..e939d57b8e 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py @@ -10,7 +10,7 @@ import torch -from deepmd.kernels.cute.neo import ( +from deepmd.pt_expt.kernels.cute.sezm import ( k1_gate_structural, ) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py index 1e348b9018..82b559180b 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py @@ -13,7 +13,7 @@ TorchDispatchMode, ) -from deepmd.kernels.cute.neo import ( +from deepmd.pt_expt.kernels.cute.sezm import ( k1_so2linear, ) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_manual_adjoints.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_manual_adjoints.py index 85d96d195a..28061d6cd2 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_manual_adjoints.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_manual_adjoints.py @@ -14,9 +14,6 @@ import pytest import torch -from deepmd.kernels.cute.neo import ( - k1, -) from deepmd.pt.model.descriptor.sezm_nn.activation import ( SwiGLU, ) @@ -32,6 +29,9 @@ from deepmd.pt.utils import ( env, ) +from deepmd.pt_expt.kernels.cute.sezm import ( + k1, +) @pytest.fixture(autouse=True) @@ -196,10 +196,10 @@ def test_sm90_message_grid_forward_and_adjoint_match_real_module() -> None: if tuple(torch.cuda.get_device_capability()) != (9, 0): pytest.skip("SM90 message-grid differential requires an SM90 GPU") - from deepmd.kernels.cute.neo.k1_message_grid_packed import ( + from deepmd.pt_expt.kernels.cute.sezm.k1_message_grid_packed import ( run_packed_message_grid_forward, ) - from deepmd.kernels.cute.neo.message_grid_readout_sm90 import ( + from deepmd.pt_expt.kernels.cute.sezm.message_grid_readout_sm90 import ( prepare_sm90_message_grid_state, run_sm90_message_grid_backward, ) @@ -301,7 +301,7 @@ def test_sm90_message_grid_forward_and_adjoint_match_real_module() -> None: reason="Q/K CuTe adjoint regression requires CUDA and CuTe DSL", ) def test_qk_manual_adjoint_matches_autograd_with_sparse_edges() -> None: - from deepmd.kernels.cute.neo.k1_kernels.cute_neo_qk_edge import ( + from deepmd.pt_expt.kernels.cute.sezm.k1_kernels.cute_neo_qk_edge import ( compile_neo_qk_edge_backward, compile_neo_qk_node_input_adjoint, ) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_packed_impl.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_packed_impl.py index acc4e94659..459cd26dfa 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_packed_impl.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_packed_impl.py @@ -39,7 +39,7 @@ def _randn( def _dense_wigner(edge_count: int, generator: torch.Generator) -> torch.Tensor: - from deepmd.kernels.cute.neo import k1_wigner_layout as layout + from deepmd.pt_expt.kernels.cute.sezm import k1_wigner_layout as layout dense = torch.zeros( edge_count, @@ -70,13 +70,13 @@ def _dense_wigner(edge_count: int, generator: torch.Generator) -> torch.Tensor: reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable", ) def test_panel_native_quaternion_backward_matches_dense_torch(): - from deepmd.kernels.cute.neo import k1_wigner_layout as layout - from deepmd.kernels.cute.neo import ( - k4_wignerd, - ) from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( WignerDCalculator, ) + from deepmd.pt_expt.kernels.cute.sezm import k1_wigner_layout as layout + from deepmd.pt_expt.kernels.cute.sezm import ( + k4_wignerd, + ) generator = torch.Generator(device="cuda").manual_seed(20260702) q_value = _randn(3, 4, generator=generator, scale=1.0) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py index 4197bf4163..e8bc32e950 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py @@ -54,7 +54,7 @@ def _phase_c_api(): pytest.importorskip("cutlass") - from deepmd.kernels.cute.neo.k1_kernels.cute_neo_phase_c_backward_layout_runner import ( + from deepmd.pt_expt.kernels.cute.sezm.k1_kernels.cute_neo_phase_c_backward_layout_runner import ( CuteNeoPhaseCBackwardLayout, NeoPhaseCBackwardLayoutOutputs, ) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_focus_fusion.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_focus_fusion.py index afa6e94f7b..bdccca901d 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_focus_fusion.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_focus_fusion.py @@ -19,7 +19,7 @@ def _has_supported_gpu() -> bool: def test_source_csr_preserves_equal_source_order() -> None: - from deepmd.kernels.cute.neo.k1_radial_phase_a_node import ( + from deepmd.pt_expt.kernels.cute.sezm.k1_radial_phase_a_node import ( build_source_csr, ) @@ -35,7 +35,7 @@ def test_source_csr_preserves_equal_source_order() -> None: @unittest.skipUnless(_has_supported_gpu(), "requires an SM80 or SM90 CUDA device") class TestRadialFocusSourceFusion(unittest.TestCase): def test_fused_load_matches_standalone_scalar_lane_add(self) -> None: - from deepmd.kernels.cute.neo.k1_radial_phase_a_node import ( + from deepmd.pt_expt.kernels.cute.sezm.k1_radial_phase_a_node import ( build_source_csr, prepare_batched_radial_projection_weight, run_neo_radial_phase_a_backward_node_tiled, diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_pitch68_safe.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_pitch68_safe.py index 70771144f1..9815909e37 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_pitch68_safe.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_pitch68_safe.py @@ -13,7 +13,7 @@ ) REPO_ROOT = Path(__file__).resolve().parents[4] -CUTE_PATH = REPO_ROOT / "deepmd/kernels/cute/neo" +CUTE_PATH = REPO_ROOT / "deepmd/pt_expt/kernels/cute/sezm" WRAPPER_PATH = CUTE_PATH / "k1_radial_phase_a_node.py" KERNEL_PATH = CUTE_PATH / "k1_kernels/cute_neo_radial_phase_a_backward_node.py" diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py index 6d382e030d..495838263f 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py @@ -16,8 +16,10 @@ import torch -from deepmd.kernels.cute.neo import k1_radial_phase_a_node as _RADIAL_PROJECTION -from deepmd.kernels.cute.neo import k1_runner as _K1_RUNNER +from deepmd.pt_expt.kernels.cute.sezm import ( + k1_radial_phase_a_node as _RADIAL_PROJECTION, +) +from deepmd.pt_expt.kernels.cute.sezm import k1_runner as _K1_RUNNER COMPACT_WIDTH = _RADIAL_PROJECTION.COMPACT_WIDTH FOCUS_COUNT = _RADIAL_PROJECTION.FOCUS_COUNT diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py index 92d153b090..b8009a1e7e 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py @@ -17,7 +17,7 @@ ) REPO_ROOT = Path(__file__).resolve().parents[4] -CUTE_ROOT = REPO_ROOT / "deepmd/kernels/cute/neo" +CUTE_ROOT = REPO_ROOT / "deepmd/pt_expt/kernels/cute/sezm" KERNEL_PATH = CUTE_ROOT / "k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py" DYNAMIC_EDGE_COUNTS = (0, 1, 7, 8, 9, 31, 32, 33) @@ -223,7 +223,9 @@ def constant_value(node: ast.expr, namespace: dict[str, object]) -> object: def _load_wrapper_module(): - return importlib.import_module("deepmd.kernels.cute.neo.k1_gate_structural") + return importlib.import_module( + "deepmd.pt_expt.kernels.cute.sezm.k1_gate_structural" + ) class TestSM80StructuralGateVec4Contract(unittest.TestCase): @@ -403,7 +405,7 @@ def test_empty_tensor_does_not_launch(self): class TestSM80StructuralGateVec4CudaDifferential(unittest.TestCase): @classmethod def setUpClass(cls): - from deepmd.kernels.cute.neo.k1_kernels.cute_neo_gate_split_structural_vec4_sm80 import ( + from deepmd.pt_expt.kernels.cute.sezm.k1_kernels.cute_neo_gate_split_structural_vec4_sm80 import ( compile_neo_gate_split_structural_vec4_sm80_backward, compile_neo_gate_split_structural_vec4_sm80_forward, ) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_memory_reuse.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_memory_reuse.py index 76b4ae5b55..f0fa392dc9 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_memory_reuse.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_memory_reuse.py @@ -25,8 +25,8 @@ ) REPO_ROOT = Path(__file__).resolve().parents[4] -STRUCTURAL_HELPER = REPO_ROOT / "deepmd/kernels/cute/neo/k1_gate_structural.py" -SO2_HELPER = REPO_ROOT / "deepmd/kernels/cute/neo/k1_so2linear.py" +STRUCTURAL_HELPER = REPO_ROOT / "deepmd/pt_expt/kernels/cute/sezm/k1_gate_structural.py" +SO2_HELPER = REPO_ROOT / "deepmd/pt_expt/kernels/cute/sezm/k1_so2linear.py" def _load_module(name: str, path: Path): diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_output_grid_product.py b/source/tests/pt/model/test_descriptor_sezm_cute_output_grid_product.py index 00335a9258..8b31958dd5 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_output_grid_product.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_output_grid_product.py @@ -58,7 +58,7 @@ def _reference(left, right, to_grid, from_grid): return out.reshape_as(left) -def test_grid_mlp_accepts_fused_middle_callback(): +def test_grid_mlp_accepts_fused_pair_callback(): from deepmd.pt.model.descriptor.sezm_nn.grid_net import ( GridMLP, ) @@ -76,7 +76,7 @@ def test_grid_mlp_accepts_fused_middle_callback(): scalar_pair = torch.empty(2, 1, 4, device="cpu") calls = 0 - def fused_middle(projected_left, projected_right): + def fused_pair(projected_left, projected_right): nonlocal calls calls += 1 return projected_left * projected_right @@ -87,7 +87,7 @@ def fused_middle(projected_left, projected_right): scalar_pair, to_grid=lambda value: value, from_grid=lambda value: value, - grid_product=fused_middle, + pair_grid=fused_pair, ) assert calls == 1 @@ -99,7 +99,7 @@ def test_dispatch_falls_back_without_exact_cuda_contract( monkeypatch: pytest.MonkeyPatch, hidden_channels: int, ): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( output_grid_product, ) @@ -133,7 +133,7 @@ def test_exact_shape_guard_accepts_only_validated_widths( hidden_channels: int, expected: int | None, ) -> None: - from deepmd.kernels.cute.neo.output_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid_product import ( _exact_hidden_channels, ) @@ -143,7 +143,7 @@ def test_exact_shape_guard_accepts_only_validated_widths( @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) def test_fake_registrations_allocate_canonical_strides(hidden_channels: int) -> None: - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( output_grid_product, ) @@ -205,12 +205,12 @@ def test_sm90_c96_asymmetric_panel_dispatch_is_explicit( policy_enabled: bool, expected: bool, ) -> None: - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( output_grid_product, ) kernel_module_name = ( - "deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product" + "deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product" ) kernel_module = types.ModuleType(kernel_module_name) calls: dict[str, dict[str, bool]] = {} @@ -339,7 +339,7 @@ def test_forward_and_first_backward_match_strict_fp32( nodes: int, hidden_channels: int, ): - from deepmd.kernels.cute.neo.output_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid_product import ( output_grid_product_cute, ) @@ -387,7 +387,7 @@ def test_sm80_c96_n48_forward_custom_op_matches_strict_fp32( monkeypatch: pytest.MonkeyPatch, nodes: int, ): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( output_grid_product, ) @@ -415,7 +415,7 @@ def test_sm80_c96_n48_panel_custom_op_matches_strict_fp32( self, monkeypatch: pytest.MonkeyPatch, ): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( output_grid_product, ) @@ -464,7 +464,7 @@ def test_dispatch_uses_only_the_master_gate( monkeypatch: pytest.MonkeyPatch, hidden_channels: int, ): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( output_grid_product, ) @@ -511,7 +511,7 @@ def test_dispatch_declines_non_strict_matmul_state( monkeypatch: pytest.MonkeyPatch, hidden_channels: int, ): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( output_grid_product, ) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_readout_l0.py b/source/tests/pt/model/test_descriptor_sezm_cute_readout_l0.py index 93fd80108e..eef81470e5 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_readout_l0.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_readout_l0.py @@ -185,19 +185,17 @@ def _reference_product(left, right, to_grid, from_grid): ) -def test_sm80_readout_input_fold_selector_uses_only_full_neo_gate(monkeypatch): - from deepmd.kernels.cute.neo import ( +def test_sm80_readout_input_fold_selector_uses_cute_gate(monkeypatch): + from deepmd.pt_expt.kernels.cute.sezm import ( runtime_policy, ) - monkeypatch.delenv("DP_NEO_CUTE_INFER", raising=False) - monkeypatch.setenv("DP_CUTE_INFER", "1") + monkeypatch.delenv("DP_CUTE_INFER", raising=False) monkeypatch.delenv("DP_CUTE_READOUT_INPUT_FOLD_SM80", raising=False) assert not runtime_policy.is_cute_infer_enabled() assert not runtime_policy.is_readout_input_fold_sm80_enabled((8, 0)) - monkeypatch.setenv("DP_CUTE_INFER", "0") - monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setenv("DP_CUTE_INFER", "1") assert runtime_policy.is_cute_infer_enabled() assert runtime_policy.is_readout_input_fold_sm80_enabled((8, 0)) monkeypatch.setenv("DP_CUTE_READOUT_INPUT_FOLD_SM80", "0") @@ -207,18 +205,16 @@ def test_sm80_readout_input_fold_selector_uses_only_full_neo_gate(monkeypatch): assert not runtime_policy.is_readout_input_fold_sm80_enabled((9, 0)) -def test_sm90_readout_input_fold_selector_uses_only_full_neo_gate(monkeypatch): - from deepmd.kernels.cute.neo import ( +def test_sm90_readout_input_fold_selector_uses_cute_gate(monkeypatch): + from deepmd.pt_expt.kernels.cute.sezm import ( runtime_policy, ) - monkeypatch.delenv("DP_NEO_CUTE_INFER", raising=False) - monkeypatch.setenv("DP_CUTE_INFER", "1") + monkeypatch.delenv("DP_CUTE_INFER", raising=False) monkeypatch.delenv("DP_CUTE_READOUT_INPUT_FOLD_SM90", raising=False) assert not runtime_policy.is_readout_input_fold_sm90_enabled((9, 0)) - monkeypatch.setenv("DP_CUTE_INFER", "0") - monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setenv("DP_CUTE_INFER", "1") assert runtime_policy.is_readout_input_fold_sm90_enabled((9, 0)) assert runtime_policy.is_readout_input_fold_enabled((9, 0)) monkeypatch.setenv("DP_CUTE_READOUT_INPUT_FOLD_SM90", "0") @@ -229,7 +225,7 @@ def test_sm90_readout_input_fold_selector_uses_only_full_neo_gate(monkeypatch): def test_readout_input_fold_guard_accepts_validated_architectures(monkeypatch): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, runtime_policy, ) @@ -257,7 +253,7 @@ def select_readout(capability): def test_sm80_readout_input_fold_matches_forward_and_input_vjp(monkeypatch): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -321,7 +317,7 @@ def compact_reference(left, right, to_grid, from_grid): def test_sm80_readout_input_fold_is_fullgraph_safe(monkeypatch): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -350,7 +346,7 @@ def folded(value): def test_sm80_readout_input_fold_cache_refreshes_after_inplace_change(monkeypatch): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -383,7 +379,7 @@ def test_sm80_readout_input_fold_cache_refreshes_after_inplace_change(monkeypatc def test_sm80_readout_input_fold_cache_refreshes_after_parameter_replacement(): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -402,7 +398,7 @@ def test_sm80_readout_input_fold_cache_refreshes_after_parameter_replacement(): @pytest.mark.parametrize("assign", [False, True]) def test_sm80_readout_input_fold_load_state_dict_invalidates_cache(assign): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -425,7 +421,7 @@ def test_sm80_readout_input_fold_load_state_dict_invalidates_cache(assign): def test_sm80_readout_input_fold_buffers_follow_device_moves(): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -445,7 +441,7 @@ def test_sm80_readout_input_fold_buffers_follow_device_moves(): def test_sm80_readout_input_fold_compile_requires_explicit_preparation( monkeypatch, ): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -458,7 +454,7 @@ def test_sm80_readout_input_fold_compile_requires_explicit_preparation( def test_sm80_readout_input_fold_compiled_lookup_returns_prepared_buffers( monkeypatch, ): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -479,7 +475,7 @@ def test_shape_guard_accepts_only_exact_c192_readout( hidden_channels: int, expected: bool, ): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -488,7 +484,7 @@ def test_shape_guard_accepts_only_exact_c192_readout( def test_fake_registrations_return_canonical_strides(): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -518,7 +514,7 @@ def test_custom_ops_use_canonical_fake_metadata(): FakeTensorMode, ) - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -545,7 +541,7 @@ def test_custom_ops_use_canonical_fake_metadata(): def test_reference_completion_matches_module_output_and_input_vjp(): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -567,7 +563,7 @@ def test_reference_completion_matches_module_output_and_input_vjp(): def test_exact_structure_and_frozen_inference_guards(): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -585,14 +581,14 @@ def test_exact_structure_and_frozen_inference_guards(): def test_module_boundary_falls_back_when_device_is_unsupported(monkeypatch): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) module = _output_ffn() value = torch.randn(2, COEFF_DIM, 1, 32, device="cpu") expected = (value + module(value))[:, 0, 0, :] - monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setenv("DP_CUTE_INFER", "1") assert readout_l0.maybe_run_neo_readout_l0(module, value) is None torch.testing.assert_close( @@ -601,7 +597,7 @@ def test_module_boundary_falls_back_when_device_is_unsupported(monkeypatch): def test_descriptor_freeze_guard_and_readout_wiring(cpu_neo_descriptor, monkeypatch): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -625,7 +621,7 @@ def record_readout(output_ffn, ffn_in, *, parameters_frozen): def test_trainable_descriptor_state_bypasses_candidate(monkeypatch): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -651,7 +647,7 @@ def test_legacy_matmul_precision_guard_is_eager_and_fullgraph_safe( precision: str, expected: bool, ): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -680,7 +676,7 @@ def test_modern_matmul_precision_guard_rejects_tf32( precision: str, expected: bool, ): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -724,7 +720,7 @@ def test_forward_and_input_vjp_match_strict_fp32( self, nodes: int, ): - from deepmd.kernels.cute.neo.readout_l0 import ( + from deepmd.pt_expt.kernels.cute.sezm.readout_l0 import ( readout_l0_product_cute, ) @@ -749,7 +745,7 @@ def test_forward_and_input_vjp_match_strict_fp32( ) def test_opcheck_and_fullgraph_preserve_metadata_and_vjp(self): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -774,7 +770,7 @@ def test_opcheck_and_fullgraph_preserve_metadata_and_vjp(self): assert grads[1].shape == right.shape def test_exact_module_fullgraph_uses_full_neo_gate(self, monkeypatch): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -785,11 +781,11 @@ def test_exact_module_fullgraph_uses_full_neo_gate(self, monkeypatch): seed = torch.randn(3, 32, device="cuda") value_ref = value.detach().clone().requires_grad_(True) value_actual = value.detach().clone().requires_grad_(True) - monkeypatch.delenv("DP_NEO_CUTE_INFER", raising=False) + monkeypatch.delenv("DP_CUTE_INFER", raising=False) expected = (value_ref + module(value_ref))[:, 0, 0, :] expected_grad = torch.autograd.grad(expected, value_ref, seed)[0] - monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setenv("DP_CUTE_INFER", "1") readout_l0.prepare_sm80_readout_input_fold(module) compiled = torch.compile( lambda tensor: readout_l0.run_neo_output_readout(module, tensor), @@ -803,7 +799,7 @@ def test_exact_module_fullgraph_uses_full_neo_gate(self, monkeypatch): torch.testing.assert_close(actual_grad, expected_grad, atol=TOL, rtol=TOL) def test_sm80_input_fold_preparation_is_cross_stream_ready(self): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -827,7 +823,7 @@ def test_sm80_input_fold_preparation_is_cross_stream_ready(self): torch.testing.assert_close(actual, reference, atol=0.0, rtol=0.0) def test_sm80_input_fold_fullgraph_matches_module_vjp(self, monkeypatch): - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( readout_l0, ) @@ -844,7 +840,7 @@ def test_sm80_input_fold_fullgraph_matches_module_vjp(self, monkeypatch): expected = (value_ref + module(value_ref))[:, 0, 0, :] expected_grad = torch.autograd.grad(expected, value_ref, seed)[0] - monkeypatch.setenv("DP_NEO_CUTE_INFER", "1") + monkeypatch.setenv("DP_CUTE_INFER", "1") monkeypatch.setenv("DP_CUTE_READOUT_INPUT_FOLD_SM80", "1") assert readout_l0._can_use_sm80_readout_input_fold(module, value_actual) readout_l0.prepare_sm80_readout_input_fold(module) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_runtime_policy.py b/source/tests/pt/model/test_descriptor_sezm_cute_runtime_policy.py index 89a7be4c92..c966f3adac 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_runtime_policy.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_runtime_policy.py @@ -17,7 +17,7 @@ ) REPO_ROOT = Path(__file__).resolve().parents[4] -POLICY_PATH = REPO_ROOT / "deepmd/kernels/cute/neo/runtime_policy.py" +POLICY_PATH = REPO_ROOT / "deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py" def _load_policy(): @@ -34,15 +34,15 @@ def _load_policy(): return module -def test_full_neo_master_is_independent_from_inner_cute_selector() -> None: +def test_cute_master_gate_controls_sezm_path() -> None: policy = _load_policy() cases = ( ({}, False), - ({"DP_CUTE_INFER": "1"}, False), - ({"DP_NEO_CUTE_INFER": "1"}, True), - ({"DP_CUTE_INFER": "0", "DP_NEO_CUTE_INFER": "1"}, True), - ({"DP_CUTE_INFER": "1", "DP_NEO_CUTE_INFER": "0"}, False), - ({"DP_NEO_CUTE_INFER": "1", "DP_TRITON_INFER": "2"}, True), + ({"DP_CUTE_INFER": "0"}, False), + ({"DP_CUTE_INFER": "1"}, True), + ({"DP_CUTE_INFER": "true"}, True), + ({"DP_CUTE_INFER": "unsupported"}, False), + ({"DP_CUTE_INFER": "1", "DP_TRITON_INFER": "2"}, True), ) for environment, expected in cases: with mock.patch.dict(os.environ, environment, clear=True): @@ -76,7 +76,7 @@ def test_master_switch_controls_every_sm80_subfeature() -> None: def test_sm80_family_profile_defaults_from_master_only() -> None: policy = _load_policy() - with mock.patch.dict(os.environ, {"DP_NEO_CUTE_INFER": "1"}, clear=True): + with mock.patch.dict(os.environ, {"DP_CUTE_INFER": "1"}, clear=True): assert policy.SM80_PROFILE_CAPABILITIES == frozenset({(8, 0), (8, 6)}) for capability in policy.SM80_PROFILE_CAPABILITIES: assert policy.is_sm80_profile_enabled(capability) @@ -107,7 +107,7 @@ def test_every_sm80_profile_feature_can_be_disabled() -> None: for name, checker in checks.items(): with mock.patch.dict( os.environ, - {"DP_NEO_CUTE_INFER": "1", name: "0"}, + {"DP_CUTE_INFER": "1", name: "0"}, clear=True, ): assert not checker() @@ -115,7 +115,7 @@ def test_every_sm80_profile_feature_can_be_disabled() -> None: def test_non_sm80_family_uses_architecture_defaults_without_sm80_features() -> None: policy = _load_policy() - with mock.patch.dict(os.environ, {"DP_NEO_CUTE_INFER": "1"}, clear=True): + with mock.patch.dict(os.environ, {"DP_CUTE_INFER": "1"}, clear=True): capability = (8, 9) assert not policy.is_sm80_profile_enabled(capability) assert not policy.is_gie_enabled(capability) @@ -130,14 +130,14 @@ def test_sm90_c96_asymmetric_panels_default_is_exact_arch_disable_only() -> None policy = _load_policy() switch = policy.OUTPUT_GRID_SM90_C96_ASYMMETRIC_PANELS_ENV - with mock.patch.dict(os.environ, {"DP_NEO_CUTE_INFER": "1"}, clear=True): + with mock.patch.dict(os.environ, {"DP_CUTE_INFER": "1"}, clear=True): assert policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((9, 0)) assert not policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((8, 9)) assert not policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((9, 1)) with mock.patch.dict( os.environ, - {"DP_NEO_CUTE_INFER": "1", switch: "0"}, + {"DP_CUTE_INFER": "1", switch: "0"}, clear=True, ): assert not policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((9, 0)) @@ -147,7 +147,7 @@ def test_sm90_c96_asymmetric_panels_default_is_exact_arch_disable_only() -> None with mock.patch.dict( os.environ, - {"DP_NEO_CUTE_INFER": "1", switch: "1"}, + {"DP_CUTE_INFER": "1", switch: "1"}, clear=True, ): assert policy.is_output_grid_sm90_c96_asymmetric_panels_enabled((9, 0)) @@ -159,7 +159,7 @@ def test_overrides_remain_architecture_safe() -> None: with mock.patch.dict( os.environ, { - "DP_NEO_CUTE_INFER": "1", + "DP_CUTE_INFER": "1", "DP_CUTE_K1_THIN_WRAPPER": "1", "DP_CUTE_OUTPUT_GRID_BWD_SM80_C96_N48_PANEL": "1", "DP_CUTE_OUTPUT_GRID_FWD_SM80_C96_N48": "1", @@ -175,7 +175,7 @@ def test_overrides_remain_architecture_safe() -> None: with mock.patch.dict( os.environ, { - "DP_NEO_CUTE_INFER": "1", + "DP_CUTE_INFER": "1", "DP_CUTE_GIE": "1", "DP_CUTE_K1_PACKED_WIGNER": "1", }, @@ -224,12 +224,12 @@ def test_neighbor_list_eager_island_policy_is_owned_by_neo_runtime() -> None: policy = _load_policy() cases = ( ({}, (8, 0), False), - ({"DP_NEO_CUTE_INFER": "1"}, (8, 0), True), - ({"DP_NEO_CUTE_INFER": "1"}, (8, 6), False), - ({"DP_NEO_CUTE_INFER": "1"}, (9, 0), False), + ({"DP_CUTE_INFER": "1"}, (8, 0), True), + ({"DP_CUTE_INFER": "1"}, (8, 6), False), + ({"DP_CUTE_INFER": "1"}, (9, 0), False), ( { - "DP_NEO_CUTE_INFER": "1", + "DP_CUTE_INFER": "1", "DP_CUTE_K1_EAGER_ISLANDS": "0", }, (8, 0), @@ -237,7 +237,7 @@ def test_neighbor_list_eager_island_policy_is_owned_by_neo_runtime() -> None: ), ( { - "DP_NEO_CUTE_INFER": "1", + "DP_CUTE_INFER": "1", "DP_CUTE_K1_EAGER_ISLANDS": "1", }, (9, 0), diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_tiled_grid_product.py b/source/tests/pt/model/test_descriptor_sezm_cute_tiled_grid_product.py index f8ef09c246..3108b5cf6a 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_tiled_grid_product.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_tiled_grid_product.py @@ -22,12 +22,12 @@ REPO_ROOT = Path(__file__).resolve().parents[4] TILED_KERNEL_PATH = ( REPO_ROOT - / "deepmd/kernels/cute/neo/output_grid_kernels" + / "deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels" / "cute_tiled_grid_product.py" ) MESSAGE_GRID_PATH = ( REPO_ROOT - / "deepmd/kernels/cute/neo/k1_kernels" + / "deepmd/pt_expt/kernels/cute/sezm/k1_kernels" / "cute_neo_message_grid_product.py" ) @@ -184,7 +184,7 @@ def _inputs(nodes: int, hidden_channels: int): def test_architecture_policy_shares_sm80_backend_with_sm86(): - from deepmd.kernels.cute.neo.runtime_policy import ( + from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import ( PORTABLE_TILED_BACKEND, PYTORCH_BACKEND, output_grid_arch_key, @@ -226,7 +226,7 @@ def test_matches_strict_fp32_reference( nodes: int, hidden_channels: int, ): - from deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product import ( run_tiled_output_grid_product, ) @@ -241,7 +241,7 @@ def test_matches_strict_fp32_reference( ) @pytest.mark.parametrize("nodes", [1, 7, 65]) def test_sm80_c96_n48_matches_strict_fp32_reference(self, nodes: int): - from deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product import ( run_tiled_output_grid_product, ) @@ -258,7 +258,7 @@ def test_sm80_c96_n48_matches_strict_fp32_reference(self, nodes: int): @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) def test_one_compile_accepts_symbolic_node_counts(self, hidden_channels: int): - from deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product import ( _compiled_tiled_forward, run_tiled_output_grid_product, ) @@ -284,7 +284,7 @@ def test_one_compile_accepts_symbolic_node_counts(self, hidden_channels: int): @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) def test_launches_on_current_non_default_stream(self, hidden_channels: int): - from deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product import ( run_tiled_output_grid_product, ) @@ -314,7 +314,7 @@ def test_first_backward_matches_strict_fp32_reference( nodes: int, hidden_channels: int, ): - from deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product import ( run_tiled_output_grid_product_backward, ) @@ -340,7 +340,7 @@ def test_sm80_c96_n48_panel_matches_strict_fp32_reference( self, nodes: int, ): - from deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product import ( run_tiled_output_grid_product_backward, ) @@ -363,7 +363,7 @@ def test_one_backward_compile_accepts_symbolic_node_counts( self, hidden_channels: int, ): - from deepmd.kernels.cute.neo.output_grid_kernels.cute_tiled_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product import ( _compiled_tiled_backward, run_tiled_output_grid_product_backward, ) diff --git a/source/tests/pt/model/test_descriptor_sezm_readout_l0_algebra.py b/source/tests/pt/model/test_descriptor_sezm_readout_l0_algebra.py index e4b83402dd..a839b7dfc8 100644 --- a/source/tests/pt/model/test_descriptor_sezm_readout_l0_algebra.py +++ b/source/tests/pt/model/test_descriptor_sezm_readout_l0_algebra.py @@ -121,7 +121,7 @@ def test_scalar_readout_matches_row_zero_of_generic_grid_product(): def test_dense_gram_matches_both_projected_input_adjoints(): - from deepmd.kernels.cute.neo.readout_l0 import ( + from deepmd.pt_expt.kernels.cute.sezm.readout_l0 import ( build_readout_l0_gram, ) @@ -184,7 +184,7 @@ def test_dense_gram_matches_both_projected_input_adjoints(): def test_dense_gram_forward_matches_row_zero_grid_projection(): - from deepmd.kernels.cute.neo.readout_l0 import ( + from deepmd.pt_expt.kernels.cute.sezm.readout_l0 import ( build_readout_l0_gram, ) diff --git a/source/tests/pt/model/test_mlp.py b/source/tests/pt/model/test_mlp.py index 1f141301a5..91a283bfd8 100644 --- a/source/tests/pt/model/test_mlp.py +++ b/source/tests/pt/model/test_mlp.py @@ -106,7 +106,7 @@ def test_thin_k1_eval_linear_traces_without_addmm(self) -> None: mlp_module.enable_neo_cute_compile_visible_linears(layer) value = torch.randn(7, 5, device=env.DEVICE, dtype=torch.float32) environment = { - "DP_NEO_CUTE_INFER": "1", + "DP_CUTE_INFER": "1", "DP_CUTE_K1_THIN_WRAPPER": "1", } @@ -139,7 +139,7 @@ def test_thin_k1_eval_linear_is_scoped_to_marked_model(self) -> None: layer.eval() value = torch.randn(7, 5, device=env.DEVICE, dtype=torch.float32) environment = { - "DP_NEO_CUTE_INFER": "1", + "DP_CUTE_INFER": "1", "DP_CUTE_K1_THIN_WRAPPER": "1", } @@ -160,7 +160,7 @@ def test_thin_k1_linear_defaults_on_only_for_sm80(self) -> None: with ( mock.patch.dict( os.environ, - {"DP_NEO_CUTE_INFER": "1"}, + {"DP_CUTE_INFER": "1"}, clear=True, ), mock.patch.object(torch.cuda, "is_available", return_value=True), @@ -172,7 +172,7 @@ def test_thin_k1_linear_defaults_on_only_for_sm80(self) -> None: mock.patch.dict( os.environ, { - "DP_NEO_CUTE_INFER": "1", + "DP_CUTE_INFER": "1", "DP_CUTE_K1_THIN_WRAPPER": "0", }, clear=True, diff --git a/source/tests/pt/model/test_sezm_cute_efs.py b/source/tests/pt/model/test_sezm_cute_efs.py index 08e42e6ba4..e8a0bf33e4 100644 --- a/source/tests/pt/model/test_sezm_cute_efs.py +++ b/source/tests/pt/model/test_sezm_cute_efs.py @@ -196,7 +196,7 @@ def _detached_outputs(outputs: dict[str, torch.Tensor]) -> dict[str, torch.Tenso def _assert_runtime_profile() -> dict[str, bool]: - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( k1, runtime_policy, ) @@ -256,7 +256,7 @@ def _assert_runtime_profile() -> dict[str, bool]: missing = sorted(name for name, enabled in selected.items() if not enabled) if missing: raise AssertionError( - "DP_NEO_CUTE_INFER did not select the expected runtime profile: " + "DP_CUTE_INFER did not select the expected runtime profile: " + ", ".join(missing) ) return selected @@ -268,7 +268,7 @@ def _run_efs_child( use_compile: bool, output_path: Path, ) -> None: - from deepmd.kernels.cute.neo import ( + from deepmd.pt_expt.kernels.cute.sezm import ( k1, ) @@ -319,7 +319,7 @@ def counted_build_runner(*args: Any, **kwargs: Any) -> Any: def _clean_child_environment() -> dict[str, str]: environment = os.environ.copy() for name in tuple(environment): - if name.startswith(("DP_CUTE_", "DP_NEO_CUTE_")): + if name.startswith("DP_CUTE_"): environment.pop(name) for name in ( "DP_COMPILE_INFER", @@ -351,7 +351,7 @@ def _child_environment(*, use_compile: bool) -> dict[str, str]: { "DP_COMPILE_INFER": "1", "DP_CUTE_STRICT": "1", - "DP_NEO_CUTE_INFER": "1", + "DP_CUTE_INFER": "1", "DP_TRITON_INFER": "2", } ) @@ -360,7 +360,6 @@ def _child_environment(*, use_compile: bool) -> dict[str, str]: { "DP_COMPILE_INFER": "0", "DP_CUTE_INFER": "0", - "DP_NEO_CUTE_INFER": "0", "DP_TRITON_INFER": "0", } ) diff --git a/source/tests/pt/model/test_sezm_cute_multigpu_dispatch.py b/source/tests/pt/model/test_sezm_cute_multigpu_dispatch.py index 3ec7830d1e..cef11b404a 100644 --- a/source/tests/pt/model/test_sezm_cute_multigpu_dispatch.py +++ b/source/tests/pt/model/test_sezm_cute_multigpu_dispatch.py @@ -18,13 +18,13 @@ import pytest import torch -from deepmd.kernels.cute.neo import ( - k1, - runtime_policy, -) from deepmd.pt.model.network import ( mlp, ) +from deepmd.pt_expt.kernels.cute.sezm import ( + k1, + runtime_policy, +) def _cuda_operand(index: int = 1) -> SimpleNamespace: @@ -125,7 +125,7 @@ def test_k1_prepare_uses_requested_device_instead_of_current_device() -> None: def test_k4_custom_ops_are_called_directly() -> None: pytest.importorskip("cutlass.cute") - k4_wignerd = importlib.import_module("deepmd.kernels.cute.neo.k4_wignerd") + k4_wignerd = importlib.import_module("deepmd.pt_expt.kernels.cute.sezm.k4_wignerd") operand = _cuda_operand() sentinel = object() @@ -161,7 +161,7 @@ def capability(device: torch.device | None = None) -> tuple[int, int]: with ( mock.patch.dict( os.environ, - {"DP_NEO_CUTE_INFER": "1"}, + {"DP_CUTE_INFER": "1"}, clear=True, ), mock.patch.object(torch.cuda, "is_available", return_value=True), @@ -208,7 +208,7 @@ def capability(device: torch.device | None = None) -> tuple[int, int]: with ( mock.patch.dict( os.environ, - {"DP_NEO_CUTE_INFER": "1"}, + {"DP_CUTE_INFER": "1"}, clear=True, ), mock.patch.object(torch.cuda, "is_available", return_value=True), diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index 7c5b1294c0..cf3daa54f4 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -465,7 +465,7 @@ def _make_frame_with_natoms( def test_neo_cute_flag_is_noop_for_ineligible_model(self) -> None: """An unsupported SeZM shape must remain bit-identical with the flag on.""" - from deepmd.kernels.cute.neo import k1 as cute_k1 + from deepmd.pt_expt.kernels.cute.sezm import k1 as cute_k1 cpu = torch.device("cpu") with ( @@ -513,14 +513,14 @@ def run_lower() -> dict[str, torch.Tensor]: with mock.patch.dict( os.environ, - {"DP_NEO_CUTE_INFER": "0"}, + {"DP_CUTE_INFER": "0"}, clear=False, ): reference = run_lower() with ( mock.patch.dict( os.environ, - {"DP_NEO_CUTE_INFER": "1"}, + {"DP_CUTE_INFER": "1"}, clear=False, ), mock.patch.object( @@ -538,7 +538,7 @@ def run_lower() -> dict[str, torch.Tensor]: continue self.assertTrue( torch.equal(actual[name], expected), - msg=f"DP_NEO_CUTE_INFER changed ineligible-model {name}", + msg=f"DP_CUTE_INFER changed ineligible-model {name}", ) def test_neo_cute_sort_guard_uses_post_cast_geometry_dtype(self) -> None: From 2344b4dd8393af876cc3ca3729f3087a753d3abc Mon Sep 17 00:00:00 2001 From: atulcthakur Date: Wed, 2 Sep 2026 20:31:30 -0700 Subject: [PATCH 04/11] fix(pt-expt): harden Neo CuTe integration --- deepmd/pt/model/descriptor/sezm.py | 2 + deepmd/pt_expt/kernels/cute/sezm/k1.py | 33 ++++------- .../cute_neo_phase_a_radial_forward.py | 6 +- .../k1_kernels/cute_neo_phase_c_onepass.py | 13 +++++ .../cute_neo_radial_phase_a_backward_node.py | 24 ++++---- .../cute/sezm/k1_message_grid_packed.py | 18 +++--- .../cute/sezm/k1_radial_phase_a_node.py | 42 ++++++++++++-- .../cute/sezm/message_grid_gaunt_sm90.py | 1 - .../output_grid_kernels/cute_readout_l0.py | 11 +++- .../cute/sezm/sm90_k1/final_phase_c.py | 5 +- .../cute/sezm/sm90_k1/phase_a_backward.py | 18 ++++++ .../sm90_k1/phase_c_attention_backward.py | 13 ++++- .../kernels/cute/sezm/sm90_k1/runner.py | 2 - doc/install/easy-install.md | 3 +- ...t_descriptor_sezm_block_ffn_grid_fusion.py | 58 +++++++------------ ...descriptor_sezm_cute_k1_manual_adjoints.py | 11 +++- ...or_sezm_cute_k1_phase_c_inplace_adjoint.py | 3 - ...riptor_sezm_cute_k1_radial_focus_fusion.py | 22 +++++++ ...riptor_sezm_cute_k1_radial_pitch68_safe.py | 20 ++++++- source/tests/pt/model/test_sezm_model.py | 14 ++++- 20 files changed, 212 insertions(+), 107 deletions(-) diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index 70bf750827..6bf8fee7f4 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/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 diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1.py b/deepmd/pt_expt/kernels/cute/sezm/k1.py index 7e227c332a..082c8e636f 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1.py +++ b/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: diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_a_radial_forward.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_a_radial_forward.py index 488806d314..15ad5e4220 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_a_radial_forward.py +++ b/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, diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_onepass.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_onepass.py index ee489af8bc..7abdfa762f 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_onepass.py +++ b/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( diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_radial_phase_a_backward_node.py b/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_radial_phase_a_backward_node.py index 603e330fb2..db12c713f1 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_radial_phase_a_backward_node.py +++ b/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] diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_message_grid_packed.py b/deepmd/pt_expt/kernels/cute/sezm/k1_message_grid_packed.py index c7d9a6ce2c..ab5d732469 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_message_grid_packed.py +++ b/deepmd/pt_expt/kernels/cute/sezm/k1_message_grid_packed.py @@ -7,8 +7,8 @@ ) from typing import ( - TYPE_CHECKING, Any, + Protocol, ) import torch @@ -16,15 +16,17 @@ 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) diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_radial_phase_a_node.py b/deepmd/pt_expt/kernels/cute/sezm/k1_radial_phase_a_node.py index ac968839ba..decdd8071a 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_radial_phase_a_node.py +++ b/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) diff --git a/deepmd/pt_expt/kernels/cute/sezm/message_grid_gaunt_sm90.py b/deepmd/pt_expt/kernels/cute/sezm/message_grid_gaunt_sm90.py index 2bbc087ac8..d7dd442e99 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/message_grid_gaunt_sm90.py +++ b/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 diff --git a/deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/cute_readout_l0.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/cute_readout_l0.py index 793367a4e1..9d78c11dbc 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/cute_readout_l0.py +++ b/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__( diff --git a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/final_phase_c.py b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/final_phase_c.py index 6475566c30..b711889c4c 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/final_phase_c.py +++ b/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 diff --git a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_a_backward.py b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_a_backward.py index 85a980799d..3b111362b7 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_a_backward.py +++ b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_a_backward.py @@ -860,6 +860,24 @@ def run_neo_phase_a_persistent_complex_backward_fp32( ) for name, tensor, shape, dtype in specs: _expect(name, tensor, shape, device, dtype) + torch._assert_async( + source_ptr[0] == 0, + "SM90 Phase-A backward requires source_ptr[0] == 0", + ) + torch._assert_async( + source_ptr[-1] == edge_count, + "SM90 Phase-A backward requires source_ptr[-1] == edge_count", + ) + if source_ptr.numel() > 1: + torch._assert_async( + torch.all(source_ptr[1:] >= source_ptr[:-1]), + "SM90 Phase-A backward requires nondecreasing source_ptr", + ) + if source_order.numel() != 0: + torch._assert_async( + torch.all((source_order >= 0) & (source_order < edge_count)), + "SM90 Phase-A backward requires source_order entries in [0, E)", + ) if outputs is None or workspace is None: if outputs is not None or workspace is not None: diff --git a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_c_attention_backward.py b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_c_attention_backward.py index f1f0430357..7ba611114a 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_c_attention_backward.py +++ b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_c_attention_backward.py @@ -40,6 +40,9 @@ from ..compile_cache import ( device_aware_lru_cache, ) +from ..k1_wigner_layout import ( + PACKED_VALUE_COUNT, +) if TYPE_CHECKING: from collections.abc import ( @@ -55,7 +58,7 @@ DEGREE_COUNT = 16 M0_WIDTH = 128 M1_WIDTH = 96 -PACKED_WIGNER_VALUES = 46 +PACKED_WIGNER_VALUES = PACKED_VALUE_COUNT MAX_EDGES_PER_NODE = 128 QK_SCALE = CHANNELS**-0.5 B0_NODE_VALUES = FOCUS_COUNT * DEGREE_COUNT * M0_WIDTH @@ -753,6 +756,14 @@ def run_grouped_expanded_final_phase_c_attention_adjoint( ) for name, tensor, shape, dtype in expected_inputs: _require_tensor(name, tensor, shape, dtype, device) + torch._assert_async( + dst_ptr[0] == 0, + "SM90 Phase-C backward requires dst_ptr[0] == 0", + ) + torch._assert_async( + dst_ptr[-1] == edge_count, + "SM90 Phase-C backward requires dst_ptr[-1] == edge_count", + ) if dst_ptr.numel() > 1: destination_degrees = dst_ptr[1:] - dst_ptr[:-1] torch._assert_async( diff --git a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/runner.py b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/runner.py index 1f8f37339e..b39b26cfdb 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/runner.py +++ b/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/runner.py @@ -492,7 +492,6 @@ def _final_manual_backward_sm90( phase = runner.phase_c_out.detach() x_wide = runner.x_wide.detach() - post_in = runner.post_mix_input.unsqueeze(2) grad_post_norm_in = _equivariant_rmsnorm_backward( block.post_so2_norm, runner.post_norm_input, @@ -500,7 +499,6 @@ def _final_manual_backward_sm90( ) grad_post_mix = _so3_linear_backward_input( so2.post_focus_mix, - post_in, grad_post_norm_in.squeeze(2).unsqueeze(2), ).squeeze(2) diff --git a/doc/install/easy-install.md b/doc/install/easy-install.md index e8a7f9cd77..ab07c24a92 100644 --- a/doc/install/easy-install.md +++ b/doc/install/easy-install.md @@ -192,7 +192,8 @@ pip install deepmd-kit :::::: Optional CuTe inference kernels are available on Linux with Python 3.11 or -newer for the PyTorch backend: +newer for the PyTorch backend on supported NVIDIA Ampere, Ada, Hopper, and +Blackwell GPUs: ```bash pip install "deepmd-kit[torch,cute]" diff --git a/source/tests/pt/model/test_descriptor_sezm_block_ffn_grid_fusion.py b/source/tests/pt/model/test_descriptor_sezm_block_ffn_grid_fusion.py index 678d8a0763..edc9d32114 100644 --- a/source/tests/pt/model/test_descriptor_sezm_block_ffn_grid_fusion.py +++ b/source/tests/pt/model/test_descriptor_sezm_block_ffn_grid_fusion.py @@ -22,7 +22,6 @@ def test_single_branch_bypasses_router_and_accepts_fused_middle( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setenv("DP_CUTE_INFER", "1") branch = ( GridBranch( channels=2, @@ -57,7 +56,7 @@ def fused_middle( torch.randn(2, 1, 4, device="cpu"), to_grid=lambda value: value, from_grid=lambda value: value, - grid_product=fused_middle, + pair_grid=fused_middle, ) assert calls == 1 @@ -67,7 +66,6 @@ def fused_middle( def test_single_branch_keeps_pytorch_middle_fallback( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setenv("DP_CUTE_INFER", "1") branch = ( GridBranch( channels=2, @@ -80,11 +78,15 @@ def test_single_branch_keeps_pytorch_middle_fallback( .to("cpu") .eval() ) - monkeypatch.setattr( - branch.router, - "forward", - lambda value: pytest.fail("single-branch router must not run"), - ) + original_forward = branch.router.forward + router_calls = 0 + + def tracked_router(value: torch.Tensor) -> torch.Tensor: + nonlocal router_calls + router_calls += 1 + return original_forward(value) + + monkeypatch.setattr(branch.router, "forward", tracked_router) calls = {"to_grid": 0, "from_grid": 0} def to_grid(value: torch.Tensor) -> torch.Tensor: @@ -105,10 +107,11 @@ def from_grid(value: torch.Tensor) -> torch.Tensor: ) assert calls == {"to_grid": 2, "from_grid": 1} + assert router_calls == 1 assert out.shape == left.shape -def test_single_branch_opt_in_matches_routed_result( +def test_single_branch_fused_pair_matches_routed_fallback( monkeypatch: pytest.MonkeyPatch, ) -> None: branch = ( @@ -135,7 +138,6 @@ def tracked_router(value: torch.Tensor) -> torch.Tensor: return original_forward(value) monkeypatch.setattr(branch.router, "forward", tracked_router) - monkeypatch.delenv("DP_CUTE_INFER", raising=False) routed = branch( left, right, @@ -143,13 +145,13 @@ def tracked_router(value: torch.Tensor) -> torch.Tensor: to_grid=lambda value: value, from_grid=lambda value: value, ) - monkeypatch.setenv("DP_CUTE_INFER", "1") shortcut = branch( left, right, scalar_pair, to_grid=lambda value: value, from_grid=lambda value: value, + pair_grid=lambda lhs, rhs: lhs * rhs, ) assert router_calls == 1 @@ -186,7 +188,7 @@ def tracked_softmax( torch.randn(2, 1, 4, device="cpu"), to_grid=lambda value: value, from_grid=lambda value: value, - grid_product=lambda left, right: pytest.fail( + pair_grid=lambda left, right: pytest.fail( "multi-branch routing must keep the generic path" ), ) @@ -196,7 +198,7 @@ def tracked_softmax( assert torch.isfinite(out).all() -def test_single_branch_training_preserves_router_zero_gradient( +def test_single_branch_training_uses_frozen_router_fallback( monkeypatch: pytest.MonkeyPatch, ) -> None: branch = ( @@ -231,20 +233,15 @@ def tracked_softmax( torch.randn(2, 1, 4, device="cpu"), to_grid=lambda value: value, from_grid=lambda value: value, - grid_product=lambda left, right: pytest.fail( - "training must keep the generic router path" - ), ) out.sum().backward() assert softmax_calls == 1 - assert branch.router.weight.grad is not None - assert torch.count_nonzero(branch.router.weight.grad) == 0 + assert branch.router.weight.grad is None + assert left.grad is not None -def test_single_branch_eval_with_trainable_parameters_preserves_router_gradient( - monkeypatch: pytest.MonkeyPatch, -) -> None: +def test_single_branch_eval_with_trainable_parameters_uses_router_fallback() -> None: branch = ( GridBranch( channels=2, @@ -264,19 +261,15 @@ def test_single_branch_eval_with_trainable_parameters_preserves_router_gradient( torch.randn(2, 1, 4, device="cpu"), to_grid=lambda value: value, from_grid=lambda value: value, - grid_product=lambda left, right: pytest.fail( - "trainable eval must preserve the generic router path" - ), ) out.sum().backward() - assert branch.router.weight.grad is not None - assert torch.count_nonzero(branch.router.weight.grad) == 0 + assert branch.router.weight.grad is None + assert left.grad is not None @pytest.mark.parametrize("training", (False, True)) -def test_grid_net_with_trainable_parameters_does_not_offer_fused_product( - monkeypatch: pytest.MonkeyPatch, +def test_grid_net_with_trainable_parameters_keeps_differentiable_path( training: bool, ) -> None: net = S2GridNet( @@ -294,13 +287,6 @@ def test_grid_net_with_trainable_parameters_does_not_offer_fused_product( seed=23, ).to("cpu") net.train(training) - monkeypatch.setattr( - net, - "_grid_product", - lambda left, right: pytest.fail( - "trainable grid net must not offer its first-order-only fused product" - ), - ) query = torch.randn(2, 4, 1, 4, device="cpu", requires_grad=True) net(query).sum().backward() @@ -335,7 +321,7 @@ def tracked_product(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: calls += 1 return net._from_grid(net._to_grid(left) * net._to_grid(right)) - monkeypatch.setattr(net, "_grid_product", tracked_product) + monkeypatch.setattr(net, "_pair_grid", tracked_product) output = net(torch.randn(2, 4, 1, 4, device="cpu")) assert calls == 1 diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_manual_adjoints.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_manual_adjoints.py index 28061d6cd2..15e54d5dcb 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_manual_adjoints.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_manual_adjoints.py @@ -337,6 +337,10 @@ def test_qk_manual_adjoint_matches_autograd_with_sparse_edges() -> None: k_node = torch.einsum("nfi,ifo->nfo", x_norm, k_weight) logits = (q_node[dst.long()] * k_node[src.long()]).sum(dim=-1) * scale reference = torch.autograd.grad(logits, x_wide, grad_logits)[0] + compile_identity = ( + torch.cuda.current_device(), + *torch.cuda.get_device_capability(device), + ) runner = SimpleNamespace( so2=SimpleNamespace( @@ -351,8 +355,11 @@ def test_qk_manual_adjoint_matches_autograd_with_sparse_edges() -> None: k_node=k_node.detach().contiguous(), src_i32=src, dst_i32=dst, - qk_edge_backward=compile_neo_qk_edge_backward(scale), - qk_node_input_adjoint=compile_neo_qk_node_input_adjoint(eps), + qk_edge_backward=compile_neo_qk_edge_backward(scale, compile_identity), + qk_node_input_adjoint=compile_neo_qk_node_input_adjoint( + eps, + compile_identity, + ), ) actual = k1._qk_manual_backward(runner, grad_logits) torch.cuda.synchronize(device) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py index e8bc32e950..42ceef92bb 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py @@ -39,7 +39,6 @@ "focus_alpha", "dst_ptr", "rotate_inv_rescale", - "logits", "edge_gate", "z_bias_raw", "group_max", @@ -92,7 +91,6 @@ def _inputs(dst_ptr_values: list[int]): "focus_alpha": torch.softmax(rand(edge_count, 2), dim=1), "dst_ptr": torch.tensor(dst_ptr_values, device=device, dtype=torch.int32), "rotate_inv_rescale": rand(16), - "logits": rand(edge_count, 2), "edge_gate": rand(edge_count).abs(), "z_bias_raw": rand(2), "group_max": rand(node_count, 2), @@ -132,7 +130,6 @@ def _call(runner, inputs, outputs): inputs["focus_alpha"], inputs["dst_ptr"], inputs["rotate_inv_rescale"], - inputs["logits"], inputs["edge_gate"], inputs["z_bias_raw"], inputs["group_max"], diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_focus_fusion.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_focus_fusion.py index bdccca901d..0d6ea568dd 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_focus_fusion.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_focus_fusion.py @@ -8,6 +8,7 @@ import unittest +import pytest import torch @@ -32,6 +33,27 @@ def test_source_csr_preserves_equal_source_order() -> None: ) +@pytest.mark.parametrize("src_values", ([-1, 0], [0, 3])) +def test_source_csr_rejects_out_of_range_sources(src_values: list[int]) -> None: + from deepmd.pt_expt.kernels.cute.sezm.k1_radial_phase_a_node import ( + build_source_csr, + ) + + src = torch.tensor(src_values, dtype=torch.int64, device="cpu") + with pytest.raises(RuntimeError, match="source indices"): + build_source_csr(src, node_count=3) + + +def test_source_csr_eager_validation_reports_value_error() -> None: + from deepmd.pt_expt.kernels.cute.sezm.k1_radial_phase_a_node import ( + build_source_csr, + ) + + src = torch.tensor([0, 3], dtype=torch.int64, device="cpu") + with pytest.raises(ValueError, match="source indices"): + build_source_csr(src, node_count=3, validate_sources=True) + + @unittest.skipUnless(_has_supported_gpu(), "requires an SM80 or SM90 CUDA device") class TestRadialFocusSourceFusion(unittest.TestCase): def test_fused_load_matches_standalone_scalar_lane_add(self) -> None: diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_pitch68_safe.py b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_pitch68_safe.py index 9815909e37..704e56707d 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_pitch68_safe.py +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_pitch68_safe.py @@ -58,8 +58,24 @@ class TestRadialPitch68SafeStatic(unittest.TestCase): def test_pitch68_specialization_contract(self) -> None: wrapper_source = WRAPPER_PATH.read_text(encoding="utf-8") kernel_source = KERNEL_PATH.read_text(encoding="utf-8") - - self.assertIn("SHARED_ROW_PITCH = HIDDEN + 4", kernel_source) + kernel_tree = ast.parse(kernel_source) + pitch_assignment = next( + node + for node in kernel_tree.body + if isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == "SHARED_ROW_PITCH" + for target in node.targets + ) + ) + pitch_expression = pitch_assignment.value + + self.assertIsInstance(pitch_expression, ast.BinOp) + self.assertIsInstance(pitch_expression.op, ast.Add) + self.assertIsInstance(pitch_expression.left, ast.Name) + self.assertEqual(pitch_expression.left.id, "HIDDEN") + self.assertIsInstance(pitch_expression.right, ast.Constant) + self.assertEqual(pitch_expression.right.value, 4) self.assertNotIn("pitch68_safe", wrapper_source) self.assertNotIn("pitch68_safe", kernel_source) diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index cf3daa54f4..d89138fcbb 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -980,12 +980,20 @@ def test_fixed_edge_geometry_matches_standard_cache(self) -> None: wraps=edge_cache_module._build_edge_wigner, ) as build_edge_wigner, ): + sorted_edge_index, sorted_edge_vec, sorted_edge_mask, _ = ( + _sort_edge_tensors_by_destination( + edge_index, + edge_vec, + edge_mask, + edge_index, + ) + ) cache_sfpg = build_edge_cache_from_edges( type_ebed=type_ebed, atype_flat=atype_loc.reshape(-1), - edge_index=edge_index, - edge_vec=edge_vec, - edge_mask=edge_mask, + edge_index=sorted_edge_index, + edge_vec=sorted_edge_vec, + edge_mask=sorted_edge_mask, compute_dtype=descriptor.compute_dtype, eps=descriptor.eps, deg_norm_floor=descriptor.deg_norm_floor, From f179f6a8f67858c52a1c3d68b3d4cdcada466bad Mon Sep 17 00:00:00 2001 From: OutisLi Date: Thu, 3 Sep 2026 13:46:36 +0800 Subject: [PATCH 05/11] perf(pt): unify DPA4 CuTe inference dispatch --- deepmd/dpmodel/descriptor/dpa4.py | 156 +++++++- .../dpmodel/descriptor/dpa4_nn/edge_cache.py | 97 ++++- .../dpmodel/descriptor/dpa4_nn/embedding.py | 96 ++++- deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py | 41 ++- deepmd/dpmodel/descriptor/dpa4_nn/so2.py | 84 +++-- .../dpmodel/utils/neighbor_graph/__init__.py | 2 + deepmd/dpmodel/utils/neighbor_graph/graph.py | 241 ++++++++++--- deepmd/pt/entrypoints/freeze_pt2.py | 8 +- deepmd/pt/model/descriptor/sezm.py | 214 ++++++----- deepmd/pt/model/descriptor/sezm_nn/block.py | 64 +--- .../pt/model/descriptor/sezm_nn/edge_cache.py | 86 +---- .../pt/model/descriptor/sezm_nn/embedding.py | 106 ++++-- .../pt/model/descriptor/sezm_nn/grid_net.py | 73 ++-- deepmd/pt/model/descriptor/sezm_nn/lora.py | 8 +- deepmd/pt/model/descriptor/sezm_nn/so2.py | 208 ++++++----- deepmd/pt/model/descriptor/sezm_nn/wignerd.py | 18 +- deepmd/pt/model/model/sezm_model.py | 48 +-- deepmd/pt/model/network/mlp.py | 6 +- deepmd/pt_expt/descriptor/dpa4.py | 276 +++++++++++++- deepmd/pt_expt/descriptor/dpa4_nn/block.py | 32 ++ .../pt_expt/descriptor/dpa4_nn/embedding.py | 85 ++++- deepmd/pt_expt/descriptor/dpa4_nn/grid_net.py | 108 +++++- deepmd/pt_expt/descriptor/dpa4_nn/so2.py | 280 +++++++------- deepmd/pt_expt/descriptor/dpa4_nn/wignerd.py | 12 +- deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py | 4 +- .../kernels/cute/sezm/output_grid/__init__.py | 3 + .../kernels}/__init__.py | 2 +- .../kernels/readout_l0.py} | 2 +- .../kernels/tiled_product.py} | 2 +- .../product.py} | 8 +- .../cute/sezm/{ => output_grid}/readout_l0.py | 70 ++-- .../kernels/cute/sezm/runtime_policy.py | 38 +- .../cute/sezm/{k1_kernels => so2}/__init__.py | 2 +- .../sezm/{sm90_k1 => so2/kernels}/__init__.py | 2 +- .../kernels/combined_gate_forward.py} | 6 +- .../kernels/envelope_softmax.py} | 2 +- .../kernels/focus_source_backward.py} | 8 +- .../kernels/gate_linear_residual_backward.py} | 2 +- .../kernels/message_grid_product.py} | 10 +- .../kernels/output_gate_backward.py} | 0 .../kernels/phase_a_radial_forward.py} | 10 +- .../kernels/phase_c_backward.py} | 2 +- .../kernels/phase_c_forward.py} | 4 +- .../kernels/qk_edge.py} | 0 .../kernels/radial_phase_a_backward.py} | 4 +- .../kernels/structural_gate_sm80.py} | 14 +- .../sezm/{k1_so2linear.py => so2/linear.py} | 8 +- .../message_grid.py} | 8 +- .../pt_expt/kernels/cute/sezm/so2/metadata.py | 96 +++++ .../cute/sezm/{k1.py => so2/operation.py} | 299 ++++++++------- .../phase_c.py} | 4 +- .../radial_phase_a.py} | 6 +- .../cute/sezm/{k1_runner.py => so2/runner.py} | 70 ++-- .../kernels/cute/sezm/so2/sm90/__init__.py | 3 + .../{sm90_k1 => so2/sm90}/final_phase_c.py | 4 +- .../sm90/message_grid_gaunt.py} | 2 +- .../sm90/message_grid_readout.py} | 14 +- .../sezm/{sm90_k1 => so2/sm90}/output_gate.py | 2 +- .../sezm/{sm90_k1 => so2/sm90}/persistent.py | 2 +- .../sezm/{sm90_k1 => so2/sm90}/phase_a.py | 4 +- .../{sm90_k1 => so2/sm90}/phase_a_backward.py | 4 +- .../sm90}/phase_c_attention_backward.py | 2 +- .../cute/sezm/{sm90_k1 => so2/sm90}/prefix.py | 0 .../cute/sezm/{sm90_k1 => so2/sm90}/radial.py | 8 +- .../cute/sezm/{sm90_k1 => so2/sm90}/runner.py | 56 +-- .../structural_gate.py} | 2 +- .../wigner_layout.py} | 2 +- .../cute/sezm/{k4_wignerd.py => wignerd.py} | 26 +- .../kernels/triton/sezm/so2_value_path.py | 2 +- deepmd/pt_expt/utils/serialization.py | 6 +- doc/install/easy-install.md | 14 +- doc/model/dpa4.md | 192 ++++++++-- pyproject.toml | 1 + .../common/dpmodel/test_dpa4_edge_cache.py | 39 ++ .../common/dpmodel/test_neighbor_graph.py | 74 ++++ source/tests/pt/model/sezm_cute/__init__.py | 1 + source/tests/pt/model/sezm_cute/_paths.py | 12 + .../test_compile_cache.py} | 33 +- .../test_efs.py} | 29 +- .../test_envelope_softmax.py} | 2 +- .../test_gie.py} | 0 .../test_multigpu_dispatch.py} | 54 +-- .../test_output_grid_product.py} | 34 +- .../test_readout_l0.py} | 74 ++-- .../test_runtime_policy.py} | 74 ++-- .../test_so2.py} | 92 ++--- .../test_so2_dst_ptr.py} | 74 ++-- .../test_so2_gate_structural.py} | 8 +- .../test_so2_inplace_residual.py} | 10 +- .../test_so2_manual_adjoints.py} | 26 +- .../test_so2_packed_impl.py} | 8 +- .../test_so2_phase_c_inplace_adjoint.py} | 2 +- .../test_so2_radial_focus_fusion.py} | 4 +- .../test_so2_radial_pitch68_safe.py} | 11 +- .../test_so2_radial_projection.py} | 8 +- .../test_so2_structural_gate_vec4_sm80.py} | 17 +- .../test_so2_structural_memory_reuse.py} | 70 ++-- .../test_tiled_grid_product.py} | 34 +- ...test_descriptor_sezm_readout_l0_algebra.py | 4 +- .../model/test_descriptor_sezm_train_paths.py | 46 +-- .../pt/model/test_descriptor_sezm_triton.py | 6 +- source/tests/pt/model/test_mlp.py | 16 +- source/tests/pt/model/test_sezm_model.py | 26 +- .../descriptor/test_dpa4_accelerated.py | 341 +++++++++++++++++- .../descriptor/test_dpa4_ckpt_triton.py | 6 +- .../descriptor/test_dpa4_train_paths.py | 44 +-- .../utils/test_serialization_kernel_levels.py | 2 +- 107 files changed, 3126 insertions(+), 1536 deletions(-) create mode 100644 deepmd/pt_expt/kernels/cute/sezm/output_grid/__init__.py rename deepmd/pt_expt/kernels/cute/sezm/{output_grid_kernels => output_grid/kernels}/__init__.py (71%) rename deepmd/pt_expt/kernels/cute/sezm/{output_grid_kernels/cute_readout_l0.py => output_grid/kernels/readout_l0.py} (99%) rename deepmd/pt_expt/kernels/cute/sezm/{output_grid_kernels/cute_tiled_grid_product.py => output_grid/kernels/tiled_product.py} (99%) rename deepmd/pt_expt/kernels/cute/sezm/{output_grid_product.py => output_grid/product.py} (98%) rename deepmd/pt_expt/kernels/cute/sezm/{ => output_grid}/readout_l0.py (94%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_kernels => so2}/__init__.py (68%) rename deepmd/pt_expt/kernels/cute/sezm/{sm90_k1 => so2/kernels}/__init__.py (72%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_kernels/cute_neo_so2_gate_combined_fwd.py => so2/kernels/combined_gate_forward.py} (99%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_kernels/cute_envelope_gated_softmax.py => so2/kernels/envelope_softmax.py} (99%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_kernels/cute_neo_focus_src_backward.py => so2/kernels/focus_source_backward.py} (97%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_kernels/cute_neo_gate_linear_residual_backward_fused.py => so2/kernels/gate_linear_residual_backward.py} (99%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_kernels/cute_neo_message_grid_product.py => so2/kernels/message_grid_product.py} (96%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_kernels/cute_neo_output_gate_backward.py => so2/kernels/output_gate_backward.py} (100%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_kernels/cute_neo_phase_a_radial_forward.py => so2/kernels/phase_a_radial_forward.py} (97%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_kernels/cute_neo_phase_c_backward_layout.py => so2/kernels/phase_c_backward.py} (99%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_kernels/cute_neo_phase_c_onepass.py => so2/kernels/phase_c_forward.py} (99%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_kernels/cute_neo_qk_edge.py => so2/kernels/qk_edge.py} (100%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_kernels/cute_neo_radial_phase_a_backward_node.py => so2/kernels/radial_phase_a_backward.py} (99%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py => so2/kernels/structural_gate_sm80.py} (99%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_so2linear.py => so2/linear.py} (98%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_message_grid_packed.py => so2/message_grid.py} (98%) create mode 100644 deepmd/pt_expt/kernels/cute/sezm/so2/metadata.py rename deepmd/pt_expt/kernels/cute/sezm/{k1.py => so2/operation.py} (89%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_kernels/cute_neo_phase_c_backward_layout_runner.py => so2/phase_c.py} (98%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_radial_phase_a_node.py => so2/radial_phase_a.py} (98%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_runner.py => so2/runner.py} (94%) create mode 100644 deepmd/pt_expt/kernels/cute/sezm/so2/sm90/__init__.py rename deepmd/pt_expt/kernels/cute/sezm/{sm90_k1 => so2/sm90}/final_phase_c.py (99%) rename deepmd/pt_expt/kernels/cute/sezm/{message_grid_gaunt_sm90.py => so2/sm90/message_grid_gaunt.py} (99%) rename deepmd/pt_expt/kernels/cute/sezm/{message_grid_readout_sm90.py => so2/sm90/message_grid_readout.py} (99%) rename deepmd/pt_expt/kernels/cute/sezm/{sm90_k1 => so2/sm90}/output_gate.py (99%) rename deepmd/pt_expt/kernels/cute/sezm/{sm90_k1 => so2/sm90}/persistent.py (99%) rename deepmd/pt_expt/kernels/cute/sezm/{sm90_k1 => so2/sm90}/phase_a.py (99%) rename deepmd/pt_expt/kernels/cute/sezm/{sm90_k1 => so2/sm90}/phase_a_backward.py (99%) rename deepmd/pt_expt/kernels/cute/sezm/{sm90_k1 => so2/sm90}/phase_c_attention_backward.py (99%) rename deepmd/pt_expt/kernels/cute/sezm/{sm90_k1 => so2/sm90}/prefix.py (100%) rename deepmd/pt_expt/kernels/cute/sezm/{sm90_k1 => so2/sm90}/radial.py (97%) rename deepmd/pt_expt/kernels/cute/sezm/{sm90_k1 => so2/sm90}/runner.py (94%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_gate_structural.py => so2/structural_gate.py} (98%) rename deepmd/pt_expt/kernels/cute/sezm/{k1_wigner_layout.py => so2/wigner_layout.py} (97%) rename deepmd/pt_expt/kernels/cute/sezm/{k4_wignerd.py => wignerd.py} (97%) create mode 100644 source/tests/pt/model/sezm_cute/__init__.py create mode 100644 source/tests/pt/model/sezm_cute/_paths.py rename source/tests/pt/model/{test_descriptor_sezm_cute_compile_cache.py => sezm_cute/test_compile_cache.py} (71%) rename source/tests/pt/model/{test_sezm_cute_efs.py => sezm_cute/test_efs.py} (94%) rename source/tests/pt/model/{test_descriptor_sezm_cute_envelope_softmax.py => sezm_cute/test_envelope_softmax.py} (96%) rename source/tests/pt/model/{test_descriptor_sezm_cute_gie.py => sezm_cute/test_gie.py} (100%) rename source/tests/pt/model/{test_sezm_cute_multigpu_dispatch.py => sezm_cute/test_multigpu_dispatch.py} (82%) rename source/tests/pt/model/{test_descriptor_sezm_cute_output_grid_product.py => sezm_cute/test_output_grid_product.py} (94%) rename source/tests/pt/model/{test_descriptor_sezm_cute_readout_l0.py => sezm_cute/test_readout_l0.py} (92%) rename source/tests/pt/model/{test_descriptor_sezm_cute_runtime_policy.py => sezm_cute/test_runtime_policy.py} (76%) rename source/tests/pt/model/{test_descriptor_sezm_cute_k1.py => sezm_cute/test_so2.py} (84%) rename source/tests/pt/model/{test_descriptor_sezm_cute_k1_dst_ptr.py => sezm_cute/test_so2_dst_ptr.py} (91%) rename source/tests/pt/model/{test_descriptor_sezm_cute_k1_gate_structural.py => sezm_cute/test_so2_gate_structural.py} (94%) rename source/tests/pt/model/{test_descriptor_sezm_cute_k1_inplace_residual.py => sezm_cute/test_so2_inplace_residual.py} (98%) rename source/tests/pt/model/{test_descriptor_sezm_cute_k1_manual_adjoints.py => sezm_cute/test_so2_manual_adjoints.py} (92%) rename source/tests/pt/model/{test_descriptor_sezm_cute_k1_packed_impl.py => sezm_cute/test_so2_packed_impl.py} (93%) rename source/tests/pt/model/{test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py => sezm_cute/test_so2_phase_c_inplace_adjoint.py} (98%) rename source/tests/pt/model/{test_descriptor_sezm_cute_k1_radial_focus_fusion.py => sezm_cute/test_so2_radial_focus_fusion.py} (96%) rename source/tests/pt/model/{test_descriptor_sezm_cute_k1_radial_pitch68_safe.py => sezm_cute/test_so2_radial_pitch68_safe.py} (94%) rename source/tests/pt/model/{test_descriptor_sezm_cute_k1_radial_projection.py => sezm_cute/test_so2_radial_projection.py} (96%) rename source/tests/pt/model/{test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py => sezm_cute/test_so2_structural_gate_vec4_sm80.py} (97%) rename source/tests/pt/model/{test_descriptor_sezm_cute_k1_structural_memory_reuse.py => sezm_cute/test_so2_structural_memory_reuse.py} (78%) rename source/tests/pt/model/{test_descriptor_sezm_cute_tiled_grid_product.py => sezm_cute/test_tiled_grid_product.py} (92%) diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py index f1106cd395..c81c289134 100644 --- a/deepmd/dpmodel/descriptor/dpa4.py +++ b/deepmd/dpmodel/descriptor/dpa4.py @@ -1199,10 +1199,10 @@ def __init__( # Accelerated backends may replace the distance-to-radial chain and the # packed Wigner-D construction. The array-API reference leaves these # hooks unbound and always retains the dense Wigner matrices. - self._cuda_radial_fn = None - self._cuda_wigner_fn = None - self._wigner_free_conv = False - self._packed_wigner_train = False + self.cuda_infer_l_1_radial = None + self.cuda_infer_l_1_wigner = None + self.cuda_infer_l_2_covers_all_blocks = False + self.cuda_train_covers_all_blocks = False # === Optional descriptor-level attention residuals === self.final_block_attn_res = None @@ -1542,9 +1542,31 @@ def _run_graph( graph = apply_pair_exclusion(graph, atype_flat, self.emask) if n_out_nodes is None: n_out_nodes = atype_flat.shape[0] + packed_wigner_graph = self.prepare_packed_wigner_graph( + graph, atype_flat.shape[0] + ) + packed_wigner = packed_wigner_graph is not None + if packed_wigner_graph is not None: + graph = packed_wigner_graph edge_index = graph.edge_index edge_vec = graph.edge_vec edge_mask = graph.edge_mask + # Graph-owned endpoint orderings remain aligned with the edge payload + # and are shared by every segmented consumer through the edge cache. + graph_csr_cache = None + if all( + value is not None + for value in ( + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + ) + ): + graph_csr_cache = { + "dst": (graph.destination_order, graph.destination_row_ptr), + "src": (graph.source_order, graph.source_row_ptr), + } xp = array_api_compat.array_namespace(edge_vec) device = array_api_compat.device(edge_vec) @@ -1592,14 +1614,18 @@ def _run_graph( bridging_switch=self.bridging_switch, edge_envelope=self.edge_envelope, radial_basis=self.radial_basis, - fused_radial=None if training else self._cuda_radial_fn, - fused_wigner=None if training else self._cuda_wigner_fn, + fused_radial=None if training else self.cuda_infer_l_1_radial, + fused_wigner=None if training else self.cuda_infer_l_1_wigner, # Random local-Z roll is a training-only augmentation; the model # is roll-equivariant, so inference fixes gamma. random_gamma=self.random_gamma and training, wigner_calc=self.wigner_calc, - build_wigner=self._build_full_wigner(), + build_wigner=self._build_full_wigner() or packed_wigner, node_partial_exchange=node_partial_exchange, + packed_wigner=packed_wigner, + destinations_sorted=graph.destination_sorted, + packed_wigner_fn=self.build_packed_wigner, + csr_cache=graph_csr_cache, ) ebed_dim_0 = self.node_init_dim # (node_init_lmax+1)^2 @@ -1719,6 +1745,9 @@ def _run_graph( edge_cache = edge_cache_to_dtype( edge_cache, get_xp_precision(xp, self.precision) ) + edge_cache.cute_infer_so2_metadata = self.prepare_cute_infer_so2_metadata( + edge_cache, n_nodes + ) x = self._forward_blocks( x, edge_cache, rad_feat_per_block, comm_dict=comm_dict ) @@ -1897,10 +1926,30 @@ def _apply_readout(self, x: Array, n_rows: int) -> Array: ) for layer in self.readout_pre_layers: x_ro = x_ro + layer(x_ro) + if not self.readout_pre_layers and self.so3_readout != "none": + accelerated = self.run_cute_infer_readout(x_ro) + if accelerated is not None: + return xp.reshape(accelerated, (n_rows, 1, 1, self.channels)) if self.so3_readout == "none": return (x_ro + self.output_ffn(x_ro))[:, 0:1, :, :] return x_ro[:, 0:1, :, :] + self.output_ffn.call_scalar(x_ro) + def run_cute_infer_readout(self, ffn_in: Array) -> Array | None: + """Run the CuTe readout when its exact inference contract matches. + + Parameters + ---------- + ffn_in : Array + Equivariant readout input with shape ``(N, D, 1, C)``. + + Returns + ------- + Array or None + Residual-inclusive scalar output with shape ``(N, C)``, or ``None`` + when the backend has no eligible implementation. + """ + return None + def _edge_quaternion(self, edge_cache: EdgeCache) -> Array: """ Return the cached global->local edge quaternion, rebuilding if absent. @@ -1931,8 +1980,73 @@ def _build_full_wigner(self) -> bool: if not self._need_full_wigner: return False if self._in_training_mode(): - return not self._packed_wigner_train - return not self._wigner_free_conv + return not self.cuda_train_covers_all_blocks + return not self.cuda_infer_l_2_covers_all_blocks + + def prepare_packed_wigner_graph( + self, + graph: NeighborGraph, + n_nodes: int, + ) -> NeighborGraph | None: + """Prepare an edge graph for backend packed-Wigner storage. + + Parameters + ---------- + graph : NeighborGraph + Edge graph supplied to the descriptor. + n_nodes : int + Number of nodes addressed by the graph. + + Returns + ------- + NeighborGraph or None + A graph satisfying the packed layout contract, or ``None`` when + the backend does not select that representation. + """ + return None + + def build_packed_wigner( + self, + edge_quat: Array, + wigner_calc: Any, + ) -> Array | None: + """Build backend-specific packed Wigner storage when available. + + Parameters + ---------- + edge_quat : Array + Global-to-local edge quaternions with shape ``(E, 4)``. + wigner_calc : Any + Wigner calculator carrying the degree and basis convention. + + Returns + ------- + Array or None + Packed per-edge Wigner storage, or ``None`` to retain dense storage. + """ + return None + + def prepare_cute_infer_so2_metadata( + self, + edge_cache: EdgeCache, + n_nodes: int, + ) -> tuple[Array, Array, Array] | None: + """Build the edge metadata consumed by the CuTe SO2 implementation. + + Parameters + ---------- + edge_cache : EdgeCache + Per-forward cache carrying the destination-major edge payload. + n_nodes : int + Number of nodes addressed by the edge payload. + + Returns + ------- + tuple[Array, Array, Array] or None + Destination row pointers, source order, and source row pointers, or + ``None`` when the backend does not select CuTe SO2. + """ + return None def _shared_wigner_runs( self, @@ -1981,6 +2095,8 @@ def _build_gie_zonal_coupling( the blocks are skipped (all-Cartesian model) the full coupling is reconstructed from the edge quaternion via the m=0-only path. """ + if edge_cache.D_packed is not None: + return self.build_cute_infer_zonal_coupling(edge_cache) if edge_cache.Dt_full is None: calc = self.gie_zonal_wigner_calc or self.wigner_calc shared = self._shared_wigner_runs(edge_cache, calc.lmax) @@ -2008,6 +2124,28 @@ def _build_gie_zonal_coupling( ) return xp.concat([mp_coupling, extra_coupling], axis=1) + def build_cute_infer_zonal_coupling(self, edge_cache: EdgeCache) -> Array: + """Extract the GIE zonal coupling from packed Wigner storage. + + Parameters + ---------- + edge_cache : EdgeCache + Per-forward cache carrying backend packed Wigner storage. + + Returns + ------- + Array + Zonal coupling with shape ``(E, D_node - 1)``. + + Raises + ------ + NotImplementedError + If a backend supplies packed Wigner storage without this extractor. + """ + raise NotImplementedError( + "packed Wigner storage requires a backend zonal-coupling implementation" + ) + def _apply_charge_spin_embedding( self, type_ebed: Array, diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py b/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py index ec51a3d23e..8fc0f0e5fa 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py @@ -44,6 +44,7 @@ ) WignerCalculatorFn = Callable[[Any], "tuple[Any, Any]"] +PackedWignerFn = Callable[[Any, WignerCalculatorFn], Any | None] # Distance and keep weight to the keep-weighted envelope and radial basis, the # fused replacement of applying the two modules separately. FusedRadialFn = Callable[[Any, Any], "tuple[Any, Any]"] @@ -94,7 +95,11 @@ class EdgeCache: csr_cache Lazy cache for endpoint CSR views used by segmented accelerated operators, keyed by endpoint role (``"dst"`` or ``"src"``). Built once - per step and shared by every consumer. + per step and shared by every consumer. Permutations and row pointers + use int64 regardless of the input graph index dtype. + cute_infer_so2_metadata + Optional ``(destination_row_ptr, source_order, source_row_ptr)`` tuple + prepared once for a backend CuTe SO2 implementation. edge_src_gate Optional per-edge Source Freeze Propagation Gate (SFPG) weight with shape (E, 1). Equals ``eta[src]`` where @@ -112,6 +117,12 @@ class EdgeCache: all slots are valid (e.g. the sparse :func:`_edge_cache_from_arrays` path, where masking is folded into the per-edge weights). This field has no pt counterpart. + destinations_sorted + Whether the edge payload is destination-major with masked entries + confined to the suffix. + D_packed + Optional backend-specific packed Wigner panel. Dense consumers observe + ``D_full`` and ``Dt_full`` as ``None`` when this storage is active. """ src: Any @@ -126,10 +137,13 @@ class EdgeCache: Dt_full: Any = None D_to_m_cache: dict[str, Any] = field(default_factory=dict) Dt_from_m_cache: dict[str, Any] = field(default_factory=dict) - csr_cache: dict[str, Any] | None = field(default_factory=dict) + csr_cache: dict[str, tuple[Any, Any]] | None = field(default_factory=dict) edge_src_gate: Any = None edge_quat: Any = None edge_mask: Any = None + cute_infer_so2_metadata: tuple[Any, Any, Any] | None = None + destinations_sorted: bool = False + D_packed: Any = None def compute_edge_src_gate( @@ -279,6 +293,10 @@ def _edge_cache_from_arrays( node_partial_exchange: Callable[[Any], Any] | None = None, fused_radial: FusedRadialFn | None = None, fused_wigner: WignerCalculatorFn | None = None, + packed_wigner: bool = False, + destinations_sorted: bool = False, + packed_wigner_fn: PackedWignerFn | None = None, + csr_cache: dict[str, tuple[Any, Any]] | None = None, ) -> EdgeCache: """ Build the global edge cache from a sparse edge list. @@ -331,6 +349,17 @@ def _edge_cache_from_arrays( fused_wigner Optional fused replacement of ``wigner_calc`` that builds the packed pair in one kernel pass. + packed_wigner + Whether the backend-specific packed Wigner contract is satisfied. + destinations_sorted + Whether the edge payload is destination-major with masked entries in + a suffix. + packed_wigner_fn + Optional backend implementation that builds a packed Wigner panel. + csr_cache + Optional endpoint CSR views inherited from the input graph. The cache + owns the same edge axis as ``edge_index``. Permutations and row pointers + are normalized to int64 at this graph-to-cache boundary. gamma Optional per-edge roll angles with shape (E,), used only when ``random_gamma`` is True. When None, drawn with the backend's RNG @@ -351,6 +380,18 @@ def _edge_cache_from_arrays( n_nodes = type_ebed.shape[0] src = xp.astype(edge_index[0, ...], xp.int64) dst = xp.astype(edge_index[1, ...], xp.int64) + if csr_cache is not None: + csr_cache = { + endpoint: ( + order if order.dtype == xp.int64 else xp.astype(order, xp.int64), + ( + row_ptr + if row_ptr.dtype == xp.int64 + else xp.astype(row_ptr, xp.int64) + ), + ) + for endpoint, (order, row_ptr) in csr_cache.items() + } # === Step 1. Normalize mask === edge_keep = xp.astype(edge_mask, xp.bool) @@ -379,7 +420,7 @@ def _edge_cache_from_arrays( edge_rbf = radial_basis(edge_len) * edge_keep_f # (E, n_radial) # === Step 4. Edge quaternion -> Wigner-D blocks === - D_full, Dt_full, edge_quat = _build_edge_wigner( + D_full, Dt_full, D_packed, edge_quat = _build_edge_wigner( edge_vec=edge_vec, edge_len=edge_len, eps=eps, @@ -387,7 +428,9 @@ def _edge_cache_from_arrays( wigner_calc=fused_wigner if fused_wigner is not None else wigner_calc, gamma=gamma, build_full=build_wigner, - ) # (E, D, D), (E, D, D), (E, 4) + packed_wigner=packed_wigner, + packed_wigner_fn=packed_wigner_fn, + ) # === Step 5. Edge type features === edge_type_feat = build_edge_type_feat(type_ebed, src, dst) @@ -419,9 +462,12 @@ def _edge_cache_from_arrays( edge_env=edge_env, D_full=D_full, Dt_full=Dt_full, + D_packed=D_packed, edge_quat=edge_quat, deg_norm_floor=deg_norm_floor, edge_src_gate=edge_src_gate, + destinations_sorted=destinations_sorted, + csr_cache=csr_cache, ) @@ -434,7 +480,9 @@ def _build_edge_wigner( wigner_calc: WignerCalculatorFn, gamma: Any = None, build_full: bool = True, -) -> tuple[Any, Any, Any]: + packed_wigner: bool = False, + packed_wigner_fn: PackedWignerFn | None = None, +) -> tuple[Any, Any, Any, Any]: """ Build packed Wigner-D blocks from edge vectors. @@ -461,13 +509,16 @@ def _build_edge_wigner( False (all message-passing blocks take the Cartesian path), only the quaternion is returned and the blocks are ``None``; the geometric initial embedding reconstructs the zonal coupling from the quaternion. + packed_wigner + Whether the backend-specific packed Wigner contract is satisfied. + packed_wigner_fn + Optional backend implementation that builds a packed Wigner panel. Returns ------- - tuple[Array, Array, Array] - Packed Wigner-D matrices ``(D_full, Dt_full)`` with shape ``(E, D, D)`` - (or ``None`` when ``build_full`` is False) and the quaternion used to - build them with shape ``(E, 4)``. + tuple[Array, Array, Array, Array] + Dense Wigner-D matrices ``(D_full, Dt_full)``, an optional packed panel, + and the quaternion used to build them. """ xp = array_api_compat.array_namespace(edge_vec) device = array_api_compat.device(edge_vec) @@ -491,9 +542,13 @@ def _build_edge_wigner( # === Step 3. Convert quaternions to packed Wigner-D blocks === if not build_full: - return None, None, edge_quat + return None, None, None, edge_quat + if packed_wigner and packed_wigner_fn is not None: + D_packed = packed_wigner_fn(edge_quat, wigner_calc) + if D_packed is not None: + return None, None, D_packed, edge_quat D_full, Dt_full = wigner_calc(edge_quat) - return D_full, Dt_full, edge_quat + return D_full, Dt_full, None, edge_quat def _finalize_edge_cache( @@ -507,9 +562,12 @@ def _finalize_edge_cache( edge_env: Any, D_full: Any, Dt_full: Any, + D_packed: Any, edge_quat: Any, deg_norm_floor: float, edge_src_gate: Any = None, + destinations_sorted: bool = False, + csr_cache: dict[str, tuple[Any, Any]] | None = None, ) -> EdgeCache: """ Assemble the shared `EdgeCache` layout. @@ -536,6 +594,8 @@ def _finalize_edge_cache( Dt_full Transposed packed Wigner-D matrices with shape (E, D, D), or None when the full Wigner-D construction is skipped. + D_packed + Optional backend-specific packed Wigner panel. edge_quat Global-to-local quaternions used to build the Wigner-D matrices with shape (E, 4). @@ -547,6 +607,10 @@ def _finalize_edge_cache( edge_src_gate Optional per-edge SFPG weight with shape (E, 1). ``None`` in non-bridging mode. + destinations_sorted + Whether ``dst`` is nondecreasing with masked entries in a suffix. + csr_cache + Optional endpoint CSR views aligned with the cache edge axis. Returns ------- @@ -574,11 +638,13 @@ def _finalize_edge_cache( inv_sqrt_deg=inv_sqrt_deg, D_full=D_full, Dt_full=Dt_full, + D_packed=D_packed, D_to_m_cache={}, Dt_from_m_cache={}, - csr_cache={}, + csr_cache={} if csr_cache is None else csr_cache, edge_src_gate=edge_src_gate, edge_quat=edge_quat, + destinations_sorted=destinations_sorted, ) @@ -639,16 +705,20 @@ def edge_cache_to_dtype(cache: EdgeCache, dtype: Any) -> EdgeCache: # Use local variables with explicit None check and assignment. _D_full = cache.D_full _Dt_full = cache.Dt_full + cached_D_packed = cache.D_packed _edge_src_gate = cache.edge_src_gate _edge_quat = cache.edge_quat D_full: Any = None Dt_full: Any = None + D_packed: Any = None edge_src_gate: Any = None edge_quat: Any = None if _D_full is not None: D_full = xp.astype(_D_full, dtype) if _Dt_full is not None: Dt_full = xp.astype(_Dt_full, dtype) + if cached_D_packed is not None: + D_packed = xp.astype(cached_D_packed, dtype) if _edge_src_gate is not None: edge_src_gate = xp.astype(_edge_src_gate, dtype) if _edge_quat is not None: @@ -667,10 +737,13 @@ def edge_cache_to_dtype(cache: EdgeCache, dtype: Any) -> EdgeCache: inv_sqrt_deg=xp.astype(cache.inv_sqrt_deg, dtype), D_full=D_full, Dt_full=Dt_full, + D_packed=D_packed, D_to_m_cache=None if cache.D_to_m_cache is None else {}, Dt_from_m_cache=None if cache.Dt_from_m_cache is None else {}, csr_cache=None if cache.csr_cache is None else dict(cache.csr_cache), edge_src_gate=edge_src_gate, edge_quat=edge_quat, edge_mask=cache.edge_mask, + cute_infer_so2_metadata=cache.cute_infer_so2_metadata, + destinations_sorted=cache.destinations_sorted, ) diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py index 9de135e641..c9e2563657 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py @@ -315,17 +315,28 @@ def call( axis=1, ) # (E, D-1) - # === Step 3. Broadcast radial features per row === + # === Step 3. Optional backend message construction and reduction === + accelerated = self.run_cute_infer_gie( + n_nodes=n_nodes, + edge_cache=edge_cache, + radial_feat=radial_feat, + zonal_coupling=zonal_coupling, + spin_l1_message=spin_l1_message, + ) + if accelerated is not None: + return accelerated + + # === Step 4. 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 5, so # it takes over whenever nothing else joins the message in between. if ( - self._can_fuse_scatter(zonal_coupling) + self.can_run_cuda_infer_l_1_scatter(zonal_coupling) and spin_l1_message is None and edge_cache.edge_src_gate is None and edge_cache.csr_cache is not None ): - return self.forward_fused_scatter( + return self.run_cuda_infer_l_1_scatter( n_nodes, edge_cache, radial_feat, zonal_coupling ) @@ -339,7 +350,7 @@ def call( zonal_coupling[:, :, None] * radial_value_for_row ) # (E, D-1, C) - # === Step 3b. Fold in the neighbor-spin l=1 message (native spin) === + # === Step 4b. Fold in the neighbor-spin l=1 message (native spin) === # The l=1 coefficients occupy the first three packed non-scalar rows, so # the neighbor-spin message joins the geometric message there and then # shares the source gate, scatter and degree normalization below. @@ -352,7 +363,7 @@ def call( non_scalar_message, 1, scatter_index, spin_l1_message ) - # === Step 4. Source Freeze Propagation Gate (optional) === + # === Step 5. Source Freeze Propagation Gate (optional) === # Mute messages emitted by nodes whose local neighborhood enters # the frozen zone. ``edge_src_gate`` is ``None`` outside bridging # mode so this is a no-op in normal training. @@ -362,7 +373,7 @@ def call( xp.reshape(src_gate, (n_edge, 1, 1)), non_scalar_message.dtype ) - # === Step 5. Scatter to nodes and normalize === + # === Step 6. Scatter to nodes and normalize === # Destination scatter-add over ``edge_cache.dst`` (pt ``index_add_``), # applied after the validity masking below. This reduction is # layout-agnostic: it is correct both for the padded ``call`` (row-major @@ -399,18 +410,83 @@ def call( out = out * xp.astype(edge_cache.inv_sqrt_deg, out.dtype) return xp.astype(out, dtype) - def _can_fuse_scatter(self, zonal_coupling: Any) -> bool: - """Return whether a backend can run the fused scatter for this input.""" + def run_cute_infer_gie( + self, + *, + n_nodes: int, + edge_cache: EdgeCache, + radial_feat: Any, + zonal_coupling: Any, + spin_l1_message: Any = None, + ) -> Any | None: + """Run a backend CuTe GIE implementation when its contract matches. + + Parameters + ---------- + n_nodes : int + Number of nodes addressed by the edge cache. + edge_cache : EdgeCache + Per-forward geometric edge cache. + radial_feat : Any + Radial features with shape ``(E, lmax, C)``. + zonal_coupling : Any + Zonal coupling with shape ``(E, D - 1)``. + spin_l1_message : Any, optional + Native-spin message with shape ``(E, 3, C)``. + + Returns + ------- + Any or None + Initial embedding with shape ``(N, D, C)``, or ``None`` when no + backend CuTe implementation serves the input. + """ + return None + + def can_run_cuda_infer_l_1_scatter(self, zonal_coupling: Any) -> bool: + """Return whether a backend can run the fused scatter for this input. + + Parameters + ---------- + zonal_coupling : Any + Zonal coupling with shape ``(E, D - 1)``. + + Returns + ------- + bool + Whether the CUDA inference level-one scatter is eligible. + """ return False - def forward_fused_scatter( + def run_cuda_infer_l_1_scatter( self, n_nodes: int, edge_cache: EdgeCache, radial_feat: Any, zonal_coupling: Any, ) -> Any: - """Build and reduce the geometric message with a backend operator.""" + """Build and reduce the geometric message with a backend operator. + + Parameters + ---------- + n_nodes : int + Number of nodes addressed by the edge cache. + edge_cache : EdgeCache + Per-forward geometric edge cache. + radial_feat : Any + Radial features with shape ``(E, lmax, C)``. + zonal_coupling : Any + Zonal coupling with shape ``(E, D - 1)``. + + Returns + ------- + Any + Initial embedding with shape ``(N, D, C)``. + + Raises + ------ + NotImplementedError + If the backend has no fused CUDA scatter implementation. + """ raise NotImplementedError( "The fused GIE scatter is unavailable in the dpmodel reference" ) diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py index 597ac02261..12c2fa131c 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py @@ -874,8 +874,8 @@ def __init__( # training hook differentiates the same expression inside the force # graph on the frame-packed operands, with analytic first and second # order. - self._grid_pair_fn = None - self._grid_pair_train_fn = None + self.cuda_infer_l_1_grid_pair = None + self.triton_train_l_1_grid_pair = None self._from_grid_t = np.ascontiguousarray(projector.from_grid_mat.T) self.scalar_act = SwiGLU() @@ -947,7 +947,9 @@ def call_scalar(self, query: Any, context: Any = None) -> Any: because materializing a scalar-only fallback grid would be slower than that fused operator. """ - if self._grid_pair_fn is not None and not getattr(self, "training", False): + if self.cuda_infer_l_1_grid_pair is not None and not getattr( + self, "training", False + ): return self._slice_scalar_layout(self.call(query, context)) return self._forward(query, context, scalar_only=True) @@ -1213,26 +1215,31 @@ def _pair_grid(self, left: Any, right: Any) -> Any | None: c_wide = left.shape[3] // self.n_frames if c_wide % 32 != 0 or left.shape != right.shape: return None - if getattr(self, "training", False): + training = getattr(self, "training", False) + if not training: + candidate = self.run_cute_infer_grid_pair(left, right) + if candidate is not None: + return candidate + if training: # Training form: frame-packed operands ride through unreshaped, # with analytic first and second order behind the call; under # autocast it runs the same reduced-precision regime as the dense # einsum composition it replaces. - if self._grid_pair_train_fn is None: + if self.triton_train_l_1_grid_pair is None: return None - return self._grid_pair_train_fn( + return self.triton_train_l_1_grid_pair( left, right, self.projector.to_grid_mat, self._from_grid_t, self.n_frames, ) - if self._grid_pair_fn is None or left.shape[2] != 1: + if self.cuda_infer_l_1_grid_pair is None or left.shape[2] != 1: return None n_batch, coeff_dim = left.shape[0], left.shape[1] flat_p = coeff_dim * self.n_frames xp = array_api_compat.array_namespace(left, right) - out = self._grid_pair_fn( + out = self.cuda_infer_l_1_grid_pair( xp.reshape(left, (n_batch, flat_p, c_wide)), xp.reshape(right, (n_batch, flat_p, c_wide)), self.projector.to_grid_mat, @@ -1240,6 +1247,24 @@ def _pair_grid(self, left: Any, right: Any) -> Any | None: ) return xp.reshape(out, (n_batch, coeff_dim, 1, self.n_frames * c_wide)) + def run_cute_infer_grid_pair(self, left: Any, right: Any) -> Any | None: + """Run a backend CuTe grid product when its contract matches. + + Parameters + ---------- + left : Any + Left coefficient operand with shape ``(N, D, F, n_frames * C)``. + right : Any + Right coefficient operand with the same shape as ``left``. + + Returns + ------- + Any or None + Coefficient result with the same shape as ``left``, or ``None`` + when the backend has no eligible CuTe implementation. + """ + return None + def _project_pair_in_one_transform( self, left: Any, diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py index a594ad9d09..68907b49f9 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/so2.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/so2.py @@ -1592,9 +1592,9 @@ def __init__( # attention layout without the optional focus-mix / value / output # projections (the deployed DPA4 configuration). Whichever of the # mutually exclusive inference gates is active then supplies the - # implementation, and ``self._flash_atten_fn`` being bound is what marks + # implementation, and ``self.flash_attention`` being bound is what marks # the fused path as live. The array-API reference leaves the hook unbound. - self._flash_atten_layout_ok = ( + self.flash_attention_layout_supported = ( self.n_atten_head > 0 and self.mmax == 1 and self.needs_local_frame @@ -1604,9 +1604,10 @@ def __init__( and self.attn_o_proj is None and self.attn_focus_mix is None ) - self._flash_atten_fn = None - self._cuda_conv_fn = None - self._cached_edge_csr_fn = None + self.flash_attention = None + self.flash_attention_supports_training = False + self.cuda_infer_l_2_conv = None + self.edge_csr = None # === Step 13. Optional fused SO(2) value-path seam === # The fused value path folds the rotate-to-local projection, radial @@ -1615,20 +1616,20 @@ def __init__( # array-API reference has no such kernel, so it never runs the fused # value path: every hook stays ``None`` and ``so2_message`` takes the # dense branch. The ``pt_expt`` backend binds the selected implementation. - self._triton_value_path = None - self._cutile_value_path = None + self.triton_infer_l_2_value = None + self.cutile_infer_value = None # === Step 14. Optional fused training seams === # Training differentiates the convolution twice under a force loss, so # its accelerated forms carry analytic backward and second-order # implementations of their own: one fused operator for the value stream - # up to the attention aggregation (``_cuda_value_train``), and the + # up to the attention aggregation (``cuda_train_value``), and the # segmented attention softmax / flash aggregation pair for the - # attention span (``_flash_atten_trains`` marks the bound aggregation + # attention span (``flash_attention_supports_training`` marks the bound + # aggregation # as training-capable). The array-API reference leaves every hook # unbound and trains through the dense expression. - self._cuda_value_train = None - self._flash_atten_trains = False + self.cuda_train_value = None self.trainable = bool(trainable) def call( @@ -1772,23 +1773,25 @@ def forward_attention( # does not serve the bridging mode, whose source gate reshapes the # softmax normalization. training = getattr(self, "training", False) - run_cuda = ( - self._cuda_conv_fn is not None + run_cuda_infer_l_2 = ( + self.cuda_infer_l_2_conv is not None and not training and edge_cache.edge_src_gate is None ) - run_flash = ( - self._flash_atten_fn is not None - and (not training or self._flash_atten_trains) - and not run_cuda + run_flash_attention = ( + self.flash_attention is not None + and (self.flash_attention_supports_training or not training) + and not run_cuda_infer_l_2 ) - if run_cuda: - return self.forward_attention_cuda(x, edge_cache, radial_feat, x_l0_node) - if run_flash: + if run_cuda_infer_l_2: + return self.forward_attention_cuda_infer_l_2( + x, edge_cache, radial_feat, x_l0_node + ) + if run_flash_attention: return self.forward_attention_flash(x, edge_cache, radial_feat, x_l0_node) return self.forward_attention_dense(x, edge_cache, radial_feat, x_l0_node) - def forward_attention_cuda( + def forward_attention_cuda_infer_l_2( self, x: Array, edge_cache: EdgeCache, @@ -1824,7 +1827,7 @@ def forward_attention_cuda( xp = array_api_compat.array_namespace(x) # === Step 1. Projected radial features === - rad_feat = self._cuda_conv_fn.radial_features( + rad_feat = self.cuda_infer_l_2_conv.radial_features( radial_feat ) # (E, lmax+1, C_wide) @@ -1835,7 +1838,9 @@ def forward_attention_cuda( head_gate = self.attention_head_gate(x_l0_node) # (N, Fa, H) # === Step 4. Fused convolution === - out = self._cuda_conv_fn(x, edge_cache, rad_feat, q_node, k_node, head_gate) + out = self.cuda_infer_l_2_conv( + x, edge_cache, rad_feat, q_node, k_node, head_gate + ) return xp.astype(out, get_xp_precision(xp, self.precision)) def forward_attention_flash( @@ -1852,7 +1857,10 @@ def forward_attention_flash( kernel folds the block-diagonal rotate-back, the inverse-rotation rescale, the per-edge weighting and the destination reduction into a single atomic-free pass, so the transient rotate-back message and - weighted value tensors are never materialized. + weighted value tensors are never materialized. The value producer is + selected independently inside ``so2_message``: training uses + ``cuda_train_value`` when it is bound, while the flash operator remains + the aggregation consumer selected by the Triton training gate. Parameters ---------- @@ -1888,14 +1896,14 @@ def forward_attention_flash( # === Step 4. Fused rotate-back and weighted destination reduction === # The destination CSR view is built once per step and shared by every # segment consumer of the graph. - if self._cached_edge_csr_fn is None: + if self.edge_csr is None: raise RuntimeError("The fused attention path requires a CSR builder") dst = edge_cache.dst - order, row_ptr = self._cached_edge_csr_fn(edge_cache, "dst", x.shape[0]) + order, row_ptr = self.edge_csr(edge_cache, "dst", x.shape[0]) rotation = edge_cache.Dt_full if rotation is None: - rotation = self._cuda_value_train.edge_runs(edge_cache) - pre_gate = self._flash_atten_fn( + rotation = self.cuda_train_value.edge_runs(edge_cache) + pre_gate = self.flash_attention( x_local, rotation, self.rotate_inv_rescale_full, @@ -2348,22 +2356,22 @@ def so2_message( n_edge = src.shape[0] training = getattr(self, "training", False) - if self._cutile_value_path is not None and not training: + if self.cutile_infer_value is not None and not training: # === Steps 1-5 (fused cuTile operators). ``rotate_mix`` folds the # rotation and the radial degree mixing into one edge-parallel # kernel writing the focus-major layout; ``mixing_stack`` runs the # whole gated stack, keeping the inter-layer activations and the # gated-layer pre-activations off the traced graph entirely. === - x_local, rad_feat = self._cutile_value_path(x, edge_cache, radial_feat) - elif self._cuda_value_train is not None and training: + x_local, rad_feat = self.cutile_infer_value(x, edge_cache, radial_feat) + elif self.cuda_train_value is not None and training: # === Steps 1-5 (one CUDA kernel, training). The whole value # stream up to the attention aggregation runs in a single launch # with analytic backward and second order; only the backward # anchors reach device memory. === - if self._cached_edge_csr_fn is not None: - self._cached_edge_csr_fn(edge_cache, "src", x.shape[0]) - x_local, rad_feat = self._cuda_value_train(x, edge_cache, radial_feat) - elif self._triton_value_path is not None and not training: + if self.edge_csr is not None: + self.edge_csr(edge_cache, "src", x.shape[0]) + x_local, rad_feat = self.cuda_train_value(x, edge_cache, radial_feat) + elif self.triton_infer_l_2_value is not None and not training: # === Steps 1-5 (fused Triton operators). ``so2_rotate_mix`` folds # the rotation and the radial degree mixing into one edge-parallel # kernel writing the focus-major layout; ``so2_mixing_stack`` runs @@ -2371,9 +2379,9 @@ def so2_message( # final store, keeping the inter-layer activations off the traced # graph. The rotate-mix backward reduces through the source CSR # view, which is built once per step and kept on the edge cache. === - if self._cached_edge_csr_fn is not None: - self._cached_edge_csr_fn(edge_cache, "src", x.shape[0]) - x_local, rad_feat = self._triton_value_path(x, edge_cache, radial_feat) + if self.edge_csr is not None: + self.edge_csr(edge_cache, "src", x.shape[0]) + x_local, rad_feat = self.triton_infer_l_2_value(x, edge_cache, radial_feat) else: # === Steps 1-3. Rotation, radial mixing and the focus-major cast === x_local, rad_feat = self._rotate_mix(x, edge_cache, radial_feat) diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py index 32c0905a2d..88e011bbc1 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py +++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py @@ -44,6 +44,7 @@ GraphLayout, NeighborGraph, apply_pair_exclusion, + compact_edges, compact_nodes, expand_node_values, frame_id_from_n_node, @@ -77,6 +78,7 @@ "build_neighbor_graph_ase", "canonicalize_neighbor_graph", "center_edge_pairs", + "compact_edges", "compact_nodes", "edge_env_mat", "edge_force_virial", diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py index 1a98504005..8debddb926 100644 --- a/deepmd/dpmodel/utils/neighbor_graph/graph.py +++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py @@ -266,6 +266,191 @@ def node_ownership_mask(n_node: Array, n_local: Array, n_total: int) -> Array: return index_in_frame < local_count +def _compact_csr( + order: Array, + row_ptr: Array, + edge_mask: Array, + new_position: Array, +) -> tuple[Array, Array]: + """Project an endpoint CSR view onto a compact edge axis. + + Parameters + ---------- + order + Edge permutation grouped by one endpoint with shape (E,). + row_ptr + CSR row pointers with shape (N + 1,). + edge_mask + Boolean survivor mask on the original edge axis with shape (E,). + new_position + Mapping from original edge indices to compact edge indices with shape + (E,). Removed edges map to -1. + + Returns + ------- + tuple[Array, Array] + The filtered edge permutation and rebuilt row pointers. + """ + xp = array_api_compat.array_namespace(order, row_ptr, edge_mask, new_position) + ordered_keep = xp.take(edge_mask, order, axis=0) + (ordered_keep_position,) = xp.nonzero(ordered_keep) + kept_edge = xp.take(order, ordered_keep_position, axis=0) + compact_order = xp.take(new_position, kept_edge, axis=0) + compact_row_ptr = _compact_csr_row_ptr(ordered_keep, row_ptr) + return compact_order, compact_row_ptr + + +def _compact_csr_row_ptr( + ordered_keep: Array, + row_ptr: Array, +) -> Array: + """Rebuild CSR row pointers after filtering the ordered edge axis.""" + xp = array_api_compat.array_namespace(ordered_keep, row_ptr) + prefix = xp.concat( + [ + xp.zeros( + (1,), + dtype=row_ptr.dtype, + device=array_api_compat.device(row_ptr), + ), + xp.cumulative_sum(xp.astype(ordered_keep, row_ptr.dtype), axis=0), + ], + axis=0, + ) + return xp.take(prefix, row_ptr, axis=0) + + +def compact_edges(graph: NeighborGraph) -> NeighborGraph: + """Drop masked edges and remap derived indices onto the compact edge axis. + + Complete destination and source CSR views are filtered in their own edge + order, remapped to the compact payload, and retained without sorting. This + preserves arbitrary mask placement inside CSR rows rather than assuming a + masked suffix. + + Parameters + ---------- + graph : NeighborGraph + Graph whose ``edge_mask`` selects the retained edge payload. Angle + fields must either both be present or both be absent. + + Returns + ------- + NeighborGraph + A graph containing only valid edges and angles. Complete CSR views are + retained; incomplete views are cleared. + + Raises + ------ + ValueError + If the angle fields are incomplete or have inconsistent shapes. + """ + import dataclasses + + xp = array_api_compat.array_namespace(graph.edge_mask) + has_angle_index = graph.angle_index is not None + has_angle_mask = graph.angle_mask is not None + if has_angle_index != has_angle_mask: + raise ValueError( + "compact_edges requires angle_index and angle_mask to both be set " + "or both be None" + ) + if has_angle_index: + if graph.angle_index.ndim != 2 or graph.angle_index.shape[0] != 2: + raise ValueError( + "compact_edges requires angle_index with shape (2, A); got " + f"{tuple(graph.angle_index.shape)}" + ) + if graph.angle_index.shape[1] != graph.angle_mask.shape[0]: + raise ValueError( + "compact_edges: angle_index and angle_mask disagree on A; got " + f"{graph.angle_index.shape[1]} and " + f"{graph.angle_mask.shape[0]}" + ) + + (keep_index,) = xp.nonzero(graph.edge_mask) + preserve_csr = all( + value is not None + for value in ( + graph.destination_order, + graph.destination_row_ptr, + graph.source_order, + graph.source_row_ptr, + ) + ) + if preserve_csr or has_angle_index: + survivor = xp.astype(graph.edge_mask, graph.edge_index.dtype) + rank = xp.cumulative_sum(survivor, axis=0) - survivor + new_position = xp.where( + graph.edge_mask, + rank, + xp.full_like(rank, -1), + ) + if preserve_csr: + if graph.destination_sorted: + destination_order = xp.arange( + keep_index.shape[0], + dtype=graph.destination_order.dtype, + device=array_api_compat.device(graph.destination_order), + ) + destination_row_ptr = _compact_csr_row_ptr( + graph.edge_mask, + graph.destination_row_ptr, + ) + else: + destination_order, destination_row_ptr = _compact_csr( + graph.destination_order, + graph.destination_row_ptr, + graph.edge_mask, + new_position, + ) + source_order, source_row_ptr = _compact_csr( + graph.source_order, + graph.source_row_ptr, + graph.edge_mask, + new_position, + ) + else: + destination_order = None + destination_row_ptr = None + source_order = None + source_row_ptr = None + fields = { + "edge_index": xp.take(graph.edge_index, keep_index, axis=1), + "edge_vec": xp.take(graph.edge_vec, keep_index, axis=0), + "edge_mask": xp.take(graph.edge_mask, keep_index, axis=0), + "destination_order": destination_order, + "destination_row_ptr": destination_row_ptr, + "source_order": source_order, + "source_row_ptr": source_row_ptr, + "destination_sorted": graph.destination_sorted and preserve_csr, + } + if has_angle_index: + # Angles address the original edge axis. An exclusive survivor count + # maps each retained edge to its compact position; angles touching a + # removed edge are discarded. + first = xp.take(new_position, graph.angle_index[0, :], axis=0) + second = xp.take(new_position, graph.angle_index[1, :], axis=0) + angle_keep = xp.logical_and( + xp.astype(graph.angle_mask, xp.bool), + xp.logical_and(first >= 0, second >= 0), + ) + (angle_keep_index,) = xp.nonzero(angle_keep) + fields["angle_index"] = xp.stack( + [ + xp.take(first, angle_keep_index, axis=0), + xp.take(second, angle_keep_index, axis=0), + ], + axis=0, + ) + fields["angle_mask"] = xp.take( + graph.angle_mask, + angle_keep_index, + axis=0, + ) + return dataclasses.replace(graph, **fields) + + def apply_pair_exclusion( graph: NeighborGraph, atype: Array, @@ -338,61 +523,7 @@ def apply_pair_exclusion( destination_sorted=False, ) if compact: - # Angle fields are a coupled pair (produced together by the angle - # builder): both present or both None. Fail fast on any inconsistent - # state — a partial or shape-mismatched pair is a caller bug that would - # otherwise remap silently wrong. - has_ai = out.angle_index is not None - has_am = out.angle_mask is not None - if has_ai != has_am: - raise ValueError( - "apply_pair_exclusion(compact=True): angle_index and angle_mask " - "must both be set or both be None; got " - f"angle_index={'set' if has_ai else 'None'}, " - f"angle_mask={'set' if has_am else 'None'}." - ) - if has_ai: - if out.angle_index.ndim != 2 or out.angle_index.shape[0] != 2: - raise ValueError( - "apply_pair_exclusion(compact=True): angle_index must have " - f"shape (2, A); got {tuple(out.angle_index.shape)}." - ) - if out.angle_index.shape[1] != out.angle_mask.shape[0]: - raise ValueError( - "apply_pair_exclusion(compact=True): angle_index (2, A) and " - f"angle_mask (A,) disagree on A: {out.angle_index.shape[1]} " - f"vs {out.angle_mask.shape[0]}." - ) - (keep_idx,) = xp.nonzero(out.edge_mask) - fields = { - "edge_index": out.edge_index[:, keep_idx], - "edge_vec": xp.take(out.edge_vec, keep_idx, axis=0), - "edge_mask": xp.take(out.edge_mask, keep_idx, axis=0), - } - if has_ai: - # Angles reference PRE-compaction edge positions; remap them to the - # compacted axis and drop any angle whose constituent edges were - # excluded. ``new_pos`` maps old edge position -> new position via an - # exclusive prefix sum over the survivors (-1 for dropped edges). - surv = xp.astype(out.edge_mask, out.edge_index.dtype) # (E,) 0/1 - rank = xp.cumulative_sum(surv, axis=0) - surv # survivors before me - new_pos = xp.where(out.edge_mask, rank, xp.full_like(rank, -1)) - a_new = xp.take(new_pos, out.angle_index[0, :], axis=0) - b_new = xp.take(new_pos, out.angle_index[1, :], axis=0) - both_survive = xp.logical_and(a_new >= 0, b_new >= 0) - angle_keep = xp.logical_and( - xp.astype(out.angle_mask, xp.bool), both_survive - ) - (angle_keep_idx,) = xp.nonzero(angle_keep) - fields["angle_index"] = xp.stack( - [ - xp.take(a_new, angle_keep_idx, axis=0), - xp.take(b_new, angle_keep_idx, axis=0), - ], - axis=0, - ) - fields["angle_mask"] = xp.take(out.angle_mask, angle_keep_idx, axis=0) - out = dataclasses.replace(out, **fields) + out = compact_edges(out) return out diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index 62693d547d..2bbd166d80 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/deepmd/pt/entrypoints/freeze_pt2.py @@ -379,7 +379,7 @@ def _tune_triton_configs(model: torch.nn.Module, target_device: torch.device) -> # tables, and must reflect the target device and any fresh registrations. for module in model.modules(): if isinstance(module, SO2Convolution) and module.triton_infer_level >= 2: - module._triton_value_path = make_triton_value_path(module) + module.triton_infer_l_2_value = make_triton_value_path(module) # The trace-time sendlist for the with-comm artifact embeds the address of a @@ -995,16 +995,16 @@ def _freeze_sezm_to_pt2( force_fused_scatter = target_device.type == "cuda" for module in model.modules(): if isinstance(module, SO2Linear): - module._force_block_diag_matmul = force_block_diag + module.force_block_diag_matmul = force_block_diag if isinstance(module, GeometricInitialEmbedding): - module._force_fused_scatter = force_fused_scatter + module.force_cuda_infer_l_1_scatter = force_fused_scatter # Sweep any Triton launch-table keys this checkpoint needs that are not # covered for the local GPU, so the traced graph bakes tuned launches. _tune_triton_configs(model, target_device) has_triton_value_path = any( - getattr(module, "_triton_value_path", None) is not None + getattr(module, "triton_infer_l_2_value", None) is not None for module in model.modules() ) if has_triton_value_path: diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index 70bf750827..8ea9409991 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -1065,19 +1065,19 @@ def __init__( # The fused convolution paths consume only the three structural rows of # each Wigner degree block. Source-gated attention bypasses that fused # convolution, so its dense per-edge rotations remain available. - self._wigner_free_conv = ( + self.cuda_infer_l_2_covers_all_blocks = ( self.bridging_switch is None and bool(self.blocks) and all( - getattr(block.so2_conv, "_cuda_conv_fn", None) is not None - and not block.so2_conv._cuda_conv_fn._compete + getattr(block.so2_conv, "cuda_infer_l_2_conv", None) is not None + and not block.so2_conv.cuda_infer_l_2_conv.focus_compete for block in self.blocks ) ) - self._packed_wigner_train = bool(self.blocks) and all( - getattr(block.so2_conv, "_cuda_value_train", None) is not None - and block.so2_conv._flash_atten_fn is not None - and block.so2_conv._flash_atten_trains + self.cuda_train_covers_all_blocks = bool(self.blocks) and all( + getattr(block.so2_conv, "cuda_train_value", None) is not None + and block.so2_conv.flash_attention is not None + and block.so2_conv.flash_attention_supports_training for block in self.blocks ) @@ -1085,8 +1085,8 @@ def __init__( # distance and are cheap enough that the compiler inlines them into # every consumer and re-evaluates them there. Behind an operator # boundary the chain runs once per step. - self._cuda_radial_fn = None - self._cuda_wigner_fn = None + self.cuda_infer_l_1_radial = None + self.cuda_infer_l_1_wigner = None if cuda_infer_level() >= 1: from deepmd.pt_expt.kernels.cuda.dpa4.edge_radial import ( make_cuda_edge_radial, @@ -1095,13 +1095,13 @@ def __init__( make_cuda_wigner_dense, ) - self._cuda_radial_fn = make_cuda_edge_radial( + self.cuda_infer_l_1_radial = make_cuda_edge_radial( self.edge_envelope, self.radial_basis ) # The dense Wigner pair otherwise costs five full-size passes # over the (E, D, D) tensors; the fused build pays only the # output writes. - self._cuda_wigner_fn = make_cuda_wigner_dense( + self.cuda_infer_l_1_wigner = make_cuda_wigner_dense( self.mp_init_lmax, self.compute_dtype ) @@ -1334,7 +1334,7 @@ def forward( # sparse-edge path, so ``forward`` keeps the original # bridging-free aggregation semantics. with nvtx_range("build_edge_cache"): - packed_wigner_candidate = self._packed_wigner_candidate( + packed_wigner_candidate = self.is_cute_infer_packed_wigner_candidate( type_ebed.device, type_ebed.dtype, extended_coord.dtype, @@ -1461,8 +1461,8 @@ def forward( x = x + force_embedding.to(dtype=block_dtype) if self.blocks and edge_cache.src.numel() > 0: edge_cache = edge_cache_to_dtype(edge_cache, block_dtype) - k1_dst_ptr, k1_source_order, k1_source_ptr = ( - self._prepare_cute_k1_sorted_metadata( + edge_cache = edge_cache._replace( + cute_infer_so2_metadata=self.prepare_cute_infer_so2_metadata( edge_cache, n_nodes, ) @@ -1473,9 +1473,6 @@ def forward( edge_cache, rad_feat_per_block, comm_dict=comm_dict, - k1_dst_ptr=k1_dst_ptr, - k1_source_order=k1_source_order, - k1_source_ptr=k1_source_ptr, ) # === Step 12. Final l=0 output mixing === @@ -1615,7 +1612,7 @@ def forward_with_edges( ) # === Step 3. Build edge cache once (sparse edges) === with nvtx_range("build_edge_cache"): - packed_wigner_candidate = self._packed_wigner_candidate( + packed_wigner_candidate = self.is_cute_infer_packed_wigner_candidate( type_ebed.device, type_ebed.dtype, extended_coord.dtype, @@ -1635,8 +1632,8 @@ def forward_with_edges( bridging_switch=self.bridging_switch, edge_envelope=self.edge_envelope, radial_basis=self.radial_basis, - fused_radial=(None if self.training else self._cuda_radial_fn), - fused_wigner=(None if self.training else self._cuda_wigner_fn), + fused_radial=(None if self.training else self.cuda_infer_l_1_radial), + fused_wigner=(None if self.training else self.cuda_infer_l_1_wigner), has_exclude_types=bool(self.exclude_types), edge_type_keep_mask=self._edge_type_keep_mask, # Random local-Z roll is a training-only augmentation; @@ -1739,8 +1736,8 @@ def forward_with_edges( x = x + force_embedding.to(dtype=block_dtype) if self.blocks: edge_cache = edge_cache_to_dtype(edge_cache, block_dtype) - k1_dst_ptr, k1_source_order, k1_source_ptr = ( - self._prepare_cute_k1_sorted_metadata( + edge_cache = edge_cache._replace( + cute_infer_so2_metadata=self.prepare_cute_infer_so2_metadata( edge_cache, n_nodes, ) @@ -1751,9 +1748,6 @@ def forward_with_edges( edge_cache, rad_feat_per_block, comm_dict=comm_dict, - k1_dst_ptr=k1_dst_ptr, - k1_source_order=k1_source_order, - k1_source_ptr=k1_source_ptr, ) # === Step 11. Keep the owned-atom rows for the read-out === @@ -1774,43 +1768,63 @@ def forward_with_edges( return descriptor.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), x.contiguous() @torch.jit.unused - def _run_output_readout(self, ffn_in: torch.Tensor) -> torch.Tensor: - """Return the residual-inclusive scalar output with guarded CuTe routing.""" - from deepmd.pt_expt.kernels.cute.sezm.readout_l0 import ( - run_neo_output_readout, + def run_cute_infer_readout( + self, + ffn_in: torch.Tensor, + ) -> torch.Tensor | None: + """Run the CuTe readout when its exact inference contract matches. + + Parameters + ---------- + ffn_in : torch.Tensor + Equivariant readout input with shape ``(N, D, 1, C)``. + + Returns + ------- + torch.Tensor or None + Residual-inclusive scalar output with shape ``(N, C)``, or ``None`` + when the exact CuTe contract is not satisfied. + """ + if self.training or not cute_runtime_policy.is_cute_infer_enabled(): + return None + from deepmd.pt_expt.kernels.cute.sezm.output_grid.readout_l0 import ( + maybe_run_neo_readout_l0, ) - return run_neo_output_readout( + return maybe_run_neo_readout_l0( self.output_ffn, ffn_in, - parameters_frozen=self._readout_parameters_are_frozen(), ) - def _readout_parameters_are_frozen(self) -> bool: - """Return whether descriptor inference can exclude every weight gradient.""" - return not self.training and not any( - parameter.requires_grad for parameter in self.parameters() - ) - - def _prepare_cute_k1_sorted_metadata( + def prepare_cute_infer_so2_metadata( self, edge_cache: EdgeFeatureCache, n_nodes: int, - ) -> tuple[ - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor | None, - ]: - """Build one per-forward CSR tensor bundle for eligible K1 blocks.""" + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: + """Build the edge metadata consumed by the CuTe SO2 implementation. + + Parameters + ---------- + edge_cache : EdgeFeatureCache + Per-forward cache carrying the destination-major edge payload. + n_nodes : int + Number of nodes addressed by the edge payload. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor, torch.Tensor] or None + Destination row pointers, source order, and source row pointers, or + ``None`` when the exact CuTe contract is not satisfied. + """ if ( self.training or not edge_cache.destinations_sorted or edge_cache.D_packed is None or edge_cache.edge_src_gate is not None ): - return None, None, None + return None if not cute_runtime_policy.is_cute_infer_enabled(): - return None, None, None + return None return build_sorted_edge_index_metadata( edge_cache.src, edge_cache.dst, @@ -1824,9 +1838,6 @@ def _forward_blocks( edge_cache: EdgeFeatureCache, radial_feat_per_block: list[torch.Tensor], comm_dict: dict[str, torch.Tensor] | None = None, - k1_dst_ptr: torch.Tensor | None = None, - k1_source_order: torch.Tensor | None = None, - k1_source_ptr: torch.Tensor | None = None, ) -> torch.Tensor: """ Run the interaction blocks with optional depth attention. @@ -1862,9 +1873,6 @@ def _forward_blocks( edge_cache, blk_radial, comm_dict=self._block_comm(i, comm_dict), - k1_dst_ptr=k1_dst_ptr, - k1_source_order=k1_source_order, - k1_source_ptr=k1_source_ptr, ) return x @@ -1893,9 +1901,6 @@ def node_l0_extractor(v: torch.Tensor) -> torch.Tensor: blk_radial, unit_history=truncated_unit_history, comm_dict=self._block_comm(i, comm_dict), - k1_dst_ptr=k1_dst_ptr, - k1_source_order=k1_source_order, - k1_source_ptr=k1_source_ptr, ) unit_history.append(so2_unit_output) unit_history.extend(ffn_unit_outputs) @@ -1929,9 +1934,6 @@ def node_l0_extractor(v: torch.Tensor) -> torch.Tensor: blk_radial, unit_history=truncated_block_history, comm_dict=self._block_comm(i, comm_dict), - k1_dst_ptr=k1_dst_ptr, - k1_source_order=k1_source_order, - k1_source_ptr=k1_source_ptr, ) block_history.append(block_summary) x = block_output @@ -1946,14 +1948,30 @@ def node_l0_extractor(v: torch.Tensor) -> torch.Tensor: ).to(dtype=self.dtype) return x - def _packed_wigner_candidate( + def is_cute_infer_packed_wigner_candidate( self, device: torch.device, dtype: torch.dtype, geometry_dtype: torch.dtype, ) -> bool: - """Return whether all blocks satisfy packed K1's stable contract.""" - from deepmd.pt_expt.kernels.cute.sezm.k1 import ( + """Return whether all blocks satisfy the packed CuTe SO2 contract. + + Parameters + ---------- + device : torch.device + Device that executes the descriptor blocks. + dtype : torch.dtype + Descriptor block dtype. + geometry_dtype : torch.dtype + Edge-geometry compute dtype. + + Returns + ------- + bool + Whether packed Wigner storage can replace dense storage for every + interaction block. + """ + from deepmd.pt_expt.kernels.cute.sezm.so2.operation import ( is_packed_wigner_candidate, ) @@ -1966,7 +1984,6 @@ def _packed_wigner_candidate( self.radial_basis, self.radial_embedding, self.edge_envelope, - self.wigner_calc, *((self.inner_clamp,) if self.inner_clamp is not None else ()), *((self.bridging_switch,) if self.bridging_switch is not None else ()), ), @@ -2012,14 +2029,14 @@ def _apply_readout(self, x: torch.Tensor, n_rows: int) -> torch.Tensor: x_ro = x[:, : self.node_readout_dim, :, :].to(dtype=self.compute_dtype) for layer in self.readout_pre_layers: x_ro = x_ro + layer(x_ro) - if not self.readout_pre_layers and self.so3_readout != "none": - if ( - not torch.jit.is_scripting() - and cute_runtime_policy.is_cute_infer_enabled() - ): - return self._run_output_readout(x_ro).reshape( - n_rows, 1, 1, self.channels - ) + if ( + not self.readout_pre_layers + and self.so3_readout != "none" + and not torch.jit.is_scripting() + ): + accelerated = self.run_cute_infer_readout(x_ro) + if accelerated is not None: + return accelerated.reshape(n_rows, 1, 1, self.channels) if self.so3_readout == "none": return (x_ro + self.output_ffn(x_ro))[:, 0:1, :, :] return x_ro[:, 0:1, :, :] + self.output_ffn.forward_scalar(x_ro) @@ -2054,8 +2071,8 @@ def _build_full_wigner(self) -> bool: if not self._need_full_wigner: return False if self.training: - return not self._packed_wigner_train - return not self._wigner_free_conv + return not self.cuda_train_covers_all_blocks + return not self.cuda_infer_l_2_covers_all_blocks def _shared_wigner_runs( self, @@ -2089,13 +2106,13 @@ def _shared_wigner_runs( if edge_cache.csr_cache is None: return None if self.training: - if not self._packed_wigner_train: + if not self.cuda_train_covers_all_blocks: return None - fused = self.blocks[0].so2_conv._cuda_value_train + fused = self.blocks[0].so2_conv.cuda_train_value else: - if not self._wigner_free_conv: + if not self.cuda_infer_l_2_covers_all_blocks: return None - fused = self.blocks[0].so2_conv._cuda_conv_fn + fused = self.blocks[0].so2_conv.cuda_infer_l_2_conv if fused is None or lmax > self.lmax: return None return fused.edge_runs(edge_cache)[:, 1 : (lmax + 1) ** 2] @@ -2117,17 +2134,7 @@ def _build_gie_zonal_coupling( reconstructed from the edge quaternion via the m=0-only path. """ if edge_cache.D_packed is not None: - if self.gie_zonal_wigner_calc is None: - return None - mp_coupling = edge_cache.D_packed.index_select( - 1, - self.gie.packed_zonal_offsets, - ) - extra_coupling = self.gie_zonal_wigner_calc.forward_zonal( - self._edge_quaternion(edge_cache), - lmin=self.lmax + 1, - ) - return torch.cat([mp_coupling, extra_coupling], dim=1) + return self.build_cute_infer_zonal_coupling(edge_cache) if edge_cache.Dt_full is None: calc = self.gie_zonal_wigner_calc or self.wigner_calc shared = self._shared_wigner_runs(edge_cache, calc.lmax) @@ -2150,6 +2157,39 @@ def _build_gie_zonal_coupling( ) return torch.cat([mp_coupling, extra_coupling], dim=1) + def build_cute_infer_zonal_coupling( + self, + edge_cache: EdgeFeatureCache, + ) -> torch.Tensor: + """Extract the GIE zonal coupling from packed Wigner storage. + + Parameters + ---------- + edge_cache : EdgeFeatureCache + Per-forward cache carrying CuTe packed Wigner storage. + + Returns + ------- + torch.Tensor + Zonal coupling with shape ``(E, D_node - 1)``. + + Raises + ------ + ValueError + If the edge cache does not carry packed Wigner storage. + """ + D_packed = edge_cache.D_packed + if D_packed is None: + raise ValueError("CuTe zonal coupling requires packed Wigner storage") + mp_coupling = D_packed.index_select(1, self.gie.packed_zonal_offsets) + if self.gie_zonal_wigner_calc is None: + return mp_coupling + extra_coupling = self.gie_zonal_wigner_calc.forward_zonal( + self._edge_quaternion(edge_cache), + lmin=self.lmax + 1, + ) + return torch.cat([mp_coupling, extra_coupling], dim=1) + def _apply_charge_spin_embedding( self, type_ebed: torch.Tensor, diff --git a/deepmd/pt/model/descriptor/sezm_nn/block.py b/deepmd/pt/model/descriptor/sezm_nn/block.py index 7ef5a97f3d..7a27ca99c3 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/block.py +++ b/deepmd/pt/model/descriptor/sezm_nn/block.py @@ -33,9 +33,6 @@ PRECISION_DICT, RESERVED_PRECISION_DICT, ) -from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import ( - is_cute_infer_enabled, -) from deepmd.utils.version import ( check_version_compatibility, ) @@ -677,9 +674,6 @@ def forward( radial_feat: torch.Tensor, unit_history: list[torch.Tensor] | None = None, comm_dict: dict[str, torch.Tensor] | None = None, - k1_dst_ptr: torch.Tensor | None = None, - k1_source_order: torch.Tensor | None = None, - k1_source_ptr: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, @@ -725,9 +719,6 @@ def forward( radial_feat, unit_history, comm_dict, - k1_dst_ptr, - k1_source_order, - k1_source_ptr, ) def _extract_l0_from_canonical(self, value: torch.Tensor) -> torch.Tensor: @@ -767,9 +758,6 @@ def _run_so2_unit( edge_cache: EdgeFeatureCache, radial_feat: torch.Tensor, comm_dict: dict[str, torch.Tensor] | None = None, - k1_dst_ptr: torch.Tensor | None = None, - k1_source_order: torch.Tensor | None = None, - k1_source_ptr: torch.Tensor | None = None, ) -> torch.Tensor: """ Run the SO(2) unit without an outer block-level residual shortcut. @@ -805,9 +793,6 @@ def _run_so2_unit( x_, edge_cache_no_proj, radial_feat_, - k1_dst_ptr, - k1_source_order, - k1_source_ptr, ), x, radial_feat, @@ -818,9 +803,6 @@ def _run_so2_unit( x, edge_cache, radial_feat, - k1_dst_ptr, - k1_source_order, - k1_source_ptr, ) def _run_so2_unit_impl( @@ -828,30 +810,30 @@ def _run_so2_unit_impl( x: torch.Tensor, edge_cache: EdgeFeatureCache, radial_feat: torch.Tensor, - k1_dst_ptr: torch.Tensor | None = None, - k1_source_order: torch.Tensor | None = None, - k1_source_ptr: torch.Tensor | None = None, ) -> torch.Tensor: """Run the SO(2) unit implementation.""" - if not self.training and is_cute_infer_enabled(): - from deepmd.pt_expt.kernels.cute.sezm.k1 import ( - maybe_run_cute_k1, + if edge_cache.D_packed is not None: + from deepmd.pt_expt.kernels.cute.sezm.so2.operation import ( + maybe_run_cute_so2, ) - cute_out = maybe_run_cute_k1( + metadata = edge_cache.cute_infer_so2_metadata + destination_row_ptr, source_order, source_row_ptr = ( + (None, None, None) if metadata is None else metadata + ) + accelerated = maybe_run_cute_so2( self, x, edge_cache, radial_feat, - dst_ptr=k1_dst_ptr, - source_order=k1_source_order, - source_ptr=k1_source_ptr, + dst_ptr=destination_row_ptr, + source_order=source_order, + source_ptr=source_row_ptr, ) - if cute_out is not None: - return cute_out - if edge_cache.D_packed is not None: + if accelerated is not None: + return accelerated raise RuntimeError( - "packed Wigner cache reached an ineligible Neo K1 dispatch" + "packed Wigner cache reached an ineligible Neo SO2 dispatch" ) n_node = x.shape[0] @@ -914,9 +896,6 @@ def _forward_with_residual_shortcuts( radial_feat: torch.Tensor, unit_history: list[torch.Tensor] | None = None, comm_dict: dict[str, torch.Tensor] | None = None, - k1_dst_ptr: torch.Tensor | None = None, - k1_source_order: torch.Tensor | None = None, - k1_source_ptr: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, @@ -952,9 +931,6 @@ def _forward_with_residual_shortcuts( edge_cache, radial_feat, comm_dict, - k1_dst_ptr, - k1_source_order, - k1_source_ptr, ) so2_state = x + so2_unit_output @@ -974,9 +950,6 @@ def _forward_with_full_attn_res( radial_feat: torch.Tensor, unit_history: list[torch.Tensor] | None = None, comm_dict: dict[str, torch.Tensor] | None = None, - k1_dst_ptr: torch.Tensor | None = None, - k1_source_order: torch.Tensor | None = None, - k1_source_ptr: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, @@ -1020,9 +993,6 @@ def _forward_with_full_attn_res( edge_cache, radial_feat, comm_dict, - k1_dst_ptr, - k1_source_order, - k1_source_ptr, ) with nvtx_range("ffn"): @@ -1051,9 +1021,6 @@ def _forward_with_block_attn_res( radial_feat: torch.Tensor, unit_history: list[torch.Tensor] | None = None, comm_dict: dict[str, torch.Tensor] | None = None, - k1_dst_ptr: torch.Tensor | None = None, - k1_source_order: torch.Tensor | None = None, - k1_source_ptr: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, torch.Tensor | None, @@ -1097,9 +1064,6 @@ def _forward_with_block_attn_res( edge_cache, radial_feat, comm_dict, - k1_dst_ptr, - k1_source_order, - k1_source_ptr, ) with nvtx_range("ffn"): diff --git a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py index 2e21b1b8be..11665b56e3 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py +++ b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py @@ -28,6 +28,9 @@ from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import ( is_cute_infer_enabled, ) +from deepmd.pt_expt.kernels.cute.sezm.so2.metadata import ( + build_sorted_edge_index_metadata as build_sorted_edge_index_metadata, +) from .utils import ( get_promoted_dtype, @@ -95,6 +98,9 @@ class EdgeFeatureCache(NamedTuple): Lazy cache for endpoint CSR views used by segmented accelerated operators, keyed by endpoint role (``"dst"`` or ``"src"``). Built once per step and shared by every consumer. + cute_infer_so2_metadata + Optional ``(destination_row_ptr, source_order, source_row_ptr)`` tuple + prepared once for the destination-sorted CuTe SO2 path. edge_src_gate Optional per-edge Source Freeze Propagation Gate (SFPG) weight with shape (E, 1). Equals ``eta[src]`` where @@ -107,7 +113,7 @@ class EdgeFeatureCache(NamedTuple): the frozen zone from propagating information along its outgoing edges. destinations_sorted - Host-side provenance indicating that ``dst`` is nondecreasing. CuTe K1 + Host-side provenance indicating that ``dst`` is nondecreasing. CuTe SO2 may only consume caches carrying this guarantee. D_packed Opt-in Neo packed Wigner panel with shape (E, 46). This is kept @@ -129,6 +135,9 @@ class EdgeFeatureCache(NamedTuple): csr_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] | None = None edge_src_gate: torch.Tensor | None = None edge_quat: torch.Tensor | None = None + cute_infer_so2_metadata: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = ( + None + ) destinations_sorted: bool = False D_packed: torch.Tensor | None = None @@ -145,64 +154,6 @@ def _separate_packed_wigner( return None, None, D_full -def build_sorted_edge_index_metadata( - src: torch.Tensor, - dst: torch.Tensor, - n_nodes: int, - *, - validate_sorted: bool = False, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Build destination and source CSR metadata for one sorted edge list.""" - if src.dim() != 1 or dst.dim() != 1: - raise ValueError("src and dst must be one-dimensional") - if src.shape != dst.shape: - raise ValueError("src and dst must have the same shape") - if src.device != dst.device: - raise ValueError("src and dst must be on the same device") - if src.dtype not in (torch.int32, torch.int64): - raise TypeError("src must have dtype int32 or int64") - if dst.dtype not in (torch.int32, torch.int64): - raise TypeError("dst must have dtype int32 or int64") - if n_nodes < 0: - raise ValueError("n_nodes must be non-negative") - if src.numel() > 2**31 - 1: - raise ValueError("sorted edge metadata requires E <= 2**31 - 1") - - src = src.contiguous() - dst = dst.contiguous() - if validate_sorted and dst.numel() > 1: - torch._assert_async( - torch.all(dst[1:] >= dst[:-1]), - "Neo K1 destinations_sorted=True requires monotonically " - "nondecreasing destination indices", - ) - dst_boundaries = torch.arange( - n_nodes + 1, - device=dst.device, - dtype=dst.dtype, - ) - dst_ptr = torch.searchsorted( - dst, - dst_boundaries, - out_int32=True, - ).contiguous() - - source_order_i64 = torch.argsort(src, stable=True) - sorted_src = src.index_select(0, source_order_i64) - src_boundaries = torch.arange( - n_nodes + 1, - device=src.device, - dtype=src.dtype, - ) - source_ptr = torch.searchsorted( - sorted_src, - src_boundaries, - out_int32=True, - ).contiguous() - source_order = source_order_i64.to(dtype=torch.int32).contiguous() - return dst_ptr, source_order, source_ptr - - def cached_edge_csr( edge_cache: EdgeFeatureCache, endpoint: str, n_node: int | torch.SymInt ) -> tuple[torch.Tensor, torch.Tensor]: @@ -446,7 +397,7 @@ def build_edge_cache( Callable that converts edge-aligned quaternions into packed Wigner-D blocks. packed_wigner_candidate - Whether descriptor-level strict-FP32 K1 checks passed. Concrete edge + Whether descriptor-level strict-FP32 SO2 checks passed. Concrete edge count and destination ordering are checked before panel generation. build_wigner Whether to materialize Wigner-D blocks for the SO(2) path. @@ -621,7 +572,7 @@ def build_edge_cache_from_edges( Callable that converts edge-aligned quaternions into packed Wigner-D blocks. packed_wigner_candidate - Whether descriptor-level strict-FP32 K1 checks passed. Concrete edge + Whether descriptor-level strict-FP32 SO2 checks passed. Concrete edge count and destination ordering are checked before panel generation. destinations_sorted Host-side provenance that ``edge_index[1]`` is nondecreasing. @@ -767,7 +718,7 @@ def _build_edge_wigner( quaternion is returned and the blocks are ``None``; the geometric initial embedding reconstructs the zonal coupling from the quaternion. packed_wigner - Whether the exact packed K1 eligibility contract has passed. + Whether the exact packed SO2 eligibility contract has passed. Returns ------- @@ -795,7 +746,7 @@ def _build_edge_wigner( if not build_full: return None, None, edge_quat if packed_wigner: - from deepmd.pt_expt.kernels.cute.sezm.k4_wignerd import ( + from deepmd.pt_expt.kernels.cute.sezm.wignerd import ( run_cute_wignerd, ) @@ -819,7 +770,7 @@ def _packed_wigner_edges_eligible( runtime_dtypes: tuple[torch.dtype, ...] = (), ) -> bool: """Finish packed eligibility from scalar shape and provenance metadata.""" - from deepmd.pt_expt.kernels.cute.sezm.k1 import ( + from deepmd.pt_expt.kernels.cute.sezm.so2.operation import ( packed_wigner_edges_eligible, ) @@ -1137,7 +1088,7 @@ def edge_cache_to_dtype( # Use local variables with explicit None check and assignment. _D_full = cache.D_full _Dt_full = cache.Dt_full - _D_packed = cache.D_packed + cached_D_packed = cache.D_packed _edge_src_gate = cache.edge_src_gate _edge_quat = cache.edge_quat D_full: torch.Tensor | None = None @@ -1149,8 +1100,8 @@ def edge_cache_to_dtype( D_full = _D_full.to(dtype=dtype) if _Dt_full is not None: Dt_full = _Dt_full.to(dtype=dtype) - if _D_packed is not None: - D_packed = _D_packed.to(dtype=dtype) + if cached_D_packed is not None: + D_packed = cached_D_packed.to(dtype=dtype) if _edge_src_gate is not None: edge_src_gate = _edge_src_gate.to(dtype=dtype) if _edge_quat is not None: @@ -1175,5 +1126,6 @@ def edge_cache_to_dtype( csr_cache=None if cache.csr_cache is None else dict(cache.csr_cache), edge_src_gate=edge_src_gate, edge_quat=edge_quat, + cute_infer_so2_metadata=cache.cute_infer_so2_metadata, destinations_sorted=cache.destinations_sorted, ) diff --git a/deepmd/pt/model/descriptor/sezm_nn/embedding.py b/deepmd/pt/model/descriptor/sezm_nn/embedding.py index 5bfe764495..017f14a4da 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/embedding.py +++ b/deepmd/pt/model/descriptor/sezm_nn/embedding.py @@ -35,7 +35,7 @@ from deepmd.pt.utils.utils import ( get_generator, ) -from deepmd.pt_expt.kernels.cute.sezm.k1_wigner_layout import ( +from deepmd.pt_expt.kernels.cute.sezm.so2.wigner_layout import ( PACKED_VALUE_COUNT, ZONAL_PANEL_OFFSETS, ) @@ -226,17 +226,17 @@ def __init__( # The reference composition materializes the per-edge message, an # (E, D-1, C) tensor that dominates the cost of this module. The fused # operator keeps it in registers and reduces through the destination CSR. - self._cuda_scatter = False + self.cuda_infer_l_1_scatter = False # ``None`` keeps the runtime ``zonal_coupling.is_cuda`` dispatch; the # freeze pins it to the AOTI target because tracing always runs on CPU. - self._force_fused_scatter: bool | None = None + self.force_cuda_infer_l_1_scatter: bool | None = None if cuda_infer_level() >= 1 and self.dtype is torch.float32: from deepmd.pt_expt.kernels.cuda.dpa4.zonal_scatter import ( op_available, supported, ) - self._cuda_scatter = op_available() and supported( + self.cuda_infer_l_1_scatter = op_available() and supported( self.lmax, self.ebed_dim - 1, self.channels ) @@ -303,35 +303,28 @@ def forward( self.zonal_m0_col_index_for_row, ] # (E, D-1) - # === Step 3. Optional fused message construction and reduction === - if not self.training: - from deepmd.pt_expt.kernels.cute.sezm.gie import ( - is_cute_gie_enabled, - maybe_run_cute_gie, - ) - - if is_cute_gie_enabled(device) and spin_l1_message is None: - cute_out = maybe_run_cute_gie( - self, - n_nodes=n_nodes, - edge_cache=edge_cache, - radial_feat=radial_feat, - zonal_coupling=zonal_coupling, - ) - if cute_out is not None: - return cute_out + # === Step 3. Optional backend message construction and reduction === + accelerated = self.run_cute_infer_gie( + n_nodes=n_nodes, + edge_cache=edge_cache, + radial_feat=radial_feat, + zonal_coupling=zonal_coupling, + spin_l1_message=spin_l1_message, + ) + if accelerated is not None: + return accelerated # === Step 4. Eager fallback: 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 5, so # it takes over whenever nothing else joins the message in between. if ( - self._can_fuse_scatter(zonal_coupling) + self.can_run_cuda_infer_l_1_scatter(zonal_coupling) and spin_l1_message is None and edge_cache.edge_src_gate is None and edge_cache.csr_cache is not None ): - return self.forward_fused_scatter( + return self.run_cuda_infer_l_1_scatter( n_nodes, edge_cache, radial_feat, zonal_coupling ) @@ -374,16 +367,71 @@ def forward( out.mul_(edge_cache.inv_sqrt_deg) return out - def _can_fuse_scatter(self, zonal_coupling: torch.Tensor) -> bool: - """Return whether the fused scatter serves the runtime or trace target.""" + def run_cute_infer_gie( + self, + *, + n_nodes: int | torch.SymInt, + edge_cache: EdgeFeatureCache, + radial_feat: torch.Tensor, + zonal_coupling: torch.Tensor, + spin_l1_message: torch.Tensor | None = None, + ) -> torch.Tensor | None: + """Run the CuTe GIE path when its exact inference contract matches. + + Parameters + ---------- + n_nodes : int or torch.SymInt + Number of nodes addressed by the edge cache. + edge_cache : EdgeFeatureCache + Per-forward geometric edge cache. + radial_feat : torch.Tensor + Radial features with shape ``(E, lmax, C)``. + zonal_coupling : torch.Tensor + Zonal coupling with shape ``(E, D - 1)``. + spin_l1_message : torch.Tensor, optional + Native-spin message with shape ``(E, 3, C)``. + + Returns + ------- + torch.Tensor or None + Initial embedding with shape ``(N, D, C)``, or ``None`` when the + exact CuTe contract is not satisfied. + """ + if self.training or spin_l1_message is not None: + return None + from deepmd.pt_expt.kernels.cute.sezm.gie import ( + maybe_run_cute_gie, + ) + + return maybe_run_cute_gie( + self, + n_nodes=n_nodes, + edge_cache=edge_cache, + radial_feat=radial_feat, + zonal_coupling=zonal_coupling, + ) + + def can_run_cuda_infer_l_1_scatter(self, zonal_coupling: torch.Tensor) -> bool: + """Return whether the fused scatter serves the runtime or trace target. + + Parameters + ---------- + zonal_coupling : torch.Tensor + Zonal coupling with shape ``(E, D - 1)``. + + Returns + ------- + bool + Whether the CUDA inference level-one scatter is eligible. + """ target_is_cuda = ( zonal_coupling.is_cuda - if self._force_fused_scatter is None - else self._force_fused_scatter + if self.force_cuda_infer_l_1_scatter is None + else self.force_cuda_infer_l_1_scatter ) - return self._cuda_scatter and not self.training and target_is_cuda + return self.cuda_infer_l_1_scatter and not self.training and target_is_cuda - def forward_fused_scatter( + def run_cuda_infer_l_1_scatter( self, n_nodes: int | torch.SymInt, edge_cache: EdgeFeatureCache, diff --git a/deepmd/pt/model/descriptor/sezm_nn/grid_net.py b/deepmd/pt/model/descriptor/sezm_nn/grid_net.py index bc6b0aa1c6..d9bc63a5c4 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/grid_net.py +++ b/deepmd/pt/model/descriptor/sezm_nn/grid_net.py @@ -714,7 +714,7 @@ def __init__( # transposed so both matrices are read row-major by grid point. # The operator is instantiated per coefficient-slot count, which this # projector fixes, so the choice is made once here rather than per call. - self._grid_pair_fn = None + self.cuda_infer_l_1_grid_pair = None if ( cuda_infer_level() >= 1 and self.projector.to_grid_mat.dtype is torch.float32 @@ -727,7 +727,7 @@ def __init__( slots = int(self.projector.to_grid_mat.shape[1]) if op_available() and slots in SUPPORTED_SLOTS: - self._grid_pair_fn = grid_pair + self.cuda_infer_l_1_grid_pair = grid_pair # The training form differentiates the same expression inside the # force graph: a Triton tensor-core sandwich (grid-axis blocks, one # resident output tile) with analytic first and second order, one @@ -738,7 +738,7 @@ def __init__( # dense section is small and the operator's dispatch chain costs # more than its kernels save on the host-bound configurations, so # the narrow grids stay with the compiler. - self._grid_pair_train_fn = None + self.triton_train_l_1_grid_pair = None if ( triton_train_level() >= 1 and self.projector.to_grid_mat.dtype is torch.float32 @@ -750,7 +750,7 @@ def __init__( ) if GRID_PAIR_TRITON_AVAILABLE: - self._grid_pair_train_fn = grid_pair_train + self.triton_train_l_1_grid_pair = grid_pair_train self.register_buffer( "_from_grid_t", self.projector.from_grid_mat.transpose(0, 1).contiguous(), @@ -839,7 +839,7 @@ def forward_scalar( because materializing a scalar-only fallback grid would be slower than that fused operator. """ - if self._grid_pair_fn is not None and not self.training: + if self.cuda_infer_l_1_grid_pair is not None and not self.training: return self._slice_scalar_layout(self.forward(query, context)) return self._forward(query, context, scalar_only=True) @@ -1057,22 +1057,8 @@ def _pair_grid( c_wide = left.shape[3] // self.n_frames if c_wide % 32 != 0 or left.shape != right.shape: return None - if ( - not self.training - and cute_runtime_policy.is_cute_infer_enabled() - and _inference_mode_is_frozen(self) - ): - from deepmd.pt_expt.kernels.cute.sezm.output_grid_product import ( - maybe_run_cute_output_grid_product, - ) - - candidate = maybe_run_cute_output_grid_product( - left, - right, - self.projector.to_grid_mat, - self.projector.from_grid_mat, - n_frames=self.n_frames, - ) + if not self.training: + candidate = self.run_cute_infer_grid_pair(left, right) if candidate is not None: return candidate if self.training: @@ -1081,20 +1067,20 @@ def _pair_grid( # operator carries an autocast rule, so under AMP it runs the # same bf16-with-fp32-accumulation regime as the dense einsum # composition it replaces. - if self._grid_pair_train_fn is None: + if self.triton_train_l_1_grid_pair is None: return None - return self._grid_pair_train_fn( + return self.triton_train_l_1_grid_pair( left, right, self.projector.to_grid_mat, self._from_grid_t, self.n_frames, ) - if self._grid_pair_fn is None or left.shape[2] != 1: + if self.cuda_infer_l_1_grid_pair is None or left.shape[2] != 1: return None n_batch, coeff_dim = left.shape[0], left.shape[1] flat_p = coeff_dim * self.n_frames - out = self._grid_pair_fn( + out = self.cuda_infer_l_1_grid_pair( left.reshape(n_batch, flat_p, c_wide), right.reshape(n_batch, flat_p, c_wide), self.projector.to_grid_mat, @@ -1102,6 +1088,43 @@ def _pair_grid( ) return out.reshape(n_batch, coeff_dim, 1, self.n_frames * c_wide) + def run_cute_infer_grid_pair( + self, + left: torch.Tensor, + right: torch.Tensor, + ) -> torch.Tensor | None: + """Run the CuTe grid product when its exact inference contract matches. + + Parameters + ---------- + left : torch.Tensor + Left coefficient operand with shape ``(N, D, F, n_frames * C)``. + right : torch.Tensor + Right coefficient operand with the same shape as ``left``. + + Returns + ------- + torch.Tensor or None + Coefficient result with the same shape as ``left``, or ``None`` + when the exact CuTe contract is not satisfied. + """ + if ( + not cute_runtime_policy.is_cute_infer_enabled() + or not _inference_mode_is_frozen(self) + ): + return None + from deepmd.pt_expt.kernels.cute.sezm.output_grid.product import ( + maybe_run_cute_output_grid_product, + ) + + return maybe_run_cute_output_grid_product( + left, + right, + self.projector.to_grid_mat, + self.projector.from_grid_mat, + n_frames=self.n_frames, + ) + def _project_pair_in_one_transform( self, left: torch.Tensor, diff --git a/deepmd/pt/model/descriptor/sezm_nn/lora.py b/deepmd/pt/model/descriptor/sezm_nn/lora.py index 1b60dfa5de..9c41d0da7a 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/lora.py +++ b/deepmd/pt/model/descriptor/sezm_nn/lora.py @@ -490,17 +490,17 @@ def _clear_sezm_compile_cache(model: nn.Module) -> None: from deepmd.pt.model.model.sezm_model import ( _clear_shared_sezm_compile_cache, ) - from deepmd.pt_expt.kernels.cute.sezm.k1 import ( - invalidate_cute_k1_state, + from deepmd.pt_expt.kernels.cute.sezm.so2.operation import ( + invalidate_cute_so2_state, ) readout_invalidator = None for m in model.modules(): - invalidate_cute_k1_state(m) + invalidate_cute_so2_state(m) if hasattr(m, "_neo_sm80_readout_input_fold_cache"): if readout_invalidator is None: - from deepmd.pt_expt.kernels.cute.sezm.readout_l0 import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid.readout_l0 import ( invalidate_neo_readout_input_fold, ) diff --git a/deepmd/pt/model/descriptor/sezm_nn/so2.py b/deepmd/pt/model/descriptor/sezm_nn/so2.py index 1c4ed431e0..4c5d7efd73 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/so2.py +++ b/deepmd/pt/model/descriptor/sezm_nn/so2.py @@ -370,7 +370,7 @@ def __init__( # Export override for the block-diagonal vs dense matmul branch below. # ``None`` keeps the runtime ``x_flat.is_cuda`` dispatch; the freeze sets # it so the AOTI graph follows the *target* device, not the CPU trace. - self._force_block_diag_matmul: bool | None = None + self.force_block_diag_matmul: bool | None = None # The assembled SO(2) weight is block-diagonal over |m| groups; the # forward contracts only the diagonal blocks (see _block_diagonal_matmul). @@ -382,7 +382,7 @@ def __init__( # Triton BN=64 block-diagonal GEMM that consumes the strided operands # without a contiguity copy. Bound only when Triton is available and every # block width aligns to BN=64; otherwise the eager path is kept. - self._block_diag_gemm = None + self.triton_l_1_block_diag_gemm = None self.triton_infer_level = triton_infer_level() self.triton_train_level = triton_train_level() if max(self.triton_infer_level, self.triton_train_level) >= 1: @@ -395,7 +395,7 @@ def __init__( if SO2_BLOCK_GEMM_TRITON_AVAILABLE and slices_supported( self._block_diag_slices ): - self._block_diag_gemm = block_diag_gemm + self.triton_l_1_block_diag_gemm = block_diag_gemm def forward(self, x: torch.Tensor) -> torch.Tensor: """ @@ -431,12 +431,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # trips an Inductor AVX2 C++ codegen bug, so only CPU needs it. Every other # device uses the block-diagonal contraction, which skips the structural # off-|m| zeros. ``make_fx`` resolves this Python branch at trace time, so - # the freeze pins ``_force_block_diag_matmul`` to the AOTI target device + # the freeze pins ``force_block_diag_matmul`` to the AOTI target device # (tracing always runs on CPU regardless of where the artifact will run). - if self._force_block_diag_matmul is None: + if self.force_block_diag_matmul is None: use_block_diag = not x_flat.is_cpu else: - use_block_diag = self._force_block_diag_matmul + use_block_diag = self.force_block_diag_matmul if use_block_diag: out_flat = self._block_diagonal_matmul(x_flat, weight) else: @@ -575,8 +575,13 @@ def _block_diagonal_matmul( Flattened output with shape ``(F, E, D_m*Cout)``. """ weight = weight.permute(1, 0, 2) # (F, D_m*Cin, D_m*Cout) - if self._block_diag_gemm is not None and active_triton_level(self) >= 1: - return self._block_diag_gemm(x_flat, weight, self._block_diag_slices) + if ( + self.triton_l_1_block_diag_gemm is not None + and active_triton_level(self) >= 1 + ): + return self.triton_l_1_block_diag_gemm( + x_flat, weight, self._block_diag_slices + ) blocks = [ torch.bmm( x_flat[:, :, in0:in1], @@ -726,7 +731,7 @@ def __init__( # of the ``degree_channel`` low-rank branch in the ``mmax == 1`` layout. self.triton_infer_level = triton_infer_level() self.triton_train_level = triton_train_level() - self._radial_mix_block = None + self.triton_l_1_radial_mix = None if ( max(self.triton_infer_level, self.triton_train_level) >= 1 and self.mode == "degree_channel" @@ -737,7 +742,7 @@ def __init__( radial_mix_block, ) - self._radial_mix_block = radial_mix_block + self.triton_l_1_radial_mix = radial_mix_block def _build_dense_scatter_indices(self) -> tuple[torch.Tensor, torch.Tensor]: compact_indices: list[int] = [] @@ -831,8 +836,11 @@ def forward(self, x_local: torch.Tensor, radial_feat: torch.Tensor) -> torch.Ten compact = kernel_flat.view( x_local.shape[0], self.degree_kernel_size, self.rank ) - if self._radial_mix_block is not None and active_triton_level(self) >= 1: - return self._radial_mix_block( + if ( + self.triton_l_1_radial_mix is not None + and active_triton_level(self) >= 1 + ): + return self.triton_l_1_radial_mix( compact, x_local, self.channel_basis, self.lmax ) kernel = self._scatter_rank_kernel(compact) @@ -1212,31 +1220,27 @@ def __init__( self.device = env.DEVICE self.precision = RESERVED_PRECISION_DICT[dtype] self.compute_dtype = get_promoted_dtype(self.dtype) - # Opt-in fused fast paths, selected by ``DP_TRITON_INFER`` / - # ``DP_TRITON_TRAIN`` (cumulative levels, see :func:`triton_infer_level` - # and :func:`triton_train_level`). Each is read - # once at construction so it becomes a compile-time constant in the - # traced (``make_fx``) graph. Level 1 replaces the dense ``bmm`` - # rotation with universal Triton kernels; level 2 additionally binds the - # table-configured fused value path; inference level 3 routes the mixing - # stack through the fp16x3 tensor-core operator on swept shapes. - # ``DP_CUTILE_INFER`` selects a complete alternative SO(2) path and is - # therefore mutually exclusive with Triton. CuTe dispatch is resolved at - # the enclosing SeZM block, where its exact-shape K1 kernel can coexist - # with these fallback paths. cuTile and the hand-written CUDA paths remain - # inference-only. Fused entries stay ``None`` when their backend is - # unavailable or the block layout is unsupported. + # Acceleration gates are read once at construction so they become + # compile-time constants in the traced (``make_fx``) graph. Triton level + # 1 replaces individual operators, level 2 binds the fused value path, + # and inference level 3 selects its fp16x3 implementation on swept + # shapes. ``DP_CUTILE_INFER`` selects a complete alternative SO(2) path + # and is mutually exclusive with Triton. ``DP_CUDA_INFER`` and + # ``DP_CUDA_TRAIN`` independently bind hand-written CUDA producers that + # take precedence over the corresponding narrower paths. CuTe dispatch + # is resolved at the enclosing block, where its exact-shape SO2 kernel + # can coexist with every fallback. Unsupported entries remain ``None``. self.triton_infer_level = triton_infer_level() self.triton_train_level = triton_train_level() - self.use_triton_infer = self.triton_infer_level >= 1 - self.use_cutile_infer = use_cutile_infer() - if self.use_triton_infer and self.use_cutile_infer: + self.cuda_infer_level = cuda_infer_level() + self.cutile_infer_enabled = use_cutile_infer() + if self.triton_infer_level >= 1 and self.cutile_infer_enabled: raise ValueError( "DP_TRITON_INFER and DP_CUTILE_INFER are mutually exclusive: " "each selects a complete accelerated SO(2) inference path." ) - self._triton_value_path = None - self._cutile_value_path = None + self.triton_infer_l_2_value = None + self.cutile_infer_value = None # === Step 1. Split deterministic seeds at the module top-level === seed_so2_stack = child_seed(seed, 0) @@ -1615,9 +1619,9 @@ def __init__( # attention layout without the optional focus-mix / value / output # projections (the deployed DPA4 configuration). Whichever of the # mutually exclusive inference gates is active then supplies the - # implementation, and ``self._flash_atten_fn`` being bound is what marks + # implementation, and ``self.flash_attention`` being bound is what marks # the fused path as live. - self._flash_atten_layout_ok = ( + self.flash_attention_layout_supported = ( self.n_atten_head > 0 and self.mmax == 1 and self.needs_local_frame @@ -1629,24 +1633,24 @@ def __init__( ) # The cuTile aggregation is inference-only; the Triton one also serves # training, so it is bound whenever either gate asks for level 1. - self._flash_atten_fn = None - self._flash_atten_trains = False - if self._flash_atten_layout_ok and self.use_cutile_infer: + self.flash_attention = None + self.flash_attention_supports_training = False + if self.flash_attention_layout_supported and self.cutile_infer_enabled: from deepmd.pt_expt.kernels.cutile.sezm.flash_atten import ( flash_atten_aggregate, ) - self._flash_atten_fn = flash_atten_aggregate + self.flash_attention = flash_atten_aggregate elif ( - self._flash_atten_layout_ok + self.flash_attention_layout_supported and max(self.triton_infer_level, self.triton_train_level) >= 1 ): from deepmd.pt_expt.kernels.triton.sezm.flash_atten import ( flash_atten_aggregate, ) - self._flash_atten_fn = flash_atten_aggregate - self._flash_atten_trains = self.triton_train_level >= 1 + self.flash_attention = flash_atten_aggregate + self.flash_attention_supports_training = self.triton_train_level >= 1 # === Step 12. Optional fused Triton SO(2) value path (inference) === # Fuses rotate-to-local, the radial degree mixing, the gated mixing @@ -1667,7 +1671,7 @@ def __init__( make_triton_value_path, ) - self._triton_value_path = make_triton_value_path(self) + self.triton_infer_l_2_value = make_triton_value_path(self) # === Step 13. Optional fused rotate-to-local + radial degree mixing === # Level-1 companion of the value path: only the rotation and the @@ -1684,7 +1688,7 @@ def __init__( # hidden widths. Below the bound the separate rotation and radial-mix # kernels win end to end (a 64-wide stack loses ~10%, a 128-wide one # gains), so the binding follows the measured crossover. - self._triton_rotate_mix = None + self.triton_l_1_rotate_mix = None if ( max(self.triton_infer_level, self.triton_train_level) >= 1 and self.hidden_channels >= 128 @@ -1693,7 +1697,7 @@ def __init__( make_triton_rotate_mix, ) - self._triton_rotate_mix = make_triton_rotate_mix(self) + self.triton_l_1_rotate_mix = make_triton_rotate_mix(self) # === Step 14. Optional fused CUDA SO(2) convolution (inference) === # One hand-written CUDA operator spans the complete per-edge path: @@ -1703,13 +1707,13 @@ def __init__( # flash aggregation, and takes precedence over them when the block # matches its supported configuration. The factory returns ``None`` # otherwise, leaving whichever narrower path is bound in charge. - self._cuda_conv_fn = None - if cuda_infer_level() >= 2 and self._flash_atten_layout_ok: + self.cuda_infer_l_2_conv = None + if self.cuda_infer_level >= 2 and self.flash_attention_layout_supported: from deepmd.pt_expt.kernels.cuda.dpa4 import ( make_cuda_so2_conv, ) - self._cuda_conv_fn = make_cuda_so2_conv(self) + self.cuda_infer_l_2_conv = make_cuda_so2_conv(self) # === Step 15. Optional fused CUDA SO(2) value path (training) === # One CUDA operator spans the training value stream up to the attention @@ -1720,13 +1724,13 @@ def __init__( # behind the same differentiable boundary. The attention span stays on # the Triton operator composition inside the traced graph. Bound under # ``DP_CUDA_TRAIN=1``; ``so2_message`` dispatches to it in training mode. - self._cuda_value_train = None + self.cuda_train_value = None if cuda_train_enabled(): from deepmd.pt_expt.kernels.cuda.dpa4.so2_conv_train import ( make_cuda_so2_value, ) - self._cuda_value_train = make_cuda_so2_value(self) + self.cuda_train_value = make_cuda_so2_value(self) # === Step 16. Optional fused destination-segmented attention softmax === # One CSR-segmented operator per direction replaces the @@ -1735,7 +1739,7 @@ def __init__( # hand-derived second order keep the force-loss trace from expanding # the chain into materialized surfaces and serialized scatters. The # source-gated (SFPG) form keeps the reference path. - self._segment_softmax_fn = None + self.triton_l_1_segment_softmax = None if ( max(self.triton_infer_level, self.triton_train_level) >= 1 and self.attn_n_focus * self.n_atten_head <= 16 @@ -1746,18 +1750,18 @@ def __init__( ) if SEGMENT_SOFTMAX_TRITON_AVAILABLE: - self._segment_softmax_fn = segment_softmax + self.triton_l_1_segment_softmax = segment_softmax # === Step 17. Optional fused cuTile SO(2) value-path operators === # Complete cuTile inference path, mutually exclusive with Triton. The # factory validates the block layout and returns ``None`` # otherwise, leaving the dense reference path in charge. - if self.use_cutile_infer: + if self.cutile_infer_enabled: from deepmd.pt_expt.kernels.cutile.sezm.so2_value_path import ( make_cutile_value_path, ) - self._cutile_value_path = make_cutile_value_path(self) + self.cutile_infer_value = make_cutile_value_path(self) def forward( self, @@ -1896,23 +1900,25 @@ def forward_attention( # The fused CUDA operator computes the attention weights itself, so it # does not serve the bridging mode, whose source gate reshapes the # softmax normalization. - run_cuda = ( - self._cuda_conv_fn is not None + run_cuda_infer_l_2 = ( + self.cuda_infer_l_2_conv is not None and not self.training and edge_cache.edge_src_gate is None ) - run_flash = ( - self._flash_atten_fn is not None - and (self._flash_atten_trains or not self.training) - and not run_cuda + run_flash_attention = ( + self.flash_attention is not None + and (self.flash_attention_supports_training or not self.training) + and not run_cuda_infer_l_2 ) - if run_cuda: - return self.forward_attention_cuda(x, edge_cache, radial_feat, x_l0_node) - if run_flash: + if run_cuda_infer_l_2: + return self.forward_attention_cuda_infer_l_2( + x, edge_cache, radial_feat, x_l0_node + ) + if run_flash_attention: return self.forward_attention_flash(x, edge_cache, radial_feat, x_l0_node) return self.forward_attention_dense(x, edge_cache, radial_feat, x_l0_node) - def forward_attention_cuda( + def forward_attention_cuda_infer_l_2( self, x: torch.Tensor, edge_cache: EdgeFeatureCache, @@ -1946,7 +1952,7 @@ def forward_attention_cuda( Node update with shape (N, D, C_wide). """ # === Step 1. Projected radial features === - rad_feat = self._cuda_conv_fn.radial_features( + rad_feat = self.cuda_infer_l_2_conv.radial_features( radial_feat ) # (E, lmax+1, C_wide) @@ -1957,7 +1963,9 @@ def forward_attention_cuda( head_gate = self.attention_head_gate(x_l0_node) # (N, Fa, H) # === Step 4. Fused convolution === - out = self._cuda_conv_fn(x, edge_cache, rad_feat, q_node, k_node, head_gate) + out = self.cuda_infer_l_2_conv( + x, edge_cache, rad_feat, q_node, k_node, head_gate + ) return out.to(dtype=self.dtype) # (N, D, C_wide) def forward_attention_flash( @@ -1974,7 +1982,10 @@ def forward_attention_flash( kernel folds the block-diagonal rotate-back, the inverse-rotation rescale, the per-edge weighting and the destination reduction into a single atomic-free pass, so the transient rotate-back message and - weighted value tensors are never materialized. + weighted value tensors are never materialized. The value producer is + selected independently inside ``so2_message``: training uses + ``cuda_train_value`` when it is bound, while the flash operator remains + the aggregation consumer selected by the Triton training gate. Parameters ---------- @@ -2013,8 +2024,8 @@ def forward_attention_flash( order, row_ptr = cached_edge_csr(edge_cache, "dst", n_node) rotation = edge_cache.Dt_full if rotation is None: - rotation = self._cuda_value_train.edge_runs(edge_cache) - pre_gate = self._flash_atten_fn( + rotation = self.cuda_train_value.edge_runs(edge_cache) + pre_gate = self.flash_attention( x_local, rotation, self.rotate_inv_rescale_full, @@ -2199,7 +2210,7 @@ def attention_weights( edge_src_gate = edge_cache.edge_src_gate n_nodes = x_l0_node.shape[0] if ( - self._segment_softmax_fn is not None + self.triton_l_1_segment_softmax is not None and edge_src_gate is None and attn_logits.is_cuda and active_triton_level(self) >= 1 @@ -2217,7 +2228,7 @@ def attention_weights( + float(self.eps) ).reshape(-1) # (F * H,) n_channel = self.attn_n_focus * self.n_atten_head - alpha = self._segment_softmax_fn( + alpha = self.triton_l_1_segment_softmax( attn_logits.reshape(n_edge, n_channel).to(dtype=torch.float32), edge_cache.edge_env.reshape(n_edge).to(dtype=torch.float32), null_logit, @@ -2419,24 +2430,21 @@ def so2_message( src, dst = edge_cache.src, edge_cache.dst n_edge = src.numel() - if self._cutile_value_path is not None and not self.training: + if self.cutile_infer_value is not None and not self.training: # === Steps 1-5 (fused cuTile operators). ``rotate_mix`` folds the # rotation and the radial degree mixing into one edge-parallel # kernel writing the focus-major layout; ``mixing_stack`` runs the # whole gated stack, keeping the inter-layer activations and the # gated-layer pre-activations off the traced graph entirely. === - x_local, rad_feat = self._cutile_value_path(x, edge_cache, radial_feat) - elif self._cuda_value_train is not None and self.training: + x_local, rad_feat = self.cutile_infer_value(x, edge_cache, radial_feat) + elif self.cuda_train_value is not None and self.training: # === Steps 1-5 (one CUDA kernel, training). The whole value - # stream up to the attention aggregation runs in a single launch; - # only the backward anchors reach device memory. === + # stream up to the attention aggregation runs in a single launch + # with analytic backward and second order; only the backward + # anchors reach device memory. === cached_edge_csr(edge_cache, "src", x.shape[0]) - x_local, rad_feat = self._cuda_value_train(x, edge_cache, radial_feat) - elif ( - self._triton_value_path is not None - and not self.training - and active_triton_level(self) >= 2 - ): + x_local, rad_feat = self.cuda_train_value(x, edge_cache, radial_feat) + elif self.triton_infer_l_2_value is not None and not self.training: # === Steps 1-5 (fused Triton operators, inference). # ``so2_rotate_mix`` folds the rotation and the radial degree # mixing into one edge-parallel kernel writing the focus-major @@ -2444,8 +2452,8 @@ def so2_message( # the competition weight fused into its final store, keeping the # inter-layer activations off the traced graph. === cached_edge_csr(edge_cache, "src", x.shape[0]) - x_local, rad_feat = self._triton_value_path(x, edge_cache, radial_feat) - elif self._triton_rotate_mix is not None and active_triton_level(self) >= 1: + x_local, rad_feat = self.triton_infer_l_2_value(x, edge_cache, radial_feat) + elif self.triton_l_1_rotate_mix is not None and active_triton_level(self) >= 1: # === Steps 1-3 (fused rotate-mix operator). One edge-parallel # kernel gathers the source features, applies the block-diagonal # Wigner rotation and the radial degree mixing, and writes the @@ -2455,7 +2463,7 @@ def so2_message( # step and kept on the edge cache. === with nvtx_range("SO2Conv/rotate_mix"): cached_edge_csr(edge_cache, "src", x.shape[0]) - u0, rad_feat = self._triton_rotate_mix(x, edge_cache, radial_feat) + u0, rad_feat = self.triton_l_1_rotate_mix(x, edge_cache, radial_feat) x_local = u0.view( self.n_focus, n_edge, self.reduced_dim, self.so2_focus_dim ) # (F, E, D_m, Cf) @@ -2474,14 +2482,14 @@ def so2_message( D_full = edge_cache.D_full x_dst_local: torch.Tensor | None = None if active_triton_level(self) >= 1: - # ``self._rotate_to_local_fn`` was bound in ``__init__`` (the + # ``self.triton_l_1_rotate_to_local`` was bound in ``__init__`` (the # block kernel for the m-major ``mmax == 1`` layout, dense # otherwise). - x_local = self._rotate_to_local_fn( + x_local = self.triton_l_1_rotate_to_local( x, src, D_full ) # (E, D_m, C_wide) if self.node_wise_grid_product is not None: - x_dst_local = self._rotate_to_local_fn( + x_dst_local = self.triton_l_1_rotate_to_local( x, dst, D_full ) # (E, D_m, C_wide) else: @@ -2712,7 +2720,9 @@ def _so2_rotate_back( if active_triton_level(self) >= 1 and self.mmax == 1: # The block kernel consumes the (E, F, D_m, Cf) focus layout in # place, folding the inverse transpose into its channel addressing. - x_message = self._rotate_back_fn(x_local, Dt_full) # (E, D, C_wide) + x_message = self.triton_l_1_rotate_back( + x_local, Dt_full + ) # (E, D, C_wide) else: # Restore reduced global layout (E, D_m, C_wide) for inverse rotation. x_local = ( @@ -2721,7 +2731,9 @@ def _so2_rotate_back( .reshape(n_edge, self.reduced_dim, self.hidden_channels) ) if active_triton_level(self) >= 1: - x_message = self._rotate_back_fn(x_local, Dt_full) # (E, D, C_wide) + x_message = self.triton_l_1_rotate_back( + x_local, Dt_full + ) # (E, D, C_wide) else: Dt_from_m = project_Dt_from_m( Dt_full=Dt_full, @@ -2872,8 +2884,8 @@ def _build_so2_mixing( self.reduced_dim = int(coeff_index_m.numel()) # === Step 2. Triton rotation kernels: block for mmax == 1, dense otherwise === - self._rotate_to_local_fn = None - self._rotate_back_fn = None + self.triton_l_1_rotate_to_local = None + self.triton_l_1_rotate_back = None if max(self.triton_infer_level, self.triton_train_level) >= 1: from deepmd.pt_expt.kernels.triton.sezm.so2_rotation import ( rotate_back_block_so2, @@ -2883,20 +2895,22 @@ def _build_so2_mixing( ) if self.mmax == 1: - self._rotate_to_local_fn = lambda x, src, wigner: rotate_to_local_block( - x, src, wigner, self.lmax + self.triton_l_1_rotate_to_local = lambda x, src, wigner: ( + rotate_to_local_block(x, src, wigner, self.lmax) ) # The block kernel reads the (E, F, D_m, Cf) focus layout directly, # so the rotate-back path passes ``x_local`` before the global # reshape and the transpose-back copy is skipped. - self._rotate_back_fn = lambda x_local, wigner: rotate_back_block_so2( - x_local, wigner, self.lmax + self.triton_l_1_rotate_back = lambda x_local, wigner: ( + rotate_back_block_so2(x_local, wigner, self.lmax) ) else: - self._rotate_to_local_fn = lambda x, src, wigner: rotate_to_local_dense( - x, src, wigner, self.coeff_index_m, self.ebed_dim_full + self.triton_l_1_rotate_to_local = lambda x, src, wigner: ( + rotate_to_local_dense( + x, src, wigner, self.coeff_index_m, self.ebed_dim_full + ) ) - self._rotate_back_fn = lambda x_local, wigner: rotate_back_dense( + self.triton_l_1_rotate_back = lambda x_local, wigner: rotate_back_dense( x_local, wigner, self.coeff_index_m, self.ebed_dim_full ) diff --git a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py index b283e3a7e5..f75efe572c 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/wignerd.py +++ b/deepmd/pt/model/descriptor/sezm_nn/wignerd.py @@ -459,9 +459,9 @@ def __init__( ] # The monomial basis routes through whichever accelerated backend # is selected; the two gates are mutually exclusive. - self._use_cutile_monomials = use_cutile_infer() - self._use_triton_monomials = triton_infer_level() >= 1 - self._use_triton_train_monomials = triton_train_level() >= 1 + self.cutile_infer_monomials = use_cutile_infer() + self.triton_infer_l_1_monomials = triton_infer_level() >= 1 + self.triton_train_l_1_monomials = triton_train_level() >= 1 # The l = 2 contraction tensor collapsed onto the 35 unique # degree-4 monomials: column m of the coefficient matrix sums # C_l2[:, :, p] over the 4^4 index tuples p whose component @@ -1150,12 +1150,12 @@ def _monomial_matrix( and ( ( not self.training - and (self._use_triton_monomials or self._use_cutile_monomials) + and (self.triton_infer_l_1_monomials or self.cutile_infer_monomials) ) - or (self.training and self._use_triton_train_monomials) + or (self.training and self.triton_train_l_1_monomials) ) ): - if not self.training and self._use_cutile_monomials: + if not self.training and self.cutile_infer_monomials: from deepmd.pt_expt.kernels.cutile.sezm.wigner_monomials import ( wigner_monomials as monomial_basis, ) @@ -1191,12 +1191,12 @@ def _compute_l2_block(self, edge_quaternion: torch.Tensor) -> torch.Tensor: and ( ( not self.training - and (self._use_triton_monomials or self._use_cutile_monomials) + and (self.triton_infer_l_1_monomials or self.cutile_infer_monomials) ) - or (self.training and self._use_triton_train_monomials) + or (self.training and self.triton_train_l_1_monomials) ) ): - if not self.training and self._use_cutile_monomials: + if not self.training and self.cutile_infer_monomials: from deepmd.pt_expt.kernels.cutile.sezm.wigner_monomials import ( wigner_monomials as monomial_basis, ) diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index 39aa8f13ed..33c40deee2 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -258,7 +258,7 @@ * ``triton.cudagraphs=False`` cudagraphs capture autograd metadata only once. Higher-order gradients need fresh metadata per call, so cudagraphs would feed - stale autograd state into the second backward. Packed K1 also retains + stale autograd state into the second backward. Packed SO2 also retains Python-owned forward state, so direct graph capture remains disabled. * ``max_fusion_size=DP_FUSION_SIZE`` (default 8) Caps kernel fusion complexity so Inductor's scheduler does not @@ -549,10 +549,10 @@ def _neo_cute_nlist_eager_island_enabled(device: torch.device) -> bool: except RuntimeError: return False from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import ( - is_k1_eager_island_enabled, + is_so2_eager_island_enabled, ) - return is_k1_eager_island_enabled(compute_capability) + return is_so2_eager_island_enabled(compute_capability) @torch.compiler.disable @@ -610,13 +610,13 @@ def _build_neo_neighbor_list( ) -def _neo_cute_k1_requires_sorted_edges( +def _neo_cute_so2_requires_sorted_edges( descriptor: Any, *, training: bool, device: torch.device, ) -> bool: - """Return whether this descriptor can consume packed destination-sorted K1. + """Return whether this descriptor can use packed Wigner-D on sorted edges. ``forward_with_edges`` promotes coordinates and edge vectors to the descriptor compute dtype before constructing Wigner data. Mirror that @@ -626,7 +626,7 @@ def _neo_cute_k1_requires_sorted_edges( return bool( not training and _neo_cute_infer_enabled() - and descriptor._packed_wigner_candidate( + and descriptor.is_cute_infer_packed_wigner_candidate( device, compute_dtype, compute_dtype, @@ -921,20 +921,20 @@ def _invalidate_compiled_state_after_load( module._core_compute_pending_compile_t0 = None module._core_compute_pending_compile_key = None module._dens_pending_compile_t0 = None - k1_invalidator = None + so2_invalidator = None for child in module.modules(): if not ( - hasattr(child, "_deepmd_cute_k1_state") + hasattr(child, "_deepmd_cute_so2_state") or hasattr(child, "_deepmd_cute_gate_expand_contract") ): continue - if k1_invalidator is None: - from deepmd.pt_expt.kernels.cute.sezm.k1 import ( - invalidate_cute_k1_state, + if so2_invalidator is None: + from deepmd.pt_expt.kernels.cute.sezm.so2.operation import ( + invalidate_cute_so2_state, ) - k1_invalidator = invalidate_cute_k1_state - k1_invalidator(child) + so2_invalidator = invalidate_cute_so2_state + so2_invalidator(child) # Shared callables can contain make_fx get_attr constants, including # prepared CuTe readout folds. A checkpoint load therefore invalidates # the process-level cache as well as this instance's local slots. @@ -1663,7 +1663,7 @@ def core_compute( descriptor_model = self.atomic_model.descriptor # === Step 1. Establish the force-autograd endpoint === - sort_edges_by_dst = _neo_cute_k1_requires_sorted_edges( + sort_edges_by_dst = _neo_cute_so2_requires_sorted_edges( descriptor_model, training=self.training, device=edge_vec.device, @@ -1928,7 +1928,7 @@ def core_compute_dens( descriptor_model = self.atomic_model.descriptor # === Step 1. Build compact sparse edges === - sort_edges_by_dst = _neo_cute_k1_requires_sorted_edges( + sort_edges_by_dst = _neo_cute_so2_requires_sorted_edges( descriptor_model, training=self.training, device=extended_coord.device, @@ -2153,7 +2153,7 @@ def trace_and_compile( ) return - # Register Python-owned K1 state before make_fx starts. The opt-in thin + # Register Python-owned SO2 state before make_fx starts. The opt-in thin # path can then keep adjacent linears in this graph while the CuTe work # remains opaque behind its existing custom op. from deepmd.pt_expt.kernels.cute.sezm import ( @@ -2166,7 +2166,7 @@ def trace_and_compile( else None ) if not self.training and cute_runtime_policy.is_cute_infer_enabled(): - from deepmd.pt_expt.kernels.cute.sezm.readout_l0 import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid.readout_l0 import ( maybe_prepare_sm80_readout_input_fold, ) @@ -2174,24 +2174,24 @@ def trace_and_compile( self.atomic_model.descriptor.output_ffn, compute_capability, ) - prepared_cute_k1 = False + prepared_cute_so2 = False if ( not self.training and compute_capability is not None and cute_runtime_policy.is_cute_infer_enabled() - and cute_runtime_policy.is_supported_k1_capability(compute_capability) + and cute_runtime_policy.is_supported_so2_capability(compute_capability) ): - from deepmd.pt_expt.kernels.cute.sezm.k1 import ( - prepare_cute_k1_blocks, + from deepmd.pt_expt.kernels.cute.sezm.so2.operation import ( + prepare_cute_so2_blocks, ) - prepared_cute_k1 = prepare_cute_k1_blocks( + prepared_cute_so2 = prepare_cute_so2_blocks( self.atomic_model.descriptor.blocks, training=self.training, device=coord.device, dtype=coord.dtype, ) - if prepared_cute_k1 and cute_runtime_policy.is_k1_thin_wrapper_enabled( + if prepared_cute_so2 and cute_runtime_policy.is_so2_thin_wrapper_enabled( compute_capability ): from ..network.mlp import ( @@ -3116,7 +3116,7 @@ def build_edge_list_from_nlist( edge_schema.edge_scatter_index, ) should_sort = ( - _neo_cute_k1_requires_sorted_edges( + _neo_cute_so2_requires_sorted_edges( self.atomic_model.descriptor, training=self.training, device=extended_coord.device, diff --git a/deepmd/pt/model/network/mlp.py b/deepmd/pt/model/network/mlp.py index 51bc476845..1f6ad02b35 100644 --- a/deepmd/pt/model/network/mlp.py +++ b/deepmd/pt/model/network/mlp.py @@ -48,7 +48,7 @@ def empty_t(shape: tuple[int, ...], precision: torch.dtype) -> torch.Tensor: @torch.compiler.assume_constant_result -def _use_k1_compile_visible_linear( +def _use_so2_compile_visible_linear( input_device: torch.device | None = None, ) -> bool: """Keep the SM80 linear topology stable for one compiled graph.""" @@ -57,7 +57,7 @@ def _use_k1_compile_visible_linear( cute_enabled = os.environ.get("DP_CUTE_INFER", "").strip().lower() if cute_enabled not in truthy: return False - thin_enabled = os.environ.get("DP_CUTE_K1_THIN_WRAPPER", "").strip().lower() + thin_enabled = os.environ.get("DP_CUTE_SO2_THIN_WRAPPER", "").strip().lower() if thin_enabled in falsy: return False if thin_enabled in truthy: @@ -250,7 +250,7 @@ def forward( and (self.bias is None or self.bias.dtype == torch.float32) and not torch.is_autocast_enabled(xx.device.type) and self._deepmd_cute_compile_visible_linear - and _use_k1_compile_visible_linear(xx.device) + and _use_so2_compile_visible_linear(xx.device) ): yy = _matmul_bias(xx, self.matrix, self.bias) else: diff --git a/deepmd/pt_expt/descriptor/dpa4.py b/deepmd/pt_expt/descriptor/dpa4.py index 2205c6b849..6d6a6626e4 100644 --- a/deepmd/pt_expt/descriptor/dpa4.py +++ b/deepmd/pt_expt/descriptor/dpa4.py @@ -16,6 +16,11 @@ C3CutoffEnvelope as C3CutoffEnvelopeDP, ) from deepmd.dpmodel.descriptor.dpa4_nn.radial import InnerClamp as InnerClampDP +from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, + attach_edge_csr, + compact_edges, +) from deepmd.pt_expt.common import ( register_dpmodel_mapping, torch_module, @@ -204,19 +209,19 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # The fused convolution paths consume only the three structural rows of # each Wigner degree block. Source-gated attention bypasses that fused # convolution, so its dense per-edge rotations remain available. - self._wigner_free_conv = ( + self.cuda_infer_l_2_covers_all_blocks = ( self.bridging_switch is None and bool(self.blocks) and all( - getattr(block.so2_conv, "_cuda_conv_fn", None) is not None - and not block.so2_conv._cuda_conv_fn._compete + getattr(block.so2_conv, "cuda_infer_l_2_conv", None) is not None + and not block.so2_conv.cuda_infer_l_2_conv.focus_compete for block in self.blocks ) ) - self._packed_wigner_train = bool(self.blocks) and all( - getattr(block.so2_conv, "_cuda_value_train", None) is not None - and block.so2_conv._flash_atten_fn is not None - and block.so2_conv._flash_atten_trains + self.cuda_train_covers_all_blocks = bool(self.blocks) and all( + getattr(block.so2_conv, "cuda_train_value", None) is not None + and block.so2_conv.flash_attention is not None + and block.so2_conv.flash_attention_supports_training for block in self.blocks ) @@ -224,8 +229,8 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # distance and are cheap enough that the compiler inlines them into # every consumer and re-evaluates them there. Behind an operator # boundary the chain runs once per step. - self._cuda_radial_fn = None - self._cuda_wigner_fn = None + self.cuda_infer_l_1_radial = None + self.cuda_infer_l_1_wigner = None if cuda_infer_level() >= 1: from deepmd.pt_expt.kernels.cuda.dpa4.edge_radial import ( make_cuda_edge_radial, @@ -234,13 +239,13 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: make_cuda_wigner_dense, ) - self._cuda_radial_fn = make_cuda_edge_radial( + self.cuda_infer_l_1_radial = make_cuda_edge_radial( self.edge_envelope, self.radial_basis ) # The dense Wigner pair otherwise costs five full-size passes # over the (E, D, D) tensors; the fused build pays only the # output writes. - self._cuda_wigner_fn = make_cuda_wigner_dense( + self.cuda_infer_l_1_wigner = make_cuda_wigner_dense( self.mp_init_lmax, get_xp_precision(torch, self.compute_precision), ) @@ -298,17 +303,258 @@ def _shared_wigner_runs(self, edge_cache: Any, lmax: int) -> torch.Tensor | None if edge_cache.csr_cache is None: return None if self.training: - if not self._packed_wigner_train: + if not self.cuda_train_covers_all_blocks: return None - fused = self.blocks[0].so2_conv._cuda_value_train + fused = self.blocks[0].so2_conv.cuda_train_value else: - if not self._wigner_free_conv: + if not self.cuda_infer_l_2_covers_all_blocks: return None - fused = self.blocks[0].so2_conv._cuda_conv_fn + fused = self.blocks[0].so2_conv.cuda_infer_l_2_conv if fused is None or lmax > self.lmax: return None return fused.edge_runs(edge_cache)[:, 1 : (lmax + 1) ** 2] + def is_cute_infer_packed_wigner_candidate( + self, + device: torch.device, + dtype: torch.dtype, + geometry_dtype: torch.dtype, + ) -> bool: + """Return whether all blocks satisfy the packed CuTe SO2 contract. + + Parameters + ---------- + device : torch.device + Device that executes the descriptor blocks. + dtype : torch.dtype + Descriptor block dtype. + geometry_dtype : torch.dtype + Edge-geometry compute dtype. + + Returns + ------- + bool + Whether packed Wigner storage can replace dense storage for every + interaction block. + """ + from deepmd.pt_expt.kernels.cute.sezm.so2.operation import ( + is_packed_wigner_candidate, + ) + + return is_packed_wigner_candidate( + blocks=self.blocks, + training=self.training, + device=device, + dtype=dtype, + producer_modules=( + self.radial_basis, + self.radial_embedding, + self.edge_envelope, + *((self.inner_clamp,) if self.inner_clamp is not None else ()), + *((self.bridging_switch,) if self.bridging_switch is not None else ()), + ), + producer_dtypes=(dtype, geometry_dtype), + has_edge_src_gate=self.bridging_switch is not None, + ) + + def prepare_packed_wigner_graph( + self, + graph: NeighborGraph, + n_nodes: int, + ) -> NeighborGraph | None: + """Prepare a compact destination-major graph for CuTe packed SO2. + + Parameters + ---------- + graph : NeighborGraph + Edge graph supplied to the descriptor. + n_nodes : int + Number of nodes addressed by the graph. + + Returns + ------- + NeighborGraph or None + Canonical graph without masked edges, or ``None`` when the exact + CuTe contract is not satisfied. + """ + dtype = get_xp_precision(torch, self.precision) + if not self.is_cute_infer_packed_wigner_candidate( + graph.edge_vec.device, + dtype, + graph.edge_vec.dtype, + ): + return None + # Graph-native inputs retain canonical CSR through compaction. Dense + # adapters and exclusion transforms construct it once at this boundary. + graph = compact_edges(graph) + if not graph.destination_sorted: + graph = attach_edge_csr(graph, n_nodes, canonicalize=True) + from deepmd.pt_expt.kernels.cute.sezm.so2.operation import ( + packed_wigner_edges_eligible, + ) + + if not packed_wigner_edges_eligible( + candidate=True, + edge_count=graph.edge_index.shape[1], + node_count=n_nodes, + destinations_sorted=graph.destination_sorted, + runtime_dtypes=( + dtype, + get_xp_precision(torch, self.compute_precision), + ), + ): + return None + return graph + + def build_packed_wigner( + self, + edge_quat: torch.Tensor, + wigner_calc: Any, + ) -> torch.Tensor | None: + """Build the CuTe packed Wigner panel for an eligible graph. + + Parameters + ---------- + edge_quat : torch.Tensor + Global-to-local edge quaternions with shape ``(E, 4)``. + wigner_calc : Any + Wigner calculator carrying the degree and basis convention. + + Returns + ------- + torch.Tensor or None + CuTe packed Wigner storage with shape ``(E, 46)``, or ``None`` when + the kernel declines the runtime input. + """ + from deepmd.pt_expt.kernels.cute.sezm.wignerd import ( + run_cute_wignerd, + ) + + result = run_cute_wignerd( + edge_quat, + wigner_calc, + packed_wigner=True, + ) + return None if result is None else result[0] + + def prepare_cute_infer_so2_metadata( + self, + edge_cache: Any, + n_nodes: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: + """Adapt graph-owned CSR metadata for the CuTe SO2 implementation. + + Parameters + ---------- + edge_cache : EdgeCache + Per-forward cache carrying the destination-major edge payload. + n_nodes : int + Number of nodes addressed by the edge payload. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor, torch.Tensor] or None + Destination row pointers, source order, and source row pointers, or + ``None`` when the exact CuTe contract is not satisfied. + """ + if ( + self.training + or not edge_cache.destinations_sorted + or edge_cache.D_packed is None + or edge_cache.edge_src_gate is not None + ): + return None + from deepmd.pt_expt.kernels.cute.sezm import ( + runtime_policy, + ) + + if not runtime_policy.is_cute_infer_enabled(): + return None + csr_cache = edge_cache.csr_cache + if csr_cache is None or "dst" not in csr_cache or "src" not in csr_cache: + raise RuntimeError("packed CuTe SO2 requires destination/source graph CSR") + destination_row_ptr = csr_cache["dst"][1] + source_order, source_row_ptr = csr_cache["src"] + if ( + destination_row_ptr.shape != (n_nodes + 1,) + or source_order.shape != edge_cache.src.shape + or source_row_ptr.shape != (n_nodes + 1,) + ): + raise RuntimeError("packed CuTe SO2 received inconsistent graph CSR shapes") + if runtime_policy.is_cute_strict_enabled() and edge_cache.dst.numel() > 1: + torch._assert_async( + torch.all(edge_cache.dst[1:] >= edge_cache.dst[:-1]), + "Neo SO2 destinations_sorted=True requires monotonically " + "nondecreasing destination indices", + ) + return ( + destination_row_ptr.to(dtype=torch.int32).contiguous(), + source_order.to(dtype=torch.int32).contiguous(), + source_row_ptr.to(dtype=torch.int32).contiguous(), + ) + + def build_cute_infer_zonal_coupling(self, edge_cache: Any) -> torch.Tensor: + """Extract the GIE zonal coupling from packed Wigner storage. + + Parameters + ---------- + edge_cache : EdgeCache + Per-forward cache carrying CuTe packed Wigner storage. + + Returns + ------- + torch.Tensor + Zonal coupling with shape ``(E, D_node - 1)``. + + Raises + ------ + ValueError + If the edge cache does not carry packed Wigner storage. + """ + D_packed = edge_cache.D_packed + if D_packed is None: + raise ValueError("CuTe zonal coupling requires packed Wigner storage") + mp_coupling = D_packed.index_select(1, self.gie.packed_zonal_offsets) + if self.gie_zonal_wigner_calc is None: + return mp_coupling + extra_coupling = self.gie_zonal_wigner_calc.forward_zonal( + self._edge_quaternion(edge_cache), + lmin=self.lmax + 1, + ) + return torch.cat([mp_coupling, extra_coupling], dim=1) + + def run_cute_infer_readout( + self, + ffn_in: torch.Tensor, + ) -> torch.Tensor | None: + """Run the CuTe readout when its exact inference contract matches. + + Parameters + ---------- + ffn_in : torch.Tensor + Equivariant readout input with shape ``(N, D, 1, C)``. + + Returns + ------- + torch.Tensor or None + Residual-inclusive scalar output with shape ``(N, C)``, or ``None`` + when the exact CuTe contract is not satisfied. + """ + from deepmd.pt_expt.kernels.cute.sezm import ( + runtime_policy, + ) + + if self.training or not runtime_policy.is_cute_infer_enabled(): + return None + from deepmd.pt_expt.kernels.cute.sezm.output_grid.readout_l0 import ( + maybe_run_neo_readout_l0, + ) + + return maybe_run_neo_readout_l0( + self.output_ffn, + ffn_in, + ) + @classmethod def deserialize(cls, data: dict) -> "DescrptDPA4": # deserialize assigns numpy arrays after __init__, which demotes diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/block.py b/deepmd/pt_expt/descriptor/dpa4_nn/block.py index af14de1b49..095ebfad7f 100644 --- a/deepmd/pt_expt/descriptor/dpa4_nn/block.py +++ b/deepmd/pt_expt/descriptor/dpa4_nn/block.py @@ -150,6 +150,38 @@ def _run_so2_unit( ) return self._run_so2_unit_impl(x, edge_cache, radial_feat) + def _run_so2_unit_impl( + self, + x: torch.Tensor, + edge_cache: EdgeCache, + radial_feat: torch.Tensor, + ) -> torch.Tensor: + """Run the SO(2) unit implementation.""" + if edge_cache.D_packed is not None: + from deepmd.pt_expt.kernels.cute.sezm.so2.operation import ( + maybe_run_cute_so2, + ) + + metadata = edge_cache.cute_infer_so2_metadata + destination_row_ptr, source_order, source_row_ptr = ( + (None, None, None) if metadata is None else metadata + ) + accelerated = maybe_run_cute_so2( + self, + x, + edge_cache, + radial_feat, + dst_ptr=destination_row_ptr, + source_order=source_order, + source_ptr=source_row_ptr, + ) + if accelerated is not None: + return accelerated + raise RuntimeError( + "packed Wigner cache reached an ineligible Neo SO2 dispatch" + ) + return super()._run_so2_unit_impl(x, edge_cache, radial_feat) + def _run_ffn_unit(self, x: torch.Tensor, unit_idx: int) -> torch.Tensor: if self._use_infer_activation_checkpoint(x): return checkpoint( diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/embedding.py b/deepmd/pt_expt/descriptor/dpa4_nn/embedding.py index 3dd60102d2..4fc9d156cf 100644 --- a/deepmd/pt_expt/descriptor/dpa4_nn/embedding.py +++ b/deepmd/pt_expt/descriptor/dpa4_nn/embedding.py @@ -20,6 +20,9 @@ from deepmd.pt_expt.common import ( torch_module, ) +from deepmd.pt_expt.kernels.cute.sezm.so2.wigner_layout import ( + ZONAL_PANEL_OFFSETS, +) from deepmd.pt_expt.kernels.utils import ( cuda_infer_level, ) @@ -35,14 +38,23 @@ class GeometricInitialEmbedding(GeometricInitialEmbeddingDP): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) + self.register_buffer( + "packed_zonal_offsets", + torch.tensor( + ZONAL_PANEL_OFFSETS if self.lmax == 3 else (), + device=self.non_scalar_row_index.device, + dtype=torch.long, + ), + persistent=False, + ) # === Fused message-and-scatter operator === # The reference composition materializes the per-edge message, an # (E, D-1, C) tensor that dominates the cost of this module. The fused # operator keeps it in registers and reduces through the destination CSR. - self._cuda_scatter = False + self.cuda_infer_l_1_scatter = False # ``None`` keeps the runtime ``zonal_coupling.is_cuda`` dispatch; the # freeze pins it to the AOTI target because tracing always runs on CPU. - self._force_fused_scatter: bool | None = None + self.force_cuda_infer_l_1_scatter: bool | None = None if ( cuda_infer_level() >= 1 and get_xp_precision(torch, self.precision) is torch.float32 @@ -52,20 +64,75 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: supported, ) - self._cuda_scatter = op_available() and supported( + self.cuda_infer_l_1_scatter = op_available() and supported( self.lmax, self.ebed_dim - 1, self.channels ) - def _can_fuse_scatter(self, zonal_coupling: torch.Tensor) -> bool: - """Return whether the fused scatter serves the runtime or trace target.""" + def run_cute_infer_gie( + self, + *, + n_nodes: int | torch.SymInt, + edge_cache: Any, + radial_feat: torch.Tensor, + zonal_coupling: torch.Tensor, + spin_l1_message: torch.Tensor | None = None, + ) -> torch.Tensor | None: + """Run the CuTe GIE path when its exact inference contract matches. + + Parameters + ---------- + n_nodes : int or torch.SymInt + Number of nodes addressed by the edge cache. + edge_cache : EdgeCache + Per-forward geometric edge cache. + radial_feat : torch.Tensor + Radial features with shape ``(E, lmax, C)``. + zonal_coupling : torch.Tensor + Zonal coupling with shape ``(E, D - 1)``. + spin_l1_message : torch.Tensor, optional + Native-spin message with shape ``(E, 3, C)``. + + Returns + ------- + torch.Tensor or None + Initial embedding with shape ``(N, D, C)``, or ``None`` when the + exact CuTe contract is not satisfied. + """ + if self.training or spin_l1_message is not None: + return None + from deepmd.pt_expt.kernels.cute.sezm.gie import ( + maybe_run_cute_gie, + ) + + return maybe_run_cute_gie( + self, + n_nodes=n_nodes, + edge_cache=edge_cache, + radial_feat=radial_feat, + zonal_coupling=zonal_coupling, + ) + + def can_run_cuda_infer_l_1_scatter(self, zonal_coupling: torch.Tensor) -> bool: + """Return whether the fused scatter serves the runtime or trace target. + + Parameters + ---------- + zonal_coupling : torch.Tensor + Zonal coupling with shape ``(E, D - 1)``. + + Returns + ------- + bool + Whether the CUDA inference level-one scatter is eligible. + """ target_is_cuda = ( zonal_coupling.is_cuda - if self._force_fused_scatter is None - else self._force_fused_scatter + if self.force_cuda_infer_l_1_scatter is None + else self.force_cuda_infer_l_1_scatter ) - return self._cuda_scatter and not self.training and target_is_cuda + return self.cuda_infer_l_1_scatter and not self.training and target_is_cuda - def forward_fused_scatter( + def run_cuda_infer_l_1_scatter( self, n_nodes: int | torch.SymInt, edge_cache: Any, diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/grid_net.py b/deepmd/pt_expt/descriptor/dpa4_nn/grid_net.py index 909ea7ca82..b88da2b9bc 100644 --- a/deepmd/pt_expt/descriptor/dpa4_nn/grid_net.py +++ b/deepmd/pt_expt/descriptor/dpa4_nn/grid_net.py @@ -7,6 +7,9 @@ import torch +from deepmd.dpmodel.common import ( + get_xp_precision, +) from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import S2GridNet as S2GridNetDP from deepmd.dpmodel.descriptor.dpa4_nn.grid_net import SO3GridNet as SO3GridNetDP from deepmd.pt_expt.common import ( @@ -18,7 +21,7 @@ ) -def _bind_grid_pair(module: Any) -> None: +def bind_grid_pair_operators(module: Any) -> None: """Bind the fused coefficient-grid pair operators that serve the layout. The inference operator (CUDA, register-resident walk) and the training @@ -28,7 +31,13 @@ def _bind_grid_pair(module: Any) -> None: below 75 slots the dense section is small and the operator's dispatch chain costs more than its kernels save on the host-bound configurations, so the narrow grids stay with the compiler. + + Parameters + ---------- + module : Any + PT-expt grid module receiving the eligible operator bindings. """ + module.dtype = get_xp_precision(torch, module.precision) if module.projector.to_grid_mat.dtype is not torch.float32: return slots = int(module.projector.to_grid_mat.shape[1]) @@ -40,7 +49,7 @@ def _bind_grid_pair(module: Any) -> None: ) if op_available() and slots in SUPPORTED_SLOTS: - module._grid_pair_fn = grid_pair + module.cuda_infer_l_1_grid_pair = grid_pair if triton_train_level() >= 1 and slots >= 75: from deepmd.pt_expt.kernels.triton.sezm.grid_pair import ( GRID_PAIR_TRITON_AVAILABLE, @@ -48,7 +57,52 @@ def _bind_grid_pair(module: Any) -> None: ) if GRID_PAIR_TRITON_AVAILABLE: - module._grid_pair_train_fn = grid_pair_train + module.triton_train_l_1_grid_pair = grid_pair_train + + +def run_cute_infer_grid_pair( + module: Any, + left: torch.Tensor, + right: torch.Tensor, +) -> torch.Tensor | None: + """Run the CuTe grid product when its exact inference contract matches. + + Parameters + ---------- + module : Any + Grid module carrying the projector and frozen inference state. + left : torch.Tensor + Left coefficient operand with shape ``(N, D, F, n_frames * C)``. + right : torch.Tensor + Right coefficient operand with the same shape as ``left``. + + Returns + ------- + torch.Tensor or None + Coefficient result with the same shape as ``left``, or ``None`` when + the exact CuTe contract is not satisfied. + """ + from deepmd.pt_expt.kernels.cute.sezm import ( + runtime_policy, + ) + + if ( + not runtime_policy.is_cute_infer_enabled() + or module.training + or any(parameter.requires_grad for parameter in module.parameters()) + ): + return None + from deepmd.pt_expt.kernels.cute.sezm.output_grid.product import ( + maybe_run_cute_output_grid_product, + ) + + return maybe_run_cute_output_grid_product( + left, + right, + module.projector.to_grid_mat, + module.projector.from_grid_mat, + n_frames=module.n_frames, + ) @torch_module @@ -57,7 +111,29 @@ class S2GridNet(S2GridNetDP): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - _bind_grid_pair(self) + bind_grid_pair_operators(self) + + def run_cute_infer_grid_pair( + self, + left: torch.Tensor, + right: torch.Tensor, + ) -> torch.Tensor | None: + """Run the CuTe grid product when its exact inference contract matches. + + Parameters + ---------- + left : torch.Tensor + Left coefficient operand with shape ``(N, D, F, n_frames * C)``. + right : torch.Tensor + Right coefficient operand with the same shape as ``left``. + + Returns + ------- + torch.Tensor or None + Coefficient result with the same shape as ``left``, or ``None`` + when the exact CuTe contract is not satisfied. + """ + return run_cute_infer_grid_pair(self, left, right) @torch_module @@ -66,4 +142,26 @@ class SO3GridNet(SO3GridNetDP): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - _bind_grid_pair(self) + bind_grid_pair_operators(self) + + def run_cute_infer_grid_pair( + self, + left: torch.Tensor, + right: torch.Tensor, + ) -> torch.Tensor | None: + """Run the CuTe grid product when its exact inference contract matches. + + Parameters + ---------- + left : torch.Tensor + Left coefficient operand with shape ``(N, D, F, n_frames * C)``. + right : torch.Tensor + Right coefficient operand with the same shape as ``left``. + + Returns + ------- + torch.Tensor or None + Coefficient result with the same shape as ``left``, or ``None`` + when the exact CuTe contract is not satisfied. + """ + return run_cute_infer_grid_pair(self, left, right) diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/so2.py b/deepmd/pt_expt/descriptor/dpa4_nn/so2.py index 595169e1e3..30cbf88715 100644 --- a/deepmd/pt_expt/descriptor/dpa4_nn/so2.py +++ b/deepmd/pt_expt/descriptor/dpa4_nn/so2.py @@ -1,19 +1,19 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""pt_expt SO(2) runtime bindings for accelerated inference kernels. +"""PT-expt SO(2) runtime bindings for accelerated kernels. -The dpmodel SO(2) modules are array-API only. These wrappers inject the -reference PT inference paths around three hot paths, mirroring +The array-API dpmodel modules define the shared execution seams. These wrappers +bind the same accelerated implementations and dispatch names used by ``deepmd.pt.model.descriptor.sezm_nn.so2``: - the block-diagonal GEMM of :class:`SO2Linear`, -- the two rotation hot paths of :class:`SO2Convolution`, and +- the rotation, value, and attention paths of :class:`SO2Convolution`, and - the low-rank branch of :class:`DynamicRadialDegreeMixer`. -Triton and cuTile are mutually exclusive complete SO(2) paths. The hand-written -CUDA operators form an independent cumulative layer and take precedence where -their factories bind. Every gate is resolved at construction so export records a -static dispatch choice; training and unsupported layouts retain the dpmodel -reference path. +Triton and cuTile are mutually exclusive complete SO(2) paths. Hand-written +CUDA operators form an independent cumulative layer, while CuTe replaces exact +Neo spans at the enclosing block. Every gate is resolved at construction so +tracing records a static dispatch choice; unsupported layouts retain the +dpmodel reference path. """ from __future__ import ( @@ -27,6 +27,9 @@ import torch +from deepmd.dpmodel.common import ( + get_xp_precision, +) from deepmd.dpmodel.descriptor.dpa4_nn.so2 import ( DynamicRadialDegreeMixer as DynamicRadialDegreeMixerDP, ) @@ -48,13 +51,16 @@ ) -def _active_triton_level(module: Any) -> int: - """Return the Triton dispatch level governing the module's current mode. +def active_triton_level(module: Any) -> int: + """ + Return the Triton acceleration level that applies to a module's mode. - The levels are read at construction to decide which kernels to bind, but - consulted here at call time: inference and training are separate gates, - and a module built with both bound must follow whichever one matches the - mode it is being called in. + Inference and training are gated independently: an operator qualifies for + inference as soon as it reproduces the forward and the coordinate gradient + with the parameters held fixed, whereas training additionally requires the + gradient of every parameter it consumes and a second derivative of its own + backward, which the force loss traverses. Both levels are captured at + construction, so the branch this drives resolves at trace time. Parameters ---------- @@ -64,7 +70,7 @@ def _active_triton_level(module: Any) -> int: Returns ------- int - The training level in training mode, the inference level otherwise. + The level in effect for the module's current mode. """ return module.triton_train_level if module.training else module.triton_infer_level @@ -84,7 +90,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # Export override for the block-diagonal vs dense matmul branch below. # ``None`` keeps the runtime ``x_flat.is_cuda`` dispatch; the freeze sets # it so the AOTI graph follows the *target* device, not the CPU trace. - self._force_block_diag_matmul: bool | None = None + self.force_block_diag_matmul: bool | None = None # Fast path (``DP_TRITON_INFER >= 1`` or ``DP_TRITON_TRAIN >= 1``): # the per-|m|-block batched bmm + cat of ``_block_diagonal_matmul`` is @@ -97,7 +103,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # constants in the traced (``make_fx``) graph. self.triton_infer_level = triton_infer_level() self.triton_train_level = triton_train_level() - self._block_diag_gemm = None + self.triton_l_1_block_diag_gemm = None if max(self.triton_infer_level, self.triton_train_level) >= 1: from deepmd.pt_expt.kernels.triton.sezm.so2_block_gemm import ( SO2_BLOCK_GEMM_TRITON_AVAILABLE, @@ -108,7 +114,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: if SO2_BLOCK_GEMM_TRITON_AVAILABLE and slices_supported( self._block_diag_slices ): - self._block_diag_gemm = block_diag_gemm + self.triton_l_1_block_diag_gemm = block_diag_gemm def _block_diagonal_matmul( self, x_flat: torch.Tensor, weight: torch.Tensor @@ -117,22 +123,27 @@ def _block_diagonal_matmul( # trips an Inductor AVX2 C++ codegen bug, so only CPU needs it. Every other # device uses the block-diagonal contraction, which skips the structural # off-|m| zeros. ``make_fx`` resolves this Python branch at trace time, so - # the freeze pins ``_force_block_diag_matmul`` to the AOTI target device + # the freeze pins ``force_block_diag_matmul`` to the AOTI target device # (tracing always runs on CPU regardless of where the artifact will run). - if self._force_block_diag_matmul is None: + if self.force_block_diag_matmul is None: use_block_diag = not x_flat.is_cpu else: - use_block_diag = self._force_block_diag_matmul + use_block_diag = self.force_block_diag_matmul if not use_block_diag: return torch.einsum("fei,ifo->feo", x_flat, weight) - if self._block_diag_gemm is not None and _active_triton_level(self) >= 1: + if ( + self.triton_l_1_block_diag_gemm is not None + and active_triton_level(self) >= 1 + ): # The fused GEMM consumes the ``(F, D_m*Cin, D_m*Cout)`` presentation # directly from the strided weight, so the permute is applied here and # the contiguity copy the dpmodel ``bmm`` cat path would need is # skipped. The eager fallback permutes ``weight`` internally, so it is # passed the stored ``(D_m*Cin, F, D_m*Cout)`` layout untouched. weight = weight.permute(1, 0, 2) # (F, D_m*Cin, D_m*Cout) - return self._block_diag_gemm(x_flat, weight, self._block_diag_slices) + return self.triton_l_1_block_diag_gemm( + x_flat, weight, self._block_diag_slices + ) return super()._block_diagonal_matmul(x_flat, weight) @@ -151,7 +162,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # constants in the traced (``make_fx``) graph. self.triton_infer_level = triton_infer_level() self.triton_train_level = triton_train_level() - self._radial_mix_block = None + self.triton_l_1_radial_mix = None if ( max(self.triton_infer_level, self.triton_train_level) >= 1 and self.mode == "degree_channel" @@ -162,13 +173,13 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: radial_mix_block, ) - self._radial_mix_block = radial_mix_block + self.triton_l_1_radial_mix = radial_mix_block def _mix_rank_compact( self, compact: torch.Tensor, x_local: torch.Tensor ) -> torch.Tensor: - if self._radial_mix_block is not None and _active_triton_level(self) >= 1: - return self._radial_mix_block( + if self.triton_l_1_radial_mix is not None and active_triton_level(self) >= 1: + return self.triton_l_1_radial_mix( compact, x_local, self.channel_basis, self.lmax ) return super()._mix_rank_compact(compact, x_local) @@ -176,34 +187,41 @@ def _mix_rank_compact( @torch_module class SO2Convolution(SO2ConvolutionDP): - """SO(2) convolution with opt-in accelerated inference kernels.""" + """SO(2) convolution with opt-in inference and training kernels.""" def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - # The inference gates are read once at construction so they become - # compile-time constants in the traced (``make_fx``) graph. Triton and - # cuTile claim the same SO(2) value path and are mutually exclusive; the - # hand-written CUDA operators form an independent, cumulative layer and - # take precedence where their factories bind. + self.dtype = get_xp_precision(torch, self.precision) + self.compute_dtype = get_xp_precision(torch, self.compute_precision) + # Acceleration gates are read once at construction so they become + # compile-time constants in the traced (``make_fx``) graph. Triton level + # 1 replaces individual operators, level 2 binds the fused value path, + # and inference level 3 selects its fp16x3 implementation on swept + # shapes. ``DP_CUTILE_INFER`` selects a complete alternative SO(2) path + # and is mutually exclusive with Triton. ``DP_CUDA_INFER`` and + # ``DP_CUDA_TRAIN`` independently bind hand-written CUDA producers that + # take precedence over the corresponding narrower paths. CuTe dispatch + # is resolved at the enclosing block, where its exact-shape SO2 kernel + # can coexist with every fallback. Unsupported entries remain ``None``. self.triton_infer_level = triton_infer_level() self.triton_train_level = triton_train_level() - self.use_triton_infer = self.triton_infer_level >= 1 - self.use_cutile_infer = use_cutile_infer() - if self.use_triton_infer and self.use_cutile_infer: + self.cuda_infer_level = cuda_infer_level() + self.cutile_infer_enabled = use_cutile_infer() + if self.triton_infer_level >= 1 and self.cutile_infer_enabled: raise ValueError( "DP_TRITON_INFER and DP_CUTILE_INFER are mutually exclusive: " "each selects a complete accelerated SO(2) inference path." ) - self._triton_value_path = None - self._cutile_value_path = None - self._cached_edge_csr_fn = cached_edge_csr + self.triton_infer_l_2_value = None + self.cutile_infer_value = None + self.edge_csr = cached_edge_csr # === Triton rotation kernels: block for mmax == 1, dense otherwise === # The rotation operators carry differentiable backwards (the force # loss traverses them twice), so the training gate binds them as well; - # ``_active_triton_level`` then selects the path per mode. - self._rotate_to_local_fn = None - self._rotate_back_fn = None + # ``active_triton_level`` then selects the path per mode. + self.triton_l_1_rotate_to_local = None + self.triton_l_1_rotate_back = None if max(self.triton_infer_level, self.triton_train_level) >= 1: from deepmd.pt_expt.kernels.triton.sezm.so2_rotation import ( rotate_back_block_so2, @@ -213,24 +231,26 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: ) if self.mmax == 1: - self._rotate_to_local_fn = lambda x, src, wigner: rotate_to_local_block( - x, src, wigner, self.lmax + self.triton_l_1_rotate_to_local = lambda x, src, wigner: ( + rotate_to_local_block(x, src, wigner, self.lmax) ) # The block kernel reads the (E, F, D_m, Cf) focus layout directly, # so the rotate-back path passes ``x_local`` before the global # reshape and the transpose-back copy is skipped. - self._rotate_back_fn = lambda x_local, wigner: rotate_back_block_so2( - x_local, wigner, self.lmax + self.triton_l_1_rotate_back = lambda x_local, wigner: ( + rotate_back_block_so2(x_local, wigner, self.lmax) ) else: - self._rotate_to_local_fn = lambda x, src, wigner: rotate_to_local_dense( - x, src, wigner, self.coeff_index_m, self.ebed_dim_full + self.triton_l_1_rotate_to_local = lambda x, src, wigner: ( + rotate_to_local_dense( + x, src, wigner, self.coeff_index_m, self.ebed_dim_full + ) ) - self._rotate_back_fn = lambda x_local, wigner: rotate_back_dense( + self.triton_l_1_rotate_back = lambda x_local, wigner: rotate_back_dense( x_local, wigner, self.coeff_index_m, self.ebed_dim_full ) - # === Step 12. Optional fused flash-attention aggregation kernel === + # === Step 11. Optional fused flash-attention aggregation kernel === # Folds the entire ``n_atten_head > 0`` value aggregation -- block-diagonal # rotate-back, inverse-rotation rescale, envelope-gated softmax weighting, # and the destination scatter -- into a single destination-segmented @@ -244,29 +264,30 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # attention layout without the optional focus-mix / value / output # projections (the deployed DPA4 configuration). Whichever of the # mutually exclusive inference gates is active then supplies the - # implementation, and ``self._flash_atten_fn`` being bound is what marks + # implementation, and ``self.flash_attention`` being bound is what marks # the fused path as live. # The cuTile aggregation is inference-only; the Triton one also serves - # training (analytic backward and second order), so it is bound - # whenever either gate asks for level 1 and ``_flash_atten_trains`` - # marks it as training-capable for the dpmodel dispatch. - if self._flash_atten_layout_ok and self.use_cutile_infer: + # training, so it is bound whenever either gate asks for level 1. + self.flash_attention = None + self.flash_attention_supports_training = False + if self.flash_attention_layout_supported and self.cutile_infer_enabled: from deepmd.pt_expt.kernels.cutile.sezm.flash_atten import ( flash_atten_aggregate, ) - self._flash_atten_fn = flash_atten_aggregate - elif self._flash_atten_layout_ok and ( - self.use_triton_infer or self.triton_train_level >= 1 + self.flash_attention = flash_atten_aggregate + elif ( + self.flash_attention_layout_supported + and max(self.triton_infer_level, self.triton_train_level) >= 1 ): from deepmd.pt_expt.kernels.triton.sezm.flash_atten import ( flash_atten_aggregate, ) - self._flash_atten_fn = flash_atten_aggregate - self._flash_atten_trains = self.triton_train_level >= 1 + self.flash_attention = flash_atten_aggregate + self.flash_attention_supports_training = self.triton_train_level >= 1 - # === Step 13. Optional fused Triton SO(2) value-path operators === + # === Step 12. Optional fused Triton SO(2) value path (inference) === # Fuses rotate-to-local, the radial degree mixing, the gated mixing # stack, and the focus competition of ``so2_message`` into the # ``sezm_triton::so2_rotate_mix`` / ``so2_mixing_stack`` operators. @@ -277,56 +298,32 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # it engages at ``DP_TRITON_INFER >= 2``; at level 3 the factory # additionally routes the mixing stack through the fp16x3 tensor-core # operator on shapes whose configuration passed the fp64 validation - # sweep. + # sweep. Training never composes this path: the training value stream + # is either the level-1 operator composition (step 13 and the dense + # mixing stack) or the fused CUDA value path (step 15). if self.triton_infer_level >= 2: from deepmd.pt_expt.kernels.triton.sezm.so2_value_path import ( make_triton_value_path, ) - self._triton_value_path = make_triton_value_path(self) - - # === Step 13b. Optional fused CUDA SO(2) convolution === - # One hand-written CUDA operator spans the complete per-edge path: - # rotate-to-local, the radial degree mixing, the gated mixing stack, the - # inverse rotation, the attention weighting and the destination - # reduction. It therefore supersedes both the fused value path and the - # flash aggregation, and takes precedence over them when the block - # matches its supported configuration. The factory returns ``None`` - # otherwise, leaving whichever narrower path is bound in charge. - self._cuda_conv_fn = None - if cuda_infer_level() >= 2 and self._flash_atten_layout_ok: - from deepmd.pt_expt.kernels.cuda.dpa4 import ( - make_cuda_so2_conv, - ) - - self._cuda_conv_fn = make_cuda_so2_conv(self) - - # === Step 14. Optional fused cuTile SO(2) value-path operators === - # Complete cuTile inference path, mutually exclusive with Triton. The - # factory validates the block layout and returns ``None`` - # otherwise, leaving the dense reference path in charge. - if self.use_cutile_infer: - from deepmd.pt_expt.kernels.cutile.sezm.so2_value_path import ( - make_cutile_value_path, - ) - - self._cutile_value_path = make_cutile_value_path(self) + self.triton_infer_l_2_value = make_triton_value_path(self) - # === Step 16. Optional fused rotate-mix operator === - # One edge-parallel kernel gathers the source features, applies the - # block-diagonal Wigner rotation and the radial degree mixing, and - # writes the focus-major mixing input directly; the degree-expanded - # local intermediate and its relayout never reach the traced graph. - # The operator carries a differentiable backward and a hand-derived - # second order, so it serves force-loss training. + # === Step 13. Optional fused rotate-to-local + radial degree mixing === + # Level-1 companion of the value path: only the rotation and the + # degree mixing fuse into one edge-parallel operator writing the + # focus-major mixing input directly, while the mixing stack itself + # stays with the compiler. This removes the degree-expanded local + # intermediate and its relayout from the traced graph; the operator's + # backward reduces through the source CSR view and carries a + # hand-derived second order, so it serves force-loss training. # # The operator is quadrilinear, so a force loss re-enters its forward # and backward several times for the second order. That fixed cost is # repaid only where the materialization it removes is large: the wide # hidden widths. Below the bound the separate rotation and radial-mix - # kernels win end to end, so the binding follows the measured - # crossover (see :func:`_rotate_mix_supported`). - self._triton_rotate_mix = None + # kernels win end to end (a 64-wide stack loses ~10%, a 128-wide one + # gains), so the binding follows the measured crossover. + self.triton_l_1_rotate_mix = None if ( max(self.triton_infer_level, self.triton_train_level) >= 1 and self.hidden_channels >= 128 @@ -335,9 +332,25 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: make_triton_rotate_mix, ) - self._triton_rotate_mix = make_triton_rotate_mix(self) + self.triton_l_1_rotate_mix = make_triton_rotate_mix(self) - # === Step 17. Optional fused CUDA SO(2) value path (training) === + # === Step 14. Optional fused CUDA SO(2) convolution (inference) === + # One hand-written CUDA operator spans the complete per-edge path: + # rotate-to-local, the radial degree mixing, the gated mixing stack, the + # inverse rotation, the attention weighting and the destination + # reduction. It therefore supersedes both the fused value path and the + # flash aggregation, and takes precedence over them when the block + # matches its supported configuration. The factory returns ``None`` + # otherwise, leaving whichever narrower path is bound in charge. + self.cuda_infer_l_2_conv = None + if self.cuda_infer_level >= 2 and self.flash_attention_layout_supported: + from deepmd.pt_expt.kernels.cuda.dpa4 import ( + make_cuda_so2_conv, + ) + + self.cuda_infer_l_2_conv = make_cuda_so2_conv(self) + + # === Step 15. Optional fused CUDA SO(2) value path (training) === # One CUDA operator spans the training value stream up to the attention # aggregation: rotate-to-local, radial degree mixing, the cross-focus # competition weight, the whole gated mixing stack and the final @@ -346,21 +359,22 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: # behind the same differentiable boundary. The attention span stays on # the Triton operator composition inside the traced graph. Bound under # ``DP_CUDA_TRAIN=1``; ``so2_message`` dispatches to it in training mode. + self.cuda_train_value = None if cuda_train_enabled(): from deepmd.pt_expt.kernels.cuda.dpa4.so2_conv_train import ( make_cuda_so2_value, ) - self._cuda_value_train = make_cuda_so2_value(self) + self.cuda_train_value = make_cuda_so2_value(self) - # === Step 18. Optional fused destination-segmented attention softmax === + # === Step 16. Optional fused destination-segmented attention softmax === # One CSR-segmented operator per direction replaces the # scatter/gather softmax chain of the attention weights, sharing the # destination-sorted view with the flash aggregation; its backward and # hand-derived second order keep the force-loss trace from expanding # the chain into materialized surfaces and serialized scatters. The # source-gated (SFPG) form keeps the reference path. - self._segment_softmax_fn = None + self.triton_l_1_segment_softmax = None if ( max(self.triton_infer_level, self.triton_train_level) >= 1 and self.attn_n_focus * self.n_atten_head <= 16 @@ -371,7 +385,18 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: ) if SEGMENT_SOFTMAX_TRITON_AVAILABLE: - self._segment_softmax_fn = segment_softmax + self.triton_l_1_segment_softmax = segment_softmax + + # === Step 17. Optional fused cuTile SO(2) value-path operators === + # Complete cuTile inference path, mutually exclusive with Triton. The + # factory validates the block layout and returns ``None`` + # otherwise, leaving the dense reference path in charge. + if self.cutile_infer_enabled: + from deepmd.pt_expt.kernels.cutile.sezm.so2_value_path import ( + make_cutile_value_path, + ) + + self.cutile_infer_value = make_cutile_value_path(self) def _rotate_mix( self, @@ -379,11 +404,11 @@ def _rotate_mix( edge_cache: EdgeCache, radial_feat: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - if self._triton_rotate_mix is not None and _active_triton_level(self) >= 1: + if self.triton_l_1_rotate_mix is not None and active_triton_level(self) >= 1: # The operator's backward reduces through the source CSR view, # built once per step and kept on the edge cache. cached_edge_csr(edge_cache, "src", x.shape[0]) - u0, rad_feat = self._triton_rotate_mix(x, edge_cache, radial_feat) + u0, rad_feat = self.triton_l_1_rotate_mix(x, edge_cache, radial_feat) x_local = u0.view( self.n_focus, edge_cache.src.shape[0], @@ -399,9 +424,9 @@ def _attention_softmax( edge_cache: EdgeCache, n_nodes: int, ) -> torch.Tensor: - active_level = _active_triton_level(self) + active_level = active_triton_level(self) if ( - self._segment_softmax_fn is not None + self.triton_l_1_segment_softmax is not None and edge_cache.edge_src_gate is None and attn_logits.is_cuda and active_level >= 1 @@ -420,7 +445,7 @@ def _attention_softmax( + float(self.eps) ).reshape(-1) # (F * H,) n_channel = self.attn_n_focus * self.n_atten_head - alpha = self._segment_softmax_fn( + alpha = self.triton_l_1_segment_softmax( attn_logits.reshape(n_edge, n_channel).to(dtype=torch.float32), edge_cache.edge_env.reshape(n_edge).to(dtype=torch.float32), null_logit, @@ -433,38 +458,47 @@ def _attention_softmax( ) return super()._attention_softmax(attn_logits, edge_cache, n_nodes) - def _rotation_active(self) -> bool: - """Whether the bound rotation kernels serve the current mode.""" - return self._rotate_to_local_fn is not None and _active_triton_level(self) >= 1 + def rotation_kernel_active(self) -> bool: + """Return whether the bound rotation kernels serve the current mode. + + Returns + ------- + bool + Whether level-one Triton rotation is active. + """ + return ( + self.triton_l_1_rotate_to_local is not None + and active_triton_level(self) >= 1 + ) def _rotate_to_local( self, x: torch.Tensor, edge_cache: EdgeCache ) -> tuple[torch.Tensor, torch.Tensor | None]: - if self._rotation_active(): - # ``self._rotate_to_local_fn`` was bound in ``__init__`` (the block - # kernel for the m-major ``mmax == 1`` layout, dense otherwise). + if self.rotation_kernel_active(): + # The block kernel serves the m-major ``mmax == 1`` layout; the + # dense kernel serves every other supported layout. D_full = edge_cache.D_full - x_local = self._rotate_to_local_fn(x, edge_cache.src, D_full) + x_local = self.triton_l_1_rotate_to_local(x, edge_cache.src, D_full) x_dst_local: torch.Tensor | None = None if self.node_wise_grid_product is not None: - x_dst_local = self._rotate_to_local_fn(x, edge_cache.dst, D_full) + x_dst_local = self.triton_l_1_rotate_to_local(x, edge_cache.dst, D_full) return x_local, x_dst_local return super()._rotate_to_local(x, edge_cache) def _rotate_back( self, x_local: torch.Tensor, edge_cache: EdgeCache, n_edge: int ) -> torch.Tensor: - if self._rotation_active(): + if self.rotation_kernel_active(): Dt_full = edge_cache.Dt_full if self.mmax == 1: # The block kernel consumes the (E, F, D_m, Cf) focus layout in # place, folding the inverse transpose into its channel addressing. - return self._rotate_back_fn(x_local, Dt_full) + return self.triton_l_1_rotate_back(x_local, Dt_full) # Restore reduced global layout (E, D_m, C_wide) for the dense kernel. x_std = ( x_local.transpose(1, 2) .contiguous() .reshape(n_edge, self.reduced_dim, self.hidden_channels) ) - return self._rotate_back_fn(x_std, Dt_full) + return self.triton_l_1_rotate_back(x_std, Dt_full) return super()._rotate_back(x_local, edge_cache, n_edge) diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/wignerd.py b/deepmd/pt_expt/descriptor/dpa4_nn/wignerd.py index aa6be46d21..64ea10c738 100644 --- a/deepmd/pt_expt/descriptor/dpa4_nn/wignerd.py +++ b/deepmd/pt_expt/descriptor/dpa4_nn/wignerd.py @@ -98,8 +98,8 @@ def __init__( ] # The monomial basis routes through whichever accelerated backend # is selected; the two gates are mutually exclusive. - self._use_cutile_monomials = use_cutile_infer() - self._use_triton_monomials = triton_infer_level() >= 1 + self.cutile_infer_monomials = use_cutile_infer() + self.triton_infer_l_1_monomials = triton_infer_level() >= 1 # The l = 2 contraction tensor collapsed onto the 35 unique # degree-4 monomials: column m of the coefficient matrix sums # C_l2[:, :, p] over the 4^4 index tuples p whose component @@ -185,9 +185,9 @@ def _monomial_matrix( exponents is not None and edge_quaternion.is_cuda and not self.training - and (self._use_triton_monomials or self._use_cutile_monomials) + and (self.triton_infer_l_1_monomials or self.cutile_infer_monomials) ): - if self._use_cutile_monomials: + if self.cutile_infer_monomials: from deepmd.pt_expt.kernels.cutile.sezm.wigner_monomials import ( wigner_monomials as monomial_basis, ) @@ -212,9 +212,9 @@ def _compute_l2_block(self, edge_quaternion: torch.Tensor) -> torch.Tensor: exponents is not None and edge_quaternion.is_cuda and not self.training - and (self._use_triton_monomials or self._use_cutile_monomials) + and (self.triton_infer_l_1_monomials or self.cutile_infer_monomials) ): - if self._use_cutile_monomials: + if self.cutile_infer_monomials: from deepmd.pt_expt.kernels.cutile.sezm.wigner_monomials import ( wigner_monomials as monomial_basis, ) diff --git a/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py b/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py index 7a86943666..dcc104e92f 100644 --- a/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py +++ b/deepmd/pt_expt/kernels/cuda/dpa4/so2_conv.py @@ -563,7 +563,7 @@ def __init__(self, conv: Any) -> None: self._conv = conv mixer = conv.radial_degree_mixer self._rank = 0 if mixer is None else int(mixer.rank) - self._compete = bool(conv.focus_compete and conv.n_focus > 1) + self.focus_compete = bool(conv.focus_compete and conv.n_focus > 1) # Packed-run polynomial tables, fitted once per degree and materialized # on the compute device at the first call. self._tables_cpu = wigner_run_tables(conv.lmax) @@ -743,7 +743,7 @@ def __call__( conv = self._conv n_node = x.shape[0] kc, cb = self._degree_kernel(rad_feat) - if self._compete: + if self.focus_compete: fscale = self._focus_scale(x, edge_cache, rad_feat, kc, cb) # (E, F) else: fscale = x.new_empty(0) diff --git a/deepmd/pt_expt/kernels/cute/sezm/output_grid/__init__.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid/__init__.py new file mode 100644 index 0000000000..bae1a1c36e --- /dev/null +++ b/deepmd/pt_expt/kernels/cute/sezm/output_grid/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Output-grid orchestration and kernels for the CuTe SeZM path.""" diff --git a/deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/__init__.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/__init__.py similarity index 71% rename from deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/__init__.py rename to deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/__init__.py index 3412625810..4cde7864c0 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/__init__.py +++ b/deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/__init__.py @@ -1,3 +1,3 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""CuTe kernels for guarded non-K1 DeePMD operations.""" +"""CuTe DSL kernels for SeZM output-grid contractions.""" diff --git a/deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/cute_readout_l0.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/readout_l0.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/cute_readout_l0.py rename to deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/readout_l0.py index 793367a4e1..d16a2571b6 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/cute_readout_l0.py +++ b/deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/readout_l0.py @@ -23,7 +23,7 @@ make_fake_stream, ) -from .cute_tiled_grid_product import ( +from .tiled_product import ( FAKE_TENSOR_KW, GRID_SIZE, PACKED_COEFF_DIM, diff --git a/deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/cute_tiled_grid_product.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/tiled_product.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/cute_tiled_grid_product.py rename to deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/tiled_product.py index 582fc8ef22..ccdaf239f7 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels/cute_tiled_grid_product.py +++ b/deepmd/pt_expt/kernels/cute/sezm/output_grid/kernels/tiled_product.py @@ -24,7 +24,7 @@ make_fake_stream, ) -from .. import ( +from ... import ( runtime_policy, ) diff --git a/deepmd/pt_expt/kernels/cute/sezm/output_grid_product.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid/product.py similarity index 98% rename from deepmd/pt_expt/kernels/cute/sezm/output_grid_product.py rename to deepmd/pt_expt/kernels/cute/sezm/output_grid/product.py index 4d106f0a4e..de4e343bf6 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/output_grid_product.py +++ b/deepmd/pt_expt/kernels/cute/sezm/output_grid/product.py @@ -12,10 +12,10 @@ import torch -from . import ( +from .. import ( runtime_policy, ) -from .runtime_policy import ( +from ..runtime_policy import ( PORTABLE_TILED_BACKEND, SUPPORTED_HIDDEN_CHANNELS, select_output_grid_backend, @@ -121,7 +121,7 @@ def _output_grid_product_impl( from_grid, n_frames, ) - from .output_grid_kernels.cute_tiled_grid_product import ( + from .kernels.tiled_product import ( run_tiled_output_grid_product, ) @@ -173,7 +173,7 @@ def _output_grid_product_bwd_impl( or grad_out.device != left.device ): raise ValueError("grad_out must match the fused output GridMLP output") - from .output_grid_kernels.cute_tiled_grid_product import ( + from .kernels.tiled_product import ( run_tiled_output_grid_product_backward, ) diff --git a/deepmd/pt_expt/kernels/cute/sezm/readout_l0.py b/deepmd/pt_expt/kernels/cute/sezm/output_grid/readout_l0.py similarity index 94% rename from deepmd/pt_expt/kernels/cute/sezm/readout_l0.py rename to deepmd/pt_expt/kernels/cute/sezm/output_grid/readout_l0.py index 3ba6e4403b..0a0af62277 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/readout_l0.py +++ b/deepmd/pt_expt/kernels/cute/sezm/output_grid/readout_l0.py @@ -13,10 +13,10 @@ import torch -from . import ( +from .. import ( runtime_policy, ) -from .runtime_policy import ( +from ..runtime_policy import ( PORTABLE_TILED_BACKEND, select_output_grid_backend, ) @@ -126,19 +126,26 @@ def build_readout_l0_gram( ).contiguous() -def _has_exact_neo_readout_structure(output_ffn: Any) -> bool: - from deepmd.pt.model.descriptor.sezm_nn.ffn import ( - EquivariantFFN, - ) - from deepmd.pt.model.descriptor.sezm_nn.grid_net import ( - GridMLP, - SO3GridNet, - ) +def has_neo_readout_contract(output_ffn: Any) -> bool: + """Return whether an output FFN matches the fused Neo readout contract. - if type(output_ffn) is not EquivariantFFN: + Parameters + ---------- + output_ffn : Any + Candidate equivariant output FFN. + + Returns + ------- + bool + Whether the FFN and its grid projection expose the required structure. + """ + if type(output_ffn).__name__ != "EquivariantFFN": return False grid_net = output_ffn.act - if type(grid_net) is not SO3GridNet or type(grid_net.grid_op) is not GridMLP: + if ( + type(grid_net).__name__ != "SO3GridNet" + or type(grid_net.grid_op).__name__ != "GridMLP" + ): return False grid_op = grid_net.grid_op projector = grid_net.projector @@ -215,7 +222,7 @@ def _has_exact_neo_readout_contract( ffn_in: torch.Tensor, ) -> bool: if ( - not _has_exact_neo_readout_structure(output_ffn) + not has_neo_readout_contract(output_ffn) or not _inference_mode_is_frozen(output_ffn) or ffn_in.ndim != 4 or ffn_in.shape[0] <= 0 @@ -257,6 +264,19 @@ def _readout_input_fold_sources(output_ffn: Any) -> tuple[torch.Tensor, ...]: ) +def _project_frames( + coeff: torch.Tensor, + projection: Any, + n_frames: int, +) -> torch.Tensor: + """Apply one channel projection independently to every frame.""" + n_batch, coefficient_dim, n_focus, _ = coeff.shape + projected = projection( + coeff.reshape(n_batch, coefficient_dim, n_focus, n_frames, -1) + ) + return projected.reshape(n_batch, coefficient_dim, n_focus, -1) + + def _readout_input_fold_cache_key( sources: tuple[torch.Tensor, ...], ) -> tuple[tuple[Any, ...], ...]: @@ -410,9 +430,9 @@ def prepare_sm80_readout_input_fold( """Prepare immutable folded weights before compiling the frozen readout.""" if torch.compiler.is_compiling(): raise RuntimeError(_READOUT_INPUT_FOLD_PREPARE_ERROR) - if not _has_exact_neo_readout_structure( + if not has_neo_readout_contract(output_ffn) or not _inference_mode_is_frozen( output_ffn - ) or not _inference_mode_is_frozen(output_ffn): + ): raise ValueError("readout input folding requires the exact frozen Neo FFN") sources = _readout_input_fold_sources(output_ffn) @@ -459,9 +479,9 @@ def maybe_prepare_sm80_readout_input_fold( compute_capability: tuple[int, int] | None = None, ) -> bool: """Prepare only when a supported frozen Neo readout contract matches.""" - if not _has_exact_neo_readout_structure( + if not has_neo_readout_contract(output_ffn) or not _inference_mode_is_frozen( output_ffn - ) or not _inference_mode_is_frozen(output_ffn): + ): return False sources = _readout_input_fold_sources(output_ffn) device = sources[0].device @@ -570,12 +590,8 @@ def _run_neo_readout_l0( ], ) -> torch.Tensor: """Complete the exact Neo readout around one row-zero grid product.""" - if not _has_exact_neo_readout_structure(output_ffn): + if not has_neo_readout_contract(output_ffn): raise ValueError("readout l=0 completion requires the exact Neo output FFN") - from deepmd.pt.model.descriptor.sezm_nn.grid_net import ( - _project_frames, - ) - grid_net = output_ffn.act grid_op = grid_net.grid_op @@ -638,7 +654,10 @@ def run_neo_output_readout( candidate = maybe_run_neo_readout_l0(output_ffn, ffn_in) if candidate is not None: return candidate - return (ffn_in[:, 0:1, :, :] + output_ffn.forward_scalar(ffn_in)).reshape( + forward_scalar = getattr(output_ffn, "forward_scalar", None) + if forward_scalar is None: + forward_scalar = output_ffn.call_scalar + return (ffn_in[:, 0:1, :, :] + forward_scalar(ffn_in)).reshape( ffn_in.shape[0], output_ffn.channels ) @@ -650,7 +669,7 @@ def _readout_l0_impl( from_grid: torch.Tensor, ) -> torch.Tensor: _validate_product_contract(left, right, to_grid, from_grid) - from .output_grid_kernels.cute_readout_l0 import ( + from .kernels.readout_l0 import ( run_readout_l0, ) @@ -679,7 +698,7 @@ def _readout_l0_bwd_impl( or not dq0.is_contiguous() ): raise ValueError("dq0 must be contiguous and have shape (N,192)") - from .output_grid_kernels.cute_readout_l0 import ( + from .kernels.readout_l0 import ( run_readout_l0_backward, ) @@ -774,6 +793,7 @@ def readout_l0_product_cute( __all__ = [ "build_readout_l0_gram", + "has_neo_readout_contract", "invalidate_neo_readout_input_fold", "maybe_prepare_sm80_readout_input_fold", "maybe_run_neo_readout_l0", diff --git a/deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py b/deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py index c0bf24e754..3cdd1730bf 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py +++ b/deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py @@ -23,13 +23,13 @@ "DP_CUTE_OUTPUT_GRID_SM90_C96_ASYMMETRIC_PANELS" ) READOUT_INPUT_FOLD_SM90_ENV = "DP_CUTE_READOUT_INPUT_FOLD_SM90" -SUPPORTED_K1_CAPABILITIES = SM80_PROFILE_CAPABILITIES | frozenset( +SUPPORTED_SO2_CAPABILITIES = SM80_PROFILE_CAPABILITIES | frozenset( {(8, 9), (9, 0), (10, 0), (12, 0)} ) _GIE_DEFAULT_CAPABILITIES = SM80_PROFILE_CAPABILITIES INT32_MAX = (1 << 31) - 1 -K1_VALUES_PER_EDGE = 10 * 64 -K1_VALUES_PER_NODE = 16 * 64 +SO2_VALUES_PER_EDGE = 10 * 64 +SO2_VALUES_PER_NODE = 16 * 64 PORTABLE_TILED_BACKEND = "portable_tiled" PYTORCH_BACKEND = "pytorch" SUPPORTED_HIDDEN_CHANNELS = (96, 192) @@ -129,12 +129,12 @@ def _profile_or_explicit_feature( @torch.compiler.assume_constant_result -def is_k1_thin_wrapper_enabled( +def is_so2_thin_wrapper_enabled( compute_capability: tuple[int, int] | None = None, ) -> bool: - """Select compile-visible K1 dispatch for the SM80 profile.""" + """Select compile-visible SO2 dispatch for the SM80 profile.""" return _profile_or_explicit_feature( - "DP_CUTE_K1_THIN_WRAPPER", + "DP_CUTE_SO2_THIN_WRAPPER", compute_capability, ) @@ -213,22 +213,22 @@ def is_readout_input_fold_enabled( ) or is_readout_input_fold_sm90_enabled(compute_capability) -def is_supported_k1_capability(compute_capability: tuple[int, int]) -> bool: - """Return whether K1 supports this compute capability.""" - return tuple(compute_capability) in SUPPORTED_K1_CAPABILITIES +def is_supported_so2_capability(compute_capability: tuple[int, int]) -> bool: + """Return whether SO2 supports this compute capability.""" + return tuple(compute_capability) in SUPPORTED_SO2_CAPABILITIES -def k1_int32_indexing_is_safe( +def so2_int32_indexing_is_safe( *, edge_count: int, node_count: int, ) -> bool: - """Check every flattened K1 offset represented with signed Int32.""" + """Check every flattened SO2 offset represented with signed Int32.""" if edge_count < 0 or node_count < 0: return False return ( - edge_count <= INT32_MAX // K1_VALUES_PER_EDGE - and node_count <= INT32_MAX // K1_VALUES_PER_NODE + edge_count <= INT32_MAX // SO2_VALUES_PER_EDGE + and node_count <= INT32_MAX // SO2_VALUES_PER_NODE ) @@ -266,21 +266,21 @@ def is_gie_enabled(compute_capability: tuple[int, int]) -> bool: def is_packed_wigner_enabled(compute_capability: tuple[int, int]) -> bool: - """Select packed Wigner storage required by the optimized K1 profiles.""" + """Select packed Wigner storage required by the optimized SO2 profiles.""" return _master_gated_feature( - "DP_CUTE_K1_PACKED_WIGNER", - default=compute_capability in SUPPORTED_K1_CAPABILITIES, + "DP_CUTE_SO2_PACKED_WIGNER", + default=compute_capability in SUPPORTED_SO2_CAPABILITIES, ) @torch.compiler.assume_constant_result -def is_k1_eager_island_enabled( +def is_so2_eager_island_enabled( compute_capability: tuple[int, int] | None = None, ) -> bool: - """Select the SM80 neighbor-list eager island for the Neo K1 path.""" + """Select the SM80 neighbor-list eager island for the Neo SO2 path.""" if not is_cute_infer_enabled(): return False - override = _env_override("DP_CUTE_K1_EAGER_ISLANDS") + override = _env_override("DP_CUTE_SO2_EAGER_ISLANDS") if override is not None: return override if compute_capability is None: diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/__init__.py b/deepmd/pt_expt/kernels/cute/sezm/so2/__init__.py similarity index 68% rename from deepmd/pt_expt/kernels/cute/sezm/k1_kernels/__init__.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/__init__.py index 77198c0277..896b4aa2d9 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/__init__.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/__init__.py @@ -1,3 +1,3 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""Private CuTe kernels for the Neo SO2Conv/K1 inference path.""" +"""SO(2) orchestration and kernels for the CuTe SeZM path.""" diff --git a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/__init__.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/__init__.py similarity index 72% rename from deepmd/pt_expt/kernels/cute/sezm/sm90_k1/__init__.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/kernels/__init__.py index 353bc4a488..910740da8d 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/__init__.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/__init__.py @@ -1,3 +1,3 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""SM90 split-complex Neo K1 implementation.""" +"""Portable CuTe DSL kernels for the Neo SO(2) path.""" diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_so2_gate_combined_fwd.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/combined_gate_forward.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_so2_gate_combined_fwd.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/kernels/combined_gate_forward.py index 9a1de48779..c1654c898c 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_so2_gate_combined_fwd.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/combined_gate_forward.py @@ -21,10 +21,10 @@ make_fake_tensor, ) -from ..compile_cache import ( +from ...compile_cache import ( device_aware_lru_cache, ) -from ..runtime_policy import ( +from ...runtime_policy import ( FUSED_SO2_GATE_CAPABILITIES, ) @@ -95,7 +95,7 @@ def prepare_neo_so2_gate_combined_weights( class CuteNeoSO2GateCombined: - """M64/K16/T256/S3 SIMT SGEMMs with one CTA-resident gate epilogue.""" + """M64/SO26/T256/S3 SIMT SGEMMs with one CTA-resident gate epilogue.""" def __init__(self) -> None: self.cta_sync_barrier = pipeline.NamedBarrier( diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_envelope_gated_softmax.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/envelope_softmax.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_envelope_gated_softmax.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/kernels/envelope_softmax.py index 45a996e816..d709e3eff2 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_envelope_gated_softmax.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/envelope_softmax.py @@ -30,7 +30,7 @@ make_fake_stream, ) -from ..compile_cache import ( +from ...compile_cache import ( device_aware_lru_cache, ) diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_focus_src_backward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/focus_source_backward.py similarity index 97% rename from deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_focus_src_backward.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/kernels/focus_source_backward.py index 8aee494844..c3d6cdde81 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_focus_src_backward.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/focus_source_backward.py @@ -8,7 +8,7 @@ focus RMSNorm -> two-focus logits -> softmax -> label smoothing scalar Q/K RMSNorm -> Q projection + K projection -The kernel is specialized to the Neo K1 shape `(E, F=2, C=32)`. +The kernel is specialized to the Neo SO2 shape `(E, F=2, C=32)`. """ from __future__ import ( @@ -155,18 +155,18 @@ def neo_attention_prelude_forward_kernel( q0 = cutlass.Float32(0.0) k0 = cutlass.Float32(0.0) q1 = cutlass.Float32(0.0) - k1 = cutlass.Float32(0.0) + so2 = cutlass.Float32(0.0) for input_channel in cutlass.range_constexpr(32): value0 = cute.arch.shuffle_sync(qk_norm0, input_channel) value1 = cute.arch.shuffle_sync(qk_norm1, input_channel) q0 += value0 * q_weight[input_channel, 0, lane].to(cutlass.Float32) k0 += value0 * k_weight[input_channel, 0, lane].to(cutlass.Float32) q1 += value1 * q_weight[input_channel, 1, lane].to(cutlass.Float32) - k1 += value1 * k_weight[input_channel, 1, lane].to(cutlass.Float32) + so2 += value1 * k_weight[input_channel, 1, lane].to(cutlass.Float32) q_node[node, 0, lane] = q0.to(q_node.element_type) k_node[node, 0, lane] = k0.to(k_node.element_type) q_node[node, 1, lane] = q1.to(q_node.element_type) - k_node[node, 1, lane] = k1.to(k_node.element_type) + k_node[node, 1, lane] = so2.to(k_node.element_type) @lru_cache(maxsize=8) diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/gate_linear_residual_backward.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/kernels/gate_linear_residual_backward.py index 27d356caa0..f0b21e1c18 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/gate_linear_residual_backward.py @@ -21,7 +21,7 @@ make_fake_stream, ) -from ..compile_cache import ( +from ...compile_cache import ( device_aware_lru_cache, ) diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_message_grid_product.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/message_grid_product.py similarity index 96% rename from deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_message_grid_product.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/kernels/message_grid_product.py index 54d66b9d99..72cc3bbe3d 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_message_grid_product.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/message_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 -"""Packed strict-FP32 CuTe grid product for Neo's F=2 K1 branch.""" +"""Packed strict-FP32 CuTe grid product for Neo's F=2 SO2 branch.""" from __future__ import ( annotations, @@ -24,7 +24,7 @@ make_fake_stream, ) -from .. import ( +from ... import ( runtime_policy, ) @@ -46,7 +46,7 @@ def _make_grid_operation( *, panel_adjoint: bool = False, ) -> Any: - from ..output_grid_kernels import cute_tiled_grid_product as tiled + from ...output_grid.kernels import tiled_product as tiled # F=2 and C=32 form one complete 64-channel panel. Bypass only the # shared readout policy; the tiled implementation itself is unchanged. @@ -111,7 +111,7 @@ def compile_message_grid_product_forward( import torch _validate_compile_target(device_index, compute_capability) - from ..output_grid_kernels.cute_tiled_grid_product import ( + from ...output_grid.kernels.tiled_product import ( TiledOutputGridProductForward, ) @@ -138,7 +138,7 @@ def compile_message_grid_product_backward( import torch _validate_compile_target(device_index, compute_capability) - from ..output_grid_kernels.cute_tiled_grid_product import ( + from ...output_grid.kernels.tiled_product import ( TiledOutputGridProductBackward, ) diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_output_gate_backward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/output_gate_backward.py similarity index 100% rename from deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_output_gate_backward.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/kernels/output_gate_backward.py diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_a_radial_forward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_a_radial_forward.py similarity index 97% rename from deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_a_radial_forward.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_a_radial_forward.py index 488806d314..e64e6b06e0 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_a_radial_forward.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_a_radial_forward.py @@ -25,10 +25,10 @@ make_fake_stream, ) -from ..compile_cache import ( +from ...compile_cache import ( device_aware_lru_cache, ) -from ..k1_wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES +from ..wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} @@ -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 one rank of 64 channel weights" + ) combined_weight = _combined_radial_weight( radial_hidden_proj, diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_backward_layout.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_backward.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_backward_layout.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_backward.py index 43e3f7ef14..6280257bb5 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_backward_layout.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_backward.py @@ -44,7 +44,7 @@ make_fake_stream, ) -from ..k1_wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES +from ..wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES if TYPE_CHECKING: from collections.abc import ( diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_onepass.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_forward.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_onepass.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_forward.py index ee489af8bc..d4adb1243a 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_onepass.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_forward.py @@ -41,10 +41,10 @@ make_fake_stream, ) -from ..compile_cache import ( +from ...compile_cache import ( device_aware_lru_cache, ) -from ..k1_wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES +from ..wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES if TYPE_CHECKING: from collections.abc import ( diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_qk_edge.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/qk_edge.py similarity index 100% rename from deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_qk_edge.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/kernels/qk_edge.py diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_radial_phase_a_backward_node.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/radial_phase_a_backward.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_radial_phase_a_backward_node.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/kernels/radial_phase_a_backward.py index 603e330fb2..e05aca3c51 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_radial_phase_a_backward_node.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/radial_phase_a_backward.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Source-CSR node-tiled Neo radial and Phase-A backward. -This is the contention-free node-tiled implementation for the exact Neo K1 shape +This is the contention-free node-tiled implementation for the exact Neo SO2 shape ``lmax=3, D=16, Dm=10, Cwide=64, F=2``. Edge tensors retain their physical destination-sorted order. ``source_ptr`` delimits intervals in ``source_order``, whose slots hold the corresponding physical edge ids. @@ -37,7 +37,7 @@ make_fake_stream, ) -from ..k1_wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES +from ..wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES FAKE_TENSOR_KW = {"assumed_align": 16, "use_32bit_stride": True} diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/structural_gate_sm80.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/kernels/structural_gate_sm80.py index c0a4fc058e..7ad1616496 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/structural_gate_sm80.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later # ruff: noqa: ANN001, ANN201, ANN202, TC002, TC003 -"""Strict-FP32 vectorized structural-gate forward/backward for Neo K1. +"""Strict-FP32 vectorized structural-gate forward/backward for Neo SO2. The split path keeps both neighboring SO2 products and the two gate projections in cuBLAS. Its elementwise consumer assigns eight threads to each 32-channel @@ -30,7 +30,7 @@ make_fake_stream, ) -from .. import ( +from ... import ( runtime_policy, ) @@ -51,7 +51,7 @@ def _guard_vec4_dispatch( kernel: Callable, tensor_names: tuple[str, ...], ) -> Callable: - from ..k1_gate_structural import ( + from ..structural_gate import ( _dispatch_aligned_vec4_kernel, ) @@ -392,10 +392,10 @@ def compile_neo_gate_split_structural_vec4_sm80_forward( ) -> Callable: if ( compile_identity is not None - and compile_identity[1:] not in runtime_policy.SUPPORTED_K1_CAPABILITIES + and compile_identity[1:] not in runtime_policy.SUPPORTED_SO2_CAPABILITIES ): raise ValueError( - "vectorized structural gate forward requires a supported K1 device" + "vectorized structural gate forward requires a supported SO2 device" ) rows = cute.sym_int64() edges = cute.sym_int64() @@ -444,10 +444,10 @@ def compile_neo_gate_split_structural_vec4_sm80_backward( ) -> Callable: if ( compile_identity is not None - and compile_identity[1:] not in runtime_policy.SUPPORTED_K1_CAPABILITIES + and compile_identity[1:] not in runtime_policy.SUPPORTED_SO2_CAPABILITIES ): raise ValueError( - "vectorized structural gate backward requires a supported K1 device" + "vectorized structural gate backward requires a supported SO2 device" ) rows = cute.sym_int64() edges = cute.sym_int64() diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_so2linear.py b/deepmd/pt_expt/kernels/cute/sezm/so2/linear.py similarity index 98% rename from deepmd/pt_expt/kernels/cute/sezm/k1_so2linear.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/linear.py index 8f35d4afd2..8912823e9a 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_so2linear.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/linear.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""SO2Linear helpers for the Neo CuTe K1 path. +"""SO2Linear helpers for the Neo CuTe SO2 path. The in-place residual adjoint enforces FP32 operand dtypes. Strict FP32 GEMM also requires the caller to select highest float32 matmul precision and disable @@ -37,9 +37,7 @@ def _validate_neo_so2_linear(so2_linear: Any) -> None: or so2_linear.n_focus != 2 or so2_linear.mlp_bias ): - raise NotImplementedError( - "Neo CuTe K1 SO2Linear expects lmax=3,mmax=1,F=2,C=32" - ) + raise NotImplementedError("Neo CuTe SO2Linear expects lmax=3,mmax=1,F=2,C=32") def run_neo_so2_linear_manual( @@ -191,7 +189,7 @@ def _validate_edge_focus(tensor: Tensor, name: str) -> None: if tensor.shape[0] <= 0: raise ValueError(f"{name} requires E > 0") if not tensor.is_contiguous(): - raise ValueError(f"{name} must use the canonical contiguous K1 layout") + raise ValueError(f"{name} must use the canonical contiguous SO2 layout") def _validate_weight( diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_message_grid_packed.py b/deepmd/pt_expt/kernels/cute/sezm/so2/message_grid.py similarity index 98% rename from deepmd/pt_expt/kernels/cute/sezm/k1_message_grid_packed.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/message_grid.py index c7d9a6ce2c..38d5b3d414 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_message_grid_packed.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/message_grid.py @@ -17,7 +17,7 @@ ) if TYPE_CHECKING: - from .message_grid_readout_sm90 import ( + from .sm90.message_grid_readout import ( Sm90MessageGridState, ) @@ -209,7 +209,7 @@ def run_packed_message_grid_forward( ) -> 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) - from .k1_kernels.cute_neo_message_grid_product import ( + from .kernels.message_grid_product import ( run_message_grid_product, ) @@ -228,7 +228,7 @@ def run_packed_message_grid_forward( net.projector.from_grid_mat, ) else: - from .message_grid_gaunt_sm90 import ( + from .sm90.message_grid_gaunt import ( run_sm90_gaunt_forward, ) @@ -268,7 +268,7 @@ def run_packed_message_grid_backward( ) -> tuple[Tensor, Tensor]: """Return query/context adjoints while retaining packed grid intermediates.""" query, context = _validate_contract(net, query_flat, context_flat) - from .k1_kernels.cute_neo_message_grid_product import ( + from .kernels.message_grid_product import ( run_message_grid_product, run_message_grid_product_backward, ) diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/metadata.py b/deepmd/pt_expt/kernels/cute/sezm/so2/metadata.py new file mode 100644 index 0000000000..c47cf609f1 --- /dev/null +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/metadata.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Build CuTe SO2 metadata when the input does not carry graph-owned CSR.""" + +from __future__ import ( + annotations, +) + +import torch + + +def build_sorted_edge_index_metadata( + src: torch.Tensor, + dst: torch.Tensor, + n_nodes: int, + *, + validate_sorted: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build destination and source CSR metadata for one sorted edge list. + + Parameters + ---------- + src : torch.Tensor + Source indices with shape ``(E,)``. + dst : torch.Tensor + Nondecreasing destination indices with shape ``(E,)``. + n_nodes : int + Number of nodes addressed by the edge list. + validate_sorted : bool, default: False + Whether to assert the destination-order contract on the device. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor, torch.Tensor] + Destination row pointers, the source-major edge permutation, and source + row pointers. All three tensors use int32 storage. + + Raises + ------ + ValueError + If the edge arrays have incompatible shapes or devices, or if their + edge count exceeds int32 indexing. + TypeError + If either edge-index tensor is not int32 or int64. + """ + if src.dim() != 1 or dst.dim() != 1: + raise ValueError("src and dst must be one-dimensional") + if src.shape != dst.shape: + raise ValueError("src and dst must have the same shape") + if src.device != dst.device: + raise ValueError("src and dst must be on the same device") + if src.dtype not in (torch.int32, torch.int64): + raise TypeError("src must have dtype int32 or int64") + if dst.dtype not in (torch.int32, torch.int64): + raise TypeError("dst must have dtype int32 or int64") + if n_nodes < 0: + raise ValueError("n_nodes must be non-negative") + if src.numel() > 2**31 - 1: + raise ValueError("sorted edge metadata requires E <= 2**31 - 1") + + src = src.contiguous() + dst = dst.contiguous() + if validate_sorted and dst.numel() > 1: + torch._assert_async( + torch.all(dst[1:] >= dst[:-1]), + "Neo SO2 destinations_sorted=True requires monotonically " + "nondecreasing destination indices", + ) + destination_boundaries = torch.arange( + n_nodes + 1, + device=dst.device, + dtype=dst.dtype, + ) + destination_row_ptr = torch.searchsorted( + dst, + destination_boundaries, + out_int32=True, + ).contiguous() + + source_order_i64 = torch.argsort(src, stable=True) + sorted_src = src.index_select(0, source_order_i64) + source_boundaries = torch.arange( + n_nodes + 1, + device=src.device, + dtype=src.dtype, + ) + source_row_ptr = torch.searchsorted( + sorted_src, + source_boundaries, + out_int32=True, + ).contiguous() + source_order = source_order_i64.to(dtype=torch.int32).contiguous() + return destination_row_ptr, source_order, source_row_ptr + + +__all__ = ["build_sorted_edge_index_metadata"] diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1.py b/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py similarity index 89% rename from deepmd/pt_expt/kernels/cute/sezm/k1.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/operation.py index 7e227c332a..90fc8d9d65 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""Opt-in CuTe Neo K1 custom op for SeZM/DPA4 inference.""" +"""Opt-in CuTe Neo SO2 custom op for SeZM/DPA4 inference.""" from __future__ import ( annotations, @@ -27,21 +27,17 @@ Tensor, ) -from deepmd.pt.model.descriptor.sezm_nn.norm import ( - EquivariantRMSNorm, -) - -from . import ( +from .. import ( runtime_policy, ) -from .k1_message_grid_packed import ( +from .message_grid import ( is_supported_message_grid, ) @dataclass(frozen=True) -class NeoK1RuntimeConfig: - """Architecture-specific choices for the compact Neo K1 path.""" +class NeoSO2RuntimeConfig: + """Architecture-specific choices for the compact Neo SO(2) path.""" native_sm90_path: bool = False per_focus_so2_fwd_pair: bool = False @@ -49,8 +45,8 @@ class NeoK1RuntimeConfig: @dataclass(frozen=True) -class NeoK1Spec: - """Shape and branch contract for one Neo K1 replacement.""" +class NeoSO2Spec: + """Shape and branch contract for one Neo SO(2) unit replacement.""" lmax: int node_lmax: int @@ -107,11 +103,6 @@ def is_current_neo_target(self) -> bool: ) -def is_cute_k1_enabled() -> bool: - """Return whether the Neo CuTe K1 path is enabled by environment.""" - return runtime_policy.is_cute_infer_enabled() - - def _is_identity(module: Any) -> bool: return ( module.__class__.__name__ == "Identity" @@ -119,10 +110,10 @@ def _is_identity(module: Any) -> bool: ) -def get_neo_k1_spec(block: Any) -> NeoK1Spec: - """Extract the K1 contract from a SeZM interaction block.""" +def get_neo_so2_spec(block: Any) -> NeoSO2Spec: + """Extract the SO2 contract from a SeZM interaction block.""" so2 = block.so2_conv - return NeoK1Spec( + return NeoSO2Spec( lmax=int(so2.lmax), node_lmax=int(block.node_lmax), mmax=int(so2.mmax), @@ -150,15 +141,43 @@ def get_neo_k1_spec(block: Any) -> NeoK1Spec: ) -def is_supported_neo_k1_block(block: Any) -> bool: - """Return whether this block can use the current Neo CuTe K1 path.""" +def has_equivariant_rms_norm_contract(module: Any) -> bool: + """Return whether a norm exposes the exact state consumed by Neo SO2. + + Parameters + ---------- + module : Any + Candidate equivariant normalization module. + + Returns + ------- + bool + Whether the module's layout and state match the fused kernel contract. + """ + return ( + type(module).__name__ == "EquivariantRMSNorm" + and int(getattr(module, "lmax", -1)) == 3 + and int(getattr(module, "channels", -1)) == 32 + and int(getattr(module, "n_focus", -1)) == 1 + and tuple(getattr(getattr(module, "adam_scale", None), "shape", ())) + == (4, 1, 32) + and tuple(getattr(getattr(module, "bias", None), "shape", ())) == (1, 32) + and tuple(getattr(getattr(module, "balance_weight", None), "shape", ())) + == (16,) + and tuple(getattr(getattr(module, "expand_index", None), "shape", ())) == (16,) + and hasattr(module, "eps") + ) + + +def is_supported_neo_so2_block(block: Any) -> bool: + """Return whether this block can use the current Neo CuTe SO2 path.""" so2 = block.so2_conv return ( - get_neo_k1_spec(block).is_current_neo_target + get_neo_so2_spec(block).is_current_neo_target and bool(so2.focus_norm) and not bool(so2.edge_cartesian) and getattr(so2, "node_cartesian_tp", None) is None - and type(block.post_so2_norm) is EquivariantRMSNorm + and has_equivariant_rms_norm_contract(block.post_so2_norm) and is_supported_message_grid(so2.message_node_grid_product) ) @@ -292,11 +311,11 @@ def _dtypes_use_strict_fp32(dtypes: Any) -> bool: return all(dtype == torch.float32 for dtype in dtypes) -def is_supported_k1_compute_capability( +def is_supported_so2_compute_capability( compute_capability: tuple[int, int], ) -> bool: - """Return whether K1 supports this compute capability.""" - return runtime_policy.is_supported_k1_capability(compute_capability) + """Return whether SO2 supports this compute capability.""" + return runtime_policy.is_supported_so2_capability(compute_capability) def _device_compute_capability(device: torch.device) -> tuple[int, int]: @@ -314,10 +333,10 @@ def _tensor_compute_capability(tensor: Any) -> tuple[int, int] | None: return _device_compute_capability(device) -def _device_is_supported_for_k1(device: torch.device) -> bool: +def _device_is_supported_for_so2(device: torch.device) -> bool: # 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( + return not torch.cuda.is_available() or is_supported_so2_compute_capability( _device_compute_capability(device) ) @@ -336,27 +355,27 @@ def packed_wigner_edges_eligible( and edge_count > 0 and destinations_sorted and _dtypes_use_strict_fp32(runtime_dtypes) - and runtime_policy.k1_int32_indexing_is_safe( + and runtime_policy.so2_int32_indexing_is_safe( edge_count=edge_count, node_count=node_count, ) ) -def is_neo_k1_static_eligible( +def is_neo_so2_static_eligible( block: Any, *, training: bool, device: torch.device, dtype: torch.dtype, ) -> bool: - """Check K1 conditions that are stable during one descriptor forward.""" + """Check SO2 conditions that are stable during one descriptor forward.""" return ( - is_cute_k1_enabled() + runtime_policy.is_cute_infer_enabled() and not training and not bool(getattr(block, "training", False)) and device.type == "cuda" - and _device_is_supported_for_k1(device) + and _device_is_supported_for_so2(device) and dtype == torch.float32 and not torch.is_autocast_enabled(device.type) and runtime_policy.uses_strict_fp32_matmul() @@ -364,12 +383,12 @@ def is_neo_k1_static_eligible( and _module_is_frozen(block) and _module_state_is_aligned(block) and _gate_expand_index_structure_is_supported(block) - and getattr(block, "_deepmd_cute_k1_state", None) is not False - and is_supported_neo_k1_block(block) + and getattr(block, "_deepmd_cute_so2_state", None) is not False + and is_supported_neo_so2_block(block) ) -def is_neo_k1_runtime_eligible( +def is_neo_so2_runtime_eligible( block: Any, *, training: bool, @@ -379,11 +398,11 @@ def is_neo_k1_runtime_eligible( node_count: int, destinations_sorted: bool, ) -> bool: - """Return the exact strict-FP32 inference contract for K1 dispatch.""" + """Return the exact strict-FP32 inference contract for SO2 dispatch.""" return ( edge_count > 0 and destinations_sorted - and is_neo_k1_static_eligible( + and is_neo_so2_static_eligible( block, training=training, device=device, @@ -404,7 +423,7 @@ def is_packed_wigner_candidate( ) -> bool: """Check packed-Wigner conditions known before edge construction.""" if has_edge_src_gate: - # K1 currently falls back for SFPG so eager must receive dense Wigner data. + # SO2 currently falls back for SFPG so eager must receive dense Wigner data. return False block_tuple = tuple(blocks) if device.type == "cuda": @@ -414,8 +433,8 @@ def is_packed_wigner_candidate( packed_wigner_enabled = False else: packed_wigner_enabled = ( - is_cute_k1_enabled() - and runtime_policy.is_supported_k1_capability(compute_capability) + runtime_policy.is_cute_infer_enabled() + and runtime_policy.is_supported_so2_capability(compute_capability) ) else: packed_wigner_enabled = False @@ -428,7 +447,7 @@ def is_packed_wigner_candidate( ): return False return all( - is_neo_k1_static_eligible( + is_neo_so2_static_eligible( block, training=training, device=device, @@ -458,15 +477,15 @@ def __init__( def block(self) -> Any: block = self._block_ref() if block is None: - raise RuntimeError("the registered Neo K1 block has been released") + raise RuntimeError("the registered Neo SO2 block has been released") return block @dataclass(frozen=True) -class _RegisteredK1State: +class _RegisteredSO2State: device_index: int handle: int - config: NeoK1RuntimeConfig + config: NeoSO2RuntimeConfig @dataclass @@ -510,7 +529,7 @@ def _compile_on_runner_device( def _compile_output_gate_backward( compile_identity: tuple[int, int, int], eps: float ) -> Any: - from .k1_kernels.cute_neo_output_gate_backward import ( + from .kernels.output_gate_backward import ( compile_neo_output_gate_backward, ) @@ -521,7 +540,7 @@ def _compile_output_gate_backward( ) -def register_cute_k1_block(block: Any, config: Any) -> int: +def register_cute_so2_block(block: Any, config: Any) -> int: """Register a DeePMD block/config pair and return a stable integer handle.""" global _NEXT_HANDLE @@ -542,14 +561,14 @@ def remove_collected_entry(block_ref: weakref.ReferenceType[Any]) -> None: return handle -def invalidate_cute_k1_state(block: Any) -> None: +def invalidate_cute_so2_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): + state = getattr(block, "_deepmd_cute_so2_state", None) + if isinstance(state, _RegisteredSO2State): 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_so2_state"): + delattr(block, "_deepmd_cute_so2_state") if hasattr(block, "_deepmd_cute_gate_expand_contract"): delattr(block, "_deepmd_cute_gate_expand_contract") @@ -558,20 +577,20 @@ def _validate_gate_expand_index(block: Any) -> None: """Pin the degree-to-gate map assumed by fused Neo gate kernels.""" if not _gate_expand_index_is_supported(block): raise ValueError( - "Neo K1 fused gate kernels require expand_index=[0,1,2,0,1,2,0,1,2]" + "Neo SO2 fused gate kernels require expand_index=[0,1,2,0,1,2,0,1,2]" ) @torch.compiler.disable -def _register_cute_k1_state( +def _register_cute_so2_state( block: Any, device_index: int, - config: NeoK1RuntimeConfig, -) -> _RegisteredK1State | None: - """Publish cold K1 state eagerly before its handle enters a custom op.""" + config: NeoSO2RuntimeConfig, +) -> _RegisteredSO2State | None: + """Publish cold SO2 state eagerly before its handle enters a custom op.""" _validate_gate_expand_index(block) - old_state = getattr(block, "_deepmd_cute_k1_state", None) - if isinstance(old_state, _RegisteredK1State): + old_state = getattr(block, "_deepmd_cute_so2_state", None) + if isinstance(old_state, _RegisteredSO2State): with _REGISTRY_LOCK: old_entry = _REGISTRY.get(old_state.handle) if ( @@ -585,26 +604,26 @@ def _register_cute_k1_state( 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. - block._deepmd_cute_k1_state = False + block._deepmd_cute_so2_state = False return None - state = _RegisteredK1State( + state = _RegisteredSO2State( device_index=device_index, - handle=register_cute_k1_block(block, config), + handle=register_cute_so2_block(block, config), config=config, ) - block._deepmd_cute_k1_state = state + block._deepmd_cute_so2_state = state return state @torch.compiler.disable -def prepare_cute_k1_blocks( +def prepare_cute_so2_blocks( blocks: Any, *, training: bool, device: torch.device, dtype: torch.dtype, ) -> bool: - """Validate and register K1 state before model graph capture begins.""" + """Validate and register SO2 state before model graph capture begins.""" if training or device.type != "cuda" or dtype != torch.float32: return False block_tuple = tuple(blocks) @@ -614,11 +633,11 @@ def prepare_cute_k1_blocks( if device_index is None: device_index = torch.cuda.current_device() compute_capability = _cuda_compute_capability(device_index) - if not is_supported_k1_compute_capability(compute_capability): + if not is_supported_so2_compute_capability(compute_capability): return False config = _architecture_default_config(compute_capability) if not all( - is_neo_k1_static_eligible( + is_neo_so2_static_eligible( block, training=training, device=device, @@ -628,7 +647,7 @@ def prepare_cute_k1_blocks( ): return False return all( - _register_cute_k1_state(block, device_index, config) is not None + _register_cute_so2_state(block, device_index, config) is not None for block in block_tuple ) @@ -637,7 +656,7 @@ def _runner_token_key( runner_token: Tensor, *, path: str ) -> tuple[str, int | None, int]: if runner_token.dtype != torch.uint8 or runner_token.numel() != 1: - raise ValueError(f"{path} Neo K1 runner token must be one uint8 value") + raise ValueError(f"{path} Neo SO2 runner token must be one uint8 value") return ( runner_token.device.type, runner_token.device.index, @@ -664,7 +683,7 @@ def _store_packed_runner(runner_token: Tensor, runner: Any) -> None: with _PACKED_RUNNER_CACHE_LOCK: old_state = _PACKED_RUNNER_CACHE.get(key) if old_state is not None: - raise RuntimeError("packed Neo K1 runner token is already outstanding") + raise RuntimeError("packed Neo SO2 runner token is already outstanding") else: _PACKED_RUNNER_CACHE[key] = state weakref.finalize(runner_token, _release_packed_runner, key, state) @@ -675,7 +694,7 @@ def _borrow_packed_runner(runner_token: Tensor) -> Any | None: with _PACKED_RUNNER_CACHE_LOCK: state = _PACKED_RUNNER_CACHE.get(key) if state is None: - raise RuntimeError("packed Neo K1 runner token is not outstanding") + raise RuntimeError("packed Neo SO2 runner token is not outstanding") with state.reservation_lock: # Transfer the forward runner once; retained VJPs rebuild isolated workspaces. runner = state.runner @@ -723,7 +742,7 @@ def _assert_grad_meta_contract( or actual.stride() != stride ): raise RuntimeError( - f"Neo K1 {name} gradient violates the custom-op meta contract: " + f"Neo SO2 {name} gradient violates the custom-op meta contract: " f"got shape={tuple(actual.shape)} stride={actual.stride()}, expected " f"shape={tuple(expected_like.shape)} stride={stride}" ) @@ -733,14 +752,14 @@ def _assert_not_cuda_graph_capturing(path: str, tensor: Tensor) -> None: """Reject stateful paths that cannot be safely captured and replayed.""" if tensor.is_cuda and torch.cuda.is_current_stream_capturing(): raise RuntimeError( - f"Neo K1 {path} cannot run inside direct CUDA graph capture; " + f"Neo SO2 {path} cannot run inside direct CUDA graph capture; " "use torch.compile without wrapping this stateful custom op in " "torch.cuda.graph" ) def _fake_x_wide_grad_like(x: Tensor, *, skip: bool) -> Tensor: - """Fake the native K1 grad-x layout when layout restore is skipped. + """Fake the native SO2 grad-x layout when layout restore is skipped. The manual SO3 pre-mix backward returns a degree-major view with shape ``(N, D, 1, C)`` and stride ``(C, N*C, C, 1)``. If the fake kernel @@ -765,8 +784,9 @@ def _equivariant_rmsnorm_backward(norm: Any, x: Tensor, grad_out: Tensor) -> Ten # Matches EquivariantRMSNorm.forward for the inference case. Parameters are # frozen; only the input gradient is needed for forces/stress. in_dtype = x.dtype - xf = x.to(dtype=norm.dtype) - gf = grad_out.to(dtype=norm.dtype) + norm_dtype = norm.adam_scale.dtype + xf = x.to(dtype=norm_dtype) + gf = grad_out.to(dtype=norm_dtype) x0_in = xf[:, :1, :, :] xt = xf[:, 1:, :, :] x0 = x0_in - x0_in.mean(dim=-1, keepdim=True) @@ -776,7 +796,8 @@ def _equivariant_rmsnorm_backward(norm: Any, x: Tensor, grad_out: Tensor) -> Ten mean_variance = mean_variance + torch.einsum( "ndfc,d->nf", xt * xt, norm.balance_weight[1:] ) - inv = torch.rsqrt(mean_variance + norm.eps_tensor).unsqueeze(1).unsqueeze(-1) + eps = getattr(norm, "eps_tensor", norm.eps) + inv = torch.rsqrt(mean_variance + eps).unsqueeze(1).unsqueeze(-1) expanded_scale = torch.index_select( norm.adam_scale, dim=0, index=norm.expand_index ).unsqueeze(0) @@ -861,7 +882,7 @@ def _neo_so2_linear_backward_input_with_residual( inplace_residual: bool = False, out: Tensor | None = None, ) -> Tensor: - from .k1_so2linear import ( + from .linear import ( cached_neo_so2_linear_weights, ) @@ -888,7 +909,7 @@ def _neo_so2_linear_backward_input_with_residual( if inplace_residual: if out is not None: raise ValueError("SO2 backward cannot select both in-place and out storage") - from .k1_so2linear import ( + from .linear import ( neo_so2_linear_backward_residual_inplace, ) @@ -899,7 +920,7 @@ def _neo_so2_linear_backward_input_with_residual( wpair_t, ) if grad_out is residual: - from .k1_gate_structural import ( + from .structural_gate import ( focus_major_so2_backward_with_folded_residual_out, ) @@ -941,7 +962,7 @@ def _neo_so2_linear_backward_input_with_residual( out=out, ) - raise RuntimeError("Neo K1 SO2 backward reached an unsupported storage pattern") + raise RuntimeError("Neo SO2 backward reached an unsupported storage pattern") def _so3_grid_cross_glu_flat_backward( @@ -1083,7 +1104,7 @@ def _final_manual_backward(runner: Any, grad_out: Tensor) -> tuple[Tensor, Tenso if runner.packed_message_grid: message_grid_product = runner.message_grid_product runner.message_grid_product = None - from .k1_message_grid_packed import ( + from .message_grid import ( run_packed_message_grid_backward, ) @@ -1361,9 +1382,9 @@ def _build_runner( entry = _REGISTRY[int(handle)] config = entry.config if config.native_sm90_path: - from .sm90_k1.runner import NeoSm90K1Runner as Runner + from .sm90.runner import NeoSm90SO2Runner as Runner else: - from .k1_runner import NeoFullCuteBackward as Runner + from .runner import NeoFullCuteBackward as Runner x = _aligned_contiguous(x) shared_wigner_storage = d_full.data_ptr() == dt_full.data_ptr() @@ -1403,7 +1424,7 @@ def _build_runner( return runner -def _k1_packed_direct_forward_impl( +def _so2_packed_direct_forward_impl( handle: int, x: Tensor, d_full: Tensor, @@ -1441,7 +1462,7 @@ def _k1_packed_direct_forward_impl( return _layout_like(runner.final.detach(), x), runner_token -def _k1_backward_from_runner_current_device( +def _so2_backward_from_runner_current_device( handle: int, grad_out: Tensor, x: Tensor, @@ -1499,7 +1520,7 @@ def _k1_backward_from_runner_current_device( ) -def _k1_backward_from_runner( +def _so2_backward_from_runner( handle: int, grad_out: Tensor, x: Tensor, @@ -1512,7 +1533,7 @@ def _k1_backward_from_runner( ) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: """Run all compilation and launches on the operand's CUDA device.""" with torch.cuda.device(grad_out.device): - return _k1_backward_from_runner_current_device( + return _so2_backward_from_runner_current_device( handle, grad_out, x, @@ -1525,7 +1546,7 @@ def _k1_backward_from_runner( ) -def _k1_packed_direct_backward_impl( +def _so2_packed_direct_backward_impl( handle: int, grad_out: Tensor, x: Tensor, @@ -1558,7 +1579,7 @@ def _k1_packed_direct_backward_impl( source_ptr, edge_src_gate, ) - return _k1_backward_from_runner( + return _so2_backward_from_runner( handle, grad_out, x, @@ -1582,10 +1603,10 @@ def _runner_backward_manual( n_edge = runner.edge_count grad_phase_c_out, grad_x_wide_down = _final_manual_backward(runner, grad_out) - # Keep the large K1 slabs out of the readout-backward allocation crest. + # Keep the large SO2 slabs out of the readout-backward allocation crest. runner.ensure_backward_workspace() - from .k1_kernels.cute_neo_phase_c_backward_layout_runner import ( + from .phase_c import ( NeoPhaseCBackwardLayoutOutputs, ) @@ -1638,7 +1659,7 @@ def _runner_backward_manual( grad_mixed_focus = grad_mixed_focus.contiguous() x_wide_flat = runner.x_wide.detach().contiguous().view(runner.node_count, 16 * 64) radial_state = runner.radial_compact - from .k1_radial_phase_a_node import ( + from .radial_phase_a import ( run_neo_radial_phase_a_backward_node_tiled, ) @@ -1678,16 +1699,16 @@ def _stateful_custom_op_tags() -> tuple[Any, ...] | None: return (cudagraph_unsafe,) -_K1_CUSTOM_OP_TAGS = _stateful_custom_op_tags() -_k1_packed_direct_op = torch.library.custom_op( - "sezm_cute::k1_packed_direct", mutates_args=(), tags=_K1_CUSTOM_OP_TAGS -)(_k1_packed_direct_forward_impl) -_k1_packed_direct_bwd_op = torch.library.custom_op( - "sezm_cute::k1_packed_direct_bwd", mutates_args=(), tags=_K1_CUSTOM_OP_TAGS -)(_k1_packed_direct_backward_impl) +_SO2_CUSTOM_OP_TAGS = _stateful_custom_op_tags() +_so2_packed_direct_op = torch.library.custom_op( + "sezm_cute::so2_packed_direct", mutates_args=(), tags=_SO2_CUSTOM_OP_TAGS +)(_so2_packed_direct_forward_impl) +_so2_packed_direct_bwd_op = torch.library.custom_op( + "sezm_cute::so2_packed_direct_bwd", mutates_args=(), tags=_SO2_CUSTOM_OP_TAGS +)(_so2_packed_direct_backward_impl) -@_k1_packed_direct_op.register_fake +@_so2_packed_direct_op.register_fake def _( handle: int, x: Tensor, @@ -1711,7 +1732,7 @@ def _( ) -@_k1_packed_direct_bwd_op.register_fake +@_so2_packed_direct_bwd_op.register_fake def _( handle: int, grad_out: Tensor, @@ -1739,7 +1760,7 @@ def _( ) -def _k1_packed_direct_setup_context( +def _so2_packed_direct_setup_context( ctx: Any, inputs: tuple[Any, ...], output: tuple[Tensor, Tensor], @@ -1776,7 +1797,7 @@ def _k1_packed_direct_setup_context( ) -def _k1_packed_direct_registered_backward_impl( +def _so2_packed_direct_registered_backward_impl( ctx: Any, grad_out: Tensor, ) -> tuple[Any, ...]: @@ -1794,7 +1815,7 @@ def _k1_packed_direct_registered_backward_impl( edge_src_gate, runner_token, ) = ctx.saved_tensors - grad_x, grad_d, grad_dt, grad_radial, grad_edge_env = _k1_packed_direct_bwd_op( + grad_x, grad_d, grad_dt, grad_radial, grad_edge_env = _so2_packed_direct_bwd_op( ctx.handle, grad_out, x, @@ -1826,23 +1847,23 @@ def _k1_packed_direct_registered_backward_impl( ) -def _k1_packed_direct_backward( +def _so2_packed_direct_backward( ctx: Any, grad_out: Tensor, grad_runner_token: Tensor | None, ) -> tuple[Any, ...]: """Run packed-direct backward with its custom op visible to compilation.""" del grad_runner_token - return _k1_packed_direct_registered_backward_impl(ctx, grad_out) + return _so2_packed_direct_registered_backward_impl(ctx, grad_out) -_k1_packed_direct_op.register_autograd( - _k1_packed_direct_backward, - setup_context=_k1_packed_direct_setup_context, +_so2_packed_direct_op.register_autograd( + _so2_packed_direct_backward, + setup_context=_so2_packed_direct_setup_context, ) -def _cute_k1_impl( +def _cute_so2_impl( handle: int, x: Tensor, d_full: Tensor, @@ -1856,7 +1877,7 @@ def _cute_k1_impl( source_ptr: Tensor, edge_src_gate: Tensor, ) -> Tensor: - output, _runner_token = _k1_packed_direct_op( + output, _runner_token = _so2_packed_direct_op( int(handle), x, d_full, @@ -1873,7 +1894,7 @@ def _cute_k1_impl( return output -def cute_k1( +def cute_so2( handle: int, x: Tensor, d_full: Tensor, @@ -1887,12 +1908,12 @@ def cute_k1( source_order: Tensor | None = None, source_ptr: Tensor | None = None, ) -> Tensor: - """Run K1 through its registered custom-op boundary.""" + """Run SO2 through its registered custom-op boundary.""" if source_order is None: source_order = src.new_empty((0,), dtype=torch.int32) if source_ptr is None: source_ptr = src.new_empty((0,), dtype=torch.int32) - return _cute_k1_impl( + return _cute_so2_impl( handle, x, d_full, @@ -1920,7 +1941,7 @@ def _dst_ptr_from_sorted( if runtime_policy.is_cute_strict_enabled() and dst.numel() > 1: torch_module._assert_async( torch_module.all(dst[1:] >= dst[:-1]), - "Neo K1 destinations_sorted=True requires monotonically " + "Neo SO2 destinations_sorted=True requires monotonically " "nondecreasing destination indices", ) boundaries = torch_module.arange( @@ -1970,8 +1991,8 @@ def _cuda_compute_capability(device_index: int) -> tuple[int, int]: def _architecture_default_config( compute_capability: tuple[int, int], -) -> NeoK1RuntimeConfig: - return NeoK1RuntimeConfig( +) -> NeoSO2RuntimeConfig: + return NeoSO2RuntimeConfig( native_sm90_path=compute_capability == runtime_policy.SM90_CAPABILITY, per_focus_so2_fwd_pair=( compute_capability in runtime_policy.SM80_PROFILE_CAPABILITIES @@ -1982,7 +2003,7 @@ def _architecture_default_config( ) -def _maybe_run_prepared_cute_k1( +def _maybe_run_prepared_cute_so2( block: Any, x: Tensor, edge_cache: Any, @@ -1991,9 +2012,9 @@ def _maybe_run_prepared_cute_k1( source_order: Tensor | None = None, source_ptr: Tensor | None = None, ) -> Tensor | None: - """Dispatch prevalidated packed K1 state without a Python graph break.""" - state = getattr(block, "_deepmd_cute_k1_state", None) - if not isinstance(state, _RegisteredK1State): + """Dispatch prevalidated packed SO2 state without a Python graph break.""" + state = getattr(block, "_deepmd_cute_so2_state", None) + if not isinstance(state, _RegisteredSO2State): return None with _REGISTRY_LOCK: entry = _REGISTRY.get(state.handle) @@ -2028,7 +2049,7 @@ def _maybe_run_prepared_cute_k1( ): return None edge_count = edge_cache.src.numel() - if not runtime_policy.k1_int32_indexing_is_safe( + if not runtime_policy.so2_int32_indexing_is_safe( edge_count=edge_count, node_count=x.shape[0], ): @@ -2057,7 +2078,7 @@ def _maybe_run_prepared_cute_k1( # canonicalization in ``_build_runner``. Keep this wrapper graph-visible. empty_edge_src_gate = edge_cache.edge_env.new_empty((0,)) d_arg = d_full.contiguous() - output = cute_k1( + output = cute_so2( state.handle, x.contiguous(), d_arg, @@ -2075,7 +2096,7 @@ def _maybe_run_prepared_cute_k1( @torch.compiler.disable -def _maybe_run_cute_k1_fallback( +def _maybe_run_cute_so2_fallback( block: Any, x: Tensor, edge_cache: Any, @@ -2084,7 +2105,7 @@ def _maybe_run_cute_k1_fallback( source_order: Tensor | None = None, source_ptr: Tensor | None = None, ) -> Tensor | None: - """Run the opt-in Neo CuTe K1 path, or return ``None`` for fallback.""" + """Run the opt-in Neo CuTe SO2 path, or return ``None`` for fallback.""" if edge_cache.D_packed is None: return None # The optimized backward does not expose the differentiable SFPG @@ -2092,7 +2113,7 @@ def _maybe_run_cute_k1_fallback( if edge_cache.edge_src_gate is not None: return None destinations_sorted = bool(getattr(edge_cache, "destinations_sorted", False)) - if not is_neo_k1_runtime_eligible( + if not is_neo_so2_runtime_eligible( block, training=bool(block.training), device=x.device, @@ -2111,19 +2132,19 @@ def _maybe_run_cute_k1_fallback( ): return None - state = getattr(block, "_deepmd_cute_k1_state", None) + state = getattr(block, "_deepmd_cute_so2_state", None) if state is False: return None - if state is not None and not isinstance(state, _RegisteredK1State): + if state is not None and not isinstance(state, _RegisteredSO2State): state = None device_index = x.device.index if device_index is None: device_index = torch.cuda.current_device() compute_capability = _cuda_compute_capability(device_index) - if not is_supported_k1_compute_capability(compute_capability): + if not is_supported_so2_compute_capability(compute_capability): return None config = _architecture_default_config(compute_capability) - if not runtime_policy.k1_int32_indexing_is_safe( + if not runtime_policy.so2_int32_indexing_is_safe( edge_count=edge_cache.src.numel(), node_count=x.shape[0], ): @@ -2131,7 +2152,7 @@ def _maybe_run_cute_k1_fallback( if edge_cache.D_packed.dim() != 2: return None if state is None or state.device_index != device_index or state.config != config: - state = _register_cute_k1_state(block, device_index, config) + state = _register_cute_so2_state(block, device_index, config) if state is None: return None @@ -2169,7 +2190,7 @@ def _maybe_run_cute_k1_fallback( if edge_src_gate is None: edge_src_gate = edge_cache.edge_env.new_empty((0,)) edge_src_gate_arg = _aligned_contiguous(edge_src_gate) - output = cute_k1( + output = cute_so2( state.handle, x_arg, d_arg, @@ -2186,7 +2207,7 @@ def _maybe_run_cute_k1_fallback( return _layout_like(output, x) -def maybe_run_cute_k1( +def maybe_run_cute_so2( block: Any, x: Tensor, edge_cache: Any, @@ -2196,13 +2217,13 @@ def maybe_run_cute_k1( source_ptr: Tensor | None = None, ) -> Tensor | None: """Use prevalidated opaque dispatch, or the conservative eager fallback.""" - state = getattr(block, "_deepmd_cute_k1_state", None) + state = getattr(block, "_deepmd_cute_so2_state", None) use_prepared = isinstance( state, - _RegisteredK1State, - ) or runtime_policy.is_k1_thin_wrapper_enabled(_tensor_compute_capability(x)) + _RegisteredSO2State, + ) or runtime_policy.is_so2_thin_wrapper_enabled(_tensor_compute_capability(x)) if use_prepared: - output = _maybe_run_prepared_cute_k1( + output = _maybe_run_prepared_cute_so2( block, x, edge_cache, @@ -2213,7 +2234,7 @@ def maybe_run_cute_k1( ) if output is not None: return output - return _maybe_run_cute_k1_fallback( + return _maybe_run_cute_so2_fallback( block, x, edge_cache, diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_backward_layout_runner.py b/deepmd/pt_expt/kernels/cute/sezm/so2/phase_c.py similarity index 98% rename from deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_backward_layout_runner.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/phase_c.py index f33059e9be..2d4ed1a104 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_kernels/cute_neo_phase_c_backward_layout_runner.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/phase_c.py @@ -18,8 +18,7 @@ from ..compile_cache import ( device_aware_lru_cache, ) -from ..k1_wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES -from .cute_neo_phase_c_backward_layout import ( +from .kernels.phase_c_backward import ( DEGREE_COUNT, FOCUS_CHANNELS, HIDDEN, @@ -27,6 +26,7 @@ REDUCED_COUNT, compile_neo_phase_c_backward_layout, ) +from .wigner_layout import PACKED_VALUE_COUNT as PACKED_WIGNER_VALUES REQUIRED_ALIGNMENT = 16 diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_radial_phase_a_node.py b/deepmd/pt_expt/kernels/cute/sezm/so2/radial_phase_a.py similarity index 98% rename from deepmd/pt_expt/kernels/cute/sezm/k1_radial_phase_a_node.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/radial_phase_a.py index ac968839ba..4086bbc12f 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_radial_phase_a_node.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/radial_phase_a.py @@ -14,10 +14,10 @@ Any, ) -from .compile_cache import ( +from ..compile_cache import ( device_aware_lru_cache, ) -from .k1_wigner_layout import ( +from .wigner_layout import ( PACKED_VALUE_COUNT, ) @@ -103,7 +103,7 @@ def build_source_csr( @device_aware_lru_cache(maxsize=4) def _compile_node_tiled() -> Any: - from .k1_kernels.cute_neo_radial_phase_a_backward_node import ( + from .kernels.radial_phase_a_backward import ( compile_neo_radial_phase_a_backward_node_tiled, ) diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_runner.py b/deepmd/pt_expt/kernels/cute/sezm/so2/runner.py similarity index 94% rename from deepmd/pt_expt/kernels/cute/sezm/k1_runner.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/runner.py index af793d5a5d..321c6fe871 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_runner.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/runner.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""Runtime runner for the Neo CuTe K1 SO2 unit.""" +"""Runtime runner for the Neo CuTe SO2 unit.""" from __future__ import ( annotations, @@ -13,11 +13,11 @@ Any, ) -from .runtime_policy import ( +from ..runtime_policy import ( FUSED_SO2_GATE_CAPABILITIES, SM80_PROFILE_CAPABILITIES, SM90_CAPABILITY, - SUPPORTED_K1_CAPABILITIES, + SUPPORTED_SO2_CAPABILITIES, ) @@ -34,10 +34,10 @@ 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 supported compute capability") + if compute_capability not in SUPPORTED_SO2_CAPABILITIES: + raise RuntimeError("Neo SO2 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") + raise RuntimeError("the native SM90 SO2 path must be selected only on sm_90") if runtime_config.per_focus_so2_fwd_pair != ( compute_capability in SM80_PROFILE_CAPABILITIES ): @@ -125,7 +125,7 @@ def _batched_radial_projection_weight( if cache is not None and cache[0] == cache_key: return cache[1] - from .k1_radial_phase_a_node import ( + from .radial_phase_a import ( prepare_batched_radial_projection_weight, ) @@ -164,7 +164,7 @@ def __post_init__(self) -> None: self.logits = self.logits.detach() -class _NeoK1BackwardWorkspaceBase: +class _NeoSO2BackwardWorkspaceBase: _EXPORTED_NAMES = ( "grad_stack_focus", "grad_focus_alpha", @@ -256,8 +256,8 @@ def _structural_memory_views( ) -class NeoK1BackwardWorkspace(_NeoK1BackwardWorkspaceBase): - """Lazily allocated K1 backward scratch with lifetime-based slab reuse.""" +class NeoSO2BackwardWorkspace(_NeoSO2BackwardWorkspaceBase): + """Lazily allocated SO2 backward scratch with lifetime-based slab reuse.""" def __init__( self, @@ -282,7 +282,7 @@ def __init__( for size in radial.shape[1:]: radial_values_per_edge *= size if d_values_per_edge > 10 * 64: - raise ValueError("Neo K1 grad_D does not fit the Phase-C scratch slab") + raise ValueError("Neo SO2 grad_D does not fit the Phase-C scratch slab") if phase_c_single_input_reuse and not structural_memory_reuse: raise ValueError( @@ -404,16 +404,16 @@ def __init__( *, runtime_config: Any, ) -> None: - from .k1_kernels.cute_envelope_gated_softmax import ( + from .kernels.envelope_softmax import ( compile_envelope_softmax_forward, ) - from .k1_kernels.cute_neo_phase_a_radial_forward import ( + from .kernels.phase_a_radial_forward import ( run_neo_phase_a_radial_forward_packed_direct, ) - from .k1_so2linear import ( + from .linear import ( run_neo_so2_linear_manual, ) - from .k1_wigner_layout import ( + from .wigner_layout import ( PACKED_VALUE_COUNT, ) @@ -446,10 +446,10 @@ def __init__( or tuple(dt_full.shape) != expected_shape ): raise ValueError( - f"packed Neo K1 Wigner tensors must have shape {expected_shape}" + f"packed Neo SO2 Wigner tensors must have shape {expected_shape}" ) if d_full.data_ptr() != dt_full.data_ptr(): - raise ValueError("packed Neo K1 D and Dt must share one storage") + raise ValueError("packed Neo SO2 D and Dt must share one storage") self.radial = radial_feat self.dst_ptr_i32 = dst_ptr.to(device=x.device, dtype=torch.int32).contiguous() self.edge_count = record.edge_cache.src.numel() @@ -471,7 +471,7 @@ def __init__( dtype=torch.int32, ).contiguous() else: - from .k1_radial_phase_a_node import ( + from .radial_phase_a import ( build_source_csr, ) @@ -480,7 +480,7 @@ def __init__( self.source_ptr_i32 = source_csr.source_ptr self.rotate = self.so2.rotate_inv_rescale_full.contiguous() self.edge_gate = _edge_gate(torch, record.edge_cache) - from .k1_kernels.cute_neo_focus_src_backward import ( + from .kernels.focus_source_backward import ( compile_neo_attention_prelude_forward, ) @@ -492,7 +492,7 @@ def __init__( float(self.so2.focus_label_smoothing), self.compile_identity, ) - from .k1_kernels.cute_neo_qk_edge import ( + from .kernels.qk_edge import ( compile_neo_qk_edge_backward, compile_neo_qk_edge_forward, ) @@ -507,7 +507,7 @@ def __init__( qk_scale, self.compile_identity, ) - from .k1_kernels.cute_neo_qk_edge import ( + from .kernels.qk_edge import ( compile_neo_qk_node_input_adjoint, ) @@ -516,7 +516,7 @@ def __init__( float(self.so2.attn_qk_norm.eps), self.compile_identity, ) - from .k1_kernels.cute_neo_phase_c_backward_layout_runner import ( + from .phase_c import ( CuteNeoPhaseCBackwardLayout, ) @@ -538,7 +538,7 @@ def __init__( self._run_structural_gate_forward = None self._run_structural_gate_backward = None if runtime_config.combined_so2_gate: - from .k1_kernels.cute_neo_gate_linear_residual_backward_fused import ( + from .kernels.gate_linear_residual_backward import ( compile_neo_gate_linear_residual_backward_fused, ) @@ -547,16 +547,16 @@ def __init__( compile_neo_gate_linear_residual_backward_fused() ) elif not runtime_config.native_sm90_path: - from .k1_gate_structural import ( + from .kernels.structural_gate_sm80 import ( + compile_neo_gate_split_structural_vec4_sm80_backward, + compile_neo_gate_split_structural_vec4_sm80_forward, + ) + from .structural_gate import ( focus_major_gate_linear_backward_add_, focus_major_gate_linear_forward, run_structural_gate_backward, run_structural_gate_forward, ) - from .k1_kernels.cute_neo_gate_split_structural_vec4_sm80 import ( - compile_neo_gate_split_structural_vec4_sm80_backward, - compile_neo_gate_split_structural_vec4_sm80_forward, - ) self.structural_gate_forward = ( compile_neo_gate_split_structural_vec4_sm80_forward( @@ -609,7 +609,7 @@ def __init__( def ensure_backward_workspace( self, - ) -> NeoK1BackwardWorkspace: + ) -> NeoSO2BackwardWorkspace: if self._backward_workspace is None: radial_scratch = getattr(self, "structural_scratch", None) if radial_scratch is None: @@ -623,14 +623,14 @@ def ensure_backward_workspace( ) if radial_scratch is None: raise RuntimeError( - "Neo K1 backward requires a reusable radial scratch buffer; " + "Neo SO2 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) self.phase_c_stack.record_stream(stream) radial_scratch.record_stream(stream) - workspace = NeoK1BackwardWorkspace( + workspace = NeoSO2BackwardWorkspace( self.torch, edge_count=self.edge_count, node_count=self.node_count, @@ -679,11 +679,11 @@ def _build_forward_graph(self) -> None: self.stack_caches: list[StackCache] = [] if self.config.combined_so2_gate: - from .k1_kernels.cute_neo_so2_gate_combined_fwd import ( + from .kernels.combined_gate_forward import ( CuteNeoSO2GateCombinedFwdRunner, prepare_neo_so2_gate_combined_weights, ) - from .k1_so2linear import ( + from .linear import ( cached_neo_so2_linear_weights, ) @@ -853,7 +853,7 @@ def _build_forward_graph(self) -> None: ) self.attn_logits = None x_wide_down = self.x_wide.detach() - from .k1_kernels.cute_neo_phase_c_onepass import ( + from .kernels.phase_c_forward import ( run_neo_phase_c_onepass_output_gate, ) @@ -880,7 +880,7 @@ def _build_forward_graph(self) -> None: self.message_grid_product = None if so2.message_node_grid_product is not None: if self.packed_message_grid: - from .k1_message_grid_packed import ( + from .message_grid import ( run_packed_message_grid_forward, ) diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/__init__.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/__init__.py new file mode 100644 index 0000000000..74481d7e7e --- /dev/null +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Native SM90 split-complex Neo SO(2) implementation.""" diff --git a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/final_phase_c.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/final_phase_c.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/sm90_k1/final_phase_c.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/sm90/final_phase_c.py index 6475566c30..6b3c769aad 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/final_phase_c.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/final_phase_c.py @@ -34,7 +34,7 @@ make_fake_stream, ) -from ..compile_cache import ( +from ...compile_cache import ( device_aware_lru_cache, ) @@ -430,7 +430,7 @@ def run_direct_statistics_forward( """Build final node statistics directly and apply node-scale SO2Linear. ``edge_chunk`` and ``chunk_slots`` retain the chunked-forward argument - contract used by the SM90 K1 runner. + contract used by the SM90 SO2 runner. Only the 64-edge internal schedule is supported; ``chunk_slots`` is not an allocation dimension in this implementation. """ diff --git a/deepmd/pt_expt/kernels/cute/sezm/message_grid_gaunt_sm90.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/message_grid_gaunt.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/message_grid_gaunt_sm90.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/sm90/message_grid_gaunt.py index 2bbc087ac8..978bbed5d1 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/message_grid_gaunt_sm90.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/message_grid_gaunt.py @@ -25,7 +25,7 @@ make_fake_stream, ) -from .compile_cache import ( +from ...compile_cache import ( device_aware_lru_cache, ) diff --git a/deepmd/pt_expt/kernels/cute/sezm/message_grid_readout_sm90.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/message_grid_readout.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/message_grid_readout_sm90.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/sm90/message_grid_readout.py index 1ed8b0d680..abf3cb2774 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/message_grid_readout_sm90.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/message_grid_readout.py @@ -26,17 +26,17 @@ make_fake_stream, ) -from .compile_cache import ( +from ...compile_cache import ( device_aware_lru_cache, ) -from .message_grid_gaunt_sm90 import ( - Sm90GauntSchedule, - build_sm90_gaunt_schedule, -) -from .runtime_policy import ( +from ...runtime_policy import ( SM90_CAPABILITY, uses_strict_fp32_matmul, ) +from .message_grid_gaunt import ( + Sm90GauntSchedule, + build_sm90_gaunt_schedule, +) if TYPE_CHECKING: from collections.abc import ( @@ -751,7 +751,7 @@ def run_sm90_message_grid_backward( state: Sm90MessageGridState, ) -> tuple[torch.Tensor, torch.Tensor]: """Return query/context adjoints without expanded global adjoint slabs.""" - from .k1_message_grid_packed import ( + from ..message_grid import ( _as_flat_ndfc, _focus_linear_backward_input, _frame_expand_packed, diff --git a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/output_gate.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/output_gate.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/sm90_k1/output_gate.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/sm90/output_gate.py index 5ba25e3f41..4020acb4a6 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/output_gate.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/output_gate.py @@ -22,7 +22,7 @@ make_fake_stream, ) -from ..compile_cache import ( +from ...compile_cache import ( device_aware_lru_cache, ) diff --git a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/persistent.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/persistent.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/sm90_k1/persistent.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/sm90/persistent.py index 3c1a00641c..921d403399 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/persistent.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/persistent.py @@ -32,7 +32,7 @@ Tensor, ) -from ..compile_cache import ( +from ...compile_cache import ( device_aware_lru_cache, ) diff --git a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_a.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_a.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_a.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_a.py index e0f3b481a7..5341d1f9fb 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_a.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_a.py @@ -35,10 +35,10 @@ make_fake_stream, ) -from ..compile_cache import ( +from ...compile_cache import ( device_aware_lru_cache, ) -from ..k1_wigner_layout import ( +from ..wigner_layout import ( PACKED_VALUE_COUNT, ) from .persistent import ( diff --git a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_a_backward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_a_backward.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_a_backward.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_a_backward.py index 85a980799d..b351789950 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_a_backward.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_a_backward.py @@ -40,10 +40,10 @@ Tensor, ) -from ..compile_cache import ( +from ...compile_cache import ( device_aware_lru_cache, ) -from ..k1_kernels.cute_neo_radial_phase_a_backward_node import ( +from ..kernels.radial_phase_a_backward import ( COMPACT_WIDTH, DEGREE_COUNT, FOCUS_COUNT, diff --git a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_c_attention_backward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_c_attention_backward.py similarity index 99% rename from deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_c_attention_backward.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_c_attention_backward.py index f1f0430357..083c24a7cd 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/phase_c_attention_backward.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_c_attention_backward.py @@ -37,7 +37,7 @@ make_fake_stream, ) -from ..compile_cache import ( +from ...compile_cache import ( device_aware_lru_cache, ) diff --git a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/prefix.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/prefix.py similarity index 100% rename from deepmd/pt_expt/kernels/cute/sezm/sm90_k1/prefix.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/sm90/prefix.py diff --git a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/radial.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/radial.py similarity index 97% rename from deepmd/pt_expt/kernels/cute/sezm/sm90_k1/radial.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/sm90/radial.py index cb7cb017f7..d8c256f474 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/radial.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/radial.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""Compact radial state for the SM90 persistent-complex Neo K1 path. +"""Compact radial state for the SM90 persistent-complex Neo SO2 path. The fused Phase-A kernel computes three logically separate values in one edge CTA: the packed-Wigner rotation, the 25-value compact radial map, and the @@ -32,7 +32,7 @@ make_fake_stream, ) -from ..compile_cache import ( +from ...compile_cache import ( device_aware_lru_cache, ) @@ -205,9 +205,9 @@ def run_neo_radial_state_forward_fp32( device, ) if edge_count <= 0: - raise ValueError("persistent-complex K1 requires E > 0") + raise ValueError("persistent-complex SO2 requires E > 0") if tuple(torch.cuda.get_device_capability(device)) != (9, 0): - raise RuntimeError("persistent-complex K1 requires SM90") + raise RuntimeError("persistent-complex SO2 requires SM90") if torch.backends.cuda.matmul.allow_tf32: raise RuntimeError("strict FP32 requires allow_tf32=False") if torch.get_float32_matmul_precision() != "highest": diff --git a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/runner.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/runner.py similarity index 94% rename from deepmd/pt_expt/kernels/cute/sezm/sm90_k1/runner.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/sm90/runner.py index 1f8f37339e..80b78519cb 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/sm90_k1/runner.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/runner.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""Split-complex strict-FP32 SM90 implementation of the complete Neo K1. +"""Split-complex strict-FP32 SM90 implementation of the complete Neo SO2. The value path uses one split-complex edge representation throughout: @@ -30,28 +30,28 @@ Tensor, ) -from ..k1 import ( +from ..linear import ( + cached_neo_so2_linear_weights, +) +from ..operation import ( _compile_output_gate_backward, _equivariant_rmsnorm_backward, _runner_compile_identity, _so3_linear_backward_input, _x_wide_manual_backward, ) -from ..k1_runner import ( +from ..runner import ( NeoFullCuteBackward, ) -from ..k1_so2linear import ( - cached_neo_so2_linear_weights, -) -from ..message_grid_readout_sm90 import ( - prepare_sm90_message_grid_state, - run_sm90_message_grid_backward, -) from .final_phase_c import ( ExpandedFinalWeights, prepare_expanded_final_weights, run_direct_statistics_forward, ) +from .message_grid_readout import ( + prepare_sm90_message_grid_state, + run_sm90_message_grid_backward, +) from .output_gate import ( run_chunked_final_output_gate, ) @@ -107,21 +107,21 @@ def _prepare_persistent_weights(so2: Any) -> NeoPersistentComplexWeights: linears = tuple(so2.so2_linears) nonlinearities = tuple(so2.non_linearities) if len(linears) != 3 or len(nonlinearities) != 3: - raise NotImplementedError("SM90 K1 requires three SO2 layers") + raise NotImplementedError("SM90 SO2 requires three SO2 layers") if any(type(norm).__name__ != "Identity" for norm in so2.so2_inter_norms): - raise NotImplementedError("SM90 K1 requires disabled inter-layer norms") + raise NotImplementedError("SM90 SO2 requires disabled inter-layer norms") for nonlinearity in nonlinearities[:GATED_LAYERS]: if getattr(nonlinearity, "layout", None) != "fndc": - raise NotImplementedError("SM90 K1 requires fndc gates") + raise NotImplementedError("SM90 SO2 requires fndc gates") activation = getattr( getattr(nonlinearity, "scalar_act", None), "activation", None, ) if str(activation).lower() != "silu": - raise NotImplementedError("SM90 K1 requires SiLU gates") + raise NotImplementedError("SM90 SO2 requires SiLU gates") if getattr(nonlinearity.gate_linear, "bias", None) is not None: - raise NotImplementedError("SM90 K1 does not support gate bias") + raise NotImplementedError("SM90 SO2 does not support gate bias") sources = tuple( tensor @@ -169,9 +169,9 @@ def _prepare_final_weights(so2: Any) -> ExpandedFinalWeights: w0, wp = cached_neo_so2_linear_weights(linear) if tuple(w0.shape) != (FOCUS_COUNT, M0_WIDTH, M0_WIDTH): - raise ValueError("SM90 K1 requires final m=0 weights (2,128,128)") + raise ValueError("SM90 SO2 requires final m=0 weights (2,128,128)") if tuple(wp.shape) != (FOCUS_COUNT, 2 * M1_WIDTH, 2 * M1_WIDTH): - raise ValueError("SM90 K1 requires final pair weights (2,192,192)") + raise ValueError("SM90 SO2 requires final pair weights (2,192,192)") u = wp[:, :M1_WIDTH, :M1_WIDTH] v = wp[:, :M1_WIDTH, M1_WIDTH:] if not torch.equal(wp[:, M1_WIDTH:, :M1_WIDTH], -v): @@ -189,7 +189,7 @@ def _prepare_final_weights(so2: Any) -> ExpandedFinalWeights: return weights -__all__ = ["NeoSm90K1Runner"] +__all__ = ["NeoSm90SO2Runner"] @dataclass(frozen=True) @@ -246,7 +246,7 @@ def _run_final_reverse_panels( def _qk_node_input_adjoint( - runner: NeoSm90K1Runner, + runner: NeoSm90SO2Runner, grad_q_node: Tensor, grad_k_node: Tensor, ) -> Tensor: @@ -273,8 +273,8 @@ def _qk_node_input_adjoint( return grad_x_wide -class NeoSm90K1Runner(NeoFullCuteBackward): - """Complete Neo K1 runner with one native split representation.""" +class NeoSm90SO2Runner(NeoFullCuteBackward): + """Complete Neo SO2 runner with one native split representation.""" uses_native_sm90_path = True @@ -285,7 +285,7 @@ def _build_forward_graph(self) -> None: node_count = self.node_count edge_count = self.edge_count if self.compute_capability != (9, 0): - raise RuntimeError("split-complex K1 requires SM90") + raise RuntimeError("split-complex SO2 requires SM90") self.structural_scratch = None self.use_full_node = block.node_lmax == block.lmax @@ -436,7 +436,7 @@ def _build_forward_graph(self) -> None: self.message_grid_sm90_state = None if so2.message_node_grid_product is not None: if self.packed_message_grid: - from ..k1_message_grid_packed import ( + from ..message_grid import ( run_packed_message_grid_forward, ) @@ -467,21 +467,21 @@ def _build_forward_graph(self) -> None: self.final = final def input_adjoint(self, grad_out: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]: - """Return K1 input adjoints through the split-complex reverse path.""" + """Return SO2 input adjoints through the split-complex reverse path.""" return _runner_backward(self, grad_out) def _final_manual_backward_sm90( - runner: NeoSm90K1Runner, + runner: NeoSm90SO2Runner, grad_out: Tensor, ) -> tuple[Tensor, Tensor]: - """Use the one-slab tiled message-grid adjoint in the final K1 boundary.""" + """Use the one-slab tiled message-grid adjoint in the final SO2 boundary.""" if not ( runner.packed_message_grid and runner.so2.message_node_grid_product is not None and runner.message_grid_sm90_state is not None ): - raise RuntimeError("SM90 K1 requires the packed message-grid path") + raise RuntimeError("SM90 SO2 requires the packed message-grid path") so2 = runner.so2 block = runner.block @@ -546,7 +546,7 @@ def _final_manual_backward_sm90( def _runner_backward( - runner: NeoSm90K1Runner, + runner: NeoSm90SO2Runner, grad_out: Tensor, ) -> tuple[Tensor, Tensor, Tensor, Tensor]: so2 = runner.so2 diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_gate_structural.py b/deepmd/pt_expt/kernels/cute/sezm/so2/structural_gate.py similarity index 98% rename from deepmd/pt_expt/kernels/cute/sezm/k1_gate_structural.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/structural_gate.py index caafbf17ff..4ea1c7cda9 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_gate_structural.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/structural_gate.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""PyTorch/cuBLAS glue for the structural Neo K1 split-gate path. +"""PyTorch/cuBLAS glue for the structural Neo SO2 split-gate path. The two focus panels stay contiguous across the gate-linear boundary. Forward uses direct cuBLAS matrix products, while backward accumulates into the scalar diff --git a/deepmd/pt_expt/kernels/cute/sezm/k1_wigner_layout.py b/deepmd/pt_expt/kernels/cute/sezm/so2/wigner_layout.py similarity index 97% rename from deepmd/pt_expt/kernels/cute/sezm/k1_wigner_layout.py rename to deepmd/pt_expt/kernels/cute/sezm/so2/wigner_layout.py index 0283f23a2f..ff26fa0cc9 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k1_wigner_layout.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/wigner_layout.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""Fixed packed Wigner layout for the Neo ``lmax=3, mmax=1`` K1 path. +"""Fixed packed Wigner layout for the Neo ``lmax=3, mmax=1`` SO2 path. The panel stores only rows selected by ``coeff_index_m`` and only columns from the matching Wigner block. Phase A reads ``D[coeff, degree]`` while Phase C diff --git a/deepmd/pt_expt/kernels/cute/sezm/k4_wignerd.py b/deepmd/pt_expt/kernels/cute/sezm/wignerd.py similarity index 97% rename from deepmd/pt_expt/kernels/cute/sezm/k4_wignerd.py rename to deepmd/pt_expt/kernels/cute/sezm/wignerd.py index b02e66fa52..cb70a10573 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/k4_wignerd.py +++ b/deepmd/pt_expt/kernels/cute/sezm/wignerd.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later # pyright: reportMissingImports=false # ruff: noqa: ANN001, ANN201, ANN202, ANN204, TC002, UP035 -"""Packed CuTe Wigner-D panel for the Neo K1 inference path.""" +"""Packed CuTe Wigner-D panel for the Neo SO2 inference path.""" from __future__ import ( annotations, @@ -34,7 +34,7 @@ from .compile_cache import ( device_aware_lru_cache, ) -from .k1_wigner_layout import PACKED_VALUE_COUNT as K1_PANEL_VALUES +from .so2.wigner_layout import PACKED_VALUE_COUNT as SO2_PANEL_VALUES L2_SPARSE_TERMS = 10 L3_SPARSE_TERMS = 20 @@ -691,7 +691,7 @@ def compile_wignerd_panel_forward( fake_q = make_fake_compact_tensor(cutlass.Float32, (e, 4), stride_order=(1, 0)) fake_panel = make_fake_compact_tensor( cutlass.Float32, - (e, K1_PANEL_VALUES), + (e, SO2_PANEL_VALUES), stride_order=(1, 0), ) fake_exp_l2 = make_fake_compact_tensor(cutlass.Int32, (35, 4), stride_order=(1, 0)) @@ -742,7 +742,7 @@ def compile_wignerd_panel_backward( fake_q = make_fake_compact_tensor(cutlass.Float32, (e, 4), stride_order=(1, 0)) fake_grad_panel = make_fake_compact_tensor( cutlass.Float32, - (e, K1_PANEL_VALUES), + (e, SO2_PANEL_VALUES), stride_order=(1, 0), ) fake_grad_q = make_fake_compact_tensor(cutlass.Float32, (e, 4), stride_order=(1, 0)) @@ -824,7 +824,7 @@ def _wignerd_panel_impl(edge_quat: torch.Tensor) -> torch.Tensor: panel = torch.empty( q.shape[0], - K1_PANEL_VALUES, + SO2_PANEL_VALUES, device=q.device, dtype=q.dtype, ) @@ -859,9 +859,9 @@ def _wignerd_panel_bwd_impl( edge_quat: torch.Tensor, ) -> torch.Tensor: q = edge_quat.detach().contiguous() - if tuple(grad_panel.shape) != (q.shape[0], K1_PANEL_VALUES): + if tuple(grad_panel.shape) != (q.shape[0], SO2_PANEL_VALUES): raise ValueError( - f"packed Wigner gradient must have shape ({q.shape[0]}, {K1_PANEL_VALUES})" + f"packed Wigner gradient must have shape ({q.shape[0]}, {SO2_PANEL_VALUES})" ) if q.dtype != torch.float32: raise TypeError(f"packed WignerD requires float32, got {q.dtype}") @@ -902,18 +902,18 @@ def _stateful_custom_op_tags() -> tuple[Any, ...] | None: return (cudagraph_unsafe,) -_K4_CUSTOM_OP_TAGS = _stateful_custom_op_tags() +_WIGNERD_CUSTOM_OP_TAGS = _stateful_custom_op_tags() _wignerd_panel_op = torch.library.custom_op( - "sezm_cute::wignerd_k1_panel", mutates_args=(), tags=_K4_CUSTOM_OP_TAGS + "sezm_cute::wignerd_so2_panel", mutates_args=(), tags=_WIGNERD_CUSTOM_OP_TAGS )(_wignerd_panel_impl) _wignerd_panel_bwd_op = torch.library.custom_op( - "sezm_cute::wignerd_k1_panel_bwd", mutates_args=(), tags=_K4_CUSTOM_OP_TAGS + "sezm_cute::wignerd_so2_panel_bwd", mutates_args=(), tags=_WIGNERD_CUSTOM_OP_TAGS )(_wignerd_panel_bwd_impl) @_wignerd_panel_op.register_fake def _(edge_quat: torch.Tensor) -> torch.Tensor: - return edge_quat.new_empty((edge_quat.shape[0], K1_PANEL_VALUES)) + return edge_quat.new_empty((edge_quat.shape[0], SO2_PANEL_VALUES)) @_wignerd_panel_bwd_op.register_fake @@ -953,7 +953,7 @@ def _wignerd_panel_backward( ctx: Any, grad_panel: torch.Tensor, ) -> tuple[torch.Tensor]: - """Run packed K4 backward with its custom op visible to compilation.""" + """Run packed Wigner-D backward with its custom op visible to compilation.""" return _wignerd_panel_registered_backward_impl(ctx, grad_panel) @@ -1001,7 +1001,7 @@ def run_cute_wignerd( *, packed_wigner: bool = False, ) -> tuple[torch.Tensor, torch.Tensor] | None: - """Run K4 through its registered custom-op boundary.""" + """Run packed Wigner-D through its registered custom-op boundary.""" return _run_cute_wignerd_impl( edge_quat, wigner_calc, diff --git a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py index c0932ba1c4..c0547127a3 100644 --- a/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py +++ b/deepmd/pt_expt/kernels/triton/sezm/so2_value_path.py @@ -5373,7 +5373,7 @@ def __call__( def prepare_triton_value_path_weights(model: torch.nn.Module) -> None: """Prepare fixed SO(2) weight layouts for every bound Triton value path.""" for module in model.modules(): - value_path = getattr(module, "_triton_value_path", None) + value_path = getattr(module, "triton_infer_l_2_value", None) if isinstance(value_path, _TritonSO2ValuePath): value_path.prepare_inference_weights() diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index de0be3f023..3255dfc335 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -1456,7 +1456,7 @@ def _prepare_dpa4_triton_value_path_weights( if not _uses_dpa4_kernel_defaults(model_data): return if not any( - getattr(module, "_triton_value_path", None) is not None + getattr(module, "triton_infer_l_2_value", None) is not None for module in model.modules() ): return @@ -1766,9 +1766,9 @@ def _trace_and_export_impl( force_fused_scatter = target_device.type == "cuda" for module in model.modules(): if isinstance(module, SO2Linear): - module._force_block_diag_matmul = force_block_diag + module.force_block_diag_matmul = force_block_diag if isinstance(module, GeometricInitialEmbedding): - module._force_fused_scatter = force_fused_scatter + module.force_cuda_infer_l_1_scatter = force_fused_scatter if lower_kind == "graph" and not _supports_graph_export(model): raise NotImplementedError( diff --git a/doc/install/easy-install.md b/doc/install/easy-install.md index e8a7f9cd77..2c3baed870 100644 --- a/doc/install/easy-install.md +++ b/doc/install/easy-install.md @@ -192,18 +192,18 @@ pip install deepmd-kit :::::: Optional CuTe inference kernels are available on Linux with Python 3.11 or -newer for the PyTorch backend: +newer for the PT and PT-expt backends: ```bash pip install "deepmd-kit[torch,cute]" ``` -The `cute` extra installs the CUTLASS CuTe DSL and NVIDIA Alchemi Toolkit-Ops -runtime used by the optional CuTe paths. Set `DP_CUTE_INFER=1` to enable the -DPA4-Neo implementation in PyTorch inference. It may coexist with -`DP_TRITON_INFER`: CuTe replaces exact supported spans, while unsupported model -shapes retain the selected Triton or reference path. CuTe is not captured in a -frozen `.pt2`; requesting it without the runtime dependencies raises an error +The `cute` extra installs the CUTLASS CuTe DSL, TVM FFI, and NVIDIA Alchemi +Toolkit-Ops runtimes used by the optional CuTe paths. Set `DP_CUTE_INFER=1` to +enable the DPA4-Neo implementation in PT or PT-expt inference. It may coexist +with `DP_TRITON_INFER`: CuTe replaces exact supported spans, while unsupported +model shapes retain the selected Triton or reference path. CuTe is not captured +in a frozen `.pt2`; requesting it without the runtime dependencies raises an error rather than silently selecting another implementation. The supported platform includes Linux x86-64 and aarch64 with GNU C Library 2.28 or above, macOS x86-64 and arm64, and Windows x86-64. diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index ef87a26e6f..2631e96db5 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -413,30 +413,167 @@ Three options control training precision and the compiled path: through the energy gradient) and can speed training markedly on supported setups. +Two environment variables independently select the fused training kernels: + +| Environment variable | Default | Effect | +| -------------------- | ------- | ------------------------------------------------------------------------------------------------------------------ | +| `DP_TRITON_TRAIN` | `0` | Triton training level `0` or `1`. Level `1` enables the eligible second-order-complete stage kernels. | +| `DP_CUDA_TRAIN` | off | Enable the fused CUDA SO(2) value path on supported blocks. Accepted true values are `1`, `true`, `yes`, and `on`. | + +The corresponding false values for `DP_CUDA_TRAIN` are `0`, `false`, `no`, +and `off`. + +#### Training kernel composition + +Training and inference use separate gates because their derivative contracts +are different. Inference holds model parameters fixed and requires the forward +and coordinate gradient. A force-loss training step additionally requires +gradients of every consumed parameter and differentiates the coordinate +gradient once more. The training kernels therefore provide their own parameter +backward and second-order formulas rather than reusing inference-only kernels. + +`DP_TRITON_TRAIN=1` enables a composition of eligible Triton operators, +including Wigner monomials, the local-frame rotations, radial degree mixing, +SO(2) block GEMMs, gated activations, segmented attention softmax, flash +attention aggregation, and the wider SO(3) grid products. Unsupported stages +remain ordinary PyTorch operations. There are no training levels `2` or `3`: +the fused Triton value path and fp16x3 mixing selected by +`DP_TRITON_INFER>=2` are inference-only. + +`DP_CUDA_TRAIN=1` enables one fused CUDA operator for each supported SO(2) +block. It owns the value stream from the source-node rotation through radial +degree mixing, focus competition, and the gated mixing stack, stopping before +the attention aggregation. The operator supplies parameter gradients and the +second derivative required by force-loss training. A block outside its +supported layout or an installation without the CUDA operator falls through to +the enabled Triton stages or to PyTorch. + +The CUDA and Triton training gates are complementary, not mutually exclusive. +For a span covered by both, CUDA takes precedence only over the SO(2) value +stream; Triton remains active for the attention span and other non-overlapping +operators. Their combinations are: + +| `DP_TRITON_TRAIN` | `DP_CUDA_TRAIN` | Result | +| ----------------- | --------------- | ----------------------------------------------------------------------------------------------------------------------------------------- | +| `0` | `0` | PyTorch implements the complete training graph. | +| `1` | `0` | Eligible stages use the Triton composition. | +| `0` | `1` | Eligible value streams use CUDA; attention and unsupported blocks use PyTorch. | +| `1` | `1` | Eligible value streams use CUDA, Triton serves attention and the other supported stages, and unsupported CUDA blocks fall back to Triton. | + +When every block supports the CUDA value path and Triton flash attention, the +descriptor reuses the CUDA operator's packed Wigner runs instead of +materializing dense per-edge Wigner-D matrices. Enabling both gates can +therefore reduce memory as well as compose the two faster execution paths. +The current inference-only force and virial assembly kernels are not selected +during training. The training path keeps this assembly in the differentiable +PyTorch graph so the force loss can traverse it. + +The intended combined setting is: + +```bash +export DP_TRITON_TRAIN=1 +export DP_CUDA_TRAIN=1 +``` + +Both gates serve PT and PT-expt and are read when the model is constructed. +Set them before loading or constructing the model and before `torch.compile`. +They do not change inference dispatch; likewise, `DP_CUTE_INFER`, +`DP_TRITON_INFER`, and `DP_CUDA_INFER` do not change training dispatch. + ### Inference and deployment settings Inference behavior is controlled by environment variables, each with an equivalent input-file option used during training validation: -| Environment variable | Input-file option | Default | Effect | -| -------------------- | --------------------------- | ------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `DP_COMPILE_INFER` | `validating.compiled_infer` | off | Use the compile path for evaluation/inference. Same `torch==2.11` / CUDA ≥ 12.6 requirements as `model.use_compile`. | -| `DP_TF32_INFER` | `validating.tf32_infer` | `0` (highest) | float32 matmul precision for inference: `0` highest, `1` high, `2` medium. Higher values improve throughput but make the potential energy surface less smooth. | -| `DP_AMP_INFER` | `validating.amp_infer` | off | bf16 autocast inside the descriptor interaction blocks for inference, independently of `descriptor.use_amp`. Training AMP remains controlled by `descriptor.use_amp`. Usually keeps aggregate MAE similar but can make the potential energy surface less smooth. | -| `DP_TRITON_INFER` | — | `0` | Triton inference kernel level `0`-`3` (CUDA eval only, compatible with `DP_COMPILE_INFER`). Levels `1` and `2` are exact float32; level `3` trades a small accuracy margin for a substantial speedup. Detailed below. | -| `DP_CUTE_INFER` | — | off | CuTe inference kernels for supported SeZM shapes (PyTorch CUDA eval only, compatible with `DP_TRITON_INFER` and `DP_COMPILE_INFER`). Unsupported shapes retain the selected fallback. Python inference only, and **not** captured in a frozen `.pt2`. Detailed below. | -| `DP_CUTILE_INFER` | — | off | cuTile inference path (CUDA eval only, compatible with `DP_COMPILE_INFER`, mutually exclusive with `DP_TRITON_INFER`). Python inference only, and **not** captured in a frozen `.pt2`. Detailed below. | -| `DP_CUDA_INFER` | — | `0` | Hand-written CUDA operator level `0`-`2` (CUDA eval only, stacks on top of `DP_TRITON_INFER`). Level `1` is faster on every GPU and checkpoint measured; level `2` additionally offers the fused convolution, which routes itself per checkpoint and falls back to the level-`1` behaviour where it would not pay. Detailed below. | - -Accepted boolean values for the other switches are `1`/`true`/`yes`/`on` and -`0`/`false`/`no`/`off`; `DP_TRITON_INFER` accepts only the numeric levels. -`DP_TRITON_INFER` and `DP_CUTILE_INFER` select mutually exclusive complete -SO(2) paths. `DP_CUTE_INFER` is an exact-shape overlay and may coexist with -either fallback. +| Environment variable | Input-file option | Default | Effect | +| -------------------- | --------------------------- | ------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `DP_COMPILE_INFER` | `validating.compiled_infer` | off | Use the compile path for evaluation/inference. Same `torch==2.11` / CUDA ≥ 12.6 requirements as `model.use_compile`. | +| `DP_TF32_INFER` | `validating.tf32_infer` | `0` (highest) | float32 matmul precision for inference: `0` highest, `1` high, `2` medium. Higher values improve throughput but make the potential energy surface less smooth. | +| `DP_AMP_INFER` | `validating.amp_infer` | off | bf16 autocast inside the descriptor interaction blocks for inference, independently of `descriptor.use_amp`. Training AMP remains controlled by `descriptor.use_amp`. Usually keeps aggregate MAE similar but can make the potential energy surface less smooth. | +| `DP_TRITON_INFER` | — | `0` | Triton inference kernel level `0`-`3` (CUDA eval only, compatible with `DP_COMPILE_INFER`). Levels `1` and `2` are exact float32; level `3` trades a small accuracy margin for a substantial speedup. Detailed below. | +| `DP_CUTE_INFER` | — | off | CuTe inference kernels for supported SeZM shapes (PT and PT-expt CUDA eval only, compatible with `DP_TRITON_INFER`, `DP_CUDA_INFER`, and `DP_COMPILE_INFER`). Unsupported shapes retain the selected fallback. Python inference only, and **not** captured in a frozen `.pt2`. Detailed below. | +| `DP_CUTILE_INFER` | — | off | cuTile inference path (CUDA eval only, compatible with `DP_COMPILE_INFER`, mutually exclusive with `DP_TRITON_INFER`). Python inference only, and **not** captured in a frozen `.pt2`. Detailed below. | +| `DP_CUDA_INFER` | — | `0` | Hand-written CUDA operator level `0`-`2` (CUDA eval only, composed with the CuTe and Triton paths). Level `1` is faster on every GPU and checkpoint measured; level `2` additionally offers the fused convolution, which routes itself per checkpoint and falls back to the level-`1` behaviour where it would not pay. Detailed below. | + +Accepted boolean values are `1`/`true`/`yes`/`on` and +`0`/`false`/`no`/`off`. `DP_TRITON_INFER` and `DP_CUDA_INFER` instead accept +only their documented numeric levels. Shell exports take precedence over the input-file options and over values written in the input; they are read when the model is constructed and changing them afterward has no effect. +#### How the inference kernels compose + +`DP_CUTE_INFER`, `DP_TRITON_INFER`, and `DP_CUDA_INFER` are permissions for +different implementations, not mutually exclusive backend selectors. The +model selects one implementation independently for each computational span; +enabling several paths never evaluates the same span more than once. A larger +eligible fused span takes precedence, and an unsupported span falls through to +the next enabled implementation. + +For the main SO(2) interaction block, the dispatch order is: + +1. the exact-shape CuTe block when its complete eligibility contract holds; +1. the level-2 CUDA fused convolution when its structure and profitability + checks hold; +1. the selected Triton composition; +1. the dense PyTorch implementation. + +The surrounding operations are selected separately: + +| Computational span | First available implementation | +| -------------------------------------------------- | ------------------------------------------------------------------ | +| cutoff envelope and radial basis | CUDA level 1, then PyTorch | +| Wigner-D construction | CuTe packed Wigner-D, CUDA level 1, Triton monomials, then PyTorch | +| complete SO(2) interaction block | CuTe, CUDA level 2, Triton, then PyTorch | +| geometric initial embedding and SO(3) grid product | CuTe, CUDA level 1, then PyTorch | +| force and virial assembly | CUDA level 1, Triton level 1, then PyTorch | + +The packed Wigner-D layout used by the CuTe SO(2) path is descriptor-wide. +Consequently, every interaction block must satisfy the CuTe contract before +any block enters that path. Other CuTe kernels, such as the grid product, apply +their own local eligibility checks. CUDA level 1 remains useful around a CuTe +block because it covers independent spans such as the radial basis and force +assembly. + +With `DP_CUTILE_INFER=0`, the common combinations behave as follows, where `T` +is a nonzero Triton level: + +| `DP_CUTE_INFER` | `DP_TRITON_INFER` | `DP_CUDA_INFER` | Result | +| --------------- | ----------------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------- | +| `0` | `0` | `0` | The dense PyTorch path is used. | +| `0` | `T` | `0` | Triton serves its supported spans; other spans remain in PyTorch. | +| `0` | `T` | `1` | CUDA level 1 serves its spans, while Triton serves the remaining accelerated spans. | +| `0` | `T` | `2` | Eligible SO(2) blocks use the CUDA fused convolution; ineligible blocks fall back to Triton. CUDA level 1 remains enabled. | +| `1` | any | any | Eligible exact-shape spans use CuTe; the configured CUDA, Triton, and PyTorch paths remain as fallbacks and for non-overlapping spans. | + +When CUDA level 2 takes over an SO(2) block, its fused convolution also takes +over the mixing stack. The fp16x3 GEMMs added by Triton level 3 therefore do not +run inside that block, although they remain available to a block rejected by +the CUDA gate. The same principle applies to a block selected by CuTe. + +Among these three variables there is no prohibited combination. Do not confuse +`DP_CUTE_INFER` with `DP_CUTILE_INFER`: cuTile and Triton are mutually exclusive +complete SO(2) paths, while CuTe is an exact-shape overlay that may coexist with +either path and with the CUDA operators. + +With the default `DP_TF32_INFER=0` and `DP_AMP_INFER=0` precision policy, a +conservative exact-float32 Python inference configuration enables CuTe with +both general fallbacks: + +```bash +export DP_CUTE_INFER=1 +export DP_TRITON_INFER=2 +export DP_CUDA_INFER=1 +``` + +CUDA level 2 can replace level 1 for narrow checkpoints on GPUs where its fused +SO(2) convolution is profitable. See +[Training kernel composition](#training-kernel-composition) for the independent +training gates. + +#### Triton levels + `DP_TRITON_INFER` selects how much of the descriptor runs in fused Triton kernels. Level `1` adds universal fused kernels, numerically equivalent to the dense path with full float32 accumulation. Level `2` adds the table-configured @@ -450,14 +587,19 @@ shipping built in. On other GPUs the kernels fall back to conservative configurations, and `dp --pt freeze` tunes the missing entries on the local GPU before exporting, a one-off sweep of a few minutes baked into the `.pt2`. +#### CuTe overlay + `DP_CUTE_INFER` enables CuTe kernels for the current DPA4-Neo inference -contract. The path fuses the eligible K1 interaction block, packed Wigner-D -construction, geometric initial embedding, output grid product, and scalar -readout while retaining strict float32 arithmetic. Eligibility checks cover the -model structure, tensor layout, dtype, device, and GPU capability. Any -unsupported configuration falls through to the enabled Triton, CUDA, cuTile, or -reference implementation rather than entering a partially compatible kernel. -The kernels are JIT compiled at runtime and are not captured in a frozen `.pt2`. +contract in both PT and PT-expt. The path fuses the eligible SO(2) interaction +block, packed Wigner-D construction, geometric initial embedding, output grid +product, and scalar readout while retaining strict float32 arithmetic. +Eligibility checks cover the model structure, tensor layout, dtype, device, and +GPU capability. Any unsupported configuration falls through to the enabled +Triton, CUDA, cuTile, or reference implementation rather than entering a +partially compatible kernel. The kernels are JIT compiled at runtime and are +not captured in a frozen `.pt2`. + +#### cuTile path `DP_CUTILE_INFER` replaces the whole SeZM edge pipeline — Wigner monomials, rotate-and-mix, the gated SO(2) mixing stack, the attention aggregation and the @@ -473,6 +615,8 @@ whose layout it does not support falls back to the dense reference rather than to Triton. The kernels are JIT compiled at runtime, so this path serves Python inference only and is not captured in a frozen `.pt2`. +#### CUDA levels + `DP_CUDA_INFER` enables hand-written CUDA operators that fuse spans of the SeZM descriptor. Unlike the paths above it is not an alternative backend: it stacks on top of `DP_TRITON_INFER`, taking over the spans it covers and leaving the @@ -536,6 +680,8 @@ Both levels are precompiled custom operators that `make_fx` traces, so unlike `DP_CUTILE_INFER` they are baked into a frozen `.pt2` and keep their effect when it is later loaded by ASE or LAMMPS. +#### Precision guidance + For molecular dynamics and other workflows sensitive to the smoothness of the potential energy surface, keep `DP_TF32_INFER=0` and `DP_AMP_INFER=0`. `DP_AMP_INFER` can coexist with `DP_TF32_INFER`, but bf16 autocast dominates @@ -548,6 +694,8 @@ for the target system. `DP_CUTILE_INFER` inherits that same rounding scale through its mixing stack and is the faster of the two on Blackwell, at the cost of being unavailable to the frozen `.pt2` route. +#### Frozen `.pt2` kernel selection + > [!IMPORTANT] > Set these variables **before** running `dp --pt freeze` or > `dp --pt-expt freeze`. The exported `.pt2` is an AOTInductor artifact, so the diff --git a/pyproject.toml b/pyproject.toml index 31452aa183..967349a193 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -147,6 +147,7 @@ cu12 = [ ] cute = [ 'nvidia-cutlass-dsl[cu13]>=4.6.1,<5; python_version >= "3.11" and platform_system == "Linux"', + 'apache-tvm-ffi; python_version >= "3.11" and platform_system == "Linux"', 'nvalchemi-toolkit-ops>=0.3.1; python_version >= "3.11" and platform_system == "Linux"', ] jax = [ diff --git a/source/tests/common/dpmodel/test_dpa4_edge_cache.py b/source/tests/common/dpmodel/test_dpa4_edge_cache.py index 17f823f69f..9cdf3cb98f 100644 --- a/source/tests/common/dpmodel/test_dpa4_edge_cache.py +++ b/source/tests/common/dpmodel/test_dpa4_edge_cache.py @@ -66,3 +66,42 @@ def fused_wigner(quaternion: np.ndarray) -> tuple[np.ndarray, np.ndarray]: cache.csr_cache = None assert edge_cache_to_dtype(cache, np.float32).csr_cache is None + + +def test_edge_cache_normalizes_graph_csr_to_int64() -> None: + def wigner_calc(quaternion: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + block = np.ones((quaternion.shape[0], 1, 1), dtype=quaternion.dtype) + return block, block + + cache = _edge_cache_from_arrays( + type_ebed=np.array([[1.0], [2.0]]), + edge_index=np.array([[1, 0], [0, 1]], dtype=np.int32), + edge_vec=np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]), + edge_mask=np.array([True, True]), + compute_dtype=np.float64, + eps=1.0e-12, + deg_norm_floor=1.0, + inner_clamp=None, + bridging_switch=None, + edge_envelope=np.ones_like, + radial_basis=np.ones_like, + random_gamma=False, + wigner_calc=wigner_calc, + csr_cache={ + "dst": ( + np.array([0, 1], dtype=np.int32), + np.array([0, 1, 2], dtype=np.int32), + ), + "src": ( + np.array([1, 0], dtype=np.int32), + np.array([0, 1, 2], dtype=np.int32), + ), + }, + ) + + assert cache.src.dtype == np.int64 + assert cache.dst.dtype == np.int64 + assert cache.csr_cache is not None + for order, row_ptr in cache.csr_cache.values(): + assert order.dtype == np.int64 + assert row_ptr.dtype == np.int64 diff --git a/source/tests/common/dpmodel/test_neighbor_graph.py b/source/tests/common/dpmodel/test_neighbor_graph.py index d147ab3dd3..5ba771be86 100644 --- a/source/tests/common/dpmodel/test_neighbor_graph.py +++ b/source/tests/common/dpmodel/test_neighbor_graph.py @@ -8,6 +8,7 @@ NeighborGraph, build_edge_csr, canonicalize_neighbor_graph, + compact_edges, ) @@ -141,6 +142,79 @@ def test_canonicalization_rebuilds_an_untrusted_canonical_claim(self) -> None: np.array([[1, 0], [0, 1]], dtype=np.int64), ) + def test_compaction_preserves_canonical_csr(self) -> None: + graph = NeighborGraph( + n_node=np.array([3], dtype=np.int64), + edge_index=np.array([[2, 0, 0, 0], [0, 1, 0, 0]], dtype=np.int64), + edge_vec=np.arange(12, dtype=np.float64).reshape(4, 3), + edge_mask=np.array([True, True, False, False]), + destination_order=np.arange(4, dtype=np.int64), + destination_row_ptr=np.array([0, 1, 2, 2], dtype=np.int64), + source_order=np.array([1, 0, 2, 3], dtype=np.int64), + source_row_ptr=np.array([0, 1, 1, 2], dtype=np.int64), + destination_sorted=True, + ) + + result = compact_edges(graph) + + np.testing.assert_array_equal(result.edge_index, graph.edge_index[:, :2]) + np.testing.assert_array_equal(result.edge_vec, graph.edge_vec[:2]) + np.testing.assert_array_equal(result.edge_mask, [True, True]) + np.testing.assert_array_equal(result.destination_order, [0, 1]) + np.testing.assert_array_equal(result.destination_row_ptr, [0, 1, 2, 2]) + np.testing.assert_array_equal(result.source_order, [1, 0]) + np.testing.assert_array_equal(result.source_row_ptr, [0, 1, 1, 2]) + self.assertTrue(result.destination_sorted) + + def test_compaction_preserves_csr_with_internal_mask(self) -> None: + graph = NeighborGraph( + n_node=np.array([2], dtype=np.int64), + edge_index=np.array([[1, 0, 0], [0, 0, 1]], dtype=np.int64), + edge_vec=np.arange(9, dtype=np.float64).reshape(3, 3), + edge_mask=np.array([True, False, True]), + destination_order=np.array([0, 1, 2], dtype=np.int64), + destination_row_ptr=np.array([0, 2, 3], dtype=np.int64), + source_order=np.array([1, 2, 0], dtype=np.int64), + source_row_ptr=np.array([0, 2, 3], dtype=np.int64), + destination_sorted=True, + ) + + result = compact_edges(graph) + + np.testing.assert_array_equal(result.edge_index, graph.edge_index[:, [0, 2]]) + np.testing.assert_array_equal(result.edge_vec, graph.edge_vec[[0, 2]]) + np.testing.assert_array_equal(result.edge_mask, [True, True]) + np.testing.assert_array_equal(result.destination_order, [0, 1]) + np.testing.assert_array_equal(result.destination_row_ptr, [0, 1, 2]) + np.testing.assert_array_equal(result.source_order, [1, 0]) + np.testing.assert_array_equal(result.source_row_ptr, [0, 1, 2]) + self.assertTrue(result.destination_sorted) + + def test_compaction_preserves_permuted_csr_with_internal_mask(self) -> None: + graph = NeighborGraph( + n_node=np.array([3], dtype=np.int64), + edge_index=np.array( + [[2, 0, 1, 0], [1, 2, 0, 1]], + dtype=np.int32, + ), + edge_vec=np.arange(12, dtype=np.float64).reshape(4, 3), + edge_mask=np.array([True, False, True, True]), + destination_order=np.array([2, 0, 3, 1], dtype=np.int32), + destination_row_ptr=np.array([0, 1, 3, 4], dtype=np.int64), + source_order=np.array([1, 3, 2, 0], dtype=np.int32), + source_row_ptr=np.array([0, 2, 3, 4], dtype=np.int64), + destination_sorted=False, + ) + + result = compact_edges(graph) + + np.testing.assert_array_equal(result.edge_index, graph.edge_index[:, [0, 2, 3]]) + np.testing.assert_array_equal(result.destination_order, [1, 0, 2]) + np.testing.assert_array_equal(result.destination_row_ptr, [0, 1, 3, 3]) + np.testing.assert_array_equal(result.source_order, [2, 1, 0]) + np.testing.assert_array_equal(result.source_row_ptr, [0, 1, 2, 3]) + self.assertFalse(result.destination_sorted) + from deepmd.dpmodel.utils.neighbor_graph import ( node_ownership_mask, diff --git a/source/tests/pt/model/sezm_cute/__init__.py b/source/tests/pt/model/sezm_cute/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/pt/model/sezm_cute/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/pt/model/sezm_cute/_paths.py b/source/tests/pt/model/sezm_cute/_paths.py new file mode 100644 index 0000000000..c635e9f780 --- /dev/null +++ b/source/tests/pt/model/sezm_cute/_paths.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Source paths for CuTe kernels that cannot be imported without CUDA.""" + +from pathlib import ( + Path, +) + +from deepmd.pt_expt.kernels.cute import ( + sezm, +) + +CUTE_ROOT = Path(sezm.__path__[0]) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_compile_cache.py b/source/tests/pt/model/sezm_cute/test_compile_cache.py similarity index 71% rename from source/tests/pt/model/test_descriptor_sezm_cute_compile_cache.py rename to source/tests/pt/model/sezm_cute/test_compile_cache.py index 560c31ca2b..5c62379106 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_compile_cache.py +++ b/source/tests/pt/model/sezm_cute/test_compile_cache.py @@ -6,11 +6,7 @@ annotations, ) -import importlib.util import sys -from pathlib import ( - Path, -) from types import ( SimpleNamespace, ) @@ -18,30 +14,16 @@ mock, ) -REPO_ROOT = Path(__file__).resolve().parents[4] -CACHE_PATH = REPO_ROOT / "deepmd/pt_expt/kernels/cute/sezm/compile_cache.py" - - -def _load_cache_module(): - assert CACHE_PATH.is_file(), f"CuTe compile cache is missing: {CACHE_PATH}" - name = "sezm_cute_compile_cache_test" - spec = importlib.util.spec_from_file_location(name, CACHE_PATH) - assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) - sys.modules[name] = module - try: - spec.loader.exec_module(module) - finally: - sys.modules.pop(name, None) - return module +from deepmd.pt_expt.kernels.cute.sezm import ( + compile_cache, +) def test_cache_separates_device_and_compute_capability() -> None: - cache_module = _load_cache_module() identity = [(0, 8, 0)] calls: list[tuple[str, tuple[int, int, int]]] = [] - @cache_module.device_aware_lru_cache( + @compile_cache.device_aware_lru_cache( maxsize=8, identity_getter=lambda: identity[0], ) @@ -65,9 +47,7 @@ def compile_kernel(mode: str): def test_cache_exposes_standard_cache_controls() -> None: - cache_module = _load_cache_module() - - @cache_module.device_aware_lru_cache( + @compile_cache.device_aware_lru_cache( maxsize=2, identity_getter=lambda: (0, 8, 6), ) @@ -81,7 +61,6 @@ def compile_kernel(rows: int): def test_cache_compiles_inside_the_keyed_cuda_device() -> None: - cache_module = _load_cache_module() entered: list[tuple[str, int]] = [] class DeviceContext: @@ -101,7 +80,7 @@ def __exit__(self, *_args) -> None: ) ) - @cache_module.device_aware_lru_cache( + @compile_cache.device_aware_lru_cache( maxsize=2, identity_getter=lambda: (3, 9, 0), ) diff --git a/source/tests/pt/model/test_sezm_cute_efs.py b/source/tests/pt/model/sezm_cute/test_efs.py similarity index 94% rename from source/tests/pt/model/test_sezm_cute_efs.py rename to source/tests/pt/model/sezm_cute/test_efs.py index e8a0bf33e4..b89fa5bfd3 100644 --- a/source/tests/pt/model/test_sezm_cute_efs.py +++ b/source/tests/pt/model/sezm_cute/test_efs.py @@ -34,7 +34,6 @@ SUPPORTED_CAPABILITIES = SM80_CAPABILITIES | frozenset( {(8, 9), (9, 0), (10, 0), (12, 0)} ) -REPO_ROOT = Path(__file__).resolve().parents[4] _CHILD_MODE = "--neo-cute-efs-child" @@ -47,6 +46,7 @@ def _cute_runtime_skip_reason() -> str | None: try: importlib.import_module("cutlass.cute") importlib.import_module("cuda.bindings.driver") + importlib.import_module("tvm_ffi") except Exception as exc: # pragma: no cover - runtime dependent return f"Neo CuTe E/F/S acceptance requires the CuTe DSL runtime: {exc}" return None @@ -197,25 +197,25 @@ def _detached_outputs(outputs: dict[str, torch.Tensor]) -> dict[str, torch.Tenso def _assert_runtime_profile() -> dict[str, bool]: from deepmd.pt_expt.kernels.cute.sezm import ( - k1, runtime_policy, ) + from deepmd.pt_expt.kernels.cute.sezm.so2 import operation as so2 capability = tuple(torch.cuda.get_device_capability()) if capability not in SUPPORTED_CAPABILITIES: - raise AssertionError(f"unsupported K1 capability {capability}") + raise AssertionError(f"unsupported SO2 capability {capability}") policy_checks = { "master": runtime_policy.is_cute_infer_enabled(), "triton_infer_2": os.environ.get("DP_TRITON_INFER") == "2", - "supported": runtime_policy.is_supported_k1_capability(capability), + "supported": runtime_policy.is_supported_so2_capability(capability), "packed_wigner": runtime_policy.is_packed_wigner_enabled(capability), } - config = k1._architecture_default_config(capability) + config = so2._architecture_default_config(capability) if capability in SM80_CAPABILITIES: architecture_checks = { "sm80_profile": runtime_policy.is_sm80_profile_enabled(capability), "gie": runtime_policy.is_gie_enabled(capability), - "thin_wrapper": runtime_policy.is_k1_thin_wrapper_enabled(capability), + "thin_wrapper": runtime_policy.is_so2_thin_wrapper_enabled(capability), "output_grid_bwd_panel": ( runtime_policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled(capability) ), @@ -268,14 +268,12 @@ def _run_efs_child( use_compile: bool, output_path: Path, ) -> None: - from deepmd.pt_expt.kernels.cute.sezm import ( - k1, - ) + from deepmd.pt_expt.kernels.cute.sezm.so2 import operation as so2 torch.manual_seed(20260726) torch.cuda.manual_seed_all(20260726) runner_calls = 0 - original_build_runner = k1._build_runner + original_build_runner = so2._build_runner def counted_build_runner(*args: Any, **kwargs: Any) -> Any: nonlocal runner_calls @@ -283,7 +281,7 @@ def counted_build_runner(*args: Any, **kwargs: Any) -> Any: return original_build_runner(*args, **kwargs) if use_compile: - k1._build_runner = counted_build_runner + so2._build_runner = counted_build_runner try: model = _neo_model(use_compile=use_compile) state_keys_before = tuple(model.state_dict()) @@ -299,12 +297,12 @@ def counted_build_runner(*args: Any, **kwargs: Any) -> Any: state_keys_after = tuple(model.state_dict()) finally: if use_compile: - k1._build_runner = original_build_runner + so2._build_runner = original_build_runner if state_keys_after != state_keys_before: raise AssertionError("CuTe warmup changed the model state_dict contract") if use_compile and runner_calls == 0: - raise AssertionError("compiled CuTe run did not instantiate the K1 runner") + raise AssertionError("compiled CuTe run did not instantiate the SO2 runner") torch.save( { "outputs": runs, @@ -329,10 +327,6 @@ def _clean_child_environment() -> dict[str, str]: "NVIDIA_TF32_OVERRIDE", ): environment.pop(name, None) - pythonpath = environment.get("PYTHONPATH") - environment["PYTHONPATH"] = ( - str(REPO_ROOT) if not pythonpath else f"{REPO_ROOT}{os.pathsep}{pythonpath}" - ) environment.update( { "DP_INTERFACE_PREC": "low", @@ -381,7 +375,6 @@ def _run_child_process( "compiled" if use_compile else "eager", str(output_path), ], - cwd=REPO_ROOT, env=_child_environment(use_compile=use_compile), capture_output=True, text=True, diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_envelope_softmax.py b/source/tests/pt/model/sezm_cute/test_envelope_softmax.py similarity index 96% rename from source/tests/pt/model/test_descriptor_sezm_cute_envelope_softmax.py rename to source/tests/pt/model/sezm_cute/test_envelope_softmax.py index 4ae651ba60..e6c27b5d05 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_envelope_softmax.py +++ b/source/tests/pt/model/sezm_cute/test_envelope_softmax.py @@ -44,7 +44,7 @@ def test_cute_envelope_softmax_matches_eager_across_extreme_frames( logit_scale: float, edge_gate: float, ) -> None: - from deepmd.pt_expt.kernels.cute.sezm.k1_kernels.cute_envelope_gated_softmax import ( + from deepmd.pt_expt.kernels.cute.sezm.so2.kernels.envelope_softmax import ( compile_envelope_softmax_forward, ) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_gie.py b/source/tests/pt/model/sezm_cute/test_gie.py similarity index 100% rename from source/tests/pt/model/test_descriptor_sezm_cute_gie.py rename to source/tests/pt/model/sezm_cute/test_gie.py diff --git a/source/tests/pt/model/test_sezm_cute_multigpu_dispatch.py b/source/tests/pt/model/sezm_cute/test_multigpu_dispatch.py similarity index 82% rename from source/tests/pt/model/test_sezm_cute_multigpu_dispatch.py rename to source/tests/pt/model/sezm_cute/test_multigpu_dispatch.py index cef11b404a..f8863d777d 100644 --- a/source/tests/pt/model/test_sezm_cute_multigpu_dispatch.py +++ b/source/tests/pt/model/sezm_cute/test_multigpu_dispatch.py @@ -22,9 +22,9 @@ mlp, ) from deepmd.pt_expt.kernels.cute.sezm import ( - k1, runtime_policy, ) +from deepmd.pt_expt.kernels.cute.sezm.so2 import operation as so2 def _cuda_operand(index: int = 1) -> SimpleNamespace: @@ -34,12 +34,12 @@ def _cuda_operand(index: int = 1) -> SimpleNamespace: ) -def test_k1_custom_ops_are_direct_and_thin_path_uses_operand_device() -> None: +def test_so2_custom_ops_are_direct_and_thin_path_uses_operand_device() -> None: operand = _cuda_operand() sentinel = object() - with mock.patch.object(k1, "_cute_k1_impl", return_value=sentinel) as direct: - actual = k1.cute_k1( + with mock.patch.object(so2, "_cute_so2_impl", return_value=sentinel) as direct: + actual = so2.cute_so2( 1, operand, object(), @@ -59,33 +59,33 @@ def test_k1_custom_ops_are_direct_and_thin_path_uses_operand_device() -> None: context = object() with mock.patch.object( - k1, - "_k1_packed_direct_registered_backward_impl", + so2, + "_so2_packed_direct_registered_backward_impl", return_value=sentinel, ) as backward: - actual = k1._k1_packed_direct_backward(context, operand, None) + actual = so2._so2_packed_direct_backward(context, operand, None) assert actual is sentinel backward.assert_called_once_with(context, operand) with ( mock.patch.object( - k1, + so2, "_cuda_compute_capability", return_value=(9, 0), ) as capability, mock.patch.object( runtime_policy, - "is_k1_thin_wrapper_enabled", + "is_so2_thin_wrapper_enabled", return_value=False, ) as thin_selector, mock.patch.object( - k1, - "_maybe_run_cute_k1_fallback", + so2, + "_maybe_run_cute_so2_fallback", return_value=sentinel, ) as fallback, ): - actual = k1.maybe_run_cute_k1(object(), operand, object(), object()) + actual = so2.maybe_run_cute_so2(object(), operand, object(), object()) assert actual is sentinel capability.assert_called_once_with(1) @@ -93,17 +93,17 @@ def test_k1_custom_ops_are_direct_and_thin_path_uses_operand_device() -> None: fallback.assert_called_once() -def test_k1_prepare_uses_requested_device_instead_of_current_device() -> None: +def test_so2_prepare_uses_requested_device_instead_of_current_device() -> None: device = torch.device("cuda", 1) with ( mock.patch.object( - k1, + so2, "_cuda_compute_capability", return_value=(8, 6), ) as capability, mock.patch.object( - k1, - "is_supported_k1_compute_capability", + so2, + "is_supported_so2_compute_capability", return_value=False, ) as supported, mock.patch.object( @@ -112,7 +112,7 @@ def test_k1_prepare_uses_requested_device_instead_of_current_device() -> None: side_effect=AssertionError("must not read the current CUDA device"), ), ): - assert not k1.prepare_cute_k1_blocks( + assert not so2.prepare_cute_so2_blocks( [object()], training=False, device=device, @@ -123,18 +123,18 @@ def test_k1_prepare_uses_requested_device_instead_of_current_device() -> None: supported.assert_called_once_with((8, 6)) -def test_k4_custom_ops_are_called_directly() -> None: +def test_wignerd_custom_ops_are_called_directly() -> None: pytest.importorskip("cutlass.cute") - k4_wignerd = importlib.import_module("deepmd.pt_expt.kernels.cute.sezm.k4_wignerd") + wignerd = importlib.import_module("deepmd.pt_expt.kernels.cute.sezm.wignerd") operand = _cuda_operand() sentinel = object() with mock.patch.object( - k4_wignerd, + wignerd, "_run_cute_wignerd_impl", return_value=sentinel, ) as forward: - actual = k4_wignerd.run_cute_wignerd(operand, object()) + actual = wignerd.run_cute_wignerd(operand, object()) assert actual is sentinel forward.assert_called_once_with(operand, mock.ANY, packed_wigner=False) @@ -142,11 +142,11 @@ def test_k4_custom_ops_are_called_directly() -> None: context = SimpleNamespace(saved_tensors=(operand,)) grad_panel = object() with mock.patch.object( - k4_wignerd, + wignerd, "_wignerd_panel_registered_backward_impl", return_value=sentinel, ) as backward: - actual = k4_wignerd._wignerd_panel_backward(context, grad_panel) + actual = wignerd._wignerd_panel_backward(context, grad_panel) assert actual is sentinel backward.assert_called_once_with(context, grad_panel) @@ -171,9 +171,9 @@ def capability(device: torch.device | None = None) -> tuple[int, int]: side_effect=capability, ) as get_capability, ): - assert mlp._use_k1_compile_visible_linear(cuda1) - assert not mlp._use_k1_compile_visible_linear(torch.device("cuda", 0)) - assert not mlp._use_k1_compile_visible_linear(torch.device("cpu")) + assert mlp._use_so2_compile_visible_linear(cuda1) + assert not mlp._use_so2_compile_visible_linear(torch.device("cuda", 0)) + assert not mlp._use_so2_compile_visible_linear(torch.device("cpu")) assert get_capability.call_args_list == [ mock.call(cuda1), @@ -191,7 +191,7 @@ def capability(device: torch.device | None = None) -> tuple[int, int]: value = torch.randn(3, 4, device="cpu") with mock.patch.object( mlp, - "_use_k1_compile_visible_linear", + "_use_so2_compile_visible_linear", return_value=False, ) as selector: layer(value) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_output_grid_product.py b/source/tests/pt/model/sezm_cute/test_output_grid_product.py similarity index 94% rename from source/tests/pt/model/test_descriptor_sezm_cute_output_grid_product.py rename to source/tests/pt/model/sezm_cute/test_output_grid_product.py index 8b31958dd5..f3531794d3 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_output_grid_product.py +++ b/source/tests/pt/model/sezm_cute/test_output_grid_product.py @@ -99,8 +99,8 @@ def test_dispatch_falls_back_without_exact_cuda_contract( monkeypatch: pytest.MonkeyPatch, hidden_channels: int, ): - from deepmd.pt_expt.kernels.cute.sezm import ( - output_grid_product, + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( + product as output_grid_product, ) monkeypatch.setattr( @@ -133,7 +133,7 @@ def test_exact_shape_guard_accepts_only_validated_widths( hidden_channels: int, expected: int | None, ) -> None: - from deepmd.pt_expt.kernels.cute.sezm.output_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid.product import ( _exact_hidden_channels, ) @@ -143,8 +143,8 @@ def test_exact_shape_guard_accepts_only_validated_widths( @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) def test_fake_registrations_allocate_canonical_strides(hidden_channels: int) -> None: - from deepmd.pt_expt.kernels.cute.sezm import ( - output_grid_product, + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( + product as output_grid_product, ) width = N_FRAMES * hidden_channels @@ -205,12 +205,12 @@ def test_sm90_c96_asymmetric_panel_dispatch_is_explicit( policy_enabled: bool, expected: bool, ) -> None: - from deepmd.pt_expt.kernels.cute.sezm import ( - output_grid_product, + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( + product as output_grid_product, ) kernel_module_name = ( - "deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product" + "deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product" ) kernel_module = types.ModuleType(kernel_module_name) calls: dict[str, dict[str, bool]] = {} @@ -339,7 +339,7 @@ def test_forward_and_first_backward_match_strict_fp32( nodes: int, hidden_channels: int, ): - from deepmd.pt_expt.kernels.cute.sezm.output_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid.product import ( output_grid_product_cute, ) @@ -387,8 +387,8 @@ def test_sm80_c96_n48_forward_custom_op_matches_strict_fp32( monkeypatch: pytest.MonkeyPatch, nodes: int, ): - from deepmd.pt_expt.kernels.cute.sezm import ( - output_grid_product, + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( + product as output_grid_product, ) monkeypatch.setattr( @@ -415,8 +415,8 @@ def test_sm80_c96_n48_panel_custom_op_matches_strict_fp32( self, monkeypatch: pytest.MonkeyPatch, ): - from deepmd.pt_expt.kernels.cute.sezm import ( - output_grid_product, + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( + product as output_grid_product, ) monkeypatch.setattr( @@ -464,8 +464,8 @@ def test_dispatch_uses_only_the_master_gate( monkeypatch: pytest.MonkeyPatch, hidden_channels: int, ): - from deepmd.pt_expt.kernels.cute.sezm import ( - output_grid_product, + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( + product as output_grid_product, ) left, right, to_grid, from_grid = self._inputs(2, hidden_channels) @@ -511,8 +511,8 @@ def test_dispatch_declines_non_strict_matmul_state( monkeypatch: pytest.MonkeyPatch, hidden_channels: int, ): - from deepmd.pt_expt.kernels.cute.sezm import ( - output_grid_product, + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( + product as output_grid_product, ) left, right, to_grid, from_grid = self._inputs(2, hidden_channels) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_readout_l0.py b/source/tests/pt/model/sezm_cute/test_readout_l0.py similarity index 92% rename from source/tests/pt/model/test_descriptor_sezm_cute_readout_l0.py rename to source/tests/pt/model/sezm_cute/test_readout_l0.py index eef81470e5..8adaac0f81 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_readout_l0.py +++ b/source/tests/pt/model/sezm_cute/test_readout_l0.py @@ -226,9 +226,11 @@ def test_sm90_readout_input_fold_selector_uses_cute_gate(monkeypatch): def test_readout_input_fold_guard_accepts_validated_architectures(monkeypatch): from deepmd.pt_expt.kernels.cute.sezm import ( - readout_l0, runtime_policy, ) + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( + readout_l0, + ) module = _output_ffn() value = torch.randn(2, COEFF_DIM, 1, 32, device="cpu") @@ -253,7 +255,7 @@ def select_readout(capability): def test_sm80_readout_input_fold_matches_forward_and_input_vjp(monkeypatch): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -317,7 +319,7 @@ def compact_reference(left, right, to_grid, from_grid): def test_sm80_readout_input_fold_is_fullgraph_safe(monkeypatch): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -346,7 +348,7 @@ def folded(value): def test_sm80_readout_input_fold_cache_refreshes_after_inplace_change(monkeypatch): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -379,7 +381,7 @@ def test_sm80_readout_input_fold_cache_refreshes_after_inplace_change(monkeypatc def test_sm80_readout_input_fold_cache_refreshes_after_parameter_replacement(): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -398,7 +400,7 @@ def test_sm80_readout_input_fold_cache_refreshes_after_parameter_replacement(): @pytest.mark.parametrize("assign", [False, True]) def test_sm80_readout_input_fold_load_state_dict_invalidates_cache(assign): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -421,7 +423,7 @@ def test_sm80_readout_input_fold_load_state_dict_invalidates_cache(assign): def test_sm80_readout_input_fold_buffers_follow_device_moves(): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -441,7 +443,7 @@ def test_sm80_readout_input_fold_buffers_follow_device_moves(): def test_sm80_readout_input_fold_compile_requires_explicit_preparation( monkeypatch, ): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -454,7 +456,7 @@ def test_sm80_readout_input_fold_compile_requires_explicit_preparation( def test_sm80_readout_input_fold_compiled_lookup_returns_prepared_buffers( monkeypatch, ): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -475,7 +477,7 @@ def test_shape_guard_accepts_only_exact_c192_readout( hidden_channels: int, expected: bool, ): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -484,7 +486,7 @@ def test_shape_guard_accepts_only_exact_c192_readout( def test_fake_registrations_return_canonical_strides(): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -514,7 +516,7 @@ def test_custom_ops_use_canonical_fake_metadata(): FakeTensorMode, ) - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -541,7 +543,7 @@ def test_custom_ops_use_canonical_fake_metadata(): def test_reference_completion_matches_module_output_and_input_vjp(): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -563,16 +565,14 @@ def test_reference_completion_matches_module_output_and_input_vjp(): def test_exact_structure_and_frozen_inference_guards(): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) module = _output_ffn() - assert readout_l0._has_exact_neo_readout_structure(module) + assert readout_l0.has_neo_readout_contract(module) assert readout_l0._inference_mode_is_frozen(module) - assert not readout_l0._has_exact_neo_readout_structure( - _output_ffn(hidden_channels=64) - ) + assert not readout_l0.has_neo_readout_contract(_output_ffn(hidden_channels=64)) module.train() assert not readout_l0._inference_mode_is_frozen(module) module.eval() @@ -581,7 +581,7 @@ def test_exact_structure_and_frozen_inference_guards(): def test_module_boundary_falls_back_when_device_is_unsupported(monkeypatch): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -596,8 +596,8 @@ def test_module_boundary_falls_back_when_device_is_unsupported(monkeypatch): ) -def test_descriptor_freeze_guard_and_readout_wiring(cpu_neo_descriptor, monkeypatch): - from deepmd.pt_expt.kernels.cute.sezm import ( +def test_descriptor_readout_wiring(cpu_neo_descriptor, monkeypatch): + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -606,22 +606,18 @@ def test_descriptor_freeze_guard_and_readout_wiring(cpu_neo_descriptor, monkeypa expected = torch.randn(2, 32, device="cpu") calls = [] - def record_readout(output_ffn, ffn_in, *, parameters_frozen): - calls.append((output_ffn, ffn_in, parameters_frozen)) + def record_readout(output_ffn, ffn_in): + calls.append((output_ffn, ffn_in)) return expected - monkeypatch.setattr(readout_l0, "run_neo_output_readout", record_readout) - assert descriptor._readout_parameters_are_frozen() - assert descriptor._run_output_readout(value) is expected - assert calls[-1] == (descriptor.output_ffn, value, True) - - next(descriptor.parameters()).requires_grad_(True) - descriptor._run_output_readout(value) - assert calls[-1] == (descriptor.output_ffn, value, False) + monkeypatch.setenv("DP_CUTE_INFER", "1") + monkeypatch.setattr(readout_l0, "maybe_run_neo_readout_l0", record_readout) + assert descriptor.run_cute_infer_readout(value) is expected + assert calls[-1] == (descriptor.output_ffn, value) def test_trainable_descriptor_state_bypasses_candidate(monkeypatch): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -647,7 +643,7 @@ def test_legacy_matmul_precision_guard_is_eager_and_fullgraph_safe( precision: str, expected: bool, ): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -676,7 +672,7 @@ def test_modern_matmul_precision_guard_rejects_tf32( precision: str, expected: bool, ): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -720,7 +716,7 @@ def test_forward_and_input_vjp_match_strict_fp32( self, nodes: int, ): - from deepmd.pt_expt.kernels.cute.sezm.readout_l0 import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid.readout_l0 import ( readout_l0_product_cute, ) @@ -745,7 +741,7 @@ def test_forward_and_input_vjp_match_strict_fp32( ) def test_opcheck_and_fullgraph_preserve_metadata_and_vjp(self): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -770,7 +766,7 @@ def test_opcheck_and_fullgraph_preserve_metadata_and_vjp(self): assert grads[1].shape == right.shape def test_exact_module_fullgraph_uses_full_neo_gate(self, monkeypatch): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -799,7 +795,7 @@ def test_exact_module_fullgraph_uses_full_neo_gate(self, monkeypatch): torch.testing.assert_close(actual_grad, expected_grad, atol=TOL, rtol=TOL) def test_sm80_input_fold_preparation_is_cross_stream_ready(self): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) @@ -823,7 +819,7 @@ def test_sm80_input_fold_preparation_is_cross_stream_ready(self): torch.testing.assert_close(actual, reference, atol=0.0, rtol=0.0) def test_sm80_input_fold_fullgraph_matches_module_vjp(self, monkeypatch): - from deepmd.pt_expt.kernels.cute.sezm import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( readout_l0, ) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_runtime_policy.py b/source/tests/pt/model/sezm_cute/test_runtime_policy.py similarity index 76% rename from source/tests/pt/model/test_descriptor_sezm_cute_runtime_policy.py rename to source/tests/pt/model/sezm_cute/test_runtime_policy.py index c966f3adac..5276b842e8 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_runtime_policy.py +++ b/source/tests/pt/model/sezm_cute/test_runtime_policy.py @@ -6,36 +6,15 @@ annotations, ) -import importlib.util import os -import sys -from pathlib import ( - Path, -) from unittest import ( mock, ) -REPO_ROOT = Path(__file__).resolve().parents[4] -POLICY_PATH = REPO_ROOT / "deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py" - - -def _load_policy(): - assert POLICY_PATH.is_file(), f"CuTe runtime policy is missing: {POLICY_PATH}" - name = "sezm_cute_runtime_policy_test" - spec = importlib.util.spec_from_file_location(name, POLICY_PATH) - assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) - sys.modules[name] = module - try: - spec.loader.exec_module(module) - finally: - sys.modules.pop(name, None) - return module +from deepmd.pt_expt.kernels.cute.sezm import runtime_policy as policy def test_cute_master_gate_controls_sezm_path() -> None: - policy = _load_policy() cases = ( ({}, False), ({"DP_CUTE_INFER": "0"}, False), @@ -50,13 +29,12 @@ def test_cute_master_gate_controls_sezm_path() -> None: def test_master_switch_controls_every_sm80_subfeature() -> None: - policy = _load_policy() with mock.patch.dict( os.environ, { "DP_CUTE_GIE": "1", - "DP_CUTE_K1_PACKED_WIGNER": "1", - "DP_CUTE_K1_THIN_WRAPPER": "1", + "DP_CUTE_SO2_PACKED_WIGNER": "1", + "DP_CUTE_SO2_THIN_WRAPPER": "1", "DP_CUTE_OUTPUT_GRID_BWD_SM80_C96_N48_PANEL": "1", "DP_CUTE_OUTPUT_GRID_FWD_SM80_C96_N48": "1", "DP_CUTE_READOUT_INPUT_FOLD_SM80": "1", @@ -68,32 +46,30 @@ def test_master_switch_controls_every_sm80_subfeature() -> None: assert not policy.is_sm80_profile_enabled(capability) assert not policy.is_gie_enabled(capability) assert not policy.is_packed_wigner_enabled(capability) - assert not policy.is_k1_thin_wrapper_enabled(capability) + assert not policy.is_so2_thin_wrapper_enabled(capability) assert not policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled(capability) assert not policy.is_output_grid_fwd_sm80_c96_n48_enabled(capability) assert not policy.is_readout_input_fold_sm80_enabled(capability) def test_sm80_family_profile_defaults_from_master_only() -> None: - policy = _load_policy() with mock.patch.dict(os.environ, {"DP_CUTE_INFER": "1"}, clear=True): assert policy.SM80_PROFILE_CAPABILITIES == frozenset({(8, 0), (8, 6)}) for capability in policy.SM80_PROFILE_CAPABILITIES: assert policy.is_sm80_profile_enabled(capability) assert policy.is_gie_enabled(capability) assert policy.is_packed_wigner_enabled(capability) - assert policy.is_k1_thin_wrapper_enabled(capability) + assert policy.is_so2_thin_wrapper_enabled(capability) assert policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled(capability) assert policy.is_output_grid_fwd_sm80_c96_n48_enabled(capability) assert policy.is_readout_input_fold_sm80_enabled(capability) def test_every_sm80_profile_feature_can_be_disabled() -> None: - policy = _load_policy() checks = { "DP_CUTE_GIE": lambda: policy.is_gie_enabled((8, 0)), - "DP_CUTE_K1_PACKED_WIGNER": lambda: policy.is_packed_wigner_enabled((8, 0)), - "DP_CUTE_K1_THIN_WRAPPER": lambda: policy.is_k1_thin_wrapper_enabled((8, 0)), + "DP_CUTE_SO2_PACKED_WIGNER": lambda: policy.is_packed_wigner_enabled((8, 0)), + "DP_CUTE_SO2_THIN_WRAPPER": lambda: policy.is_so2_thin_wrapper_enabled((8, 0)), "DP_CUTE_OUTPUT_GRID_BWD_SM80_C96_N48_PANEL": lambda: ( policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled((8, 0)) ), @@ -114,20 +90,18 @@ def test_every_sm80_profile_feature_can_be_disabled() -> None: def test_non_sm80_family_uses_architecture_defaults_without_sm80_features() -> None: - policy = _load_policy() with mock.patch.dict(os.environ, {"DP_CUTE_INFER": "1"}, clear=True): capability = (8, 9) assert not policy.is_sm80_profile_enabled(capability) assert not policy.is_gie_enabled(capability) assert policy.is_packed_wigner_enabled(capability) - assert not policy.is_k1_thin_wrapper_enabled(capability) + assert not policy.is_so2_thin_wrapper_enabled(capability) assert not policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled(capability) assert not policy.is_output_grid_fwd_sm80_c96_n48_enabled(capability) assert not policy.is_readout_input_fold_sm80_enabled(capability) def test_sm90_c96_asymmetric_panels_default_is_exact_arch_disable_only() -> None: - policy = _load_policy() switch = policy.OUTPUT_GRID_SM90_C96_ASYMMETRIC_PANELS_ENV with mock.patch.dict(os.environ, {"DP_CUTE_INFER": "1"}, clear=True): @@ -155,19 +129,18 @@ def test_sm90_c96_asymmetric_panels_default_is_exact_arch_disable_only() -> None def test_overrides_remain_architecture_safe() -> None: - policy = _load_policy() with mock.patch.dict( os.environ, { "DP_CUTE_INFER": "1", - "DP_CUTE_K1_THIN_WRAPPER": "1", + "DP_CUTE_SO2_THIN_WRAPPER": "1", "DP_CUTE_OUTPUT_GRID_BWD_SM80_C96_N48_PANEL": "1", "DP_CUTE_OUTPUT_GRID_FWD_SM80_C96_N48": "1", "DP_CUTE_READOUT_INPUT_FOLD_SM80": "1", }, clear=True, ): - assert policy.is_k1_thin_wrapper_enabled((9, 0)) + assert policy.is_so2_thin_wrapper_enabled((9, 0)) assert not policy.is_output_grid_bwd_sm80_c96_n48_panel_enabled((9, 0)) assert not policy.is_output_grid_fwd_sm80_c96_n48_enabled((9, 0)) assert not policy.is_readout_input_fold_sm80_enabled((9, 0)) @@ -177,7 +150,7 @@ def test_overrides_remain_architecture_safe() -> None: { "DP_CUTE_INFER": "1", "DP_CUTE_GIE": "1", - "DP_CUTE_K1_PACKED_WIGNER": "1", + "DP_CUTE_SO2_PACKED_WIGNER": "1", }, clear=True, ): @@ -185,35 +158,33 @@ def test_overrides_remain_architecture_safe() -> None: assert not policy.is_packed_wigner_enabled((10, 1)) -def test_int32_k1_capacity_checks_every_flattened_axis() -> None: - policy = _load_policy() - max_edges = policy.INT32_MAX // policy.K1_VALUES_PER_EDGE - max_nodes = policy.INT32_MAX // policy.K1_VALUES_PER_NODE +def test_int32_so2_capacity_checks_every_flattened_axis() -> None: + max_edges = policy.INT32_MAX // policy.SO2_VALUES_PER_EDGE + max_nodes = policy.INT32_MAX // policy.SO2_VALUES_PER_NODE - assert policy.k1_int32_indexing_is_safe( + assert policy.so2_int32_indexing_is_safe( edge_count=max_edges, node_count=max_nodes, ) - assert not policy.k1_int32_indexing_is_safe( + assert not policy.so2_int32_indexing_is_safe( edge_count=max_edges + 1, node_count=1, ) - assert not policy.k1_int32_indexing_is_safe( + assert not policy.so2_int32_indexing_is_safe( edge_count=1, node_count=max_nodes + 1, ) - assert not policy.k1_int32_indexing_is_safe( + assert not policy.so2_int32_indexing_is_safe( edge_count=-1, node_count=1, ) - assert not policy.k1_int32_indexing_is_safe( + assert not policy.so2_int32_indexing_is_safe( edge_count=1, node_count=-1, ) def test_strict_mode_is_explicitly_opt_in() -> None: - policy = _load_policy() with mock.patch.dict(os.environ, {}, clear=True): assert not policy.is_cute_strict_enabled() with mock.patch.dict(os.environ, {"DP_CUTE_STRICT": "1"}, clear=True): @@ -221,7 +192,6 @@ def test_strict_mode_is_explicitly_opt_in() -> None: def test_neighbor_list_eager_island_policy_is_owned_by_neo_runtime() -> None: - policy = _load_policy() cases = ( ({}, (8, 0), False), ({"DP_CUTE_INFER": "1"}, (8, 0), True), @@ -230,7 +200,7 @@ def test_neighbor_list_eager_island_policy_is_owned_by_neo_runtime() -> None: ( { "DP_CUTE_INFER": "1", - "DP_CUTE_K1_EAGER_ISLANDS": "0", + "DP_CUTE_SO2_EAGER_ISLANDS": "0", }, (8, 0), False, @@ -238,7 +208,7 @@ def test_neighbor_list_eager_island_policy_is_owned_by_neo_runtime() -> None: ( { "DP_CUTE_INFER": "1", - "DP_CUTE_K1_EAGER_ISLANDS": "1", + "DP_CUTE_SO2_EAGER_ISLANDS": "1", }, (9, 0), True, @@ -246,4 +216,4 @@ def test_neighbor_list_eager_island_policy_is_owned_by_neo_runtime() -> None: ) for environment, capability, expected in cases: with mock.patch.dict(os.environ, environment, clear=True): - assert policy.is_k1_eager_island_enabled(capability) is expected + assert policy.is_so2_eager_island_enabled(capability) is expected diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1.py b/source/tests/pt/model/sezm_cute/test_so2.py similarity index 84% rename from source/tests/pt/model/test_descriptor_sezm_cute_k1.py rename to source/tests/pt/model/sezm_cute/test_so2.py index 02c97b0cc5..1fc07f52b2 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1.py +++ b/source/tests/pt/model/sezm_cute/test_so2.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""Contracts for the opt-in CuTe Neo K1 inference path.""" +"""Contracts for the opt-in CuTe Neo SO2 inference path.""" from __future__ import ( annotations, @@ -25,18 +25,18 @@ from deepmd.pt.model.descriptor.sezm_nn.norm import ( EquivariantRMSNorm, ) -from deepmd.pt_expt.kernels.cute.sezm import k1 as _K1 -from deepmd.pt_expt.kernels.cute.sezm import k1_runner as _K1_RUNNER from deepmd.pt_expt.kernels.cute.sezm import ( runtime_policy, ) +from deepmd.pt_expt.kernels.cute.sezm.so2 import operation as _SO2 +from deepmd.pt_expt.kernels.cute.sezm.so2 import runner as _SO2_RUNNER -NeoK1BackwardWorkspace = _K1_RUNNER.NeoK1BackwardWorkspace -NeoK1RuntimeConfig = _K1.NeoK1RuntimeConfig -NeoFullCuteBackward = _K1_RUNNER.NeoFullCuteBackward -StackCache = _K1_RUNNER.StackCache -_validate_runtime_config = _K1_RUNNER._validate_runtime_config -_uses_packed_message_grid = _K1_RUNNER._uses_packed_message_grid +NeoSO2BackwardWorkspace = _SO2_RUNNER.NeoSO2BackwardWorkspace +NeoSO2RuntimeConfig = _SO2.NeoSO2RuntimeConfig +NeoFullCuteBackward = _SO2_RUNNER.NeoFullCuteBackward +StackCache = _SO2_RUNNER.StackCache +_validate_runtime_config = _SO2_RUNNER._validate_runtime_config +_uses_packed_message_grid = _SO2_RUNNER._uses_packed_message_grid class _Identity: @@ -110,8 +110,8 @@ def _neo_like_block(**overrides): def test_exact_neo_contract_is_supported() -> None: block = _neo_like_block() - assert _K1.get_neo_k1_spec(block).is_current_neo_target - assert _K1.is_supported_neo_k1_block(block) + assert _SO2.get_neo_so2_spec(block).is_current_neo_target + assert _SO2.is_supported_neo_so2_block(block) @pytest.mark.parametrize( @@ -127,21 +127,21 @@ def test_exact_neo_contract_is_supported() -> None: ], ) def test_non_neo_contracts_fall_back(field: str, value: object) -> None: - assert not _K1.is_supported_neo_k1_block(_neo_like_block(**{field: value})) + assert not _SO2.is_supported_neo_so2_block(_neo_like_block(**{field: value})) def test_unsupported_message_grid_contract_falls_back() -> None: block = _neo_like_block() block.so2_conv.message_node_grid_product.op_type = "mlp" - assert not _K1.is_supported_neo_k1_block(block) + assert not _SO2.is_supported_neo_so2_block(block) def test_unsupported_post_norm_falls_back() -> None: block = _neo_like_block() block.post_so2_norm = torch.nn.LayerNorm(32, device="cpu") - assert not _K1.is_supported_neo_k1_block(block) + assert not _SO2.is_supported_neo_so2_block(block) def test_gate_expand_contract_cache_tracks_buffer_versions() -> None: @@ -153,21 +153,21 @@ def test_gate_expand_contract_cache_tracks_buffer_versions() -> None: ) block.so2_conv.non_linearities = [SimpleNamespace(expand_index=expand_index)] - assert _K1._gate_expand_index_structure_is_supported(block) - assert _K1._gate_expand_index_is_supported(block) + assert _SO2._gate_expand_index_structure_is_supported(block) + assert _SO2._gate_expand_index_is_supported(block) expand_index[0] = 2 - assert _K1._gate_expand_index_structure_is_supported(block) - assert not _K1._gate_expand_index_is_supported(block) + assert _SO2._gate_expand_index_structure_is_supported(block) + assert not _SO2._gate_expand_index_is_supported(block) @pytest.mark.parametrize( "capability", - sorted(runtime_policy.SUPPORTED_K1_CAPABILITIES), + sorted(runtime_policy.SUPPORTED_SO2_CAPABILITIES), ) def test_every_supported_architecture_has_a_validated_config( capability: tuple[int, int], ) -> None: - config = _K1._architecture_default_config(capability) + config = _SO2._architecture_default_config(capability) assert _validate_runtime_config(config, compute_capability=capability) is None @@ -176,7 +176,7 @@ def test_every_supported_architecture_has_a_validated_config( def test_sm80_family_uses_per_focus_so2_forward( capability: tuple[int, int], ) -> None: - config = _K1._architecture_default_config(capability) + config = _SO2._architecture_default_config(capability) assert config.per_focus_so2_fwd_pair assert not config.native_sm90_path @@ -184,7 +184,7 @@ def test_sm80_family_uses_per_focus_so2_forward( def test_sm90_uses_native_split_complex_path() -> None: - config = _K1._architecture_default_config((9, 0)) + config = _SO2._architecture_default_config((9, 0)) assert config.native_sm90_path assert not config.per_focus_so2_fwd_pair @@ -192,7 +192,7 @@ def test_sm90_uses_native_split_complex_path() -> None: def test_sm100_uses_shared_default_profile() -> None: - config = _K1._architecture_default_config((10, 0)) + config = _SO2._architecture_default_config((10, 0)) assert not config.native_sm90_path assert not config.per_focus_so2_fwd_pair @@ -203,7 +203,7 @@ def test_sm100_uses_shared_default_profile() -> None: def test_sm89_and_sm120_use_combined_so2_gate( capability: tuple[int, int], ) -> None: - config = _K1._architecture_default_config(capability) + config = _SO2._architecture_default_config(capability) assert config.combined_so2_gate assert not config.native_sm90_path @@ -213,7 +213,7 @@ def test_sm89_and_sm120_use_combined_so2_gate( def test_combined_so2_gate_rejects_misaligned_contiguous_tensor() -> None: pytest.importorskip("cutlass.cute") pytest.importorskip("cuda.bindings.driver") - from deepmd.pt_expt.kernels.cute.sezm.k1_kernels.cute_neo_so2_gate_combined_fwd import ( + from deepmd.pt_expt.kernels.cute.sezm.so2.kernels.combined_gate_forward import ( _require_16_byte_alignment, ) @@ -228,7 +228,7 @@ def test_combined_so2_gate_rejects_misaligned_contiguous_tensor() -> None: def test_runtime_config_contains_only_reached_selectors() -> None: - assert {field.name for field in fields(NeoK1RuntimeConfig)} == { + assert {field.name for field in fields(NeoSO2RuntimeConfig)} == { "native_sm90_path", "per_focus_so2_fwd_pair", "combined_so2_gate", @@ -236,26 +236,26 @@ def test_runtime_config_contains_only_reached_selectors() -> None: def test_runtime_config_rejects_incompatible_capability() -> None: - config = _K1._architecture_default_config((8, 0)) + config = _SO2._architecture_default_config((8, 0)) with pytest.raises(RuntimeError, match="supported compute capability"): _validate_runtime_config(config, compute_capability=(7, 5)) def test_runtime_config_rejects_wrong_architecture_profile() -> None: - sm80_config = _K1._architecture_default_config((8, 0)) + sm80_config = _SO2._architecture_default_config((8, 0)) with pytest.raises(RuntimeError, match="per-focus SO2"): _validate_runtime_config(sm80_config, compute_capability=(12, 0)) - sm90_config = _K1._architecture_default_config((9, 0)) - with pytest.raises(RuntimeError, match="native SM90 K1"): + sm90_config = _SO2._architecture_default_config((9, 0)) + with pytest.raises(RuntimeError, match="native SM90 SO2"): _validate_runtime_config(sm90_config, compute_capability=(8, 0)) - sm120_config = _K1._architecture_default_config((12, 0)) + sm120_config = _SO2._architecture_default_config((12, 0)) with pytest.raises(RuntimeError, match="combined SO2/gate"): _validate_runtime_config(sm120_config, compute_capability=(10, 0)) with pytest.raises(RuntimeError, match="combined SO2/gate"): - _validate_runtime_config(NeoK1RuntimeConfig(), compute_capability=(12, 0)) + _validate_runtime_config(NeoSO2RuntimeConfig(), compute_capability=(12, 0)) @pytest.mark.parametrize( @@ -285,14 +285,14 @@ def test_packed_edge_eligibility_requires_sorted_strict_fp32() -> None: "runtime_dtypes": (torch.float32, torch.float32), } - assert _K1.packed_wigner_edges_eligible(**kwargs) - assert not _K1.packed_wigner_edges_eligible( + assert _SO2.packed_wigner_edges_eligible(**kwargs) + assert not _SO2.packed_wigner_edges_eligible( **{**kwargs, "destinations_sorted": False} ) - assert not _K1.packed_wigner_edges_eligible( + assert not _SO2.packed_wigner_edges_eligible( **{**kwargs, "runtime_dtypes": (torch.float64,)} ) - assert _K1.packed_wigner_edges_eligible( + assert _SO2.packed_wigner_edges_eligible( **{**kwargs, "edge_count": 3, "node_count": 4} ) @@ -358,7 +358,7 @@ def test_ineligible_runtime_contracts_fall_back( monkeypatch.setenv("DP_CUTE_INFER", "1") block = _neo_like_block().eval() - assert not _K1.is_neo_k1_runtime_eligible( + assert not _SO2.is_neo_so2_runtime_eligible( block, training=training, device=torch.device("cuda", 0), @@ -375,7 +375,7 @@ def test_runtime_contract_allows_fewer_edges_than_nodes( monkeypatch.setenv("DP_CUTE_INFER", "1") block = _neo_like_block().eval() - assert _K1.is_neo_k1_runtime_eligible( + assert _SO2.is_neo_so2_runtime_eligible( block, training=False, device=torch.device("cuda", 0), @@ -402,7 +402,7 @@ def _cute_cuda_runtime_available() -> bool: reason="CuTe attention-prelude regression requires CUDA and CuTe DSL", ) def test_attention_prelude_initializes_all_qk_rows_when_edges_are_sparse() -> None: - from deepmd.pt_expt.kernels.cute.sezm.k1_kernels.cute_neo_focus_src_backward import ( + from deepmd.pt_expt.kernels.cute.sezm.so2.kernels.focus_source_backward import ( compile_neo_attention_prelude_forward, ) @@ -490,9 +490,9 @@ def test_exact_runtime_contract_is_eligible_without_autocast( "destinations_sorted": True, } - assert _K1.is_neo_k1_runtime_eligible(block, **kwargs) + assert _SO2.is_neo_so2_runtime_eligible(block, **kwargs) monkeypatch.setattr(torch, "is_autocast_enabled", lambda device_type: True) - assert not _K1.is_neo_k1_runtime_eligible(block, **kwargs) + assert not _SO2.is_neo_so2_runtime_eligible(block, **kwargs) def test_runtime_contract_rejects_tf32_matmul_policy( @@ -511,16 +511,16 @@ def test_runtime_contract_rejects_tf32_matmul_policy( } monkeypatch.setattr(runtime_policy, "uses_strict_fp32_matmul", lambda: True) - assert _K1.is_neo_k1_runtime_eligible(block, **kwargs) + assert _SO2.is_neo_so2_runtime_eligible(block, **kwargs) monkeypatch.setattr(runtime_policy, "uses_strict_fp32_matmul", lambda: False) - assert not _K1.is_neo_k1_runtime_eligible(block, **kwargs) + assert not _SO2.is_neo_so2_runtime_eligible(block, **kwargs) def test_fake_native_gradient_layout_matches_runtime_contract() -> None: x = torch.empty(5, 16, 1, 32, device="cpu") - grad = _K1._fake_x_wide_grad_like(x, skip=True) + grad = _SO2._fake_x_wide_grad_like(x, skip=True) assert grad.shape == x.shape assert grad.stride() == (32, 5 * 32, 32, 1) @@ -532,7 +532,7 @@ def test_custom_op_runtime_canonicalizes_misaligned_contiguous_view() -> None: assert offset_view.is_contiguous() assert offset_view.storage_offset() == 1 - actual = _K1._aligned_contiguous(offset_view) + actual = _SO2._aligned_contiguous(offset_view) torch.testing.assert_close(actual, offset_view) assert actual.is_contiguous() @@ -543,7 +543,7 @@ def test_custom_op_runtime_canonicalizes_misaligned_contiguous_view() -> None: def test_custom_op_runtime_preserves_aligned_compact_tensor() -> None: tensor = torch.arange(16, dtype=torch.float32, device="cpu") - assert _K1._aligned_contiguous(tensor) is tensor + assert _SO2._aligned_contiguous(tensor) is tensor def test_stack_cache_retains_only_backward_state() -> None: diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_dst_ptr.py b/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py similarity index 91% rename from source/tests/pt/model/test_descriptor_sezm_cute_k1_dst_ptr.py rename to source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py index ea4f573c89..9c4c1e3212 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_dst_ptr.py +++ b/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""Contracts for asynchronous K1 destination-pointer construction.""" +"""Contracts for asynchronous SO2 destination-pointer construction.""" from __future__ import ( annotations, @@ -9,9 +9,6 @@ import ast import importlib import unittest -from pathlib import ( - Path, -) from types import ( SimpleNamespace, ) @@ -22,15 +19,18 @@ mock, ) +from ._paths import ( + CUTE_ROOT, +) + try: import torch except ModuleNotFoundError: # pragma: no cover - lightweight source-test host torch = None -REPO_ROOT = Path(__file__).resolve().parents[4] -K1_PATH = REPO_ROOT / "deepmd/pt_expt/kernels/cute/sezm/k1.py" -EDGE_CACHE_PATH = REPO_ROOT / "deepmd/pt/model/descriptor/sezm_nn/edge_cache.py" +SO2_PATH = CUTE_ROOT / "so2" / "operation.py" +SO2_METADATA_PATH = CUTE_ROOT / "so2" / "metadata.py" class _Tensor: @@ -45,7 +45,7 @@ def _function(tree: ast.AST, name: str) -> ast.FunctionDef: def _load_extracted_dst_ptr_function(*, strict: bool = False): - source = K1_PATH.read_text(encoding="utf-8") + source = SO2_PATH.read_text(encoding="utf-8") tree = ast.parse(source) function = _function(tree, "_dst_ptr_from_sorted") module = ast.Module(body=[function], type_ignores=[]) @@ -57,18 +57,18 @@ def _load_extracted_dst_ptr_function(*, strict: bool = False): is_cute_strict_enabled=lambda: strict, ), } - exec(compile(module, str(K1_PATH), "exec"), namespace) + exec(compile(module, str(SO2_PATH), "exec"), namespace) return namespace["_dst_ptr_from_sorted"] def _load_extracted_sorted_metadata_function(): - source = EDGE_CACHE_PATH.read_text(encoding="utf-8") + source = SO2_METADATA_PATH.read_text(encoding="utf-8") tree = ast.parse(source) function = _function(tree, "build_sorted_edge_index_metadata") module = ast.Module(body=[function], type_ignores=[]) ast.fix_missing_locations(module) namespace = {"torch": torch} - exec(compile(module, str(EDGE_CACHE_PATH), "exec"), namespace) + exec(compile(module, str(SO2_METADATA_PATH), "exec"), namespace) return namespace["build_sorted_edge_index_metadata"] @@ -104,7 +104,7 @@ def bincount(self, *args, **kwargs): raise AssertionError("CUDA bincount must not construct sorted dst_ptr") -class TestK1DstPtrExtracted(unittest.TestCase): +class TestSO2DstPtrExtracted(unittest.TestCase): def test_sorted_path_operator_and_output_contract(self): helper = _load_extracted_dst_ptr_function() torch_module = _RecordingTorch() @@ -143,7 +143,7 @@ def test_unsorted_path_returns_before_tensor_work(self): @unittest.skipIf(torch is None, "destination-pointer differential requires PyTorch") -class TestK1DstPtrTorch(unittest.TestCase): +class TestSO2DstPtrTorch(unittest.TestCase): @classmethod def setUpClass(cls): cls.helper = staticmethod(_load_extracted_dst_ptr_function()) @@ -260,27 +260,27 @@ def build_ptr(dst): @unittest.skipUnless( torch is not None and torch.cuda.is_available(), - "K1 custom-op registry lifetime regression requires CUDA", + "SO2 custom-op registry lifetime regression requires CUDA", ) def test_compile_cold_registration_survives_into_packed_custom_op_runtime(self): assert torch is not None - k1 = importlib.import_module("deepmd.pt_expt.kernels.cute.sezm.k1") - prior_registry = dict(k1._REGISTRY) - prior_next_handle = k1._NEXT_HANDLE + so2 = importlib.import_module("deepmd.pt_expt.kernels.cute.sezm.so2.operation") + prior_registry = dict(so2._REGISTRY) + prior_next_handle = so2._NEXT_HANDLE def cleanup(): torch._dynamo.reset() - k1._PACKED_RUNNER_CACHE.clear() - k1._REGISTRY.clear() - k1._REGISTRY.update(prior_registry) - k1._NEXT_HANDLE = prior_next_handle + so2._PACKED_RUNNER_CACHE.clear() + so2._REGISTRY.clear() + so2._REGISTRY.update(prior_registry) + so2._NEXT_HANDLE = prior_next_handle self.addCleanup(cleanup) - k1._REGISTRY.clear() - k1._PACKED_RUNNER_CACHE.clear() + so2._REGISTRY.clear() + so2._PACKED_RUNNER_CACHE.clear() torch._dynamo.reset() block = torch.nn.Module() - config = k1.NeoK1RuntimeConfig() + config = so2.NeoSO2RuntimeConfig() runtime_handles = [] class FakeRunner: @@ -288,7 +288,7 @@ class FakeRunner: def fake_build(handle, x_arg, *args): del args - entry = k1._REGISTRY[int(handle)] + entry = so2._REGISTRY[int(handle)] self.assertIs(entry.block, block) runtime_handles.append(int(handle)) runner = FakeRunner() @@ -304,21 +304,21 @@ def invoke( src_arg, dst_arg, ): - state = getattr(block, "_deepmd_cute_k1_state", None) + state = getattr(block, "_deepmd_cute_so2_state", None) if state is None: - state = k1._register_cute_k1_state( + state = so2._register_cute_so2_state( block, x_arg.device.index, config, ) - dst_ptr = k1._dst_ptr_from_sorted( + dst_ptr = so2._dst_ptr_from_sorted( torch, dst_arg, x_arg.shape[0], destinations_sorted=True, ) edge_src_gate = edge_arg.new_empty((0,)) - return k1.cute_k1( + return so2.cute_so2( state.handle, x_arg, d_arg, @@ -345,15 +345,15 @@ def compiled_entry(*args): src = torch.tensor((0, 1, 0), dtype=torch.int64, device=device) dst = torch.tensor((0, 0, 1), dtype=torch.int64, device=device) - with mock.patch.object(k1, "_build_runner", new=fake_build): + with mock.patch.object(so2, "_build_runner", new=fake_build): compiled = torch.compile(compiled_entry, backend="eager", dynamic=True) first = compiled(x, d_full, dt_full, radial, edge_env, src, dst) second = compiled(x, d_full, dt_full, radial, edge_env, src, dst) - state = block._deepmd_cute_k1_state + state = block._deepmd_cute_so2_state self.assertEqual(runtime_handles, [state.handle, state.handle]) - self.assertEqual(list(k1._REGISTRY), [state.handle]) - self.assertIs(k1._REGISTRY[state.handle].block, block) + self.assertEqual(list(so2._REGISTRY), [state.handle]) + self.assertIs(so2._REGISTRY[state.handle].block, block) torch.testing.assert_close(first, x, rtol=0.0, atol=0.0) torch.testing.assert_close(second, x, rtol=0.0, atol=0.0) @@ -551,7 +551,7 @@ def test_dynamic_node_and_edge_counts_build_independent_metadata(self): self.assertNotEqual(first[1].data_ptr(), second[1].data_ptr()) self.assertNotEqual(first[2].data_ptr(), second[2].data_ptr()) - def test_explicit_tensor_metadata_survives_disabled_k1_boundary(self): + def test_explicit_tensor_metadata_survives_disabled_so2_boundary(self): assert torch is not None edge_cache = importlib.import_module( "deepmd.pt.model.descriptor.sezm_nn.edge_cache" @@ -560,7 +560,7 @@ def test_explicit_tensor_metadata_survives_disabled_k1_boundary(self): if dynamo is None: self.skipTest("torch._dynamo is unavailable") - def k1_boundary( + def so2_boundary( x, cache, radial, @@ -580,11 +580,11 @@ def k1_boundary( + source_ptr.sum() ) - disabled_k1_boundary = torch.compiler.disable(k1_boundary) + disabled_so2_boundary = torch.compiler.disable(so2_boundary) def consume(cache, x, radial, dst_ptr, source_order, source_ptr): before_break = x + cache.edge_env.sum() - return disabled_k1_boundary( + return disabled_so2_boundary( before_break, cache, radial, diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py b/source/tests/pt/model/sezm_cute/test_so2_gate_structural.py similarity index 94% rename from source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py rename to source/tests/pt/model/sezm_cute/test_so2_gate_structural.py index e939d57b8e..d8ef04185f 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py +++ b/source/tests/pt/model/sezm_cute/test_so2_gate_structural.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""Numerical differentials for the structural Neo K1 gate.""" +"""Numerical differentials for the structural Neo SO2 gate.""" from __future__ import ( annotations, @@ -10,15 +10,13 @@ import torch -from deepmd.pt_expt.kernels.cute.sezm import ( - k1_gate_structural, -) +from deepmd.pt_expt.kernels.cute.sezm.so2 import structural_gate as so2_gate_structural class TestStructuralGateHelpers(unittest.TestCase): @classmethod def setUpClass(cls): - cls.helper = k1_gate_structural + cls.helper = so2_gate_structural def setUp(self): torch.manual_seed(20260703) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py b/source/tests/pt/model/sezm_cute/test_so2_inplace_residual.py similarity index 98% rename from source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py rename to source/tests/pt/model/sezm_cute/test_so2_inplace_residual.py index 82b559180b..65e81cd758 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py +++ b/source/tests/pt/model/sezm_cute/test_so2_inplace_residual.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""Contracts for the in-place K1 SO2 residual/addmm helper.""" +"""Contracts for the in-place CuTe SO2 residual/addmm helper.""" from __future__ import ( annotations, @@ -13,13 +13,11 @@ TorchDispatchMode, ) -from deepmd.pt_expt.kernels.cute.sezm import ( - k1_so2linear, -) +from deepmd.pt_expt.kernels.cute.sezm.so2 import linear as so2_linear def _load_candidate(): - return k1_so2linear + return so2_linear def _out_of_place_reference(grad_out, residual, w0_t, wpair_t): @@ -55,7 +53,7 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None): return func(*args, **kwargs) -class TestK1InplaceResidual(unittest.TestCase): +class TestSO2InplaceResidual(unittest.TestCase): def setUp(self): torch.manual_seed(20260630) self.edge_count = 7 diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_manual_adjoints.py b/source/tests/pt/model/sezm_cute/test_so2_manual_adjoints.py similarity index 92% rename from source/tests/pt/model/test_descriptor_sezm_cute_k1_manual_adjoints.py rename to source/tests/pt/model/sezm_cute/test_so2_manual_adjoints.py index 28061d6cd2..c1412af664 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_manual_adjoints.py +++ b/source/tests/pt/model/sezm_cute/test_so2_manual_adjoints.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""Parity tests for the real-module adjoints used by Neo K1 backward.""" +"""Parity tests for the real-module adjoints used by Neo SO2 backward.""" from __future__ import ( annotations, @@ -29,9 +29,7 @@ from deepmd.pt.utils import ( env, ) -from deepmd.pt_expt.kernels.cute.sezm import ( - k1, -) +from deepmd.pt_expt.kernels.cute.sezm.so2 import operation as so2 @pytest.fixture(autouse=True) @@ -80,7 +78,7 @@ def test_equivariant_rmsnorm_manual_adjoint_matches_real_module() -> None: grad_out = _randn(*x.shape) reference_grad = torch.autograd.grad(norm(x), x, grad_out)[0] - actual_grad = k1._equivariant_rmsnorm_backward( + actual_grad = so2._equivariant_rmsnorm_backward( norm, x.detach(), grad_out, @@ -105,9 +103,9 @@ def test_focus_linear_manual_adjoint_matches_real_module() -> None: reference = linear(x) reference_grad = torch.autograd.grad(reference, x, grad_out)[0] - torch.testing.assert_close(k1._focus_linear_forward(linear, x), reference) + torch.testing.assert_close(so2._focus_linear_forward(linear, x), reference) torch.testing.assert_close( - k1._focus_linear_backward_input(linear, grad_out), + so2._focus_linear_backward_input(linear, grad_out), reference_grad, atol=1.0e-12, rtol=1.0e-12, @@ -122,9 +120,9 @@ def test_swiglu_manual_adjoint_matches_real_module() -> None: reference = activation(x) reference_grad = torch.autograd.grad(reference, x, grad_out)[0] - torch.testing.assert_close(k1._swiglu_forward(x), reference) + torch.testing.assert_close(so2._swiglu_forward(x), reference) torch.testing.assert_close( - k1._swiglu_backward_input(x.detach(), grad_out), + so2._swiglu_backward_input(x.detach(), grad_out), reference_grad, atol=1.0e-12, rtol=1.0e-12, @@ -156,7 +154,7 @@ def test_grid_cross_glu_flat_manual_adjoint_matches_real_module() -> None: (query, context), grad_out, ) - actual_query, actual_context = k1._so3_grid_cross_glu_flat_backward( + actual_query, actual_context = so2._so3_grid_cross_glu_flat_backward( net, query.detach(), context.detach(), @@ -196,10 +194,10 @@ def test_sm90_message_grid_forward_and_adjoint_match_real_module() -> None: if tuple(torch.cuda.get_device_capability()) != (9, 0): pytest.skip("SM90 message-grid differential requires an SM90 GPU") - from deepmd.pt_expt.kernels.cute.sezm.k1_message_grid_packed import ( + from deepmd.pt_expt.kernels.cute.sezm.so2.message_grid import ( run_packed_message_grid_forward, ) - from deepmd.pt_expt.kernels.cute.sezm.message_grid_readout_sm90 import ( + from deepmd.pt_expt.kernels.cute.sezm.so2.sm90.message_grid_readout import ( prepare_sm90_message_grid_state, run_sm90_message_grid_backward, ) @@ -301,7 +299,7 @@ def test_sm90_message_grid_forward_and_adjoint_match_real_module() -> None: reason="Q/K CuTe adjoint regression requires CUDA and CuTe DSL", ) def test_qk_manual_adjoint_matches_autograd_with_sparse_edges() -> None: - from deepmd.pt_expt.kernels.cute.sezm.k1_kernels.cute_neo_qk_edge import ( + from deepmd.pt_expt.kernels.cute.sezm.so2.kernels.qk_edge import ( compile_neo_qk_edge_backward, compile_neo_qk_node_input_adjoint, ) @@ -354,7 +352,7 @@ def test_qk_manual_adjoint_matches_autograd_with_sparse_edges() -> None: qk_edge_backward=compile_neo_qk_edge_backward(scale), qk_node_input_adjoint=compile_neo_qk_node_input_adjoint(eps), ) - actual = k1._qk_manual_backward(runner, grad_logits) + actual = so2._qk_manual_backward(runner, grad_logits) torch.cuda.synchronize(device) torch.testing.assert_close(actual, reference, atol=5.0e-5, rtol=5.0e-5) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_packed_impl.py b/source/tests/pt/model/sezm_cute/test_so2_packed_impl.py similarity index 93% rename from source/tests/pt/model/test_descriptor_sezm_cute_k1_packed_impl.py rename to source/tests/pt/model/sezm_cute/test_so2_packed_impl.py index 459cd26dfa..95099e323c 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_packed_impl.py +++ b/source/tests/pt/model/sezm_cute/test_so2_packed_impl.py @@ -39,7 +39,7 @@ def _randn( def _dense_wigner(edge_count: int, generator: torch.Generator) -> torch.Tensor: - from deepmd.pt_expt.kernels.cute.sezm import k1_wigner_layout as layout + from deepmd.pt_expt.kernels.cute.sezm.so2 import wigner_layout as layout dense = torch.zeros( edge_count, @@ -73,10 +73,10 @@ def test_panel_native_quaternion_backward_matches_dense_torch(): from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( WignerDCalculator, ) - from deepmd.pt_expt.kernels.cute.sezm import k1_wigner_layout as layout from deepmd.pt_expt.kernels.cute.sezm import ( - k4_wignerd, + wignerd, ) + from deepmd.pt_expt.kernels.cute.sezm.so2 import wigner_layout as layout generator = torch.Generator(device="cuda").manual_seed(20260702) q_value = _randn(3, 4, generator=generator, scale=1.0) @@ -84,7 +84,7 @@ def test_panel_native_quaternion_backward_matches_dense_torch(): grad_panel = _randn(3, 46, generator=generator) q_panel = q_value.detach().requires_grad_(True) - panel = k4_wignerd._wignerd_panel_op(q_panel) + panel = wignerd._wignerd_panel_op(q_panel) grad_q_panel = torch.autograd.grad((panel * grad_panel).sum(), q_panel)[0] q_dense = q_value.detach().requires_grad_(True) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py b/source/tests/pt/model/sezm_cute/test_so2_phase_c_inplace_adjoint.py similarity index 98% rename from source/tests/pt/model/test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py rename to source/tests/pt/model/sezm_cute/test_so2_phase_c_inplace_adjoint.py index e8bc32e950..3778130025 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_phase_c_inplace_adjoint.py +++ b/source/tests/pt/model/sezm_cute/test_so2_phase_c_inplace_adjoint.py @@ -54,7 +54,7 @@ def _phase_c_api(): pytest.importorskip("cutlass") - from deepmd.pt_expt.kernels.cute.sezm.k1_kernels.cute_neo_phase_c_backward_layout_runner import ( + from deepmd.pt_expt.kernels.cute.sezm.so2.phase_c import ( CuteNeoPhaseCBackwardLayout, NeoPhaseCBackwardLayoutOutputs, ) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_focus_fusion.py b/source/tests/pt/model/sezm_cute/test_so2_radial_focus_fusion.py similarity index 96% rename from source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_focus_fusion.py rename to source/tests/pt/model/sezm_cute/test_so2_radial_focus_fusion.py index bdccca901d..996a7749d5 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_focus_fusion.py +++ b/source/tests/pt/model/sezm_cute/test_so2_radial_focus_fusion.py @@ -19,7 +19,7 @@ def _has_supported_gpu() -> bool: def test_source_csr_preserves_equal_source_order() -> None: - from deepmd.pt_expt.kernels.cute.sezm.k1_radial_phase_a_node import ( + from deepmd.pt_expt.kernels.cute.sezm.so2.radial_phase_a import ( build_source_csr, ) @@ -35,7 +35,7 @@ def test_source_csr_preserves_equal_source_order() -> None: @unittest.skipUnless(_has_supported_gpu(), "requires an SM80 or SM90 CUDA device") class TestRadialFocusSourceFusion(unittest.TestCase): def test_fused_load_matches_standalone_scalar_lane_add(self) -> None: - from deepmd.pt_expt.kernels.cute.sezm.k1_radial_phase_a_node import ( + from deepmd.pt_expt.kernels.cute.sezm.so2.radial_phase_a import ( build_source_csr, prepare_batched_radial_projection_weight, run_neo_radial_phase_a_backward_node_tiled, diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_pitch68_safe.py b/source/tests/pt/model/sezm_cute/test_so2_radial_pitch68_safe.py similarity index 94% rename from source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_pitch68_safe.py rename to source/tests/pt/model/sezm_cute/test_so2_radial_pitch68_safe.py index 9815909e37..8c3ea3358c 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_pitch68_safe.py +++ b/source/tests/pt/model/sezm_cute/test_so2_radial_pitch68_safe.py @@ -8,14 +8,13 @@ import ast import unittest -from pathlib import ( - Path, + +from ._paths import ( + CUTE_ROOT, ) -REPO_ROOT = Path(__file__).resolve().parents[4] -CUTE_PATH = REPO_ROOT / "deepmd/pt_expt/kernels/cute/sezm" -WRAPPER_PATH = CUTE_PATH / "k1_radial_phase_a_node.py" -KERNEL_PATH = CUTE_PATH / "k1_kernels/cute_neo_radial_phase_a_backward_node.py" +WRAPPER_PATH = CUTE_ROOT / "so2" / "radial_phase_a.py" +KERNEL_PATH = CUTE_ROOT / "so2" / "kernels" / "radial_phase_a_backward.py" WARP_HELPERS = {"_warp_owned_grad_compact", "_warp_owned_grad_d"} diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py b/source/tests/pt/model/sezm_cute/test_so2_radial_projection.py similarity index 96% rename from source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py rename to source/tests/pt/model/sezm_cute/test_so2_radial_projection.py index 495838263f..13f84f8af9 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py +++ b/source/tests/pt/model/sezm_cute/test_so2_radial_projection.py @@ -16,10 +16,8 @@ import torch -from deepmd.pt_expt.kernels.cute.sezm import ( - k1_radial_phase_a_node as _RADIAL_PROJECTION, -) -from deepmd.pt_expt.kernels.cute.sezm import k1_runner as _K1_RUNNER +from deepmd.pt_expt.kernels.cute.sezm.so2 import radial_phase_a as _RADIAL_PROJECTION +from deepmd.pt_expt.kernels.cute.sezm.so2 import runner as _SO2_RUNNER COMPACT_WIDTH = _RADIAL_PROJECTION.COMPACT_WIDTH FOCUS_COUNT = _RADIAL_PROJECTION.FOCUS_COUNT @@ -30,7 +28,7 @@ prepare_batched_radial_projection_weight = ( _RADIAL_PROJECTION.prepare_batched_radial_projection_weight ) -_batched_radial_projection_weight = _K1_RUNNER._batched_radial_projection_weight +_batched_radial_projection_weight = _SO2_RUNNER._batched_radial_projection_weight class TestBatchedRadialProjection(unittest.TestCase): diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py b/source/tests/pt/model/sezm_cute/test_so2_structural_gate_vec4_sm80.py similarity index 97% rename from source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py rename to source/tests/pt/model/sezm_cute/test_so2_structural_gate_vec4_sm80.py index b8009a1e7e..e3dd389da1 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py +++ b/source/tests/pt/model/sezm_cute/test_so2_structural_gate_vec4_sm80.py @@ -12,13 +12,12 @@ import random import struct import unittest -from pathlib import ( - Path, + +from ._paths import ( + CUTE_ROOT, ) -REPO_ROOT = Path(__file__).resolve().parents[4] -CUTE_ROOT = REPO_ROOT / "deepmd/pt_expt/kernels/cute/sezm" -KERNEL_PATH = CUTE_ROOT / "k1_kernels/cute_neo_gate_split_structural_vec4_sm80.py" +KERNEL_PATH = CUTE_ROOT / "so2" / "kernels" / "structural_gate_sm80.py" DYNAMIC_EDGE_COUNTS = (0, 1, 7, 8, 9, 31, 32, 33) @@ -224,7 +223,7 @@ def constant_value(node: ast.expr, namespace: dict[str, object]) -> object: def _load_wrapper_module(): return importlib.import_module( - "deepmd.pt_expt.kernels.cute.sezm.k1_gate_structural" + "deepmd.pt_expt.kernels.cute.sezm.so2.structural_gate" ) @@ -257,7 +256,7 @@ def test_kernel_is_strict_fp32_edge_major_only(self): self.assertIn("cutlass.Float32", self.kernel_source) self.assertIn('"assumed_align": 16', self.kernel_source) self.assertIn("_guard_vec4_dispatch", self.kernel_source) - self.assertIn("runtime_policy.SUPPORTED_K1_CAPABILITIES", self.kernel_source) + self.assertIn("runtime_policy.SUPPORTED_SO2_CAPABILITIES", self.kernel_source) class TestSM80StructuralGateVec4Arithmetic(unittest.TestCase): @@ -400,12 +399,12 @@ def test_empty_tensor_does_not_launch(self): and torch.cuda.is_available() and tuple(torch.cuda.get_device_capability()) in {(8, 0), (8, 6), (8, 9), (9, 0), (10, 0), (12, 0)}, - "A supported Neo K1 CUDA runtime is required", + "A supported Neo SO2 CUDA runtime is required", ) class TestSM80StructuralGateVec4CudaDifferential(unittest.TestCase): @classmethod def setUpClass(cls): - from deepmd.pt_expt.kernels.cute.sezm.k1_kernels.cute_neo_gate_split_structural_vec4_sm80 import ( + from deepmd.pt_expt.kernels.cute.sezm.so2.kernels.structural_gate_sm80 import ( compile_neo_gate_split_structural_vec4_sm80_backward, compile_neo_gate_split_structural_vec4_sm80_forward, ) diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_memory_reuse.py b/source/tests/pt/model/sezm_cute/test_so2_structural_memory_reuse.py similarity index 78% rename from source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_memory_reuse.py rename to source/tests/pt/model/sezm_cute/test_so2_structural_memory_reuse.py index f0fa392dc9..92ec6ded3e 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_memory_reuse.py +++ b/source/tests/pt/model/sezm_cute/test_so2_structural_memory_reuse.py @@ -1,49 +1,30 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: LGPL-3.0-or-later -"""Behavioral contracts for K1 structural storage reuse.""" +"""Behavioral contracts for SO2 structural storage reuse.""" from __future__ import ( annotations, ) -import importlib.util -import sys -from pathlib import ( - Path, -) from types import ( SimpleNamespace, ) -from .test_descriptor_sezm_cute_k1 import ( - _K1, +from deepmd.pt_expt.kernels.cute.sezm.so2 import linear as so2_linear +from deepmd.pt_expt.kernels.cute.sezm.so2 import structural_gate as so2_gate_structural + +from .test_so2 import ( + _SO2, NeoFullCuteBackward, - NeoK1BackwardWorkspace, - NeoK1RuntimeConfig, + NeoSO2BackwardWorkspace, + NeoSO2RuntimeConfig, StackCache, _validate_runtime_config, ) -REPO_ROOT = Path(__file__).resolve().parents[4] -STRUCTURAL_HELPER = REPO_ROOT / "deepmd/pt_expt/kernels/cute/sezm/k1_gate_structural.py" -SO2_HELPER = REPO_ROOT / "deepmd/pt_expt/kernels/cute/sezm/k1_so2linear.py" - - -def _load_module(name: str, path: Path): - spec = importlib.util.spec_from_file_location(name, path) - if spec is None or spec.loader is None: - raise RuntimeError(f"could not load {path}") - module = importlib.util.module_from_spec(spec) - sys.modules[name] = module - try: - spec.loader.exec_module(module) - finally: - sys.modules.pop(name, None) - return module - -def _valid_structural_config() -> NeoK1RuntimeConfig: - return NeoK1RuntimeConfig(per_focus_so2_fwd_pair=True) +def _valid_structural_config() -> NeoSO2RuntimeConfig: + return NeoSO2RuntimeConfig(per_focus_so2_fwd_pair=True) def _storage_ptr(tensor) -> int: @@ -51,7 +32,7 @@ def _storage_ptr(tensor) -> int: def _workspace_inputs(edge_count: int): - torch = _K1.torch + torch = _SO2.torch return { "edge_count": edge_count, "node_count": 3, @@ -80,14 +61,13 @@ def test_structural_memory_reuse_profile_is_valid() -> None: def test_final_so2_single_input_fold_matches_residual_plus_linear(): - torch = _K1.torch - helper = _load_module("sezm_cute_k1_so2_test", SO2_HELPER) + torch = _SO2.torch torch.manual_seed(20260705) x_local = torch.randn(3, 2, 10, 32, dtype=torch.float32, device="cpu") linear = _so2_linear(torch) - linear_only = helper.run_neo_so2_linear_manual(linear, x_local) - folded = helper.run_neo_so2_linear_manual( + linear_only = so2_linear.run_neo_so2_linear_manual(linear, x_local) + folded = so2_linear.run_neo_so2_linear_manual( linear, x_local, add_residual=True, @@ -97,14 +77,13 @@ def test_final_so2_single_input_fold_matches_residual_plus_linear(): def test_per_focus_pair_forward_matches_batched_forward(): - torch = _K1.torch - helper = _load_module("sezm_cute_k1_so2_per_focus_test", SO2_HELPER) + torch = _SO2.torch torch.manual_seed(20260718) x_local = torch.randn(17, 2, 10, 32, dtype=torch.float32, device="cpu") linear = _so2_linear(torch) - batched = helper.run_neo_so2_linear_manual(linear, x_local) - per_focus = helper.run_neo_so2_linear_manual( + batched = so2_linear.run_neo_so2_linear_manual(linear, x_local) + per_focus = so2_linear.run_neo_so2_linear_manual( linear, x_local, per_focus_pair=True, @@ -114,7 +93,7 @@ def test_per_focus_pair_forward_matches_batched_forward(): def test_structural_workspace_reuses_saved_slabs_and_gate_panel(): - torch = _K1.torch + torch = _SO2.torch edge_count = 5 phase_c_stack = torch.empty( edge_count, 2, 10, 32, dtype=torch.float32, device="cpu" @@ -124,7 +103,7 @@ def test_structural_workspace_reuses_saved_slabs_and_gate_panel(): 2, edge_count, 3 * 32, dtype=torch.float32, device="cpu" ) - workspace = NeoK1BackwardWorkspace( + workspace = NeoSO2BackwardWorkspace( torch, **_workspace_inputs(edge_count), phase_c_stack=phase_c_stack, @@ -144,14 +123,14 @@ def test_structural_workspace_reuses_saved_slabs_and_gate_panel(): def test_single_input_workspace_omits_phase_c_y_and_keeps_mixed_output_distinct(): - torch = _K1.torch + torch = _SO2.torch edge_count = 5 phase_c_stack = torch.empty( edge_count, 2, 10, 32, dtype=torch.float32, device="cpu" ) radial_scratch = torch.empty(edge_count, 4, 32, dtype=torch.float32, device="cpu") - workspace = NeoK1BackwardWorkspace( + workspace = NeoSO2BackwardWorkspace( torch, **_workspace_inputs(edge_count), phase_c_stack=phase_c_stack, @@ -168,7 +147,7 @@ def test_single_input_workspace_omits_phase_c_y_and_keeps_mixed_output_distinct( def test_runner_routes_saved_structural_buffers_into_lazy_workspace(): - torch = _K1.torch + torch = _SO2.torch edge_count = 5 runner = object.__new__(NeoFullCuteBackward) runner.torch = torch @@ -208,8 +187,7 @@ def test_runner_routes_saved_structural_buffers_into_lazy_workspace(): def test_structural_backward_can_overwrite_consumed_logits(): - torch = _K1.torch - helper = _load_module("sezm_cute_k1_structural_test", STRUCTURAL_HELPER) + torch = _SO2.torch edge_count = 3 grad_out = torch.randn(edge_count, 2, 10, 32, dtype=torch.float32, device="cpu") y = torch.randn_like(grad_out) @@ -222,7 +200,7 @@ def fake_kernel(grad_out_flat, y_flat, logits_in, grad_y_flat, grad_logits): grad_y_flat.copy_(grad_out_flat + y_flat) grad_logits.fill_(7.0) - result = helper.run_structural_gate_backward( + result = so2_gate_structural.run_structural_gate_backward( fake_kernel, grad_out, y, diff --git a/source/tests/pt/model/test_descriptor_sezm_cute_tiled_grid_product.py b/source/tests/pt/model/sezm_cute/test_tiled_grid_product.py similarity index 92% rename from source/tests/pt/model/test_descriptor_sezm_cute_tiled_grid_product.py rename to source/tests/pt/model/sezm_cute/test_tiled_grid_product.py index 3108b5cf6a..3453dc1a85 100644 --- a/source/tests/pt/model/test_descriptor_sezm_cute_tiled_grid_product.py +++ b/source/tests/pt/model/sezm_cute/test_tiled_grid_product.py @@ -8,28 +8,20 @@ import ast import importlib -from pathlib import ( - Path, -) import pytest import torch +from ._paths import ( + CUTE_ROOT, +) + TOL = 5.0e-5 PACKED_COEFF_DIM = 48 GRID_SIZE = 152 SUPPORTED_HIDDEN_CHANNELS = (96, 192) -REPO_ROOT = Path(__file__).resolve().parents[4] -TILED_KERNEL_PATH = ( - REPO_ROOT - / "deepmd/pt_expt/kernels/cute/sezm/output_grid_kernels" - / "cute_tiled_grid_product.py" -) -MESSAGE_GRID_PATH = ( - REPO_ROOT - / "deepmd/pt_expt/kernels/cute/sezm/k1_kernels" - / "cute_neo_message_grid_product.py" -) +TILED_KERNEL_PATH = CUTE_ROOT / "output_grid" / "kernels" / "tiled_product.py" +MESSAGE_GRID_PATH = CUTE_ROOT / "so2" / "kernels" / "message_grid_product.py" def _cute_runtime_skip_reason() -> str | None: @@ -226,7 +218,7 @@ def test_matches_strict_fp32_reference( nodes: int, hidden_channels: int, ): - from deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product import ( run_tiled_output_grid_product, ) @@ -241,7 +233,7 @@ def test_matches_strict_fp32_reference( ) @pytest.mark.parametrize("nodes", [1, 7, 65]) def test_sm80_c96_n48_matches_strict_fp32_reference(self, nodes: int): - from deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product import ( run_tiled_output_grid_product, ) @@ -258,7 +250,7 @@ def test_sm80_c96_n48_matches_strict_fp32_reference(self, nodes: int): @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) def test_one_compile_accepts_symbolic_node_counts(self, hidden_channels: int): - from deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product import ( _compiled_tiled_forward, run_tiled_output_grid_product, ) @@ -284,7 +276,7 @@ def test_one_compile_accepts_symbolic_node_counts(self, hidden_channels: int): @pytest.mark.parametrize("hidden_channels", SUPPORTED_HIDDEN_CHANNELS) def test_launches_on_current_non_default_stream(self, hidden_channels: int): - from deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product import ( run_tiled_output_grid_product, ) @@ -314,7 +306,7 @@ def test_first_backward_matches_strict_fp32_reference( nodes: int, hidden_channels: int, ): - from deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product import ( run_tiled_output_grid_product_backward, ) @@ -340,7 +332,7 @@ def test_sm80_c96_n48_panel_matches_strict_fp32_reference( self, nodes: int, ): - from deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product import ( run_tiled_output_grid_product_backward, ) @@ -363,7 +355,7 @@ def test_one_backward_compile_accepts_symbolic_node_counts( self, hidden_channels: int, ): - from deepmd.pt_expt.kernels.cute.sezm.output_grid_kernels.cute_tiled_grid_product import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid.kernels.tiled_product import ( _compiled_tiled_backward, run_tiled_output_grid_product_backward, ) diff --git a/source/tests/pt/model/test_descriptor_sezm_readout_l0_algebra.py b/source/tests/pt/model/test_descriptor_sezm_readout_l0_algebra.py index a839b7dfc8..02327ab143 100644 --- a/source/tests/pt/model/test_descriptor_sezm_readout_l0_algebra.py +++ b/source/tests/pt/model/test_descriptor_sezm_readout_l0_algebra.py @@ -121,7 +121,7 @@ def test_scalar_readout_matches_row_zero_of_generic_grid_product(): def test_dense_gram_matches_both_projected_input_adjoints(): - from deepmd.pt_expt.kernels.cute.sezm.readout_l0 import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid.readout_l0 import ( build_readout_l0_gram, ) @@ -184,7 +184,7 @@ def test_dense_gram_matches_both_projected_input_adjoints(): def test_dense_gram_forward_matches_row_zero_grid_projection(): - from deepmd.pt_expt.kernels.cute.sezm.readout_l0 import ( + from deepmd.pt_expt.kernels.cute.sezm.output_grid.readout_l0 import ( build_readout_l0_gram, ) diff --git a/source/tests/pt/model/test_descriptor_sezm_train_paths.py b/source/tests/pt/model/test_descriptor_sezm_train_paths.py index 263e74ba08..00150ed826 100644 --- a/source/tests/pt/model/test_descriptor_sezm_train_paths.py +++ b/source/tests/pt/model/test_descriptor_sezm_train_paths.py @@ -141,34 +141,34 @@ def test_triton_mode_gate_binds_each_stage( assert conv.triton_infer_level == infer_level # Rotation and flash wrappers retain their eager implementations when # Triton is unavailable, so their binding follows the mode gates alone. - assert (conv._rotate_to_local_fn is not None) is requested - assert (conv._rotate_back_fn is not None) is requested - assert (conv._flash_atten_fn is not None) is requested - assert conv._flash_atten_trains is (requested and training) + assert (conv.triton_l_1_rotate_to_local is not None) is requested + assert (conv.triton_l_1_rotate_back is not None) is requested + assert (conv.flash_attention is not None) is requested + assert conv.flash_attention_supports_training is (requested and training) # Segment softmax has no wrapper-level fallback and binds only when its # own Triton implementation is importable. - assert (conv._segment_softmax_fn is not None) is ( + assert (conv.triton_l_1_segment_softmax is not None) is ( requested and SEGMENT_SOFTMAX_TRITON_AVAILABLE ) # The rotate-mix front end is bound by a profitability bound on the # hidden width, which this narrow block sits below. assert conv.hidden_channels < 128 - assert conv._triton_rotate_mix is None + assert conv.triton_l_1_rotate_mix is None # The CUDA gate is off, so the value stream stays on the stages. - assert conv._cuda_value_train is None + assert conv.cuda_train_value is None for module in descriptor.modules(): if isinstance(module, SO2Linear): # The fused GEMM additionally needs every |m| block width to align # to its BN=64 tile, which a narrow block does not satisfy. aligned = slices_supported(module._block_diag_slices) - assert (module._block_diag_gemm is not None) is ( + assert (module.triton_l_1_block_diag_gemm is not None) is ( requested and SO2_BLOCK_GEMM_TRITON_AVAILABLE and aligned ) if isinstance(module, DynamicRadialDegreeMixer): # The callable contains its eager fallback, so construction binds it # whenever either mode requests the stage. - assert (module._radial_mix_block is not None) is requested + assert (module.triton_l_1_radial_mix is not None) is requested if isinstance(module, GatedActivation): assert module.triton_train_level == train_level assert module.triton_infer_level == infer_level @@ -210,13 +210,13 @@ def test_cuda_train_gate_binds_the_value_stream(monkeypatch) -> None: ] assert convolutions for conv in convolutions: - assert conv._cuda_value_train is not None + assert conv.cuda_train_value is not None # The two layers are independent: the CUDA value stream does not # switch on any Triton stage, and the attention span stays dense # until the Triton gate asks for it. assert conv.triton_train_level == 0 - assert conv._segment_softmax_fn is None - assert conv._flash_atten_trains is False + assert conv.triton_l_1_segment_softmax is None + assert conv.flash_attention_supports_training is False def test_cuda_triton_train_reuses_packed_wigner_runs(monkeypatch) -> None: @@ -230,14 +230,14 @@ def test_cuda_triton_train_reuses_packed_wigner_runs(monkeypatch) -> None: monkeypatch.setenv("DP_TRITON_TRAIN", "1") descriptor = _make_descriptor(2, [20], 4.0).train() - assert descriptor._packed_wigner_train + assert descriptor.cuda_train_covers_all_blocks assert not descriptor._build_full_wigner() for block in descriptor.blocks: conv = block.so2_conv - assert conv._cuda_value_train is not None - assert conv._flash_atten_fn is not None - assert conv._flash_atten_trains + assert conv.cuda_train_value is not None + assert conv.flash_attention is not None + assert conv.flash_attention_supports_training @pytest.mark.parametrize("gate_name", ["DP_TRITON_TRAIN", "DP_TRITON_INFER"]) @@ -262,7 +262,7 @@ def test_grid_pair_train_follows_its_gate_and_slot_bound( and GRID_PAIR_TRITON_AVAILABLE and slots >= 75 ) - assert (net._grid_pair_train_fn is not None) is expected + assert (net.triton_train_l_1_grid_pair is not None) is expected @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @@ -326,11 +326,11 @@ def test_source_gated_flash_retains_dense_rotations(self, monkeypatch) -> None: for module in accelerated.modules() if isinstance(module, SO2Convolution) ) - if conv._cuda_conv_fn is None: + if conv.cuda_infer_l_2_conv is None: pytest.skip("the descriptor layout has no fused CUDA convolution") - assert conv._flash_atten_fn is not None - assert conv._cuda_value_train is None - assert not accelerated._wigner_free_conv + assert conv.flash_attention is not None + assert conv.cuda_train_value is None + assert not accelerated.cuda_infer_l_2_covers_all_blocks assert accelerated._build_full_wigner() output, gradient = self._inference_step(accelerated) @@ -362,9 +362,9 @@ def test_training_step_matches_the_dense_path(self, monkeypatch, path) -> None: module for module in fused.modules() if isinstance(module, SO2Convolution) ) if path in ("cuda", "cuda-triton"): - assert conv._cuda_value_train is not None + assert conv.cuda_train_value is not None if path in ("triton", "cuda-triton"): - assert conv._segment_softmax_fn is not None + assert conv.triton_l_1_segment_softmax is not None fused_objective, fused_gradient = self._step(fused) np.testing.assert_allclose( diff --git a/source/tests/pt/model/test_descriptor_sezm_triton.py b/source/tests/pt/model/test_descriptor_sezm_triton.py index 796d1c8178..6f71c6a673 100644 --- a/source/tests/pt/model/test_descriptor_sezm_triton.py +++ b/source/tests/pt/model/test_descriptor_sezm_triton.py @@ -566,7 +566,7 @@ def test_reference_matches_module_eager_path(self): with self.subTest(lmax=lmax): mixer = self._mixer(lmax) # Force the dense scatter path regardless of the ambient flag. - mixer._radial_mix_block = None + mixer.triton_l_1_radial_mix = None x_local, radial_feat, compact = self._inputs(mixer, seed=lmax) with torch.no_grad(): module_out = mixer(x_local, radial_feat) @@ -898,7 +898,7 @@ def test_prepared_weights_follow_device_without_entering_state_dict(self) -> Non conv = self._build_conv(*self.CASES[1]).cpu() value_path = make_triton_value_path(conv) self.assertIsNotNone(value_path) - conv._triton_value_path = value_path + conv.triton_infer_l_2_value = value_path buffer_names = ( "_triton_w0_all", @@ -1056,7 +1056,7 @@ def test_wigner_calculator_matches_reference_chain(self): .to("cuda") .eval() ) - self.assertTrue(fused_calc._use_triton_monomials) + self.assertTrue(fused_calc.triton_infer_l_1_monomials) got, got_t = fused_calc(q) want, want_t = ref_calc(q) torch.testing.assert_close(got, want, atol=1e-5, rtol=1e-5) diff --git a/source/tests/pt/model/test_mlp.py b/source/tests/pt/model/test_mlp.py index 91a283bfd8..55492b5f7f 100644 --- a/source/tests/pt/model/test_mlp.py +++ b/source/tests/pt/model/test_mlp.py @@ -89,7 +89,7 @@ def test_jit(self) -> None: ml1 = MLPLayer.deserialize(ml.serialize()) model = torch.jit.script(ml1) - def test_thin_k1_eval_linear_traces_without_addmm(self) -> None: + def test_thin_so2_eval_linear_traces_without_addmm(self) -> None: from torch.fx.experimental.proxy_tensor import ( make_fx, ) @@ -107,7 +107,7 @@ def test_thin_k1_eval_linear_traces_without_addmm(self) -> None: value = torch.randn(7, 5, device=env.DEVICE, dtype=torch.float32) environment = { "DP_CUTE_INFER": "1", - "DP_CUTE_K1_THIN_WRAPPER": "1", + "DP_CUTE_SO2_THIN_WRAPPER": "1", } with mock.patch.dict(os.environ, environment, clear=False): @@ -127,7 +127,7 @@ def test_thin_k1_eval_linear_traces_without_addmm(self) -> None: self.assertIn(torch.ops.aten.mm.default, targets) self.assertNotIn(torch.ops.aten.addmm.default, targets) - def test_thin_k1_eval_linear_is_scoped_to_marked_model(self) -> None: + def test_thin_so2_eval_linear_is_scoped_to_marked_model(self) -> None: layer = MLPLayer( 5, 8, @@ -140,7 +140,7 @@ def test_thin_k1_eval_linear_is_scoped_to_marked_model(self) -> None: value = torch.randn(7, 5, device=env.DEVICE, dtype=torch.float32) environment = { "DP_CUTE_INFER": "1", - "DP_CUTE_K1_THIN_WRAPPER": "1", + "DP_CUTE_SO2_THIN_WRAPPER": "1", } with ( @@ -156,7 +156,7 @@ def test_thin_k1_eval_linear_is_scoped_to_marked_model(self) -> None: expected = mlp_module.F.linear(value, layer.matrix.t(), layer.bias) torch.testing.assert_close(actual, expected) - def test_thin_k1_linear_defaults_on_only_for_sm80(self) -> None: + def test_thin_so2_linear_defaults_on_only_for_sm80(self) -> None: with ( mock.patch.dict( os.environ, @@ -166,21 +166,21 @@ def test_thin_k1_linear_defaults_on_only_for_sm80(self) -> None: mock.patch.object(torch.cuda, "is_available", return_value=True), mock.patch.object(torch.cuda, "get_device_capability", return_value=(8, 0)), ): - self.assertTrue(mlp_module._use_k1_compile_visible_linear()) + self.assertTrue(mlp_module._use_so2_compile_visible_linear()) with ( mock.patch.dict( os.environ, { "DP_CUTE_INFER": "1", - "DP_CUTE_K1_THIN_WRAPPER": "0", + "DP_CUTE_SO2_THIN_WRAPPER": "0", }, clear=True, ), mock.patch.object(torch.cuda, "is_available", return_value=True), mock.patch.object(torch.cuda, "get_device_capability", return_value=(8, 0)), ): - self.assertFalse(mlp_module._use_k1_compile_visible_linear()) + self.assertFalse(mlp_module._use_so2_compile_visible_linear()) class TestMLP(unittest.TestCase): diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index cf3daa54f4..0bd9c22881 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -463,9 +463,9 @@ def _make_frame_with_natoms( ) return coord, atype, box - def test_neo_cute_flag_is_noop_for_ineligible_model(self) -> None: - """An unsupported SeZM shape must remain bit-identical with the flag on.""" - from deepmd.pt_expt.kernels.cute.sezm import k1 as cute_k1 + def test_neo_cute_flag_skips_ineligible_kernel(self) -> None: + """Reject an unsupported model before entering the per-block CuTe path.""" + from deepmd.pt_expt.kernels.cute.sezm.so2 import operation as cute_so2 cpu = torch.device("cpu") with ( @@ -524,14 +524,14 @@ def run_lower() -> dict[str, torch.Tensor]: clear=False, ), mock.patch.object( - cute_k1, - "maybe_run_cute_k1", - wraps=cute_k1.maybe_run_cute_k1, + cute_so2, + "maybe_run_cute_so2", + wraps=cute_so2.maybe_run_cute_so2, ) as maybe_run, ): actual = run_lower() - self.assertGreater(maybe_run.call_count, 0) + self.assertEqual(maybe_run.call_count, 0) self.assertEqual(actual.keys(), reference.keys()) for name, expected in reference.items(): if not isinstance(expected, torch.Tensor): @@ -542,27 +542,27 @@ def run_lower() -> dict[str, torch.Tensor]: ) def test_neo_cute_sort_guard_uses_post_cast_geometry_dtype(self) -> None: - """Raw FP64 coordinates do not hide an otherwise eligible FP32 K1.""" + """Raw FP64 coordinates do not hide an otherwise eligible FP32 SO2.""" from deepmd.pt.model.model import ( sezm_model, ) descriptor = mock.Mock() descriptor.compute_dtype = torch.float32 - descriptor._packed_wigner_candidate.return_value = True + descriptor.is_cute_infer_packed_wigner_candidate.return_value = True device = torch.device("cuda") with mock.patch.object( sezm_model, "_neo_cute_infer_enabled", return_value=True ): - actual = sezm_model._neo_cute_k1_requires_sorted_edges( + actual = sezm_model._neo_cute_so2_requires_sorted_edges( descriptor, training=False, device=device, ) self.assertTrue(actual) - descriptor._packed_wigner_candidate.assert_called_once_with( + descriptor.is_cute_infer_packed_wigner_candidate.assert_called_once_with( device, torch.float32, torch.float32, @@ -748,7 +748,7 @@ def test_load_state_dict_invalidates_local_and_shared_compile_caches(self) -> No object.__setattr__(model, "_embedding_task_buf_order", ("task",)) object.__setattr__(model, "compiled_dens_compute", local_callable) model._dens_compiled = True - model._deepmd_cute_k1_state = False + model._deepmd_cute_so2_state = False with ( mock.patch.dict( @@ -769,7 +769,7 @@ def test_load_state_dict_invalidates_local_and_shared_compile_caches(self) -> No self.assertIsNone(model._embedding_task_buf_order) self.assertIsNone(model.compiled_dens_compute) self.assertFalse(model._dens_compiled) - self.assertFalse(hasattr(model, "_deepmd_cute_k1_state")) + self.assertFalse(hasattr(model, "_deepmd_cute_so2_state")) self.assertEqual(sezm_model_module._SEZM_COMPILE_CACHE, {}) self.assertEqual(sezm_model_module._SEZM_TASK_BUF_ORDER, {}) diff --git a/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py b/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py index 820a29834e..e3ae46c1c4 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py +++ b/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py @@ -1,6 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """End-to-end parity of the pt_expt DPA4 accelerated inference paths.""" +from types import ( + SimpleNamespace, +) +from typing import ( + Any, +) + import numpy as np import pytest import torch @@ -10,6 +17,9 @@ except ImportError: pass +from deepmd.dpmodel.utils.neighbor_graph import ( + NeighborGraph, +) from deepmd.pt_expt.descriptor.dpa4 import ( DescrptDPA4, ) @@ -32,6 +42,13 @@ so2_conv, zonal_scatter, ) +from deepmd.pt_expt.kernels.cute.sezm import ( + runtime_policy, +) +from deepmd.pt_expt.kernels.cute.sezm.output_grid import ( + readout_l0, +) +from deepmd.pt_expt.kernels.cute.sezm.so2 import operation as so2 from deepmd.pt_expt.kernels.cutile import ( CUTILE_AVAILABLE, ) @@ -81,6 +98,229 @@ def _make_descriptor( ) +def make_neo_descriptor(sel: int | list[int] = 4) -> DescrptDPA4: + """Build the exact Neo layout served by the current CuTe kernels.""" + return DescrptDPA4( + ntypes=2, + sel=sel, + channels=32, + lmax=3, + mmax=1, + n_blocks=2, + so2_layers=3, + n_focus=2, + message_node_so3=True, + ffn_neurons=0, + ffn_so3_grid=True, + grid_branch=[0, 0, 1], + ffn_blocks=1, + so3_readout="mlp", + use_amp=False, + random_gamma=False, + precision="float32", + trainable=False, + seed=42, + ).eval() + + +def test_cute_contract_recognizes_pt_expt_neo_modules(monkeypatch) -> None: + """Keep CuTe eligibility independent of the concrete PT module classes.""" + for name in ( + "DP_TRITON_INFER", + "DP_CUDA_INFER", + "DP_CUTILE_INFER", + "DP_CUTE_INFER", + ): + monkeypatch.setenv(name, "0") + + descriptor = make_neo_descriptor() + + assert descriptor.blocks + assert all(so2.is_supported_neo_so2_block(block) for block in descriptor.blocks) + assert readout_l0.has_neo_readout_contract(descriptor.output_ffn) + + +def test_packed_wigner_graph_compacts_and_reuses_csr(monkeypatch) -> None: + """Reuse canonical graph metadata without sorting the edge stream again.""" + monkeypatch.setenv("DP_CUTE_INFER", "1") + descriptor = _make_descriptor(2, [4], 4.0).eval() + monkeypatch.setattr( + descriptor, + "is_cute_infer_packed_wigner_candidate", + lambda *args: True, + ) + graph = NeighborGraph( + n_node=torch.tensor([4], dtype=torch.int64), + edge_index=torch.tensor( + [[3, 0, 2, 0, 0], [0, 1, 2, 0, 0]], + dtype=torch.int64, + ), + edge_vec=torch.arange(15, dtype=torch.float32).reshape(5, 3), + edge_mask=torch.tensor([True, True, True, False, False]), + destination_order=torch.arange(5, dtype=torch.int64), + destination_row_ptr=torch.tensor([0, 1, 2, 3, 3], dtype=torch.int64), + source_order=torch.tensor([1, 2, 0, 3, 4], dtype=torch.int64), + source_row_ptr=torch.tensor([0, 1, 1, 2, 3], dtype=torch.int64), + destination_sorted=True, + ) + monkeypatch.setattr( + torch, + "argsort", + lambda *args, **kwargs: pytest.fail("canonical CSR must not be rebuilt"), + ) + + packed_graph = descriptor.prepare_packed_wigner_graph(graph, n_nodes=4) + + assert packed_graph is not None + torch.testing.assert_close( + packed_graph.edge_index, + torch.tensor([[3, 0, 2], [0, 1, 2]], dtype=torch.int64), + ) + torch.testing.assert_close( + packed_graph.edge_vec, + torch.arange(9, dtype=torch.float32).reshape(3, 3), + ) + assert torch.all(packed_graph.edge_mask) + assert packed_graph.destination_sorted + torch.testing.assert_close( + packed_graph.destination_row_ptr, + torch.tensor([0, 1, 2, 3, 3], dtype=torch.int64), + ) + torch.testing.assert_close( + packed_graph.source_order, + torch.tensor([1, 2, 0], dtype=torch.int64), + ) + torch.testing.assert_close( + packed_graph.source_row_ptr, + torch.tensor([0, 1, 1, 2, 3], dtype=torch.int64), + ) + + edge_cache = SimpleNamespace( + src=packed_graph.edge_index[0], + dst=packed_graph.edge_index[1], + csr_cache={ + "dst": ( + packed_graph.destination_order, + packed_graph.destination_row_ptr, + ), + "src": (packed_graph.source_order, packed_graph.source_row_ptr), + }, + destinations_sorted=True, + D_packed=torch.empty(3, 46), + edge_src_gate=None, + ) + metadata = descriptor.prepare_cute_infer_so2_metadata( + edge_cache, + n_nodes=4, + ) + assert metadata is not None + destination_row_ptr, source_order, source_row_ptr = metadata + torch.testing.assert_close( + destination_row_ptr, + packed_graph.destination_row_ptr.to(torch.int32), + ) + torch.testing.assert_close( + source_order, + packed_graph.source_order.to(torch.int32), + ) + torch.testing.assert_close( + source_row_ptr, + packed_graph.source_row_ptr.to(torch.int32), + ) + + +def test_packed_wigner_graph_without_csr_builds_metadata_once(monkeypatch) -> None: + """Build each CSR ordering once when the graph does not provide one.""" + monkeypatch.setenv("DP_CUTE_INFER", "1") + descriptor = _make_descriptor(2, [4], 4.0).eval() + monkeypatch.setattr( + descriptor, + "is_cute_infer_packed_wigner_candidate", + lambda *args: True, + ) + graph = NeighborGraph( + n_node=torch.tensor([2], dtype=torch.int64), + edge_index=torch.tensor([[0, 1], [1, 0]], dtype=torch.int64), + edge_vec=torch.ones(2, 3), + edge_mask=torch.ones(2, dtype=torch.bool), + ) + argsort = torch.argsort + sort_count = 0 + + def count_argsort(*args: Any, **kwargs: Any) -> torch.Tensor: + nonlocal sort_count + sort_count += 1 + return argsort(*args, **kwargs) + + monkeypatch.setattr(torch, "argsort", count_argsort) + + packed_graph = descriptor.prepare_packed_wigner_graph(graph, n_nodes=2) + + assert packed_graph is not None + edge_cache = SimpleNamespace( + src=packed_graph.edge_index[0], + dst=packed_graph.edge_index[1], + csr_cache={ + "dst": ( + packed_graph.destination_order, + packed_graph.destination_row_ptr, + ), + "src": (packed_graph.source_order, packed_graph.source_row_ptr), + }, + destinations_sorted=True, + D_packed=torch.empty(2, 46), + edge_src_gate=None, + ) + + assert descriptor.prepare_cute_infer_so2_metadata(edge_cache, n_nodes=2) is not None + assert sort_count == 2 + + +def test_packed_wigner_cache_routes_pt_expt_block_to_cute(monkeypatch) -> None: + """Use packed Wigner storage as the prevalidated per-block dispatch token.""" + block = make_neo_descriptor().blocks[0] + destination_row_ptr = torch.tensor([0, 1], dtype=torch.int32) + source_order = torch.tensor([0], dtype=torch.int32) + source_row_ptr = torch.tensor([0, 1], dtype=torch.int32) + edge_cache = SimpleNamespace( + D_packed=torch.empty(1, 46), + cute_infer_so2_metadata=( + destination_row_ptr, + source_order, + source_row_ptr, + ), + ) + expected = torch.empty(1) + + def run_cute( + candidate_block: Any, + x: torch.Tensor, + candidate_edge_cache: Any, + radial_feat: torch.Tensor, + *, + dst_ptr: torch.Tensor, + source_order: torch.Tensor, + source_ptr: torch.Tensor, + ) -> torch.Tensor: + del x, radial_feat + assert candidate_block is block + assert candidate_edge_cache is edge_cache + assert dst_ptr is destination_row_ptr + assert source_order is edge_cache.cute_infer_so2_metadata[1] + assert source_ptr is source_row_ptr + return expected + + monkeypatch.setattr(so2, "maybe_run_cute_so2", run_cute) + + actual = block._run_so2_unit_impl( + torch.empty(1), + edge_cache, + torch.empty(1), + ) + + assert actual is expected + + @pytest.mark.parametrize( ("precision", "expected_bound"), [("float32", True), ("float64", False)], @@ -117,17 +357,23 @@ def test_fp32_only_cuda_bindings( assert len(initial_embeddings) == 1 assert grid_nets - assert (descriptor._cuda_radial_fn is not None) is expected_bound - assert all(module._cuda_scatter is expected_bound for module in initial_embeddings) + assert (descriptor.cuda_infer_l_1_radial is not None) is expected_bound + assert all( + module.cuda_infer_l_1_scatter is expected_bound for module in initial_embeddings + ) assert all( - (module._grid_pair_fn is not None) is expected_bound for module in grid_nets + (module.cuda_infer_l_1_grid_pair is not None) is expected_bound + for module in grid_nets ) cpu_zonal = torch.empty(1) - assert all(not module._can_fuse_scatter(cpu_zonal) for module in initial_embeddings) + assert all( + not module.can_run_cuda_infer_l_1_scatter(cpu_zonal) + for module in initial_embeddings + ) for module in initial_embeddings: - module._force_fused_scatter = True + module.force_cuda_infer_l_1_scatter = True assert all( - module._can_fuse_scatter(cpu_zonal) is expected_bound + module.can_run_cuda_infer_l_1_scatter(cpu_zonal) is expected_bound for module in initial_embeddings ) @@ -219,17 +465,74 @@ def test_source_gated_flash_retains_dense_rotations(self, monkeypatch) -> None: for module in accelerated.modules() if isinstance(module, SO2Convolution) ) - if conv._cuda_conv_fn is None: + if conv.cuda_infer_l_2_conv is None: pytest.skip("the descriptor layout has no fused CUDA convolution") - assert conv._flash_atten_fn is not None - assert conv._cuda_value_train is None - assert not accelerated._wigner_free_conv + assert conv.flash_attention is not None + assert conv.cuda_train_value is None + assert not accelerated.cuda_infer_l_2_covers_all_blocks assert accelerated._build_full_wigner() output, gradient = self._step(accelerated) np.testing.assert_allclose(output, dense_output, rtol=2e-4, atol=2e-5) np.testing.assert_allclose(gradient, dense_gradient, rtol=2e-4, atol=2e-5) + def test_cute_neo_forward_and_coordinate_gradient(self, monkeypatch) -> None: + """Exercise packed Wigner and SO2 through the PT-expt descriptor.""" + try: + import cuda.bindings.driver # noqa: F401 + import cutlass.cute # noqa: F401 + import tvm_ffi # noqa: F401 + + from deepmd.pt_expt.kernels.cute.sezm import ( + wignerd, + ) + except ImportError: + pytest.skip("the CuTe DSL runtime is unavailable") + capability = tuple(torch.cuda.get_device_capability(self.device)) + if not runtime_policy.is_supported_so2_capability(capability): + pytest.skip(f"the CuTe SO2 path does not support {capability}") + + for name in ( + "DP_TRITON_INFER", + "DP_CUDA_INFER", + "DP_CUTILE_INFER", + "DP_CUTE_INFER", + ): + monkeypatch.setenv(name, "0") + data = make_neo_descriptor(self.sel_mix).serialize() + reference = DescrptDPA4.deserialize(data).to(self.device).eval() + reference_output, reference_gradient = self._step(reference) + + monkeypatch.setenv("DP_CUTE_INFER", "1") + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + accelerated = DescrptDPA4.deserialize(data).to(self.device).eval() + dispatch_count = {"wigner": 0, "so2": 0} + run_wignerd = wignerd.run_cute_wignerd + run_so2 = so2.maybe_run_cute_so2 + + def count_wignerd(*args: Any, **kwargs: Any) -> Any: + result = run_wignerd(*args, **kwargs) + dispatch_count["wigner"] += result is not None + return result + + def count_so2(*args: Any, **kwargs: Any) -> Any: + result = run_so2(*args, **kwargs) + dispatch_count["so2"] += result is not None + return result + + monkeypatch.setattr(wignerd, "run_cute_wignerd", count_wignerd) + monkeypatch.setattr(so2, "maybe_run_cute_so2", count_so2) + output, gradient = self._step(accelerated) + + assert dispatch_count == {"wigner": 1, "so2": len(accelerated.blocks)} + np.testing.assert_allclose(output, reference_output, rtol=5e-5, atol=5e-5) + np.testing.assert_allclose( + gradient, + reference_gradient, + rtol=5e-5, + atol=5e-5, + ) + @pytest.mark.parametrize("backend", ["triton", "cuda", "cutile"]) def test_forward_and_coordinate_gradient(self, monkeypatch, backend) -> None: if backend == "triton" and not FORCE_ASSEMBLY_TRITON_AVAILABLE: @@ -268,18 +571,18 @@ def test_forward_and_coordinate_gradient(self, monkeypatch, backend) -> None: if isinstance(module, WignerDCalculator) ) if backend == "triton": - assert so2._flash_atten_fn is not None - assert so2._triton_value_path is not None - assert wigner._use_triton_monomials + assert so2.flash_attention is not None + assert so2.triton_infer_l_2_value is not None + assert wigner.triton_infer_l_1_monomials elif backend == "cuda": - if so2._cuda_conv_fn is None: + if so2.cuda_infer_l_2_conv is None: pytest.skip("DPA4 CUDA operators are unavailable") - assert accelerated._cuda_radial_fn is not None - assert accelerated._cuda_wigner_fn is not None + assert accelerated.cuda_infer_l_1_radial is not None + assert accelerated.cuda_infer_l_1_wigner is not None else: - assert so2._flash_atten_fn is not None - assert so2._cutile_value_path is not None - assert wigner._use_cutile_monomials + assert so2.flash_attention is not None + assert so2.cutile_infer_value is not None + assert wigner.cutile_infer_monomials coord_ref, atype, nlist = self._inputs() output_ref = reference(coord_ref, atype, nlist)[0] diff --git a/source/tests/pt_expt/descriptor/test_dpa4_ckpt_triton.py b/source/tests/pt_expt/descriptor/test_dpa4_ckpt_triton.py index 200a947aa6..930d0a957f 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4_ckpt_triton.py +++ b/source/tests/pt_expt/descriptor/test_dpa4_ckpt_triton.py @@ -120,9 +120,9 @@ def test_triton_eager_fallback_parity(self, monkeypatch) -> None: monkeypatch.setenv("DP_TRITON_INFER", "1") m = DescrptDPA4.deserialize(self.data).to(self.device).eval() so2 = next(x for x in m.modules() if isinstance(x, SO2Convolution)) - assert so2.use_triton_infer - assert so2._rotate_to_local_fn is not None - assert so2._rotate_back_fn is not None + assert so2.triton_infer_level >= 1 + assert so2.triton_l_1_rotate_to_local is not None + assert so2.triton_l_1_rotate_back is not None mixers = [x for x in m.modules() if isinstance(x, DynamicRadialDegreeMixer)] assert mixers and all(x.triton_infer_level >= 1 for x in mixers) diff --git a/source/tests/pt_expt/descriptor/test_dpa4_train_paths.py b/source/tests/pt_expt/descriptor/test_dpa4_train_paths.py index c29e89d622..4d6c1fc1b3 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4_train_paths.py +++ b/source/tests/pt_expt/descriptor/test_dpa4_train_paths.py @@ -44,7 +44,7 @@ DynamicRadialDegreeMixer, SO2Convolution, SO2Linear, - _active_triton_level, + active_triton_level, ) from deepmd.pt_expt.kernels.cuda.dpa4.so2_conv_train import ( op_available as cuda_value_available, @@ -140,34 +140,34 @@ def test_triton_mode_gate_binds_each_stage( assert conv.triton_infer_level == infer_level # Rotation and flash wrappers retain their eager implementations when # Triton is unavailable, so their binding follows the mode gates alone. - assert (conv._rotate_to_local_fn is not None) is requested - assert (conv._rotate_back_fn is not None) is requested - assert (conv._flash_atten_fn is not None) is requested - assert conv._flash_atten_trains is (requested and training) + assert (conv.triton_l_1_rotate_to_local is not None) is requested + assert (conv.triton_l_1_rotate_back is not None) is requested + assert (conv.flash_attention is not None) is requested + assert conv.flash_attention_supports_training is (requested and training) # Segment softmax has no wrapper-level fallback and binds only when its # own Triton implementation is importable. - assert (conv._segment_softmax_fn is not None) is ( + assert (conv.triton_l_1_segment_softmax is not None) is ( requested and SEGMENT_SOFTMAX_TRITON_AVAILABLE ) # The rotate-mix front end is bound by a profitability bound on the # hidden width, which this narrow block sits below. assert conv.hidden_channels < 128 - assert conv._triton_rotate_mix is None + assert conv.triton_l_1_rotate_mix is None # The CUDA gate is off, so the value stream stays on the stages. - assert conv._cuda_value_train is None + assert conv.cuda_train_value is None for module in descriptor.modules(): if isinstance(module, SO2Linear): # The fused GEMM additionally needs every |m| block width to align # to its BN=64 tile, which a narrow block does not satisfy. aligned = slices_supported(module._block_diag_slices) - assert (module._block_diag_gemm is not None) is ( + assert (module.triton_l_1_block_diag_gemm is not None) is ( requested and SO2_BLOCK_GEMM_TRITON_AVAILABLE and aligned ) if isinstance(module, DynamicRadialDegreeMixer): # The callable contains its eager fallback, so construction binds it # whenever either mode requests the stage. - assert (module._radial_mix_block is not None) is requested + assert (module.triton_l_1_radial_mix is not None) is requested if isinstance(module, GatedActivation): assert module.triton_train_level == train_level assert module.triton_infer_level == infer_level @@ -188,9 +188,9 @@ def test_triton_mode_gate_binds_each_stage( if isinstance( module, (SO2Convolution, SO2Linear, DynamicRadialDegreeMixer) ): - assert _active_triton_level(module) == active_level + assert active_triton_level(module) == active_level if isinstance(module, SO2Convolution): - assert module._rotation_active() is bool(active_level) + assert module.rotation_kernel_active() is bool(active_level) if isinstance(module, GatedActivation): level = ( module.triton_train_level @@ -213,13 +213,13 @@ def test_cuda_train_gate_binds_the_value_stream(monkeypatch) -> None: ] assert convolutions for conv in convolutions: - assert conv._cuda_value_train is not None + assert conv.cuda_train_value is not None # The two layers are independent: the CUDA value stream does not # switch on any Triton stage, and the attention span stays dense # until the Triton gate asks for it. assert conv.triton_train_level == 0 - assert conv._segment_softmax_fn is None - assert conv._flash_atten_trains is False + assert conv.triton_l_1_segment_softmax is None + assert conv.flash_attention_supports_training is False def test_cuda_triton_train_reuses_packed_wigner_runs(monkeypatch) -> None: @@ -233,14 +233,14 @@ def test_cuda_triton_train_reuses_packed_wigner_runs(monkeypatch) -> None: monkeypatch.setenv("DP_TRITON_TRAIN", "1") descriptor = _make_descriptor(2, [20], 4.0).train() - assert descriptor._packed_wigner_train + assert descriptor.cuda_train_covers_all_blocks assert not descriptor._build_full_wigner() for block in descriptor.blocks: conv = block.so2_conv - assert conv._cuda_value_train is not None - assert conv._flash_atten_fn is not None - assert conv._flash_atten_trains + assert conv.cuda_train_value is not None + assert conv.flash_attention is not None + assert conv.flash_attention_supports_training @pytest.mark.parametrize("gate_name", ["DP_TRITON_TRAIN", "DP_TRITON_INFER"]) @@ -267,7 +267,7 @@ def test_grid_pair_train_follows_its_gate_and_slot_bound( and GRID_PAIR_TRITON_AVAILABLE and slots >= 75 ) - assert (net._grid_pair_train_fn is not None) is expected + assert (net.triton_train_l_1_grid_pair is not None) is expected @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @@ -327,9 +327,9 @@ def test_training_step_matches_the_dense_path(self, monkeypatch, path) -> None: module for module in fused.modules() if isinstance(module, SO2Convolution) ) if path in ("cuda", "cuda-triton"): - assert conv._cuda_value_train is not None + assert conv.cuda_train_value is not None if path in ("triton", "cuda-triton"): - assert conv._segment_softmax_fn is not None + assert conv.triton_l_1_segment_softmax is not None fused_objective, fused_gradient = self._step(fused) np.testing.assert_allclose( diff --git a/source/tests/pt_expt/utils/test_serialization_kernel_levels.py b/source/tests/pt_expt/utils/test_serialization_kernel_levels.py index 53651103ca..ad929b5a4b 100644 --- a/source/tests/pt_expt/utils/test_serialization_kernel_levels.py +++ b/source/tests/pt_expt/utils/test_serialization_kernel_levels.py @@ -225,7 +225,7 @@ def test_packed_weights_only_prepare_for_bound_dpa4_value_path( ) -> None: model = torch.nn.Module() conv = torch.nn.Module() - conv._triton_value_path = object() if has_value_path else None + conv.triton_infer_l_2_value = object() if has_value_path else None model.add_module("conv", conv) calls = [] From 105633b75e8544b34f64f67cef15035cf68f8237 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Thu, 3 Sep 2026 14:10:49 +0800 Subject: [PATCH 06/11] fix(pt): remove stale CuTe metadata re-export --- deepmd/pt/model/descriptor/sezm.py | 6 +-- .../pt/model/descriptor/sezm_nn/edge_cache.py | 3 -- .../pt/model/sezm_cute/test_so2_dst_ptr.py | 41 ++++++++----------- 3 files changed, 20 insertions(+), 30 deletions(-) diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index baf8090759..b0b446826b 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -66,6 +66,9 @@ UpdateSel, ) from deepmd.pt_expt.kernels.cute.sezm import runtime_policy as cute_runtime_policy +from deepmd.pt_expt.kernels.cute.sezm.so2.metadata import ( + build_sorted_edge_index_metadata, +) from deepmd.pt_expt.kernels.utils import ( cuda_infer_level, use_amp_infer, @@ -108,9 +111,6 @@ safe_norm, safe_numpy_to_tensor, ) -from .sezm_nn.edge_cache import ( - build_sorted_edge_index_metadata, -) if TYPE_CHECKING: from collections.abc import ( diff --git a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py index 11665b56e3..bf7df7fb3e 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py +++ b/deepmd/pt/model/descriptor/sezm_nn/edge_cache.py @@ -28,9 +28,6 @@ from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import ( is_cute_infer_enabled, ) -from deepmd.pt_expt.kernels.cute.sezm.so2.metadata import ( - build_sorted_edge_index_metadata as build_sorted_edge_index_metadata, -) from .utils import ( get_promoted_dtype, diff --git a/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py b/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py index 9c4c1e3212..d2efab4cfa 100644 --- a/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py +++ b/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py @@ -72,6 +72,11 @@ def _load_extracted_sorted_metadata_function(): return namespace["build_sorted_edge_index_metadata"] +def _load_sorted_metadata_function() -> Any: + module = importlib.import_module("deepmd.pt_expt.kernels.cute.sezm.so2.metadata") + return module.build_sorted_edge_index_metadata + + class _FakeDst: device = "cuda:0" dtype = "int64" @@ -384,9 +389,7 @@ def _cache(src, dst, node_count): def test_builds_destination_and_indirect_source_csr(self): assert torch is not None - edge_cache = importlib.import_module( - "deepmd.pt.model.descriptor.sezm_nn.edge_cache" - ) + builder = _load_sorted_metadata_function() src = torch.tensor( [2, 0, 3, 0, 1, 2], dtype=torch.int64, @@ -397,9 +400,7 @@ def test_builds_destination_and_indirect_source_csr(self): dtype=torch.int64, device="cpu", ) - dst_ptr, source_order, source_ptr = edge_cache.build_sorted_edge_index_metadata( - src, dst, 4 - ) + dst_ptr, source_order, source_ptr = builder(src, dst, 4) self.assertEqual(dst_ptr.dtype, torch.int32) self.assertEqual(source_order.dtype, torch.int32) @@ -437,10 +438,8 @@ def test_builds_destination_and_indirect_source_csr(self): def test_empty_edges_build_valid_zero_csr(self): assert torch is not None - edge_cache = importlib.import_module( - "deepmd.pt.model.descriptor.sezm_nn.edge_cache" - ) - dst_ptr, source_order, source_ptr = edge_cache.build_sorted_edge_index_metadata( + builder = _load_sorted_metadata_function() + dst_ptr, source_order, source_ptr = builder( torch.empty(0, dtype=torch.int64, device="cpu"), torch.empty(0, dtype=torch.int64, device="cpu"), 4, @@ -509,15 +508,13 @@ def build_ptr(src, dst): def test_dynamic_node_and_edge_counts_build_independent_metadata(self): assert torch is not None - edge_cache = importlib.import_module( - "deepmd.pt.model.descriptor.sezm_nn.edge_cache" - ) - first = edge_cache.build_sorted_edge_index_metadata( + builder = _load_sorted_metadata_function() + first = builder( torch.tensor([0, 1, 2], dtype=torch.int64, device="cpu"), torch.tensor([0, 1, 2], dtype=torch.int64, device="cpu"), 4, ) - second = edge_cache.build_sorted_edge_index_metadata( + second = builder( torch.tensor( [0, 2, 4, 1, 3], dtype=torch.int64, @@ -553,9 +550,7 @@ def test_dynamic_node_and_edge_counts_build_independent_metadata(self): def test_explicit_tensor_metadata_survives_disabled_so2_boundary(self): assert torch is not None - edge_cache = importlib.import_module( - "deepmd.pt.model.descriptor.sezm_nn.edge_cache" - ) + builder = _load_sorted_metadata_function() dynamo = getattr(torch, "_dynamo", None) if dynamo is None: self.skipTest("torch._dynamo is unavailable") @@ -616,12 +611,10 @@ def consume(cache, x, radial, dst_ptr, source_order, source_ptr): ) for src, dst, node_count in cases: cache = self._cache(src, dst, node_count) - dst_ptr, source_order, source_ptr = ( - edge_cache.build_sorted_edge_index_metadata( - src, - dst, - node_count, - ) + dst_ptr, source_order, source_ptr = builder( + src, + dst, + node_count, ) radial = torch.ones(src.numel(), 4, 1, device="cpu") args = ( From 879337778a8956937161335e05f8c8f2a3ec81a8 Mon Sep 17 00:00:00 2001 From: atulcthakur Date: Thu, 3 Sep 2026 22:34:02 -0700 Subject: [PATCH 07/11] fix(pt-expt): support Neo CuTe without focus normalization --- .../sezm/so2/kernels/focus_source_backward.py | 20 ++++++--- .../cute/sezm/so2/kernels/phase_c_backward.py | 45 +++++++++++-------- .../kernels/cute/sezm/so2/operation.py | 18 +++++++- .../pt_expt/kernels/cute/sezm/so2/phase_c.py | 4 ++ .../pt_expt/kernels/cute/sezm/so2/runner.py | 20 +++++++-- .../so2/sm90/phase_c_attention_backward.py | 34 +++++++++----- .../kernels/cute/sezm/so2/sm90/runner.py | 22 ++++----- source/tests/pt/model/sezm_cute/test_so2.py | 36 +++++++++++---- .../descriptor/test_dpa4_accelerated.py | 16 +++++-- 9 files changed, 152 insertions(+), 63 deletions(-) diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/focus_source_backward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/focus_source_backward.py index c3d6cdde81..061e6c2179 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/focus_source_backward.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/focus_source_backward.py @@ -5,7 +5,7 @@ The forward kernel fuses two independent PyTorch producer chains in one launch: - focus RMSNorm -> two-focus logits -> softmax -> label smoothing + optional focus RMSNorm -> two-focus logits -> softmax -> label smoothing scalar Q/K RMSNorm -> Q projection + K projection The kernel is specialized to the Neo SO2 shape `(E, F=2, C=32)`. @@ -57,6 +57,7 @@ def neo_attention_prelude_forward_jit( qk_eps: cutlass.Float32, tau: cutlass.Float32, label_smoothing: cutlass.Float32, + use_focus_norm: cutlass.Constexpr[bool], ): edges, _ = focus.shape nodes, _, _ = x_l0.shape @@ -75,6 +76,7 @@ def neo_attention_prelude_forward_jit( qk_eps, tau, label_smoothing, + use_focus_norm, ).launch( # The launch covers two independent domains. Using their sum avoids a # host-side branch on symbolic E/N while guaranteeing coverage of both. @@ -100,6 +102,7 @@ def neo_attention_prelude_forward_kernel( qk_eps: cutlass.Float32, tau: cutlass.Float32, label_smoothing: cutlass.Float32, + use_focus_norm: cutlass.Constexpr[bool], ): tid, _, _ = cute.arch.thread_idx() block, _, _ = cute.arch.block_idx() @@ -110,10 +113,14 @@ def neo_attention_prelude_forward_kernel( if edge < edges: x0 = focus[edge, lane].to(cutlass.Float32) x1 = focus[edge, 32 + lane].to(cutlass.Float32) - inv0 = cute.rsqrt(_warp_sum(x0 * x0) / cutlass.Float32(32.0) + focus_eps) - inv1 = cute.rsqrt(_warp_sum(x1 * x1) / cutlass.Float32(32.0) + focus_eps) - norm0 = x0 * inv0 * focus_scale[0, lane].to(cutlass.Float32) - norm1 = x1 * inv1 * focus_scale[1, lane].to(cutlass.Float32) + if cutlass.const_expr(use_focus_norm): + inv0 = cute.rsqrt(_warp_sum(x0 * x0) / cutlass.Float32(32.0) + focus_eps) + inv1 = cute.rsqrt(_warp_sum(x1 * x1) / cutlass.Float32(32.0) + focus_eps) + norm0 = x0 * inv0 * focus_scale[0, lane].to(cutlass.Float32) + norm1 = x1 * inv1 * focus_scale[1, lane].to(cutlass.Float32) + else: + norm0 = x0 + norm1 = x1 logit0 = _warp_sum(norm0 * focus_weight[lane, 0].to(cutlass.Float32)) logit1 = _warp_sum(norm1 * focus_weight[lane, 1].to(cutlass.Float32)) @@ -176,6 +183,8 @@ def compile_neo_attention_prelude_forward( tau: float, label_smoothing: float, compile_identity: tuple[int, int, int] | None = None, + *, + use_focus_norm: bool = True, ) -> Callable: # The identity keeps independently compiled device/architecture binaries in # distinct cache entries. Compilation itself runs under the runner's device. @@ -230,6 +239,7 @@ def compile_neo_attention_prelude_forward( cutlass.Float32(qk_eps), cutlass.Float32(tau), cutlass.Float32(label_smoothing), + bool(use_focus_norm), options="--enable-tvm-ffi", ) diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_backward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_backward.py index 6280257bb5..b14b4c3402 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_backward.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/kernels/phase_c_backward.py @@ -138,18 +138,11 @@ def _focus_source_backward( eps: cutlass.Constexpr[float], tau: cutlass.Constexpr[float], label_smoothing: cutlass.Constexpr[float], + use_focus_norm: cutlass.Constexpr[bool], ): """Consume the Phase-C focus-alpha reduction without a global temporary.""" - value = params.focus_src[edge, head, lane].to(cutlass.Float32) - scale = params.focus_scale[head, lane].to(cutlass.Float32) weight = params.focus_weight[lane, head].to(cutlass.Float32) - square_sum = _warp_sum(value * value) - if lane == 0: - focus_inv[head] = cute.rsqrt( - square_sum / cutlass.Float32(FOCUS_CHANNELS) + cutlass.Float32(eps) - ) - if tidx == 0: probability_keep = cutlass.Float32(1.0 - label_smoothing) smooth = cutlass.Float32(label_smoothing / N_FOCUS) @@ -168,15 +161,25 @@ def _focus_source_backward( focus_grad_logits[1] = probability1 * (grad1 - dot) * inv_tau cute.arch.sync_threads() - grad_scaled = focus_grad_logits[head] * weight * scale - coeff_sum = _warp_sum(grad_scaled * value) - if lane == 0: - focus_coeff[head] = coeff_sum / cutlass.Float32(FOCUS_CHANNELS) - cute.arch.sync_warp() - - inv = focus_inv[head] - grad_value = grad_scaled * inv - grad_value -= value * inv * inv * inv * focus_coeff[head] + if cutlass.const_expr(use_focus_norm): + value = params.focus_src[edge, head, lane].to(cutlass.Float32) + scale = params.focus_scale[head, lane].to(cutlass.Float32) + square_sum = _warp_sum(value * value) + if lane == 0: + focus_inv[head] = cute.rsqrt( + square_sum / cutlass.Float32(FOCUS_CHANNELS) + cutlass.Float32(eps) + ) + grad_scaled = focus_grad_logits[head] * weight * scale + coeff_sum = _warp_sum(grad_scaled * value) + if lane == 0: + focus_coeff[head] = coeff_sum / cutlass.Float32(FOCUS_CHANNELS) + cute.arch.sync_warp() + + inv = focus_inv[head] + grad_value = grad_scaled * inv + grad_value -= value * inv * inv * inv * focus_coeff[head] + else: + grad_value = focus_grad_logits[head] * weight params.grad_focus_src[head, edge, lane] = grad_value.to( params.grad_focus_src.element_type ) @@ -190,6 +193,7 @@ def neo_phase_c_backward_layout_kernel( focus_eps: cutlass.Constexpr[float], focus_tau: cutlass.Constexpr[float], focus_label_smoothing: cutlass.Constexpr[float], + use_focus_norm: cutlass.Constexpr[bool], ): tidx, _, _ = cute.arch.thread_idx() node, _, _ = cute.arch.block_idx() @@ -359,6 +363,7 @@ def neo_phase_c_backward_layout_kernel( focus_eps, focus_tau, focus_label_smoothing, + use_focus_norm, ) cute.arch.sync_threads() @@ -462,6 +467,7 @@ def neo_phase_c_backward_layout_jit( focus_eps: cutlass.Constexpr[float], focus_tau: cutlass.Constexpr[float], focus_label_smoothing: cutlass.Constexpr[float], + use_focus_norm: cutlass.Constexpr[bool], ): params = NeoPhaseCBackwardLayoutParams( grad_out=grad_out, @@ -507,6 +513,7 @@ def neo_phase_c_backward_layout_jit( focus_eps, focus_tau, focus_label_smoothing, + use_focus_norm, ).launch( grid=[node_count, 1, 1], block=[HIDDEN, 1, 1], @@ -524,9 +531,10 @@ def compile_neo_phase_c_backward_layout( focus_eps: float, focus_tau: float, focus_label_smoothing: float, + use_focus_norm: bool = True, ) -> Callable: """Compile the exact Neo Phase-C layout-boundary backward callable.""" - if focus_eps <= 0.0: + if use_focus_norm and focus_eps <= 0.0: raise ValueError("focus_eps must be positive") if focus_tau <= 0.0: raise ValueError("focus_tau must be positive") @@ -689,5 +697,6 @@ def compile_neo_phase_c_backward_layout( focus_eps, focus_tau, focus_label_smoothing, + bool(use_focus_norm), options="--enable-tvm-ffi", ) diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py b/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py index 0ad1d2de12..dbbe1f5f9f 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py @@ -172,12 +172,26 @@ def has_equivariant_rms_norm_contract(module: Any) -> bool: ) +def has_focus_compete_norm_contract(so2: Any) -> bool: + """Return whether focus competition matches the optional norm contract.""" + module = getattr(so2, "focus_compete_norm", None) + if not bool(getattr(so2, "focus_norm", False)): + return module is None or _is_identity(module) + return ( + type(module).__name__ == "ScalarRMSNorm" + and int(getattr(module, "channels", -1)) == 32 + and int(getattr(module, "n_focus", -1)) == 2 + and tuple(getattr(getattr(module, "adam_scale", None), "shape", ())) == (2, 32) + and hasattr(module, "eps") + ) + + def is_supported_neo_so2_block(block: Any) -> bool: """Return whether this block can use the current Neo CuTe SO2 path.""" so2 = block.so2_conv return ( get_neo_so2_spec(block).is_current_neo_target - and bool(so2.focus_norm) + and has_focus_compete_norm_contract(so2) and not bool(so2.edge_cartesian) and getattr(so2, "node_cartesian_tp", None) is None and has_equivariant_rms_norm_contract(block.post_so2_norm) @@ -1629,7 +1643,7 @@ def _runner_backward_manual( runner.denom, runner.focus_gate_src.detach().contiguous(), so2.adamw_focus_compete_w.detach().float().contiguous(), - so2.focus_compete_norm.adam_scale.detach().float().reshape(2, 32).contiguous(), + runner.focus_norm_scale, layout_outputs, ) grad_stack_out = _native_edge_major_stack_grad(layout_outputs.grad_stack) diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/phase_c.py b/deepmd/pt_expt/kernels/cute/sezm/so2/phase_c.py index 2d4ed1a104..471ec21d5a 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/so2/phase_c.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/phase_c.py @@ -133,11 +133,13 @@ def _compile_layout_boundary( focus_eps: float, focus_tau: float, focus_label_smoothing: float, + use_focus_norm: bool, ) -> Any: return compile_neo_phase_c_backward_layout( focus_eps=focus_eps, focus_tau=focus_tau, focus_label_smoothing=focus_label_smoothing, + use_focus_norm=use_focus_norm, ) @@ -159,11 +161,13 @@ def __init__( focus_eps: float, focus_tau: float, focus_label_smoothing: float, + use_focus_norm: bool = True, ) -> None: self._compiled = _compile_layout_boundary( float(focus_eps), float(focus_tau), float(focus_label_smoothing), + bool(use_focus_norm), ) def __call__( diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/runner.py b/deepmd/pt_expt/kernels/cute/sezm/so2/runner.py index 321c6fe871..32e92e79fc 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/so2/runner.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/runner.py @@ -480,17 +480,30 @@ def __init__( self.source_ptr_i32 = source_csr.source_ptr self.rotate = self.so2.rotate_inv_rescale_full.contiguous() self.edge_gate = _edge_gate(torch, record.edge_cache) + self.focus_norm_enabled = bool(self.so2.focus_norm) + if self.focus_norm_enabled: + focus_norm = self.so2.focus_compete_norm + self.focus_norm_eps = float(focus_norm.eps) + self.focus_norm_scale = focus_norm.adam_scale.detach().float().contiguous() + else: + self.focus_norm_eps = 0.0 + # The no-norm specialization does not read this fixed-ABI argument; + # reuse an existing shape-compatible tensor instead of allocating one. + self.focus_norm_scale = ( + self.so2.attn_qk_norm.adam_scale.detach().float().contiguous() + ) from .kernels.focus_source_backward import ( compile_neo_attention_prelude_forward, ) with torch.cuda.device(self.device_index): self.attention_prelude_forward = compile_neo_attention_prelude_forward( - float(self.so2.focus_compete_norm.eps), + self.focus_norm_eps, float(self.so2.attn_qk_norm.eps), float(self.so2.focus_softmax_tau), float(self.so2.focus_label_smoothing), self.compile_identity, + use_focus_norm=self.focus_norm_enabled, ) from .kernels.qk_edge import ( compile_neo_qk_edge_backward, @@ -521,9 +534,10 @@ def __init__( ) self.phase_c_layout_backward = CuteNeoPhaseCBackwardLayout( - focus_eps=float(self.so2.focus_compete_norm.eps), + focus_eps=self.focus_norm_eps, focus_tau=float(self.so2.focus_softmax_tau), focus_label_smoothing=float(self.so2.focus_label_smoothing), + use_focus_norm=self.focus_norm_enabled, ) with torch.cuda.device(self.device_index): self.softmax_fwd = compile_envelope_softmax_forward( @@ -815,7 +829,7 @@ def _build_forward_graph(self) -> None: self.focus_gate_src.view(n_edge, 64), x_l0_node.contiguous(), so2.adamw_focus_compete_w.detach().float().contiguous(), - so2.focus_compete_norm.adam_scale.detach().float().contiguous(), + self.focus_norm_scale, so2.attn_q_proj.weight.detach().float().view(32, 2, 32).contiguous(), so2.attn_k_proj.weight.detach().float().view(32, 2, 32).contiguous(), so2.attn_qk_norm.adam_scale.detach().float().contiguous(), diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_c_attention_backward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_c_attention_backward.py index e7cbb703a9..582b705bdc 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_c_attention_backward.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_c_attention_backward.py @@ -162,6 +162,7 @@ def _grouped_expanded_adjoint_jit( focus_tau: cutlass.Constexpr[float], label_smoothing: cutlass.Constexpr[float], qk_scale: cutlass.Constexpr[float], + use_focus_norm: cutlass.Constexpr[bool], stream: CUstream, ): _grouped_expanded_adjoint_kernel( @@ -193,6 +194,7 @@ def _grouped_expanded_adjoint_jit( focus_tau, label_smoothing, qk_scale, + use_focus_norm, ).launch( grid=[dst_ptr.shape[0] - 1, 1, 1], block=[THREADS, 1, 1], @@ -230,6 +232,7 @@ def _grouped_expanded_adjoint_kernel( focus_tau: cutlass.Constexpr[float], label_smoothing: cutlass.Constexpr[float], qk_scale: cutlass.Constexpr[float], + use_focus_norm: cutlass.Constexpr[bool], ): tidx, _, _ = cute.arch.thread_idx() node, _, _ = cute.arch.block_idx() @@ -470,19 +473,22 @@ def _grouped_expanded_adjoint_kernel( focus_grad_probability = focus_grad1 focus_grad_logit = probability * (focus_grad_probability - focus_dot) * inv_tau - focus_value_raw = focus_src[edge, focus, lane].to(cutlass.Float32) - scale = focus_scale[focus, lane].to(cutlass.Float32) weight = focus_weight[lane, focus].to(cutlass.Float32) - inv_rms = cute.rsqrt( - _warp_sum(focus_value_raw * focus_value_raw) / cutlass.Float32(CHANNELS) - + cutlass.Float32(focus_eps) - ) - grad_scaled = focus_grad_logit * weight * scale - coeff = _warp_sum(grad_scaled * focus_value_raw) / cutlass.Float32(CHANNELS) - grad_focus_src[focus, edge, lane] = ( - grad_scaled * inv_rms - - focus_value_raw * inv_rms * inv_rms * inv_rms * coeff - ) + if cutlass.const_expr(use_focus_norm): + focus_value_raw = focus_src[edge, focus, lane].to(cutlass.Float32) + scale = focus_scale[focus, lane].to(cutlass.Float32) + inv_rms = cute.rsqrt( + _warp_sum(focus_value_raw * focus_value_raw) / cutlass.Float32(CHANNELS) + + cutlass.Float32(focus_eps) + ) + grad_scaled = focus_grad_logit * weight * scale + coeff = _warp_sum(grad_scaled * focus_value_raw) / cutlass.Float32(CHANNELS) + grad_focus_src[focus, edge, lane] = ( + grad_scaled * inv_rms + - focus_value_raw * inv_rms * inv_rms * inv_rms * coeff + ) + else: + grad_focus_src[focus, edge, lane] = focus_grad_logit * weight if local == 0: grad_attention0 = beta_adjoint[edge_slot * FOCUS_COUNT] * focus_alpha[ @@ -518,6 +524,7 @@ def compile_grouped_expanded_final_phase_c_attention_adjoint( focus_tau: float, label_smoothing: float, qk_scale: float = QK_SCALE, + use_focus_norm: bool = True, ) -> Callable: """Compile the fixed four-group SM90 schedule.""" edges = cute.sym_int64() @@ -594,6 +601,7 @@ def compile_grouped_expanded_final_phase_c_attention_adjoint( float(focus_tau), float(label_smoothing), float(qk_scale), + bool(use_focus_norm), make_fake_stream(use_tvm_ffi_env_stream=True), options="--enable-tvm-ffi", ) @@ -685,6 +693,7 @@ def run_grouped_expanded_final_phase_c_attention_adjoint( focus_tau: float, label_smoothing: float, qk_scale: float = QK_SCALE, + use_focus_norm: bool = True, outputs: GroupedExpandedFinalPhaseCAttentionAdjointOutputs | None = None, ) -> GroupedExpandedFinalPhaseCAttentionAdjointOutputs: """Run the adjoint; reusable outputs require clearing ``grad_k_node``.""" @@ -824,6 +833,7 @@ def run_grouped_expanded_final_phase_c_attention_adjoint( float(focus_tau), float(label_smoothing), float(qk_scale), + bool(use_focus_norm), )( b0, torch.view_as_real(b1), diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/runner.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/runner.py index 5e1f5862ab..a937402c65 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/runner.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/runner.py @@ -113,11 +113,13 @@ def _prepare_persistent_weights(so2: Any) -> NeoPersistentComplexWeights: for nonlinearity in nonlinearities[:GATED_LAYERS]: if getattr(nonlinearity, "layout", None) != "fndc": raise NotImplementedError("SM90 SO2 requires fndc gates") - activation = getattr( - getattr(nonlinearity, "scalar_act", None), - "activation", - None, - ) + activation = getattr(nonlinearity, "activation_function", None) + if activation is None: + activation = getattr( + getattr(nonlinearity, "scalar_act", None), + "activation", + None, + ) if str(activation).lower() != "silu": raise NotImplementedError("SM90 SO2 requires SiLU gates") if getattr(nonlinearity.gate_linear, "bias", None) is not None: @@ -353,7 +355,7 @@ def _build_forward_graph(self) -> None: self.focus_gate_src.view(edge_count, HIDDEN), x_l0_node.contiguous(), so2.adamw_focus_compete_w.detach().float().contiguous(), - so2.focus_compete_norm.adam_scale.detach().float().contiguous(), + self.focus_norm_scale, so2.attn_q_proj.weight.detach() .float() .view(CHANNELS, FOCUS_COUNT, CHANNELS) @@ -586,19 +588,17 @@ def _runner_backward( focus_alpha=runner.focus_alpha, focus_src=runner.focus_gate_src, focus_weight=so2.adamw_focus_compete_w.detach().float().contiguous(), - focus_scale=so2.focus_compete_norm.adam_scale.detach() - .float() - .reshape(FOCUS_COUNT, CHANNELS) - .contiguous(), + focus_scale=runner.focus_norm_scale, q_node=runner.q_node, k_node=runner.k_node, edge_gate=runner.edge_gate, src=runner.src_i32, dst_ptr=runner.dst_ptr_i32, - focus_eps=float(so2.focus_compete_norm.eps), + focus_eps=runner.focus_norm_eps, focus_tau=float(so2.focus_softmax_tau), label_smoothing=float(so2.focus_label_smoothing), qk_scale=CHANNELS**-0.5, + use_focus_norm=runner.focus_norm_enabled, outputs=fused_outputs, ) grad_stack = NeoPersistentComplexState( diff --git a/source/tests/pt/model/sezm_cute/test_so2.py b/source/tests/pt/model/sezm_cute/test_so2.py index 1fc07f52b2..2a830e3e30 100644 --- a/source/tests/pt/model/sezm_cute/test_so2.py +++ b/source/tests/pt/model/sezm_cute/test_so2.py @@ -24,6 +24,7 @@ ) from deepmd.pt.model.descriptor.sezm_nn.norm import ( EquivariantRMSNorm, + ScalarRMSNorm, ) from deepmd.pt_expt.kernels.cute.sezm import ( runtime_policy, @@ -89,6 +90,16 @@ def _neo_like_block(**overrides): ) for key, value in overrides.items(): setattr(so2, key, value) + so2.focus_compete_norm = ( + ScalarRMSNorm( + channels=32, + n_focus=2, + dtype=torch.float32, + trainable=False, + ) + if so2.focus_norm + else None + ) block = torch.nn.Module() block.so2_conv = so2 block.lmax = 3 @@ -107,8 +118,9 @@ def _neo_like_block(**overrides): return block -def test_exact_neo_contract_is_supported() -> None: - block = _neo_like_block() +@pytest.mark.parametrize("focus_norm", [True, False]) +def test_exact_neo_contract_is_supported(focus_norm: bool) -> None: + block = _neo_like_block(focus_norm=focus_norm) assert _SO2.get_neo_so2_spec(block).is_current_neo_target assert _SO2.is_supported_neo_so2_block(block) @@ -401,7 +413,10 @@ def _cute_cuda_runtime_available() -> bool: not _cute_cuda_runtime_available(), reason="CuTe attention-prelude regression requires CUDA and CuTe DSL", ) -def test_attention_prelude_initializes_all_qk_rows_when_edges_are_sparse() -> None: +@pytest.mark.parametrize("use_focus_norm", [True, False]) +def test_attention_prelude_initializes_all_qk_rows_when_edges_are_sparse( + use_focus_norm: bool, +) -> None: from deepmd.pt_expt.kernels.cute.sezm.so2.kernels.focus_source_backward import ( compile_neo_attention_prelude_forward, ) @@ -426,14 +441,16 @@ def test_attention_prelude_initializes_all_qk_rows_when_edges_are_sparse() -> No label_smoothing = 0.1 focus_view = focus.view(edge_count, 2, 32) - focus_norm = ( - focus_view - * torch.rsqrt(focus_view.square().mean(dim=-1, keepdim=True) + focus_eps) - * focus_scale.unsqueeze(0) - ) + focus_input = focus_view + if use_focus_norm: + focus_input = ( + focus_view + * torch.rsqrt(focus_view.square().mean(dim=-1, keepdim=True) + focus_eps) + * focus_scale.unsqueeze(0) + ) focus_logits = torch.stack( [ - (focus_norm[:, focus_idx] * focus_weight[:, focus_idx]).sum(dim=-1) + (focus_input[:, focus_idx] * focus_weight[:, focus_idx]).sum(dim=-1) for focus_idx in range(2) ], dim=-1, @@ -453,6 +470,7 @@ def test_attention_prelude_initializes_all_qk_rows_when_edges_are_sparse() -> No qk_eps, tau, label_smoothing, + use_focus_norm=use_focus_norm, ) run( focus, diff --git a/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py b/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py index e3ae46c1c4..32dce863c7 100644 --- a/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py +++ b/source/tests/pt_expt/descriptor/test_dpa4_accelerated.py @@ -98,7 +98,11 @@ def _make_descriptor( ) -def make_neo_descriptor(sel: int | list[int] = 4) -> DescrptDPA4: +def make_neo_descriptor( + sel: int | list[int] = 4, + *, + edge_norm: bool = True, +) -> DescrptDPA4: """Build the exact Neo layout served by the current CuTe kernels.""" return DescrptDPA4( ntypes=2, @@ -117,6 +121,7 @@ def make_neo_descriptor(sel: int | list[int] = 4) -> DescrptDPA4: so3_readout="mlp", use_amp=False, random_gamma=False, + edge_norm=edge_norm, precision="float32", trainable=False, seed=42, @@ -476,7 +481,12 @@ def test_source_gated_flash_retains_dense_rotations(self, monkeypatch) -> None: np.testing.assert_allclose(output, dense_output, rtol=2e-4, atol=2e-5) np.testing.assert_allclose(gradient, dense_gradient, rtol=2e-4, atol=2e-5) - def test_cute_neo_forward_and_coordinate_gradient(self, monkeypatch) -> None: + @pytest.mark.parametrize("edge_norm", [True, False]) + def test_cute_neo_forward_and_coordinate_gradient( + self, + monkeypatch, + edge_norm: bool, + ) -> None: """Exercise packed Wigner and SO2 through the PT-expt descriptor.""" try: import cuda.bindings.driver # noqa: F401 @@ -499,7 +509,7 @@ def test_cute_neo_forward_and_coordinate_gradient(self, monkeypatch) -> None: "DP_CUTE_INFER", ): monkeypatch.setenv(name, "0") - data = make_neo_descriptor(self.sel_mix).serialize() + data = make_neo_descriptor(self.sel_mix, edge_norm=edge_norm).serialize() reference = DescrptDPA4.deserialize(data).to(self.device).eval() reference_output, reference_gradient = self._step(reference) From 895193022aa0758b7ca0add5ddfa994d433b07ea Mon Sep 17 00:00:00 2001 From: atulcthakur Date: Tue, 15 Sep 2026 00:05:11 -0700 Subject: [PATCH 08/11] fix(pt): harden Neo CuTe SO2 dispatch --- deepmd/pt/model/model/sezm_model.py | 2 +- .../pt_expt/kernels/cute/sezm/so2/metadata.py | 46 +++++++-- .../kernels/cute/sezm/so2/operation.py | 30 +++++- .../pt_expt/kernels/cute/sezm/so2/runner.py | 2 +- .../so2/sm90/phase_c_attention_backward.py | 2 +- source/tests/pt/model/sezm_cute/test_efs.py | 93 +++++++++++++++++-- source/tests/pt/model/sezm_cute/test_so2.py | 28 ++++++ .../pt/model/sezm_cute/test_so2_dst_ptr.py | 68 +++++++++++++- 8 files changed, 244 insertions(+), 27 deletions(-) diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py index 33c40deee2..0b07de4d11 100644 --- a/deepmd/pt/model/model/sezm_model.py +++ b/deepmd/pt/model/model/sezm_model.py @@ -2189,7 +2189,7 @@ def trace_and_compile( self.atomic_model.descriptor.blocks, training=self.training, device=coord.device, - dtype=coord.dtype, + dtype=self.atomic_model.descriptor.dtype, ) if prepared_cute_so2 and cute_runtime_policy.is_so2_thin_wrapper_enabled( compute_capability diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/metadata.py b/deepmd/pt_expt/kernels/cute/sezm/so2/metadata.py index c47cf609f1..56a59d04bb 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/so2/metadata.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/metadata.py @@ -9,6 +9,41 @@ import torch +def _destination_row_ptr_impl( + dst: torch.Tensor, + n_nodes: int, +) -> torch.Tensor: + """Build destination CSR where real storage offsets remain visible.""" + boundaries = torch.arange( + n_nodes + 1, + device=dst.device, + dtype=dst.dtype, + ) + return torch.searchsorted( + dst.contiguous(), + boundaries, + out_int32=True, + ).contiguous() + + +_destination_row_ptr_op = torch.library.custom_op( + "sezm_cute::destination_row_ptr", + mutates_args=(), +)(_destination_row_ptr_impl) + + +@_destination_row_ptr_op.register_fake +def _( + dst: torch.Tensor, + n_nodes: int, +) -> torch.Tensor: + return torch.empty( + (n_nodes + 1,), + device=dst.device, + dtype=torch.int32, + ) + + def build_sorted_edge_index_metadata( src: torch.Tensor, dst: torch.Tensor, @@ -66,16 +101,7 @@ def build_sorted_edge_index_metadata( "Neo SO2 destinations_sorted=True requires monotonically " "nondecreasing destination indices", ) - destination_boundaries = torch.arange( - n_nodes + 1, - device=dst.device, - dtype=dst.dtype, - ) - destination_row_ptr = torch.searchsorted( - dst, - destination_boundaries, - out_int32=True, - ).contiguous() + destination_row_ptr = _destination_row_ptr_op(dst, n_nodes) source_order_i64 = torch.argsort(src, stable=True) sorted_src = src.index_select(0, source_order_i64) diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py b/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py index dbbe1f5f9f..a7882b687a 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py @@ -11,6 +11,7 @@ from dataclasses import ( dataclass, field, + replace, ) from functools import ( lru_cache, @@ -1358,6 +1359,20 @@ def _make_edge_cache( ) +def _destination_degrees_fit_limit( + destination_row_ptr: Tensor, + max_edges_per_node: int, +) -> bool: + """Return whether every destination row fits a bounded native kernel.""" + destination_degrees = destination_row_ptr[1:] - destination_row_ptr[:-1] + return bool( + torch.all( + (destination_degrees >= 0) + & (destination_degrees <= max_edges_per_node) + ).item() + ) + + def _build_runner( handle: int, x: Tensor, @@ -1383,7 +1398,20 @@ def _build_runner( entry = _REGISTRY[int(handle)] config = entry.config if config.native_sm90_path: - from .sm90.runner import NeoSm90SO2Runner as Runner + from .sm90.phase_c_attention_backward import ( + MAX_EDGES_PER_NODE, + ) + + native_degree_eligible = _destination_degrees_fit_limit( + dst_ptr, + MAX_EDGES_PER_NODE, + ) + if native_degree_eligible: + from .sm90.runner import NeoSm90SO2Runner as Runner + else: + from .runner import NeoFullCuteBackward as Runner + + config = replace(config, native_sm90_path=False) else: from .runner import NeoFullCuteBackward as Runner diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/runner.py b/deepmd/pt_expt/kernels/cute/sezm/so2/runner.py index 32e92e79fc..8cf05fe469 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/so2/runner.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/runner.py @@ -36,7 +36,7 @@ def _validate_runtime_config( ) -> None: if compute_capability not in SUPPORTED_SO2_CAPABILITIES: raise RuntimeError("Neo SO2 requires a supported compute capability") - if runtime_config.native_sm90_path != (compute_capability == SM90_CAPABILITY): + if runtime_config.native_sm90_path and compute_capability != SM90_CAPABILITY: raise RuntimeError("the native SM90 SO2 path must be selected only on sm_90") if runtime_config.per_focus_so2_fwd_pair != ( compute_capability in SM80_PROFILE_CAPABILITIES diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_c_attention_backward.py b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_c_attention_backward.py index 582b705bdc..bb8f93eda2 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_c_attention_backward.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/sm90/phase_c_attention_backward.py @@ -59,7 +59,7 @@ M0_WIDTH = 128 M1_WIDTH = 96 PACKED_WIGNER_VALUES = PACKED_VALUE_COUNT -MAX_EDGES_PER_NODE = 128 +MAX_EDGES_PER_NODE = 256 QK_SCALE = CHANNELS**-0.5 B0_NODE_VALUES = FOCUS_COUNT * DEGREE_COUNT * M0_WIDTH B1_NODE_VALUES = FOCUS_COUNT * (DEGREE_COUNT - 1) * M1_WIDTH * 2 diff --git a/source/tests/pt/model/sezm_cute/test_efs.py b/source/tests/pt/model/sezm_cute/test_efs.py index b89fa5bfd3..7dc6d855f4 100644 --- a/source/tests/pt/model/sezm_cute/test_efs.py +++ b/source/tests/pt/model/sezm_cute/test_efs.py @@ -73,7 +73,7 @@ def _strict_fp32() -> Iterator[None]: torch.set_float32_matmul_precision(prior_precision) -def _neo_model(*, use_compile: bool) -> torch.nn.Module: +def _neo_model(*, use_compile: bool, sel: int = 32) -> torch.nn.Module: from deepmd.pt.model.model import ( get_sezm_model, ) @@ -84,7 +84,7 @@ def _neo_model(*, use_compile: bool) -> torch.nn.Module: "type_map": ["O", "H"], "descriptor": { "type": "SeZM", - "sel": 32, + "sel": sel, "rcut": 3.0, "channels": 32, "n_radial": 16, @@ -124,7 +124,7 @@ def _neo_model(*, use_compile: bool) -> torch.nn.Module: return model -def _water_frame() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +def _water_frame() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: coord = torch.tensor( [ [ @@ -137,7 +137,7 @@ def _water_frame() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ] ], device="cuda", - dtype=torch.float32, + dtype=torch.float64, ) atype = torch.tensor( [[0, 1, 1, 0, 1, 1]], @@ -147,12 +147,12 @@ def _water_frame() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: box = torch.tensor( [[5.4, 0.0, 0.0, 0.0, 5.2, 0.0, 0.0, 0.0, 5.0]], device="cuda", - dtype=torch.float32, + dtype=torch.float64, ) return coord, atype, box -def _triclinic_frame() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +def _triclinic_frame() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: box_matrix = torch.tensor( [ [5.2, 0.0, 0.0], @@ -185,6 +185,24 @@ def _triclinic_frame() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: return coord, atype, box_matrix.reshape(1, 9) +def _high_degree_frame() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + node_count = 258 + coord = torch.zeros( + (1, node_count, 3), + device="cuda", + dtype=torch.float32, + ) + coord[0, :, 0] = torch.arange( + node_count, + device="cuda", + dtype=torch.float32, + ) * 0.01 + atype = (torch.arange(node_count, device="cuda") % 2).to( + dtype=torch.int32 + )[None, :] + return coord, atype, None + + def _detached_outputs(outputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: return {name: outputs[name].detach().cpu().clone() for name in OUTPUT_KEYS} @@ -192,7 +210,9 @@ def _detached_outputs(outputs: dict[str, torch.Tensor]) -> dict[str, torch.Tenso _FRAME_FACTORIES = { "water": _water_frame, "triclinic": _triclinic_frame, + "high_degree": _high_degree_frame, } +_STANDARD_FRAME_NAMES = ("water", "triclinic") def _assert_runtime_profile() -> dict[str, bool]: @@ -273,17 +293,28 @@ def _run_efs_child( torch.manual_seed(20260726) torch.cuda.manual_seed_all(20260726) runner_calls = 0 + runner_types: set[str] = set() + max_destination_degree = 0 original_build_runner = so2._build_runner def counted_build_runner(*args: Any, **kwargs: Any) -> Any: - nonlocal runner_calls + nonlocal runner_calls, max_destination_degree + runner = original_build_runner(*args, **kwargs) runner_calls += 1 - return original_build_runner(*args, **kwargs) + runner_types.add(f"{type(runner).__module__}.{type(runner).__qualname__}") + if runner_calls == 1: + destination_row_ptr = args[8] + degrees = destination_row_ptr[1:] - destination_row_ptr[:-1] + max_destination_degree = int(degrees.max().item()) + return runner if use_compile: so2._build_runner = counted_build_runner try: - model = _neo_model(use_compile=use_compile) + model = _neo_model( + use_compile=use_compile, + sel=288 if frame_name == "high_degree" else 32, + ) state_keys_before = tuple(model.state_dict()) coord, atype, box = _FRAME_FACTORIES[frame_name]() with _strict_fp32(): @@ -308,6 +339,8 @@ def counted_build_runner(*args: Any, **kwargs: Any) -> Any: "outputs": runs, "profile": profile, "runner_calls": runner_calls, + "runner_types": sorted(runner_types), + "max_destination_degree": max_destination_degree, "use_compile": use_compile, }, output_path, @@ -394,7 +427,7 @@ def _run_child_process( _CUTE_SKIP_REASON is not None, reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable", ) -@pytest.mark.parametrize("frame_name", tuple(_FRAME_FACTORIES)) +@pytest.mark.parametrize("frame_name", _STANDARD_FRAME_NAMES) def test_neo_cute_compiled_profile_matches_eager_at_shared_5e5( tmp_path: Path, frame_name: str, @@ -430,6 +463,46 @@ def test_neo_cute_compiled_profile_matches_eager_at_shared_5e5( ) +@pytest.mark.skipif( + _CUTE_SKIP_REASON is not None, + reason=_CUTE_SKIP_REASON or "CuTe runtime unavailable", +) +def test_sm90_high_degree_force_uses_portable_cute_fallback( + tmp_path: Path, +) -> None: + if tuple(torch.cuda.get_device_capability()) != (9, 0): + pytest.skip("the bounded native SO2 runner exists only on SM90") + from deepmd.pt_expt.kernels.cute.sezm.so2.sm90.phase_c_attention_backward import ( + MAX_EDGES_PER_NODE, + ) + + eager = _run_child_process( + frame_name="high_degree", + use_compile=False, + output_path=tmp_path / "high-degree-eager.pt", + ) + compiled = _run_child_process( + frame_name="high_degree", + use_compile=True, + output_path=tmp_path / "high-degree-compiled.pt", + ) + + assert compiled["runner_calls"] > 0 + assert compiled["max_destination_degree"] > MAX_EDGES_PER_NODE + assert any("NeoFullCuteBackward" in name for name in compiled["runner_types"]) + assert all("NeoSm90SO2Runner" not in name for name in compiled["runner_types"]) + expected = eager["outputs"][0] + for actual in compiled["outputs"]: + for name in OUTPUT_KEYS: + torch.testing.assert_close( + actual[name], + expected[name], + atol=TOL, + rtol=TOL, + msg=f"high-degree SM90 portable CuTe {name} differs from eager", + ) + + def _main() -> None: if len(sys.argv) != 5 or sys.argv[1] != _CHILD_MODE: raise SystemExit("this test module is only executable in E/F/S child mode") diff --git a/source/tests/pt/model/sezm_cute/test_so2.py b/source/tests/pt/model/sezm_cute/test_so2.py index 2a830e3e30..98fdd846eb 100644 --- a/source/tests/pt/model/sezm_cute/test_so2.py +++ b/source/tests/pt/model/sezm_cute/test_so2.py @@ -9,6 +9,7 @@ import importlib from dataclasses import ( fields, + replace, ) from types import ( SimpleNamespace, @@ -38,6 +39,7 @@ StackCache = _SO2_RUNNER.StackCache _validate_runtime_config = _SO2_RUNNER._validate_runtime_config _uses_packed_message_grid = _SO2_RUNNER._uses_packed_message_grid +_destination_degrees_fit_limit = _SO2._destination_degrees_fit_limit class _Identity: @@ -203,6 +205,32 @@ def test_sm90_uses_native_split_complex_path() -> None: assert not config.combined_so2_gate +def test_sm90_portable_fallback_is_a_validated_config() -> None: + config = replace( + _SO2._architecture_default_config((9, 0)), + native_sm90_path=False, + ) + + assert _validate_runtime_config(config, compute_capability=(9, 0)) is None + + +@pytest.mark.parametrize( + ("row_ptr", "expected"), + [ + ([0, 256, 512], True), + ([0, 257, 512], False), + ([0, 4, 3], False), + ], +) +def test_destination_degree_limit( + row_ptr: list[int], + expected: bool, +) -> None: + destination_row_ptr = torch.tensor(row_ptr, dtype=torch.int32, device="cpu") + + assert _destination_degrees_fit_limit(destination_row_ptr, 256) is expected + + def test_sm100_uses_shared_default_profile() -> None: config = _SO2._architecture_default_config((10, 0)) diff --git a/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py b/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py index d2efab4cfa..b199903645 100644 --- a/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py +++ b/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py @@ -9,6 +9,9 @@ import ast import importlib import unittest +from itertools import ( + count, +) from types import ( SimpleNamespace, ) @@ -67,7 +70,18 @@ def _load_extracted_sorted_metadata_function(): function = _function(tree, "build_sorted_edge_index_metadata") module = ast.Module(body=[function], type_ignores=[]) ast.fix_missing_locations(module) - namespace = {"torch": torch} + def destination_row_ptr(dst, n_nodes): + boundaries = torch.arange( + n_nodes + 1, + device=dst.device, + dtype=dst.dtype, + ) + return torch.searchsorted(dst, boundaries, out_int32=True).contiguous() + + namespace = { + "torch": torch, + "_destination_row_ptr_op": destination_row_ptr, + } exec(compile(module, str(SO2_METADATA_PATH), "exec"), namespace) return namespace["build_sorted_edge_index_metadata"] @@ -271,14 +285,15 @@ def test_compile_cold_registration_survives_into_packed_custom_op_runtime(self): assert torch is not None so2 = importlib.import_module("deepmd.pt_expt.kernels.cute.sezm.so2.operation") prior_registry = dict(so2._REGISTRY) - prior_next_handle = so2._NEXT_HANDLE + prior_next_handle = next(so2._HANDLE_COUNTER) + so2._HANDLE_COUNTER = count(prior_next_handle) def cleanup(): torch._dynamo.reset() so2._PACKED_RUNNER_CACHE.clear() so2._REGISTRY.clear() so2._REGISTRY.update(prior_registry) - so2._NEXT_HANDLE = prior_next_handle + so2._HANDLE_COUNTER = count(prior_next_handle) self.addCleanup(cleanup) so2._REGISTRY.clear() @@ -506,6 +521,53 @@ def build_ptr(src, dst): atol=0.0, ) + @unittest.skipUnless( + torch is not None and torch.cuda.is_available(), + "offset destination compile regression requires CUDA", + ) + def test_inductor_handles_offset_destination_view(self): + assert torch is not None + builder = _load_sorted_metadata_function() + + def build_metadata(edge_index): + return builder(edge_index[0], edge_index[1], 4) + + compiled = torch.compile( + build_metadata, + backend="inductor", + dynamic=True, + fullgraph=True, + ) + edge_index = torch.tensor( + [ + [2, 0, 3, 0, 1, 2], + [0, 0, 1, 2, 2, 3], + ], + dtype=torch.int64, + device="cuda", + ) + self.assertGreater(edge_index[1].storage_offset(), 0) + + dst_ptr, source_order, source_ptr = compiled(edge_index) + torch.testing.assert_close( + dst_ptr, + torch.tensor([0, 2, 3, 5, 6], dtype=torch.int32, device="cuda"), + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + edge_index[0].index_select(0, source_order.to(torch.int64)), + torch.tensor([0, 0, 1, 2, 2, 3], device="cuda"), + rtol=0.0, + atol=0.0, + ) + torch.testing.assert_close( + source_ptr, + torch.tensor([0, 2, 3, 5, 6], dtype=torch.int32, device="cuda"), + rtol=0.0, + atol=0.0, + ) + def test_dynamic_node_and_edge_counts_build_independent_metadata(self): assert torch is not None builder = _load_sorted_metadata_function() From b89d94228f37d05b7055f7098af08c7d9d4e45da Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 07:06:38 +0000 Subject: [PATCH 09/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../pt_expt/kernels/cute/sezm/so2/operation.py | 3 +-- source/tests/pt/model/sezm_cute/test_efs.py | 17 +++++++++-------- .../pt/model/sezm_cute/test_so2_dst_ptr.py | 1 + 3 files changed, 11 insertions(+), 10 deletions(-) diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py b/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py index a7882b687a..8f23181ede 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py @@ -1367,8 +1367,7 @@ def _destination_degrees_fit_limit( destination_degrees = destination_row_ptr[1:] - destination_row_ptr[:-1] return bool( torch.all( - (destination_degrees >= 0) - & (destination_degrees <= max_edges_per_node) + (destination_degrees >= 0) & (destination_degrees <= max_edges_per_node) ).item() ) diff --git a/source/tests/pt/model/sezm_cute/test_efs.py b/source/tests/pt/model/sezm_cute/test_efs.py index 7dc6d855f4..a729ba5fd1 100644 --- a/source/tests/pt/model/sezm_cute/test_efs.py +++ b/source/tests/pt/model/sezm_cute/test_efs.py @@ -192,14 +192,15 @@ def _high_degree_frame() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | Non device="cuda", dtype=torch.float32, ) - coord[0, :, 0] = torch.arange( - node_count, - device="cuda", - dtype=torch.float32, - ) * 0.01 - atype = (torch.arange(node_count, device="cuda") % 2).to( - dtype=torch.int32 - )[None, :] + coord[0, :, 0] = ( + torch.arange( + node_count, + device="cuda", + dtype=torch.float32, + ) + * 0.01 + ) + atype = (torch.arange(node_count, device="cuda") % 2).to(dtype=torch.int32)[None, :] return coord, atype, None diff --git a/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py b/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py index b199903645..e1646699d9 100644 --- a/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py +++ b/source/tests/pt/model/sezm_cute/test_so2_dst_ptr.py @@ -70,6 +70,7 @@ def _load_extracted_sorted_metadata_function(): function = _function(tree, "build_sorted_edge_index_metadata") module = ast.Module(body=[function], type_ignores=[]) ast.fix_missing_locations(module) + def destination_row_ptr(dst, n_nodes): boundaries = torch.arange( n_nodes + 1, From 9395ff87820851ee82f47a2a7370c64ea61d6c36 Mon Sep 17 00:00:00 2001 From: atulcthakur Date: Tue, 15 Sep 2026 10:47:46 -0700 Subject: [PATCH 10/11] fix(pt): preserve destination offsets in CuTe GIE metadata Reuse the existing destination CSR custom op so Inductor does not lose the storage offset of edge_index[1]. Keep GIE enabled and add compiled forward/backward regression coverage with dynamic graph sizes and optional gates. --- deepmd/pt_expt/kernels/cute/sezm/gie.py | 11 +- .../pt_expt/kernels/cute/sezm/so2/metadata.py | 22 +++- source/tests/pt/model/sezm_cute/test_gie.py | 103 ++++++++++++++++++ 3 files changed, 129 insertions(+), 7 deletions(-) diff --git a/deepmd/pt_expt/kernels/cute/sezm/gie.py b/deepmd/pt_expt/kernels/cute/sezm/gie.py index 8208164483..972ece4f88 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/gie.py +++ b/deepmd/pt_expt/kernels/cute/sezm/gie.py @@ -30,6 +30,10 @@ Tensor, ) +from .so2.metadata import ( + build_destination_row_ptr, +) + if TYPE_CHECKING: from collections.abc import ( Callable, @@ -678,12 +682,7 @@ def gie_fused_cuda( """Run the fused CUDA path after the caller has validated its contract.""" if not SEZM_CUTE_GIE_AVAILABLE: raise RuntimeError("CuTe DSL is unavailable") - boundaries = torch.arange( - n_nodes + 1, - device=dst.device, - dtype=dst.dtype, - ) - dst_ptr = torch.searchsorted(dst, boundaries) + dst_ptr = build_destination_row_ptr(dst, n_nodes) has_gate = gate.numel() != 0 kernel_gate = gate if has_gate else radial.new_ones((1,)) return _gie_op( diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/metadata.py b/deepmd/pt_expt/kernels/cute/sezm/so2/metadata.py index 56a59d04bb..99cb83220e 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/so2/metadata.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/metadata.py @@ -44,6 +44,26 @@ def _( ) +def build_destination_row_ptr(dst: torch.Tensor, n_nodes: int) -> torch.Tensor: + """Build integer CSR for sorted destinations through an opaque runtime op. + + Parameters + ---------- + dst : torch.Tensor + Nondecreasing int32 or int64 destinations with shape ``(E,)``. + n_nodes : int + Number of nodes addressed by the edge list. + + Returns + ------- + torch.Tensor + Contiguous int32 row pointers with shape ``(n_nodes + 1,)`` and no + autograd history. The runtime op preserves nonzero storage offsets + when ``dst`` is a view such as ``edge_index[1]`` under Inductor. + """ + return _destination_row_ptr_op(dst, n_nodes) + + def build_sorted_edge_index_metadata( src: torch.Tensor, dst: torch.Tensor, @@ -119,4 +139,4 @@ def build_sorted_edge_index_metadata( return destination_row_ptr, source_order, source_row_ptr -__all__ = ["build_sorted_edge_index_metadata"] +__all__ = ["build_destination_row_ptr", "build_sorted_edge_index_metadata"] diff --git a/source/tests/pt/model/sezm_cute/test_gie.py b/source/tests/pt/model/sezm_cute/test_gie.py index 9a132c997e..0dce88913a 100644 --- a/source/tests/pt/model/sezm_cute/test_gie.py +++ b/source/tests/pt/model/sezm_cute/test_gie.py @@ -14,6 +14,9 @@ import torch from deepmd.pt_expt.kernels.cute.sezm import gie as gie_module +from deepmd.pt_expt.kernels.cute.sezm.so2.metadata import ( + build_destination_row_ptr, +) def _load_gie_module(): @@ -222,6 +225,106 @@ def test_cuda_forward_and_wigner_radial_degree_gate_gradients(self): actual_grad, expected_grad, rtol=2e-5, atol=2e-5 ) + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required") + def test_inductor_gie_handles_offset_destination_view(self): + gie = _load_gie_module() + if not gie.SEZM_CUTE_GIE_AVAILABLE: + self.skipTest("CuTe DSL is not available") + + def run_gie(edge_index, radial, zonal, inv, gate): + dst = edge_index[1] + n_nodes = inv.shape[0] + out = gie.gie_fused_cuda( + radial, + zonal, + inv, + dst, + gate, + n_nodes=n_nodes, + lmax=3, + ) + return out, build_destination_row_ptr(dst, n_nodes) + + compiled = torch.compile( + run_gie, backend="inductor", dynamic=True, fullgraph=True + ) + for with_gate in (False, True): + for n_nodes, src_values, dst_values, row_ptr in ( + (4, (2, 0, 3, 0, 1, 2), (0, 0, 1, 2, 2, 3), (0, 2, 3, 5, 6)), + ( + 7, + (6, 0, 5, 2, 0, 4, 1, 3, 2), + (0, 1, 1, 1, 4, 6, 6, 6, 6), + (0, 1, 4, 4, 4, 5, 5, 9), + ), + ): + with self.subTest( + gate=with_gate, n_nodes=n_nodes, edges=len(dst_values) + ): + expected_inputs = _inputs( + n_nodes=n_nodes, + dst_values=dst_values, + lmax=3, + channels=32, + device=torch.device("cuda"), + with_gate=with_gate, + ) + radial, dense_dt, zonal, inv, dst, gate = expected_inputs + weight = torch.randn_like(radial.new_empty(n_nodes, 16, 32)) + expected = _materialized_reference( + radial, + zonal, + inv, + dst, + gate, + n_nodes=n_nodes, + lmax=3, + ) + expected_grads = torch.autograd.grad( + (expected * weight).sum(), + (radial, dense_dt, zonal, inv, *([gate] if with_gate else [])), + ) + + actual_inputs = _inputs( + n_nodes=n_nodes, + dst_values=dst_values, + lmax=3, + channels=32, + device=torch.device("cuda"), + with_gate=with_gate, + ) + radial, dense_dt, zonal, inv, _dst, gate = actual_inputs + # The first row is deliberately not sorted: losing the second + # row's offset makes searchsorted operate on unrelated sources. + edge_index = torch.tensor( + (src_values, dst_values), dtype=torch.int64, device="cuda" + ) + self.assertTrue(edge_index[1].is_contiguous()) + self.assertEqual(edge_index[1].storage_offset(), len(dst_values)) + self.assertFalse(edge_index.requires_grad) + actual, dst_ptr = compiled(edge_index, radial, zonal, inv, gate) + self.assertEqual(dst_ptr.dtype, torch.int32) + self.assertTrue(dst_ptr.is_contiguous()) + self.assertFalse(dst_ptr.requires_grad) + self.assertIsNone(dst_ptr.grad_fn) + torch.testing.assert_close( + dst_ptr, + torch.tensor(row_ptr, dtype=torch.int32, device="cuda"), + rtol=0, + atol=0, + ) + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + actual_grads = torch.autograd.grad( + (actual * weight).sum(), + (radial, dense_dt, zonal, inv, *([gate] if with_gate else [])), + ) + for actual_grad, expected_grad in zip( + actual_grads, expected_grads, strict=True + ): + torch.testing.assert_close( + actual_grad, expected_grad, rtol=2e-5, atol=2e-5 + ) + if __name__ == "__main__": unittest.main() From 5d720facc05d69ec0cdeaa8ee28bce8b8f2c8dd0 Mon Sep 17 00:00:00 2001 From: atulcthakur Date: Thu, 17 Sep 2026 13:31:26 -0700 Subject: [PATCH 11/11] fix(pt): address CuTe inference review findings Recover declined packed-Wigner dispatch without retaining dense rotations, protect the local-only neighbor-list path from ghost exchange, and enforce frozen strict-FP32 state after preparation. Share the thin-wrapper device policy with MLP dispatch, validate runtime overrides consistently, and add focused regression tests. Keep labeled CUDA CI jobs enabled on subsequent pushes. --- .github/workflows/test_cuda.yml | 4 +- .../dpmodel/descriptor/dpa4_nn/edge_cache.py | 2 + .../dpmodel/descriptor/dpa4_nn/embedding.py | 2 +- deepmd/pt/model/descriptor/sezm.py | 1 - deepmd/pt/model/descriptor/sezm_nn/block.py | 19 ++- deepmd/pt/model/network/mlp.py | 35 ++--- deepmd/pt/utils/nv_nlist.py | 21 --- deepmd/pt_expt/descriptor/dpa4_nn/block.py | 20 ++- .../kernels/cute/sezm/runtime_policy.py | 5 +- .../kernels/cute/sezm/so2/operation.py | 25 ++- .../kernels/cute/sezm/so2/wigner_layout.py | 44 ++++++ .../common/dpmodel/test_dpa4_edge_cache.py | 21 +++ .../pt/model/sezm_cute/test_runtime_policy.py | 10 ++ source/tests/pt/model/sezm_cute/test_so2.py | 142 ++++++++++++++++++ source/tests/pt/model/test_mlp.py | 21 ++- source/tests/pt/model/test_sezm_parallel.py | 26 ++++ 16 files changed, 342 insertions(+), 56 deletions(-) diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index 0f3532fb45..bc37eea241 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -28,7 +28,7 @@ jobs: # container: # image: nvidia/cuda:12.9.1-cudnn-devel-ubuntu22.04 # options: --gpus all - if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group') + if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && contains(github.event.pull_request.labels.*.name, 'Test CUDA') || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group') steps: - uses: actions/checkout@v7 - uses: actions/setup-python@v7 @@ -79,7 +79,7 @@ jobs: test_cc: name: Test C++ on CUDA runs-on: gpu - if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group') + if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && contains(github.event.pull_request.labels.*.name, 'Test CUDA') || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group') steps: # Jobs run on separate runners, so the C++ job needs its own complete # CUDA toolchain and Python dependency installation. diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py b/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py index 8fc0f0e5fa..2c8a569099 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/edge_cache.py @@ -617,6 +617,8 @@ def _finalize_edge_cache( EdgeCache Finalized per-edge cache shared by eager and compile paths. """ + if deg_norm_floor <= 0.0: + raise ValueError("deg_norm_floor must be positive") xp = array_api_compat.array_namespace(edge_vec, dst) device = array_api_compat.device(edge_vec) # === Step 1. Build smooth destination degrees === diff --git a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py index c9e2563657..ba138049e7 100644 --- a/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py +++ b/deepmd/dpmodel/descriptor/dpa4_nn/embedding.py @@ -328,7 +328,7 @@ def call( # === Step 4. 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 5, so + # The fused operator spans this broadcast and the scatter of Step 6, so # it takes over whenever nothing else joins the message in between. if ( self.can_run_cuda_infer_l_1_scatter(zonal_coupling) diff --git a/deepmd/pt/model/descriptor/sezm.py b/deepmd/pt/model/descriptor/sezm.py index b0b446826b..f915e15e4b 100644 --- a/deepmd/pt/model/descriptor/sezm.py +++ b/deepmd/pt/model/descriptor/sezm.py @@ -1472,7 +1472,6 @@ def forward( x, edge_cache, rad_feat_per_block, - comm_dict=comm_dict, ) # === Step 12. Final l=0 output mixing === diff --git a/deepmd/pt/model/descriptor/sezm_nn/block.py b/deepmd/pt/model/descriptor/sezm_nn/block.py index 7a27ca99c3..e09038a3b2 100644 --- a/deepmd/pt/model/descriptor/sezm_nn/block.py +++ b/deepmd/pt/model/descriptor/sezm_nn/block.py @@ -832,8 +832,23 @@ def _run_so2_unit_impl( ) if accelerated is not None: return accelerated - raise RuntimeError( - "packed Wigner cache reached an ineligible Neo SO2 dispatch" + from deepmd.pt_expt.kernels.cute.sezm.so2.wigner_layout import ( + dense_wigner_for_fallback, + ) + + if edge_cache.edge_quat is None: + raise ValueError("packed Wigner fallback requires edge quaternions") + d_full, dt_full = dense_wigner_for_fallback( + edge_cache.edge_quat, + lmax=self.lmax, + eps=self.so2_conv.eps, + ) + edge_cache = edge_cache._replace( + D_full=d_full, + Dt_full=dt_full, + D_packed=None, + D_to_m_cache=None, + Dt_from_m_cache=None, ) n_node = x.shape[0] diff --git a/deepmd/pt/model/network/mlp.py b/deepmd/pt/model/network/mlp.py index 1f6ad02b35..705d05c362 100644 --- a/deepmd/pt/model/network/mlp.py +++ b/deepmd/pt/model/network/mlp.py @@ -1,5 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import os from typing import ( Any, ClassVar, @@ -51,25 +50,21 @@ def empty_t(shape: tuple[int, ...], precision: torch.dtype) -> torch.Tensor: def _use_so2_compile_visible_linear( input_device: torch.device | None = None, ) -> bool: - """Keep the SM80 linear topology stable for one compiled graph.""" - truthy = {"1", "true", "yes", "on"} - falsy = {"0", "false", "no", "off"} - cute_enabled = os.environ.get("DP_CUTE_INFER", "").strip().lower() - if cute_enabled not in truthy: - return False - thin_enabled = os.environ.get("DP_CUTE_SO2_THIN_WRAPPER", "").strip().lower() - if thin_enabled in falsy: - return False - if thin_enabled in truthy: - return True - if input_device is not None and input_device.type != "cuda": - return False - if not torch.cuda.is_available(): - return False - try: - return tuple(torch.cuda.get_device_capability(input_device)) == (8, 0) - except RuntimeError: - return False + """Use the same device-aware thin-wrapper policy as SO2 dispatch.""" + from deepmd.pt_expt.kernels.cute.sezm.runtime_policy import ( + is_so2_thin_wrapper_enabled, + ) + + capability = None + if ( + input_device is None or input_device.type == "cuda" + ) and torch.cuda.is_available(): + try: + capability = tuple(torch.cuda.get_device_capability(input_device)) + except RuntimeError: + pass + # A sentinel avoids querying another device when the input is on the CPU. + return is_so2_thin_wrapper_enabled(capability or (-1, -1)) def _matmul_bias( diff --git a/deepmd/pt/utils/nv_nlist.py b/deepmd/pt/utils/nv_nlist.py index 5b2b13ef33..3c68251a7d 100644 --- a/deepmd/pt/utils/nv_nlist.py +++ b/deepmd/pt/utils/nv_nlist.py @@ -178,27 +178,6 @@ def build( ``return_mode='edges'`` does not support ``pair_excl``; a :class:`NotImplementedError` is raised in that combination. """ - return self._build_impl( - coord, - atype, - box, - rcut, - sel, - return_mode=return_mode, - pair_excl=pair_excl, - ) - - def _build_impl( - self, - coord: Any, - atype: Any, - box: Any, - rcut: float, - sel: list[int], - return_mode: str = "extended", - pair_excl: PairExcludeMask | None = None, - ) -> tuple[Any, Any, Any, Any] | EdgeNeighborList: - """Implement the complete Toolkit-Ops build behind the dispatch guard.""" if return_mode == "edges" and pair_excl is not None: raise NotImplementedError( "pair_excl is not supported with return_mode='edges'; " diff --git a/deepmd/pt_expt/descriptor/dpa4_nn/block.py b/deepmd/pt_expt/descriptor/dpa4_nn/block.py index 095ebfad7f..732c4154f1 100644 --- a/deepmd/pt_expt/descriptor/dpa4_nn/block.py +++ b/deepmd/pt_expt/descriptor/dpa4_nn/block.py @@ -177,8 +177,24 @@ def _run_so2_unit_impl( ) if accelerated is not None: return accelerated - raise RuntimeError( - "packed Wigner cache reached an ineligible Neo SO2 dispatch" + from deepmd.pt_expt.kernels.cute.sezm.so2.wigner_layout import ( + dense_wigner_for_fallback, + ) + + if edge_cache.edge_quat is None: + raise ValueError("packed Wigner fallback requires edge quaternions") + d_full, dt_full = dense_wigner_for_fallback( + edge_cache.edge_quat, + lmax=self.lmax, + eps=self.so2_conv.eps, + ) + edge_cache = dataclasses.replace( + edge_cache, + D_full=d_full, + Dt_full=dt_full, + D_packed=None, + D_to_m_cache=None, + Dt_from_m_cache=None, ) return super()._run_so2_unit_impl(x, edge_cache, radial_feat) diff --git a/deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py b/deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py index 3cdd1730bf..320381d149 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py +++ b/deepmd/pt_expt/kernels/cute/sezm/runtime_policy.py @@ -54,7 +54,10 @@ def _env_override(name: str) -> bool | None: return True if normalized in _FALSE_VALUES: return False - return False + raise ValueError( + f"{name} must be a boolean value (0/1, false/true, no/yes, off/on); " + f"got {value!r}" + ) def is_cute_infer_enabled() -> bool: diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py b/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py index 8f23181ede..50c8d73927 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/operation.py @@ -396,9 +396,7 @@ def is_neo_so2_static_eligible( and _device_is_supported_for_so2(device) and dtype == torch.float32 and not torch.is_autocast_enabled(device.type) - and runtime_policy.uses_strict_fp32_matmul() - and _module_uses_strict_fp32(block) - and _module_is_frozen(block) + and _has_frozen_fp32_contract(block) and _module_state_is_aligned(block) and _gate_expand_index_structure_is_supported(block) and getattr(block, "_deepmd_cute_so2_state", None) is not False @@ -406,6 +404,16 @@ def is_neo_so2_static_eligible( ) +def _has_frozen_fp32_contract(block: Any) -> bool: + """Recheck mutable precision and parameter state before inference dispatch.""" + return ( + not block.training + and runtime_policy.uses_strict_fp32_matmul() + and _module_uses_strict_fp32(block) + and _module_is_frozen(block) + ) + + def is_neo_so2_runtime_eligible( block: Any, *, @@ -1395,6 +1403,13 @@ def _build_runner( """ with _REGISTRY_LOCK: entry = _REGISTRY[int(handle)] + # Compiled graphs keep the handle, not guards on the registered parameters. + # Enforce the contract here too if state changes after graph capture. + if not _has_frozen_fp32_contract(entry.block): + raise RuntimeError( + "CuTe SO2 requires frozen FP32 parameters and strict FP32 matmul; " + "model or precision state changed after preparation" + ) config = entry.config if config.native_sm90_path: from .sm90.phase_c_attention_backward import ( @@ -2044,6 +2059,10 @@ def _maybe_run_prepared_cute_so2( state = getattr(block, "_deepmd_cute_so2_state", None) if not isinstance(state, _RegisteredSO2State): return None + if not runtime_policy.is_cute_infer_enabled() or not _has_frozen_fp32_contract( + block + ): + return None with _REGISTRY_LOCK: entry = _REGISTRY.get(state.handle) if entry is None or entry.block is not block: diff --git a/deepmd/pt_expt/kernels/cute/sezm/so2/wigner_layout.py b/deepmd/pt_expt/kernels/cute/sezm/so2/wigner_layout.py index ff26fa0cc9..f6d977ae18 100644 --- a/deepmd/pt_expt/kernels/cute/sezm/so2/wigner_layout.py +++ b/deepmd/pt_expt/kernels/cute/sezm/so2/wigner_layout.py @@ -14,15 +14,24 @@ from dataclasses import ( dataclass, ) +from functools import ( + lru_cache, +) from typing import ( TYPE_CHECKING, ) +import torch + if TYPE_CHECKING: from collections.abc import ( Iterator, ) + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + WignerDCalculator, + ) + BLOCK_WIDTHS = (1, 3, 5, 7) FULL_BLOCK_OFFSETS = (0, 1, 4, 9, 16) @@ -32,6 +41,41 @@ PANEL_BLOCK_OFFSETS = (0, 1, 10, 25, 46) PACKED_VALUE_COUNT = PANEL_BLOCK_OFFSETS[-1] + +@lru_cache(maxsize=8) +def _reference_wigner_calculator( + lmax: int, eps: float, dtype: torch.dtype, device: torch.device +) -> WignerDCalculator: + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + WignerDCalculator, + ) + + calculator = WignerDCalculator(lmax, eps=eps, dtype=dtype).to(device) + # The recovery path must not depend on an optional accelerated backend. + calculator.cutile_infer_monomials = False + calculator.triton_infer_l_1_monomials = False + calculator.triton_train_l_1_monomials = False + return calculator.eval() + + +@torch.compiler.disable +def dense_wigner_for_fallback( + quaternion: torch.Tensor, + *, + lmax: int, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """Rebuild dense rotations only when a packed SO2 invocation declines. + + The quaternion retains its geometry autograd history. Only constant + coefficient tables are cached, never edge tensors or their gradients. + """ + calculator = _reference_wigner_calculator( + lmax, eps, quaternion.dtype, quaternion.device + ) + return calculator(quaternion) + + # DeePMD's m-major reduced ordering for lmax=3, mmax=1. COEFF_INDEX_M = (0, 2, 6, 12, 1, 5, 11, 3, 7, 13) REDUCED_DEGREES = (0, 1, 2, 3, 1, 2, 3, 1, 2, 3) diff --git a/source/tests/common/dpmodel/test_dpa4_edge_cache.py b/source/tests/common/dpmodel/test_dpa4_edge_cache.py index 9cdf3cb98f..0cdcae6856 100644 --- a/source/tests/common/dpmodel/test_dpa4_edge_cache.py +++ b/source/tests/common/dpmodel/test_dpa4_edge_cache.py @@ -2,13 +2,34 @@ """Tests for the backend-neutral DPA4 edge-cache acceleration seams.""" import numpy as np +import pytest from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import ( _edge_cache_from_arrays, + _finalize_edge_cache, edge_cache_to_dtype, ) +@pytest.mark.parametrize("floor", [0.0, -1.0]) +def test_degree_normalization_rejects_nonpositive_floor(floor: float) -> None: + with pytest.raises(ValueError, match="deg_norm_floor must be positive"): + _finalize_edge_cache( + n_nodes=1, + src=np.zeros(1, dtype=np.int64), + dst=np.zeros(1, dtype=np.int64), + edge_type_feat=np.ones((1, 1)), + edge_vec=np.ones((1, 3)), + edge_rbf=np.ones((1, 1)), + edge_env=np.ones((1, 1)), + D_full=None, + Dt_full=None, + D_packed=None, + edge_quat=np.ones((1, 4)), + deg_norm_floor=floor, + ) + + def test_fused_builders_replace_reference_and_initialize_step_cache() -> None: calls: dict[str, int] = {"radial": 0, "wigner": 0} keep_seen: list[np.ndarray] = [] diff --git a/source/tests/pt/model/sezm_cute/test_runtime_policy.py b/source/tests/pt/model/sezm_cute/test_runtime_policy.py index 5276b842e8..aca4b0938d 100644 --- a/source/tests/pt/model/sezm_cute/test_runtime_policy.py +++ b/source/tests/pt/model/sezm_cute/test_runtime_policy.py @@ -11,9 +11,19 @@ mock, ) +import pytest + from deepmd.pt_expt.kernels.cute.sezm import runtime_policy as policy +@pytest.mark.parametrize("value", ["2", "unsupported", "tru"]) +@pytest.mark.parametrize("name", ["DP_CUTE_STRICT", "DP_CUTE_SO2_THIN_WRAPPER"]) +def test_invalid_boolean_override_raises(name: str, value: str) -> None: + with mock.patch.dict(os.environ, {name: value}, clear=True): + with pytest.raises(ValueError, match=name): + policy._env_override(name) + + def test_cute_master_gate_controls_sezm_path() -> None: cases = ( ({}, False), diff --git a/source/tests/pt/model/sezm_cute/test_so2.py b/source/tests/pt/model/sezm_cute/test_so2.py index 98fdd846eb..69d898eb77 100644 --- a/source/tests/pt/model/sezm_cute/test_so2.py +++ b/source/tests/pt/model/sezm_cute/test_so2.py @@ -563,6 +563,148 @@ def test_runtime_contract_rejects_tf32_matmul_policy( assert not _SO2.is_neo_so2_runtime_eligible(block, **kwargs) +@pytest.mark.parametrize("change", ["tf32", "trainable", "dtype", "training"]) +def test_prepared_dispatch_rechecks_mutable_model_state( + monkeypatch: pytest.MonkeyPatch, change: str +) -> None: + monkeypatch.setenv("DP_CUTE_INFER", "1") + monkeypatch.setattr(runtime_policy, "uses_strict_fp32_matmul", lambda: True) + block = _neo_like_block().eval() + config = NeoSO2RuntimeConfig() + handle = _SO2.register_cute_so2_block(block, config) + block._deepmd_cute_so2_state = _SO2._RegisteredSO2State(0, handle, config) + try: + assert _SO2._has_frozen_fp32_contract(block) + if change == "tf32": + monkeypatch.setattr( + runtime_policy, "uses_strict_fp32_matmul", lambda: False + ) + elif change == "trainable": + block.runtime_weight.requires_grad_(True) + elif change == "dtype": + block.double() + else: + block.train() + + # Rejection must happen before inspecting edge data or entering a kernel. + assert _SO2._maybe_run_prepared_cute_so2(block, None, None, None) is None + with pytest.raises(RuntimeError, match="model or precision state changed"): + _SO2._build_runner(handle, *([None] * 11)) + finally: + _SO2.invalidate_cute_so2_state(block) + + +@pytest.mark.parametrize("backend", ["pt", "pt_expt"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_packed_dispatch_decline_recovers_outputs_and_gradients( + monkeypatch: pytest.MonkeyPatch, backend: str, dtype: torch.dtype +) -> None: + from deepmd.dpmodel.descriptor.dpa4_nn.edge_cache import ( + EdgeCache, + ) + from deepmd.pt.model.descriptor.sezm_nn.block import ( + SeZMInteractionBlock, + ) + from deepmd.pt.model.descriptor.sezm_nn.wignerd import ( + WignerDCalculator, + ) + from deepmd.pt_expt.descriptor.dpa4_nn.block import ( + SeZMInteractionBlock as ExptBlock, + ) + from deepmd.pt_expt.kernels.cute.sezm.so2.wigner_layout import ( + iter_packed_entries, + ) + + monkeypatch.setenv("DP_TRITON_INFER", "0") + monkeypatch.setenv("DP_CUTE_INFER", "0") + generator = torch.Generator(device="cpu").manual_seed(71) + block = ( + SeZMInteractionBlock( + lmax=3, + mmax=1, + channels=4, + mixing_layers=3, + radial_so2_mode="degree_channel", + radial_so2_rank=1, + ffn_activation_function="silu", + dtype=dtype, + seed=11, + trainable=True, + ) + .cpu() + .eval() + ) + with torch.no_grad(): + for parameter in block.parameters(): + parameter.add_( + 0.05 + * torch.randn( + parameter.shape, dtype=dtype, device="cpu", generator=generator + ) + ) + if backend == "pt_expt": + block = ExptBlock.deserialize(block.serialize()).cpu().eval() + calculator = WignerDCalculator(3, dtype=dtype).cpu() + x_base = torch.randn(3, 16, 1, 4, dtype=dtype, device="cpu", generator=generator) + radial_base = torch.randn(4, 4, 4, dtype=dtype, device="cpu", generator=generator) + quat_base = torch.randn(4, 4, dtype=dtype, device="cpu", generator=generator) + src = torch.tensor([1, 2, 0, 1], dtype=torch.int64, device="cpu") + dst = torch.tensor([0, 0, 1, 2], dtype=torch.int64, device="cpu") + parameter = block.so2_conv.so2_linears[0].weight_m0.requires_grad_(True) + + def run(*, packed: bool): + x, radial, quaternion = ( + value.clone().requires_grad_(True) + for value in (x_base, radial_base, quat_base) + ) + d_full, dt_full = calculator(quaternion) + panel = torch.stack( + [ + d_full[:, entry.full_row, entry.full_col] + for entry in iter_packed_entries() + ], + dim=1, + ) + cache_type = EdgeFeatureCache if backend == "pt" else EdgeCache + cache = cache_type( + src=src, + dst=dst, + edge_type_feat=radial[:, 0], + edge_vec=quaternion[:, 1:], + edge_rbf=radial[:, 0], + edge_env=torch.ones(4, 1, dtype=dtype, device="cpu"), + deg=torch.tensor([2, 1, 1], dtype=dtype, device="cpu"), + inv_sqrt_deg=torch.ones(3, 1, 1, dtype=dtype, device="cpu"), + edge_quat=quaternion, + D_full=None if packed else d_full, + Dt_full=None if packed else dt_full, + D_packed=panel if packed else None, + destinations_sorted=True, + ) + output = block._run_so2_unit_impl(x, cache, radial) + gradients = torch.autograd.grad( + output.square().sum(), (x, radial, quaternion, parameter) + ) + if packed: + # Recovery is local to the declined block, not a mutation of shared state. + assert cache.D_full is None and cache.D_packed is panel + return output, gradients + + expected, expected_grads = run(packed=False) + declined_calls = [] + + def decline(*args, **kwargs): + declined_calls.append(True) + + monkeypatch.setattr(_SO2, "maybe_run_cute_so2", decline) + actual, actual_grads = run(packed=True) + assert declined_calls == [True] + tol = 5e-5 if dtype == torch.float32 else 1e-10 + torch.testing.assert_close(actual, expected, atol=tol, rtol=tol) + for actual_grad, expected_grad in zip(actual_grads, expected_grads, strict=True): + torch.testing.assert_close(actual_grad, expected_grad, atol=tol, rtol=tol) + + def test_fake_native_gradient_layout_matches_runtime_contract() -> None: x = torch.empty(5, 16, 1, 32, device="cpu") diff --git a/source/tests/pt/model/test_mlp.py b/source/tests/pt/model/test_mlp.py index 55492b5f7f..def412d0be 100644 --- a/source/tests/pt/model/test_mlp.py +++ b/source/tests/pt/model/test_mlp.py @@ -156,7 +156,11 @@ def test_thin_so2_eval_linear_is_scoped_to_marked_model(self) -> None: expected = mlp_module.F.linear(value, layer.matrix.t(), layer.bias) torch.testing.assert_close(actual, expected) - def test_thin_so2_linear_defaults_on_only_for_sm80(self) -> None: + def test_thin_so2_linear_matches_shared_device_policy(self) -> None: + from deepmd.pt_expt.kernels.cute.sezm import ( + runtime_policy, + ) + with ( mock.patch.dict( os.environ, @@ -164,9 +168,20 @@ def test_thin_so2_linear_defaults_on_only_for_sm80(self) -> None: clear=True, ), mock.patch.object(torch.cuda, "is_available", return_value=True), - mock.patch.object(torch.cuda, "get_device_capability", return_value=(8, 0)), ): - self.assertTrue(mlp_module._use_so2_compile_visible_linear()) + for capability in ((8, 0), (8, 6), (8, 9), (9, 0), (12, 0)): + with ( + self.subTest(capability=capability), + mock.patch.object( + torch.cuda, "get_device_capability", return_value=capability + ) as get_capability, + ): + device = torch.device("cuda", 1) + self.assertEqual( + mlp_module._use_so2_compile_visible_linear(device), + runtime_policy.is_so2_thin_wrapper_enabled(capability), + ) + get_capability.assert_called_once_with(device) with ( mock.patch.dict( diff --git a/source/tests/pt/model/test_sezm_parallel.py b/source/tests/pt/model/test_sezm_parallel.py index be3b093fc1..1c2bb68238 100644 --- a/source/tests/pt/model/test_sezm_parallel.py +++ b/source/tests/pt/model/test_sezm_parallel.py @@ -17,6 +17,9 @@ import ctypes import unittest +from unittest import ( + mock, +) import numpy as np import torch @@ -44,6 +47,29 @@ _SENDLIST_KEEPALIVE: list[np.ndarray] = [] +class TestSeZMDenseCommIsolation(unittest.TestCase): + def test_local_nlist_features_are_not_passed_to_ghost_exchange(self) -> None: + model = _build_model(torch.device("cpu")) + descriptor = model.atomic_model.descriptor + coord = torch.tensor( + [[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]], + dtype=torch.float64, + device="cpu", + ) + atype = torch.tensor([[0, 1, 0]], dtype=torch.int64, device="cpu") + # Two local atoms and a mapped ghost atom. + nlist = torch.tensor([[[1, 2], [0, 2]]], dtype=torch.int64, device="cpu") + mapping = torch.tensor([[0, 1, 0]], dtype=torch.int64, device="cpu") + expected = descriptor(coord, atype, nlist, mapping=mapping)[0] + with mock.patch.object( + sezm_block, + "exchange_ghost_features", + side_effect=AssertionError("local-only features reached ghost exchange"), + ): + actual = descriptor(coord, atype, nlist, mapping=mapping, comm_dict={})[0] + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + def _tiny_parallel_model_params(**overrides) -> dict: """Minimal fp64 SeZM config exercising message passing and FiLM/GIE seeds.""" descriptor = {