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..7e227c332a --- /dev/null +++ b/deepmd/kernels/cute/neo/k1.py @@ -0,0 +1,2224 @@ +# 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 +_REGISTRY_LOCK = threading.Lock() +_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 + + def remove_collected_entry(block_ref: weakref.ReferenceType[Any]) -> None: + with _REGISTRY_LOCK: + entry = _REGISTRY.get(handle) + if entry is not None and entry._block_ref is block_ref: + _REGISTRY.pop(handle, None) + + with _REGISTRY_LOCK: + handle = _NEXT_HANDLE + _NEXT_HANDLE += 1 + _REGISTRY[handle] = _RegistryEntry( + block=block, + config=config, + on_collect=remove_collected_entry, + ) + return handle + + +def invalidate_cute_k1_state(block: Any) -> None: + """Release a block's registered CuTe state after its modules change.""" + state = getattr(block, "_deepmd_cute_k1_state", None) + if isinstance(state, _RegisteredK1State): + 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"): + 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): + 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. + 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) + + 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) + 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. + """ + with _REGISTRY_LOCK: + 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.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 + 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)) + 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..45a996e816 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_envelope_gated_softmax.py @@ -0,0 +1,238 @@ +# 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 +import cutlass.utils +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..27d356caa0 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_gate_linear_residual_backward_fused.py @@ -0,0 +1,169 @@ +# 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 +import cutlass.utils +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..488806d314 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_a_radial_forward.py @@ -0,0 +1,392 @@ +# 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(), + hidden_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] + + 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)}" + ) + 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, + 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..43e3f7ef14 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout.py @@ -0,0 +1,693 @@ +# 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 + 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, + 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, + 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_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_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..f33059e9be --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_backward_layout_runner.py @@ -0,0 +1,376 @@ +# 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, + 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("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), + ("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, + 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..ee489af8bc --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_phase_c_onepass.py @@ -0,0 +1,552 @@ +# 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 +import cutlass.utils +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 + +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 +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..9a1de48779 --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_kernels/cute_neo_so2_gate_combined_fwd.py @@ -0,0 +1,780 @@ +# 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 +DEFAULT_STREAM = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT) + + +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 = DEFAULT_STREAM, + ): + 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..ac968839ba --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_radial_phase_a_node.py @@ -0,0 +1,400 @@ +# 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, +) +from .k1_wigner_layout import ( + PACKED_VALUE_COUNT, +) + +if TYPE_CHECKING: + from torch import ( + Tensor, + ) + + +DEGREE_COUNT = 16 +REDUCED_COUNT = 10 +HIDDEN = 64 +FOCUS_COUNT = 2 +FOCUS_HIDDEN = 32 +PACKED_WIGNER_VALUES = PACKED_VALUE_COUNT +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. + ``grad_out_focus`` is repacked in place as projection workspace and must + not be reused after this function returns. + """ + 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..af793d5a5d --- /dev/null +++ b/deepmd/kernels/cute/neo/k1_runner.py @@ -0,0 +1,904 @@ +# 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 + ), + 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) + 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, + ) + 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, + ) + + 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..b02e66fa52 --- /dev/null +++ b/deepmd/kernels/cute/neo/k4_wignerd.py @@ -0,0 +1,1009 @@ +# 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, +) + +import threading +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, ...]] = {} +_TABLE_CACHE_LOCK = threading.Lock() + + +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: + 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 + + +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..2bbc087ac8 --- /dev/null +++ b/deepmd/kernels/cute/neo/message_grid_gaunt_sm90.py @@ -0,0 +1,324 @@ +# 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) + 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() + + 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.""" + 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 + + +__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..0d51e6c1fa --- /dev/null +++ b/deepmd/kernels/cute/neo/readout_l0.py @@ -0,0 +1,783 @@ +# 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, strict=True) + ) + 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, + strict=True, + ): + _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, 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, strict=True) + ) + 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..5ba25e3f41 --- /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 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..e0f3b481a7 --- /dev/null +++ b/deepmd/kernels/cute/neo/sm90_k1/phase_a.py @@ -0,0 +1,555 @@ +# 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} +DEFAULT_STREAM = cuda.CUstream(cuda.CUstream_flags.CU_STREAM_DEFAULT) + +__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 = DEFAULT_STREAM, + ): + 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..f1f0430357 --- /dev/null +++ b/deepmd/kernels/cute/neo/sm90_k1/phase_c_attention_backward.py @@ -0,0 +1,842 @@ +# 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, + 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'") + + 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 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( + 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..cb7cb017f7 --- /dev/null +++ b/deepmd/kernels/cute/neo/sm90_k1/radial.py @@ -0,0 +1,310 @@ +# 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. + """ + 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( + "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..1f8f37339e --- /dev/null +++ b/deepmd/kernels/cute/neo/sm90_k1/runner.py @@ -0,0 +1,660 @@ +# 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), + 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..eccaae84b2 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 @@ -351,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 @@ -361,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) ---------------------------------------------------------- @@ -526,6 +528,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 +711,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 +865,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 +903,41 @@ 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 + 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. + _clear_shared_sezm_compile_cache() + # ========================================================================= # Forward Methods # ========================================================================= @@ -1489,6 +1661,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 +1746,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 +1927,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 +1954,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 +2152,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 +3008,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 +3037,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 +3054,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 +3079,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 ------- @@ -2843,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``. @@ -2860,12 +3111,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..a77d2b35d7 100644 --- a/doc/install/easy-install.md +++ b/doc/install/easy-install.md @@ -189,6 +189,21 @@ pip install deepmd-kit :::::: +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]" +``` + +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..ec79b7ca64 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; python_version >= "3.11" and 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..c03da7bcc9 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_gie.py @@ -0,0 +1,227 @@ +# 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 unittest +from types import ( + SimpleNamespace, +) + +import torch + +from deepmd.kernels.cute.neo import gie as gie_module + + +def _load_gie_module(): + return gie_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..662cb29ffa --- /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 + + 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) + k1._REGISTRY.clear() + k1._PACKED_RUNNER_CACHE.clear() + torch._dynamo.reset() + 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..de39840bd7 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_gate_structural.py @@ -0,0 +1,98 @@ +# 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 unittest + +import torch + +from deepmd.kernels.cute.neo import ( + k1_gate_structural, +) + + +class TestStructuralGateHelpers(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.helper = k1_gate_structural + + 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..1e348b9018 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_inplace_residual.py @@ -0,0 +1,366 @@ +# 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 unittest + +import torch +from torch.utils._python_dispatch import ( + TorchDispatchMode, +) + +from deepmd.kernels.cute.neo import ( + k1_so2linear, +) + + +def _load_candidate(): + return k1_so2linear + + +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..6d382e030d --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_radial_projection.py @@ -0,0 +1,204 @@ +# 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 unittest +from types import ( + SimpleNamespace, +) +from unittest import ( + mock, +) + +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 + +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..92d153b090 --- /dev/null +++ b/source/tests/pt/model/test_descriptor_sezm_cute_k1_structural_gate_vec4_sm80.py @@ -0,0 +1,549 @@ +# 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 +import math +import random +import struct +import unittest +from pathlib import ( + Path, +) + +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" +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 _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): + 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(): + return importlib.import_module("deepmd.kernels.cute.neo.k1_gate_structural") + + +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) + ] + 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(actual_grad_y, expected_grad_y) + self.assertEqual(actual_grad_logits, expected_grad_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..7f5cc5c7ad 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,105 @@ 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 + model._deepmd_cute_k1_state = False + + 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.assertFalse(hasattr(model, "_deepmd_cute_k1_state")) + 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 +970,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."""