From c0c65ee1f1b5b2088c3eda0b7ca3d091eb919109 Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Tue, 7 Jul 2026 09:45:47 -0500 Subject: [PATCH 01/13] Add framework-agnostic kernel APIs with JAX support --- docs/fe-oss-apis/rmsnorm_rht_amax.md | 28 ++ pyproject.toml | 5 + python/cudnn/__init__.py | 7 + python/cudnn/_jax/__init__.py | 36 ++ python/cudnn/_jax/api_base.py | 277 +++++++++++++ python/cudnn/_jax/datatypes.py | 64 +++ python/cudnn/_jax/layout.py | 112 +++++ python/cudnn/_op_kernel.py | 47 +++ python/cudnn/_operation_api.py | 54 +++ python/cudnn/_tensor_desc.py | 138 +++++++ python/cudnn/api_base.py | 134 ++++-- .../grouped_gemm/grouped_gemm_quant/api.py | 1 + .../grouped_gemm/grouped_gemm_srelu/api.py | 1 + .../grouped_gemm/grouped_gemm_wgrad/api.py | 1 + python/cudnn/jax/__init__.py | 17 + python/cudnn/rmsnorm_rht_amax/__init__.py | 30 +- python/cudnn/rmsnorm_rht_amax/api.py | 273 ++++++------ python/cudnn/rmsnorm_rht_amax/jax.py | 78 ++++ python/cudnn/rmsnorm_rht_amax/kernel.py | 232 ++++++++++- test/python/fe_api/test_jax_api_base.py | 390 ++++++++++++++++++ test/python/fe_api/test_jax_datatypes.py | 84 ++++ test/python/fe_api/test_jax_layout.py | 89 ++++ .../fe_api/test_jax_rmsnorm_rht_amax.py | 126 ++++++ test/python/fe_api/test_rmsnorm_rht_amax.py | 53 ++- .../fe_api/test_rmsnorm_rht_amax_imports.py | 298 +++++++++++++ .../fe_api/test_rmsnorm_rht_amax_kernel.py | 280 +++++++++++++ test/python/fe_api/test_shared_tensor_desc.py | 159 +++++++ test/python/fe_api/test_tensor_desc.py | 153 +++++++ 28 files changed, 2977 insertions(+), 190 deletions(-) create mode 100644 python/cudnn/_jax/__init__.py create mode 100644 python/cudnn/_jax/api_base.py create mode 100644 python/cudnn/_jax/datatypes.py create mode 100644 python/cudnn/_jax/layout.py create mode 100644 python/cudnn/_op_kernel.py create mode 100644 python/cudnn/_operation_api.py create mode 100644 python/cudnn/_tensor_desc.py create mode 100644 python/cudnn/jax/__init__.py create mode 100644 python/cudnn/rmsnorm_rht_amax/jax.py create mode 100644 test/python/fe_api/test_jax_api_base.py create mode 100644 test/python/fe_api/test_jax_datatypes.py create mode 100644 test/python/fe_api/test_jax_layout.py create mode 100644 test/python/fe_api/test_jax_rmsnorm_rht_amax.py create mode 100644 test/python/fe_api/test_rmsnorm_rht_amax_imports.py create mode 100644 test/python/fe_api/test_rmsnorm_rht_amax_kernel.py create mode 100644 test/python/fe_api/test_shared_tensor_desc.py create mode 100644 test/python/fe_api/test_tensor_desc.py diff --git a/docs/fe-oss-apis/rmsnorm_rht_amax.md b/docs/fe-oss-apis/rmsnorm_rht_amax.md index ba875fef8..720331182 100644 --- a/docs/fe-oss-apis/rmsnorm_rht_amax.md +++ b/docs/fe-oss-apis/rmsnorm_rht_amax.md @@ -9,6 +9,7 @@ This frontend integration exposes the kernel as a standard FE-OSS Python API with: - a class API (`RmsNormRhtAmaxSm100`) - a wrapper API (`rmsnorm_rht_amax_wrapper_sm100`) +- an explicit JAX API under `cudnn.jax` - grouped-gemm-style regression coverage for compile/execute, wrapper use, and cache reuse ## Shapes @@ -97,6 +98,32 @@ op.execute( ) ``` +### JAX API + +The CUDA 13 JAX packages require Python 3.11 or newer. Install the JAX integration explicitly: + +```bash +pip install 'nvidia-cudnn-frontend[jax]' +``` + +```python +from cudnn.jax import rmsnorm_rht_amax_sm100 + +o, amax = rmsnorm_rht_amax_sm100( + x, + w, + eps=1e-5, + num_threads=128, + rows_per_cta=2, +) +``` + +The function API is JIT-compiled and treats `eps`, `num_threads`, and +`rows_per_cta` as static compilation options. Use `RmsNormRhtAmaxSm100` +directly when the caller needs to control JAX transformations explicitly. + +JAX buffers remain row-major. Framework-specific mode metadata maps public JAX axes into the canonical kernel axis order without materializing a transpose. + ## Parameters ### Input and output tensors @@ -149,3 +176,4 @@ Tuple unpacking order is `(o_tensor, amax_tensor)`. Focused correctness and cache coverage live in: - `test/python/fe_api/test_rmsnorm_rht_amax.py` +- `test/python/fe_api/test_jax_rmsnorm_rht_amax.py` diff --git a/pyproject.toml b/pyproject.toml index ca8243c38..2c5a14d30 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -76,3 +76,8 @@ cutedsl = [ "apache-tvm-ffi", "torch-c-dlpack-ext", ] +jax = [ + "nvidia-cutlass-dsl[cu13]==4.5.0", + "cuda-python", + "jax[cuda13]>=0.9.1", +] diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index af48c763f..e77cef0e5 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -51,6 +51,8 @@ def is_windows(): if hasattr(_pybind_module, _optional_symbol): globals()[_optional_symbol] = getattr(_pybind_module, _optional_symbol) +from ._op_kernel import OpKernel +from ._tensor_desc import TensorDesc from .datatypes import _library_type, _is_torch_tensor __version__ = "1.26.0" @@ -334,6 +336,11 @@ def __getattr__(name: str) -> Any: globals()["ops"] = _ops return _ops + if name == "jax": + _jax = importlib.import_module(".jax", __name__) + globals()["jax"] = _jax + return _jax + if name == "experimental": from . import experimental as _experimental diff --git a/python/cudnn/_jax/__init__.py b/python/cudnn/_jax/__init__.py new file mode 100644 index 000000000..2e3af9dfa --- /dev/null +++ b/python/cudnn/_jax/__init__.py @@ -0,0 +1,36 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional-dependency boundary shared by all JAX operation adapters.""" + +from importlib import import_module + +_INSTALL_HINT = "pip install 'nvidia-cudnn-frontend[jax]'" +_OPTIONAL_DEPENDENCY_PREFIXES = ("cuda", "cutlass", "jax") + + +def _require_dependency(module_name: str): + try: + return import_module(module_name) + except ModuleNotFoundError as error: + missing_name = error.name + if missing_name is None or not any(missing_name == prefix or missing_name.startswith(f"{prefix}.") for prefix in _OPTIONAL_DEPENDENCY_PREFIXES): + raise + raise ImportError(f"cuDNN JAX APIs require {module_name!r}. " f"Install the optional JAX dependencies with `{_INSTALL_HINT}`.") from error + + +jax = _require_dependency("jax") +cutlass_jax = _require_dependency("cutlass.jax") + +if not cutlass_jax.is_available(): + minimum_version = ".".join(str(part) for part in cutlass_jax.CUTE_DSL_MIN_SUPPORTED_JAX_VERSION) + installed_version = getattr(jax, "__version__", "unknown") + raise ImportError( + f"CUTLASS JAX support is unavailable with JAX {installed_version}; " + f"the minimum supported JAX version is {minimum_version}. " + f"Install the optional JAX dependencies with `{_INSTALL_HINT}`." + ) + +from .api_base import JaxApiBase, JaxTensorDesc + +__all__ = ["JaxApiBase", "JaxTensorDesc"] diff --git a/python/cudnn/_jax/api_base.py b/python/cudnn/_jax/api_base.py new file mode 100644 index 000000000..18447f3c3 --- /dev/null +++ b/python/cudnn/_jax/api_base.py @@ -0,0 +1,277 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Base class, tensor metadata, and CuTe binding for optional JAX adapters.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable +from typing import Any + +from .. import data_type +from .._op_kernel import OpKernel +from .._tensor_desc import TensorDesc +from .layout import compact_stride, normalize_mode, to_canonical_axes, to_cutlass_layout, to_public_axes + + +class JaxTensorDesc(TensorDesc[Any]): + """Framework-neutral tensor metadata backed by a JAX dtype.""" + + @property + def cudnn_dtype(self) -> data_type: + from .datatypes import jax_to_cudnn_dtype + + return jax_to_cudnn_dtype(self.dtype) + + +def _require_array_metadata(value: Any, name: str) -> tuple[Any, ...]: + if not hasattr(value, "shape") or not hasattr(value, "dtype"): + raise TypeError(f"{name} must have shape and dtype metadata") + return tuple(value.shape) + + +class JaxApiBase(ABC): + """Common tensor metadata, validation, and kernel binding for JAX adapters.""" + + kernel: OpKernel + + @staticmethod + def _to_tensor_desc( + value: Any, + name: str, + *, + mode: tuple[int, ...] | None = None, + ) -> JaxTensorDesc: + """Describe a public JAX array in canonical kernel-axis order. + + ``mode[kernel_axis]`` selects the corresponding public array axis. + Public JAX arrays are modeled as compact row-major buffers. + """ + + public_shape = _require_array_metadata(value, name) + rank = len(public_shape) + mode = normalize_mode(rank, mode) + public_stride_order = tuple(reversed(range(rank))) + public_stride = compact_stride(public_shape, public_stride_order) + canonical_axis_by_public_axis = to_public_axes(tuple(range(rank)), mode) + + return JaxTensorDesc( + dtype=value.dtype, + shape=to_canonical_axes(public_shape, mode), + stride=to_canonical_axes(public_stride, mode), + stride_order=tuple(canonical_axis_by_public_axis[axis] for axis in public_stride_order), + name=name, + ) + + @staticmethod + def _check_tensor_signature( + value: Any, + expected: TensorDesc[Any], + *, + mode: tuple[int, ...] | None = None, + ) -> None: + """Validate a public JAX value against a canonical descriptor.""" + + name = expected.name or "value" + public_shape = _require_array_metadata(value, name) + mode = normalize_mode(expected.ndim, mode) + if len(public_shape) != expected.ndim: + raise ValueError(f"{name} tensor shape mismatch: expected {expected.shape}, got public shape {public_shape}") + actual_shape = to_canonical_axes(public_shape, mode) + if actual_shape != expected.shape: + raise ValueError(f"{name} tensor shape mismatch: expected {expected.shape}, got {actual_shape}") + + from .datatypes import jax_to_cudnn_dtype + + actual_dtype = jax_to_cudnn_dtype(value.dtype) + if actual_dtype != expected.cudnn_dtype: + raise ValueError(f"{name} tensor dtype mismatch: expected {expected.cudnn_dtype}, got {actual_dtype}") + + @staticmethod + def _to_tensor_spec( + desc: TensorDesc[Any], + *, + mode: tuple[int, ...] | None = None, + divisibility: tuple[int | None, ...] | None = None, + ) -> Any: + """Build a CUTLASS TensorSpec indexed by public JAX array axes. + + ``desc`` and ``divisibility`` use canonical kernel axes. TensorSpec + layout and divisibility use public axes, while TensorSpec ``mode`` + records the canonical-to-public binding. + """ + + mode = normalize_mode(desc.ndim, mode) + public_layout = to_cutlass_layout( + desc.shape, + desc.stride, + desc.stride_order, + mode=mode, + name=desc.name or "tensor", + ) + if divisibility is None: + public_divisibility = None + else: + divisibility = tuple(divisibility) + if len(divisibility) != desc.ndim: + raise ValueError(f"divisibility rank mismatch: expected {desc.ndim}, got {len(divisibility)}") + public_divisibility = to_public_axes(divisibility, mode) + + from cutlass.jax import TensorSpec + + return TensorSpec( + layout=public_layout, + mode=mode, + divisibility=public_divisibility, + ) + + @staticmethod + def _materialize_tensor_desc( + desc: TensorDesc[Any], + *, + mode: tuple[int, ...] | None = None, + ) -> Any: + """Declare a canonical descriptor as a public JAX output buffer.""" + + if not isinstance(desc, TensorDesc): + raise TypeError(f"desc must be a TensorDesc, got {type(desc).__name__}") + mode = normalize_mode(desc.ndim, mode) + + from jax import ShapeDtypeStruct + + from .datatypes import cudnn_to_jax_dtype + + return ShapeDtypeStruct( + to_public_axes(desc.shape, mode), + cudnn_to_jax_dtype(desc.cudnn_dtype), + ) + + def _call_kernel( + self, + inputs: tuple[Any, ...], + *, + input_spec: tuple[Any | None, ...] | None = None, + output_spec: tuple[Any | None, ...] | None = None, + workspace_spec: tuple[Any | None, ...] | None = None, + allow_cuda_graph: bool = True, + compile_options: Any = None, + use_static_tensors: bool = True, + ) -> tuple[Any, ...]: + """Bind the owned kernel to JAX and return its public outputs. + + CUTLASS JAX supplies the stream first; the launcher adapts that call to + the kernel's ``inputs, outputs, workspaces, stream`` order. Workspaces + are declared as custom-call results so XLA owns their lifetime, then + omitted from this method's return value. Descriptors with a non-``None`` + ``init_value`` are materialized as JAX inputs and aliased to their + corresponding custom-call results. Inputs are a flat tuple of arrays, + and all explicit specs are native CUTLASS ``TensorSpec`` values. + """ + + import cutlass.jax as cutlass_jax + import jax.numpy as jnp + + inputs = tuple(inputs) + for index, value in enumerate(inputs): + _require_array_metadata(value, f"input #{index}") + outputs = tuple(self.kernel.infer_output()) + workspaces = tuple(self.kernel.infer_workspace()) + if not outputs: + raise ValueError("A JAX operation must infer at least one output") + + for label, descs in (("output", outputs), ("workspace", workspaces)): + for desc in descs: + if not isinstance(desc, TensorDesc): + raise TypeError(f"{label} inference must return TensorDesc values, got {type(desc).__name__}") + + if input_spec is None: + input_specs = (None,) * len(inputs) + else: + input_specs = tuple(input_spec) + if len(input_specs) != len(inputs): + raise ValueError(f"Expected {len(inputs)} input specs, got {len(input_specs)}") + if any(spec is not None and not isinstance(spec, cutlass_jax.TensorSpec) for spec in input_specs): + raise TypeError("input_spec must contain only TensorSpec or None values") + + def resolve_specs( + descs: tuple[TensorDesc[Any], ...], + supplied: tuple[Any | None, ...] | None, + label: str, + ) -> tuple[Any, ...]: + if supplied is None: + supplied = (None,) * len(descs) + else: + supplied = tuple(supplied) + if len(supplied) != len(descs): + raise ValueError(f"Expected {len(descs)} {label} specs, got {len(supplied)}") + if any(spec is not None and not isinstance(spec, cutlass_jax.TensorSpec) for spec in supplied): + raise TypeError(f"{label}_spec must contain only TensorSpec or None values") + return tuple(self._to_tensor_spec(desc) if spec is None else spec for desc, spec in zip(descs, supplied)) + + output_specs = resolve_specs(outputs, output_spec, "output") + workspace_specs = resolve_specs(workspaces, workspace_spec, "workspace") + buffers = outputs + workspaces + buffer_specs = output_specs + workspace_specs + buffer_metadata = tuple(self._materialize_tensor_desc(desc, mode=getattr(spec, "mode", None)) for desc, spec in zip(buffers, buffer_specs)) + + initialized = tuple(index for index, desc in enumerate(buffers) if desc.init_value is not None) + uninitialized = tuple(index for index, desc in enumerate(buffers) if desc.init_value is None) + seed_inputs = tuple( + jnp.full( + buffer_metadata[index].shape, + buffers[index].init_value, + dtype=buffer_metadata[index].dtype, + ) + for index in initialized + ) + aliases = {len(inputs) + seed_index: result_index for seed_index, result_index in enumerate(initialized)} + call_input_specs = input_specs + tuple(buffer_specs[index] for index in initialized) + + input_count = len(inputs) + seed_count = len(seed_inputs) + + def launcher(stream, *args): + kernel_inputs = args[:input_count] + seeded_buffers = args[input_count : input_count + seed_count] + allocated_buffers = args[input_count + seed_count :] + if len(allocated_buffers) != len(uninitialized): + raise RuntimeError(f"Kernel received {len(allocated_buffers)} allocated buffers; expected {len(uninitialized)}") + + ordered_buffers: list[Any] = [None] * len(buffers) + for index, value in zip(initialized, seeded_buffers): + ordered_buffers[index] = value + for index, value in zip(uninitialized, allocated_buffers): + ordered_buffers[index] = value + self.kernel(*kernel_inputs, *ordered_buffers, stream) + + call = cutlass_jax.cutlass_call( + launcher, + output_shape_dtype=buffer_metadata, + input_spec=call_input_specs, + output_spec=buffer_specs, + input_output_aliases=aliases, + allow_cuda_graph=allow_cuda_graph, + compile_options=compile_options, + use_static_tensors=use_static_tensors, + ) + results = call(*inputs, *seed_inputs) + if not isinstance(results, (tuple, list)): + results = (results,) + return tuple(results[: len(outputs)]) + + @abstractmethod + def check_support(self) -> bool: + """Validate the operation signature and static configuration.""" + + @abstractmethod + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """Trace or execute the operation with JAX array arguments.""" + + def get_jax_callable(self) -> Callable[..., Any]: + """Return this callable without imposing a JIT policy.""" + + return self + + +__all__ = ["JaxApiBase", "JaxTensorDesc"] diff --git a/python/cudnn/_jax/datatypes.py b/python/cudnn/_jax/datatypes.py new file mode 100644 index 000000000..eac294c5e --- /dev/null +++ b/python/cudnn/_jax/datatypes.py @@ -0,0 +1,64 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Conversions between JAX and canonical cuDNN data types.""" + +from __future__ import annotations + +from typing import Any + +import jax.numpy as jnp + +from .. import data_type + +_DTYPE_NAMES = ( + ("bool_", "BOOLEAN"), + ("float16", "HALF"), + ("bfloat16", "BFLOAT16"), + ("float32", "FLOAT"), + ("float64", "DOUBLE"), + ("int8", "INT8"), + ("int32", "INT32"), + ("int64", "INT64"), + ("uint8", "UINT8"), + ("float8_e4m3fn", "FP8_E4M3"), + ("float8_e5m2", "FP8_E5M2"), + ("float8_e8m0fnu", "FP8_E8M0"), + ("float4_e2m1fn", "FP4_E2M1"), + ("int4", "INT4"), +) + + +def _make_mappings(): + jax_to_cudnn = {} + cudnn_to_jax = {} + for jax_name, cudnn_name in _DTYPE_NAMES: + jax_scalar_type = getattr(jnp, jax_name, None) + cudnn_dtype = getattr(data_type, cudnn_name, None) + if jax_scalar_type is None or cudnn_dtype is None: + continue + jax_dtype = jnp.dtype(jax_scalar_type) + jax_to_cudnn[jax_dtype] = cudnn_dtype + cudnn_to_jax[cudnn_dtype] = jax_dtype + return jax_to_cudnn, cudnn_to_jax + + +_JAX_TO_CUDNN_DATA_TYPE, _CUDNN_TO_JAX_DATA_TYPE = _make_mappings() + + +def jax_to_cudnn_dtype(dtype: Any) -> data_type: + """Return the canonical cuDNN type for a JAX dtype-like value.""" + + return _JAX_TO_CUDNN_DATA_TYPE.get(jnp.dtype(dtype), data_type.NOT_SET) + + +def cudnn_to_jax_dtype(dtype: data_type) -> Any: + """Return the JAX dtype corresponding to a canonical cuDNN type.""" + + try: + return _CUDNN_TO_JAX_DATA_TYPE[dtype] + except KeyError as error: + raise ValueError(f"Unsupported JAX data type {dtype}") from error + + +__all__ = ["cudnn_to_jax_dtype", "jax_to_cudnn_dtype"] diff --git a/python/cudnn/_jax/layout.py b/python/cudnn/_jax/layout.py new file mode 100644 index 000000000..f2be19940 --- /dev/null +++ b/python/cudnn/_jax/layout.py @@ -0,0 +1,112 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free axis and layout conversion for CUTLASS JAX calls.""" + +from __future__ import annotations + +from operator import index +from typing import TypeVar + +AxisValueT = TypeVar("AxisValueT") + + +def normalize_mode(rank: int, mode: tuple[int, ...] | None = None) -> tuple[int, ...]: + """Return and validate a kernel-axis to public-axis mapping. + + ``mode[kernel_axis]`` is the corresponding axis in the public JAX array. + Omitting ``mode`` keeps the public and kernel axis orders identical. + """ + + if isinstance(rank, bool): + raise TypeError(f"rank must be an integer, got {rank!r}") + try: + rank = index(rank) + except TypeError as error: + raise TypeError(f"rank must be an integer, got {rank!r}") from error + if rank < 0: + raise ValueError(f"rank must be non-negative, got {rank}") + + if mode is None: + return tuple(range(rank)) + + normalized = [] + for axis in mode: + if isinstance(axis, bool): + raise TypeError(f"mode entries must be integers, got {axis!r}") + try: + normalized.append(index(axis)) + except TypeError as error: + raise TypeError(f"mode entries must be integers, got {axis!r}") from error + normalized_mode = tuple(normalized) + if tuple(sorted(normalized_mode)) != tuple(range(rank)): + raise ValueError(f"mode must be a permutation of [0, {rank - 1}], got {normalized_mode}") + return normalized_mode + + +def to_canonical_axes( + public_values: tuple[AxisValueT, ...], + mode: tuple[int, ...] | None = None, +) -> tuple[AxisValueT, ...]: + """Reorder public-axis values into canonical kernel-axis order.""" + + public_values = tuple(public_values) + mode = normalize_mode(len(public_values), mode) + return tuple(public_values[public_axis] for public_axis in mode) + + +def to_public_axes( + canonical_values: tuple[AxisValueT, ...], + mode: tuple[int, ...] | None = None, +) -> tuple[AxisValueT, ...]: + """Reorder canonical kernel-axis values into public JAX axis order.""" + + canonical_values = tuple(canonical_values) + mode = normalize_mode(len(canonical_values), mode) + canonical_axis_by_public_axis = [0] * len(mode) + for canonical_axis, public_axis in enumerate(mode): + canonical_axis_by_public_axis[public_axis] = canonical_axis + return tuple(canonical_values[canonical_axis_by_public_axis[public_axis]] for public_axis in range(len(mode))) + + +def compact_stride(shape: tuple[int, ...], stride_order: tuple[int, ...]) -> tuple[int, ...]: + """Return compact strides for dimensions ordered fastest to slowest.""" + + stride = [0] * len(shape) + running = 1 + for dimension in stride_order: + stride[dimension] = running + running *= max(shape[dimension], 1) + return tuple(stride) + + +def to_cutlass_layout( + shape: tuple[int, ...], + stride: tuple[int, ...], + stride_order: tuple[int, ...], + *, + mode: tuple[int, ...] | None = None, + name: str = "tensor", +) -> tuple[int, ...]: + """Return CUTLASS stride ranks indexed by public JAX array axes. + + The descriptor arguments are in canonical kernel-axis order. ``mode`` + maps those axes back to the public array axes consumed by TensorSpec. + """ + + if stride != compact_stride(shape, stride_order): + raise ValueError(f"JAX TensorSpec cannot represent non-compact stride {stride} for {name}") + + layout = [0] * len(stride_order) + for rank, dimension in enumerate(stride_order): + layout[dimension] = rank + return to_public_axes(tuple(layout), mode) + + +__all__ = [ + "compact_stride", + "normalize_mode", + "to_canonical_axes", + "to_cutlass_layout", + "to_public_axes", +] diff --git a/python/cudnn/_op_kernel.py b/python/cudnn/_op_kernel.py new file mode 100644 index 000000000..9e448196c --- /dev/null +++ b/python/cudnn/_op_kernel.py @@ -0,0 +1,47 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral contract for operation kernels.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from . import data_type +from ._tensor_desc import TensorDesc + + +class OpKernel(ABC): + """Host-side contract implemented by framework-agnostic kernels. + + Implementations own logical validation and tensor inference. Framework + adapters remain responsible for turning the inferred descriptors into + framework tensors or abstract output declarations and for compiling and + executing the kernel. Callable implementations receive inputs, outputs, + workspaces, and finally the CUDA stream. Framework adapters may reorder + those arguments to match a framework-specific launch ABI. + """ + + @abstractmethod + def check_support(self) -> bool: + """Validate the logical signature and resolve static kernel state.""" + + @abstractmethod + def infer_output(self) -> tuple[TensorDesc[data_type], ...]: + """Return framework-neutral descriptors for the operation outputs. + + A descriptor's ``init_value`` requests scalar initialization before + launch; ``None`` leaves the output uninitialized. + """ + + def infer_workspace(self) -> tuple[TensorDesc[data_type], ...]: + """Return framework-neutral workspace descriptors, if any. + + Workspace descriptors follow the same ``init_value`` convention as + outputs but are not part of the framework-visible result. + """ + + return () + + +__all__ = ["OpKernel"] diff --git a/python/cudnn/_operation_api.py b/python/cudnn/_operation_api.py new file mode 100644 index 000000000..e2a395bd5 --- /dev/null +++ b/python/cudnn/_operation_api.py @@ -0,0 +1,54 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Reusable lazy exports for framework-specific operation packages.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, MutableMapping, Sequence +from importlib import import_module +from typing import Any + + +def make_operation_api( + module_globals: MutableMapping[str, Any], + *, + exports: Mapping[str, Sequence[str]], + submodules: Sequence[str] = (), +) -> tuple[list[str], Callable[[str], Any], Callable[[], list[str]]]: + """Create ``__all__``, ``__getattr__``, and ``__dir__`` for an operation. + + ``exports`` maps sibling module names to the public symbols provided by + each module. ``submodules`` lists sibling modules that are themselves + available as lazy attributes. Resolved values are cached in the package's + globals, following Python's normal module import behavior. + """ + + package_name = module_globals["__name__"] + export_routes: dict[str, str] = {} + for module_name, names in exports.items(): + for name in names: + if name in export_routes: + raise ValueError(f"Operation export {name!r} is provided by more than one module") + export_routes[name] = module_name + + lazy_submodules = frozenset(submodules) + + def get_attribute(name: str) -> Any: + if name in lazy_submodules: + value = import_module(f".{name}", package_name) + elif name in export_routes: + value = getattr(import_module(f".{export_routes[name]}", package_name), name) + else: + raise AttributeError(f"module {package_name!r} has no attribute {name!r}") + + module_globals[name] = value + return value + + def list_attributes() -> list[str]: + return sorted((*module_globals, *export_routes, *lazy_submodules)) + + return sorted(export_routes), get_attribute, list_attributes + + +__all__ = ["make_operation_api"] diff --git a/python/cudnn/_tensor_desc.py b/python/cudnn/_tensor_desc.py new file mode 100644 index 000000000..1080e2d11 --- /dev/null +++ b/python/cudnn/_tensor_desc.py @@ -0,0 +1,138 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral tensor metadata for operation kernels.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from numbers import Real +from operator import index +from typing import Generic, TypeVar + +from . import data_type + +DTypeT = TypeVar("DTypeT") + + +@dataclass(frozen=True) +class TensorDesc(Generic[DTypeT]): + """Framework-neutral tensor signature used by operation kernels. + + ``stride_order`` lists dimension indices from the smallest stride to the + largest. Framework adapters may extend this descriptor with storage or + lowering metadata, but operation kernels only consume the fields defined + here and :attr:`cudnn_dtype`. ``init_value`` is a scalar used to initialize + inferred output or workspace storage; ``None`` leaves the storage + uninitialized. + """ + + dtype: DTypeT + shape: tuple[int, ...] + stride: tuple[int, ...] + stride_order: tuple[int, ...] + ndim: int = field(init=False) + name: str = "" + init_value: bool | int | float | None = None + + def __post_init__(self) -> None: + shape = self._integer_tuple(self.shape, "shape") + stride = self._integer_tuple(self.stride, "stride") + stride_order = self._integer_tuple(self.stride_order, "stride_order") + ndim = len(shape) + + if self.init_value is not None and not isinstance(self.init_value, (bool, Real)): + raise TypeError(f"TensorDesc.init_value must be a real scalar or None, got {self.init_value!r}") + + if any(size < 0 for size in shape): + raise ValueError(f"TensorDesc.shape entries must be non-negative, got {shape}") + if any(value < 0 for value in stride): + raise ValueError(f"TensorDesc.stride entries must be non-negative, got {stride}") + if len(stride) != ndim: + raise ValueError(f"Stride rank mismatch: expected {ndim}, got {len(stride)}") + if len(stride_order) != ndim: + raise ValueError(f"Stride order rank mismatch: expected {ndim}, got {len(stride_order)}") + if tuple(sorted(stride_order)) != tuple(range(ndim)): + raise ValueError(f"Stride order must be a permutation of [0, {ndim - 1}], got {stride_order}") + ordered_strides = tuple(stride[dimension] for dimension in stride_order) + if any(left > right for left, right in zip(ordered_strides, ordered_strides[1:])): + raise ValueError(f"Stride order {stride_order} is inconsistent with stride {stride}") + + object.__setattr__(self, "shape", shape) + object.__setattr__(self, "stride", stride) + object.__setattr__(self, "stride_order", stride_order) + object.__setattr__(self, "ndim", ndim) + + @staticmethod + def _integer_tuple(values, field_name: str) -> tuple[int, ...]: + normalized = [] + for value in values: + if isinstance(value, bool): + raise TypeError(f"TensorDesc.{field_name} entries must be integers, got {value!r}") + try: + normalized.append(index(value)) + except TypeError as error: + raise TypeError(f"TensorDesc.{field_name} entries must be integers, got {value!r}") from error + return tuple(normalized) + + @property + def cudnn_dtype(self) -> data_type: + """Return the canonical dtype consumed by operation kernels. + + Framework-specific subclasses override this property when ``dtype`` is + native to that framework. + """ + + if not isinstance(self.dtype, data_type): + raise TypeError(f"{type(self).__name__}.dtype must provide a cudnn.data_type mapping, " f"got {type(self.dtype).__name__}") + return self.dtype + + +def make_compact_tensor_desc( + *, + dtype: data_type, + shape: tuple[int, ...], + stride_order: tuple[int, ...] | None = None, + name: str = "", + init_value: bool | int | float | None = None, +) -> TensorDesc[data_type]: + """Construct a framework-neutral descriptor for a compact tensor. + + ``stride_order`` lists dimensions from fastest varying to slowest varying. + It defaults to the conventional row-major order. + """ + + if not isinstance(dtype, data_type): + raise TypeError(f"dtype must be a cudnn.data_type, got {type(dtype).__name__}") + + shape = TensorDesc._integer_tuple(shape, "shape") + if any(size < 0 for size in shape): + raise ValueError(f"TensorDesc.shape entries must be non-negative, got {shape}") + + if stride_order is None: + stride_order = tuple(reversed(range(len(shape)))) + else: + stride_order = TensorDesc._integer_tuple(stride_order, "stride_order") + + if len(stride_order) != len(shape): + raise ValueError(f"Stride order rank mismatch: expected {len(shape)}, got {len(stride_order)}") + if tuple(sorted(stride_order)) != tuple(range(len(shape))): + raise ValueError(f"Stride order must be a permutation of [0, {len(shape) - 1}], got {stride_order}") + + stride = [0] * len(shape) + running = 1 + for dimension in stride_order: + stride[dimension] = running + running *= max(shape[dimension], 1) + + return TensorDesc( + dtype=dtype, + shape=shape, + stride=tuple(stride), + stride_order=stride_order, + name=name, + init_value=init_value, + ) + + +__all__ = ["TensorDesc", "make_compact_tensor_desc"] diff --git a/python/cudnn/api_base.py b/python/cudnn/api_base.py index 2a2a1741a..30895efc7 100644 --- a/python/cudnn/api_base.py +++ b/python/cudnn/api_base.py @@ -11,7 +11,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any, List, Tuple, Optional import logging import threading @@ -20,8 +20,14 @@ import torch import cutlass.cute as cute +from cudnn import data_type from cudnn._experimental_warnings import warn_experimental_api_once -from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn._tensor_desc import TensorDesc as _TensorDesc +from cudnn.datatypes import ( + _convert_to_cutlass_data_type, + _cudnn_to_torch_data_type, + _torch_to_cudnn_data_type, +) def ceil_div(a: int, b: int) -> int: @@ -53,23 +59,37 @@ def _reset_experimental_api_warning_registry() -> None: _experimental_api_warnings_emitted.clear() -@dataclass(frozen=True) -class TensorDesc: - """Metadata needed to validate/compile tensor signatures without storage.""" +@dataclass(frozen=True, init=False) +class TensorDesc(_TensorDesc[torch.dtype]): + """Torch tensor signature with device and packed-storage metadata.""" - dtype: torch.dtype - shape: Tuple[int, ...] - stride: Tuple[int, ...] - stride_order: Tuple[int, ...] device: torch.device interpret_uint8_as_fp4x2: bool = False - ndim: int = field(init=False) - name: str = "" + + def __init__( + self, + dtype: torch.dtype, + shape: Tuple[int, ...], + stride: Tuple[int, ...], + stride_order: Tuple[int, ...], + device: torch.device, + interpret_uint8_as_fp4x2: bool = False, + name: str = "", + init_value: bool | int | float | None = None, + ) -> None: + object.__setattr__(self, "device", device) + object.__setattr__(self, "interpret_uint8_as_fp4x2", interpret_uint8_as_fp4x2) + super().__init__( + dtype=dtype, + shape=shape, + stride=stride, + stride_order=stride_order, + name=name, + init_value=init_value, + ) def __post_init__(self): - shape = tuple(self.shape) - stride = tuple(self.stride) - stride_order = tuple(self.stride_order) + super().__post_init__() device = self.device if not isinstance(device, torch.device): try: @@ -77,19 +97,14 @@ def __post_init__(self): except (TypeError, ValueError, RuntimeError) as exc: raise TypeError(f"Invalid device for TensorDesc: {self.device!r}") from exc - ndim = len(shape) - if len(stride) != ndim: - raise ValueError(f"Stride rank mismatch: expected {ndim}, got {len(stride)}") - if len(stride_order) != ndim: - raise ValueError(f"Stride order rank mismatch: expected {ndim}, got {len(stride_order)}") - if tuple(sorted(stride_order)) != tuple(range(ndim)): - raise ValueError(f"Stride order must be a permutation of [0, {ndim - 1}], got {stride_order}") - - object.__setattr__(self, "shape", shape) - object.__setattr__(self, "stride", stride) - object.__setattr__(self, "stride_order", stride_order) object.__setattr__(self, "device", device) - object.__setattr__(self, "ndim", ndim) + + @property + def cudnn_dtype(self) -> data_type: + if self.interpret_uint8_as_fp4x2 and self.dtype == torch.uint8: + return data_type.FP4_E2M1 + converted = _torch_to_cudnn_data_type(self.dtype) + return data_type.NOT_SET if converted is None else converted @staticmethod def _normalize_dim(dim: int, ndim: int, *, allow_new_dim: bool = False) -> int: @@ -181,6 +196,7 @@ def _with_layout(self, shape: Tuple[int, ...], stride: Tuple[int, ...]) -> "Tens device=self.device, interpret_uint8_as_fp4x2=self.interpret_uint8_as_fp4x2, name=self.name, + init_value=self.init_value, ) def __len__(self) -> int: @@ -910,6 +926,63 @@ def _make_fake_cute_tensor_like( assumed_align=assumed_align, ) + @staticmethod + def _to_tensor_desc( + tensor: torch.Tensor, + name: str = "", + *, + shape: Optional[Tuple[int, ...]] = None, + stride: Optional[Tuple[int, ...]] = None, + interpret_uint8_as_fp4x2: bool = False, + ) -> TensorDesc: + """Create a Torch descriptor from tensor and optional logical layout metadata.""" + + tensor_shape = tuple(tensor.shape) if shape is None else tuple(shape) + tensor_stride = tuple(tensor.stride()) if stride is None else tuple(stride) + tensor_stride_order = tuple(i for i, _ in sorted(enumerate(tensor_stride), key=lambda item: (item[1], tensor_shape[item[0]]))) + return TensorDesc( + dtype=tensor.dtype, + shape=tensor_shape, + stride=tensor_stride, + stride_order=tensor_stride_order, + device=tensor.device, + interpret_uint8_as_fp4x2=interpret_uint8_as_fp4x2, + name=name, + ) + + @staticmethod + def _materialize_tensor_desc( + desc: _TensorDesc[Any], + *, + device: torch.device, + stream: Optional[cuda.CUstream] = None, + ) -> torch.Tensor: + """Allocate a Torch tensor from a canonical tensor descriptor.""" + + if not isinstance(desc, _TensorDesc): + raise TypeError(f"desc must be a TensorDesc, got {type(desc).__name__}") + dtype = _cudnn_to_torch_data_type(desc.cudnn_dtype) + if dtype is None: + raise ValueError(f"Unsupported Torch output dtype {desc.cudnn_dtype}") + + def allocate() -> torch.Tensor: + tensor = torch.empty_strided( + desc.shape, + desc.stride, + dtype=dtype, + device=device, + ) + if desc.init_value is not None: + tensor.fill_(desc.init_value) + return tensor + + if stream is None: + return allocate() + if torch.device(device).type != "cuda": + raise ValueError("A CUDA stream can only be used to materialize a CUDA tensor") + with torch.cuda.stream(torch.cuda.ExternalStream(int(stream), device=device)): + return allocate() + def _make_tensor_desc( self, tensor: Optional[torch.Tensor], @@ -928,15 +1001,12 @@ def _make_tensor_desc( tensor_stride = self._tensor_stride(tensor, name=name) finally: self._interpret_uint8_as_fp4x2 = prev_interpret - tensor_stride_order = tuple(i for i, s in sorted(enumerate(tensor_stride), key=lambda x: (x[1], tensor_shape[x[0]]))) - return TensorDesc( - dtype=tensor.dtype, + return self._to_tensor_desc( + tensor, + name, shape=tensor_shape, stride=tensor_stride, - stride_order=tensor_stride_order, - device=tensor.device, interpret_uint8_as_fp4x2=interpret_uint8_as_fp4x2, - name=name, ) def _make_fake_cute_tensor_from_desc( diff --git a/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py b/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py index eaf7d7d98..8bc17b253 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py @@ -168,6 +168,7 @@ def __init__( stride_order=self.d_desc.stride_order, device=self.d_desc.device, name="sample_d_col", + init_value=self.d_desc.init_value, ) self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row") self.sfd_col_desc = self._make_tensor_desc(sample_sfd_col, name="sample_sfd_col") diff --git a/python/cudnn/grouped_gemm/grouped_gemm_srelu/api.py b/python/cudnn/grouped_gemm/grouped_gemm_srelu/api.py index 2ff670e08..97d7d26da 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_srelu/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_srelu/api.py @@ -179,6 +179,7 @@ def __init__( stride_order=self.d_desc.stride_order, device=self.d_desc.device, name="sample_d_col", + init_value=self.d_desc.init_value, ) self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row") self.sfd_col_desc = self._make_tensor_desc(sample_sfd_col, name="sample_sfd_col") diff --git a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py index 29b3d7328..d7fd32071 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py @@ -101,6 +101,7 @@ def __init__( stride_order=tuple(i for i, s in sorted(enumerate(self.wgrad_desc.stride[1:]), key=lambda x: x[1])), device=self.wgrad_desc.device, name="single_expert_wgrad", + init_value=self.wgrad_desc.init_value, ) else: # MoEWeightMode.DISCRETE self.expert_cnt = num_experts diff --git a/python/cudnn/jax/__init__.py b/python/cudnn/jax/__init__.py new file mode 100644 index 000000000..8d31c8e3c --- /dev/null +++ b/python/cudnn/jax/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Public facade for framework-specific JAX operation APIs.""" + +from .._jax import JaxApiBase, JaxTensorDesc +from ..rmsnorm_rht_amax.jax import ( + RmsNormRhtAmaxSm100, + rmsnorm_rht_amax_sm100, +) + +__all__ = [ + "JaxApiBase", + "JaxTensorDesc", + "RmsNormRhtAmaxSm100", + "rmsnorm_rht_amax_sm100", +] diff --git a/python/cudnn/rmsnorm_rht_amax/__init__.py b/python/cudnn/rmsnorm_rht_amax/__init__.py index 9da5cae15..3cab20e81 100644 --- a/python/cudnn/rmsnorm_rht_amax/__init__.py +++ b/python/cudnn/rmsnorm_rht_amax/__init__.py @@ -1,16 +1,22 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - RmsNormRhtAmaxSm100, - best_num_threads, - pick_rows_per_cta, - rmsnorm_rht_amax_wrapper_sm100, -) +"""Lazy framework adapters for RMSNorm + RHT + amax.""" + +from .._operation_api import make_operation_api -__all__ = [ - "RmsNormRhtAmaxSm100", - "best_num_threads", - "pick_rows_per_cta", - "rmsnorm_rht_amax_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "kernel": ( + "RMSNormRHTAmaxKernel", + "best_num_threads", + "pick_rows_per_cta", + ), + "api": ( + "RmsNormRhtAmaxSm100", + "rmsnorm_rht_amax_wrapper_sm100", + ), + }, + submodules=("api", "jax", "kernel"), +) diff --git a/python/cudnn/rmsnorm_rht_amax/api.py b/python/cudnn/rmsnorm_rht_amax/api.py index 0294a3a22..2f21964e7 100644 --- a/python/cudnn/rmsnorm_rht_amax/api.py +++ b/python/cudnn/rmsnorm_rht_amax/api.py @@ -7,57 +7,51 @@ from typing import Optional from cuda.bindings import driver as cuda -import cutlass import cutlass.cute as cute import torch -from cutlass import Float32 from cutlass.cute.runtime import make_fake_stream from cudnn.api_base import APIBase, TupleDict -from .kernel import RMSNormRHTAmaxKernel +from .kernel import ( + DEFAULT_NUM_THREADS_BY_N, + RPC_CANDIDATES, + TARGET_MIN_CTAS, + RMSNormRHTAmaxKernel, + best_num_threads, + pick_rows_per_cta, +) -DEFAULT_NUM_THREADS_BY_N = { - 2048: 128, - 4096: 256, - 7168: 128, - 8192: 512, - 16384: 1024, - 32768: 512, -} -RPC_CANDIDATES = (2, 4, 8) -TARGET_MIN_CTAS = 148 +_TENSOR_ALIGNMENT = RMSNormRHTAmaxKernel.TENSOR_ALIGNMENT -def best_num_threads(n: int) -> Optional[int]: - for num_threads in (1024, 512, 256, 128, 64): - if n % num_threads != 0: - continue - ept = n // num_threads - if ept >= 8 and ept % 8 == 0: - return num_threads - return None +def _validate_torch_tensors(*named_tensors: tuple[str, torch.Tensor]) -> None: + """Validate framework properties that are absent from TensorDesc.""" - -def pick_rows_per_cta(m: int) -> int: - for rows_per_cta in reversed(RPC_CANDIDATES): - if m % rows_per_cta != 0: - continue - num_ctas = m // rows_per_cta - if num_ctas >= TARGET_MIN_CTAS: - return rows_per_cta - return RPC_CANDIDATES[0] + if not named_tensors: + return + expected_device = named_tensors[0][1].device + for name, tensor in named_tensors: + if tensor.device != expected_device: + raise ValueError(f"{name} must be on {expected_device}, got {tensor.device}") + if tensor.data_ptr() % _TENSOR_ALIGNMENT != 0: + raise ValueError(f"{name} must be {_TENSOR_ALIGNMENT}-byte aligned") class RmsNormRhtAmaxSm100(APIBase): - """Class API for the RMSNorm + RHT + amax kernel.""" + """Class API for the RMSNorm + RHT + amax kernel. + + Logical validation, output inference, and resolved launch state are owned + by :attr:`kernel`. Output samples remain optional for callers that want + the framework adapter to allocate inferred outputs. + """ def __init__( self, sample_x: torch.Tensor, sample_w: torch.Tensor, - sample_o: torch.Tensor, - sample_amax: torch.Tensor, + sample_o: Optional[torch.Tensor] = None, + sample_amax: Optional[torch.Tensor] = None, eps: float = 1e-5, num_threads: Optional[int] = None, rows_per_cta: Optional[int] = None, @@ -65,66 +59,98 @@ def __init__( super().__init__() self._warn_experimental_api() + if (sample_o is None) != (sample_amax is None): + raise ValueError("sample_o and sample_amax must either both be provided or both be omitted") + + named_tensors = [("sample_x", sample_x), ("sample_w", sample_w)] + if sample_o is not None and sample_amax is not None: + named_tensors.extend((("sample_o", sample_o), ("sample_amax", sample_amax))) + _validate_torch_tensors(*named_tensors) self.x_desc = self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_x, name="sample_x"), 2, "sample_x") self.w_desc = self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_w, name="sample_w"), 1, "sample_w") - self.o_desc = self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_o, name="sample_o"), 2, "sample_o") - self.amax_desc = self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_amax, name="sample_amax"), 1, "sample_amax") + self.o_desc = None if sample_o is None else self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_o, name="sample_o"), 2, "sample_o") + self.amax_desc = None if sample_amax is None else self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_amax, name="sample_amax"), 1, "sample_amax") + + self.kernel = RMSNormRHTAmaxKernel( + x=self.x_desc, + weight=self.w_desc, + output=self.o_desc, + amax=self.amax_desc, + eps=eps, + num_threads=num_threads, + rows_per_cta=rows_per_cta, + ) - self.eps = eps - self.requested_num_threads = num_threads - self.requested_rows_per_cta = rows_per_cta - self.num_threads = None - self.rows_per_cta = None - self.n = None + @property + def eps(self) -> float: + return self.kernel.eps + + @eps.setter + def eps(self, value: float) -> None: + self.kernel.eps = value + self._is_supported = False + self._compiled_kernel = None + + @property + def requested_num_threads(self) -> Optional[int]: + return self.kernel.requested_num_threads + + @requested_num_threads.setter + def requested_num_threads(self, value: Optional[int]) -> None: + self.kernel.requested_num_threads = value + self._is_supported = False + self._compiled_kernel = None + + @property + def requested_rows_per_cta(self) -> Optional[int]: + return self.kernel.requested_rows_per_cta + + @requested_rows_per_cta.setter + def requested_rows_per_cta(self, value: Optional[int]) -> None: + self.kernel.requested_rows_per_cta = value + self._is_supported = False + self._compiled_kernel = None + + @property + def num_threads(self) -> Optional[int]: + return self.kernel.num_threads + + @property + def rows_per_cta(self) -> Optional[int]: + return self.kernel.rows_per_cta + + @property + def n(self) -> Optional[int]: + return self.kernel.n + + def _materialize_outputs( + self, + *, + current_stream: Optional[cuda.CUstream] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Infer and allocate outputs using this API's kernel instance.""" - def check_support(self) -> bool: - m, n = self._tensor_shape(self.x_desc, name="sample_x") - w_n = self._tensor_shape(self.w_desc, name="sample_w")[0] - o_m, o_n = self._tensor_shape(self.o_desc, name="sample_o") - - self._check_tensor_shape(self.x_desc, (m, n), "X") - self._check_tensor_shape(self.w_desc, (n,), "W") - self._check_tensor_shape(self.o_desc, (m, n), "O") - self._value_error_if(w_n != n, f"W length must match X hidden dimension, got {w_n} and {n}") - self._value_error_if((n % 16) != 0, f"N must be divisible by 16 for the Hadamard block size, got {n}") - self._value_error_if(o_m != m or o_n != n, f"O shape must match X shape, got {(o_m, o_n)} and {(m, n)}") - - self._check_tensor_stride(self.x_desc, stride=(n, 1), name="X", extra_error_msg="X must be row-major contiguous") - self._check_tensor_stride(self.w_desc, stride=(1,), name="W", extra_error_msg="W must be contiguous") - self._check_tensor_stride(self.o_desc, stride=(n, 1), name="O", extra_error_msg="O must be row-major contiguous") - - self._check_dtype(self.x_desc, dtype=torch.bfloat16, name="X") - self._check_dtype(self.w_desc, dtype=torch.bfloat16, name="W") - self._check_dtype(self.o_desc, dtype=torch.bfloat16, name="O") - self._check_dtype(self.amax_desc, dtype=torch.float32, name="Amax") - - resolved_num_threads = self.requested_num_threads - if resolved_num_threads is None: - resolved_num_threads = DEFAULT_NUM_THREADS_BY_N.get(n, best_num_threads(n)) - self._value_error_if(resolved_num_threads is None, f"No valid num_threads found for N={n}") - self._value_error_if(resolved_num_threads <= 0, f"num_threads must be positive, got {resolved_num_threads}") - self._value_error_if( - (resolved_num_threads % 32) != 0, - f"num_threads must be warp-aligned, got {resolved_num_threads}", - ) - self._value_error_if( - resolved_num_threads > 1024, - f"num_threads must not exceed the CUDA block size limit, got {resolved_num_threads}", - ) + self._ensure_support_checked() + output_desc, amax_desc = self.kernel.infer_output() + o_tensor = self._materialize_tensor_desc(output_desc, device=self.x_desc.device, stream=current_stream) + amax_tensor = self._materialize_tensor_desc(amax_desc, device=self.x_desc.device, stream=current_stream) - resolved_rows_per_cta = self.requested_rows_per_cta - if resolved_rows_per_cta is None: - resolved_rows_per_cta = pick_rows_per_cta(m) + if self.o_desc is None and self.amax_desc is None: + self.o_desc = self._to_tensor_desc(o_tensor, "sample_o") + self.amax_desc = self._to_tensor_desc(amax_tensor, "sample_amax") + self.kernel.output = self.o_desc + self.kernel.amax = self.amax_desc - self._value_error_if(m % resolved_rows_per_cta != 0, f"M must be divisible by rows_per_cta, got M={m}, rows_per_cta={resolved_rows_per_cta}") - self._value_error_if(n % resolved_num_threads != 0, f"N={n} must be divisible by num_threads={resolved_num_threads}") + return o_tensor, amax_tensor - ept = n // resolved_num_threads - self._value_error_if(ept < 8 or ept % 8 != 0, f"EPT={ept} must be >= 8 and divisible by 8") + def check_support(self) -> bool: + self._is_supported = False + self.kernel.check_support() - expected_num_ctas = m // resolved_rows_per_cta - self._check_tensor_shape(self.amax_desc, (expected_num_ctas,), "Amax") + for name, desc in (("W", self.w_desc), ("O", self.o_desc), ("Amax", self.amax_desc)): + if desc is not None: + self._value_error_if(desc.device != self.x_desc.device, f"{name} must be on {self.x_desc.device}, got {desc.device}") self._runtime_error_if(not torch.cuda.is_available(), "CUDA is not available") major, minor = torch.cuda.get_device_capability(self.x_desc.device) @@ -134,9 +160,6 @@ def check_support(self) -> bool: f"RmsNormRhtAmaxSm100 requires SM100+, found SM{compute_capability}", ) - self.num_threads = resolved_num_threads - self.rows_per_cta = resolved_rows_per_cta - self.n = n self._is_supported = True return True @@ -144,47 +167,42 @@ def compile(self) -> None: self._ensure_support_checked() if self._compiled_kernel is not None: return + self._runtime_error_if(self.o_desc is None or self.amax_desc is None, "Output tensors must be provided or materialized before compile()") - kernel = RMSNormRHTAmaxKernel( - n=self.n, - num_threads=self.num_threads, - eps=self.eps, - rows_per_cta=self.rows_per_cta, - ) - - valid_m = cute.sym_int(divisibility=self.rows_per_cta) + valid_m = cute.sym_int(divisibility=self.kernel.rows_per_cta) fake_x_tensor = self._make_fake_cute_compact_tensor( dtype=self.x_desc.dtype, - shape=(valid_m, self.n), + shape=(valid_m, self.kernel.n), stride_order=self.x_desc.stride_order, + assumed_align=_TENSOR_ALIGNMENT, dynamic_mode=None, - divisibility=self.rows_per_cta, + divisibility=self.kernel.rows_per_cta, ) - fake_w_tensor = self._make_fake_cute_tensor_from_desc(self.w_desc, assumed_align=16) + fake_w_tensor = self._make_fake_cute_tensor_from_desc(self.w_desc, assumed_align=_TENSOR_ALIGNMENT) fake_o_tensor = self._make_fake_cute_compact_tensor( dtype=self.o_desc.dtype, - shape=(valid_m, self.n), + shape=(valid_m, self.kernel.n), stride_order=self.o_desc.stride_order, + assumed_align=_TENSOR_ALIGNMENT, dynamic_mode=None, - divisibility=self.rows_per_cta, + divisibility=self.kernel.rows_per_cta, ) fake_num_ctas = cute.sym_int() fake_amax_tensor = self._make_fake_cute_tensor( dtype=self.amax_desc.dtype, shape=(fake_num_ctas,), stride=self.amax_desc.stride, - assumed_align=16, + assumed_align=_TENSOR_ALIGNMENT, ) fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) compiled_kernel = cute.compile( - kernel, + self.kernel, fake_x_tensor, fake_w_tensor, fake_o_tensor, fake_amax_tensor, - Float32(self.eps), fake_stream, options="--enable-tvm-ffi", ) @@ -201,7 +219,6 @@ def tensor_api( w_tensor, o_tensor, amax_tensor, - Float32(self.eps), stream, ) @@ -221,6 +238,12 @@ def execute( w_tensor = self._unpad_tensor_to_ndim(w_tensor, 1, "w_tensor") o_tensor = self._unpad_tensor_to_ndim(o_tensor, 2, "o_tensor") amax_tensor = self._unpad_tensor_to_ndim(amax_tensor, 1, "amax_tensor") + _validate_torch_tensors( + ("x_tensor", x_tensor), + ("w_tensor", w_tensor), + ("o_tensor", o_tensor), + ("amax_tensor", amax_tensor), + ) if current_stream is None: current_stream = cuda.CUstream(torch.cuda.current_stream(x_tensor.device).cuda_stream) @@ -250,47 +273,41 @@ def rmsnorm_rht_amax_wrapper_sm100( x_tensor = x_tensor.squeeze(-1) if x_tensor.ndim == 3 and x_tensor.shape[-1] == 1 else x_tensor w_tensor = w_tensor.squeeze(-1) if w_tensor.ndim == 2 and w_tensor.shape[-1] == 1 else w_tensor - - m, n = x_tensor.shape - resolved_num_threads = num_threads if num_threads is not None else DEFAULT_NUM_THREADS_BY_N.get(n, best_num_threads(n)) - if resolved_num_threads is None: - raise ValueError(f"No valid num_threads found for N={n}") - resolved_rows_per_cta = rows_per_cta if rows_per_cta is not None else pick_rows_per_cta(m) - if m % resolved_rows_per_cta != 0: - raise ValueError(f"M must be divisible by rows_per_cta, got M={m}, rows_per_cta={resolved_rows_per_cta}") - - o_tensor = torch.empty_like(x_tensor) - amax_tensor = torch.full((m // resolved_rows_per_cta,), float("-inf"), dtype=torch.float32, device=x_tensor.device) + _validate_torch_tensors(("x_tensor", x_tensor), ("w_tensor", w_tensor)) cache_key = ( - n, + tuple(x_tensor.shape), + tuple(w_tensor.shape), x_tensor.dtype, w_tensor.dtype, - o_tensor.dtype, tuple(x_tensor.stride()), tuple(w_tensor.stride()), - tuple(o_tensor.stride()), + x_tensor.device, eps, - resolved_num_threads, - resolved_rows_per_cta, + num_threads, + rows_per_cta, ) - if cache_key in _cache_of_RmsNormRhtAmaxSm100Objects: - api = _cache_of_RmsNormRhtAmaxSm100Objects[cache_key] - else: + api = _cache_of_RmsNormRhtAmaxSm100Objects.get(cache_key) + cache_miss = api is None + if cache_miss: api = RmsNormRhtAmaxSm100( sample_x=x_tensor, sample_w=w_tensor, - sample_o=o_tensor, - sample_amax=amax_tensor, eps=eps, - num_threads=resolved_num_threads, - rows_per_cta=resolved_rows_per_cta, + num_threads=num_threads, + rows_per_cta=rows_per_cta, ) - assert api.check_support(), "Unsupported configuration" + api.check_support() + + o_tensor, amax_tensor = api._materialize_outputs(current_stream=current_stream) + if cache_miss: api.compile() _cache_of_RmsNormRhtAmaxSm100Objects[cache_key] = api + if current_stream is None and x_tensor.is_cuda: + current_stream = cuda.CUstream(torch.cuda.current_stream(x_tensor.device).cuda_stream) + api.execute( x_tensor=x_tensor, w_tensor=w_tensor, diff --git a/python/cudnn/rmsnorm_rht_amax/jax.py b/python/cudnn/rmsnorm_rht_amax/jax.py new file mode 100644 index 000000000..fdf80b28e --- /dev/null +++ b/python/cudnn/rmsnorm_rht_amax/jax.py @@ -0,0 +1,78 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX API for fused RMSNorm + RHT + per-CTA amax.""" + +from __future__ import annotations + +from typing import Any, Optional + +from .._jax import JaxApiBase, jax +from .kernel import RMSNormRHTAmaxKernel + + +class RmsNormRhtAmaxSm100(JaxApiBase): + """JAX callable specialized from sample shape and dtype metadata.""" + + def __init__( + self, + sample_x: Any, + sample_w: Any, + eps: float = 1e-5, + num_threads: Optional[int] = None, + rows_per_cta: Optional[int] = None, + ) -> None: + self.x_desc = self._to_tensor_desc(sample_x, "sample_x") + self.w_desc = self._to_tensor_desc(sample_w, "sample_w") + self.kernel = RMSNormRHTAmaxKernel( + x=self.x_desc, + weight=self.w_desc, + eps=eps, + num_threads=num_threads, + rows_per_cta=rows_per_cta, + ) + + def check_support(self) -> bool: + return self.kernel.check_support() + + def __call__(self, x: Any, weight: Any) -> tuple[Any, Any]: + self.check_support() + self._check_tensor_signature(x, self.x_desc) + self._check_tensor_signature(weight, self.w_desc) + + x_spec = self._to_tensor_spec( + self.x_desc, + divisibility=(self.kernel.rows_per_cta, 16), + ) + weight_spec = self._to_tensor_spec(self.w_desc, divisibility=(16,)) + return self._call_kernel( + (x, weight), + input_spec=(x_spec, weight_spec), + output_spec=(x_spec, None), + ) + + +@jax.jit(static_argnames=("eps", "num_threads", "rows_per_cta")) +def rmsnorm_rht_amax_sm100( + x: Any, + weight: Any, + *, + eps: float = 1e-5, + num_threads: Optional[int] = None, + rows_per_cta: Optional[int] = None, +) -> tuple[Any, Any]: + """Apply fused RMSNorm, 16-wide RHT, and per-CTA amax from JAX.""" + + return RmsNormRhtAmaxSm100( + x, + weight, + eps=eps, + num_threads=num_threads, + rows_per_cta=rows_per_cta, + )(x, weight) + + +__all__ = [ + "RmsNormRhtAmaxSm100", + "rmsnorm_rht_amax_sm100", +] diff --git a/python/cudnn/rmsnorm_rht_amax/kernel.py b/python/cudnn/rmsnorm_rht_amax/kernel.py index 2445cae94..9d6382572 100644 --- a/python/cudnn/rmsnorm_rht_amax/kernel.py +++ b/python/cudnn/rmsnorm_rht_amax/kernel.py @@ -3,8 +3,11 @@ """CUTE DSL kernel for fused RMSNorm + RHT + per-CTA amax.""" +from __future__ import annotations + import math import operator +from typing import Any, Optional import cuda.bindings.driver as cuda import cutlass @@ -15,6 +18,40 @@ from cutlass.cute.arch import shuffle_sync_bfly from cutlass.cutlass_dsl import T, dsl_user_op +from .. import data_type +from .._op_kernel import OpKernel +from .._tensor_desc import TensorDesc, make_compact_tensor_desc + +DEFAULT_NUM_THREADS_BY_N = { + 2048: 128, + 4096: 256, + 7168: 128, + 8192: 512, + 16384: 1024, + 32768: 512, +} +RPC_CANDIDATES = (2, 4, 8) +TARGET_MIN_CTAS = 148 + + +def best_num_threads(n: int) -> Optional[int]: + for num_threads in (1024, 512, 256, 128, 64): + if n % num_threads != 0: + continue + ept = n // num_threads + if ept >= 8 and ept % 8 == 0: + return num_threads + return None + + +def pick_rows_per_cta(m: int) -> int: + for rows_per_cta in reversed(RPC_CANDIDATES): + if m % rows_per_cta != 0: + continue + if m // rows_per_cta >= TARGET_MIN_CTAS: + return rows_per_cta + return RPC_CANDIDATES[0] + @dsl_user_op def fabs_f32(val, *, loc=None, ip=None): @@ -68,35 +105,181 @@ def redux_sync_max_f32(val, *, loc=None, ip=None): return Float32(result) -class RMSNormRHTAmaxKernel: - """Fused RMSNorm + block-diagonal Hadamard + running per-CTA amax.""" +class RMSNormRHTAmaxKernel(OpKernel): + """Fused RMSNorm + block-diagonal Hadamard + running per-CTA amax. + + Framework adapters bind generic tensor descriptors, call + :meth:`check_support`, and materialize the descriptors returned by + :meth:`infer_output` in their own framework. + """ COPY_BITS = 128 HAD_BLOCK = 16 - - def __init__(self, n, num_threads=256, eps=1e-5, rows_per_cta=8): + TENSOR_ALIGNMENT = COPY_BITS // 8 + output_dtype = data_type.BFLOAT16 + amax_dtype = data_type.FLOAT + + def __init__( + self, + *, + x: TensorDesc[Any], + weight: TensorDesc[Any], + output: Optional[TensorDesc[Any]] = None, + amax: Optional[TensorDesc[Any]] = None, + eps: float = 1e-5, + num_threads: Optional[int] = None, + rows_per_cta: Optional[int] = None, + ) -> None: + for name, desc in (("x", x), ("weight", weight)): + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + for name, desc in (("output", output), ("amax", amax)): + if desc is not None and not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + + self.x = x + self.weight = weight + self.output = output + self.amax = amax + self.eps = eps + self.requested_num_threads = num_threads + self.requested_rows_per_cta = rows_per_cta + + self.m: Optional[int] = None + self.n: Optional[int] = None + self.num_threads: Optional[int] = None + self.rows_per_cta: Optional[int] = None + + def check_support(self) -> bool: + """Validate descriptors and resolve the kernel launch configuration.""" + + self.m = None + self.n = None + self.num_threads = None + self.rows_per_cta = None + + if self.x.ndim != 2: + raise ValueError(f"X must have rank 2, got shape {self.x.shape}") + if self.weight.ndim != 1: + raise ValueError(f"W must have rank 1, got shape {self.weight.shape}") + if self.x.cudnn_dtype != data_type.BFLOAT16: + raise ValueError(f"X must have dtype bfloat16, got {self.x.dtype}") + if self.weight.cudnn_dtype != data_type.BFLOAT16: + raise ValueError(f"W must have dtype bfloat16, got {self.weight.dtype}") + + m, n = self.x.shape + if m <= 0: + raise ValueError(f"M must be positive, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if self.weight.shape != (n,): + raise ValueError(f"W must have shape {(n,)}, got {self.weight.shape}") + if self.x.stride != (n, 1) or self.x.stride_order != (1, 0): + raise ValueError(f"X must be row-major contiguous, got stride {self.x.stride} " f"and stride order {self.x.stride_order}") + if self.weight.stride != (1,) or self.weight.stride_order != (0,): + raise ValueError(f"W must be contiguous, got stride {self.weight.stride} " f"and stride order {self.weight.stride_order}") + if n % self.HAD_BLOCK != 0: + raise ValueError(f"N must be divisible by {self.HAD_BLOCK} for the Hadamard block size, got {n}") + + num_threads = self.requested_num_threads + if num_threads is None: + num_threads = DEFAULT_NUM_THREADS_BY_N.get(n, best_num_threads(n)) + if num_threads is None: + raise ValueError(f"No valid num_threads found for N={n}") + + rows_per_cta = self.requested_rows_per_cta + if rows_per_cta is None: + rows_per_cta = pick_rows_per_cta(m) + + self._validate_launch_configuration(n, num_threads, rows_per_cta, m=m) + + expected_output_shape = (m, n) + expected_amax_shape = (m // rows_per_cta,) + if self.output is not None: + if self.output.shape != expected_output_shape: + raise ValueError(f"O must have shape {expected_output_shape}, got {self.output.shape}") + if self.output.cudnn_dtype != self.output_dtype: + raise ValueError(f"O must have dtype bfloat16, got {self.output.dtype}") + if self.output.stride != (n, 1) or self.output.stride_order != (1, 0): + raise ValueError(f"O must be row-major contiguous, got stride {self.output.stride} " f"and stride order {self.output.stride_order}") + if self.amax is not None: + if self.amax.shape != expected_amax_shape: + raise ValueError(f"Amax must have shape {expected_amax_shape}, got {self.amax.shape}") + if self.amax.cudnn_dtype != self.amax_dtype: + raise ValueError(f"Amax must have dtype float32, got {self.amax.dtype}") + + self.m = m self.n = n self.num_threads = num_threads - self.eps = eps self.rows_per_cta = rows_per_cta - self.vec_size = self.COPY_BITS // 16 - self.ept = n // num_threads + self._configure_lowering_state() + return True + + def infer_output(self) -> tuple[TensorDesc[data_type], ...]: + """Infer compact output descriptors from validated input metadata.""" + + if self.m is None or self.n is None or self.rows_per_cta is None: + raise RuntimeError("check_support() must be called before inferring outputs") + + return ( + make_compact_tensor_desc( + dtype=self.output_dtype, + shape=(self.m, self.n), + name="output", + ), + make_compact_tensor_desc( + dtype=self.amax_dtype, + shape=(self.m // self.rows_per_cta,), + name="amax", + ), + ) - assert n % num_threads == 0, f"N={n} must be divisible by num_threads={num_threads}" - assert self.ept % self.vec_size == 0, f"EPT={self.ept} must be a multiple of vec_size={self.vec_size}" - assert self.ept >= self.vec_size, f"EPT={self.ept} must be >= vec_size={self.vec_size}" + @staticmethod + def _validate_launch_configuration( + n: int, + num_threads: int, + rows_per_cta: int, + *, + m: Optional[int] = None, + ) -> None: + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if num_threads <= 0: + raise ValueError(f"num_threads must be positive, got {num_threads}") + if num_threads % 32 != 0: + raise ValueError(f"num_threads must be warp-aligned, got {num_threads}") + if num_threads > 1024: + raise ValueError(f"num_threads must not exceed the CUDA block size limit, got {num_threads}") + if n % num_threads != 0: + raise ValueError(f"N={n} must be divisible by num_threads={num_threads}") + + ept = n // num_threads + if ept < 8 or ept % 8 != 0: + raise ValueError(f"EPT={ept} must be >= 8 and divisible by 8") + if rows_per_cta <= 0: + raise ValueError(f"rows_per_cta must be positive, got {rows_per_cta}") + if m is not None and m % rows_per_cta != 0: + raise ValueError(f"M must be divisible by rows_per_cta, got M={m}, rows_per_cta={rows_per_cta}") + + def _configure_lowering_state(self) -> None: + if self.n is None or self.num_threads is None or self.rows_per_cta is None: + raise RuntimeError("Kernel launch configuration has not been resolved") + + self._validate_launch_configuration(self.n, self.num_threads, self.rows_per_cta, m=self.m) + self.vec_size = self.COPY_BITS // 16 + self.ept = self.n // self.num_threads self.num_vec_blocks = self.ept // self.vec_size - self.warps_per_row = num_threads // 32 + self.warps_per_row = self.num_threads // 32 self.inv_sqrt_had = 1.0 / math.sqrt(self.HAD_BLOCK) self.num_intra_stages = int(math.log2(self.vec_size)) self.num_cross_stages = 1 - self.tv_shape = ((num_threads, 1), (self.vec_size, self.num_vec_blocks)) - self.tv_stride = ((self.vec_size, 1), (1, self.vec_size * num_threads)) - self.tiler_mn = (1, n) + self.tv_shape = ((self.num_threads, 1), (self.vec_size, self.num_vec_blocks)) + self.tv_stride = ((self.vec_size, 1), (1, self.vec_size * self.num_threads)) + self.tiler_mn = (1, self.n) - tile_bytes = n * 2 + tile_bytes = self.n * 2 reduce_bytes = self.warps_per_row * 4 amax_bytes = self.warps_per_row * 4 self.smem_bytes = tile_bytes + reduce_bytes + amax_bytes + 128 @@ -241,11 +424,26 @@ def kernel(self, m_x: cute.Tensor, m_w: cute.Tensor, m_o: cute.Tensor, m_amax: c m_amax[bid] = cta_max @cute.jit - def __call__(self, x_tensor: cute.Tensor, w_tensor: cute.Tensor, o_tensor: cute.Tensor, amax_tensor: cute.Tensor, eps: Float32, stream: cuda.CUstream): + def __call__( + self, + x_tensor: cute.Tensor, + w_tensor: cute.Tensor, + o_tensor: cute.Tensor, + amax_tensor: cute.Tensor, + stream: cuda.CUstream, + ): m = x_tensor.shape[0] num_ctas = m // self.rows_per_cta tv_layout = cute.make_layout(self.tv_shape, stride=self.tv_stride) - self.kernel(x_tensor, w_tensor, o_tensor, amax_tensor, eps, tv_layout, self.tiler_mn).launch( + self.kernel( + x_tensor, + w_tensor, + o_tensor, + amax_tensor, + Float32(self.eps), + tv_layout, + self.tiler_mn, + ).launch( grid=(num_ctas, 1, 1), block=(self.num_threads, 1, 1), smem=self.smem_bytes, diff --git a/test/python/fe_api/test_jax_api_base.py b/test/python/fe_api/test_jax_api_base.py new file mode 100644 index 000000000..4df53307e --- /dev/null +++ b/test/python/fe_api/test_jax_api_base.py @@ -0,0 +1,390 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for the JAX API base.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_PACKAGE = "cudnn_jax_api_base_test" + + +class _DataType(Enum): + NOT_SET = auto() + FLOAT = auto() + BFLOAT16 = auto() + + +class _ArrayMetadata: + def __init__(self, shape, dtype, label=None): + self.shape = shape + self.dtype = dtype + self.label = label + + +class _TensorSpec: + def __init__(self, *, layout=None, mode=None, divisibility=None): + self.layout = layout + self.mode = mode + self.divisibility = divisibility + + +class _ShapeDtypeStruct: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + +class JaxApiBaseTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + datatypes_name = f"{internal_name}.datatypes" + datatypes = types.ModuleType(datatypes_name) + datatypes.jax_to_cudnn_dtype = lambda dtype: { + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, + }.get(dtype, _DataType.NOT_SET) + datatypes.cudnn_to_jax_dtype = lambda dtype: { + _DataType.BFLOAT16: "bfloat16", + _DataType.FLOAT: "float32", + }[dtype] + sys.modules[datatypes_name] = datatypes + + try: + with mock.patch.dict( + sys.modules, + { + "jax": None, + "jax.numpy": None, + "cutlass": None, + "cutlass.jax": None, + }, + ): + cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + cls.module = importlib.import_module(f"{internal_name}.api_base") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def test_converts_array_metadata_to_shared_tensor_desc(self): + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3, 4), "bfloat16"), + "sample", + ) + + self.assertIsInstance(desc, self.tensor_module.TensorDesc) + self.assertEqual(desc.dtype, "bfloat16") + self.assertEqual(desc.shape, (2, 3, 4)) + self.assertEqual(desc.stride, (12, 4, 1)) + self.assertEqual(desc.stride_order, (2, 1, 0)) + self.assertEqual(desc.name, "sample") + self.assertIsNone(desc.init_value) + self.assertEqual(desc.cudnn_dtype, _DataType.BFLOAT16) + + def test_converts_public_array_axes_to_canonical_descriptor_axes(self): + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3, 4), "bfloat16"), + "sample", + mode=(2, 0, 1), + ) + + self.assertEqual(desc.shape, (4, 2, 3)) + self.assertEqual(desc.stride, (1, 12, 4)) + self.assertEqual(desc.stride_order, (0, 2, 1)) + + def test_checks_invocation_signature(self): + desc = self.module.JaxApiBase._to_tensor_desc(_ArrayMetadata((2, 3), "bfloat16"), "sample") + self.module.JaxApiBase._check_tensor_signature(_ArrayMetadata((2, 3), "bfloat16"), desc) + + with self.assertRaisesRegex(ValueError, "sample tensor shape mismatch"): + self.module.JaxApiBase._check_tensor_signature(_ArrayMetadata((1, 3), "bfloat16"), desc) + with self.assertRaisesRegex(ValueError, "sample tensor dtype mismatch"): + self.module.JaxApiBase._check_tensor_signature(_ArrayMetadata((2, 3), "float32"), desc) + + def test_checks_public_invocation_against_canonical_descriptor(self): + mode = (2, 0, 1) + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3, 4), "bfloat16"), + "sample", + mode=mode, + ) + + self.module.JaxApiBase._check_tensor_signature( + _ArrayMetadata((2, 3, 4), "bfloat16"), + desc, + mode=mode, + ) + with self.assertRaisesRegex(ValueError, "sample tensor shape mismatch"): + self.module.JaxApiBase._check_tensor_signature( + _ArrayMetadata((1, 3, 4), "bfloat16"), + desc, + mode=mode, + ) + + def test_builds_public_tensor_spec_from_canonical_metadata(self): + mode = (2, 0, 1) + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3, 4), "bfloat16"), + "sample", + mode=mode, + ) + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass_jax = types.ModuleType("cutlass.jax") + cutlass_jax.TensorSpec = _TensorSpec + cutlass.jax = cutlass_jax + + with mock.patch.dict(sys.modules, {"cutlass": cutlass, "cutlass.jax": cutlass_jax}): + spec = self.module.JaxApiBase._to_tensor_spec( + desc, + mode=mode, + divisibility=(8, 2, 4), + ) + + self.assertEqual(spec.layout, (2, 1, 0)) + self.assertEqual(spec.mode, mode) + self.assertEqual(spec.divisibility, (2, 4, 8)) + + def test_tensor_spec_defaults_to_identity_axis_binding(self): + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3), "bfloat16"), + "sample", + ) + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass_jax = types.ModuleType("cutlass.jax") + cutlass_jax.TensorSpec = _TensorSpec + cutlass.jax = cutlass_jax + + with mock.patch.dict(sys.modules, {"cutlass": cutlass, "cutlass.jax": cutlass_jax}): + spec = self.module.JaxApiBase._to_tensor_spec(desc) + + self.assertEqual(spec.layout, (1, 0)) + self.assertEqual(spec.mode, (0, 1)) + self.assertIsNone(spec.divisibility) + + def test_tensor_spec_rejects_divisibility_with_the_wrong_rank(self): + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3), "bfloat16"), + "sample", + ) + + with self.assertRaisesRegex(ValueError, "divisibility rank mismatch"): + self.module.JaxApiBase._to_tensor_spec(desc, divisibility=(2,)) + + def test_materializes_canonical_output_in_public_axis_order(self): + desc = self.tensor_module.TensorDesc( + dtype=_DataType.BFLOAT16, + shape=(4, 2, 3), + stride=(1, 12, 4), + stride_order=(0, 2, 1), + ) + jax = types.ModuleType("jax") + jax.ShapeDtypeStruct = _ShapeDtypeStruct + + with mock.patch.dict(sys.modules, {"jax": jax}): + output = self.module.JaxApiBase._materialize_tensor_desc( + desc, + mode=(2, 0, 1), + ) + + self.assertEqual(output.shape, (2, 3, 4)) + self.assertEqual(output.dtype, "bfloat16") + + def test_calls_kernel_with_initialized_outputs_and_hidden_workspaces(self): + tensor_module = self.tensor_module + seen = {} + + class Kernel: + def infer_output(self): + return ( + tensor_module.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2,), + name="output", + ), + tensor_module.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(4, 2, 3), + name="amax", + init_value=float("-inf"), + ), + ) + + def infer_workspace(self): + return ( + tensor_module.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(3,), + name="counter", + init_value=0, + ), + tensor_module.make_compact_tensor_desc( + dtype=_DataType.BFLOAT16, + shape=(5,), + name="scratch", + ), + ) + + def __call__(self, *args): + seen["kernel_args"] = args + + class Adapter(self.module.JaxApiBase): + def __init__(self): + self.kernel = Kernel() + + def check_support(self): + return True + + def __call__(self, value): + return self._call_kernel((value,)) + + full_calls = [] + jax = types.ModuleType("jax") + jax.__path__ = [] + jax.ShapeDtypeStruct = _ShapeDtypeStruct + jnp = types.ModuleType("jax.numpy") + + def full(shape, value, dtype): + result = _ArrayMetadata(tuple(shape), dtype, label=f"full({value})") + full_calls.append((tuple(shape), value, dtype, result)) + return result + + jnp.full = full + jax.numpy = jnp + + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass_jax = types.ModuleType("cutlass.jax") + cutlass_jax.TensorSpec = _TensorSpec + + def cutlass_call(launcher, **options): + seen["call_options"] = options + + def invoke(*args): + aliases = options["input_output_aliases"] + result_metadata = options["output_shape_dtype"] + results = [None] * len(result_metadata) + for input_index, result_index in aliases.items(): + results[result_index] = args[input_index] + for result_index, metadata in enumerate(result_metadata): + if results[result_index] is None: + results[result_index] = _ArrayMetadata( + metadata.shape, + metadata.dtype, + label=f"allocated({result_index})", + ) + + unaliased_results = [results[index] for index in range(len(results)) if index not in aliases.values()] + launcher("stream", *args, *unaliased_results) + return tuple(results) + + return invoke + + cutlass_jax.cutlass_call = cutlass_call + cutlass.jax = cutlass_jax + + input_value = _ArrayMetadata((2,), "float32", label="input") + input_tensor_spec = _TensorSpec(layout=(0,), mode=(0,)) + initialized_output_spec = _TensorSpec(layout=(2, 1, 0), mode=(2, 0, 1)) + api = Adapter() + with mock.patch.dict( + sys.modules, + { + "jax": jax, + "jax.numpy": jnp, + "cutlass": cutlass, + "cutlass.jax": cutlass_jax, + }, + ): + results = api._call_kernel( + (input_value,), + input_spec=(input_tensor_spec,), + output_spec=(None, initialized_output_spec), + ) + with self.assertRaisesRegex(TypeError, "input #0 must have shape and dtype metadata"): + api._call_kernel(([input_value],)) + with self.assertRaisesRegex(TypeError, "output_spec must contain only TensorSpec or None"): + api._call_kernel( + (input_value,), + output_spec=((0,), initialized_output_spec), + ) + + self.assertEqual( + [(shape, value, dtype) for shape, value, dtype, _ in full_calls], + [ + ((2, 3, 4), float("-inf"), "float32"), + ((3,), 0, "float32"), + ], + ) + self.assertEqual(seen["call_options"]["input_output_aliases"], {1: 1, 2: 2}) + self.assertEqual(len(seen["call_options"]["input_spec"]), 3) + self.assertIs(seen["call_options"]["input_spec"][1], initialized_output_spec) + self.assertIs(seen["call_options"]["input_spec"][2], seen["call_options"]["output_spec"][2]) + + kernel_args = seen["kernel_args"] + self.assertIs(kernel_args[0], input_value) + self.assertEqual(kernel_args[1].label, "allocated(0)") + self.assertIs(kernel_args[2], full_calls[0][3]) + self.assertIs(kernel_args[3], full_calls[1][3]) + self.assertEqual(kernel_args[4].label, "allocated(3)") + self.assertEqual(kernel_args[5], "stream") + self.assertEqual([result.label for result in results], ["allocated(0)", "full(-inf)"]) + self.assertNotIn(full_calls[0][3], vars(api).values()) + self.assertNotIn(full_calls[1][3], vars(api).values()) + + def test_returns_callable_without_owning_jit_policy(self): + base = self.module.JaxApiBase + + class Adapter(base): + def check_support(self): + return True + + def __call__(self): + return self.check_support() + + api = Adapter() + api.option = 1 + self.assertIs(api.get_jax_callable(), api) + self.assertTrue(api()) + api.option = 2 + self.assertEqual(api.option, 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_datatypes.py b/test/python/fe_api/test_jax_datatypes.py new file mode 100644 index 000000000..54d39d705 --- /dev/null +++ b/test/python/fe_api/test_jax_datatypes.py @@ -0,0 +1,84 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free tests for common JAX/cuDNN dtype conversion.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_PACKAGE = "cudnn_jax_datatypes_test" + + +class _DataType(Enum): + NOT_SET = auto() + FLOAT = auto() + BFLOAT16 = auto() + + +class JaxDataTypesTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + jax = types.ModuleType("jax") + jax.__path__ = [] + jnp = types.ModuleType("jax.numpy") + jnp.bfloat16 = "bfloat16" + jnp.float32 = "float32" + jnp.dtype = lambda value: value + jax.numpy = jnp + + try: + with mock.patch.dict(sys.modules, {"jax": jax, "jax.numpy": jnp}): + cls.module = importlib.import_module(f"{internal_name}.datatypes") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def test_supported_types_round_trip(self): + self.assertEqual(self.module.jax_to_cudnn_dtype("bfloat16"), _DataType.BFLOAT16) + self.assertEqual(self.module.jax_to_cudnn_dtype("float32"), _DataType.FLOAT) + self.assertEqual(self.module.cudnn_to_jax_dtype(_DataType.BFLOAT16), "bfloat16") + self.assertEqual(self.module.cudnn_to_jax_dtype(_DataType.FLOAT), "float32") + + def test_unsupported_types_have_explicit_behavior(self): + self.assertEqual(self.module.jax_to_cudnn_dtype("unsupported"), _DataType.NOT_SET) + with self.assertRaisesRegex(ValueError, "Unsupported JAX data type"): + self.module.cudnn_to_jax_dtype(_DataType.NOT_SET) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_layout.py b/test/python/fe_api/test_jax_layout.py new file mode 100644 index 000000000..57ad2429b --- /dev/null +++ b/test/python/fe_api/test_jax_layout.py @@ -0,0 +1,89 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free tests for CUTLASS JAX layout conversion.""" + +import importlib.util +from pathlib import Path +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_LAYOUT_FILE = Path(__file__).resolve().parents[3] / "python" / "cudnn" / "_jax" / "layout.py" + + +def _load_layout_module(): + spec = importlib.util.spec_from_file_location("cudnn_jax_layout_test", _LAYOUT_FILE) + if spec is None or spec.loader is None: + raise RuntimeError("Unable to load CUTLASS JAX layout helpers") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class JaxLayoutTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.layout = _load_layout_module() + + def test_non_self_inverse_stride_order_maps_to_cutlass_axis_ranks(self): + shape = (7, 3, 5) + stride_order = (1, 2, 0) + stride = self.layout.compact_stride(shape, stride_order) + + self.assertEqual(stride, (15, 1, 3)) + self.assertEqual( + self.layout.to_cutlass_layout(shape, stride, stride_order), + (2, 0, 1), + ) + + def test_noncompact_stride_is_rejected(self): + with self.assertRaisesRegex(ValueError, "cannot represent non-compact stride"): + self.layout.to_cutlass_layout( + shape=(7, 3, 5), + stride=(16, 1, 3), + stride_order=(1, 2, 0), + name="sample", + ) + + def test_compact_stride_preserves_order_for_zero_extents(self): + self.assertEqual( + self.layout.compact_stride((2, 0, 3), (2, 1, 0)), + (3, 3, 1), + ) + + def test_identity_mode_is_the_default(self): + self.assertEqual(self.layout.normalize_mode(3), (0, 1, 2)) + self.assertEqual(self.layout.to_canonical_axes((2, 3, 4)), (2, 3, 4)) + self.assertEqual(self.layout.to_public_axes((2, 3, 4)), (2, 3, 4)) + + def test_mode_maps_kernel_axes_to_public_axes(self): + mode = (2, 0, 1) + + self.assertEqual(self.layout.to_canonical_axes((2, 3, 4), mode), (4, 2, 3)) + self.assertEqual(self.layout.to_public_axes((4, 2, 3), mode), (2, 3, 4)) + + def test_cutlass_layout_is_indexed_by_public_axes(self): + self.assertEqual( + self.layout.to_cutlass_layout( + shape=(4, 2, 3), + stride=(1, 12, 4), + stride_order=(0, 2, 1), + mode=(2, 0, 1), + ), + (2, 1, 0), + ) + + def test_invalid_mode_is_rejected(self): + with self.assertRaisesRegex(ValueError, "mode must be a permutation"): + self.layout.normalize_mode(3, (0, 0, 2)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_rmsnorm_rht_amax.py b/test/python/fe_api/test_jax_rmsnorm_rht_amax.py new file mode 100644 index 000000000..6866b405e --- /dev/null +++ b/test/python/fe_api/test_jax_rmsnorm_rht_amax.py @@ -0,0 +1,126 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX integration tests for RMSNorm + RHT + amax.""" + +import math + +import pytest + + +def _jax_runtime(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + cutlass_jax = pytest.importorskip("cutlass.jax") + if not cutlass_jax.is_available(): + pytest.skip("Installed JAX version is unsupported by CUTLASS JAX") + return jax, jnp + + +def _hadamard_16(jnp): + values = [[1.0 if ((row & column).bit_count() % 2 == 0) else -1.0 for column in range(16)] for row in range(16)] + return jnp.asarray(values, dtype=jnp.float32) / math.sqrt(16) + + +@pytest.mark.L0 +def test_jax_rmsnorm_rht_amax_abstract_contract(): + jax, jnp = _jax_runtime() + + from cudnn import OpKernel, TensorDesc + from cudnn.jax import JaxApiBase, RmsNormRhtAmaxSm100, rmsnorm_rht_amax_sm100 + from cudnn.rmsnorm_rht_amax.kernel import RMSNormRHTAmaxKernel + + sample_x = jax.ShapeDtypeStruct((256, 2048), jnp.bfloat16) + sample_weight = jax.ShapeDtypeStruct((2048,), jnp.bfloat16) + api = RmsNormRhtAmaxSm100(sample_x, sample_weight) + + assert isinstance(api, JaxApiBase) + assert isinstance(api.kernel, OpKernel) + assert isinstance(api.kernel, RMSNormRHTAmaxKernel) + assert isinstance(api.x_desc, TensorDesc) + assert isinstance(api.w_desc, TensorDesc) + assert api.kernel.x is api.x_desc + assert api.kernel.weight is api.w_desc + assert api.x_desc.dtype == jnp.dtype(jnp.bfloat16) + assert api.x_desc.stride == (2048, 1) + assert api.x_desc.stride_order == (1, 0) + assert api.x_desc.name == "sample_x" + assert all(value is not sample_x and value is not sample_weight for value in vars(api).values()) + output, amax = jax.eval_shape(api, sample_x, sample_weight) + assert output.shape == (256, 2048) + assert output.dtype == jnp.bfloat16 + assert amax.shape == (128,) + assert amax.dtype == jnp.float32 + + functional_output, functional_amax = jax.eval_shape( + lambda x, weight: rmsnorm_rht_amax_sm100( + x, + weight, + eps=1e-4, + num_threads=128, + rows_per_cta=2, + ), + sample_x, + sample_weight, + ) + assert functional_output.shape == (256, 2048) + assert functional_output.dtype == jnp.bfloat16 + assert functional_amax.shape == (128,) + assert functional_amax.dtype == jnp.float32 + + wrong_x = jax.ShapeDtypeStruct((128, 2048), jnp.bfloat16) + with pytest.raises(ValueError, match="sample_x tensor shape mismatch"): + jax.eval_shape(api, wrong_x, sample_weight) + + api.kernel.requested_rows_per_cta = 3 + with pytest.raises(ValueError, match="M must be divisible by rows_per_cta"): + jax.eval_shape(lambda x, weight: api(x, weight), sample_x, sample_weight) + + +@pytest.mark.L0 +def test_jax_rmsnorm_rht_amax_jit(): + jax, jnp = _jax_runtime() + gpu_devices = [device for device in jax.local_devices() if device.platform == "gpu"] + capable_devices = [] + for device in gpu_devices: + capability = getattr(device, "compute_capability", None) + if capability is None: + continue + major, minor = (int(value) for value in str(capability).split(".", 1)) + if major * 10 + minor >= 100: + capable_devices.append(device) + if not capable_devices: + pytest.skip("RMSNorm + RHT requires an SM100+ JAX device") + + from cudnn.jax import rmsnorm_rht_amax_sm100 + + device = capable_devices[0] + m, n = 256, 2048 + rows_per_cta = 2 + eps = 1e-4 + x = jax.device_put(jax.random.normal(jax.random.key(0), (m, n), dtype=jnp.bfloat16), device) + weight = jax.device_put(jax.random.normal(jax.random.key(1), (n,), dtype=jnp.bfloat16), device) + lowered = rmsnorm_rht_amax_sm100.lower( + x, + weight, + eps=eps, + num_threads=128, + rows_per_cta=rows_per_cta, + ) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + output, amax = lowered.compile()(x, weight) + output.block_until_ready() + x_f32 = x.astype(jnp.float32) + normalized = x_f32 * jax.lax.rsqrt(jnp.mean(jnp.square(x_f32), axis=-1, keepdims=True) + eps) + normalized *= weight.astype(jnp.float32)[None, :] + reference = (normalized.reshape(m, n // 16, 16) @ _hadamard_16(jnp)).reshape(m, n) + amax_reference = jnp.max( + jnp.abs(reference).reshape(m // rows_per_cta, rows_per_cta, n), + axis=(1, 2), + ) + + assert jnp.allclose(output.astype(jnp.float32), reference, atol=4e-2, rtol=1e-2) + assert jnp.allclose(amax, amax_reference, atol=2e-3, rtol=1e-3) diff --git a/test/python/fe_api/test_rmsnorm_rht_amax.py b/test/python/fe_api/test_rmsnorm_rht_amax.py index c00d04b02..0128606f6 100644 --- a/test/python/fe_api/test_rmsnorm_rht_amax.py +++ b/test/python/fe_api/test_rmsnorm_rht_amax.py @@ -61,7 +61,8 @@ def _assert_ref_close(x, w, o, amax, *, eps: float, rows_per_cta: int, skip_ref: @pytest.mark.parametrize("n,num_threads", SUPPORTED_N_NUM_THREADS) def test_rmsnorm_rht_amax_compile_execute(n, num_threads, request): try: - from cudnn import RmsNormRhtAmaxSm100 + from cudnn import OpKernel, RmsNormRhtAmaxSm100, TensorDesc + from cudnn.rmsnorm_rht_amax.kernel import RMSNormRHTAmaxKernel except ImportError: pytest.skip("Environment not supported: cudnn optional dependencies not installed") @@ -82,12 +83,27 @@ def test_rmsnorm_rht_amax_compile_execute(n, num_threads, request): num_threads=num_threads, rows_per_cta=rows_per_cta, ) + assert isinstance(api.kernel, OpKernel) + assert isinstance(api.kernel, RMSNormRHTAmaxKernel) + assert isinstance(api.kernel.x, TensorDesc) + assert isinstance(api.kernel.weight, TensorDesc) + assert api.kernel.x is api.x_desc + assert api.kernel.weight is api.w_desc + assert api.kernel.output is api.o_desc + assert api.kernel.amax is api.amax_desc + assert api.eps == eps + assert api.requested_num_threads == num_threads + assert api.requested_rows_per_cta == rows_per_cta try: assert api.check_support(), "Unsupported testcase" except (ValueError, RuntimeError) as exc: pytest.skip(f"Unsupported testcase: {exc}") + assert api.n == n + assert api.num_threads == num_threads + assert api.rows_per_cta == rows_per_cta + api.compile() api.execute(x_tensor=x, w_tensor=w, o_tensor=o, amax_tensor=amax) _assert_ref_close(x, w, o, amax, eps=eps, rows_per_cta=rows_per_cta, skip_ref=skip_ref) @@ -122,3 +138,38 @@ def test_rmsnorm_rht_amax_wrapper(n, num_threads, rows_per_cta, request): assert outputs["o_tensor"].shape == (m, n) assert outputs["amax_tensor"].shape == (m // rows_per_cta,) _assert_ref_close(x, w, outputs["o_tensor"], outputs["amax_tensor"], eps=eps, rows_per_cta=rows_per_cta, skip_ref=skip_ref) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_rmsnorm_rht_amax_wrapper_reuses_owned_kernel(monkeypatch): + try: + from cudnn.rmsnorm_rht_amax import api as rmsnorm_api + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + construction_count = 0 + original_init = rmsnorm_api.RMSNormRHTAmaxKernel.__init__ + + def counted_init(self, *args, **kwargs): + nonlocal construction_count + construction_count += 1 + original_init(self, *args, **kwargs) + + monkeypatch.setattr(rmsnorm_api.RMSNormRHTAmaxKernel, "__init__", counted_init) + rmsnorm_api._cache_of_RmsNormRhtAmaxSm100Objects.clear() + + m, n = 256, 2048 + rows_per_cta = 2 + x, w = _make_inputs(m=m, n=n) + try: + first = rmsnorm_api.rmsnorm_rht_amax_wrapper_sm100(x, w, num_threads=128, rows_per_cta=rows_per_cta) + second = rmsnorm_api.rmsnorm_rht_amax_wrapper_sm100(x, w, num_threads=128, rows_per_cta=rows_per_cta) + except (ValueError, RuntimeError) as exc: + pytest.skip(f"Unsupported testcase: {exc}") + finally: + rmsnorm_api._cache_of_RmsNormRhtAmaxSm100Objects.clear() + + assert construction_count == 1 + assert first["o_tensor"].shape == second["o_tensor"].shape == (m, n) + assert first["amax_tensor"].shape == second["amax_tensor"].shape == (m // rows_per_cta,) diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_imports.py b/test/python/fe_api/test_rmsnorm_rht_amax_imports.py new file mode 100644 index 000000000..9c4b2eb49 --- /dev/null +++ b/test/python/fe_api/test_rmsnorm_rht_amax_imports.py @@ -0,0 +1,298 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Import-boundary contracts for the RMSNorm framework adapters.""" + +import ast +import importlib +import importlib.util +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "rmsnorm_rht_amax" + + +def _load_cudnn_submodule(name: str, submodule: str): + package = types.ModuleType(name) + package.__path__ = [str(_CUDNN_ROOT)] + package.__package__ = name + sys.modules[name] = package + return importlib.import_module(f"{name}.{submodule}") + + +def _load_operation_package(name: str): + parent_name, separator, _ = name.rpartition(".") + if not separator: + raise ValueError("Synthetic operation package must have a parent") + parent = types.ModuleType(parent_name) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = parent_name + sys.modules[parent_name] = parent + + spec = importlib.util.spec_from_file_location( + name, + _OPERATION_ROOT / "__init__.py", + submodule_search_locations=[str(_OPERATION_ROOT)], + ) + if spec is None or spec.loader is None: + raise RuntimeError("Unable to load RMSNorm package") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _remove_package(name: str) -> None: + root_name = name.split(".", 1)[0] + for module_name in tuple(sys.modules): + if module_name == root_name or module_name.startswith(f"{root_name}."): + sys.modules.pop(module_name, None) + + +class RmsNormImportContractTest(unittest.TestCase): + def test_internal_jax_package_reports_missing_jax_dependency(self): + name = "cudnn_jax_missing_dependency_test" + try: + with mock.patch.dict(sys.modules, {"jax": None}): + with self.assertRaisesRegex(ImportError, r"nvidia-cudnn-frontend\[jax\]"): + _load_cudnn_submodule(name, "_jax") + finally: + _remove_package(name) + + def test_internal_jax_package_reports_missing_cutlass_dependency(self): + name = "cudnn_jax_missing_cutlass_test" + jax = types.ModuleType("jax") + try: + with mock.patch.dict(sys.modules, {"jax": jax, "cutlass": None, "cutlass.jax": None}): + with self.assertRaisesRegex(ImportError, r"nvidia-cudnn-frontend\[jax\]"): + _load_cudnn_submodule(name, "_jax") + finally: + _remove_package(name) + + def test_internal_jax_package_reports_incompatible_jax(self): + name = "cudnn_jax_incompatible_test" + jax = types.ModuleType("jax") + jax.__version__ = "0.8.0" + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass_jax = types.ModuleType("cutlass.jax") + cutlass_jax.is_available = lambda: False + cutlass_jax.CUTE_DSL_MIN_SUPPORTED_JAX_VERSION = (0, 9, 1) + cutlass.jax = cutlass_jax + try: + with mock.patch.dict(sys.modules, {"jax": jax, "cutlass": cutlass, "cutlass.jax": cutlass_jax}): + with self.assertRaisesRegex(ImportError, r"JAX 0\.8\.0.*minimum supported JAX version is 0\.9\.1.*nvidia-cudnn-frontend\[jax\]"): + _load_cudnn_submodule(name, "_jax") + finally: + _remove_package(name) + + def test_jax_facade_uses_internal_dependency_boundary(self): + name = "cudnn_jax_facade_dependency_test" + try: + with mock.patch.dict(sys.modules, {"jax": None}): + with self.assertRaisesRegex(ImportError, r"nvidia-cudnn-frontend\[jax\]"): + _load_cudnn_submodule(name, "jax") + finally: + _remove_package(name) + + def test_jax_operator_uses_internal_dependency_boundary(self): + name = "cudnn_jax_operator_dependency_test.rmsnorm_rht_amax" + try: + package = _load_operation_package(name) + with mock.patch.dict(sys.modules, {"jax": None}): + with self.assertRaisesRegex(ImportError, r"nvidia-cudnn-frontend\[jax\]"): + package.jax + finally: + _remove_package(name) + + def test_package_import_loads_no_framework_adapter(self): + name = "cudnn_rmsnorm_import_test.rmsnorm_rht_amax" + blocked = {module_name: None for module_name in ("torch", "jax", "cutlass", "cuda")} + try: + with mock.patch.dict(sys.modules, blocked): + package = _load_operation_package(name) + self.assertNotIn(f"{name}.api", sys.modules) + self.assertNotIn(f"{name}.jax", sys.modules) + self.assertNotIn(f"{name}.kernel", sys.modules) + self.assertTrue( + { + "RMSNormRHTAmaxKernel", + "RmsNormRhtAmaxSm100", + "api", + "jax", + "kernel", + }.issubset(dir(package)) + ) + finally: + _remove_package(name) + + def test_unqualified_exports_route_only_to_torch_api(self): + name = "cudnn_rmsnorm_torch_route_test.rmsnorm_rht_amax" + sentinel = object() + api = types.ModuleType(f"{name}.api") + api.RmsNormRhtAmaxSm100 = sentinel + try: + sys.modules[api.__name__] = api + package = _load_operation_package(name) + self.assertIs(package.RmsNormRhtAmaxSm100, sentinel) + self.assertNotIn(f"{name}.jax", sys.modules) + finally: + _remove_package(name) + + def test_kernel_exports_do_not_load_an_adapter(self): + name = "cudnn_rmsnorm_kernel_route_test.rmsnorm_rht_amax" + sentinel = object() + kernel = types.ModuleType(f"{name}.kernel") + kernel.RMSNormRHTAmaxKernel = sentinel + try: + sys.modules[kernel.__name__] = kernel + package = _load_operation_package(name) + self.assertIs(package.RMSNormRHTAmaxKernel, sentinel) + self.assertNotIn(f"{name}.api", sys.modules) + self.assertNotIn(f"{name}.jax", sys.modules) + finally: + _remove_package(name) + + def test_jax_adapter_is_explicit(self): + name = "cudnn_rmsnorm_jax_route_test.rmsnorm_rht_amax" + jax_module = types.ModuleType(f"{name}.jax") + try: + sys.modules[jax_module.__name__] = jax_module + package = _load_operation_package(name) + self.assertIs(package.jax, jax_module) + self.assertNotIn(f"{name}.api", sys.modules) + finally: + _remove_package(name) + + def test_static_import_directions(self): + def all_imports(filename): + path = Path(filename) + if not path.is_absolute(): + path = _OPERATION_ROOT / path + tree = ast.parse(path.read_text(), filename=str(path)) + return [node for node in ast.walk(tree) if isinstance(node, (ast.Import, ast.ImportFrom))] + + def imports_framework(node, framework: str): + if isinstance(node, ast.Import): + return any(alias.name == framework or alias.name.startswith(f"{framework}.") for alias in node.names) + return node.level == 0 and (node.module == framework or (node.module or "").startswith(f"{framework}.")) + + op_kernel_imports = all_imports(_CUDNN_ROOT / "_op_kernel.py") + self.assertFalse(any(imports_framework(node, framework) for node in op_kernel_imports for framework in ("torch", "jax", "cutlass", "cuda"))) + + kernel_imports = all_imports("kernel.py") + self.assertFalse(any(imports_framework(node, framework) for node in kernel_imports for framework in ("torch", "jax"))) + + api_imports = all_imports("api.py") + self.assertFalse(any(imports_framework(node, "jax") for node in api_imports)) + + jax_imports = all_imports("jax.py") + self.assertFalse(any(imports_framework(node, "torch") for node in jax_imports)) + + api_tree = ast.parse((_OPERATION_ROOT / "api.py").read_text(), filename="api.py") + jax_tree = ast.parse((_OPERATION_ROOT / "jax.py").read_text(), filename="jax.py") + jax_package_tree = ast.parse((_CUDNN_ROOT / "_jax" / "__init__.py").read_text(), filename="_jax/__init__.py") + jax_facade_tree = ast.parse((_CUDNN_ROOT / "jax" / "__init__.py").read_text(), filename="jax/__init__.py") + jax_base_tree = ast.parse((_CUDNN_ROOT / "_jax" / "api_base.py").read_text(), filename="_jax/api_base.py") + kernel_tree = ast.parse((_OPERATION_ROOT / "kernel.py").read_text(), filename="kernel.py") + torch_adapter = next(node for node in api_tree.body if isinstance(node, ast.ClassDef) and node.name == "RmsNormRhtAmaxSm100") + self.assertEqual([base.id for base in torch_adapter.bases if isinstance(base, ast.Name)], ["APIBase"]) + constructor = next(node for node in torch_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "__init__") + required_argument_count = len(constructor.args.args) - len(constructor.args.defaults) + self.assertEqual( + [argument.arg for argument in constructor.args.args[:required_argument_count]], + ["self", "sample_x", "sample_w"], + ) + self.assertNotIn( + "_kernel", + [argument.arg for argument in (*constructor.args.args, *constructor.args.kwonlyargs)], + ) + kernel_constructors = [ + node for node in ast.walk(api_tree) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RMSNormRHTAmaxKernel" + ] + self.assertEqual(len(kernel_constructors), 1) + self.assertIn(kernel_constructors[0], ast.walk(constructor)) + + jax_adapter = next(node for node in jax_tree.body if isinstance(node, ast.ClassDef) and node.name == "RmsNormRhtAmaxSm100") + self.assertEqual([base.id for base in jax_adapter.bases if isinstance(base, ast.Name)], ["JaxApiBase"]) + self.assertFalse(jax_adapter.decorator_list) + self.assertTrue( + any(isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RMSNormRHTAmaxKernel" for node in ast.walk(jax_adapter)) + ) + + jax_wrapper = next(node for node in jax_tree.body if isinstance(node, ast.FunctionDef) and node.name == "rmsnorm_rht_amax_sm100") + jit_decorator = next( + decorator + for decorator in jax_wrapper.decorator_list + if isinstance(decorator, ast.Call) and isinstance(decorator.func, ast.Attribute) and decorator.func.attr == "jit" + ) + static_argnames = next(keyword.value for keyword in jit_decorator.keywords if keyword.arg == "static_argnames") + self.assertEqual( + tuple(element.value for element in static_argnames.elts), + ("eps", "num_threads", "rows_per_cta"), + ) + + torch_wrapper = next(node for node in api_tree.body if isinstance(node, ast.FunctionDef) and node.name == "rmsnorm_rht_amax_wrapper_sm100") + self.assertFalse( + any(isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RMSNormRHTAmaxKernel" for node in ast.walk(torch_wrapper)) + ) + self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "_materialize_outputs" for node in ast.walk(torch_wrapper))) + + kernel_class = next(node for node in kernel_tree.body if isinstance(node, ast.ClassDef) and node.name == "RMSNormRHTAmaxKernel") + self.assertEqual([base.id for base in kernel_class.bases if isinstance(base, ast.Name)], ["OpKernel"]) + launcher = next(node for node in kernel_class.body if isinstance(node, ast.FunctionDef) and node.name == "__call__") + self.assertEqual( + [argument.arg for argument in launcher.args.args], + ["self", "x_tensor", "w_tensor", "o_tensor", "amax_tensor", "stream"], + ) + + compile_method = next(node for node in torch_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "compile") + tensor_api = next(node for node in ast.walk(compile_method) if isinstance(node, ast.FunctionDef) and node.name == "tensor_api") + compiled_call = next( + node for node in ast.walk(tensor_api) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "compiled_kernel" + ) + self.assertEqual( + [argument.id for argument in compiled_call.args if isinstance(argument, ast.Name)], + ["x_tensor", "w_tensor", "o_tensor", "amax_tensor", "stream"], + ) + + self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "_call_kernel" for node in ast.walk(jax_adapter))) + self.assertFalse( + any(isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "cutlass_call" for node in ast.walk(jax_adapter)) + ) + self.assertIn("_INSTALL_HINT", {node.id for node in ast.walk(jax_package_tree) if isinstance(node, ast.Name)}) + self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "is_available" for node in ast.walk(jax_package_tree))) + for tree in (jax_tree, jax_facade_tree): + self.assertNotIn("_INSTALL_HINT", {node.id for node in ast.walk(tree) if isinstance(node, ast.Name)}) + self.assertFalse(any(isinstance(node, ast.Try) for node in tree.body)) + self.assertFalse(any(isinstance(node, ast.Attribute) and node.attr == "is_available" for node in ast.walk(tree))) + jax_base = next(node for node in jax_base_tree.body if isinstance(node, ast.ClassDef) and node.name == "JaxApiBase") + base_call = next(node for node in jax_base.body if isinstance(node, ast.FunctionDef) and node.name == "_call_kernel") + cutlass_call = next( + node for node in ast.walk(base_call) if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "cutlass_call" + ) + self.assertTrue(cutlass_call.args) + self.assertNotIn("_call_cache", {node.attr for node in ast.walk(jax_adapter) if isinstance(node, ast.Attribute)}) + self.assertNotIn("_trace_lock", {node.attr for node in ast.walk(jax_adapter) if isinstance(node, ast.Attribute)}) + + root_source = (_CUDNN_ROOT / "__init__.py").read_text() + self.assertIn("from ._op_kernel import OpKernel", root_source) + self.assertIn("from ._tensor_desc import TensorDesc", root_source) + self.assertIn('if name == "jax":', root_source) + self.assertIn('importlib.import_module(".jax", __name__)', root_source) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py b/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py new file mode 100644 index 000000000..eb89130ff --- /dev/null +++ b/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py @@ -0,0 +1,280 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for the RMSNorm operation kernel.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "rmsnorm_rht_amax" +_PACKAGE = "cudnn_rmsnorm_kernel_test" + + +class _DataType(Enum): + NOT_SET = auto() + FLOAT = auto() + BFLOAT16 = auto() + + +def _identity_decorator(function=None, **_kwargs): + if function is None: + return lambda decorated: decorated + return function + + +def _module(name: str, *, package: bool = False, **attributes): + module = types.ModuleType(name) + if package: + module.__path__ = [] + for attribute, value in attributes.items(): + setattr(module, attribute, value) + return module + + +def _cutlass_stubs() -> dict[str, types.ModuleType]: + driver = _module("cuda.bindings.driver", CUstream=object) + bindings = _module("cuda.bindings", package=True, driver=driver) + cuda = _module("cuda", package=True, bindings=bindings) + + arch = _module("cutlass.cute.arch", shuffle_sync_bfly=lambda value, **_kwargs: value) + cute = _module( + "cutlass.cute", + package=True, + arch=arch, + kernel=_identity_decorator, + jit=_identity_decorator, + ) + utils = _module("cutlass.utils") + llvm = _module( + "cutlass._mlir.dialects.llvm", + AsmDialect=types.SimpleNamespace(AD_ATT=object()), + ) + dialects = _module("cutlass._mlir.dialects", package=True, llvm=llvm) + mlir = _module("cutlass._mlir", package=True, dialects=dialects) + cutlass_dsl = _module( + "cutlass.cutlass_dsl", + T=types.SimpleNamespace(f32=lambda: None), + dsl_user_op=_identity_decorator, + ) + cutlass = _module( + "cutlass", + package=True, + cute=cute, + utils=utils, + Float32=object, + Int32=object, + BFloat16=object, + ) + + return { + module.__name__: module + for module in ( + cuda, + bindings, + driver, + cutlass, + cute, + arch, + utils, + mlir, + dialects, + llvm, + cutlass_dsl, + ) + } + + +def _install_test_package() -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + operation_name = f"{_PACKAGE}.rmsnorm_rht_amax" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + +def _remove_test_package() -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + +class RmsNormRhtAmaxKernelContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + _install_test_package() + try: + with mock.patch.dict(sys.modules, _cutlass_stubs()): + cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}._op_kernel") + cls.kernel_module = importlib.import_module(f"{_PACKAGE}.rmsnorm_rht_amax.kernel") + except Exception: + _remove_test_package() + raise + + @classmethod + def tearDownClass(cls) -> None: + _remove_test_package() + + def _desc(self, shape, dtype=_DataType.BFLOAT16, *, name=""): + shape = tuple(shape) + stride_order = tuple(reversed(range(len(shape)))) + stride = [0] * len(shape) + running = 1 + for dimension in stride_order: + stride[dimension] = running + running *= shape[dimension] + return self.tensor_module.TensorDesc( + dtype=dtype, + shape=shape, + stride=tuple(stride), + stride_order=stride_order, + name=name, + ) + + def _kernel(self, **overrides): + arguments = { + "x": self._desc((256, 2048)), + "weight": self._desc((2048,)), + } + arguments.update(overrides) + return self.kernel_module.RMSNormRHTAmaxKernel(**arguments) + + def test_kernel_implements_neutral_contract_and_infers_buffers(self): + kernel = self._kernel() + self.assertIsInstance(kernel, self.base_module.OpKernel) + with self.assertRaisesRegex(RuntimeError, r"check_support\(\)"): + kernel.infer_output() + + self.assertTrue(kernel.check_support()) + self.assertEqual((kernel.m, kernel.n), (256, 2048)) + self.assertEqual((kernel.num_threads, kernel.rows_per_cta), (128, 2)) + + output, amax = kernel.infer_output() + self.assertIs(type(output), self.tensor_module.TensorDesc) + self.assertEqual( + (output.name, output.dtype, output.shape, output.stride, output.stride_order, output.init_value), + ("output", _DataType.BFLOAT16, (256, 2048), (2048, 1), (1, 0), None), + ) + self.assertEqual( + (amax.name, amax.dtype, amax.shape, amax.stride, amax.stride_order, amax.init_value), + ("amax", _DataType.FLOAT, (128,), (1,), (0,), None), + ) + self.assertEqual(kernel.infer_workspace(), ()) + + self.assertNotIn(f"{_PACKAGE}.rmsnorm_rht_amax.api", sys.modules) + self.assertNotIn(f"{_PACKAGE}.rmsnorm_rht_amax.jax", sys.modules) + + def test_requested_configuration_is_recomputed(self): + kernel = self._kernel(num_threads=256, rows_per_cta=4) + kernel.check_support() + self.assertEqual((kernel.num_threads, kernel.rows_per_cta), (256, 4)) + + kernel.requested_num_threads = 128 + kernel.requested_rows_per_cta = 8 + kernel.amax = self._desc((32,), _DataType.FLOAT) + kernel.check_support() + self.assertEqual((kernel.num_threads, kernel.rows_per_cta), (128, 8)) + self.assertEqual(kernel.infer_output()[1].shape, (32,)) + + def test_framework_descriptor_extension_maps_native_dtype(self): + class FrameworkTensorDesc(self.tensor_module.TensorDesc): + @property + def cudnn_dtype(self): + return { + "framework.bfloat16": _DataType.BFLOAT16, + "framework.float32": _DataType.FLOAT, + }.get(self.dtype, _DataType.NOT_SET) + + def desc(shape, dtype): + base = self._desc(shape) + return FrameworkTensorDesc( + dtype=dtype, + shape=base.shape, + stride=base.stride, + stride_order=base.stride_order, + name="framework_tensor", + ) + + kernel = self.kernel_module.RMSNormRHTAmaxKernel( + x=desc((256, 2048), "framework.bfloat16"), + weight=desc((2048,), "framework.bfloat16"), + output=desc((256, 2048), "framework.bfloat16"), + amax=desc((128,), "framework.float32"), + ) + self.assertTrue(kernel.check_support()) + + def test_invalid_logical_signatures_and_configuration(self): + cases = ( + ({"x": self._desc((2048,))}, "X must have rank 2"), + ({"weight": self._desc((1, 2048))}, "W must have rank 1"), + ({"x": self._desc((256, 2048), _DataType.FLOAT)}, "X must have dtype bfloat16"), + ({"weight": self._desc((2048,), _DataType.FLOAT)}, "W must have dtype bfloat16"), + ({"x": self._desc((0, 2048))}, "M must be positive"), + ( + {"x": self._desc((256, 2049)), "weight": self._desc((2049,))}, + "N must be divisible by 16", + ), + ({"weight": self._desc((1024,))}, "W must have shape"), + ( + { + "x": self.tensor_module.TensorDesc( + dtype=_DataType.BFLOAT16, + shape=(256, 2048), + stride=(1, 256), + stride_order=(0, 1), + ) + }, + "X must be row-major contiguous", + ), + ({"num_threads": 512}, "EPT=4 must be >= 8 and divisible by 8"), + ({"num_threads": 2048}, "must not exceed the CUDA block size limit"), + ({"rows_per_cta": 0}, "rows_per_cta must be positive"), + ({"rows_per_cta": 3}, "M must be divisible by rows_per_cta"), + ({"output": self._desc((256, 1024))}, "O must have shape"), + ({"output": self._desc((256, 2048), _DataType.FLOAT)}, "O must have dtype bfloat16"), + ({"amax": self._desc((64,), _DataType.FLOAT)}, "Amax must have shape"), + ({"amax": self._desc((128,))}, "Amax must have dtype float32"), + ) + + for overrides, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self._kernel(**overrides).check_support() + + with self.assertRaisesRegex(TypeError, "x must be a TensorDesc"): + self.kernel_module.RMSNormRHTAmaxKernel( + x=types.SimpleNamespace(shape=(256, 2048)), + weight=self._desc((2048,)), + ) + with self.assertRaisesRegex(TypeError, "weight must be a TensorDesc"): + self.kernel_module.RMSNormRHTAmaxKernel( + x=self._desc((256, 2048)), + weight=None, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_shared_tensor_desc.py b/test/python/fe_api/test_shared_tensor_desc.py new file mode 100644 index 000000000..332e91025 --- /dev/null +++ b/test/python/fe_api/test_shared_tensor_desc.py @@ -0,0 +1,159 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free tests for framework-neutral tensor descriptors.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_PACKAGE = "cudnn_tensor_desc_test" + + +class _DataType(Enum): + FLOAT = auto() + BFLOAT16 = auto() + + +def _load_tensor_desc_module(): + package = types.ModuleType(_PACKAGE) + package.__path__ = [str(_CUDNN_ROOT)] + package.__package__ = _PACKAGE + package.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + package.data_type = _DataType + sys.modules[_PACKAGE] = package + return importlib.import_module(f"{_PACKAGE}._tensor_desc") + + +class SharedTensorDescTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.tensor_desc = _load_tensor_desc_module() + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def test_compact_descriptor_defaults_to_row_major(self): + desc = self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.BFLOAT16, + shape=(2, 3, 4), + name="output", + ) + + self.assertIsInstance(desc, self.tensor_desc.TensorDesc) + self.assertEqual(desc.dtype, _DataType.BFLOAT16) + self.assertEqual(desc.cudnn_dtype, _DataType.BFLOAT16) + self.assertEqual(desc.shape, (2, 3, 4)) + self.assertEqual(desc.stride, (12, 4, 1)) + self.assertEqual(desc.stride_order, (2, 1, 0)) + self.assertEqual(desc.ndim, 3) + self.assertEqual(desc.name, "output") + self.assertIsNone(desc.init_value) + + def test_compact_descriptor_preserves_initial_value(self): + for init_value in (0, False, float("-inf")): + with self.subTest(init_value=init_value): + desc = self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2, 3), + init_value=init_value, + ) + + if init_value is False: + self.assertIs(desc.init_value, False) + else: + self.assertEqual(desc.init_value, init_value) + + with self.assertRaisesRegex(TypeError, "init_value must be a real scalar or None"): + self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2, 3), + init_value=[], + ) + + def test_compact_descriptor_supports_an_explicit_dimension_order(self): + desc = self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(7, 3, 5), + stride_order=(1, 2, 0), + ) + + self.assertEqual(desc.stride, (15, 1, 3)) + self.assertEqual(desc.stride_order, (1, 2, 0)) + + def test_compact_descriptor_supports_scalars(self): + desc = self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(), + ) + + self.assertEqual(desc.shape, ()) + self.assertEqual(desc.stride, ()) + self.assertEqual(desc.stride_order, ()) + self.assertEqual(desc.ndim, 0) + + def test_compact_descriptor_preserves_order_for_zero_extents(self): + desc = self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2, 0, 3), + ) + + self.assertEqual(desc.stride, (3, 3, 1)) + self.assertEqual(desc.stride_order, (2, 1, 0)) + + def test_compact_descriptor_rejects_noncanonical_dtype(self): + with self.assertRaisesRegex(TypeError, "dtype must be a cudnn.data_type"): + self.tensor_desc.make_compact_tensor_desc(dtype="float32", shape=(4,)) + + def test_compact_descriptor_rejects_invalid_shapes(self): + for shape, message in ( + ((-1, 4), "must be non-negative"), + ((2.5, 4), "must be integers"), + ((True, 4), "must be integers"), + ): + with self.subTest(shape=shape): + with self.assertRaisesRegex((TypeError, ValueError), message): + self.tensor_desc.make_compact_tensor_desc(dtype=_DataType.FLOAT, shape=shape) + + def test_compact_descriptor_rejects_invalid_dimension_orders(self): + for stride_order, message in ( + ((1,), "rank mismatch"), + ((0, 0), "must be a permutation"), + ((0, True), "must be integers"), + ): + with self.subTest(stride_order=stride_order): + with self.assertRaisesRegex((TypeError, ValueError), message): + self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2, 3), + stride_order=stride_order, + ) + + def test_descriptor_rejects_stride_order_that_contradicts_strides(self): + with self.assertRaisesRegex(ValueError, "is inconsistent with stride"): + self.tensor_desc.TensorDesc( + dtype=_DataType.FLOAT, + shape=(2, 3), + stride=(3, 1), + stride_order=(0, 1), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_tensor_desc.py b/test/python/fe_api/test_tensor_desc.py new file mode 100644 index 000000000..2c43e7f68 --- /dev/null +++ b/test/python/fe_api/test_tensor_desc.py @@ -0,0 +1,153 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Compatibility tests for shared and Torch tensor descriptors.""" + +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +try: + import torch + + from cudnn import TensorDesc as SharedTensorDesc, data_type + from cudnn.api_base import APIBase, TensorDesc as TorchTensorDesc +except ImportError as error: + _IMPORT_ERROR = error +else: + _IMPORT_ERROR = None + + +@unittest.skipIf(_IMPORT_ERROR is not None, f"Optional FE dependencies unavailable: {_IMPORT_ERROR}") +class TensorDescTest(unittest.TestCase): + def _assert_torch_metadata(self, desc, *, name="sample", init_value=None): + self.assertIsInstance(desc, SharedTensorDesc) + self.assertEqual(desc.dtype, torch.bfloat16) + self.assertEqual(desc.device, torch.device("cpu")) + self.assertFalse(desc.interpret_uint8_as_fp4x2) + self.assertEqual(desc.name, name) + self.assertEqual(desc.init_value, init_value) + self.assertEqual(desc.ndim, len(desc.shape)) + self.assertEqual(desc.cudnn_dtype, data_type.BFLOAT16) + + def test_torch_descriptor_preserves_constructor_and_layout_operations(self): + desc = TorchTensorDesc( + torch.bfloat16, + (2, 3, 4), + (12, 4, 1), + (2, 1, 0), + "cpu", + False, + "sample", + ) + self._assert_torch_metadata(desc) + self.assertEqual(desc.size(), torch.Size((2, 3, 4))) + + transformed = ( + desc.transpose(0, 1), + desc.unsqueeze(-1), + desc.unsqueeze(-1).squeeze(-1), + desc.view(6, 4), + desc.as_strided((2, 3, 4), (12, 4, 1)), + desc.transpose(0, 1).contiguous(), + ) + for result in transformed: + with self.subTest(shape=result.shape, stride=result.stride): + self._assert_torch_metadata(result) + + def test_packed_uint8_maps_to_canonical_fp4(self): + desc = TorchTensorDesc( + dtype=torch.uint8, + shape=(16,), + stride=(1,), + stride_order=(0,), + device="cpu", + interpret_uint8_as_fp4x2=True, + name="packed_fp4", + ) + self.assertEqual(desc.cudnn_dtype, data_type.FP4_E2M1) + + def test_layout_operations_preserve_initial_value(self): + desc = TorchTensorDesc( + dtype=torch.float32, + shape=(2, 3, 4), + stride=(12, 4, 1), + stride_order=(2, 1, 0), + device="cpu", + name="workspace", + init_value=float("-inf"), + ) + + transformed = ( + desc.permute(1, 0, 2), + desc.transpose(0, 1), + desc.unsqueeze(-1), + desc.unsqueeze(-1).squeeze(-1), + desc.view(6, 4), + desc.as_strided((2, 3, 4), (12, 4, 1)), + desc.transpose(0, 1).contiguous(), + ) + for result in transformed: + with self.subTest(shape=result.shape, stride=result.stride): + self.assertEqual(result.init_value, float("-inf")) + + def test_api_base_converts_torch_tensor_metadata(self): + tensor = torch.empty_strided((2, 3, 4), (12, 1, 3), dtype=torch.bfloat16) + desc = APIBase._to_tensor_desc(tensor, "sample") + + self._assert_torch_metadata(desc) + self.assertEqual(desc.shape, (2, 3, 4)) + self.assertEqual(desc.stride, (12, 1, 3)) + self.assertEqual(desc.stride_order, (1, 2, 0)) + + def test_api_base_materializes_a_canonical_descriptor(self): + desc = SharedTensorDesc( + dtype=data_type.FLOAT, + shape=(2, 3), + stride=(1, 2), + stride_order=(0, 1), + name="output", + init_value=float("-inf"), + ) + + tensor = APIBase._materialize_tensor_desc(desc, device=torch.device("cpu")) + + self.assertEqual(tensor.dtype, torch.float32) + self.assertEqual(tuple(tensor.shape), desc.shape) + self.assertEqual(tuple(tensor.stride()), desc.stride) + self.assertTrue(torch.isneginf(tensor).all()) + + def test_api_base_initializes_on_the_requested_cuda_stream(self): + if not torch.cuda.is_available(): + self.skipTest("CUDA is unavailable") + + from cuda.bindings import driver as cuda + + device = torch.device("cuda", torch.cuda.current_device()) + stream = torch.cuda.Stream(device=device) + desc = SharedTensorDesc( + dtype=data_type.FLOAT, + shape=(8,), + stride=(1,), + stride_order=(0,), + init_value=7, + ) + + tensor = APIBase._materialize_tensor_desc( + desc, + device=device, + stream=cuda.CUstream(stream.cuda_stream), + ) + stream.synchronize() + + self.assertTrue(torch.equal(tensor.cpu(), torch.full((8,), 7.0))) + + +if __name__ == "__main__": + unittest.main() From 6581a363e8da066db79a780115cb655764642b8e Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Tue, 7 Jul 2026 11:12:12 -0500 Subject: [PATCH 02/13] Align RMSNorm framework APIs around inferred outputs --- docs/fe-oss-apis/rmsnorm_rht_amax.md | 44 +++-- python/cudnn/_jax/__init__.py | 30 +--- python/cudnn/_jax/api_base.py | 30 ++++ python/cudnn/_op_kernel.py | 10 +- python/cudnn/_tensor_desc.py | 26 ++- python/cudnn/api_base.py | 111 +++++++++++-- python/cudnn/jax/__init__.py | 23 ++- python/cudnn/rmsnorm_rht_amax/__init__.py | 7 +- python/cudnn/rmsnorm_rht_amax/api.py | 151 ++++-------------- python/cudnn/rmsnorm_rht_amax/jax.py | 4 +- python/cudnn/rmsnorm_rht_amax/kernel.py | 58 +++---- test/python/fe_api/test_jax_api_base.py | 23 +++ test/python/fe_api/test_rmsnorm_rht_amax.py | 88 +++++++--- .../fe_api/test_rmsnorm_rht_amax_imports.py | 101 ++++++++---- .../fe_api/test_rmsnorm_rht_amax_kernel.py | 48 +++++- test/python/fe_api/test_shared_tensor_desc.py | 25 +++ test/python/fe_api/test_tensor_desc.py | 125 +++++++++++++++ 17 files changed, 615 insertions(+), 289 deletions(-) diff --git a/docs/fe-oss-apis/rmsnorm_rht_amax.md b/docs/fe-oss-apis/rmsnorm_rht_amax.md index 720331182..0f3a6d666 100644 --- a/docs/fe-oss-apis/rmsnorm_rht_amax.md +++ b/docs/fe-oss-apis/rmsnorm_rht_amax.md @@ -81,23 +81,18 @@ from cudnn import RmsNormRhtAmaxSm100 op = RmsNormRhtAmaxSm100( sample_x=x, sample_w=w, - sample_o=o, - sample_amax=amax, eps=1e-5, num_threads=128, rows_per_cta=2, ) -assert op.check_support() -op.compile() -op.execute( - x_tensor=x, - w_tensor=w, - o_tensor=o, - amax_tensor=amax, - current_stream=None, -) +o, amax = op(x, w, current_stream=None) ``` +The class infers and allocates its outputs from the validated sample input +descriptors. Runtime inputs must match the sample shape, layout, dtype, and +device. `check_support()` and `compile()` remain available when the caller +wants to run those steps explicitly; the call operator performs them lazily. + ### JAX API The CUDA 13 JAX packages require Python 3.11 or newer. Install the JAX integration explicitly: @@ -120,13 +115,15 @@ o, amax = rmsnorm_rht_amax_sm100( The function API is JIT-compiled and treats `eps`, `num_threads`, and `rows_per_cta` as static compilation options. Use `RmsNormRhtAmaxSm100` -directly when the caller needs to control JAX transformations explicitly. +directly when the caller needs to control JAX transformations explicitly. Its +class API follows the same `RmsNormRhtAmaxSm100(sample_x, sample_w)(x, w)` +pattern as PyTorch. JAX buffers remain row-major. Framework-specific mode metadata maps public JAX axes into the canonical kernel axis order without materializing a transpose. ## Parameters -### Input and output tensors +### Input tensors - `x_tensor` / `sample_x` - Shape: `(M, N)` @@ -136,13 +133,6 @@ JAX buffers remain row-major. Framework-specific mode metadata maps public JAX a - Shape: `(N,)` - Layout: contiguous - Dtype: `torch.bfloat16` -- `o_tensor` / `sample_o` - - Shape: `(M, N)` - - Layout: row-major contiguous - - Dtype: `torch.bfloat16` -- `amax_tensor` / `sample_amax` - - Shape: `(M / rows_per_cta,)` - - Dtype: `torch.float32` ### Common parameters @@ -154,11 +144,17 @@ JAX buffers remain row-major. Framework-specific mode metadata maps public JAX a - Rows processed by each CTA. If omitted, the wrapper uses the upstream-style heuristic over `{2, 4, 8}`. - CUDA stream (`current_stream`) -### Wrapper return values +### Return values + +The PyTorch class and wrapper return a `TupleDict` with keys: -Returns a `TupleDict` with keys: - `o_tensor` + - Shape: `(M, N)` + - Layout: row-major contiguous + - Dtype: `torch.bfloat16` - `amax_tensor` + - Shape: `(M / rows_per_cta,)` + - Dtype: `torch.float32` Tuple unpacking order is `(o_tensor, amax_tensor)`. @@ -169,11 +165,13 @@ Tuple unpacking order is `(o_tensor, amax_tensor)`. - `N` must be divisible by the resolved `num_threads`. - `EPT = N / num_threads` must be at least `8` and divisible by `8`. - `M` must be divisible by `rows_per_cta`. -- Inputs and output are currently bf16 only. +- `X`, `W`, and `O` use bf16; `amax` uses float32. +- PyTorch input tensors must be 16-byte aligned. - The frontend integration matches the upstream RMSNorm kernel semantics; it does not expose full LayerNorm mean/bias behavior. ## Verification Focused correctness and cache coverage live in: + - `test/python/fe_api/test_rmsnorm_rht_amax.py` - `test/python/fe_api/test_jax_rmsnorm_rht_amax.py` diff --git a/python/cudnn/_jax/__init__.py b/python/cudnn/_jax/__init__.py index 2e3af9dfa..3ff2b4ce8 100644 --- a/python/cudnn/_jax/__init__.py +++ b/python/cudnn/_jax/__init__.py @@ -1,35 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Optional-dependency boundary shared by all JAX operation adapters.""" - -from importlib import import_module - -_INSTALL_HINT = "pip install 'nvidia-cudnn-frontend[jax]'" -_OPTIONAL_DEPENDENCY_PREFIXES = ("cuda", "cutlass", "jax") - - -def _require_dependency(module_name: str): - try: - return import_module(module_name) - except ModuleNotFoundError as error: - missing_name = error.name - if missing_name is None or not any(missing_name == prefix or missing_name.startswith(f"{prefix}.") for prefix in _OPTIONAL_DEPENDENCY_PREFIXES): - raise - raise ImportError(f"cuDNN JAX APIs require {module_name!r}. " f"Install the optional JAX dependencies with `{_INSTALL_HINT}`.") from error - - -jax = _require_dependency("jax") -cutlass_jax = _require_dependency("cutlass.jax") - -if not cutlass_jax.is_available(): - minimum_version = ".".join(str(part) for part in cutlass_jax.CUTE_DSL_MIN_SUPPORTED_JAX_VERSION) - installed_version = getattr(jax, "__version__", "unknown") - raise ImportError( - f"CUTLASS JAX support is unavailable with JAX {installed_version}; " - f"the minimum supported JAX version is {minimum_version}. " - f"Install the optional JAX dependencies with `{_INSTALL_HINT}`." - ) +"""Internal implementation helpers for optional JAX operation adapters.""" from .api_base import JaxApiBase, JaxTensorDesc diff --git a/python/cudnn/_jax/api_base.py b/python/cudnn/_jax/api_base.py index 18447f3c3..72e9551c8 100644 --- a/python/cudnn/_jax/api_base.py +++ b/python/cudnn/_jax/api_base.py @@ -24,6 +24,36 @@ def cudnn_dtype(self) -> data_type: return jax_to_cudnn_dtype(self.dtype) + def compact_like( + self, + *, + cudnn_dtype: data_type, + shape: tuple[int, ...], + stride_order: tuple[int, ...] | None = None, + name: str = "", + init_value: bool | int | float | None = None, + ) -> "JaxTensorDesc": + """Create a compact JAX descriptor from canonical output metadata.""" + + canonical = super().compact_like( + cudnn_dtype=cudnn_dtype, + shape=shape, + stride_order=stride_order, + name=name, + init_value=init_value, + ) + + from .datatypes import cudnn_to_jax_dtype + + return JaxTensorDesc( + dtype=cudnn_to_jax_dtype(cudnn_dtype), + shape=canonical.shape, + stride=canonical.stride, + stride_order=canonical.stride_order, + name=canonical.name, + init_value=canonical.init_value, + ) + def _require_array_metadata(value: Any, name: str) -> tuple[Any, ...]: if not hasattr(value, "shape") or not hasattr(value, "dtype"): diff --git a/python/cudnn/_op_kernel.py b/python/cudnn/_op_kernel.py index 9e448196c..350fdd961 100644 --- a/python/cudnn/_op_kernel.py +++ b/python/cudnn/_op_kernel.py @@ -6,8 +6,8 @@ from __future__ import annotations from abc import ABC, abstractmethod +from typing import Any -from . import data_type from ._tensor_desc import TensorDesc @@ -27,15 +27,15 @@ def check_support(self) -> bool: """Validate the logical signature and resolve static kernel state.""" @abstractmethod - def infer_output(self) -> tuple[TensorDesc[data_type], ...]: - """Return framework-neutral descriptors for the operation outputs. + def infer_output(self) -> tuple[TensorDesc[Any], ...]: + """Return operation outputs through the framework-neutral descriptor interface. A descriptor's ``init_value`` requests scalar initialization before launch; ``None`` leaves the output uninitialized. """ - def infer_workspace(self) -> tuple[TensorDesc[data_type], ...]: - """Return framework-neutral workspace descriptors, if any. + def infer_workspace(self) -> tuple[TensorDesc[Any], ...]: + """Return workspaces through the framework-neutral descriptor interface. Workspace descriptors follow the same ``init_value`` convention as outputs but are not part of the framework-visible result. diff --git a/python/cudnn/_tensor_desc.py b/python/cudnn/_tensor_desc.py index 1080e2d11..5c519397e 100644 --- a/python/cudnn/_tensor_desc.py +++ b/python/cudnn/_tensor_desc.py @@ -8,7 +8,7 @@ from dataclasses import dataclass, field from numbers import Real from operator import index -from typing import Generic, TypeVar +from typing import Any, Generic, TypeVar from . import data_type @@ -87,6 +87,30 @@ def cudnn_dtype(self) -> data_type: raise TypeError(f"{type(self).__name__}.dtype must provide a cudnn.data_type mapping, " f"got {type(self.dtype).__name__}") return self.dtype + def compact_like( + self, + *, + cudnn_dtype: data_type, + shape: tuple[int, ...], + stride_order: tuple[int, ...] | None = None, + name: str = "", + init_value: bool | int | float | None = None, + ) -> "TensorDesc[Any]": + """Create a compact descriptor derived from this descriptor. + + The framework-neutral implementation returns a canonical descriptor. + Framework-specific descriptors override this method to preserve their + native dtype and allocation metadata. + """ + + return make_compact_tensor_desc( + dtype=cudnn_dtype, + shape=shape, + stride_order=stride_order, + name=name, + init_value=init_value, + ) + def make_compact_tensor_desc( *, diff --git a/python/cudnn/api_base.py b/python/cudnn/api_base.py index 30895efc7..f15c763d2 100644 --- a/python/cudnn/api_base.py +++ b/python/cudnn/api_base.py @@ -99,6 +99,31 @@ def __post_init__(self): object.__setattr__(self, "device", device) + @classmethod + def from_tensor( + cls, + tensor: torch.Tensor, + name: str = "", + *, + shape: Optional[Tuple[int, ...]] = None, + stride: Optional[Tuple[int, ...]] = None, + interpret_uint8_as_fp4x2: bool = False, + ) -> "TensorDesc": + """Capture a Torch tensor's physical or supplied logical layout.""" + + tensor_shape = tuple(tensor.shape) if shape is None else tuple(shape) + tensor_stride = tuple(tensor.stride()) if stride is None else tuple(stride) + tensor_stride_order = cls._compute_stride_order(tensor_shape, tensor_stride) + return cls( + dtype=tensor.dtype, + shape=tensor_shape, + stride=tensor_stride, + stride_order=tensor_stride_order, + device=tensor.device, + interpret_uint8_as_fp4x2=interpret_uint8_as_fp4x2, + name=name, + ) + @property def cudnn_dtype(self) -> data_type: if self.interpret_uint8_as_fp4x2 and self.dtype == torch.uint8: @@ -106,6 +131,61 @@ def cudnn_dtype(self) -> data_type: converted = _torch_to_cudnn_data_type(self.dtype) return data_type.NOT_SET if converted is None else converted + def compact_like( + self, + *, + cudnn_dtype: data_type, + shape: tuple[int, ...], + stride_order: tuple[int, ...] | None = None, + name: str = "", + init_value: bool | int | float | None = None, + ) -> "TensorDesc": + """Create a compact Torch descriptor on this descriptor's device.""" + + canonical = super().compact_like( + cudnn_dtype=cudnn_dtype, + shape=shape, + stride_order=stride_order, + name=name, + init_value=init_value, + ) + dtype = _cudnn_to_torch_data_type(cudnn_dtype) + if dtype is None: + raise ValueError(f"Unsupported Torch output dtype {cudnn_dtype}") + return TensorDesc( + dtype=dtype, + shape=canonical.shape, + stride=canonical.stride, + stride_order=canonical.stride_order, + device=self.device, + name=canonical.name, + init_value=canonical.init_value, + ) + + def materialize(self, stream: Optional[cuda.CUstream] = None) -> torch.Tensor: + """Allocate and optionally initialize a tensor described by this object.""" + + if self.interpret_uint8_as_fp4x2 and self.dtype == torch.uint8: + raise ValueError("Materializing a logical FP4 descriptor backed by uint8 storage is unsupported") + + def allocate() -> torch.Tensor: + tensor = torch.empty_strided( + self.shape, + self.stride, + dtype=self.dtype, + device=self.device, + ) + if self.init_value is not None: + tensor.fill_(self.init_value) + return tensor + + if stream is None: + return allocate() + if self.device.type != "cuda": + raise ValueError("A CUDA stream can only be used to materialize a CUDA tensor") + with torch.cuda.stream(torch.cuda.ExternalStream(int(stream), device=self.device)): + return allocate() + @staticmethod def _normalize_dim(dim: int, ndim: int, *, allow_new_dim: bool = False) -> int: min_dim = -ndim - (1 if allow_new_dim else 0) @@ -869,6 +949,22 @@ def _check_dtype( raise ValueError(f"Expected dtype to be a torch.dtype or list, got {type(dtype)}") return tensor_dtype + def _check_tensor_signature(self, tensor: torch.Tensor, expected: TensorDesc) -> None: + """Validate a Torch tensor against a captured logical descriptor.""" + + if not isinstance(expected, TensorDesc): + raise TypeError(f"expected must be a TensorDesc, got {type(expected).__name__}") + name = expected.name or "tensor" + actual = self._make_tensor_desc( + tensor, + name=name, + interpret_uint8_as_fp4x2=expected.interpret_uint8_as_fp4x2, + ) + self._check_tensor_shape(actual, expected.shape, name) + self._check_tensor_stride(actual, stride=expected.stride, stride_order=expected.stride_order, name=name) + self._check_dtype(actual, expected.dtype, name) + self._value_error_if(actual.device != expected.device, f"{name} tensor device mismatch: expected {expected.device}, got {actual.device}") + def _value_error_if(self, condition: bool, error_msg: str) -> None: """Raise a ValueError if the condition is true. @@ -937,17 +1033,12 @@ def _to_tensor_desc( ) -> TensorDesc: """Create a Torch descriptor from tensor and optional logical layout metadata.""" - tensor_shape = tuple(tensor.shape) if shape is None else tuple(shape) - tensor_stride = tuple(tensor.stride()) if stride is None else tuple(stride) - tensor_stride_order = tuple(i for i, _ in sorted(enumerate(tensor_stride), key=lambda item: (item[1], tensor_shape[item[0]]))) - return TensorDesc( - dtype=tensor.dtype, - shape=tensor_shape, - stride=tensor_stride, - stride_order=tensor_stride_order, - device=tensor.device, + return TensorDesc.from_tensor( + tensor, + name, + shape=shape, + stride=stride, interpret_uint8_as_fp4x2=interpret_uint8_as_fp4x2, - name=name, ) @staticmethod diff --git a/python/cudnn/jax/__init__.py b/python/cudnn/jax/__init__.py index 8d31c8e3c..20ab5a1a7 100644 --- a/python/cudnn/jax/__init__.py +++ b/python/cudnn/jax/__init__.py @@ -1,7 +1,28 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Public facade for framework-specific JAX operation APIs.""" +"""Public optional-dependency boundary for JAX operation APIs.""" + +_INSTALL_HINT = "pip install 'nvidia-cudnn-frontend[jax]'" +_OPTIONAL_DEPENDENCY_PREFIXES = ("cuda", "cutlass", "jax") + +try: + import cutlass.jax as _cutlass_jax + import jax as _jax +except ModuleNotFoundError as error: + missing_name = error.name + if missing_name is None or not any(missing_name == prefix or missing_name.startswith(f"{prefix}.") for prefix in _OPTIONAL_DEPENDENCY_PREFIXES): + raise + raise ImportError(f"cuDNN JAX APIs require the optional JAX dependencies. Install them with `{_INSTALL_HINT}`.") from error + +if not _cutlass_jax.is_available(): + minimum_version = ".".join(str(part) for part in _cutlass_jax.CUTE_DSL_MIN_SUPPORTED_JAX_VERSION) + installed_version = getattr(_jax, "__version__", "unknown") + raise ImportError( + f"CUTLASS JAX support is unavailable with JAX {installed_version}; " + f"the minimum supported JAX version is {minimum_version}. " + f"Install the optional JAX dependencies with `{_INSTALL_HINT}`." + ) from .._jax import JaxApiBase, JaxTensorDesc from ..rmsnorm_rht_amax.jax import ( diff --git a/python/cudnn/rmsnorm_rht_amax/__init__.py b/python/cudnn/rmsnorm_rht_amax/__init__.py index 3cab20e81..9f11062da 100644 --- a/python/cudnn/rmsnorm_rht_amax/__init__.py +++ b/python/cudnn/rmsnorm_rht_amax/__init__.py @@ -1,7 +1,10 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Lazy framework adapters for RMSNorm + RHT + amax.""" +"""Lazy Torch API and framework-neutral kernel exports. + +JAX APIs are exported through :mod:`cudnn.jax`. +""" from .._operation_api import make_operation_api @@ -18,5 +21,5 @@ "rmsnorm_rht_amax_wrapper_sm100", ), }, - submodules=("api", "jax", "kernel"), + submodules=("api", "kernel"), ) diff --git a/python/cudnn/rmsnorm_rht_amax/api.py b/python/cudnn/rmsnorm_rht_amax/api.py index 2f21964e7..acab89a6a 100644 --- a/python/cudnn/rmsnorm_rht_amax/api.py +++ b/python/cudnn/rmsnorm_rht_amax/api.py @@ -39,19 +39,12 @@ def _validate_torch_tensors(*named_tensors: tuple[str, torch.Tensor]) -> None: class RmsNormRhtAmaxSm100(APIBase): - """Class API for the RMSNorm + RHT + amax kernel. - - Logical validation, output inference, and resolved launch state are owned - by :attr:`kernel`. Output samples remain optional for callers that want - the framework adapter to allocate inferred outputs. - """ + """Class API for the RMSNorm + RHT + amax kernel.""" def __init__( self, sample_x: torch.Tensor, sample_w: torch.Tensor, - sample_o: Optional[torch.Tensor] = None, - sample_amax: Optional[torch.Tensor] = None, eps: float = 1e-5, num_threads: Optional[int] = None, rows_per_cta: Optional[int] = None, @@ -59,98 +52,29 @@ def __init__( super().__init__() self._warn_experimental_api() - if (sample_o is None) != (sample_amax is None): - raise ValueError("sample_o and sample_amax must either both be provided or both be omitted") - - named_tensors = [("sample_x", sample_x), ("sample_w", sample_w)] - if sample_o is not None and sample_amax is not None: - named_tensors.extend((("sample_o", sample_o), ("sample_amax", sample_amax))) - _validate_torch_tensors(*named_tensors) + _validate_torch_tensors( + ("sample_x", sample_x), + ("sample_w", sample_w), + ) self.x_desc = self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_x, name="sample_x"), 2, "sample_x") self.w_desc = self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_w, name="sample_w"), 1, "sample_w") - self.o_desc = None if sample_o is None else self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_o, name="sample_o"), 2, "sample_o") - self.amax_desc = None if sample_amax is None else self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_amax, name="sample_amax"), 1, "sample_amax") - self.kernel = RMSNormRHTAmaxKernel( + self.eps = eps + self.requested_num_threads = num_threads + self.requested_rows_per_cta = rows_per_cta + + self._kernel = RMSNormRHTAmaxKernel( x=self.x_desc, weight=self.w_desc, - output=self.o_desc, - amax=self.amax_desc, eps=eps, num_threads=num_threads, rows_per_cta=rows_per_cta, ) - @property - def eps(self) -> float: - return self.kernel.eps - - @eps.setter - def eps(self, value: float) -> None: - self.kernel.eps = value - self._is_supported = False - self._compiled_kernel = None - - @property - def requested_num_threads(self) -> Optional[int]: - return self.kernel.requested_num_threads - - @requested_num_threads.setter - def requested_num_threads(self, value: Optional[int]) -> None: - self.kernel.requested_num_threads = value - self._is_supported = False - self._compiled_kernel = None - - @property - def requested_rows_per_cta(self) -> Optional[int]: - return self.kernel.requested_rows_per_cta - - @requested_rows_per_cta.setter - def requested_rows_per_cta(self, value: Optional[int]) -> None: - self.kernel.requested_rows_per_cta = value - self._is_supported = False - self._compiled_kernel = None - - @property - def num_threads(self) -> Optional[int]: - return self.kernel.num_threads - - @property - def rows_per_cta(self) -> Optional[int]: - return self.kernel.rows_per_cta - - @property - def n(self) -> Optional[int]: - return self.kernel.n - - def _materialize_outputs( - self, - *, - current_stream: Optional[cuda.CUstream] = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Infer and allocate outputs using this API's kernel instance.""" - - self._ensure_support_checked() - output_desc, amax_desc = self.kernel.infer_output() - o_tensor = self._materialize_tensor_desc(output_desc, device=self.x_desc.device, stream=current_stream) - amax_tensor = self._materialize_tensor_desc(amax_desc, device=self.x_desc.device, stream=current_stream) - - if self.o_desc is None and self.amax_desc is None: - self.o_desc = self._to_tensor_desc(o_tensor, "sample_o") - self.amax_desc = self._to_tensor_desc(amax_tensor, "sample_amax") - self.kernel.output = self.o_desc - self.kernel.amax = self.amax_desc - - return o_tensor, amax_tensor - def check_support(self) -> bool: self._is_supported = False - self.kernel.check_support() - - for name, desc in (("W", self.w_desc), ("O", self.o_desc), ("Amax", self.amax_desc)): - if desc is not None: - self._value_error_if(desc.device != self.x_desc.device, f"{name} must be on {self.x_desc.device}, got {desc.device}") + self._kernel.check_support() self._runtime_error_if(not torch.cuda.is_available(), "CUDA is not available") major, minor = torch.cuda.get_device_capability(self.x_desc.device) @@ -167,38 +91,38 @@ def compile(self) -> None: self._ensure_support_checked() if self._compiled_kernel is not None: return - self._runtime_error_if(self.o_desc is None or self.amax_desc is None, "Output tensors must be provided or materialized before compile()") - valid_m = cute.sym_int(divisibility=self.kernel.rows_per_cta) + o_desc, amax_desc = self._kernel.infer_output() + valid_m = cute.sym_int(divisibility=self._kernel.rows_per_cta) fake_x_tensor = self._make_fake_cute_compact_tensor( dtype=self.x_desc.dtype, - shape=(valid_m, self.kernel.n), + shape=(valid_m, self._kernel.n), stride_order=self.x_desc.stride_order, assumed_align=_TENSOR_ALIGNMENT, dynamic_mode=None, - divisibility=self.kernel.rows_per_cta, + divisibility=self._kernel.rows_per_cta, ) fake_w_tensor = self._make_fake_cute_tensor_from_desc(self.w_desc, assumed_align=_TENSOR_ALIGNMENT) fake_o_tensor = self._make_fake_cute_compact_tensor( - dtype=self.o_desc.dtype, - shape=(valid_m, self.kernel.n), - stride_order=self.o_desc.stride_order, + dtype=o_desc.dtype, + shape=(valid_m, self._kernel.n), + stride_order=o_desc.stride_order, assumed_align=_TENSOR_ALIGNMENT, dynamic_mode=None, - divisibility=self.kernel.rows_per_cta, + divisibility=self._kernel.rows_per_cta, ) fake_num_ctas = cute.sym_int() fake_amax_tensor = self._make_fake_cute_tensor( - dtype=self.amax_desc.dtype, + dtype=amax_desc.dtype, shape=(fake_num_ctas,), - stride=self.amax_desc.stride, + stride=amax_desc.stride, assumed_align=_TENSOR_ALIGNMENT, ) fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) compiled_kernel = cute.compile( - self.kernel, + self._kernel, fake_x_tensor, fake_w_tensor, fake_o_tensor, @@ -228,26 +152,26 @@ def execute( self, x_tensor: torch.Tensor, w_tensor: torch.Tensor, - o_tensor: torch.Tensor, - amax_tensor: torch.Tensor, current_stream: Optional[cuda.CUstream] = None, - ) -> None: + ) -> TupleDict: self._runtime_error_if(self._compiled_kernel is None, "RmsNormRhtAmaxSm100 kernel not compiled; call compile() first") x_tensor = self._unpad_tensor_to_ndim(x_tensor, 2, "x_tensor") w_tensor = self._unpad_tensor_to_ndim(w_tensor, 1, "w_tensor") - o_tensor = self._unpad_tensor_to_ndim(o_tensor, 2, "o_tensor") - amax_tensor = self._unpad_tensor_to_ndim(amax_tensor, 1, "amax_tensor") _validate_torch_tensors( ("x_tensor", x_tensor), ("w_tensor", w_tensor), - ("o_tensor", o_tensor), - ("amax_tensor", amax_tensor), ) + self._check_tensor_signature(x_tensor, self.x_desc) + self._check_tensor_signature(w_tensor, self.w_desc) if current_stream is None: current_stream = cuda.CUstream(torch.cuda.current_stream(x_tensor.device).cuda_stream) + o_desc, amax_desc = self._kernel.infer_output() + o_tensor = o_desc.materialize(stream=current_stream) + amax_tensor = amax_desc.materialize(stream=current_stream) + self._compiled_kernel( x_tensor=x_tensor, w_tensor=w_tensor, @@ -255,6 +179,7 @@ def execute( amax_tensor=amax_tensor, stream=current_stream, ) + return TupleDict(o_tensor=o_tensor, amax_tensor=amax_tensor) _logger = logging.getLogger(__name__) @@ -273,7 +198,6 @@ def rmsnorm_rht_amax_wrapper_sm100( x_tensor = x_tensor.squeeze(-1) if x_tensor.ndim == 3 and x_tensor.shape[-1] == 1 else x_tensor w_tensor = w_tensor.squeeze(-1) if w_tensor.ndim == 2 and w_tensor.shape[-1] == 1 else w_tensor - _validate_torch_tensors(("x_tensor", x_tensor), ("w_tensor", w_tensor)) cache_key = ( tuple(x_tensor.shape), @@ -289,8 +213,7 @@ def rmsnorm_rht_amax_wrapper_sm100( ) api = _cache_of_RmsNormRhtAmaxSm100Objects.get(cache_key) - cache_miss = api is None - if cache_miss: + if api is None: api = RmsNormRhtAmaxSm100( sample_x=x_tensor, sample_w=w_tensor, @@ -299,21 +222,11 @@ def rmsnorm_rht_amax_wrapper_sm100( rows_per_cta=rows_per_cta, ) api.check_support() - - o_tensor, amax_tensor = api._materialize_outputs(current_stream=current_stream) - if cache_miss: api.compile() _cache_of_RmsNormRhtAmaxSm100Objects[cache_key] = api - if current_stream is None and x_tensor.is_cuda: - current_stream = cuda.CUstream(torch.cuda.current_stream(x_tensor.device).cuda_stream) - - api.execute( + return api.execute( x_tensor=x_tensor, w_tensor=w_tensor, - o_tensor=o_tensor, - amax_tensor=amax_tensor, current_stream=current_stream, ) - - return TupleDict(o_tensor=o_tensor, amax_tensor=amax_tensor) diff --git a/python/cudnn/rmsnorm_rht_amax/jax.py b/python/cudnn/rmsnorm_rht_amax/jax.py index fdf80b28e..4fbcc0eb2 100644 --- a/python/cudnn/rmsnorm_rht_amax/jax.py +++ b/python/cudnn/rmsnorm_rht_amax/jax.py @@ -7,7 +7,9 @@ from typing import Any, Optional -from .._jax import JaxApiBase, jax +import jax + +from .._jax import JaxApiBase from .kernel import RMSNormRHTAmaxKernel diff --git a/python/cudnn/rmsnorm_rht_amax/kernel.py b/python/cudnn/rmsnorm_rht_amax/kernel.py index 9d6382572..f00c90e81 100644 --- a/python/cudnn/rmsnorm_rht_amax/kernel.py +++ b/python/cudnn/rmsnorm_rht_amax/kernel.py @@ -20,7 +20,7 @@ from .. import data_type from .._op_kernel import OpKernel -from .._tensor_desc import TensorDesc, make_compact_tensor_desc +from .._tensor_desc import TensorDesc DEFAULT_NUM_THREADS_BY_N = { 2048: 128, @@ -108,9 +108,8 @@ def redux_sync_max_f32(val, *, loc=None, ip=None): class RMSNormRHTAmaxKernel(OpKernel): """Fused RMSNorm + block-diagonal Hadamard + running per-CTA amax. - Framework adapters bind generic tensor descriptors, call - :meth:`check_support`, and materialize the descriptors returned by - :meth:`infer_output` in their own framework. + Output inference preserves the framework descriptor specialization of the + input while validation and launch configuration remain framework-neutral. """ COPY_BITS = 128 @@ -124,8 +123,6 @@ def __init__( *, x: TensorDesc[Any], weight: TensorDesc[Any], - output: Optional[TensorDesc[Any]] = None, - amax: Optional[TensorDesc[Any]] = None, eps: float = 1e-5, num_threads: Optional[int] = None, rows_per_cta: Optional[int] = None, @@ -133,14 +130,9 @@ def __init__( for name, desc in (("x", x), ("weight", weight)): if not isinstance(desc, TensorDesc): raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") - for name, desc in (("output", output), ("amax", amax)): - if desc is not None and not isinstance(desc, TensorDesc): - raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") self.x = x self.weight = weight - self.output = output - self.amax = amax self.eps = eps self.requested_num_threads = num_threads self.requested_rows_per_cta = rows_per_cta @@ -149,6 +141,7 @@ def __init__( self.n: Optional[int] = None self.num_threads: Optional[int] = None self.rows_per_cta: Optional[int] = None + self._validated_x: Optional[TensorDesc[Any]] = None def check_support(self) -> bool: """Validate descriptors and resolve the kernel launch configuration.""" @@ -157,6 +150,7 @@ def check_support(self) -> bool: self.n = None self.num_threads = None self.rows_per_cta = None + self._validated_x = None if self.x.ndim != 2: raise ValueError(f"X must have rank 2, got shape {self.x.shape}") @@ -187,53 +181,47 @@ def check_support(self) -> bool: if num_threads is None: raise ValueError(f"No valid num_threads found for N={n}") - rows_per_cta = self.requested_rows_per_cta - if rows_per_cta is None: - rows_per_cta = pick_rows_per_cta(m) + rows_per_cta = self._resolve_rows_per_cta(m, self.requested_rows_per_cta) self._validate_launch_configuration(n, num_threads, rows_per_cta, m=m) - expected_output_shape = (m, n) - expected_amax_shape = (m // rows_per_cta,) - if self.output is not None: - if self.output.shape != expected_output_shape: - raise ValueError(f"O must have shape {expected_output_shape}, got {self.output.shape}") - if self.output.cudnn_dtype != self.output_dtype: - raise ValueError(f"O must have dtype bfloat16, got {self.output.dtype}") - if self.output.stride != (n, 1) or self.output.stride_order != (1, 0): - raise ValueError(f"O must be row-major contiguous, got stride {self.output.stride} " f"and stride order {self.output.stride_order}") - if self.amax is not None: - if self.amax.shape != expected_amax_shape: - raise ValueError(f"Amax must have shape {expected_amax_shape}, got {self.amax.shape}") - if self.amax.cudnn_dtype != self.amax_dtype: - raise ValueError(f"Amax must have dtype float32, got {self.amax.dtype}") - self.m = m self.n = n self.num_threads = num_threads self.rows_per_cta = rows_per_cta self._configure_lowering_state() + self._validated_x = self.x return True - def infer_output(self) -> tuple[TensorDesc[data_type], ...]: + def infer_output(self) -> tuple[TensorDesc[Any], ...]: """Infer compact output descriptors from validated input metadata.""" - if self.m is None or self.n is None or self.rows_per_cta is None: + if self.m is None or self.n is None or self.rows_per_cta is None or self._validated_x is None: raise RuntimeError("check_support() must be called before inferring outputs") return ( - make_compact_tensor_desc( - dtype=self.output_dtype, + self._validated_x.compact_like( + cudnn_dtype=self.output_dtype, shape=(self.m, self.n), name="output", ), - make_compact_tensor_desc( - dtype=self.amax_dtype, + self._validated_x.compact_like( + cudnn_dtype=self.amax_dtype, shape=(self.m // self.rows_per_cta,), name="amax", ), ) + @staticmethod + def _resolve_rows_per_cta(m: int, rows_per_cta: Optional[int]) -> int: + if rows_per_cta is None: + rows_per_cta = pick_rows_per_cta(m) + if rows_per_cta <= 0: + raise ValueError(f"rows_per_cta must be positive, got {rows_per_cta}") + if m % rows_per_cta != 0: + raise ValueError(f"M must be divisible by rows_per_cta, got M={m}, rows_per_cta={rows_per_cta}") + return rows_per_cta + @staticmethod def _validate_launch_configuration( n: int, diff --git a/test/python/fe_api/test_jax_api_base.py b/test/python/fe_api/test_jax_api_base.py index 4df53307e..6966e8b85 100644 --- a/test/python/fe_api/test_jax_api_base.py +++ b/test/python/fe_api/test_jax_api_base.py @@ -127,6 +127,29 @@ def test_converts_public_array_axes_to_canonical_descriptor_axes(self): self.assertEqual(desc.stride, (1, 12, 4)) self.assertEqual(desc.stride_order, (0, 2, 1)) + def test_jax_descriptor_derives_compact_output_metadata(self): + source = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3), "bfloat16"), + "input", + ) + + output = source.compact_like( + cudnn_dtype=_DataType.FLOAT, + shape=(5, 7), + stride_order=(0, 1), + name="output", + init_value=-2.0, + ) + + self.assertIsInstance(output, self.module.JaxTensorDesc) + self.assertEqual(output.dtype, "float32") + self.assertEqual(output.cudnn_dtype, _DataType.FLOAT) + self.assertEqual(output.shape, (5, 7)) + self.assertEqual(output.stride, (1, 5)) + self.assertEqual(output.stride_order, (0, 1)) + self.assertEqual(output.name, "output") + self.assertEqual(output.init_value, -2.0) + def test_checks_invocation_signature(self): desc = self.module.JaxApiBase._to_tensor_desc(_ArrayMetadata((2, 3), "bfloat16"), "sample") self.module.JaxApiBase._check_tensor_signature(_ArrayMetadata((2, 3), "bfloat16"), desc) diff --git a/test/python/fe_api/test_rmsnorm_rht_amax.py b/test/python/fe_api/test_rmsnorm_rht_amax.py index 0128606f6..abec7c9a3 100644 --- a/test/python/fe_api/test_rmsnorm_rht_amax.py +++ b/test/python/fe_api/test_rmsnorm_rht_amax.py @@ -59,10 +59,9 @@ def _assert_ref_close(x, w, o, amax, *, eps: float, rows_per_cta: int, skip_ref: @pytest.mark.L0 @torch_fork_set_rng(seed=0) @pytest.mark.parametrize("n,num_threads", SUPPORTED_N_NUM_THREADS) -def test_rmsnorm_rht_amax_compile_execute(n, num_threads, request): +def test_rmsnorm_rht_amax_class_call(n, num_threads, request): try: - from cudnn import OpKernel, RmsNormRhtAmaxSm100, TensorDesc - from cudnn.rmsnorm_rht_amax.kernel import RMSNormRHTAmaxKernel + from cudnn import RmsNormRhtAmaxSm100 except ImportError: pytest.skip("Environment not supported: cudnn optional dependencies not installed") @@ -71,26 +70,14 @@ def test_rmsnorm_rht_amax_compile_execute(n, num_threads, request): m = 256 rows_per_cta = 2 x, w = _make_inputs(m=m, n=n) - o = torch.empty_like(x) - amax = torch.full((m // rows_per_cta,), float("-inf"), dtype=torch.float32, device="cuda") api = RmsNormRhtAmaxSm100( sample_x=x, sample_w=w, - sample_o=o, - sample_amax=amax, eps=eps, num_threads=num_threads, rows_per_cta=rows_per_cta, ) - assert isinstance(api.kernel, OpKernel) - assert isinstance(api.kernel, RMSNormRHTAmaxKernel) - assert isinstance(api.kernel.x, TensorDesc) - assert isinstance(api.kernel.weight, TensorDesc) - assert api.kernel.x is api.x_desc - assert api.kernel.weight is api.w_desc - assert api.kernel.output is api.o_desc - assert api.kernel.amax is api.amax_desc assert api.eps == eps assert api.requested_num_threads == num_threads assert api.requested_rows_per_cta == rows_per_cta @@ -100,15 +87,32 @@ def test_rmsnorm_rht_amax_compile_execute(n, num_threads, request): except (ValueError, RuntimeError) as exc: pytest.skip(f"Unsupported testcase: {exc}") - assert api.n == n - assert api.num_threads == num_threads - assert api.rows_per_cta == rows_per_cta - - api.compile() - api.execute(x_tensor=x, w_tensor=w, o_tensor=o, amax_tensor=amax) + outputs = api(x, w) + o, amax = outputs + assert outputs["o_tensor"] is o + assert outputs["amax_tensor"] is amax _assert_ref_close(x, w, o, amax, eps=eps, rows_per_cta=rows_per_cta, skip_ref=skip_ref) +@pytest.mark.L0 +def test_rmsnorm_rht_amax_class_rejects_runtime_signature_mismatch(monkeypatch): + try: + from cudnn import RmsNormRhtAmaxSm100 + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + sample_x = torch.empty((4, 256), dtype=torch.bfloat16) + sample_w = torch.empty((256,), dtype=torch.bfloat16) + api = RmsNormRhtAmaxSm100(sample_x, sample_w, num_threads=32, rows_per_cta=2) + api._compiled_kernel = lambda **_kwargs: None + monkeypatch.setattr(api._kernel, "infer_output", lambda: pytest.fail("output inference must follow runtime signature validation")) + + with pytest.raises(ValueError, match="sample_x tensor shape mismatch"): + api.execute(torch.empty((2, 256), dtype=torch.bfloat16), sample_w) + with pytest.raises(ValueError, match="sample_w dtype mismatch"): + api.execute(sample_x, torch.empty((256,), dtype=torch.float32)) + + @pytest.mark.L0 @torch_fork_set_rng(seed=0) @pytest.mark.parametrize("n,num_threads", SUPPORTED_N_NUM_THREADS) @@ -173,3 +177,45 @@ def counted_init(self, *args, **kwargs): assert construction_count == 1 assert first["o_tensor"].shape == second["o_tensor"].shape == (m, n) assert first["amax_tensor"].shape == second["amax_tensor"].shape == (m // rows_per_cta,) + + +@pytest.mark.L0 +def test_rmsnorm_rht_amax_wrapper_delegates_output_allocation_to_class(monkeypatch): + try: + from cudnn.rmsnorm_rht_amax import api as rmsnorm_api + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + constructed = {} + output = torch.empty((3, 5), dtype=torch.bfloat16) + amax = torch.full((7,), -2.0, dtype=torch.float32) + + class FakeApi: + def __init__(self, **kwargs): + constructed.update(kwargs) + + def check_support(self): + return True + + def compile(self): + pass + + def execute(self, **kwargs): + constructed["execute"] = kwargs + return rmsnorm_api.TupleDict(o_tensor=output, amax_tensor=amax) + + monkeypatch.setattr(rmsnorm_api, "RmsNormRhtAmaxSm100", FakeApi) + rmsnorm_api._cache_of_RmsNormRhtAmaxSm100Objects.clear() + + x = torch.empty((4, 16), dtype=torch.bfloat16) + w = torch.empty((16,), dtype=torch.bfloat16) + try: + result = rmsnorm_api.rmsnorm_rht_amax_wrapper_sm100(x, w, rows_per_cta=4) + finally: + rmsnorm_api._cache_of_RmsNormRhtAmaxSm100Objects.clear() + + assert set(constructed) == {"sample_x", "sample_w", "eps", "num_threads", "rows_per_cta", "execute"} + assert result["o_tensor"] is output + assert result["amax_tensor"] is amax + assert torch.all(result["amax_tensor"] == -2.0) + assert constructed["execute"] == {"x_tensor": x, "w_tensor": w, "current_stream": None} diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_imports.py b/test/python/fe_api/test_rmsnorm_rht_amax_imports.py index 9c4b2eb49..769f1e794 100644 --- a/test/python/fe_api/test_rmsnorm_rht_amax_imports.py +++ b/test/python/fe_api/test_rmsnorm_rht_amax_imports.py @@ -62,26 +62,26 @@ def _remove_package(name: str) -> None: class RmsNormImportContractTest(unittest.TestCase): - def test_internal_jax_package_reports_missing_jax_dependency(self): + def test_jax_facade_reports_missing_jax_dependency(self): name = "cudnn_jax_missing_dependency_test" try: with mock.patch.dict(sys.modules, {"jax": None}): with self.assertRaisesRegex(ImportError, r"nvidia-cudnn-frontend\[jax\]"): - _load_cudnn_submodule(name, "_jax") + _load_cudnn_submodule(name, "jax") finally: _remove_package(name) - def test_internal_jax_package_reports_missing_cutlass_dependency(self): + def test_jax_facade_reports_missing_cutlass_dependency(self): name = "cudnn_jax_missing_cutlass_test" jax = types.ModuleType("jax") try: with mock.patch.dict(sys.modules, {"jax": jax, "cutlass": None, "cutlass.jax": None}): with self.assertRaisesRegex(ImportError, r"nvidia-cudnn-frontend\[jax\]"): - _load_cudnn_submodule(name, "_jax") + _load_cudnn_submodule(name, "jax") finally: _remove_package(name) - def test_internal_jax_package_reports_incompatible_jax(self): + def test_jax_facade_reports_incompatible_jax(self): name = "cudnn_jax_incompatible_test" jax = types.ModuleType("jax") jax.__version__ = "0.8.0" @@ -94,26 +94,15 @@ def test_internal_jax_package_reports_incompatible_jax(self): try: with mock.patch.dict(sys.modules, {"jax": jax, "cutlass": cutlass, "cutlass.jax": cutlass_jax}): with self.assertRaisesRegex(ImportError, r"JAX 0\.8\.0.*minimum supported JAX version is 0\.9\.1.*nvidia-cudnn-frontend\[jax\]"): - _load_cudnn_submodule(name, "_jax") - finally: - _remove_package(name) - - def test_jax_facade_uses_internal_dependency_boundary(self): - name = "cudnn_jax_facade_dependency_test" - try: - with mock.patch.dict(sys.modules, {"jax": None}): - with self.assertRaisesRegex(ImportError, r"nvidia-cudnn-frontend\[jax\]"): _load_cudnn_submodule(name, "jax") finally: _remove_package(name) - def test_jax_operator_uses_internal_dependency_boundary(self): - name = "cudnn_jax_operator_dependency_test.rmsnorm_rht_amax" + def test_operation_package_does_not_advertise_jax_submodule(self): + name = "cudnn_jax_operator_route_test.rmsnorm_rht_amax" try: package = _load_operation_package(name) - with mock.patch.dict(sys.modules, {"jax": None}): - with self.assertRaisesRegex(ImportError, r"nvidia-cudnn-frontend\[jax\]"): - package.jax + self.assertNotIn("jax", dir(package)) finally: _remove_package(name) @@ -131,7 +120,6 @@ def test_package_import_loads_no_framework_adapter(self): "RMSNormRHTAmaxKernel", "RmsNormRhtAmaxSm100", "api", - "jax", "kernel", }.issubset(dir(package)) ) @@ -165,17 +153,6 @@ def test_kernel_exports_do_not_load_an_adapter(self): finally: _remove_package(name) - def test_jax_adapter_is_explicit(self): - name = "cudnn_rmsnorm_jax_route_test.rmsnorm_rht_amax" - jax_module = types.ModuleType(f"{name}.jax") - try: - sys.modules[jax_module.__name__] = jax_module - package = _load_operation_package(name) - self.assertIs(package.jax, jax_module) - self.assertNotIn(f"{name}.api", sys.modules) - finally: - _remove_package(name) - def test_static_import_directions(self): def all_imports(filename): path = Path(filename) @@ -224,6 +201,24 @@ def imports_framework(node, framework: str): ] self.assertEqual(len(kernel_constructors), 1) self.assertIn(kernel_constructors[0], ast.walk(constructor)) + self.assertFalse( + any( + isinstance(node, ast.FunctionDef) and any(isinstance(decorator, ast.Name) and decorator.id == "property" for decorator in node.decorator_list) + for node in torch_adapter.body + ) + ) + torch_check_support = next(node for node in torch_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "check_support") + self.assertFalse( + any( + isinstance(node, ast.Attribute) + and isinstance(node.ctx, ast.Store) + and isinstance(node.value, ast.Attribute) + and isinstance(node.value.value, ast.Name) + and node.value.value.id == "self" + and node.value.attr == "_kernel" + for node in ast.walk(torch_check_support) + ) + ) jax_adapter = next(node for node in jax_tree.body if isinstance(node, ast.ClassDef) and node.name == "RmsNormRhtAmaxSm100") self.assertEqual([base.id for base in jax_adapter.bases if isinstance(base, ast.Name)], ["JaxApiBase"]) @@ -248,7 +243,27 @@ def imports_framework(node, framework: str): self.assertFalse( any(isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RMSNormRHTAmaxKernel" for node in ast.walk(torch_wrapper)) ) - self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "_materialize_outputs" for node in ast.walk(torch_wrapper))) + self.assertFalse(any(isinstance(node, ast.Attribute) and node.attr == "_materialize_outputs" for node in ast.walk(torch_wrapper))) + self.assertFalse( + any( + isinstance(node, ast.Attribute) + and node.attr in {"infer_output_from", "from_tensor", "materialize", "_to_tensor_desc", "_materialize_tensor_desc"} + for node in ast.walk(torch_wrapper) + ) + ) + self.assertFalse(any(isinstance(node, ast.Name) and node.id in {"best_num_threads", "pick_rows_per_cta"} for node in ast.walk(torch_wrapper))) + torch_api_construction = next( + node for node in ast.walk(torch_wrapper) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RmsNormRhtAmaxSm100" + ) + self.assertFalse({"sample_o", "sample_amax"} & {keyword.arg for keyword in torch_api_construction.keywords}) + + torch_execute = next(node for node in torch_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "execute") + self.assertEqual( + [argument.arg for argument in torch_execute.args.args], + ["self", "x_tensor", "w_tensor", "current_stream"], + ) + self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "infer_output" for node in ast.walk(torch_execute))) + self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "materialize" for node in ast.walk(torch_execute))) kernel_class = next(node for node in kernel_tree.body if isinstance(node, ast.ClassDef) and node.name == "RMSNormRHTAmaxKernel") self.assertEqual([base.id for base in kernel_class.bases if isinstance(base, ast.Name)], ["OpKernel"]) @@ -259,6 +274,19 @@ def imports_framework(node, framework: str): ) compile_method = next(node for node in torch_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "compile") + self.assertTrue( + {"n", "rows_per_cta"}.issubset( + { + node.attr + for node in ast.walk(compile_method) + if isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Attribute) + and isinstance(node.value.value, ast.Name) + and node.value.value.id == "self" + and node.value.attr == "_kernel" + } + ) + ) tensor_api = next(node for node in ast.walk(compile_method) if isinstance(node, ast.FunctionDef) and node.name == "tensor_api") compiled_call = next( node for node in ast.walk(tensor_api) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "compiled_kernel" @@ -272,9 +300,12 @@ def imports_framework(node, framework: str): self.assertFalse( any(isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "cutlass_call" for node in ast.walk(jax_adapter)) ) - self.assertIn("_INSTALL_HINT", {node.id for node in ast.walk(jax_package_tree) if isinstance(node, ast.Name)}) - self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "is_available" for node in ast.walk(jax_package_tree))) - for tree in (jax_tree, jax_facade_tree): + self.assertIn("_INSTALL_HINT", {node.id for node in ast.walk(jax_facade_tree) if isinstance(node, ast.Name)}) + self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "is_available" for node in ast.walk(jax_facade_tree))) + self.assertTrue(any(isinstance(node, ast.Try) for node in jax_facade_tree.body)) + self.assertTrue(any(isinstance(node, ast.Import) and any(alias.name == "jax" for alias in node.names) for node in jax_tree.body)) + self.assertFalse(any(isinstance(node, ast.ImportFrom) and any(alias.name == "jax" for alias in node.names) for node in jax_tree.body)) + for tree in (jax_package_tree, jax_tree): self.assertNotIn("_INSTALL_HINT", {node.id for node in ast.walk(tree) if isinstance(node, ast.Name)}) self.assertFalse(any(isinstance(node, ast.Try) for node in tree.body)) self.assertFalse(any(isinstance(node, ast.Attribute) and node.attr == "is_available" for node in ast.walk(tree))) diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py b/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py index eb89130ff..4a0dc56a6 100644 --- a/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py +++ b/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py @@ -182,11 +182,52 @@ def test_kernel_implements_neutral_contract_and_infers_buffers(self): (amax.name, amax.dtype, amax.shape, amax.stride, amax.stride_order, amax.init_value), ("amax", _DataType.FLOAT, (128,), (1,), (0,), None), ) + + kernel.x = self._desc((128, 2048)) + output, amax = kernel.infer_output() + self.assertEqual((output.shape, amax.shape), ((256, 2048), (128,))) self.assertEqual(kernel.infer_workspace(), ()) self.assertNotIn(f"{_PACKAGE}.rmsnorm_rht_amax.api", sys.modules) self.assertNotIn(f"{_PACKAGE}.rmsnorm_rht_amax.jax", sys.modules) + def test_output_inference_preserves_descriptor_specialization(self): + tensor_module = self.tensor_module + + class SpecializedTensorDesc(tensor_module.TensorDesc): + def compact_like(self, **kwargs): + canonical = super().compact_like(**kwargs) + return SpecializedTensorDesc( + dtype=canonical.dtype, + shape=canonical.shape, + stride=canonical.stride, + stride_order=canonical.stride_order, + name=canonical.name, + init_value=canonical.init_value, + ) + + x = SpecializedTensorDesc( + dtype=_DataType.BFLOAT16, + shape=(256, 2048), + stride=(2048, 1), + stride_order=(1, 0), + ) + + kernel = self.kernel_module.RMSNormRHTAmaxKernel( + x=x, + weight=SpecializedTensorDesc( + dtype=_DataType.BFLOAT16, + shape=(2048,), + stride=(1,), + stride_order=(0,), + ), + ) + kernel.check_support() + output, amax = kernel.infer_output() + + self.assertIsInstance(output, SpecializedTensorDesc) + self.assertIsInstance(amax, SpecializedTensorDesc) + def test_requested_configuration_is_recomputed(self): kernel = self._kernel(num_threads=256, rows_per_cta=4) kernel.check_support() @@ -194,7 +235,6 @@ def test_requested_configuration_is_recomputed(self): kernel.requested_num_threads = 128 kernel.requested_rows_per_cta = 8 - kernel.amax = self._desc((32,), _DataType.FLOAT) kernel.check_support() self.assertEqual((kernel.num_threads, kernel.rows_per_cta), (128, 8)) self.assertEqual(kernel.infer_output()[1].shape, (32,)) @@ -221,8 +261,6 @@ def desc(shape, dtype): kernel = self.kernel_module.RMSNormRHTAmaxKernel( x=desc((256, 2048), "framework.bfloat16"), weight=desc((2048,), "framework.bfloat16"), - output=desc((256, 2048), "framework.bfloat16"), - amax=desc((128,), "framework.float32"), ) self.assertTrue(kernel.check_support()) @@ -253,10 +291,6 @@ def test_invalid_logical_signatures_and_configuration(self): ({"num_threads": 2048}, "must not exceed the CUDA block size limit"), ({"rows_per_cta": 0}, "rows_per_cta must be positive"), ({"rows_per_cta": 3}, "M must be divisible by rows_per_cta"), - ({"output": self._desc((256, 1024))}, "O must have shape"), - ({"output": self._desc((256, 2048), _DataType.FLOAT)}, "O must have dtype bfloat16"), - ({"amax": self._desc((64,), _DataType.FLOAT)}, "Amax must have shape"), - ({"amax": self._desc((128,))}, "Amax must have dtype float32"), ) for overrides, message in cases: diff --git a/test/python/fe_api/test_shared_tensor_desc.py b/test/python/fe_api/test_shared_tensor_desc.py index 332e91025..45d8b3bc3 100644 --- a/test/python/fe_api/test_shared_tensor_desc.py +++ b/test/python/fe_api/test_shared_tensor_desc.py @@ -87,6 +87,31 @@ def test_compact_descriptor_preserves_initial_value(self): init_value=[], ) + def test_compact_like_derives_a_canonical_descriptor(self): + source = self.tensor_desc.TensorDesc( + dtype=_DataType.BFLOAT16, + shape=(7,), + stride=(1,), + stride_order=(0,), + name="input", + ) + + derived = source.compact_like( + cudnn_dtype=_DataType.FLOAT, + shape=(2, 3), + stride_order=(0, 1), + name="workspace", + init_value=0, + ) + + self.assertIs(type(derived), self.tensor_desc.TensorDesc) + self.assertEqual(derived.dtype, _DataType.FLOAT) + self.assertEqual(derived.shape, (2, 3)) + self.assertEqual(derived.stride, (1, 2)) + self.assertEqual(derived.stride_order, (0, 1)) + self.assertEqual(derived.name, "workspace") + self.assertEqual(derived.init_value, 0) + def test_compact_descriptor_supports_an_explicit_dimension_order(self): desc = self.tensor_desc.make_compact_tensor_desc( dtype=_DataType.FLOAT, diff --git a/test/python/fe_api/test_tensor_desc.py b/test/python/fe_api/test_tensor_desc.py index 2c43e7f68..69fb7995b 100644 --- a/test/python/fe_api/test_tensor_desc.py +++ b/test/python/fe_api/test_tensor_desc.py @@ -106,6 +106,121 @@ def test_api_base_converts_torch_tensor_metadata(self): self.assertEqual(desc.stride, (12, 1, 3)) self.assertEqual(desc.stride_order, (1, 2, 0)) + def test_api_base_checks_torch_tensor_signature(self): + class Adapter(APIBase): + def check_support(self): + return True + + def compile(self): + pass + + def execute(self, *args, **kwargs): + pass + + adapter = Adapter() + sample = torch.empty_strided((2, 3), (3, 1), dtype=torch.bfloat16) + expected = TorchTensorDesc.from_tensor(sample, "sample") + adapter._check_tensor_signature(sample, expected) + + cases = ( + (torch.empty((1, 3), dtype=torch.bfloat16), "sample tensor shape mismatch"), + (torch.empty_strided((2, 3), (1, 2), dtype=torch.bfloat16), "sample tensor stride mismatch"), + (torch.empty((2, 3), dtype=torch.float32), "sample dtype mismatch"), + ) + for tensor, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + adapter._check_tensor_signature(tensor, expected) + + def test_torch_descriptor_constructs_from_tensor_metadata(self): + tensor = torch.empty_strided((2, 3, 4), (12, 1, 3), dtype=torch.uint8) + desc = TorchTensorDesc.from_tensor( + tensor, + "logical_fp4", + shape=(2, 3, 8), + stride=(12, 1, 3), + interpret_uint8_as_fp4x2=True, + ) + + self.assertEqual(desc.dtype, torch.uint8) + self.assertEqual(desc.shape, (2, 3, 8)) + self.assertEqual(desc.stride, (12, 1, 3)) + self.assertEqual(desc.stride_order, (1, 2, 0)) + self.assertEqual(desc.device, tensor.device) + self.assertEqual(desc.name, "logical_fp4") + self.assertTrue(desc.interpret_uint8_as_fp4x2) + self.assertEqual(desc.cudnn_dtype, data_type.FP4_E2M1) + self.assertEqual( + APIBase._to_tensor_desc( + tensor, + "logical_fp4", + shape=(2, 3, 8), + stride=(12, 1, 3), + interpret_uint8_as_fp4x2=True, + ), + desc, + ) + + def test_torch_descriptor_derives_and_materializes_compact_output(self): + source = TorchTensorDesc.from_tensor( + torch.empty((4, 8), dtype=torch.bfloat16), + "input", + ) + output = source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(3, 5), + stride_order=(0, 1), + name="output", + init_value=float("-inf"), + ) + + self.assertIsInstance(output, TorchTensorDesc) + self.assertEqual(output.dtype, torch.float32) + self.assertEqual(output.device, source.device) + self.assertEqual(output.shape, (3, 5)) + self.assertEqual(output.stride, (1, 3)) + self.assertEqual(output.stride_order, (0, 1)) + self.assertEqual(output.name, "output") + self.assertEqual(output.init_value, float("-inf")) + self.assertFalse(output.interpret_uint8_as_fp4x2) + + tensor = output.materialize() + self.assertEqual(tensor.dtype, torch.float32) + self.assertEqual(tuple(tensor.shape), output.shape) + self.assertEqual(tuple(tensor.stride()), output.stride) + self.assertTrue(torch.isneginf(tensor).all()) + + def test_torch_descriptor_materializes_false_and_zero_initializers(self): + source = TorchTensorDesc.from_tensor(torch.empty((1,), dtype=torch.float32)) + for init_value in (False, 0): + with self.subTest(init_value=init_value): + output = source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(4,), + init_value=init_value, + ) + self.assertTrue(torch.equal(output.materialize(), torch.zeros(4))) + + def test_torch_descriptor_rejects_materializing_logical_packed_storage(self): + packed = TorchTensorDesc.from_tensor( + torch.empty((4,), dtype=torch.uint8), + interpret_uint8_as_fp4x2=True, + ) + + with self.assertRaisesRegex(ValueError, "logical FP4"): + packed.materialize() + + unpacked = TorchTensorDesc( + dtype=torch.float32, + shape=(4,), + stride=(1,), + stride_order=(0,), + device="cpu", + interpret_uint8_as_fp4x2=True, + init_value=1, + ) + self.assertTrue(torch.equal(unpacked.materialize(), torch.ones(4))) + def test_api_base_materializes_a_canonical_descriptor(self): desc = SharedTensorDesc( dtype=data_type.FLOAT, @@ -148,6 +263,16 @@ def test_api_base_initializes_on_the_requested_cuda_stream(self): self.assertTrue(torch.equal(tensor.cpu(), torch.full((8,), 7.0))) + torch_desc = TorchTensorDesc.from_tensor(torch.empty((1,), dtype=torch.float32, device=device)).compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(8,), + init_value=9, + ) + tensor = torch_desc.materialize(stream=cuda.CUstream(stream.cuda_stream)) + stream.synchronize() + + self.assertTrue(torch.equal(tensor.cpu(), torch.full((8,), 9.0))) + if __name__ == "__main__": unittest.main() From 0cf7976ed0e6e2ccc641c02523fb2ce33450f538 Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Tue, 7 Jul 2026 16:00:11 -0500 Subject: [PATCH 03/13] Refactor RMSNorm APIs around framework-neutral operations Split logical validation from the CuTe kernel while preserving the explicit Torch compile/execute lifecycle. Add JAX-owned output allocation, device compatibility controls, lazy framework boundaries, and focused contract tests. --- docs/fe-oss-apis/rmsnorm_rht_amax.md | 85 +++++- python/cudnn/__init__.py | 2 +- python/cudnn/_jax/__init__.py | 4 +- python/cudnn/_jax/api_base.py | 97 ++++++- python/cudnn/_op.py | 17 ++ python/cudnn/_op_kernel.py | 47 --- python/cudnn/api_base.py | 16 -- python/cudnn/jax/__init__.py | 3 +- python/cudnn/rmsnorm_rht_amax/__init__.py | 9 +- python/cudnn/rmsnorm_rht_amax/api.py | 122 ++++---- python/cudnn/rmsnorm_rht_amax/jax.py | 80 +++++- python/cudnn/rmsnorm_rht_amax/kernel.py | 182 +----------- python/cudnn/rmsnorm_rht_amax/op.py | 193 +++++++++++++ test/python/fe_api/test_jax_api_base.py | 196 ++++++++++--- .../fe_api/test_jax_rmsnorm_rht_amax.py | 78 +++-- test/python/fe_api/test_rmsnorm_rht_amax.py | 128 ++++++--- .../fe_api/test_rmsnorm_rht_amax_imports.py | 79 ++++-- .../fe_api/test_rmsnorm_rht_amax_kernel.py | 267 +++++------------- .../python/fe_api/test_rmsnorm_rht_amax_op.py | 186 ++++++++++++ test/python/fe_api/test_tensor_desc.py | 26 -- 20 files changed, 1143 insertions(+), 674 deletions(-) create mode 100644 python/cudnn/_op.py delete mode 100644 python/cudnn/_op_kernel.py create mode 100644 python/cudnn/rmsnorm_rht_amax/op.py create mode 100644 test/python/fe_api/test_rmsnorm_rht_amax_op.py diff --git a/docs/fe-oss-apis/rmsnorm_rht_amax.md b/docs/fe-oss-apis/rmsnorm_rht_amax.md index 0f3a6d666..398e17122 100644 --- a/docs/fe-oss-apis/rmsnorm_rht_amax.md +++ b/docs/fe-oss-apis/rmsnorm_rht_amax.md @@ -10,6 +10,7 @@ This frontend integration exposes the kernel as a standard FE-OSS Python API wit - a class API (`RmsNormRhtAmaxSm100`) - a wrapper API (`rmsnorm_rht_amax_wrapper_sm100`) - an explicit JAX API under `cudnn.jax` +- a framework-neutral operation signature in `op.py` and a CUTE DSL implementation in `kernel.py` - grouped-gemm-style regression coverage for compile/execute, wrapper use, and cache reuse ## Shapes @@ -81,17 +82,28 @@ from cudnn import RmsNormRhtAmaxSm100 op = RmsNormRhtAmaxSm100( sample_x=x, sample_w=w, + sample_o=o, + sample_amax=amax, eps=1e-5, num_threads=128, rows_per_cta=2, ) -o, amax = op(x, w, current_stream=None) +assert op.check_support() +op.compile() +op.execute( + x_tensor=x, + w_tensor=w, + o_tensor=o, + amax_tensor=amax, + current_stream=None, +) ``` -The class infers and allocates its outputs from the validated sample input -descriptors. Runtime inputs must match the sample shape, layout, dtype, and -device. `check_support()` and `compile()` remain available when the caller -wants to run those steps explicitly; the call operator performs them lazily. +The PyTorch class preserves caller-owned output buffers. During +`check_support()`, the framework-neutral operation validates the complete +input and output exemplar signature and resolves the launch configuration. +`compile()` then constructs the CUTE DSL kernel from that resolved +configuration. ### JAX API @@ -102,7 +114,8 @@ pip install 'nvidia-cudnn-frontend[jax]' ``` ```python -from cudnn.jax import rmsnorm_rht_amax_sm100 +import jax +from cudnn.jax import RmsNormRhtAmaxSm100, rmsnorm_rht_amax_sm100 o, amax = rmsnorm_rht_amax_sm100( x, @@ -111,19 +124,47 @@ o, amax = rmsnorm_rht_amax_sm100( num_threads=128, rows_per_cta=2, ) + +# The class API specializes from abstract placeholders, then accepts arrays +# when called under JAX transformations. +x_spec = jax.ShapeDtypeStruct(x.shape, x.dtype) +w_spec = jax.ShapeDtypeStruct(w.shape, w.dtype) +op = RmsNormRhtAmaxSm100(x_spec, w_spec, rows_per_cta=2) +o, amax = jax.jit(op)(x, w) ``` The function API is JIT-compiled and treats `eps`, `num_threads`, and `rows_per_cta` as static compilation options. Use `RmsNormRhtAmaxSm100` -directly when the caller needs to control JAX transformations explicitly. Its -class API follows the same `RmsNormRhtAmaxSm100(sample_x, sample_w)(x, w)` -pattern as PyTorch. +directly when the caller needs to control JAX transformations explicitly. A +JAX array or another object exposing shape and dtype may be used instead of a +`ShapeDtypeStruct` exemplar; only its metadata is retained. Unlike PyTorch, +JAX derives output descriptors in its adapter and lets XLA allocate the +results. Callers may instead provide both `sample_o` and `sample_amax` as +explicit output placeholders. In either case, the common operation validates +the complete signature through `check_support()`. + +By default, JAX support validation requires every local GPU to meet the +operation's SM100 minimum. AOT or remote compilation without a compatible +local target device requires an explicit process-wide opt-out before tracing: + +```python +import cudnn.jax + +cudnn.jax.disable_device_compatibility_checks(True) +lowered = rmsnorm_rht_amax_sm100.lower(x_spec, w_spec) +``` -JAX buffers remain row-major. Framework-specific mode metadata maps public JAX axes into the canonical kernel axis order without materializing a transpose. +Pass `False` to re-enable the checks. Changing this setting does not +invalidate existing JAX compilation-cache entries, so configure it before the +first `jax.jit` or `lower` call for a signature. + +The current JAX adapter models its buffers as row-major with the public and +kernel axes in the same order. Additional JAX-specific axis mappings can be +added without changing the framework-neutral operation signature. ## Parameters -### Input tensors +### PyTorch input and output tensors - `x_tensor` / `sample_x` - Shape: `(M, N)` @@ -133,6 +174,16 @@ JAX buffers remain row-major. Framework-specific mode metadata maps public JAX a - Shape: `(N,)` - Layout: contiguous - Dtype: `torch.bfloat16` +- `o_tensor` / `sample_o` + - Shape: `(M, N)` + - Layout: row-major contiguous + - Dtype: `torch.bfloat16` +- `amax_tensor` / `sample_amax` + - Shape: `(M / rows_per_cta,)` + - Dtype: `torch.float32` + +The JAX API uses the same shapes and layouts with `jax.numpy.bfloat16` for +`X`, `W`, and `O`, and `jax.numpy.float32` for `Amax`. ### Common parameters @@ -142,11 +193,14 @@ JAX buffers remain row-major. Framework-specific mode metadata maps public JAX a - Threads per CTA. If omitted, the API uses the upstream-tuned table when possible, otherwise a valid fallback search. - `rows_per_cta: Optional[int]` - Rows processed by each CTA. If omitted, the wrapper uses the upstream-style heuristic over `{2, 4, 8}`. -- CUDA stream (`current_stream`) + +The PyTorch `execute()` and wrapper APIs additionally accept the optional +CUDA stream argument `current_stream`. JAX execution uses the stream supplied +by XLA. ### Return values -The PyTorch class and wrapper return a `TupleDict` with keys: +The PyTorch wrapper returns a `TupleDict` with keys: - `o_tensor` - Shape: `(M, N)` @@ -158,6 +212,9 @@ The PyTorch class and wrapper return a `TupleDict` with keys: Tuple unpacking order is `(o_tensor, amax_tensor)`. +The class API writes into its caller-provided output tensors and returns +`None`. + ## Support surface and constraints - Requires SM100+. @@ -166,7 +223,7 @@ Tuple unpacking order is `(o_tensor, amax_tensor)`. - `EPT = N / num_threads` must be at least `8` and divisible by `8`. - `M` must be divisible by `rows_per_cta`. - `X`, `W`, and `O` use bf16; `amax` uses float32. -- PyTorch input tensors must be 16-byte aligned. +- PyTorch input and output tensors must be 16-byte aligned. - The frontend integration matches the upstream RMSNorm kernel semantics; it does not expose full LayerNorm mean/bias behavior. ## Verification diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index e77cef0e5..04f771209 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -51,7 +51,7 @@ def is_windows(): if hasattr(_pybind_module, _optional_symbol): globals()[_optional_symbol] = getattr(_pybind_module, _optional_symbol) -from ._op_kernel import OpKernel +from ._op import Op from ._tensor_desc import TensorDesc from .datatypes import _library_type, _is_torch_tensor diff --git a/python/cudnn/_jax/__init__.py b/python/cudnn/_jax/__init__.py index 3ff2b4ce8..e2ea61e8e 100644 --- a/python/cudnn/_jax/__init__.py +++ b/python/cudnn/_jax/__init__.py @@ -3,6 +3,6 @@ """Internal implementation helpers for optional JAX operation adapters.""" -from .api_base import JaxApiBase, JaxTensorDesc +from .api_base import JaxApiBase, JaxTensorDesc, disable_device_compatibility_checks -__all__ = ["JaxApiBase", "JaxTensorDesc"] +__all__ = ["JaxApiBase", "JaxTensorDesc", "disable_device_compatibility_checks"] diff --git a/python/cudnn/_jax/api_base.py b/python/cudnn/_jax/api_base.py index 72e9551c8..7ed7c76ff 100644 --- a/python/cudnn/_jax/api_base.py +++ b/python/cudnn/_jax/api_base.py @@ -10,10 +10,26 @@ from typing import Any from .. import data_type -from .._op_kernel import OpKernel from .._tensor_desc import TensorDesc from .layout import compact_stride, normalize_mode, to_canonical_axes, to_cutlass_layout, to_public_axes +_DEVICE_COMPATIBILITY_CHECKS_DISABLED = False +_DEVICE_COMPATIBILITY_CHECK_DISABLE_HINT = "Call cudnn.jax.disable_device_compatibility_checks(True) before tracing to allow AOT or remote compilation." + + +def disable_device_compatibility_checks(disabled: bool) -> None: + """Enable or disable local-device compatibility checks for future traces. + + This process-global setting does not invalidate existing JAX compilation + cache entries. Configure it before calling ``jax.jit`` or ``lower``. + """ + + if not isinstance(disabled, bool): + raise TypeError(f"disabled must be a bool, got {type(disabled).__name__}") + + global _DEVICE_COMPATIBILITY_CHECKS_DISABLED + _DEVICE_COMPATIBILITY_CHECKS_DISABLED = disabled + class JaxTensorDesc(TensorDesc[Any]): """Framework-neutral tensor metadata backed by a JAX dtype.""" @@ -62,9 +78,52 @@ def _require_array_metadata(value: Any, name: str) -> tuple[Any, ...]: class JaxApiBase(ABC): - """Common tensor metadata, validation, and kernel binding for JAX adapters.""" + """Common tensor metadata, validation, and CuTe binding for JAX adapters.""" + + @staticmethod + def _check_device_compatibility( + *, + minimum_compute_capability: int, + operation_name: str, + ) -> None: + """Require every local JAX GPU to satisfy an operation's minimum SM.""" + + if _DEVICE_COMPATIBILITY_CHECKS_DISABLED: + return + + import jax - kernel: OpKernel + def compatibility_error(reason: str) -> RuntimeError: + return RuntimeError(f"{operation_name} requires SM{minimum_compute_capability}+, {reason}. " f"{_DEVICE_COMPATIBILITY_CHECK_DISABLE_HINT}") + + try: + devices = tuple(jax.local_devices(backend="gpu")) + except RuntimeError as error: + raise compatibility_error("but JAX could not discover a local GPU") from error + + if not devices: + raise compatibility_error("but no local JAX GPU is available") + + minimum_major, minimum_minor = divmod(minimum_compute_capability, 10) + incompatible = [] + for device in devices: + reported_capability = getattr(device, "compute_capability", None) + try: + components = str(reported_capability).split(".") + if len(components) != 2: + raise ValueError + major, minor = (int(value) for value in components) + if major < 0 or minor < 0: + raise ValueError + except (TypeError, ValueError) as error: + raise compatibility_error(f"but JAX reported an invalid compute capability {reported_capability!r} for {device}") from error + + if (major, minor) < (minimum_major, minimum_minor): + incompatible.append((device, major, minor)) + + if incompatible: + found = ", ".join(f"{device} (SM{major}{minor})" for device, major, minor in incompatible) + raise compatibility_error(f"found {found}") @staticmethod def _to_tensor_desc( @@ -72,6 +131,7 @@ def _to_tensor_desc( name: str, *, mode: tuple[int, ...] | None = None, + init_value: bool | int | float | None = None, ) -> JaxTensorDesc: """Describe a public JAX array in canonical kernel-axis order. @@ -92,6 +152,7 @@ def _to_tensor_desc( stride=to_canonical_axes(public_stride, mode), stride_order=tuple(canonical_axis_by_public_axis[axis] for axis in public_stride_order), name=name, + init_value=init_value, ) @staticmethod @@ -181,6 +242,8 @@ def _call_kernel( self, inputs: tuple[Any, ...], *, + output_descs: tuple[TensorDesc[Any], ...], + workspace_descs: tuple[TensorDesc[Any], ...] = (), input_spec: tuple[Any | None, ...] | None = None, output_spec: tuple[Any | None, ...] | None = None, workspace_spec: tuple[Any | None, ...] | None = None, @@ -188,10 +251,10 @@ def _call_kernel( compile_options: Any = None, use_static_tensors: bool = True, ) -> tuple[Any, ...]: - """Bind the owned kernel to JAX and return its public outputs. + """Bind this adapter's launch hook to JAX and return public outputs. CUTLASS JAX supplies the stream first; the launcher adapts that call to - the kernel's ``inputs, outputs, workspaces, stream`` order. Workspaces + the adapter's ``inputs, outputs, workspaces, stream`` launch hook. Workspaces are declared as custom-call results so XLA owns their lifetime, then omitted from this method's return value. Descriptors with a non-``None`` ``init_value`` are materialized as JAX inputs and aliased to their @@ -205,15 +268,15 @@ def _call_kernel( inputs = tuple(inputs) for index, value in enumerate(inputs): _require_array_metadata(value, f"input #{index}") - outputs = tuple(self.kernel.infer_output()) - workspaces = tuple(self.kernel.infer_workspace()) + outputs = tuple(output_descs) + workspaces = tuple(workspace_descs) if not outputs: - raise ValueError("A JAX operation must infer at least one output") + raise ValueError("A JAX operation must provide at least one output descriptor") for label, descs in (("output", outputs), ("workspace", workspaces)): for desc in descs: if not isinstance(desc, TensorDesc): - raise TypeError(f"{label} inference must return TensorDesc values, got {type(desc).__name__}") + raise TypeError(f"{label}_descs must contain TensorDesc values, got {type(desc).__name__}") if input_spec is None: input_specs = (None,) * len(inputs) @@ -273,7 +336,9 @@ def launcher(stream, *args): ordered_buffers[index] = value for index, value in zip(uninitialized, allocated_buffers): ordered_buffers[index] = value - self.kernel(*kernel_inputs, *ordered_buffers, stream) + output_buffers = tuple(ordered_buffers[: len(outputs)]) + workspace_buffers = tuple(ordered_buffers[len(outputs) :]) + self._launch(tuple(kernel_inputs), output_buffers, workspace_buffers, stream) call = cutlass_jax.cutlass_call( launcher, @@ -290,6 +355,16 @@ def launcher(stream, *args): results = (results,) return tuple(results[: len(outputs)]) + @abstractmethod + def _launch( + self, + inputs: tuple[Any, ...], + outputs: tuple[Any, ...], + workspaces: tuple[Any, ...], + stream: Any, + ) -> None: + """Launch the concrete CuTe kernel for one traced custom call.""" + @abstractmethod def check_support(self) -> bool: """Validate the operation signature and static configuration.""" @@ -304,4 +379,4 @@ def get_jax_callable(self) -> Callable[..., Any]: return self -__all__ = ["JaxApiBase", "JaxTensorDesc"] +__all__ = ["JaxApiBase", "JaxTensorDesc", "disable_device_compatibility_checks"] diff --git a/python/cudnn/_op.py b/python/cudnn/_op.py new file mode 100644 index 000000000..4c5f97712 --- /dev/null +++ b/python/cudnn/_op.py @@ -0,0 +1,17 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral operation contract.""" + +from abc import ABC, abstractmethod + + +class Op(ABC): + """Logical operation specialized from a complete tensor signature.""" + + @abstractmethod + def check_support(self) -> bool: + """Validate the operation signature and resolve static configuration.""" + + +__all__ = ["Op"] diff --git a/python/cudnn/_op_kernel.py b/python/cudnn/_op_kernel.py deleted file mode 100644 index 350fdd961..000000000 --- a/python/cudnn/_op_kernel.py +++ /dev/null @@ -1,47 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Framework-neutral contract for operation kernels.""" - -from __future__ import annotations - -from abc import ABC, abstractmethod -from typing import Any - -from ._tensor_desc import TensorDesc - - -class OpKernel(ABC): - """Host-side contract implemented by framework-agnostic kernels. - - Implementations own logical validation and tensor inference. Framework - adapters remain responsible for turning the inferred descriptors into - framework tensors or abstract output declarations and for compiling and - executing the kernel. Callable implementations receive inputs, outputs, - workspaces, and finally the CUDA stream. Framework adapters may reorder - those arguments to match a framework-specific launch ABI. - """ - - @abstractmethod - def check_support(self) -> bool: - """Validate the logical signature and resolve static kernel state.""" - - @abstractmethod - def infer_output(self) -> tuple[TensorDesc[Any], ...]: - """Return operation outputs through the framework-neutral descriptor interface. - - A descriptor's ``init_value`` requests scalar initialization before - launch; ``None`` leaves the output uninitialized. - """ - - def infer_workspace(self) -> tuple[TensorDesc[Any], ...]: - """Return workspaces through the framework-neutral descriptor interface. - - Workspace descriptors follow the same ``init_value`` convention as - outputs but are not part of the framework-visible result. - """ - - return () - - -__all__ = ["OpKernel"] diff --git a/python/cudnn/api_base.py b/python/cudnn/api_base.py index f15c763d2..97f45da6f 100644 --- a/python/cudnn/api_base.py +++ b/python/cudnn/api_base.py @@ -949,22 +949,6 @@ def _check_dtype( raise ValueError(f"Expected dtype to be a torch.dtype or list, got {type(dtype)}") return tensor_dtype - def _check_tensor_signature(self, tensor: torch.Tensor, expected: TensorDesc) -> None: - """Validate a Torch tensor against a captured logical descriptor.""" - - if not isinstance(expected, TensorDesc): - raise TypeError(f"expected must be a TensorDesc, got {type(expected).__name__}") - name = expected.name or "tensor" - actual = self._make_tensor_desc( - tensor, - name=name, - interpret_uint8_as_fp4x2=expected.interpret_uint8_as_fp4x2, - ) - self._check_tensor_shape(actual, expected.shape, name) - self._check_tensor_stride(actual, stride=expected.stride, stride_order=expected.stride_order, name=name) - self._check_dtype(actual, expected.dtype, name) - self._value_error_if(actual.device != expected.device, f"{name} tensor device mismatch: expected {expected.device}, got {actual.device}") - def _value_error_if(self, condition: bool, error_msg: str) -> None: """Raise a ValueError if the condition is true. diff --git a/python/cudnn/jax/__init__.py b/python/cudnn/jax/__init__.py index 20ab5a1a7..453db9d8b 100644 --- a/python/cudnn/jax/__init__.py +++ b/python/cudnn/jax/__init__.py @@ -24,7 +24,7 @@ f"Install the optional JAX dependencies with `{_INSTALL_HINT}`." ) -from .._jax import JaxApiBase, JaxTensorDesc +from .._jax import JaxApiBase, JaxTensorDesc, disable_device_compatibility_checks from ..rmsnorm_rht_amax.jax import ( RmsNormRhtAmaxSm100, rmsnorm_rht_amax_sm100, @@ -33,6 +33,7 @@ __all__ = [ "JaxApiBase", "JaxTensorDesc", + "disable_device_compatibility_checks", "RmsNormRhtAmaxSm100", "rmsnorm_rht_amax_sm100", ] diff --git a/python/cudnn/rmsnorm_rht_amax/__init__.py b/python/cudnn/rmsnorm_rht_amax/__init__.py index 9f11062da..71ee075c9 100644 --- a/python/cudnn/rmsnorm_rht_amax/__init__.py +++ b/python/cudnn/rmsnorm_rht_amax/__init__.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Lazy Torch API and framework-neutral kernel exports. +"""Lazy Torch API and framework-neutral operation/kernel exports. JAX APIs are exported through :mod:`cudnn.jax`. """ @@ -11,15 +11,16 @@ __all__, __getattr__, __dir__ = make_operation_api( globals(), exports={ - "kernel": ( - "RMSNormRHTAmaxKernel", + "op": ( + "RmsNormRhtAmaxSm100Op", "best_num_threads", "pick_rows_per_cta", ), + "kernel": ("RMSNormRHTAmaxKernel",), "api": ( "RmsNormRhtAmaxSm100", "rmsnorm_rht_amax_wrapper_sm100", ), }, - submodules=("api", "kernel"), + submodules=("api", "kernel", "op"), ) diff --git a/python/cudnn/rmsnorm_rht_amax/api.py b/python/cudnn/rmsnorm_rht_amax/api.py index acab89a6a..933b88364 100644 --- a/python/cudnn/rmsnorm_rht_amax/api.py +++ b/python/cudnn/rmsnorm_rht_amax/api.py @@ -13,29 +13,15 @@ from cudnn.api_base import APIBase, TupleDict -from .kernel import ( +from .kernel import RMSNormRHTAmaxKernel +from .op import ( DEFAULT_NUM_THREADS_BY_N, - RPC_CANDIDATES, - TARGET_MIN_CTAS, - RMSNormRHTAmaxKernel, + RmsNormRhtAmaxSm100Op, best_num_threads, pick_rows_per_cta, ) -_TENSOR_ALIGNMENT = RMSNormRHTAmaxKernel.TENSOR_ALIGNMENT - - -def _validate_torch_tensors(*named_tensors: tuple[str, torch.Tensor]) -> None: - """Validate framework properties that are absent from TensorDesc.""" - - if not named_tensors: - return - expected_device = named_tensors[0][1].device - for name, tensor in named_tensors: - if tensor.device != expected_device: - raise ValueError(f"{name} must be on {expected_device}, got {tensor.device}") - if tensor.data_ptr() % _TENSOR_ALIGNMENT != 0: - raise ValueError(f"{name} must be {_TENSOR_ALIGNMENT}-byte aligned") +_TENSOR_ALIGNMENT = RMSNormRHTAmaxKernel.COPY_BITS // 8 class RmsNormRhtAmaxSm100(APIBase): @@ -45,6 +31,8 @@ def __init__( self, sample_x: torch.Tensor, sample_w: torch.Tensor, + sample_o: torch.Tensor, + sample_amax: torch.Tensor, eps: float = 1e-5, num_threads: Optional[int] = None, rows_per_cta: Optional[int] = None, @@ -52,21 +40,21 @@ def __init__( super().__init__() self._warn_experimental_api() - _validate_torch_tensors( - ("sample_x", sample_x), - ("sample_w", sample_w), - ) self.x_desc = self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_x, name="sample_x"), 2, "sample_x") self.w_desc = self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_w, name="sample_w"), 1, "sample_w") + self.o_desc = self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_o, name="sample_o"), 2, "sample_o") + self.amax_desc = self._unpad_tensor_to_ndim(self._make_tensor_desc(sample_amax, name="sample_amax"), 1, "sample_amax") self.eps = eps self.requested_num_threads = num_threads self.requested_rows_per_cta = rows_per_cta - self._kernel = RMSNormRHTAmaxKernel( + self._op = RmsNormRhtAmaxSm100Op( x=self.x_desc, weight=self.w_desc, + output=self.o_desc, + amax=self.amax_desc, eps=eps, num_threads=num_threads, rows_per_cta=rows_per_cta, @@ -74,7 +62,7 @@ def __init__( def check_support(self) -> bool: self._is_supported = False - self._kernel.check_support() + self._op.check_support() self._runtime_error_if(not torch.cuda.is_available(), "CUDA is not available") major, minor = torch.cuda.get_device_capability(self.x_desc.device) @@ -92,37 +80,43 @@ def compile(self) -> None: if self._compiled_kernel is not None: return - o_desc, amax_desc = self._kernel.infer_output() - valid_m = cute.sym_int(divisibility=self._kernel.rows_per_cta) + kernel = RMSNormRHTAmaxKernel( + n=self._op.n, + num_threads=self._op.num_threads, + eps=self._op.eps, + rows_per_cta=self._op.rows_per_cta, + ) + + valid_m = cute.sym_int(divisibility=self._op.rows_per_cta) fake_x_tensor = self._make_fake_cute_compact_tensor( dtype=self.x_desc.dtype, - shape=(valid_m, self._kernel.n), + shape=(valid_m, self._op.n), stride_order=self.x_desc.stride_order, assumed_align=_TENSOR_ALIGNMENT, dynamic_mode=None, - divisibility=self._kernel.rows_per_cta, + divisibility=self._op.rows_per_cta, ) fake_w_tensor = self._make_fake_cute_tensor_from_desc(self.w_desc, assumed_align=_TENSOR_ALIGNMENT) fake_o_tensor = self._make_fake_cute_compact_tensor( - dtype=o_desc.dtype, - shape=(valid_m, self._kernel.n), - stride_order=o_desc.stride_order, + dtype=self.o_desc.dtype, + shape=(valid_m, self._op.n), + stride_order=self.o_desc.stride_order, assumed_align=_TENSOR_ALIGNMENT, dynamic_mode=None, - divisibility=self._kernel.rows_per_cta, + divisibility=self._op.rows_per_cta, ) fake_num_ctas = cute.sym_int() fake_amax_tensor = self._make_fake_cute_tensor( - dtype=amax_desc.dtype, + dtype=self.amax_desc.dtype, shape=(fake_num_ctas,), - stride=amax_desc.stride, + stride=self.amax_desc.stride, assumed_align=_TENSOR_ALIGNMENT, ) fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) compiled_kernel = cute.compile( - self._kernel, + kernel, fake_x_tensor, fake_w_tensor, fake_o_tensor, @@ -152,26 +146,25 @@ def execute( self, x_tensor: torch.Tensor, w_tensor: torch.Tensor, + o_tensor: torch.Tensor, + amax_tensor: torch.Tensor, current_stream: Optional[cuda.CUstream] = None, - ) -> TupleDict: + ) -> None: self._runtime_error_if(self._compiled_kernel is None, "RmsNormRhtAmaxSm100 kernel not compiled; call compile() first") x_tensor = self._unpad_tensor_to_ndim(x_tensor, 2, "x_tensor") w_tensor = self._unpad_tensor_to_ndim(w_tensor, 1, "w_tensor") - _validate_torch_tensors( - ("x_tensor", x_tensor), - ("w_tensor", w_tensor), - ) - self._check_tensor_signature(x_tensor, self.x_desc) - self._check_tensor_signature(w_tensor, self.w_desc) + o_tensor = self._unpad_tensor_to_ndim(o_tensor, 2, "o_tensor") + amax_tensor = self._unpad_tensor_to_ndim(amax_tensor, 1, "amax_tensor") + + # TVM-FFI validates the compiled tensor ABI. The amax extent is an + # independent symbol, so its relationship to M is checked here. + if x_tensor.ndim == 2: + self._check_tensor_shape(amax_tensor, (x_tensor.shape[0] // self._op.rows_per_cta,), "Amax") if current_stream is None: current_stream = cuda.CUstream(torch.cuda.current_stream(x_tensor.device).cuda_stream) - o_desc, amax_desc = self._kernel.infer_output() - o_tensor = o_desc.materialize(stream=current_stream) - amax_tensor = amax_desc.materialize(stream=current_stream) - self._compiled_kernel( x_tensor=x_tensor, w_tensor=w_tensor, @@ -179,7 +172,6 @@ def execute( amax_tensor=amax_tensor, stream=current_stream, ) - return TupleDict(o_tensor=o_tensor, amax_tensor=amax_tensor) _logger = logging.getLogger(__name__) @@ -199,34 +191,54 @@ def rmsnorm_rht_amax_wrapper_sm100( x_tensor = x_tensor.squeeze(-1) if x_tensor.ndim == 3 and x_tensor.shape[-1] == 1 else x_tensor w_tensor = w_tensor.squeeze(-1) if w_tensor.ndim == 2 and w_tensor.shape[-1] == 1 else w_tensor + m, n = x_tensor.shape + resolved_num_threads = num_threads if num_threads is not None else DEFAULT_NUM_THREADS_BY_N.get(n, best_num_threads(n)) + if resolved_num_threads is None: + raise ValueError(f"No valid num_threads found for N={n}") + resolved_rows_per_cta = rows_per_cta if rows_per_cta is not None else pick_rows_per_cta(m) + if resolved_rows_per_cta <= 0: + raise ValueError(f"rows_per_cta must be positive, got {resolved_rows_per_cta}") + if m % resolved_rows_per_cta != 0: + raise ValueError(f"M must be divisible by rows_per_cta, got M={m}, rows_per_cta={resolved_rows_per_cta}") + + o_tensor = torch.empty_like(x_tensor) + amax_tensor = torch.full((m // resolved_rows_per_cta,), float("-inf"), dtype=torch.float32, device=x_tensor.device) + cache_key = ( - tuple(x_tensor.shape), - tuple(w_tensor.shape), + n, x_tensor.dtype, w_tensor.dtype, + o_tensor.dtype, tuple(x_tensor.stride()), tuple(w_tensor.stride()), + tuple(o_tensor.stride()), x_tensor.device, eps, - num_threads, - rows_per_cta, + resolved_num_threads, + resolved_rows_per_cta, ) api = _cache_of_RmsNormRhtAmaxSm100Objects.get(cache_key) - if api is None: + if api is None: api = RmsNormRhtAmaxSm100( sample_x=x_tensor, sample_w=w_tensor, + sample_o=o_tensor, + sample_amax=amax_tensor, eps=eps, - num_threads=num_threads, - rows_per_cta=rows_per_cta, + num_threads=resolved_num_threads, + rows_per_cta=resolved_rows_per_cta, ) api.check_support() api.compile() _cache_of_RmsNormRhtAmaxSm100Objects[cache_key] = api - return api.execute( + api.execute( x_tensor=x_tensor, w_tensor=w_tensor, + o_tensor=o_tensor, + amax_tensor=amax_tensor, current_stream=current_stream, ) + + return TupleDict(o_tensor=o_tensor, amax_tensor=amax_tensor) diff --git a/python/cudnn/rmsnorm_rht_amax/jax.py b/python/cudnn/rmsnorm_rht_amax/jax.py index 4fbcc0eb2..a64fe1da0 100644 --- a/python/cudnn/rmsnorm_rht_amax/jax.py +++ b/python/cudnn/rmsnorm_rht_amax/jax.py @@ -9,8 +9,10 @@ import jax -from .._jax import JaxApiBase +from .. import data_type +from .._jax import JaxApiBase, JaxTensorDesc from .kernel import RMSNormRHTAmaxKernel +from .op import RmsNormRhtAmaxSm100Op, pick_rows_per_cta class RmsNormRhtAmaxSm100(JaxApiBase): @@ -20,22 +22,62 @@ def __init__( self, sample_x: Any, sample_w: Any, + *, + sample_o: Any | None = None, + sample_amax: Any | None = None, eps: float = 1e-5, num_threads: Optional[int] = None, rows_per_cta: Optional[int] = None, ) -> None: self.x_desc = self._to_tensor_desc(sample_x, "sample_x") self.w_desc = self._to_tensor_desc(sample_w, "sample_w") - self.kernel = RMSNormRHTAmaxKernel( + if (sample_o is None) != (sample_amax is None): + raise ValueError("sample_o and sample_amax must be provided together") + + if sample_o is None: + self.o_desc, self.amax_desc = self._default_output_descs(rows_per_cta) + else: + self.o_desc = self._to_tensor_desc(sample_o, "sample_o") + self.amax_desc = self._to_tensor_desc(sample_amax, "sample_amax") + + self._op = RmsNormRhtAmaxSm100Op( x=self.x_desc, weight=self.w_desc, + output=self.o_desc, + amax=self.amax_desc, eps=eps, num_threads=num_threads, rows_per_cta=rows_per_cta, ) + def _default_output_descs(self, rows_per_cta: Optional[int]) -> tuple[JaxTensorDesc, JaxTensorDesc]: + if self.x_desc.ndim != 2: + raise ValueError(f"X must have rank 2, got shape {self.x_desc.shape}") + m, n = self.x_desc.shape + resolved_rows_per_cta = pick_rows_per_cta(m) if rows_per_cta is None else rows_per_cta + if resolved_rows_per_cta <= 0: + raise ValueError(f"rows_per_cta must be positive, got {resolved_rows_per_cta}") + + return ( + self.x_desc.compact_like( + cudnn_dtype=data_type.BFLOAT16, + shape=(m, n), + name="sample_o", + ), + self.x_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(m // resolved_rows_per_cta,), + name="sample_amax", + ), + ) + def check_support(self) -> bool: - return self.kernel.check_support() + self._op.check_support() + self._check_device_compatibility( + minimum_compute_capability=100, + operation_name="RmsNormRhtAmaxSm100", + ) + return True def __call__(self, x: Any, weight: Any) -> tuple[Any, Any]: self.check_support() @@ -44,14 +86,36 @@ def __call__(self, x: Any, weight: Any) -> tuple[Any, Any]: x_spec = self._to_tensor_spec( self.x_desc, - divisibility=(self.kernel.rows_per_cta, 16), + divisibility=(self._op.rows_per_cta, 16), ) weight_spec = self._to_tensor_spec(self.w_desc, divisibility=(16,)) + output_spec = self._to_tensor_spec( + self.o_desc, + divisibility=(self._op.rows_per_cta, 16), + ) return self._call_kernel( (x, weight), + output_descs=(self.o_desc, self.amax_desc), input_spec=(x_spec, weight_spec), - output_spec=(x_spec, None), + output_spec=(output_spec, None), + ) + + def _launch( + self, + inputs: tuple[Any, ...], + outputs: tuple[Any, ...], + workspaces: tuple[Any, ...], + stream: Any, + ) -> None: + x, weight = inputs + output, amax = outputs + kernel = RMSNormRHTAmaxKernel( + n=self._op.n, + num_threads=self._op.num_threads, + eps=self._op.eps, + rows_per_cta=self._op.rows_per_cta, ) + kernel(x, weight, output, amax, stream) @jax.jit(static_argnames=("eps", "num_threads", "rows_per_cta")) @@ -65,9 +129,11 @@ def rmsnorm_rht_amax_sm100( ) -> tuple[Any, Any]: """Apply fused RMSNorm, 16-wide RHT, and per-CTA amax from JAX.""" + sample_x = jax.ShapeDtypeStruct(x.shape, x.dtype) + sample_weight = jax.ShapeDtypeStruct(weight.shape, weight.dtype) return RmsNormRhtAmaxSm100( - x, - weight, + sample_x, + sample_weight, eps=eps, num_threads=num_threads, rows_per_cta=rows_per_cta, diff --git a/python/cudnn/rmsnorm_rht_amax/kernel.py b/python/cudnn/rmsnorm_rht_amax/kernel.py index f00c90e81..b0ca5c3bb 100644 --- a/python/cudnn/rmsnorm_rht_amax/kernel.py +++ b/python/cudnn/rmsnorm_rht_amax/kernel.py @@ -7,51 +7,16 @@ import math import operator -from typing import Any, Optional import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute import cutlass.utils as utils -from cutlass import Float32, Int32 +from cutlass import Float32 from cutlass._mlir.dialects import llvm from cutlass.cute.arch import shuffle_sync_bfly from cutlass.cutlass_dsl import T, dsl_user_op -from .. import data_type -from .._op_kernel import OpKernel -from .._tensor_desc import TensorDesc - -DEFAULT_NUM_THREADS_BY_N = { - 2048: 128, - 4096: 256, - 7168: 128, - 8192: 512, - 16384: 1024, - 32768: 512, -} -RPC_CANDIDATES = (2, 4, 8) -TARGET_MIN_CTAS = 148 - - -def best_num_threads(n: int) -> Optional[int]: - for num_threads in (1024, 512, 256, 128, 64): - if n % num_threads != 0: - continue - ept = n // num_threads - if ept >= 8 and ept % 8 == 0: - return num_threads - return None - - -def pick_rows_per_cta(m: int) -> int: - for rows_per_cta in reversed(RPC_CANDIDATES): - if m % rows_per_cta != 0: - continue - if m // rows_per_cta >= TARGET_MIN_CTAS: - return rows_per_cta - return RPC_CANDIDATES[0] - @dsl_user_op def fabs_f32(val, *, loc=None, ip=None): @@ -105,155 +70,24 @@ def redux_sync_max_f32(val, *, loc=None, ip=None): return Float32(result) -class RMSNormRHTAmaxKernel(OpKernel): - """Fused RMSNorm + block-diagonal Hadamard + running per-CTA amax. - - Output inference preserves the framework descriptor specialization of the - input while validation and launch configuration remain framework-neutral. - """ +class RMSNormRHTAmaxKernel: + """CuTe implementation of fused RMSNorm + RHT + per-CTA amax.""" COPY_BITS = 128 HAD_BLOCK = 16 - TENSOR_ALIGNMENT = COPY_BITS // 8 - output_dtype = data_type.BFLOAT16 - amax_dtype = data_type.FLOAT def __init__( self, *, - x: TensorDesc[Any], - weight: TensorDesc[Any], - eps: float = 1e-5, - num_threads: Optional[int] = None, - rows_per_cta: Optional[int] = None, - ) -> None: - for name, desc in (("x", x), ("weight", weight)): - if not isinstance(desc, TensorDesc): - raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") - - self.x = x - self.weight = weight - self.eps = eps - self.requested_num_threads = num_threads - self.requested_rows_per_cta = rows_per_cta - - self.m: Optional[int] = None - self.n: Optional[int] = None - self.num_threads: Optional[int] = None - self.rows_per_cta: Optional[int] = None - self._validated_x: Optional[TensorDesc[Any]] = None - - def check_support(self) -> bool: - """Validate descriptors and resolve the kernel launch configuration.""" - - self.m = None - self.n = None - self.num_threads = None - self.rows_per_cta = None - self._validated_x = None - - if self.x.ndim != 2: - raise ValueError(f"X must have rank 2, got shape {self.x.shape}") - if self.weight.ndim != 1: - raise ValueError(f"W must have rank 1, got shape {self.weight.shape}") - if self.x.cudnn_dtype != data_type.BFLOAT16: - raise ValueError(f"X must have dtype bfloat16, got {self.x.dtype}") - if self.weight.cudnn_dtype != data_type.BFLOAT16: - raise ValueError(f"W must have dtype bfloat16, got {self.weight.dtype}") - - m, n = self.x.shape - if m <= 0: - raise ValueError(f"M must be positive, got {m}") - if n <= 0: - raise ValueError(f"N must be positive, got {n}") - if self.weight.shape != (n,): - raise ValueError(f"W must have shape {(n,)}, got {self.weight.shape}") - if self.x.stride != (n, 1) or self.x.stride_order != (1, 0): - raise ValueError(f"X must be row-major contiguous, got stride {self.x.stride} " f"and stride order {self.x.stride_order}") - if self.weight.stride != (1,) or self.weight.stride_order != (0,): - raise ValueError(f"W must be contiguous, got stride {self.weight.stride} " f"and stride order {self.weight.stride_order}") - if n % self.HAD_BLOCK != 0: - raise ValueError(f"N must be divisible by {self.HAD_BLOCK} for the Hadamard block size, got {n}") - - num_threads = self.requested_num_threads - if num_threads is None: - num_threads = DEFAULT_NUM_THREADS_BY_N.get(n, best_num_threads(n)) - if num_threads is None: - raise ValueError(f"No valid num_threads found for N={n}") - - rows_per_cta = self._resolve_rows_per_cta(m, self.requested_rows_per_cta) - - self._validate_launch_configuration(n, num_threads, rows_per_cta, m=m) - - self.m = m - self.n = n - self.num_threads = num_threads - self.rows_per_cta = rows_per_cta - self._configure_lowering_state() - self._validated_x = self.x - return True - - def infer_output(self) -> tuple[TensorDesc[Any], ...]: - """Infer compact output descriptors from validated input metadata.""" - - if self.m is None or self.n is None or self.rows_per_cta is None or self._validated_x is None: - raise RuntimeError("check_support() must be called before inferring outputs") - - return ( - self._validated_x.compact_like( - cudnn_dtype=self.output_dtype, - shape=(self.m, self.n), - name="output", - ), - self._validated_x.compact_like( - cudnn_dtype=self.amax_dtype, - shape=(self.m // self.rows_per_cta,), - name="amax", - ), - ) - - @staticmethod - def _resolve_rows_per_cta(m: int, rows_per_cta: Optional[int]) -> int: - if rows_per_cta is None: - rows_per_cta = pick_rows_per_cta(m) - if rows_per_cta <= 0: - raise ValueError(f"rows_per_cta must be positive, got {rows_per_cta}") - if m % rows_per_cta != 0: - raise ValueError(f"M must be divisible by rows_per_cta, got M={m}, rows_per_cta={rows_per_cta}") - return rows_per_cta - - @staticmethod - def _validate_launch_configuration( n: int, num_threads: int, + eps: float, rows_per_cta: int, - *, - m: Optional[int] = None, ) -> None: - if n <= 0: - raise ValueError(f"N must be positive, got {n}") - if num_threads <= 0: - raise ValueError(f"num_threads must be positive, got {num_threads}") - if num_threads % 32 != 0: - raise ValueError(f"num_threads must be warp-aligned, got {num_threads}") - if num_threads > 1024: - raise ValueError(f"num_threads must not exceed the CUDA block size limit, got {num_threads}") - if n % num_threads != 0: - raise ValueError(f"N={n} must be divisible by num_threads={num_threads}") - - ept = n // num_threads - if ept < 8 or ept % 8 != 0: - raise ValueError(f"EPT={ept} must be >= 8 and divisible by 8") - if rows_per_cta <= 0: - raise ValueError(f"rows_per_cta must be positive, got {rows_per_cta}") - if m is not None and m % rows_per_cta != 0: - raise ValueError(f"M must be divisible by rows_per_cta, got M={m}, rows_per_cta={rows_per_cta}") - - def _configure_lowering_state(self) -> None: - if self.n is None or self.num_threads is None or self.rows_per_cta is None: - raise RuntimeError("Kernel launch configuration has not been resolved") - - self._validate_launch_configuration(self.n, self.num_threads, self.rows_per_cta, m=self.m) + self.n = n + self.num_threads = num_threads + self.eps = eps + self.rows_per_cta = rows_per_cta self.vec_size = self.COPY_BITS // 16 self.ept = self.n // self.num_threads diff --git a/python/cudnn/rmsnorm_rht_amax/op.py b/python/cudnn/rmsnorm_rht_amax/op.py new file mode 100644 index 000000000..2c342cd75 --- /dev/null +++ b/python/cudnn/rmsnorm_rht_amax/op.py @@ -0,0 +1,193 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral RMSNorm + RHT + per-CTA amax operation.""" + +from __future__ import annotations + +from typing import Any, Optional + +from .. import data_type +from .._op import Op +from .._tensor_desc import TensorDesc + +HAD_BLOCK = 16 + +DEFAULT_NUM_THREADS_BY_N = { + 2048: 128, + 4096: 256, + 7168: 128, + 8192: 512, + 16384: 1024, + 32768: 512, +} +RPC_CANDIDATES = (2, 4, 8) +TARGET_MIN_CTAS = 148 + + +def best_num_threads(n: int) -> Optional[int]: + for num_threads in (1024, 512, 256, 128, 64): + if n % num_threads != 0: + continue + ept = n // num_threads + if ept >= 8 and ept % 8 == 0: + return num_threads + return None + + +def pick_rows_per_cta(m: int) -> int: + for rows_per_cta in reversed(RPC_CANDIDATES): + if m % rows_per_cta != 0: + continue + if m // rows_per_cta >= TARGET_MIN_CTAS: + return rows_per_cta + return RPC_CANDIDATES[0] + + +class RmsNormRhtAmaxSm100Op(Op): + """Complete logical signature and launch configuration for the SM100 op.""" + + def __init__( + self, + *, + x: TensorDesc[Any], + weight: TensorDesc[Any], + output: TensorDesc[Any], + amax: TensorDesc[Any], + eps: float = 1e-5, + num_threads: Optional[int] = None, + rows_per_cta: Optional[int] = None, + ) -> None: + for name, desc in (("x", x), ("weight", weight), ("output", output), ("amax", amax)): + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + + self.x = x + self.weight = weight + self.output = output + self.amax = amax + self.eps = eps + self.requested_num_threads = num_threads + self.requested_rows_per_cta = rows_per_cta + + self.m: Optional[int] = None + self.n: Optional[int] = None + self.num_threads: Optional[int] = None + self.rows_per_cta: Optional[int] = None + + def check_support(self) -> bool: + """Validate the complete signature and resolve the launch configuration.""" + + self.m = None + self.n = None + self.num_threads = None + self.rows_per_cta = None + + self._check_input_descriptors() + m, n = self.x.shape + + num_threads = self.requested_num_threads + if num_threads is None: + num_threads = DEFAULT_NUM_THREADS_BY_N.get(n, best_num_threads(n)) + if num_threads is None: + raise ValueError(f"No valid num_threads found for N={n}") + + rows_per_cta = self._resolve_rows_per_cta(m, self.requested_rows_per_cta) + self._validate_launch_configuration(n, num_threads, rows_per_cta, m=m) + self._check_output_descriptors(m, n, rows_per_cta) + + self.m = m + self.n = n + self.num_threads = num_threads + self.rows_per_cta = rows_per_cta + return True + + def _check_input_descriptors(self) -> None: + if self.x.ndim != 2: + raise ValueError(f"X must have rank 2, got shape {self.x.shape}") + if self.weight.ndim != 1: + raise ValueError(f"W must have rank 1, got shape {self.weight.shape}") + if self.x.cudnn_dtype != data_type.BFLOAT16: + raise ValueError(f"X must have dtype bfloat16, got {self.x.dtype}") + if self.weight.cudnn_dtype != data_type.BFLOAT16: + raise ValueError(f"W must have dtype bfloat16, got {self.weight.dtype}") + + m, n = self.x.shape + if m <= 0: + raise ValueError(f"M must be positive, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if self.weight.shape != (n,): + raise ValueError(f"W must have shape {(n,)}, got {self.weight.shape}") + if self.x.stride != (n, 1) or self.x.stride_order != (1, 0): + raise ValueError(f"X must be row-major contiguous, got stride {self.x.stride} and stride order {self.x.stride_order}") + if self.weight.stride != (1,) or self.weight.stride_order != (0,): + raise ValueError(f"W must be contiguous, got stride {self.weight.stride} and stride order {self.weight.stride_order}") + if n % HAD_BLOCK != 0: + raise ValueError(f"N must be divisible by {HAD_BLOCK} for the Hadamard block size, got {n}") + + def _check_output_descriptors(self, m: int, n: int, rows_per_cta: int) -> None: + if self.output.ndim != 2: + raise ValueError(f"O must have rank 2, got shape {self.output.shape}") + if self.amax.ndim != 1: + raise ValueError(f"Amax must have rank 1, got shape {self.amax.shape}") + if self.output.cudnn_dtype != data_type.BFLOAT16: + raise ValueError(f"O must have dtype bfloat16, got {self.output.dtype}") + if self.amax.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"Amax must have dtype float32, got {self.amax.dtype}") + if self.output.shape != (m, n): + raise ValueError(f"O must have shape {(m, n)}, got {self.output.shape}") + expected_amax_shape = (m // rows_per_cta,) + if self.amax.shape != expected_amax_shape: + raise ValueError(f"Amax must have shape {expected_amax_shape}, got {self.amax.shape}") + if self.output.stride != (n, 1) or self.output.stride_order != (1, 0): + raise ValueError(f"O must be row-major contiguous, got stride {self.output.stride} and stride order {self.output.stride_order}") + if self.amax.stride == (0,): + raise ValueError("Amax stride must be positive") + + @staticmethod + def _resolve_rows_per_cta(m: int, rows_per_cta: Optional[int]) -> int: + if rows_per_cta is None: + rows_per_cta = pick_rows_per_cta(m) + if rows_per_cta <= 0: + raise ValueError(f"rows_per_cta must be positive, got {rows_per_cta}") + if m % rows_per_cta != 0: + raise ValueError(f"M must be divisible by rows_per_cta, got M={m}, rows_per_cta={rows_per_cta}") + return rows_per_cta + + @staticmethod + def _validate_launch_configuration( + n: int, + num_threads: int, + rows_per_cta: int, + *, + m: Optional[int] = None, + ) -> None: + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if num_threads <= 0: + raise ValueError(f"num_threads must be positive, got {num_threads}") + if num_threads % 32 != 0: + raise ValueError(f"num_threads must be warp-aligned, got {num_threads}") + if num_threads > 1024: + raise ValueError(f"num_threads must not exceed the CUDA block size limit, got {num_threads}") + if n % num_threads != 0: + raise ValueError(f"N={n} must be divisible by num_threads={num_threads}") + + ept = n // num_threads + if ept < 8 or ept % 8 != 0: + raise ValueError(f"EPT={ept} must be >= 8 and divisible by 8") + if rows_per_cta <= 0: + raise ValueError(f"rows_per_cta must be positive, got {rows_per_cta}") + if m is not None and m % rows_per_cta != 0: + raise ValueError(f"M must be divisible by rows_per_cta, got M={m}, rows_per_cta={rows_per_cta}") + + +__all__ = [ + "DEFAULT_NUM_THREADS_BY_N", + "RPC_CANDIDATES", + "TARGET_MIN_CTAS", + "RmsNormRhtAmaxSm100Op", + "best_num_threads", + "pick_rows_per_cta", +] diff --git a/test/python/fe_api/test_jax_api_base.py b/test/python/fe_api/test_jax_api_base.py index 6966e8b85..63d29bff2 100644 --- a/test/python/fe_api/test_jax_api_base.py +++ b/test/python/fe_api/test_jax_api_base.py @@ -50,6 +50,17 @@ def __init__(self, shape, dtype): self.dtype = dtype +class _Device: + platform = "gpu" + + def __init__(self, device_id, compute_capability): + self.id = device_id + self.compute_capability = compute_capability + + def __str__(self): + return f"CudaDevice(id={self.id})" + + class JaxApiBaseTest(unittest.TestCase): @classmethod def setUpClass(cls) -> None: @@ -101,6 +112,91 @@ def tearDownClass(cls) -> None: if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): sys.modules.pop(name, None) + def setUp(self) -> None: + self.module.disable_device_compatibility_checks(False) + + def tearDown(self) -> None: + self.module.disable_device_compatibility_checks(False) + + def test_checks_all_local_jax_gpu_compute_capabilities(self): + jax = types.ModuleType("jax") + + def local_devices(*, backend): + self.assertEqual(backend, "gpu") + return (_Device(0, "10.0"), _Device(1, "12.0")) + + jax.local_devices = local_devices + with mock.patch.dict(sys.modules, {"jax": jax}): + self.module.JaxApiBase._check_device_compatibility( + minimum_compute_capability=100, + operation_name="TestOp", + ) + + def test_rejects_an_incompatible_local_jax_gpu(self): + jax = types.ModuleType("jax") + jax.local_devices = lambda *, backend: (_Device(0, "10.0"), _Device(1, "9.0")) + + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaisesRegex(RuntimeError, r"TestOp requires SM100\+, found CudaDevice\(id=1\) \(SM90\)"): + self.module.JaxApiBase._check_device_compatibility( + minimum_compute_capability=100, + operation_name="TestOp", + ) + + def test_requires_an_explicit_opt_out_without_a_local_jax_gpu(self): + jax = types.ModuleType("jax") + jax.local_devices = lambda *, backend: () + + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaisesRegex(RuntimeError, r"no local JAX GPU.*disable_device_compatibility_checks\(True\)"): + self.module.JaxApiBase._check_device_compatibility( + minimum_compute_capability=100, + operation_name="TestOp", + ) + + def test_reports_device_discovery_and_capability_failures(self): + jax = types.ModuleType("jax") + + def fail_discovery(*, backend): + raise RuntimeError("GPU backend unavailable") + + jax.local_devices = fail_discovery + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaisesRegex(RuntimeError, r"could not discover a local GPU.*disable_device_compatibility_checks\(True\)"): + self.module.JaxApiBase._check_device_compatibility( + minimum_compute_capability=100, + operation_name="TestOp", + ) + + jax.local_devices = lambda *, backend: (_Device(0, "unknown"),) + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaisesRegex(RuntimeError, r"invalid compute capability 'unknown'.*disable_device_compatibility_checks\(True\)"): + self.module.JaxApiBase._check_device_compatibility( + minimum_compute_capability=100, + operation_name="TestOp", + ) + + def test_device_compatibility_checks_can_be_disabled_and_reenabled(self): + self.module.disable_device_compatibility_checks(True) + with mock.patch.dict(sys.modules, {"jax": None}): + self.module.JaxApiBase._check_device_compatibility( + minimum_compute_capability=100, + operation_name="TestOp", + ) + + self.module.disable_device_compatibility_checks(False) + jax = types.ModuleType("jax") + jax.local_devices = lambda *, backend: () + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaises(RuntimeError): + self.module.JaxApiBase._check_device_compatibility( + minimum_compute_capability=100, + operation_name="TestOp", + ) + + with self.assertRaisesRegex(TypeError, "disabled must be a bool"): + self.module.disable_device_compatibility_checks(1) + def test_converts_array_metadata_to_shared_tensor_desc(self): desc = self.module.JaxApiBase._to_tensor_desc( _ArrayMetadata((2, 3, 4), "bfloat16"), @@ -249,53 +345,50 @@ def test_materializes_canonical_output_in_public_axis_order(self): self.assertEqual(output.shape, (2, 3, 4)) self.assertEqual(output.dtype, "bfloat16") - def test_calls_kernel_with_initialized_outputs_and_hidden_workspaces(self): + def test_calls_launch_hook_with_initialized_outputs_and_hidden_workspaces(self): tensor_module = self.tensor_module seen = {} - class Kernel: - def infer_output(self): - return ( - tensor_module.make_compact_tensor_desc( - dtype=_DataType.FLOAT, - shape=(2,), - name="output", - ), - tensor_module.make_compact_tensor_desc( - dtype=_DataType.FLOAT, - shape=(4, 2, 3), - name="amax", - init_value=float("-inf"), - ), - ) - - def infer_workspace(self): - return ( - tensor_module.make_compact_tensor_desc( - dtype=_DataType.FLOAT, - shape=(3,), - name="counter", - init_value=0, - ), - tensor_module.make_compact_tensor_desc( - dtype=_DataType.BFLOAT16, - shape=(5,), - name="scratch", - ), - ) - - def __call__(self, *args): - seen["kernel_args"] = args + output_descs = ( + tensor_module.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2,), + name="output", + ), + tensor_module.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(4, 2, 3), + name="amax", + init_value=float("-inf"), + ), + ) + workspace_descs = ( + tensor_module.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(3,), + name="counter", + init_value=0, + ), + tensor_module.make_compact_tensor_desc( + dtype=_DataType.BFLOAT16, + shape=(5,), + name="scratch", + ), + ) class Adapter(self.module.JaxApiBase): - def __init__(self): - self.kernel = Kernel() - def check_support(self): return True def __call__(self, value): - return self._call_kernel((value,)) + return self._call_kernel( + (value,), + output_descs=output_descs, + workspace_descs=workspace_descs, + ) + + def _launch(self, inputs, outputs, workspaces, stream): + seen["launch_args"] = (inputs, outputs, workspaces, stream) full_calls = [] jax = types.ModuleType("jax") @@ -357,14 +450,22 @@ def invoke(*args): ): results = api._call_kernel( (input_value,), + output_descs=output_descs, + workspace_descs=workspace_descs, input_spec=(input_tensor_spec,), output_spec=(None, initialized_output_spec), ) with self.assertRaisesRegex(TypeError, "input #0 must have shape and dtype metadata"): - api._call_kernel(([input_value],)) + api._call_kernel( + ([input_value],), + output_descs=output_descs, + workspace_descs=workspace_descs, + ) with self.assertRaisesRegex(TypeError, "output_spec must contain only TensorSpec or None"): api._call_kernel( (input_value,), + output_descs=output_descs, + workspace_descs=workspace_descs, output_spec=((0,), initialized_output_spec), ) @@ -380,13 +481,13 @@ def invoke(*args): self.assertIs(seen["call_options"]["input_spec"][1], initialized_output_spec) self.assertIs(seen["call_options"]["input_spec"][2], seen["call_options"]["output_spec"][2]) - kernel_args = seen["kernel_args"] - self.assertIs(kernel_args[0], input_value) - self.assertEqual(kernel_args[1].label, "allocated(0)") - self.assertIs(kernel_args[2], full_calls[0][3]) - self.assertIs(kernel_args[3], full_calls[1][3]) - self.assertEqual(kernel_args[4].label, "allocated(3)") - self.assertEqual(kernel_args[5], "stream") + launch_inputs, launch_outputs, launch_workspaces, stream = seen["launch_args"] + self.assertEqual(launch_inputs, (input_value,)) + self.assertEqual(launch_outputs[0].label, "allocated(0)") + self.assertIs(launch_outputs[1], full_calls[0][3]) + self.assertIs(launch_workspaces[0], full_calls[1][3]) + self.assertEqual(launch_workspaces[1].label, "allocated(3)") + self.assertEqual(stream, "stream") self.assertEqual([result.label for result in results], ["allocated(0)", "full(-inf)"]) self.assertNotIn(full_calls[0][3], vars(api).values()) self.assertNotIn(full_calls[1][3], vars(api).values()) @@ -401,6 +502,9 @@ def check_support(self): def __call__(self): return self.check_support() + def _launch(self, inputs, outputs, workspaces, stream): + raise AssertionError("not called") + api = Adapter() api.option = 1 self.assertIs(api.get_jax_callable(), api) diff --git a/test/python/fe_api/test_jax_rmsnorm_rht_amax.py b/test/python/fe_api/test_jax_rmsnorm_rht_amax.py index 6866b405e..b5d23f7af 100644 --- a/test/python/fe_api/test_jax_rmsnorm_rht_amax.py +++ b/test/python/fe_api/test_jax_rmsnorm_rht_amax.py @@ -22,25 +22,41 @@ def _hadamard_16(jnp): return jnp.asarray(values, dtype=jnp.float32) / math.sqrt(16) +@pytest.fixture +def device_compatibility_checks_disabled(): + _jax_runtime() + from cudnn.jax import disable_device_compatibility_checks + + disable_device_compatibility_checks(True) + try: + yield + finally: + disable_device_compatibility_checks(False) + + @pytest.mark.L0 -def test_jax_rmsnorm_rht_amax_abstract_contract(): +def test_jax_rmsnorm_rht_amax_abstract_contract(device_compatibility_checks_disabled): jax, jnp = _jax_runtime() - from cudnn import OpKernel, TensorDesc + from cudnn import Op, TensorDesc from cudnn.jax import JaxApiBase, RmsNormRhtAmaxSm100, rmsnorm_rht_amax_sm100 - from cudnn.rmsnorm_rht_amax.kernel import RMSNormRHTAmaxKernel + from cudnn.rmsnorm_rht_amax.op import RmsNormRhtAmaxSm100Op sample_x = jax.ShapeDtypeStruct((256, 2048), jnp.bfloat16) sample_weight = jax.ShapeDtypeStruct((2048,), jnp.bfloat16) api = RmsNormRhtAmaxSm100(sample_x, sample_weight) assert isinstance(api, JaxApiBase) - assert isinstance(api.kernel, OpKernel) - assert isinstance(api.kernel, RMSNormRHTAmaxKernel) + assert isinstance(api._op, Op) + assert isinstance(api._op, RmsNormRhtAmaxSm100Op) assert isinstance(api.x_desc, TensorDesc) assert isinstance(api.w_desc, TensorDesc) - assert api.kernel.x is api.x_desc - assert api.kernel.weight is api.w_desc + assert isinstance(api.o_desc, TensorDesc) + assert isinstance(api.amax_desc, TensorDesc) + assert api._op.x is api.x_desc + assert api._op.weight is api.w_desc + assert api._op.output is api.o_desc + assert api._op.amax is api.amax_desc assert api.x_desc.dtype == jnp.dtype(jnp.bfloat16) assert api.x_desc.stride == (2048, 1) assert api.x_desc.stride_order == (1, 0) @@ -52,6 +68,30 @@ def test_jax_rmsnorm_rht_amax_abstract_contract(): assert amax.shape == (128,) assert amax.dtype == jnp.float32 + explicit_api = RmsNormRhtAmaxSm100( + sample_x, + sample_weight, + sample_o=jax.ShapeDtypeStruct((256, 2048), jnp.bfloat16), + sample_amax=jax.ShapeDtypeStruct((128,), jnp.float32), + ) + explicit_output, explicit_amax = jax.eval_shape(explicit_api, sample_x, sample_weight) + assert explicit_output.shape == (256, 2048) + assert explicit_amax.shape == (128,) + + array_api = RmsNormRhtAmaxSm100( + jnp.empty((2, 2048), jnp.bfloat16), + jnp.empty((2048,), jnp.bfloat16), + ) + assert array_api.x_desc.shape == (2, 2048) + assert array_api.w_desc.shape == (2048,) + + with pytest.raises(ValueError, match="sample_o and sample_amax must be provided together"): + RmsNormRhtAmaxSm100( + sample_x, + sample_weight, + sample_o=jax.ShapeDtypeStruct((256, 2048), jnp.bfloat16), + ) + functional_output, functional_amax = jax.eval_shape( lambda x, weight: rmsnorm_rht_amax_sm100( x, @@ -72,29 +112,33 @@ def test_jax_rmsnorm_rht_amax_abstract_contract(): with pytest.raises(ValueError, match="sample_x tensor shape mismatch"): jax.eval_shape(api, wrong_x, sample_weight) - api.kernel.requested_rows_per_cta = 3 + invalid_api = RmsNormRhtAmaxSm100(sample_x, sample_weight, rows_per_cta=3) with pytest.raises(ValueError, match="M must be divisible by rows_per_cta"): - jax.eval_shape(lambda x, weight: api(x, weight), sample_x, sample_weight) + jax.eval_shape(invalid_api, sample_x, sample_weight) @pytest.mark.L0 def test_jax_rmsnorm_rht_amax_jit(): jax, jnp = _jax_runtime() gpu_devices = [device for device in jax.local_devices() if device.platform == "gpu"] - capable_devices = [] + all_local_devices_are_supported = bool(gpu_devices) for device in gpu_devices: capability = getattr(device, "compute_capability", None) if capability is None: - continue + all_local_devices_are_supported = False + break major, minor = (int(value) for value in str(capability).split(".", 1)) - if major * 10 + minor >= 100: - capable_devices.append(device) - if not capable_devices: - pytest.skip("RMSNorm + RHT requires an SM100+ JAX device") + if (major, minor) < (10, 0): + all_local_devices_are_supported = False + break + if not all_local_devices_are_supported: + pytest.skip("RMSNorm + RHT requires every local JAX GPU to be SM100+") + + from cudnn.jax import disable_device_compatibility_checks, rmsnorm_rht_amax_sm100 - from cudnn.jax import rmsnorm_rht_amax_sm100 + disable_device_compatibility_checks(False) - device = capable_devices[0] + device = gpu_devices[0] m, n = 256, 2048 rows_per_cta = 2 eps = 1e-4 diff --git a/test/python/fe_api/test_rmsnorm_rht_amax.py b/test/python/fe_api/test_rmsnorm_rht_amax.py index abec7c9a3..e4bb15864 100644 --- a/test/python/fe_api/test_rmsnorm_rht_amax.py +++ b/test/python/fe_api/test_rmsnorm_rht_amax.py @@ -59,7 +59,7 @@ def _assert_ref_close(x, w, o, amax, *, eps: float, rows_per_cta: int, skip_ref: @pytest.mark.L0 @torch_fork_set_rng(seed=0) @pytest.mark.parametrize("n,num_threads", SUPPORTED_N_NUM_THREADS) -def test_rmsnorm_rht_amax_class_call(n, num_threads, request): +def test_rmsnorm_rht_amax_compile_execute(n, num_threads, request): try: from cudnn import RmsNormRhtAmaxSm100 except ImportError: @@ -70,10 +70,14 @@ def test_rmsnorm_rht_amax_class_call(n, num_threads, request): m = 256 rows_per_cta = 2 x, w = _make_inputs(m=m, n=n) + o = torch.empty_like(x) + amax = torch.full((m // rows_per_cta,), float("-inf"), dtype=torch.float32, device="cuda") api = RmsNormRhtAmaxSm100( sample_x=x, sample_w=w, + sample_o=o, + sample_amax=amax, eps=eps, num_threads=num_threads, rows_per_cta=rows_per_cta, @@ -87,15 +91,14 @@ def test_rmsnorm_rht_amax_class_call(n, num_threads, request): except (ValueError, RuntimeError) as exc: pytest.skip(f"Unsupported testcase: {exc}") - outputs = api(x, w) - o, amax = outputs - assert outputs["o_tensor"] is o - assert outputs["amax_tensor"] is amax + api.compile() + result = api.execute(x_tensor=x, w_tensor=w, o_tensor=o, amax_tensor=amax) + assert result is None _assert_ref_close(x, w, o, amax, eps=eps, rows_per_cta=rows_per_cta, skip_ref=skip_ref) @pytest.mark.L0 -def test_rmsnorm_rht_amax_class_rejects_runtime_signature_mismatch(monkeypatch): +def test_rmsnorm_rht_amax_class_rejects_invalid_output_exemplars(monkeypatch): try: from cudnn import RmsNormRhtAmaxSm100 except ImportError: @@ -103,14 +106,60 @@ def test_rmsnorm_rht_amax_class_rejects_runtime_signature_mismatch(monkeypatch): sample_x = torch.empty((4, 256), dtype=torch.bfloat16) sample_w = torch.empty((256,), dtype=torch.bfloat16) - api = RmsNormRhtAmaxSm100(sample_x, sample_w, num_threads=32, rows_per_cta=2) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (10, 0)) + + cases = ( + (torch.empty((2, 256), dtype=torch.bfloat16), torch.empty((2,), dtype=torch.float32), "O must have shape"), + (torch.empty((4, 256), dtype=torch.float32), torch.empty((2,), dtype=torch.float32), "O must have dtype"), + (torch.empty_strided((4, 256), (1, 4), dtype=torch.bfloat16), torch.empty((2,), dtype=torch.float32), "O must be row-major contiguous"), + (torch.empty((4, 256), dtype=torch.bfloat16), torch.empty((1,), dtype=torch.float32), "Amax must have shape"), + (torch.empty((4, 256), dtype=torch.bfloat16), torch.empty((2,), dtype=torch.bfloat16), "Amax must have dtype"), + ) + for sample_o, sample_amax, message in cases: + with pytest.raises(ValueError, match=message): + RmsNormRhtAmaxSm100( + sample_x, + sample_w, + sample_o, + sample_amax, + num_threads=32, + rows_per_cta=2, + ).check_support() + + strided_amax = torch.empty_strided((2,), (2,), dtype=torch.float32) + assert RmsNormRhtAmaxSm100( + sample_x, + sample_w, + torch.empty_like(sample_x), + strided_amax, + num_threads=32, + rows_per_cta=2, + ).check_support() + + +@pytest.mark.L0 +def test_rmsnorm_rht_amax_execute_checks_dynamic_amax_shape(): + try: + from cudnn import RmsNormRhtAmaxSm100 + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + x = torch.empty((4, 256), dtype=torch.bfloat16) + w = torch.empty((256,), dtype=torch.bfloat16) + o = torch.empty_like(x) + amax = torch.empty((2,), dtype=torch.float32) + api = RmsNormRhtAmaxSm100(x, w, o, amax, num_threads=32, rows_per_cta=2) + api._op.check_support() api._compiled_kernel = lambda **_kwargs: None - monkeypatch.setattr(api._kernel, "infer_output", lambda: pytest.fail("output inference must follow runtime signature validation")) - with pytest.raises(ValueError, match="sample_x tensor shape mismatch"): - api.execute(torch.empty((2, 256), dtype=torch.bfloat16), sample_w) - with pytest.raises(ValueError, match="sample_w dtype mismatch"): - api.execute(sample_x, torch.empty((256,), dtype=torch.float32)) + with pytest.raises(ValueError, match="Amax tensor shape mismatch"): + api.execute(x, w, o, torch.empty((1,), dtype=torch.float32)) + + dynamic_x = torch.empty((8, 256), dtype=torch.bfloat16) + dynamic_o = torch.empty_like(dynamic_x) + dynamic_amax = torch.empty((4,), dtype=torch.float32) + assert api.execute(dynamic_x, w, dynamic_o, dynamic_amax, current_stream=object()) is None @pytest.mark.L0 @@ -146,7 +195,7 @@ def test_rmsnorm_rht_amax_wrapper(n, num_threads, rows_per_cta, request): @pytest.mark.L0 @torch_fork_set_rng(seed=0) -def test_rmsnorm_rht_amax_wrapper_reuses_owned_kernel(monkeypatch): +def test_rmsnorm_rht_amax_wrapper_reuses_compiled_kernel(monkeypatch): try: from cudnn.rmsnorm_rht_amax import api as rmsnorm_api except ImportError: @@ -163,32 +212,36 @@ def counted_init(self, *args, **kwargs): monkeypatch.setattr(rmsnorm_api.RMSNormRHTAmaxKernel, "__init__", counted_init) rmsnorm_api._cache_of_RmsNormRhtAmaxSm100Objects.clear() - m, n = 256, 2048 + first_m, second_m, n = 256, 512, 2048 rows_per_cta = 2 - x, w = _make_inputs(m=m, n=n) + first_x, w = _make_inputs(m=first_m, n=n) + second_x, _ = _make_inputs(m=second_m, n=n) try: - first = rmsnorm_api.rmsnorm_rht_amax_wrapper_sm100(x, w, num_threads=128, rows_per_cta=rows_per_cta) - second = rmsnorm_api.rmsnorm_rht_amax_wrapper_sm100(x, w, num_threads=128, rows_per_cta=rows_per_cta) - except (ValueError, RuntimeError) as exc: - pytest.skip(f"Unsupported testcase: {exc}") + try: + first = rmsnorm_api.rmsnorm_rht_amax_wrapper_sm100(first_x, w, num_threads=128, rows_per_cta=rows_per_cta) + except (ValueError, RuntimeError) as exc: + pytest.skip(f"Unsupported testcase: {exc}") + + second = rmsnorm_api.rmsnorm_rht_amax_wrapper_sm100(second_x, w, num_threads=128, rows_per_cta=rows_per_cta) + _assert_ref_close(second_x, w, second["o_tensor"], second["amax_tensor"], eps=1e-5, rows_per_cta=rows_per_cta) + + assert construction_count == 1 + assert first["o_tensor"].shape == (first_m, n) + assert second["o_tensor"].shape == (second_m, n) + assert first["amax_tensor"].shape == (first_m // rows_per_cta,) + assert second["amax_tensor"].shape == (second_m // rows_per_cta,) finally: rmsnorm_api._cache_of_RmsNormRhtAmaxSm100Objects.clear() - assert construction_count == 1 - assert first["o_tensor"].shape == second["o_tensor"].shape == (m, n) - assert first["amax_tensor"].shape == second["amax_tensor"].shape == (m // rows_per_cta,) - @pytest.mark.L0 -def test_rmsnorm_rht_amax_wrapper_delegates_output_allocation_to_class(monkeypatch): +def test_rmsnorm_rht_amax_wrapper_passes_allocated_outputs_to_class(monkeypatch): try: from cudnn.rmsnorm_rht_amax import api as rmsnorm_api except ImportError: pytest.skip("Environment not supported: cudnn optional dependencies not installed") constructed = {} - output = torch.empty((3, 5), dtype=torch.bfloat16) - amax = torch.full((7,), -2.0, dtype=torch.float32) class FakeApi: def __init__(self, **kwargs): @@ -202,20 +255,25 @@ def compile(self): def execute(self, **kwargs): constructed["execute"] = kwargs - return rmsnorm_api.TupleDict(o_tensor=output, amax_tensor=amax) monkeypatch.setattr(rmsnorm_api, "RmsNormRhtAmaxSm100", FakeApi) rmsnorm_api._cache_of_RmsNormRhtAmaxSm100Objects.clear() - x = torch.empty((4, 16), dtype=torch.bfloat16) - w = torch.empty((16,), dtype=torch.bfloat16) + x = torch.empty((4, 256), dtype=torch.bfloat16) + w = torch.empty((256,), dtype=torch.bfloat16) try: - result = rmsnorm_api.rmsnorm_rht_amax_wrapper_sm100(x, w, rows_per_cta=4) + result = rmsnorm_api.rmsnorm_rht_amax_wrapper_sm100(x, w, num_threads=32, rows_per_cta=4) finally: rmsnorm_api._cache_of_RmsNormRhtAmaxSm100Objects.clear() - assert set(constructed) == {"sample_x", "sample_w", "eps", "num_threads", "rows_per_cta", "execute"} - assert result["o_tensor"] is output - assert result["amax_tensor"] is amax - assert torch.all(result["amax_tensor"] == -2.0) - assert constructed["execute"] == {"x_tensor": x, "w_tensor": w, "current_stream": None} + assert set(constructed) == {"sample_x", "sample_w", "sample_o", "sample_amax", "eps", "num_threads", "rows_per_cta", "execute"} + assert constructed["sample_o"] is result["o_tensor"] + assert constructed["sample_amax"] is result["amax_tensor"] + assert torch.isneginf(result["amax_tensor"]).all() + assert constructed["execute"] == { + "x_tensor": x, + "w_tensor": w, + "o_tensor": result["o_tensor"], + "amax_tensor": result["amax_tensor"], + "current_stream": None, + } diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_imports.py b/test/python/fe_api/test_rmsnorm_rht_amax_imports.py index 769f1e794..ff93d42f2 100644 --- a/test/python/fe_api/test_rmsnorm_rht_amax_imports.py +++ b/test/python/fe_api/test_rmsnorm_rht_amax_imports.py @@ -118,9 +118,11 @@ def test_package_import_loads_no_framework_adapter(self): self.assertTrue( { "RMSNormRHTAmaxKernel", + "RmsNormRhtAmaxSm100Op", "RmsNormRhtAmaxSm100", "api", "kernel", + "op", }.issubset(dir(package)) ) finally: @@ -153,6 +155,27 @@ def test_kernel_exports_do_not_load_an_adapter(self): finally: _remove_package(name) + def test_operation_exports_do_not_load_a_framework_adapter(self): + name = "cudnn_rmsnorm_op_route_test.rmsnorm_rht_amax" + op_sentinel = object() + best_num_threads_sentinel = object() + pick_rows_per_cta_sentinel = object() + op = types.ModuleType(f"{name}.op") + op.RmsNormRhtAmaxSm100Op = op_sentinel + op.best_num_threads = best_num_threads_sentinel + op.pick_rows_per_cta = pick_rows_per_cta_sentinel + try: + sys.modules[op.__name__] = op + package = _load_operation_package(name) + self.assertIs(package.RmsNormRhtAmaxSm100Op, op_sentinel) + self.assertIs(package.best_num_threads, best_num_threads_sentinel) + self.assertIs(package.pick_rows_per_cta, pick_rows_per_cta_sentinel) + self.assertNotIn(f"{name}.api", sys.modules) + self.assertNotIn(f"{name}.jax", sys.modules) + self.assertNotIn(f"{name}.kernel", sys.modules) + finally: + _remove_package(name) + def test_static_import_directions(self): def all_imports(filename): path = Path(filename) @@ -166,8 +189,11 @@ def imports_framework(node, framework: str): return any(alias.name == framework or alias.name.startswith(f"{framework}.") for alias in node.names) return node.level == 0 and (node.module == framework or (node.module or "").startswith(f"{framework}.")) - op_kernel_imports = all_imports(_CUDNN_ROOT / "_op_kernel.py") - self.assertFalse(any(imports_framework(node, framework) for node in op_kernel_imports for framework in ("torch", "jax", "cutlass", "cuda"))) + base_op_imports = all_imports(_CUDNN_ROOT / "_op.py") + self.assertFalse(any(imports_framework(node, framework) for node in base_op_imports for framework in ("torch", "jax", "cutlass", "cuda"))) + + operation_imports = all_imports("op.py") + self.assertFalse(any(imports_framework(node, framework) for node in operation_imports for framework in ("torch", "jax", "cutlass", "cuda"))) kernel_imports = all_imports("kernel.py") self.assertFalse(any(imports_framework(node, framework) for node in kernel_imports for framework in ("torch", "jax"))) @@ -183,6 +209,7 @@ def imports_framework(node, framework: str): jax_package_tree = ast.parse((_CUDNN_ROOT / "_jax" / "__init__.py").read_text(), filename="_jax/__init__.py") jax_facade_tree = ast.parse((_CUDNN_ROOT / "jax" / "__init__.py").read_text(), filename="jax/__init__.py") jax_base_tree = ast.parse((_CUDNN_ROOT / "_jax" / "api_base.py").read_text(), filename="_jax/api_base.py") + op_tree = ast.parse((_OPERATION_ROOT / "op.py").read_text(), filename="op.py") kernel_tree = ast.parse((_OPERATION_ROOT / "kernel.py").read_text(), filename="kernel.py") torch_adapter = next(node for node in api_tree.body if isinstance(node, ast.ClassDef) and node.name == "RmsNormRhtAmaxSm100") self.assertEqual([base.id for base in torch_adapter.bases if isinstance(base, ast.Name)], ["APIBase"]) @@ -190,17 +217,18 @@ def imports_framework(node, framework: str): required_argument_count = len(constructor.args.args) - len(constructor.args.defaults) self.assertEqual( [argument.arg for argument in constructor.args.args[:required_argument_count]], - ["self", "sample_x", "sample_w"], + ["self", "sample_x", "sample_w", "sample_o", "sample_amax"], ) self.assertNotIn( "_kernel", [argument.arg for argument in (*constructor.args.args, *constructor.args.kwonlyargs)], ) - kernel_constructors = [ - node for node in ast.walk(api_tree) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RMSNormRHTAmaxKernel" + op_constructors = [ + node for node in ast.walk(api_tree) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RmsNormRhtAmaxSm100Op" ] - self.assertEqual(len(kernel_constructors), 1) - self.assertIn(kernel_constructors[0], ast.walk(constructor)) + self.assertEqual(len(op_constructors), 1) + self.assertIn(op_constructors[0], ast.walk(constructor)) + self.assertTrue({"x", "weight", "output", "amax"}.issubset({keyword.arg for keyword in op_constructors[0].keywords})) self.assertFalse( any( isinstance(node, ast.FunctionDef) and any(isinstance(decorator, ast.Name) and decorator.id == "property" for decorator in node.decorator_list) @@ -208,6 +236,8 @@ def imports_framework(node, framework: str): ) ) torch_check_support = next(node for node in torch_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "check_support") + self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "check_support" for node in ast.walk(torch_check_support))) + self.assertFalse(any(isinstance(node, ast.Attribute) and node.attr == "check_output" for node in ast.walk(torch_check_support))) self.assertFalse( any( isinstance(node, ast.Attribute) @@ -215,7 +245,7 @@ def imports_framework(node, framework: str): and isinstance(node.value, ast.Attribute) and isinstance(node.value.value, ast.Name) and node.value.value.id == "self" - and node.value.attr == "_kernel" + and node.value.attr == "_op" for node in ast.walk(torch_check_support) ) ) @@ -224,7 +254,13 @@ def imports_framework(node, framework: str): self.assertEqual([base.id for base in jax_adapter.bases if isinstance(base, ast.Name)], ["JaxApiBase"]) self.assertFalse(jax_adapter.decorator_list) self.assertTrue( - any(isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RMSNormRHTAmaxKernel" for node in ast.walk(jax_adapter)) + any(isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RmsNormRhtAmaxSm100Op" for node in ast.walk(jax_adapter)) + ) + jax_check_support = next(node for node in jax_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "check_support") + self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "_check_device_compatibility" for node in ast.walk(jax_check_support))) + jax_launch = next(node for node in jax_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "_launch") + self.assertTrue( + any(isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RMSNormRHTAmaxKernel" for node in ast.walk(jax_launch)) ) jax_wrapper = next(node for node in jax_tree.body if isinstance(node, ast.FunctionDef) and node.name == "rmsnorm_rht_amax_sm100") @@ -251,22 +287,20 @@ def imports_framework(node, framework: str): for node in ast.walk(torch_wrapper) ) ) - self.assertFalse(any(isinstance(node, ast.Name) and node.id in {"best_num_threads", "pick_rows_per_cta"} for node in ast.walk(torch_wrapper))) torch_api_construction = next( node for node in ast.walk(torch_wrapper) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RmsNormRhtAmaxSm100" ) - self.assertFalse({"sample_o", "sample_amax"} & {keyword.arg for keyword in torch_api_construction.keywords}) + self.assertTrue({"sample_o", "sample_amax"}.issubset({keyword.arg for keyword in torch_api_construction.keywords})) torch_execute = next(node for node in torch_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "execute") self.assertEqual( [argument.arg for argument in torch_execute.args.args], - ["self", "x_tensor", "w_tensor", "current_stream"], + ["self", "x_tensor", "w_tensor", "o_tensor", "amax_tensor", "current_stream"], ) - self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "infer_output" for node in ast.walk(torch_execute))) - self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "materialize" for node in ast.walk(torch_execute))) + self.assertFalse(any(isinstance(node, ast.Attribute) and node.attr in {"infer_output", "materialize"} for node in ast.walk(torch_execute))) kernel_class = next(node for node in kernel_tree.body if isinstance(node, ast.ClassDef) and node.name == "RMSNormRHTAmaxKernel") - self.assertEqual([base.id for base in kernel_class.bases if isinstance(base, ast.Name)], ["OpKernel"]) + self.assertFalse(kernel_class.bases) launcher = next(node for node in kernel_class.body if isinstance(node, ast.FunctionDef) and node.name == "__call__") self.assertEqual( [argument.arg for argument in launcher.args.args], @@ -274,6 +308,11 @@ def imports_framework(node, framework: str): ) compile_method = next(node for node in torch_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "compile") + kernel_constructors = [ + node for node in ast.walk(api_tree) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RMSNormRHTAmaxKernel" + ] + self.assertEqual(len(kernel_constructors), 1) + self.assertIn(kernel_constructors[0], ast.walk(compile_method)) self.assertTrue( {"n", "rows_per_cta"}.issubset( { @@ -283,7 +322,7 @@ def imports_framework(node, framework: str): and isinstance(node.value, ast.Attribute) and isinstance(node.value.value, ast.Name) and node.value.value.id == "self" - and node.value.attr == "_kernel" + and node.value.attr == "_op" } ) ) @@ -296,11 +335,17 @@ def imports_framework(node, framework: str): ["x_tensor", "w_tensor", "o_tensor", "amax_tensor", "stream"], ) + operation_class = next(node for node in op_tree.body if isinstance(node, ast.ClassDef) and node.name == "RmsNormRhtAmaxSm100Op") + self.assertEqual([base.id for base in operation_class.bases if isinstance(base, ast.Name)], ["Op"]) self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "_call_kernel" for node in ast.walk(jax_adapter))) self.assertFalse( any(isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "cutlass_call" for node in ast.walk(jax_adapter)) ) self.assertIn("_INSTALL_HINT", {node.id for node in ast.walk(jax_facade_tree) if isinstance(node, ast.Name)}) + self.assertIn( + "disable_device_compatibility_checks", + {alias.name for node in ast.walk(jax_facade_tree) if isinstance(node, ast.ImportFrom) for alias in node.names}, + ) self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "is_available" for node in ast.walk(jax_facade_tree))) self.assertTrue(any(isinstance(node, ast.Try) for node in jax_facade_tree.body)) self.assertTrue(any(isinstance(node, ast.Import) and any(alias.name == "jax" for alias in node.names) for node in jax_tree.body)) @@ -319,7 +364,7 @@ def imports_framework(node, framework: str): self.assertNotIn("_trace_lock", {node.attr for node in ast.walk(jax_adapter) if isinstance(node, ast.Attribute)}) root_source = (_CUDNN_ROOT / "__init__.py").read_text() - self.assertIn("from ._op_kernel import OpKernel", root_source) + self.assertIn("from ._op import Op", root_source) self.assertIn("from ._tensor_desc import TensorDesc", root_source) self.assertIn('if name == "jax":', root_source) self.assertIn('importlib.import_module(".jax", __name__)', root_source) diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py b/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py index 4a0dc56a6..3eca62eb1 100644 --- a/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py +++ b/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py @@ -1,9 +1,9 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for the RMSNorm operation kernel.""" +"""Dependency-free contracts for the RMSNorm CuTe kernel.""" -from enum import Enum, auto +import ast import importlib from importlib.machinery import ModuleSpec from pathlib import Path @@ -25,12 +25,6 @@ _PACKAGE = "cudnn_rmsnorm_kernel_test" -class _DataType(Enum): - NOT_SET = auto() - FLOAT = auto() - BFLOAT16 = auto() - - def _identity_decorator(function=None, **_kwargs): if function is None: return lambda decorated: decorated @@ -99,215 +93,86 @@ def _cutlass_stubs() -> dict[str, types.ModuleType]: } -def _install_test_package() -> None: - root = types.ModuleType(_PACKAGE) - root.__path__ = [str(_CUDNN_ROOT)] - root.__package__ = _PACKAGE - root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) - root.data_type = _DataType - sys.modules[_PACKAGE] = root - - operation_name = f"{_PACKAGE}.rmsnorm_rht_amax" - operation = types.ModuleType(operation_name) - operation.__path__ = [str(_OPERATION_ROOT)] - operation.__package__ = operation_name - operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) - sys.modules[operation_name] = operation - - -def _remove_test_package() -> None: - for name in tuple(sys.modules): - if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): - sys.modules.pop(name, None) - - class RmsNormRhtAmaxKernelContractTest(unittest.TestCase): @classmethod def setUpClass(cls) -> None: - _install_test_package() + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + sys.modules[_PACKAGE] = root + + operation_name = f"{_PACKAGE}.rmsnorm_rht_amax" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + try: with mock.patch.dict(sys.modules, _cutlass_stubs()): - cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") - cls.base_module = importlib.import_module(f"{_PACKAGE}._op_kernel") - cls.kernel_module = importlib.import_module(f"{_PACKAGE}.rmsnorm_rht_amax.kernel") + cls.kernel_module = importlib.import_module(f"{operation_name}.kernel") except Exception: - _remove_test_package() + cls.tearDownClass() raise @classmethod def tearDownClass(cls) -> None: - _remove_test_package() - - def _desc(self, shape, dtype=_DataType.BFLOAT16, *, name=""): - shape = tuple(shape) - stride_order = tuple(reversed(range(len(shape)))) - stride = [0] * len(shape) - running = 1 - for dimension in stride_order: - stride[dimension] = running - running *= shape[dimension] - return self.tensor_module.TensorDesc( - dtype=dtype, - shape=shape, - stride=tuple(stride), - stride_order=stride_order, - name=name, - ) - - def _kernel(self, **overrides): - arguments = { - "x": self._desc((256, 2048)), - "weight": self._desc((2048,)), - } - arguments.update(overrides) - return self.kernel_module.RMSNormRHTAmaxKernel(**arguments) - - def test_kernel_implements_neutral_contract_and_infers_buffers(self): - kernel = self._kernel() - self.assertIsInstance(kernel, self.base_module.OpKernel) - with self.assertRaisesRegex(RuntimeError, r"check_support\(\)"): - kernel.infer_output() - - self.assertTrue(kernel.check_support()) - self.assertEqual((kernel.m, kernel.n), (256, 2048)) - self.assertEqual((kernel.num_threads, kernel.rows_per_cta), (128, 2)) - - output, amax = kernel.infer_output() - self.assertIs(type(output), self.tensor_module.TensorDesc) - self.assertEqual( - (output.name, output.dtype, output.shape, output.stride, output.stride_order, output.init_value), - ("output", _DataType.BFLOAT16, (256, 2048), (2048, 1), (1, 0), None), - ) - self.assertEqual( - (amax.name, amax.dtype, amax.shape, amax.stride, amax.stride_order, amax.init_value), - ("amax", _DataType.FLOAT, (128,), (1,), (0,), None), - ) - - kernel.x = self._desc((128, 2048)) - output, amax = kernel.infer_output() - self.assertEqual((output.shape, amax.shape), ((256, 2048), (128,))) - self.assertEqual(kernel.infer_workspace(), ()) - - self.assertNotIn(f"{_PACKAGE}.rmsnorm_rht_amax.api", sys.modules) - self.assertNotIn(f"{_PACKAGE}.rmsnorm_rht_amax.jax", sys.modules) - - def test_output_inference_preserves_descriptor_specialization(self): - tensor_module = self.tensor_module - - class SpecializedTensorDesc(tensor_module.TensorDesc): - def compact_like(self, **kwargs): - canonical = super().compact_like(**kwargs) - return SpecializedTensorDesc( - dtype=canonical.dtype, - shape=canonical.shape, - stride=canonical.stride, - stride_order=canonical.stride_order, - name=canonical.name, - init_value=canonical.init_value, - ) - - x = SpecializedTensorDesc( - dtype=_DataType.BFLOAT16, - shape=(256, 2048), - stride=(2048, 1), - stride_order=(1, 0), - ) + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + def test_kernel_consumes_only_resolved_configuration(self): kernel = self.kernel_module.RMSNormRHTAmaxKernel( - x=x, - weight=SpecializedTensorDesc( - dtype=_DataType.BFLOAT16, - shape=(2048,), - stride=(1,), - stride_order=(0,), - ), + n=2048, + num_threads=128, + eps=1e-4, + rows_per_cta=2, ) - kernel.check_support() - output, amax = kernel.infer_output() - - self.assertIsInstance(output, SpecializedTensorDesc) - self.assertIsInstance(amax, SpecializedTensorDesc) - - def test_requested_configuration_is_recomputed(self): - kernel = self._kernel(num_threads=256, rows_per_cta=4) - kernel.check_support() - self.assertEqual((kernel.num_threads, kernel.rows_per_cta), (256, 4)) - kernel.requested_num_threads = 128 - kernel.requested_rows_per_cta = 8 - kernel.check_support() - self.assertEqual((kernel.num_threads, kernel.rows_per_cta), (128, 8)) - self.assertEqual(kernel.infer_output()[1].shape, (32,)) - - def test_framework_descriptor_extension_maps_native_dtype(self): - class FrameworkTensorDesc(self.tensor_module.TensorDesc): - @property - def cudnn_dtype(self): - return { - "framework.bfloat16": _DataType.BFLOAT16, - "framework.float32": _DataType.FLOAT, - }.get(self.dtype, _DataType.NOT_SET) - - def desc(shape, dtype): - base = self._desc(shape) - return FrameworkTensorDesc( - dtype=dtype, - shape=base.shape, - stride=base.stride, - stride_order=base.stride_order, - name="framework_tensor", - ) - - kernel = self.kernel_module.RMSNormRHTAmaxKernel( - x=desc((256, 2048), "framework.bfloat16"), - weight=desc((2048,), "framework.bfloat16"), + self.assertEqual(kernel.n, 2048) + self.assertEqual(kernel.num_threads, 128) + self.assertEqual(kernel.eps, 1e-4) + self.assertEqual(kernel.rows_per_cta, 2) + self.assertEqual(kernel.ept, 16) + self.assertEqual(kernel.vec_size, 8) + self.assertEqual(kernel.num_vec_blocks, 2) + self.assertEqual(kernel.warps_per_row, 4) + self.assertEqual(kernel.tiler_mn, (1, 2048)) + + def test_kernel_module_has_no_operation_or_framework_dependency(self): + tree = ast.parse((_OPERATION_ROOT / "kernel.py").read_text(), filename="kernel.py") + imports = [node for node in ast.walk(tree) if isinstance(node, (ast.Import, ast.ImportFrom))] + + imported_modules = set() + for node in imports: + if isinstance(node, ast.Import): + imported_modules.update(alias.name for alias in node.names) + else: + if node.module: + imported_modules.add(node.module) + imported_modules.update(alias.name for alias in node.names) + + self.assertNotIn("torch", imported_modules) + self.assertNotIn("jax", imported_modules) + self.assertNotIn("_op", imported_modules) + self.assertNotIn("_tensor_desc", imported_modules) + self.assertNotIn("data_type", imported_modules) + + kernel_class = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "RMSNormRHTAmaxKernel") + self.assertFalse(kernel_class.bases) + self.assertFalse(any(isinstance(node, ast.FunctionDef) and node.name == "check_support" for node in kernel_class.body)) + self.assertFalse(any(isinstance(node, ast.FunctionDef) and node.name.startswith("infer_") for node in kernel_class.body)) + + def test_hadamard_block_matches_operation_contract(self): + tree = ast.parse((_OPERATION_ROOT / "op.py").read_text(), filename="op.py") + assignment = next( + node + for node in tree.body + if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == "HAD_BLOCK" for target in node.targets) ) - self.assertTrue(kernel.check_support()) - - def test_invalid_logical_signatures_and_configuration(self): - cases = ( - ({"x": self._desc((2048,))}, "X must have rank 2"), - ({"weight": self._desc((1, 2048))}, "W must have rank 1"), - ({"x": self._desc((256, 2048), _DataType.FLOAT)}, "X must have dtype bfloat16"), - ({"weight": self._desc((2048,), _DataType.FLOAT)}, "W must have dtype bfloat16"), - ({"x": self._desc((0, 2048))}, "M must be positive"), - ( - {"x": self._desc((256, 2049)), "weight": self._desc((2049,))}, - "N must be divisible by 16", - ), - ({"weight": self._desc((1024,))}, "W must have shape"), - ( - { - "x": self.tensor_module.TensorDesc( - dtype=_DataType.BFLOAT16, - shape=(256, 2048), - stride=(1, 256), - stride_order=(0, 1), - ) - }, - "X must be row-major contiguous", - ), - ({"num_threads": 512}, "EPT=4 must be >= 8 and divisible by 8"), - ({"num_threads": 2048}, "must not exceed the CUDA block size limit"), - ({"rows_per_cta": 0}, "rows_per_cta must be positive"), - ({"rows_per_cta": 3}, "M must be divisible by rows_per_cta"), - ) - - for overrides, message in cases: - with self.subTest(message=message): - with self.assertRaisesRegex(ValueError, message): - self._kernel(**overrides).check_support() - with self.assertRaisesRegex(TypeError, "x must be a TensorDesc"): - self.kernel_module.RMSNormRHTAmaxKernel( - x=types.SimpleNamespace(shape=(256, 2048)), - weight=self._desc((2048,)), - ) - with self.assertRaisesRegex(TypeError, "weight must be a TensorDesc"): - self.kernel_module.RMSNormRHTAmaxKernel( - x=self._desc((256, 2048)), - weight=None, - ) + self.assertEqual(ast.literal_eval(assignment.value), self.kernel_module.RMSNormRHTAmaxKernel.HAD_BLOCK) if __name__ == "__main__": diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_op.py b/test/python/fe_api/test_rmsnorm_rht_amax_op.py new file mode 100644 index 000000000..64815b5d5 --- /dev/null +++ b/test/python/fe_api/test_rmsnorm_rht_amax_op.py @@ -0,0 +1,186 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for the RMSNorm logical operation.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "rmsnorm_rht_amax" +_PACKAGE = "cudnn_rmsnorm_op_test" + + +class _DataType(Enum): + NOT_SET = auto() + FLOAT = auto() + BFLOAT16 = auto() + + +class RmsNormRhtAmaxOpContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + operation_name = f"{_PACKAGE}.rmsnorm_rht_amax" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + try: + cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}._op") + cls.op_module = importlib.import_module(f"{operation_name}.op") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def _desc(self, shape, dtype=_DataType.BFLOAT16, *, stride=None, name=""): + shape = tuple(shape) + if stride is None: + stride_order = tuple(reversed(range(len(shape)))) + stride_values = [0] * len(shape) + running = 1 + for dimension in stride_order: + stride_values[dimension] = running + running *= max(shape[dimension], 1) + stride = tuple(stride_values) + else: + stride = tuple(stride) + stride_order = tuple(index for index, _ in sorted(enumerate(stride), key=lambda item: (item[1], shape[item[0]]))) + return self.tensor_module.TensorDesc( + dtype=dtype, + shape=shape, + stride=stride, + stride_order=stride_order, + name=name, + ) + + def _op(self, **overrides): + arguments = { + "x": self._desc((256, 2048), name="x"), + "weight": self._desc((2048,), name="weight"), + "output": self._desc((256, 2048), name="output"), + "amax": self._desc((128,), _DataType.FLOAT, name="amax"), + } + arguments.update(overrides) + return self.op_module.RmsNormRhtAmaxSm100Op(**arguments) + + def test_operation_validates_complete_signature_and_resolves_configuration(self): + operation = self._op() + + self.assertIsInstance(operation, self.base_module.Op) + self.assertTrue(operation.check_support()) + self.assertEqual((operation.m, operation.n), (256, 2048)) + self.assertEqual((operation.num_threads, operation.rows_per_cta), (128, 2)) + self.assertFalse(hasattr(operation, "infer_output")) + self.assertFalse(hasattr(operation, "check_output")) + self.assertNotIn(f"{_PACKAGE}.rmsnorm_rht_amax.kernel", sys.modules) + self.assertNotIn(f"{_PACKAGE}.rmsnorm_rht_amax.api", sys.modules) + self.assertNotIn(f"{_PACKAGE}.rmsnorm_rht_amax.jax", sys.modules) + + def test_positive_strided_amax_is_supported_but_broadcast_amax_is_not(self): + self.assertTrue(self._op(amax=self._desc((128,), _DataType.FLOAT, stride=(2,))).check_support()) + + with self.assertRaisesRegex(ValueError, "Amax stride must be positive"): + self._op(amax=self._desc((128,), _DataType.FLOAT, stride=(0,))).check_support() + + def test_framework_descriptors_are_compared_through_cudnn_dtype(self): + tensor_module = self.tensor_module + + class FrameworkTensorDesc(tensor_module.TensorDesc): + @property + def cudnn_dtype(self): + return { + "framework.bfloat16": _DataType.BFLOAT16, + "framework.float32": _DataType.FLOAT, + }.get(self.dtype, _DataType.NOT_SET) + + def desc(shape, dtype, *, stride=None): + canonical = self._desc(shape, stride=stride) + return FrameworkTensorDesc( + dtype=dtype, + shape=canonical.shape, + stride=canonical.stride, + stride_order=canonical.stride_order, + ) + + operation = self.op_module.RmsNormRhtAmaxSm100Op( + x=desc((256, 2048), "framework.bfloat16"), + weight=desc((2048,), "framework.bfloat16"), + output=desc((256, 2048), "framework.bfloat16"), + amax=desc((128,), "framework.float32"), + ) + self.assertTrue(operation.check_support()) + + def test_invalid_complete_signatures_and_configuration(self): + cases = ( + ({"x": self._desc((2048,))}, "X must have rank 2"), + ({"weight": self._desc((1, 2048))}, "W must have rank 1"), + ({"output": self._desc((256 * 2048,))}, "O must have rank 2"), + ({"amax": self._desc((128, 1), _DataType.FLOAT)}, "Amax must have rank 1"), + ({"x": self._desc((256, 2048), _DataType.FLOAT)}, "X must have dtype bfloat16"), + ({"weight": self._desc((2048,), _DataType.FLOAT)}, "W must have dtype bfloat16"), + ({"output": self._desc((256, 2048), _DataType.FLOAT)}, "O must have dtype bfloat16"), + ({"amax": self._desc((128,), _DataType.BFLOAT16)}, "Amax must have dtype float32"), + ({"x": self._desc((0, 2048))}, "M must be positive"), + ( + { + "x": self._desc((256, 2049)), + "weight": self._desc((2049,)), + "output": self._desc((256, 2049)), + }, + "N must be divisible by 16", + ), + ({"weight": self._desc((1024,))}, "W must have shape"), + ({"output": self._desc((128, 2048))}, "O must have shape"), + ({"amax": self._desc((64,), _DataType.FLOAT)}, "Amax must have shape"), + ({"x": self._desc((256, 2048), stride=(1, 256))}, "X must be row-major contiguous"), + ({"weight": self._desc((2048,), stride=(2,))}, "W must be contiguous"), + ({"output": self._desc((256, 2048), stride=(1, 256))}, "O must be row-major contiguous"), + ({"num_threads": 512}, "EPT=4 must be >= 8 and divisible by 8"), + ({"num_threads": 2048}, "must not exceed the CUDA block size limit"), + ({"rows_per_cta": 0}, "rows_per_cta must be positive"), + ({"rows_per_cta": 3}, "M must be divisible by rows_per_cta"), + ) + + for overrides, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self._op(**overrides).check_support() + + def test_constructor_requires_tensor_descriptors(self): + for name in ("x", "weight", "output", "amax"): + with self.subTest(name=name): + with self.assertRaisesRegex(TypeError, f"{name} must be a TensorDesc"): + self._op(**{name: types.SimpleNamespace(shape=(1,))}) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_tensor_desc.py b/test/python/fe_api/test_tensor_desc.py index 69fb7995b..4d55bac01 100644 --- a/test/python/fe_api/test_tensor_desc.py +++ b/test/python/fe_api/test_tensor_desc.py @@ -106,32 +106,6 @@ def test_api_base_converts_torch_tensor_metadata(self): self.assertEqual(desc.stride, (12, 1, 3)) self.assertEqual(desc.stride_order, (1, 2, 0)) - def test_api_base_checks_torch_tensor_signature(self): - class Adapter(APIBase): - def check_support(self): - return True - - def compile(self): - pass - - def execute(self, *args, **kwargs): - pass - - adapter = Adapter() - sample = torch.empty_strided((2, 3), (3, 1), dtype=torch.bfloat16) - expected = TorchTensorDesc.from_tensor(sample, "sample") - adapter._check_tensor_signature(sample, expected) - - cases = ( - (torch.empty((1, 3), dtype=torch.bfloat16), "sample tensor shape mismatch"), - (torch.empty_strided((2, 3), (1, 2), dtype=torch.bfloat16), "sample tensor stride mismatch"), - (torch.empty((2, 3), dtype=torch.float32), "sample dtype mismatch"), - ) - for tensor, message in cases: - with self.subTest(message=message): - with self.assertRaisesRegex(ValueError, message): - adapter._check_tensor_signature(tensor, expected) - def test_torch_descriptor_constructs_from_tensor_metadata(self): tensor = torch.empty_strided((2, 3, 4), (12, 1, 3), dtype=torch.uint8) desc = TorchTensorDesc.from_tensor( From 8a63d686352b6dcf1912a80507505c91fab77715 Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Tue, 7 Jul 2026 19:21:38 -0500 Subject: [PATCH 04/13] Add JAX support for DeepSeek sparse attention --- docs/fe-oss-apis/dsa.md | 205 ++++- python/cudnn/_jax/__init__.py | 3 +- python/cudnn/_jax/api_base.py | 132 +++- python/cudnn/_jax/layout.py | 23 + python/cudnn/_jax/result.py | 40 + python/cudnn/_result.py | 34 + python/cudnn/_tensor_desc.py | 20 + python/cudnn/api_base.py | 33 +- .../cudnn/deepseek_sparse_attention/README.md | 9 +- .../indexer_backward/__init__.py | 31 +- .../indexer_backward/api.py | 150 ++-- .../dense_indexer_backward_sm100.py | 42 +- .../dense_indexer_backward_sm90.py | 52 +- .../indexer_backward_sm100.py | 50 +- .../indexer_backward/indexer_backward_sm90.py | 42 +- .../indexer_backward/jax.py | 685 +++++++++++++++++ .../indexer_backward/op.py | 438 +++++++++++ .../indexer_forward/__init__.py | 20 +- .../indexer_forward/api.py | 85 +-- .../indexer_forward/jax.py | 434 +++++++++++ .../indexer_forward/op.py | 229 ++++++ .../indexer_top_k/__init__.py | 31 +- .../indexer_top_k/api.py | 89 ++- .../indexer_top_k/compactify.py | 22 +- .../indexer_top_k_decode_varlen.py | 25 +- .../indexer_top_k_varlen_util.py | 11 +- .../indexer_top_k/jax.py | 404 ++++++++++ .../indexer_top_k/local_to_global_dsl.py | 23 +- .../indexer_top_k/op.py | 194 +++++ .../score_recompute/__init__.py | 57 +- .../score_recompute/_interface_sm100.py | 260 ++----- .../score_recompute/api.py | 165 ++-- .../score_recompute/config.py | 292 +++++++ .../score_recompute/jax.py | 714 ++++++++++++++++++ .../score_recompute/op.py | 427 +++++++++++ .../sparse_attention_backward/__init__.py | 21 +- .../_interface_sm100.py | 7 +- .../sparse_attention_backward/api.py | 106 ++- .../dsa_bwd_sm100.py | 78 +- .../sparse_attention_backward/dsa_bwd_sm90.py | 8 +- .../sparse_attention_backward/jax.py | 457 +++++++++++ .../sparse_attention_backward/op.py | 226 ++++++ .../utils/compiler.py | 47 +- python/cudnn/jax/__init__.py | 98 ++- test/python/fe_api/dsa/dsa_reference.py | 12 +- test/python/fe_api/dsa/dsa_utils.py | 2 +- .../dsa/test_DSA_sparse_attention_backward.py | 61 +- test/python/fe_api/test_dsa_compiler.py | 52 ++ .../fe_api/test_dsa_indexer_backward_op.py | 231 ++++++ test/python/fe_api/test_dsa_indexer_ops.py | 534 +++++++++++++ .../fe_api/test_dsa_score_recompute_op.py | 208 +++++ .../test_dsa_sparse_attention_backward_op.py | 193 +++++ test/python/fe_api/test_jax_api_base.py | 55 ++ test/python/fe_api/test_jax_dsa_backward.py | 347 +++++++++ .../test_jax_dsa_indexer_backward_contract.py | 551 ++++++++++++++ test/python/fe_api/test_jax_dsa_indexers.py | 278 +++++++ .../fe_api/test_jax_dsa_score_recompute.py | 360 +++++++++ .../test_jax_dsa_score_recompute_contract.py | 475 ++++++++++++ ..._dsa_sparse_attention_backward_contract.py | 385 ++++++++++ .../python/fe_api/test_jax_import_contract.py | 225 ++++++ test/python/fe_api/test_jax_layout.py | 39 + test/python/fe_api/test_result.py | 60 ++ test/python/fe_api/test_shared_tensor_desc.py | 21 + 63 files changed, 9830 insertions(+), 778 deletions(-) create mode 100644 python/cudnn/_jax/result.py create mode 100644 python/cudnn/_result.py create mode 100644 python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py create mode 100644 python/cudnn/deepseek_sparse_attention/indexer_backward/op.py create mode 100644 python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py create mode 100644 python/cudnn/deepseek_sparse_attention/indexer_forward/op.py create mode 100644 python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py create mode 100644 python/cudnn/deepseek_sparse_attention/indexer_top_k/op.py create mode 100644 python/cudnn/deepseek_sparse_attention/score_recompute/config.py create mode 100644 python/cudnn/deepseek_sparse_attention/score_recompute/jax.py create mode 100644 python/cudnn/deepseek_sparse_attention/score_recompute/op.py create mode 100644 python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py create mode 100644 python/cudnn/deepseek_sparse_attention/sparse_attention_backward/op.py create mode 100644 test/python/fe_api/test_dsa_compiler.py create mode 100644 test/python/fe_api/test_dsa_indexer_backward_op.py create mode 100644 test/python/fe_api/test_dsa_indexer_ops.py create mode 100644 test/python/fe_api/test_dsa_score_recompute_op.py create mode 100644 test/python/fe_api/test_dsa_sparse_attention_backward_op.py create mode 100644 test/python/fe_api/test_jax_dsa_backward.py create mode 100644 test/python/fe_api/test_jax_dsa_indexer_backward_contract.py create mode 100644 test/python/fe_api/test_jax_dsa_indexers.py create mode 100644 test/python/fe_api/test_jax_dsa_score_recompute.py create mode 100644 test/python/fe_api/test_jax_dsa_score_recompute_contract.py create mode 100644 test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py create mode 100644 test/python/fe_api/test_jax_import_contract.py create mode 100644 test/python/fe_api/test_result.py diff --git a/docs/fe-oss-apis/dsa.md b/docs/fe-oss-apis/dsa.md index 9ec5f3400..c5db25664 100644 --- a/docs/fe-oss-apis/dsa.md +++ b/docs/fe-oss-apis/dsa.md @@ -6,10 +6,9 @@ The DeepSeek Sparse Attention (DSA) module integrates a set of CuTe-DSL kernels that support the sparse-attention path used by DeepSeek-style models. -Most kernels target Hopper (SM90) and Blackwell (SM100+) GPUs; Indexer -Forward and Indexer Top-K remain SM100+ only. The kernels are -delivered as Python classes / wrappers that follow the same `APIBase` -pattern as other cuDNN Frontend operations. +The kernels target Hopper (SM90) and Blackwell (SM100+) GPUs and are exposed +through PyTorch and JAX classes and convenience wrappers. See the JAX support +matrix below for operation-specific layout restrictions. **Scope:** this module ships CuTe-DSL kernels for DSA backward, indexer scores/top-K, sparse/dense score recompute, and sparse/dense indexer @@ -24,7 +23,7 @@ The module packages the following operations: 2. **Indexer Forward** – CuTe-DSL score kernel (Q @ K^T, ReLU, head reduce, ratio causal mask). Non-fused; pair with **Indexer Top-K** for the top-K step. -3. **Indexer Top-K** – SM100 CuTe-DSL radix top-K kernel with per-row +3. **Indexer Top-K** – CuTe-DSL radix top-K kernel with per-row ``seq_lens``. 4. **Sparse Indexer / Attention Score Recompute** – sparse (top-K) recompute of indexer and attention scores for training loss. @@ -59,15 +58,23 @@ Training-score loss path: ## Installation +PyTorch: + +```bash +pip install 'nvidia-cudnn-frontend[cutedsl]' +``` + +JAX: + ```bash -pip install nvidia-cudnn-frontend[cutedsl] +pip install 'nvidia-cudnn-frontend[jax]' ``` --- ## API Usage -### DSA Namespace +### PyTorch namespace ```python from cudnn import DSA @@ -100,6 +107,129 @@ DSA.DenseIndexerBackward DSA.dense_indexer_backward_wrapper ``` +### JAX namespace and wrappers + +Importing `cudnn.jax` validates the optional JAX and CUTLASS JAX +dependencies. DSA operation classes and wrappers are then loaded lazily. Each +public symbol is available both directly and through `cudnn.jax.DSA`: + +```python +import cudnn.jax as cudnn_jax + +cudnn_jax.IndexerForward +cudnn_jax.indexer_forward_wrapper + +cudnn_jax.DSA.IndexerForward +cudnn_jax.DSA.indexer_forward_wrapper +``` + +The function wrappers are already decorated with `jax.jit`; their tuning and +target arguments, including maximum sequence lengths, are static compilation +options. They infer output metadata and return a `TupleDict` of JAX arrays: + +```python +result = cudnn_jax.indexer_forward_wrapper(q, k, w, ratio=4) +scores = result["scores"] +``` + +Layout strings describe the public JAX axis order and are static compilation +arguments. `IndexerForward` accepts batch-major or sequence-major fixed +inputs while keeping the head dimension contiguous: + +```python +result = cudnn_jax.indexer_forward_wrapper( + q_sbhd, + k_sbhd, + w_sbh, + q_layout="SBHD", + k_layout="SBHD", + w_layout="SBH", + output_layout="SBK", + ratio=4, +) +scores_sbk = result["scores"] +``` + +The defaults are `q_layout="BSHD"`, `k_layout="BSHD"`, +`w_layout="BSH"`, and `output_layout="BSK"`. Packed inputs use `THD` and +`TH`, with a `TK` output. The adapter maps these public axes to the kernel's +canonical axes; callers do not need to transpose arrays into canonical order. + +Use the class API when the caller needs to choose the enclosing JAX +transformation. Classes specialize from array-like exemplars that expose +`shape` and `dtype`; `jax.ShapeDtypeStruct` makes that intent explicit: + +```python +import jax + +q_spec = jax.ShapeDtypeStruct(q.shape, q.dtype) +k_spec = jax.ShapeDtypeStruct(k.shape, k.dtype) +w_spec = jax.ShapeDtypeStruct(w.shape, w.dtype) + +op = cudnn_jax.IndexerForward(q_spec, k_spec, w_spec, ratio=4) +scores = jax.jit(op)(q, k, w)["scores"] +``` + +The JAX adapters allocate outputs and private workspaces through XLA. Where a +class accepts optional `sample_*` output arguments, those exemplars validate +an explicit output signature; otherwise the adapter infers it. Backward APIs +are functional: they do not update caller-owned gradient buffers in place. +For `IndexerForward`, an explicit `sample_out` describes the physical FP32 +buffer whose last extent is rounded up to four; the returned `scores` view is +sliced back to the logical K extent. +`SparseAttentionBackward` returns `dq`, `dkv`, and `d_sink`; +`IndexerBackward` and `DenseIndexerBackward` return `d_index_q`, `d_weights`, +and `d_index_k`. + +### JAX target selection + +`target_compute_capability` is the exact compilation target, such as `90`, +`100`, `103`, or `107`. When it is omitted, the adapter infers it from a +homogeneous set of local JAX GPUs. With compatibility checks enabled, an +explicit target must match every local GPU. + +For AOT or remote compilation without the target GPU, disable the local-device +check and provide the target before tracing: + +```python +cudnn_jax.disable_device_compatibility_checks(True) +lowered = cudnn_jax.indexer_forward_wrapper.lower( + q_spec, + k_spec, + w_spec, + target_compute_capability=100, +) +``` + +Pass `False` to re-enable checks. This process-wide setting only disables the +local-device validation; it does not enable an unsupported kernel target. +Changing it does not invalidate existing JAX compilation-cache entries, so +configure it before the first `jax.jit` or `lower` call for a signature. + +### JAX support matrix + +The SM100+ column covers the implemented SM100-family targets: SM100, SM103, +and SM107. + +| Operation | JAX layout | SM90 | SM100+ | +| --- | --- | --- | --- | +| Sparse Attention Backward | Flat MQA tensors | Yes | Yes; inputs must be BF16 | +| Indexer Forward | Fixed BSHD/SBHD or packed THD | Yes; `H_kv=1` and default tuning | Yes | +| Indexer Top-K and index utilities | Flattened rows; fixed or packed index conversion | Yes | Yes | +| Sparse score recompute | Fixed BSHD/MQA | Yes | Yes | +| Dense score recompute | Fixed BSHD/MQA | Yes | Yes | +| Dense score recompute | Packed THD/MQA | No | Yes | +| Indexer Backward | Fixed BSHD; JAX-only packed THD with global indices | Yes | Yes | +| Dense Indexer Backward | Fixed BSHD or packed THD | Yes | Yes | + +SM90 packed THD dense score recompute is rejected because that backend needs +host-side cumulative-length reads, which cannot be represented during JAX +tracing. Sparse score recompute has no packed THD signature and its SM90 +kernel requires a top-K extent divisible by 128. Both indexer +backward variants require `H >= 64`, `D == 128`, and `block_I == 128`; sparse +Indexer Backward also requires its top-K extent to be divisible by 128. Its +packed THD signature requires `topk_indices_global=True` on JAX. + --- ## Components @@ -118,7 +248,12 @@ Backward pass for DeepSeek Sparse Attention. Expects the forward outputs - `topk_idxs`: `(total_S_q, topk_max)` INT32 (global) - `topk_length` (optional): `(total_S_q,)` INT32 — per-query valid count - **Outputs** — tuple `(dq, dkv, d_sink)` -- **Constraints** — SM90 or SM100; SM90 supports the FlashMLA DSA shape with `head_dim ∈ {512, 576}` +- **Constraints** — SM90 or SM100+, `head_dim ∈ {512, 576}`. SM100+ + currently requires BF16 inputs; SM90 also supports FP16. Every + `topk_length` value is interpreted on device and clamped to + `[0, topk_max]`; zero-length rows are supported. When `topk_length` is + supplied, every index before it must be in `[0, total_S_kv)`. When it is + omitted, negative entries are padding sentinels. ```python result = DSA.sparse_attention_backward_wrapper( @@ -145,13 +280,17 @@ The K columns are assumed to be a compressed-KV prefix starting at global compressed column 0. - **Inputs** - - `q`: `(B, S_q, H_q, D)` BF16 - - `k`: `(B, S_k, H_kv, D)` BF16 - - `w`: `(B, S_q, H_q)` BF16 - - `q_causal_offsets` (optional): CUDA INT32 tensor with one entry per + - Fixed: `q` `(B, S_q, H_q, D)`, `k` `(B, S_k, H_kv, D)`, and + `w` `(B, S_q, H_q)`, all BF16. + - Packed THD: `q` `(T_q, H_q, D)`, `k` `(T_k, H_kv, D)`, and + `w` `(T_q, H_q)`, with `cu_seqlens_q`, `cu_seqlens_k`, + `max_seqlen_q`, and `max_seqlen_k`. + - `q_causal_offsets` (optional): INT32 device array with one entry per batch/THD segment, on the same device as `q`. -- **Output** — `scores`: `(B, S_q, S_k)` FP32 -- **Constraints** — SM100+, `head_dim == 128`, `qhead_per_kv_head ∈ {32, 64}` +- **Output** — `scores`: `(B, S_q, S_k)` or `(T_q, max_seqlen_k)` FP32. +- **Constraints** — SM90 or SM100+, `head_dim == 128`, + `qhead_per_kv_head ∈ {32, 64}`. SM90 requires `H_kv == 1` and the + default tuning configuration. ```python result = DSA.indexer_forward_wrapper( @@ -170,7 +309,8 @@ with variable per-row effective length. - `seq_lens`: `(batch_size,)` INT32 (per-batch effective column count) - **Outputs** — tuple `(indices, values)` (values is `None` when `return_val=False`) -- **Constraints** — SM100+, `top_k ≤ 2048` +- **Constraints** — SM90 or SM100+, `top_k ≤ min(2048, num_cols)` and + `top_k + next_n - 1 ≤ seq_lens[b] ≤ num_cols` for every batch item. ```python result = DSA.indexer_top_k_wrapper( @@ -180,6 +320,10 @@ result = DSA.indexer_top_k_wrapper( indices, values = result["indices"], result["values"] ``` +JAX also exposes `local_to_global_wrapper` for fixed or packed local index +conversion and `compactify_wrapper` for packing nonnegative indices and +returning each row's valid `topk_length`. + ### 4. Sparse Indexer Score Recompute Computes softmax over top-K entries of the indexer score: @@ -207,21 +351,34 @@ L1-normalised head-summed softmax over top-K entries: Full-KV (no top-K) analogues of §4 and §5. Each returns `{'out', 'denom'}`. They apply the same ratio-causal mask as Indexer Forward; masked positions are -written as zero and excluded from `denom`. Pass the same `q_causal_offsets` to -all dense score tensors that feed the same loss path. +excluded from `denom`. JAX initializes skipped output positions to `-inf`. +Pass the same `q_causal_offsets` to all dense score tensors that feed the same +loss path. ### 7. Indexer Backward Three-stage sparse top-K pipeline that produces the training gradients for the indexer tower: -1. `ScoreGradSm90` / `ScoreGradSm100` (kernel 1) — in-place score-grad precompute from - `attn_score` (target) and `index_score` (predict). +1. `ScoreGradSm90` / `ScoreGradSm100` (kernel 1) — score-grad precompute from + `attn_score` (target) and `index_score` (predict) into a private temporary. 2. `IndexerBackwardSm90` / `IndexerBackwardSm100` (kernel 2) — three warp-specialised GEMMs produce `d_index_q`, `d_weights`, and a `dIndexK_f32` accumulator. -3. Pure-torch dtype cast (kernel 3) converts `dIndexK_f32` to the output - dtype. +3. A dtype cast converts `dIndexK_f32` to the output dtype. + +The JAX wrapper treats `grad_loss` as a runtime array operand and returns new +gradient arrays; no input or caller-owned output buffer is mutated. + +The common fixed signature uses `index_q` `(B, S_q, H, 128)`, `weights` +`(B, S_q, H)`, `index_k` `(B, S_k, 128)`, and score/index tensors +`(B, S_q, topk)`. JAX additionally accepts a packed signature with +`index_q` `(T_q, H, 128)`, `weights` `(T_q, H)`, `index_k` `(T_k, 128)`, and +score/index tensors `(T_q, topk)`. The packed adapter presents these tensors +to the kernel as a synthetic batch of one, so it requires +`topk_indices_global=True`; each top-K entry must address the flattened packed +K range directly. It does not infer original batch-local IDs from cumulative +sequence lengths. The PyTorch sparse wrapper remains fixed BSHD. **The TileLang fallback present in the upstream repo is dropped here (CuTe-DSL only).** If the CuTe-DSL path fails the wrapper raises @@ -278,11 +435,13 @@ result = DSA.dense_indexer_backward_wrapper( ## Limitations - **Architecture support** — Sparse Attention Backward, Score Recompute, and - Indexer Backward support SM90 and SM100; Indexer Forward and Indexer Top-K - remain SM100+ only. + Indexer operations support SM90 and SM100+ subject to the layout restrictions + in the JAX support matrix. - **No fused forward** — the production forward is FlashMLA (C++); this module ships only the CuTe-DSL kernels. - **Indexer Forward only supports `head_dim = 128`** and `qhead_per_kv_head ∈ {32, 64}`. - **Top-K only up to 2048**; `top_k > 2048` is not supported by the underlying radix top-K kernel. +- **JAX Top-K does not use the PyTorch row-chunking fallback** when its private + workspace would exceed the INT32 indexing range. diff --git a/python/cudnn/_jax/__init__.py b/python/cudnn/_jax/__init__.py index e2ea61e8e..4048c3437 100644 --- a/python/cudnn/_jax/__init__.py +++ b/python/cudnn/_jax/__init__.py @@ -4,5 +4,6 @@ """Internal implementation helpers for optional JAX operation adapters.""" from .api_base import JaxApiBase, JaxTensorDesc, disable_device_compatibility_checks +from .result import TupleDict -__all__ = ["JaxApiBase", "JaxTensorDesc", "disable_device_compatibility_checks"] +__all__ = ["JaxApiBase", "JaxTensorDesc", "TupleDict", "disable_device_compatibility_checks"] diff --git a/python/cudnn/_jax/api_base.py b/python/cudnn/_jax/api_base.py index 7ed7c76ff..3180b1f34 100644 --- a/python/cudnn/_jax/api_base.py +++ b/python/cudnn/_jax/api_base.py @@ -15,6 +15,7 @@ _DEVICE_COMPATIBILITY_CHECKS_DISABLED = False _DEVICE_COMPATIBILITY_CHECK_DISABLE_HINT = "Call cudnn.jax.disable_device_compatibility_checks(True) before tracing to allow AOT or remote compilation." +_TARGET_COMPUTE_CAPABILITY_HINT = "Pass target_compute_capability explicitly when tracing for AOT or on a host without the target GPU." def disable_device_compatibility_checks(disabled: bool) -> None: @@ -80,6 +81,107 @@ def _require_array_metadata(value: Any, name: str) -> tuple[Any, ...]: class JaxApiBase(ABC): """Common tensor metadata, validation, and CuTe binding for JAX adapters.""" + @staticmethod + def _device_compute_capability(device: Any, operation_name: str) -> int: + """Normalize JAX's device capability metadata to ``major * 10 + minor``.""" + + reported = getattr(device, "compute_capability", None) + try: + if isinstance(reported, (tuple, list)): + if len(reported) < 2: + raise ValueError + major, minor = int(reported[0]), int(reported[1]) + else: + text = str(reported) + if "." in text: + major_text, minor_text = text.split(".", 1) + major, minor = int(major_text), int(minor_text) + else: + capability = int(text) + major, minor = divmod(capability, 10) + if major < 0 or minor < 0 or minor > 9: + raise ValueError + except (TypeError, ValueError) as error: + raise RuntimeError(f"{operation_name}: JAX reported an invalid compute capability {reported!r} for {device}") from error + return major * 10 + minor + + @staticmethod + def _local_gpu_capabilities(operation_name: str) -> tuple[tuple[Any, int], ...]: + import jax + + try: + devices = tuple(jax.local_devices(backend="gpu")) + except RuntimeError as error: + raise RuntimeError(f"{operation_name}: JAX could not discover a local GPU") from error + return tuple((device, JaxApiBase._device_compute_capability(device, operation_name)) for device in devices) + + @staticmethod + def _compute_capability_family( + compute_capability: int, + supported_compute_capabilities: tuple[int, ...], + ) -> int | None: + compatible = tuple(capability for capability in supported_compute_capabilities if capability <= compute_capability) + return max(compatible, default=None) + + @staticmethod + def _resolve_compute_capability( + target_compute_capability: int | None, + supported_compute_capabilities: tuple[int, ...], + operation_name: str, + ) -> int: + """Resolve an exact JAX compilation target for a multi-arch adapter. + + ``supported_compute_capabilities`` lists the exact CuTe compilation + targets implemented by the adapter. The returned value remains exact + (for example ``103``) so the compiler selects the matching target. + With compatibility checks enabled, an explicit target must match every + local GPU exactly. Disable checks for cross-compilation. + """ + + supported = tuple(sorted(set(supported_compute_capabilities))) + if not supported or any(isinstance(value, bool) or not isinstance(value, int) or value <= 0 for value in supported): + raise ValueError(f"supported_compute_capabilities must contain positive integers, got {supported_compute_capabilities!r}") + + if target_compute_capability is not None: + if isinstance(target_compute_capability, bool) or not isinstance(target_compute_capability, int): + raise TypeError(f"target_compute_capability must be an int or None, got {type(target_compute_capability).__name__}") + if target_compute_capability <= 0: + raise ValueError(f"target_compute_capability must be positive, got {target_compute_capability}") + if target_compute_capability not in supported: + supported_text = ", ".join(f"SM{value}" for value in supported) + raise ValueError(f"{operation_name} has no kernel for SM{target_compute_capability}; supported targets are {supported_text}") + if _DEVICE_COMPATIBILITY_CHECKS_DISABLED: + return target_compute_capability + + local = JaxApiBase._local_gpu_capabilities(operation_name) + if not local: + raise RuntimeError( + f"{operation_name} targets SM{target_compute_capability}, but no local JAX GPU is available. " f"{_DEVICE_COMPATIBILITY_CHECK_DISABLE_HINT}" + ) + mismatched = tuple((device, capability) for device, capability in local if capability != target_compute_capability) + if mismatched: + found = ", ".join(f"{device} (SM{capability})" for device, capability in mismatched) + raise RuntimeError(f"{operation_name} targets SM{target_compute_capability}, but found {found}. " f"{_DEVICE_COMPATIBILITY_CHECK_DISABLE_HINT}") + return target_compute_capability + + try: + local = JaxApiBase._local_gpu_capabilities(operation_name) + except RuntimeError as error: + raise RuntimeError(f"{error}. {_TARGET_COMPUTE_CAPABILITY_HINT}") from error + if not local: + raise RuntimeError(f"{operation_name}: no local JAX GPU is available. {_TARGET_COMPUTE_CAPABILITY_HINT}") + + capabilities = tuple(sorted({capability for _, capability in local})) + if len(capabilities) != 1: + found = ", ".join(f"SM{capability}" for capability in capabilities) + raise RuntimeError(f"{operation_name}: local JAX GPUs have heterogeneous targets ({found}). {_TARGET_COMPUTE_CAPABILITY_HINT}") + + resolved = capabilities[0] + if resolved not in supported: + supported_text = ", ".join(f"SM{value}" for value in supported) + raise RuntimeError(f"{operation_name} has no kernel for SM{resolved}; supported targets are {supported_text}") + return resolved + @staticmethod def _check_device_compatibility( *, @@ -91,38 +193,28 @@ def _check_device_compatibility( if _DEVICE_COMPATIBILITY_CHECKS_DISABLED: return - import jax - def compatibility_error(reason: str) -> RuntimeError: return RuntimeError(f"{operation_name} requires SM{minimum_compute_capability}+, {reason}. " f"{_DEVICE_COMPATIBILITY_CHECK_DISABLE_HINT}") try: - devices = tuple(jax.local_devices(backend="gpu")) + devices = JaxApiBase._local_gpu_capabilities(operation_name) except RuntimeError as error: - raise compatibility_error("but JAX could not discover a local GPU") from error + reason = str(error) + prefix = f"{operation_name}: " + if reason.startswith(prefix): + reason = reason[len(prefix) :] + raise compatibility_error(f"but {reason}") from error if not devices: raise compatibility_error("but no local JAX GPU is available") - minimum_major, minimum_minor = divmod(minimum_compute_capability, 10) incompatible = [] - for device in devices: - reported_capability = getattr(device, "compute_capability", None) - try: - components = str(reported_capability).split(".") - if len(components) != 2: - raise ValueError - major, minor = (int(value) for value in components) - if major < 0 or minor < 0: - raise ValueError - except (TypeError, ValueError) as error: - raise compatibility_error(f"but JAX reported an invalid compute capability {reported_capability!r} for {device}") from error - - if (major, minor) < (minimum_major, minimum_minor): - incompatible.append((device, major, minor)) + for device, compute_capability in devices: + if compute_capability < minimum_compute_capability: + incompatible.append((device, compute_capability)) if incompatible: - found = ", ".join(f"{device} (SM{major}{minor})" for device, major, minor in incompatible) + found = ", ".join(f"{device} (SM{compute_capability})" for device, compute_capability in incompatible) raise compatibility_error(f"found {found}") @staticmethod diff --git a/python/cudnn/_jax/layout.py b/python/cudnn/_jax/layout.py index f2be19940..14dd0b7e9 100644 --- a/python/cudnn/_jax/layout.py +++ b/python/cudnn/_jax/layout.py @@ -11,6 +11,28 @@ AxisValueT = TypeVar("AxisValueT") +def mode_from_layout(layout: str, *, kernel_axes: str) -> tuple[int, ...]: + """Map named kernel axes to their positions in a public JAX layout. + + Axis names are single, case-sensitive characters. The returned mapping + follows ``mode[kernel_axis] = public_axis``. + """ + + if not isinstance(layout, str): + raise TypeError(f"layout must be a string, got {layout!r}") + if not isinstance(kernel_axes, str): + raise TypeError(f"kernel_axes must be a string, got {kernel_axes!r}") + if len(layout) != len(kernel_axes): + raise ValueError(f"layout rank must match kernel_axes: got {layout!r} and {kernel_axes!r}") + if len(set(layout)) != len(layout): + raise ValueError(f"layout axes must be unique, got {layout!r}") + if len(set(kernel_axes)) != len(kernel_axes): + raise ValueError(f"kernel_axes must be unique, got {kernel_axes!r}") + if set(layout) != set(kernel_axes): + raise ValueError(f"layout must contain exactly the axes in {kernel_axes!r}, got {layout!r}") + return tuple(layout.index(axis) for axis in kernel_axes) + + def normalize_mode(rank: int, mode: tuple[int, ...] | None = None) -> tuple[int, ...]: """Return and validate a kernel-axis to public-axis mapping. @@ -105,6 +127,7 @@ def to_cutlass_layout( __all__ = [ "compact_stride", + "mode_from_layout", "normalize_mode", "to_canonical_axes", "to_cutlass_layout", diff --git a/python/cudnn/_jax/result.py b/python/cudnn/_jax/result.py new file mode 100644 index 000000000..c19ddcc80 --- /dev/null +++ b/python/cudnn/_jax/result.py @@ -0,0 +1,40 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX pytree registration for framework-neutral frontend results.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import jax + +from .._result import TupleDict + + +def _flatten(value: TupleDict) -> tuple[tuple[Any, ...], tuple[Any, ...]]: + keys = tuple(dict.keys(value)) + return tuple(dict.__getitem__(value, key) for key in keys), keys + + +def _flatten_with_keys(value: TupleDict): + children, keys = _flatten(value) + return tuple((jax.tree_util.DictKey(key), child) for key, child in zip(keys, children)), keys + + +def _unflatten(keys: tuple[Any, ...], children: Iterable[Any]) -> TupleDict: + return TupleDict(zip(keys, children)) + + +if not getattr(TupleDict, "_jax_pytree_registered", False): + jax.tree_util.register_pytree_with_keys( + TupleDict, + _flatten_with_keys, + _unflatten, + _flatten, + ) + TupleDict._jax_pytree_registered = True + + +__all__ = ["TupleDict"] diff --git a/python/cudnn/_result.py b/python/cudnn/_result.py new file mode 100644 index 000000000..eca466092 --- /dev/null +++ b/python/cudnn/_result.py @@ -0,0 +1,34 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral result containers shared by frontend APIs.""" + +from __future__ import annotations + +from typing import Any + + +class TupleDict(dict): + """Dictionary result that also unpacks and indexes like a tuple. + + Values follow insertion order, matching the result order documented by + the operation wrapper. String-key access remains ordinary dictionary + access. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._keys = tuple(self.keys()) + + def __iter__(self): + return (dict.__getitem__(self, key) for key in self._keys) + + def __getitem__(self, key: Any) -> Any: + if isinstance(key, int): + if key < 0 or key >= len(self._keys): + raise IndexError(f"index {key} out of range for TupleDict with {len(self._keys)} items") + key = self._keys[key] + return dict.__getitem__(self, key) + + +__all__ = ["TupleDict"] diff --git a/python/cudnn/_tensor_desc.py b/python/cudnn/_tensor_desc.py index 5c519397e..037b4a257 100644 --- a/python/cudnn/_tensor_desc.py +++ b/python/cudnn/_tensor_desc.py @@ -111,6 +111,26 @@ def compact_like( init_value=init_value, ) + def is_compact(self, stride_order: tuple[int, ...] | None = None) -> bool: + """Return whether the descriptor is contiguous in the given order. + + Size-one dimensions may carry any non-negative stride because they do + not change the addressed storage. When no order is supplied, the + descriptor's canonical ``stride_order`` is used. + """ + + order = self.stride_order if stride_order is None else self._integer_tuple(stride_order, "stride_order") + if len(order) != self.ndim or tuple(sorted(order)) != tuple(range(self.ndim)): + raise ValueError(f"Stride order must be a permutation of [0, {self.ndim - 1}], got {order}") + + expected_stride = 1 + for dimension in order: + size = self.shape[dimension] + if size != 1 and self.stride[dimension] != expected_stride: + return False + expected_stride *= max(size, 1) + return True + def make_compact_tensor_desc( *, diff --git a/python/cudnn/api_base.py b/python/cudnn/api_base.py index 97f45da6f..9cddcb85c 100644 --- a/python/cudnn/api_base.py +++ b/python/cudnn/api_base.py @@ -23,6 +23,7 @@ from cudnn import data_type from cudnn._experimental_warnings import warn_experimental_api_once from cudnn._tensor_desc import TensorDesc as _TensorDesc +from cudnn._result import TupleDict as TupleDict from cudnn.datatypes import ( _convert_to_cutlass_data_type, _cudnn_to_torch_data_type, @@ -1164,35 +1165,3 @@ def _make_fake_cute_compact_tensor( stride_order=stride_order, assumed_align=assumed_align, ) - - -class TupleDict(dict): - """A dictionary that supports tuple unpacking. - - This class extends dict to allow unpacking like a tuple while still - providing dictionary-style key access. The unpacking order is determined - by the _keys attribute which preserves insertion order. - - Example: - >>> result = TupleDict(a=1, b=2, c=3) - >>> x, y, z = result # Unpacks as (1, 2, 3) - >>> result['a'] # Returns 1 - >>> result[0] # Returns 1 (integer indexing) - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - # Store keys in order for tuple unpacking - self._keys = list(self.keys()) - - def __iter__(self): - """Iterate over values in insertion order for tuple unpacking.""" - return (self[k] for k in self._keys) - - def __getitem__(self, key): - """Support both string keys and integer indices.""" - if isinstance(key, int): - if key < 0 or key >= len(self._keys): - raise IndexError(f"index {key} out of range for TupleDict with {len(self._keys)} items") - return super().__getitem__(self._keys[key]) - return super().__getitem__(key) diff --git a/python/cudnn/deepseek_sparse_attention/README.md b/python/cudnn/deepseek_sparse_attention/README.md index fc892210a..46b1a77ec 100644 --- a/python/cudnn/deepseek_sparse_attention/README.md +++ b/python/cudnn/deepseek_sparse_attention/README.md @@ -1,11 +1,16 @@ ## DSA module -- **Indexer Forward**: CuTe-DSL score kernel (Q @ Kᵗ, ReLU, head reduce, ratio causal mask). Supports the SM90 MQA path through `indexer_forward_wrapper` and SM100+ through `IndexerForward`; non-fused, so pair with **Indexer Top-K** for the top-K stage. +Install the JAX integration with +`pip install 'nvidia-cudnn-frontend[jax]'`. JAX classes and JIT-compiled +wrappers are available as direct `cudnn.jax` exports and through +`cudnn.jax.DSA`. + +- **Indexer Forward**: CuTe-DSL score kernel (Q @ Kᵗ, ReLU, head reduce, ratio causal mask). Supports SM90 MQA and SM100+ with fixed BSHD/SBHD or packed THD JAX inputs; non-fused, so pair with **Indexer Top-K** for the top-K stage. - **Indexer Top-K**: SM90+ CuTe-DSL radix top-K kernel with per-row ``seq_lens``. - **Sparse Attention Backward**: DSA backward (FlashMLA-shape, SM90/SM100). - **Sparse Indexer / Attention Score Recompute**: Sparse (top-K) recomputation of indexer and attention scores for training loss. - **Dense Indexer / Attention Score Recompute**: Dense (full-KV) analogues of the above. -- **Indexer Backward**: Three-stage pipeline (score-grad, three GEMMs, dtype cast) for sparse top-K score tensors. +- **Indexer Backward**: Three-stage pipeline (score-grad, three GEMMs, dtype cast) for sparse top-K score tensors. JAX additionally supports packed THD tensors when `topk_indices_global=True`; the PyTorch sparse wrapper remains fixed BSHD. - **Dense Indexer Backward**: Full-KV counterpart of Indexer Backward. ## Acknowledgements diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py index fbae12785..5d5c717b2 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py @@ -1,13 +1,20 @@ -from .api import ( - DenseIndexerBackward, - IndexerBackward, - dense_indexer_backward_wrapper, - indexer_backward_wrapper, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy Torch API and framework-neutral indexer-backward exports.""" -__all__ = [ - "DenseIndexerBackward", - "IndexerBackward", - "dense_indexer_backward_wrapper", - "indexer_backward_wrapper", -] +from ..._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ("DenseIndexerBackwardOp", "IndexerBackwardOp"), + "api": ( + "DenseIndexerBackward", + "IndexerBackward", + "dense_indexer_backward_wrapper", + "indexer_backward_wrapper", + ), + }, + submodules=("api", "op"), +) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py index 023c0e20a..0bda60ad8 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py @@ -17,6 +17,8 @@ import torch import cuda.bindings.driver as cuda +from cudnn import data_type +from cudnn._tensor_desc import make_compact_tensor_desc from cudnn.api_base import APIBase, TupleDict from cudnn.deepseek_sparse_attention.utils.runtime import ( torch_stream_context as _torch_stream_context, @@ -27,6 +29,7 @@ from .dense_indexer_backward_sm90 import dense_indexer_backward_sm90 from .indexer_backward_sm100 import indexer_backward_sm100 from .indexer_backward_sm90 import indexer_backward_sm90 +from .op import DenseIndexerBackwardOp, IndexerBackwardOp def _as_grad_loss_tensor(grad_loss: Union[float, torch.Tensor], device: torch.device) -> torch.Tensor: @@ -161,18 +164,20 @@ def __init__( self.idx_score_desc = self._make_tensor_desc(sample_index_score, name="sample_index_score") self.topk_desc = self._make_tensor_desc(sample_topk_indices, name="sample_topk_indices") - b, s_q, h, d = sample_index_q.shape - s_k = sample_index_k.shape[1] - topk = sample_topk_indices.shape[-1] - self.batch = b - self.seqlen = s_q - self.seqlen_k = s_k - self.heads = h - self.head_dim = d - self.topk = topk - self.sm_scale = float(sm_scale) - self.block_I = int(block_I) - self.topk_indices_global = bool(topk_indices_global) + self._op = IndexerBackwardOp( + index_q=self.iq_desc, + weights=self.w_desc, + index_k=self.ik_desc, + d_index_q=self.diq_desc, + d_weights=self.dw_desc, + d_index_k=self.dik_desc, + attn_score=self.attn_desc, + index_score=self.idx_score_desc, + topk_indices=self.topk_desc, + sm_scale=sm_scale, + block_i=block_I, + topk_indices_global=topk_indices_global, + ) def check_support(self) -> bool: major, _ = torch.cuda.get_device_capability() @@ -180,12 +185,7 @@ def check_support(self) -> bool: major != 9 and major < 10, f"IndexerBackward requires SM90 or SM100+, found SM{major}", ) - self._check_dtype(self.iq_desc, torch.bfloat16, name="index_q") - self._check_dtype(self.w_desc, torch.bfloat16, name="weights") - self._check_dtype(self.ik_desc, torch.bfloat16, name="index_k") - self._check_dtype(self.attn_desc, torch.float32, name="attn_score") - self._check_dtype(self.idx_score_desc, torch.float32, name="index_score") - self._check_dtype(self.topk_desc, torch.int32, name="topk_indices") + self._op.check_support() self._is_supported = True return True @@ -201,15 +201,15 @@ def compile(self) -> None: major, _ = torch.cuda.get_device_capability() kernel_factory = indexer_backward_sm90 if major == 9 else indexer_backward_sm100 self._compiled_kernel = kernel_factory( - self.batch, - self.seqlen, - self.seqlen_k, - self.heads, - self.head_dim, - self.topk, - sm_scale=self.sm_scale, - block_I=self.block_I, - topk_indices_global=self.topk_indices_global, + self._op.batch, + self._op.seqlen_q, + self._op.seqlen_k, + self._op.heads, + self._op.head_dim, + self._op.topk, + sm_scale=self._op.sm_scale, + block_I=self._op.block_i, + topk_indices_global=self._op.topk_indices_global, ) def execute( @@ -231,7 +231,7 @@ def execute( # ``grad_scale`` is forwarded as a runtime ``Float32`` arg; the # compiled kernel is reused when loss_coeff changes for the same # cached tensor shape. - grad_scale = float(loss_coeff) / (int(self.batch) * int(self.seqlen)) + grad_scale = float(loss_coeff) / (int(self._op.batch) * int(self._op.seqlen_q)) grad_loss_tensor = _as_grad_loss_tensor(grad_loss, index_q.device) self._compiled_kernel( index_q, @@ -290,31 +290,35 @@ def __init__( self.idx_score_desc = self._make_tensor_desc(sample_index_score, name="sample_index_score") self.idx_lse_desc = self._make_tensor_desc(sample_index_lse, name="sample_index_lse") - self.sm_scale = float(sm_scale) - self.block_I = int(block_I) - self.ratio = int(ratio) - self.is_thd = bool(is_thd) - self.has_q_causal_offsets = bool(has_q_causal_offsets) - - if self.is_thd: - total_q, heads, head_dim = sample_index_q.shape - total_k = sample_index_k.shape[0] - if batch is None or max_seqlen_q is None or max_seqlen_k is None: - raise ValueError("THD dense indexer backward requires batch and max_seqlen_q/k") - self.batch = int(batch) - self.normalization_tokens = int(total_q) - self.total_k = int(total_k) - self.max_seqlen_q = int(max_seqlen_q) - self.max_seqlen_k = int(max_seqlen_k) - else: - b, seqlen_q, heads, head_dim = sample_index_q.shape - self.batch = int(b) - self.normalization_tokens = int(b * seqlen_q) - self.total_k = int(b * sample_index_k.shape[1]) - self.max_seqlen_q = int(seqlen_q) - self.max_seqlen_k = int(sample_index_k.shape[1]) - self.heads = int(heads) - self.head_dim = int(head_dim) + is_thd = bool(is_thd) + if is_thd and batch is None: + raise ValueError("THD dense indexer backward requires batch") + inferred_batch = int(batch) if is_thd else int(sample_index_q.shape[0]) + cu_seqlens_q_desc = make_compact_tensor_desc(dtype=data_type.INT32, shape=(inferred_batch + 1,), name="sample_cu_seqlens_q") if is_thd else None + cu_seqlens_k_desc = make_compact_tensor_desc(dtype=data_type.INT32, shape=(inferred_batch + 1,), name="sample_cu_seqlens_k") if is_thd else None + q_causal_offsets_desc = ( + make_compact_tensor_desc(dtype=data_type.INT32, shape=(inferred_batch,), name="sample_q_causal_offsets") if has_q_causal_offsets else None + ) + self._op = DenseIndexerBackwardOp( + index_q=self.iq_desc, + weights=self.w_desc, + index_k=self.ik_desc, + d_index_q=self.diq_desc, + d_weights=self.dw_desc, + d_index_k=self.dik_desc, + attn_score=self.attn_desc, + attn_l1norm=self.attn_denom_desc, + index_score=self.idx_score_desc, + index_lse=self.idx_lse_desc, + cu_seqlens_q=cu_seqlens_q_desc, + cu_seqlens_k=cu_seqlens_k_desc, + q_causal_offsets=q_causal_offsets_desc, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + sm_scale=sm_scale, + block_i=block_I, + ratio=ratio, + ) self._uses_current_stream_pipeline = False def check_support(self) -> bool: @@ -323,19 +327,7 @@ def check_support(self) -> bool: major != 9 and major < 10, f"DenseIndexerBackward requires SM90 or SM100+, found SM{major}", ) - self._check_dtype(self.iq_desc, torch.bfloat16, name="index_q") - self._check_dtype(self.w_desc, torch.bfloat16, name="weights") - self._check_dtype(self.ik_desc, torch.bfloat16, name="index_k") - self._check_dtype(self.diq_desc, torch.bfloat16, name="d_index_q") - self._check_dtype(self.dw_desc, torch.bfloat16, name="d_weights") - self._check_dtype(self.dik_desc, [torch.bfloat16, torch.float32], name="d_index_k") - self._check_dtype(self.attn_desc, torch.float32, name="attn_score") - self._check_dtype(self.attn_denom_desc, torch.float32, name="attn_l1norm") - self._check_dtype(self.idx_score_desc, torch.float32, name="index_score") - self._check_dtype(self.idx_lse_desc, torch.float32, name="index_lse") - self._value_error_if(self.block_I <= 0, f"block_I must be positive, got {self.block_I}") - self._value_error_if(self.ratio < 1, f"ratio must be >= 1, got {self.ratio}") - self._value_error_if(self.heads < 64, f"DenseIndexerBackward requires heads >= 64, got {self.heads}") + self._op.check_support() self._is_supported = True return True @@ -347,16 +339,16 @@ def compile(self) -> None: kernel_factory = dense_indexer_backward_sm90 if major == 9 else dense_indexer_backward_sm100 self._uses_current_stream_pipeline = major == 9 self._compiled_kernel = kernel_factory( - self.batch, - self.max_seqlen_q, - self.max_seqlen_k, - self.heads, - self.head_dim, - sm_scale=self.sm_scale, - block_I=self.block_I, - ratio=self.ratio, - is_varlen=self.is_thd, - has_q_causal_offsets=self.has_q_causal_offsets, + self._op.batch, + self._op.max_seqlen_q, + self._op.max_seqlen_k, + self._op.heads, + self._op.head_dim, + sm_scale=self._op.sm_scale, + block_I=self._op.block_i, + ratio=self._op.ratio, + is_varlen=self._op.is_thd, + has_q_causal_offsets=self._op.q_causal_offsets is not None, ) def execute( @@ -381,7 +373,7 @@ def execute( backend_stream = None if self._uses_current_stream_pipeline else current_stream with _torch_stream_context(backend_stream): grad_loss_value = float(_as_grad_loss_tensor(grad_loss, index_q.device).item()) - grad_scale = float(loss_coeff) * grad_loss_value / max(int(self.normalization_tokens), 1) + grad_scale = float(loss_coeff) * grad_loss_value / max(int(self._op.normalization_tokens), 1) # Dense backward's dK path uses atomic/bulk reductions into fp32. d_index_k_target = d_index_k @@ -500,7 +492,7 @@ def indexer_backward_wrapper( block_I=block_I, topk_indices_global=topk_indices_global, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_IndexerBackwardObjects[key] = obj @@ -647,7 +639,7 @@ def dense_indexer_backward_wrapper( max_seqlen_k=max_k, has_q_causal_offsets=q_causal_offsets is not None, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_DenseIndexerBackwardObjects[key] = obj diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py index 7e8112211..c55059124 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py @@ -53,7 +53,6 @@ import math from functools import partial -import torch import cuda.bindings.driver as cuda import cutlass @@ -73,12 +72,7 @@ import cutlass.utils.blackwell_helpers as sm100_utils_basic -from cudnn.deepseek_sparse_attention.utils.compiler import compile_options - from cudnn.deepseek_sparse_attention.utils.copy import cpasync_reduce_bulk_add_f32 -from cudnn.deepseek_sparse_attention.utils.runtime import ( - resolve_stream as _resolve_stream, -) from cudnn.deepseek_sparse_attention.utils.seqlen import seqlen_info as _seqlen_info mul_packed_f32x2 = partial(cute.arch.mul_packed_f32x2, rnd="rn") @@ -151,7 +145,12 @@ def __call__( max_seqlen_q: Int32, max_seqlen_k: Int32, stream, + mGradSignal=None, + mGradLoss=None, ): + if const_expr(mGradSignal is None): + mGradSignal = mPredictRaw + # BSHD: mCuSeqlensQ/K = None; tensors carry batch dim. # PredictRaw/TargetRaw: (B, S_q, S_k) LSE/Denom: (B, S_q) # THD: mCuSeqlensQ/K provided; tensors are packed. @@ -167,6 +166,10 @@ def __call__( mPredictRaw.iterator, cute.select(mPredictRaw.layout, mode=[1, 2, 0]), ) + mGradSignal = cute.make_tensor( + mGradSignal.iterator, + cute.select(mGradSignal.layout, mode=[1, 2, 0]), + ) mTargetRaw = cute.make_tensor( mTargetRaw.iterator, cute.select(mTargetRaw.layout, mode=[1, 2, 0]), @@ -186,12 +189,14 @@ def __call__( self.kernel_score_grad( mPredictRaw, + mGradSignal, mTargetRaw, mLSE, mDenomTarget, mCuSeqlensQ, mCuSeqlensK, mQCausalOffsets, + mGradLoss, grad_scale, seqlen_k_pad, max_seqlen_q, @@ -206,12 +211,14 @@ def __call__( def kernel_score_grad( self, mPredictRaw, + mGradSignal, mTargetRaw, mLSE, mDenomTarget, mCuSeqlensQ, mCuSeqlensK, mQCausalOffsets, + mGradLoss, grad_scale: Float32, seqlen_k_pad: Int32, seqlen_q_static: Int32, @@ -252,11 +259,13 @@ class SharedStorage: # THD : (T_q, max_K) -> domain-offset by q_offset along T_q; LSE/Denom (T_q,) similarly. if const_expr(is_varlen): mPredict_b = cute.domain_offset((q_offset, Int32(0)), mPredictRaw) + mGradSignal_b = cute.domain_offset((q_offset, Int32(0)), mGradSignal) mTarget_b = cute.domain_offset((q_offset, Int32(0)), mTargetRaw) mLSE_b = cute.domain_offset((q_offset,), mLSE) mDenom_b = cute.domain_offset((q_offset,), mDenomTarget) else: mPredict_b = mPredictRaw[None, None, batch_idx] + mGradSignal_b = mGradSignal[None, None, batch_idx] mTarget_b = mTargetRaw[None, None, batch_idx] mLSE_b = mLSE[None, batch_idx] mDenom_b = mDenomTarget[None, batch_idx] @@ -264,6 +273,7 @@ class SharedStorage: # Scalar loads (shared across all threads in CTA) lse_val = mLSE_b[seq_local] denom_val = mDenom_b[seq_local] + effective_grad_scale = grad_scale if const_expr(mGradLoss is None) else grad_scale * mGradLoss[0] LOG2E = Float32(1.4426950408889634) @@ -284,7 +294,7 @@ class SharedStorage: target = tr / (denom_val + Float32(DENOM_EPS)) target_eff = target if target >= Float32(CLIP_PROB_MIN) else Float32(CLIP_PROB_MIN) log_clip_mask = Float32(1.0) if score_minus_lse >= Float32(CLIP_LOG_MIN) else Float32(0.0) - g = -target_eff * log_clip_mask * grad_scale + g = -target_eff * log_clip_mask * effective_grad_scale local_sum = local_sum + g pos = pos + Int32(128) @@ -324,14 +334,14 @@ class SharedStorage: target = tr / (denom_val + Float32(DENOM_EPS)) target_eff = target if target >= Float32(CLIP_PROB_MIN) else Float32(CLIP_PROB_MIN) log_clip_mask = Float32(1.0) if score_minus_lse >= Float32(CLIP_LOG_MIN) else Float32(0.0) - g = -target_eff * log_clip_mask * grad_scale + g = -target_eff * log_clip_mask * effective_grad_scale grad_signal = g - predict * sum_grad - mPredict_b[seq_local, pos] = grad_signal + mGradSignal_b[seq_local, pos] = grad_signal else: # masked / padding: grad_signal must be 0 — downstream # GEMM multiplies by this directly and must not contaminate dQ/dK/dW. - mPredict_b[seq_local, pos] = Float32(0.0) + mGradSignal_b[seq_local, pos] = Float32(0.0) pos = pos + Int32(128) @@ -1981,6 +1991,10 @@ def dense_indexer_backward_sm100( def _build_cute_dsl_kernel(batch, max_seqlen_q, max_seqlen_k, heads, dim, sm_scale, block_I, ratio, is_varlen, has_q_causal_offsets): + import torch + + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor if torch.cuda.get_device_capability()[0] < 10: @@ -2024,7 +2038,7 @@ def _ensure_compiled_score_grad( the runtime ``grad_scale`` differs from the value seen at compile. """ if compile_key not in _score_grad_compile_cache: - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) cuq_arg = to_cute_tensor(CuSeqlensQ) if CuSeqlensQ is not None else None cuk_arg = to_cute_tensor(CuSeqlensK) if CuSeqlensK is not None else None q_offsets_arg = to_cute_tensor(QCausalOffsets) if QCausalOffsets is not None else None @@ -2063,7 +2077,7 @@ def _ensure_compiled_gemm( int32 cu_seqlens tensors of shape (B+1,). """ if compile_key not in _gemm_compile_cache: - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) cuq_arg = to_cute_tensor(CuSeqlensQ) if CuSeqlensQ is not None else None cuk_arg = to_cute_tensor(CuSeqlensK) if CuSeqlensK is not None else None q_offsets_arg = to_cute_tensor(QCausalOffsets) if QCausalOffsets is not None else None @@ -2110,7 +2124,7 @@ def _run_gemm_only( assert QCausalOffsets is not None, "offset-compiled kernel requires q_causal_offsets at runtime" else: assert QCausalOffsets is None, "non-offset compiled kernel must not receive q_causal_offsets" - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) _ensure_compiled_gemm( IndexQ, Weights, @@ -2177,7 +2191,7 @@ def _run( assert QCausalOffsets is not None, "offset-compiled kernel requires q_causal_offsets at runtime" else: assert QCausalOffsets is None, "non-offset compiled kernel must not receive q_causal_offsets" - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) # Kernel 1 (CuTe DSL): in-place overwrites IdxScoreRaw with grad_signal. # Grid = (max_seqlen_q, B, 1); CTAs past per-batch seqlen exit early. diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm90.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm90.py index 5e7a18696..506bd8c76 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm90.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm90.py @@ -2,20 +2,12 @@ from __future__ import annotations -import torch import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute from cutlass import Float32, Int32, const_expr -from cudnn.deepseek_sparse_attention.utils.compiler import compile_options -from cudnn.deepseek_sparse_attention.utils.runtime import ( - resolve_stream as _resolve_stream, - torch_stream_context as _torch_stream_context, -) -from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor - from .indexer_backward_sm90 import ( CLIP_LOG_MIN, CLIP_PROB_MIN, @@ -71,7 +63,12 @@ def __call__( max_seqlen_q: Int32, max_seqlen_k: Int32, stream, + mGradSignal=None, + mGradLoss=None, ): + if const_expr(mGradSignal is None): + mGradSignal = mIdxScoreRaw + # BSHD: # IdxScoreRaw/AttnScoreRaw: (B, S_q, S_k), IdxLSE/AttnL1Norm: (B, S_q) # THD: @@ -85,6 +82,10 @@ def __call__( mIdxScoreRaw.iterator, cute.select(mIdxScoreRaw.layout, mode=[1, 2, 0]), ) + mGradSignal = cute.make_tensor( + mGradSignal.iterator, + cute.select(mGradSignal.layout, mode=[1, 2, 0]), + ) mAttnScoreRaw = cute.make_tensor( mAttnScoreRaw.iterator, cute.select(mAttnScoreRaw.layout, mode=[1, 2, 0]), @@ -103,12 +104,14 @@ def __call__( self.kernel_score_grad( mIdxScoreRaw, + mGradSignal, mAttnScoreRaw, mIdxLSE, mAttnL1Norm, mCuSeqlensQ, mCuSeqlensK, mQCausalOffsets, + mGradLoss, grad_scale, seqlen_k_pad, max_seqlen_q, @@ -123,12 +126,14 @@ def __call__( def kernel_score_grad( self, mIdxScoreRaw, + mGradSignal, mAttnScoreRaw, mIdxLSE, mAttnL1Norm, mCuSeqlensQ, mCuSeqlensK, mQCausalOffsets, + mGradLoss, grad_scale: Float32, seqlen_k_pad: Int32, seqlen_q_static: Int32, @@ -164,17 +169,20 @@ class SharedStorage: if const_expr(is_varlen): mIdxScore_b = cute.domain_offset((q_offset, Int32(0)), mIdxScoreRaw) + mGradSignal_b = cute.domain_offset((q_offset, Int32(0)), mGradSignal) mAttnScore_b = cute.domain_offset((q_offset, Int32(0)), mAttnScoreRaw) mIdxLSE_b = cute.domain_offset((q_offset,), mIdxLSE) mAttnL1_b = cute.domain_offset((q_offset,), mAttnL1Norm) else: mIdxScore_b = mIdxScoreRaw[None, None, batch_idx] + mGradSignal_b = mGradSignal[None, None, batch_idx] mAttnScore_b = mAttnScoreRaw[None, None, batch_idx] mIdxLSE_b = mIdxLSE[None, batch_idx] mAttnL1_b = mAttnL1Norm[None, batch_idx] idx_lse_val = mIdxLSE_b[seq_local] attn_l1_val = mAttnL1_b[seq_local] + effective_grad_scale = grad_scale if const_expr(mGradLoss is None) else grad_scale * mGradLoss[0] LOG2E = Float32(1.4426950408889634) ratio = Int32(self.ratio) @@ -195,7 +203,7 @@ class SharedStorage: target = attn_raw / (attn_l1_val + Float32(EPS)) target_eff = target if target >= Float32(CLIP_PROB_MIN) else Float32(CLIP_PROB_MIN) log_clip_mask = Float32(1.0) if score_minus_lse >= Float32(CLIP_LOG_MIN) else Float32(0.0) - local_sum = local_sum + (-target_eff * log_clip_mask * grad_scale) + local_sum = local_sum + (-target_eff * log_clip_mask * effective_grad_scale) pos = pos + Int32(128) warp_sum = cute.arch.warp_reduction_sum(local_sum) @@ -221,10 +229,10 @@ class SharedStorage: target = attn_raw / (attn_l1_val + Float32(EPS)) target_eff = target if target >= Float32(CLIP_PROB_MIN) else Float32(CLIP_PROB_MIN) log_clip_mask = Float32(1.0) if score_minus_lse >= Float32(CLIP_LOG_MIN) else Float32(0.0) - g = -target_eff * log_clip_mask * grad_scale - mIdxScore_b[seq_local, pos] = g - predict * sum_grad + g = -target_eff * log_clip_mask * effective_grad_scale + mGradSignal_b[seq_local, pos] = g - predict * sum_grad else: - mIdxScore_b[seq_local, pos] = Float32(0.0) + mGradSignal_b[seq_local, pos] = Float32(0.0) pos = pos + Int32(128) @@ -270,6 +278,12 @@ def _build_cute_dsl_dense_kernel( is_varlen, has_q_causal_offsets, ): + import torch + + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream, torch_stream_context + from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor + cap = torch.cuda.get_device_capability()[0] if cap < 9: raise RuntimeError(f"Requires SM90+ (got SM{cap}0)") @@ -295,7 +309,7 @@ def _build_cute_dsl_dense_kernel( def _get_dummy_topk(device, current_stream=None): if dummy_topk_holder[0] is None or dummy_topk_holder[0].device != device: - with _torch_stream_context(current_stream): + with torch_stream_context(current_stream): dummy_topk_holder[0] = torch.zeros(batch, seqlen, seqlen_k, device=device, dtype=torch.int32) return dummy_topk_holder[0] @@ -312,7 +326,7 @@ def _ensure_compiled( QCausalOffsets, current_stream=None, ): - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) if gemm_key not in _dense_compile_cache: dummy_topk = _get_dummy_topk(IndexQ.device, current_stream=current_stream) cuq_arg = to_cute_tensor(CuSeqlensQ) if CuSeqlensQ is not None else None @@ -356,7 +370,7 @@ def _ensure_compiled_score_grad( current_stream=None, ): if score_grad_key not in _dense_score_grad_compile_cache: - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) cuq_arg = to_cute_tensor(CuSeqlensQ) if CuSeqlensQ is not None else None cuk_arg = to_cute_tensor(CuSeqlensK) if CuSeqlensK is not None else None q_offsets_arg = to_cute_tensor(QCausalOffsets) if QCausalOffsets is not None else None @@ -395,7 +409,7 @@ def _run_score_grad_only( assert QCausalOffsets is not None, "offset-compiled score-grad kernel requires q_causal_offsets at runtime" else: assert QCausalOffsets is None, "non-offset compiled score-grad kernel must not receive q_causal_offsets" - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) _ensure_compiled_score_grad( IdxScoreRaw, IdxLSE, @@ -444,7 +458,7 @@ def _run_gemm_only( else: assert QCausalOffsets is None, "non-offset compiled kernel must not receive q_causal_offsets" dummy_topk = _get_dummy_topk(IndexQ.device, current_stream=current_stream) - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) _ensure_compiled( IndexQ, @@ -526,7 +540,7 @@ def _run( current_stream=current_stream, ) else: - with _torch_stream_context(current_stream): + with torch_stream_context(current_stream): dIndexK_f32 = torch.zeros_like(dIndexK, dtype=torch.float32) _run_gemm_only( IndexQ, @@ -541,7 +555,7 @@ def _run( QCausalOffsets, current_stream=current_stream, ) - with _torch_stream_context(current_stream): + with torch_stream_context(current_stream): dIndexK.copy_(dIndexK_f32) _run.score_grad = _run_score_grad_only diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py index 740c3578d..b0b96abc9 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py @@ -51,7 +51,8 @@ import math from functools import partial -import torch +from typing import Optional + import cuda.bindings.driver as cuda import cutlass @@ -71,12 +72,6 @@ import cutlass.utils.blackwell_helpers as sm100_utils_basic -from cudnn.deepseek_sparse_attention.utils.compiler import compile_options -from cudnn.deepseek_sparse_attention.utils.runtime import ( - resolve_stream as _resolve_stream, - torch_stream_context as _torch_stream_context, -) - mul_packed_f32x2 = partial(cute.arch.mul_packed_f32x2, rnd="rn") fma_packed_f32x2 = partial(cute.arch.fma_packed_f32x2, rnd="rn") @@ -1453,7 +1448,7 @@ def indexer_backward_sm100( class ScoreGradSm100: - """CuTe DSL kernel for in-place score_grad precompute.""" + """CuTe DSL sparse score-gradient kernel with optional outputs.""" THREADS_PER_CTA = 128 WARP_SIZE = 32 @@ -1470,14 +1465,24 @@ def __call__( mGradLoss: cute.Tensor, grad_scale: Float32 | float, stream: cuda.CUstream, + mGradSignal: Optional[cute.Tensor] = None, + mSumGrad: Optional[cute.Tensor] = None, ): + write_sum_grad = mSumGrad is not None or mGradSignal is None + if cutlass.const_expr(mGradSignal is None): + mGradSignal = mAttnScore + if cutlass.const_expr(mSumGrad is None): + mSumGrad = mIndexScore + # (b, s, t) -> (s, t, b): topk dim contiguous for per-CTA strided loops. mAttnScore = cute.make_tensor(mAttnScore.iterator, cute.select(mAttnScore.layout, mode=[1, 2, 0])) mIndexScore = cute.make_tensor(mIndexScore.iterator, cute.select(mIndexScore.layout, mode=[1, 2, 0])) + mGradSignal = cute.make_tensor(mGradSignal.iterator, cute.select(mGradSignal.layout, mode=[1, 2, 0])) + mSumGrad = cute.make_tensor(mSumGrad.iterator, cute.select(mSumGrad.layout, mode=[1, 2, 0])) seqlen = cute.size(mAttnScore.shape[0]) batch_size = cute.size(mAttnScore.shape[2]) if cute.rank(mAttnScore.shape) > 2 else 1 - self.kernel_score_grad(mAttnScore, mIndexScore, mGradLoss, grad_scale).launch( + self.kernel_score_grad(mAttnScore, mIndexScore, mGradLoss, mGradSignal, mSumGrad, grad_scale, write_sum_grad).launch( grid=(seqlen, batch_size, 1), block=[self.THREADS_PER_CTA, 1, 1], cluster=[1, 1, 1], @@ -1486,7 +1491,9 @@ def __call__( ) @cute.kernel - def kernel_score_grad(self, mAttnScore, mIndexScore, mGradLoss, grad_scale: Float32 | float): + def kernel_score_grad( + self, mAttnScore, mIndexScore, mGradLoss, mGradSignal, mSumGrad, grad_scale: Float32 | float, write_sum_grad: cutlass.Constexpr[bool] + ): tidx = cute.arch.thread_idx()[0] seq_idx = cute.arch.block_idx()[0] batch_idx = cute.arch.block_idx()[1] @@ -1528,11 +1535,14 @@ class SharedStorage: target_eff = cute.arch.fmax(target, Float32(CLIP_PROB_MIN)) log_clip_mask = Float32(1.0) if predict >= Float32(CLIP_PROB_MIN) else Float32(0.0) g_i = -target_eff * log_clip_mask * grad_scale_f32 - mAttnScore[seq_idx, pos, batch_idx] = g_i - predict * sum_grad - mIndexScore[seq_idx, pos, batch_idx] = sum_grad + mGradSignal[seq_idx, pos, batch_idx] = g_i - predict * sum_grad + if cutlass.const_expr(write_sum_grad): + mSumGrad[seq_idx, pos, batch_idx] = sum_grad def _score_grad_inplace_cute(AttnScore, IndexScore, GradLoss, grad_scale, current_stream=None): + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor # Kernel reads ``mGradLoss[0]`` so it must be at least 1-D. ``to_cute_tensor`` @@ -1544,7 +1554,7 @@ def _score_grad_inplace_cute(AttnScore, IndexScore, GradLoss, grad_scale, curren _, _, topk = AttnScore.shape compile_key = (topk,) - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) if compile_key not in _score_grad_cute_cache: kernel_obj = ScoreGradSm100(topk=topk) _score_grad_cute_cache[compile_key] = cute.compile( @@ -1582,6 +1592,8 @@ def _score_grad_inplace(AttnScore, IndexScore, GradLoss, grad_scale, block_I=128 # target = clip(log_target) # dL/dlog_predict = -exp(target) * I(log_predict >= -100) # dL/dlogits = g - predict * sum(g) + import torch + can_use_cute = ( AttnScore.is_cuda and IndexScore.is_cuda @@ -1598,6 +1610,10 @@ def _score_grad_inplace(AttnScore, IndexScore, GradLoss, grad_scale, block_I=128 def _build_cute_dsl_kernel(heads, dim, topk, sm_scale, block_I, topk_indices_global: bool = True): + import torch + + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream, torch_stream_context from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor if torch.cuda.get_device_capability()[0] < 10: @@ -1617,7 +1633,7 @@ def _build_cute_dsl_kernel(heads, dim, topk, sm_scale, block_I, topk_indices_glo def _ensure_compiled(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, AttnScore, TopkIndices, current_stream=None): """Lazy-compile the GEMM kernel (kernel 2) on first execute (needs real tensors).""" if compile_key not in _compile_cache: - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) cute_args = [to_cute_tensor(t) for t in [IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, AttnScore, TopkIndices]] _compile_cache[compile_key] = cute.compile( kernel_obj, @@ -1629,7 +1645,7 @@ def _ensure_compiled(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, At def _run_gemm_only(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, GradSignal, TopkIndices, current_stream=None): """Run only kernel 2 (GEMM). Caller must have run kernel 1 and zeroed dIndexK_f32.""" - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) _ensure_compiled(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, GradSignal, TopkIndices, current_stream=current_stream) with torch.cuda.nvtx.range("indexer_backward_dsl_gemm"): _compile_cache[compile_key]( @@ -1660,10 +1676,10 @@ def _run(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK, AttnScore, IndexSc _run_gemm_only(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK, AttnScore, TopkIndices, current_stream=current_stream) else: # Need a separate f32 buffer for atomicAdd, then cast back to output dtype. - with _torch_stream_context(current_stream): + with torch_stream_context(current_stream): dIndexK_f32 = torch.zeros_like(dIndexK, dtype=torch.float32) _run_gemm_only(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, AttnScore, TopkIndices, current_stream=current_stream) - with _torch_stream_context(current_stream): + with torch_stream_context(current_stream): dIndexK.copy_(dIndexK_f32) _run.score_grad = partial(_score_grad_inplace, block_I=block_I) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm90.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm90.py index 558e8d614..23f402c2e 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm90.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm90.py @@ -44,7 +44,6 @@ from typing import Optional import cuda.bindings.driver as cuda -import torch import cutlass import cutlass.cute as cute @@ -55,13 +54,7 @@ from cutlass.utils import LayoutEnum from cudnn.deepseek_sparse_attention.utils import copy as copy_ops -from cudnn.deepseek_sparse_attention.utils.compiler import compile_options -from cudnn.deepseek_sparse_attention.utils.runtime import ( - resolve_stream as _resolve_stream, - torch_stream_context as _torch_stream_context, -) from cudnn.deepseek_sparse_attention.utils.seqlen import seqlen_info as _seqlen_info -from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor from cudnn.deepseek_sparse_attention.utils.sm90.mma import ( gemm, gemm_zero_init, @@ -1385,7 +1378,7 @@ def indexer_backward_sm90( class ScoreGradSm90: - """CuTe DSL kernel for in-place score_grad precompute (SM90).""" + """CuTe DSL sparse score-gradient kernel with an optional output.""" THREADS_PER_CTA = 128 @@ -1401,14 +1394,19 @@ def __call__( mGradLoss: cute.Tensor, grad_scale: Float32 | float, stream: cuda.CUstream, + mGradSignal: Optional[cute.Tensor] = None, ): + if cutlass.const_expr(mGradSignal is None): + mGradSignal = mAttnScore + # (b, s, t) -> (s, t, b): contiguous topk traversal per CTA. mAttnScore = cute.make_tensor(mAttnScore.iterator, cute.select(mAttnScore.layout, mode=[1, 2, 0])) mIndexScore = cute.make_tensor(mIndexScore.iterator, cute.select(mIndexScore.layout, mode=[1, 2, 0])) + mGradSignal = cute.make_tensor(mGradSignal.iterator, cute.select(mGradSignal.layout, mode=[1, 2, 0])) seqlen = cute.size(mAttnScore.shape[0]) batch_size = cute.size(mAttnScore.shape[2]) if cute.rank(mAttnScore.shape) > 2 else 1 - self.kernel_score_grad(mAttnScore, mIndexScore, mGradLoss, grad_scale).launch( + self.kernel_score_grad(mAttnScore, mIndexScore, mGradLoss, mGradSignal, grad_scale).launch( grid=(seqlen, batch_size, 1), block=[self.THREADS_PER_CTA, 1, 1], cluster=[1, 1, 1], @@ -1417,7 +1415,7 @@ def __call__( ) @cute.kernel - def kernel_score_grad(self, mAttnScore, mIndexScore, mGradLoss, grad_scale: Float32 | float): + def kernel_score_grad(self, mAttnScore, mIndexScore, mGradLoss, mGradSignal, grad_scale: Float32 | float): tidx = cute.arch.thread_idx()[0] seq_idx = cute.arch.block_idx()[0] batch_idx = cute.arch.block_idx()[1] @@ -1473,7 +1471,7 @@ class SharedStorage: predict = Float32(mIndexScore[seq_idx, pos, batch_idx]) log_clip_mask = Float32(1.0) if predict >= Float32(CLIP_PROB_MIN) else Float32(0.0) g_i = -target_eff * log_clip_mask * grad_scale_f32 - mAttnScore[seq_idx, pos, batch_idx] = g_i - predict * sum_grad + mGradSignal[seq_idx, pos, batch_idx] = g_i - predict * sum_grad def _score_grad_inplace_cute( @@ -1484,6 +1482,10 @@ def _score_grad_inplace_cute( index_is_log: bool = False, current_stream=None, ): + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream + from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor + # Kernel reads ``mGradLoss[0]`` so it must be at least 1-D; defend direct # factory callers passing a 0-D scalar (the public wrapper reshapes already). if GradLoss.ndim == 0: @@ -1491,7 +1493,7 @@ def _score_grad_inplace_cute( _, _, topk = AttnScore.shape compile_key = (topk, bool(index_is_log)) - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) if compile_key not in _score_grad_cute_cache: kernel_obj = ScoreGradSm90( topk=topk, @@ -1538,6 +1540,8 @@ def _score_grad_inplace( # target = clip(log_target) # dL/dlog_predict = -exp(target) * I(log_predict >= -100) # dL/dlogits = g - predict * sum(g) + import torch + can_use_cute = ( AttnScore.is_cuda and IndexScore.is_cuda @@ -1571,6 +1575,12 @@ def _build_cute_dsl_kernel( score_input_is_log=True, topk_indices_global: bool = True, ): + import torch + + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream, torch_stream_context + from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor + cap = torch.cuda.get_device_capability()[0] if cap < 9: raise RuntimeError(f"Requires SM90+ (got SM{cap}0)") @@ -1603,7 +1613,7 @@ def _ensure_compiled( current_stream=None, ): """Lazy-compile the GEMM kernel (kernel 2) on first execute (needs real tensors).""" - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) if compile_key not in _compile_cache: cute_args = [ to_cute_tensor(t) @@ -1648,7 +1658,7 @@ def _run_gemm_only( current_stream=None, ): """Run only kernel 2 (GEMM). Caller must have run kernel 1 and zeroed dIndexK_f32.""" - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) _ensure_compiled( IndexQ, Weights, @@ -1719,7 +1729,7 @@ def _run( ) else: # Need separate f32 buffer for atomicAdd, then convert back - with _torch_stream_context(current_stream): + with torch_stream_context(current_stream): dIndexK_f32 = torch.zeros_like(dIndexK, dtype=torch.float32) _run_gemm_only( IndexQ, @@ -1732,7 +1742,7 @@ def _run( TopkIndices, current_stream=current_stream, ) - with _torch_stream_context(current_stream): + with torch_stream_context(current_stream): dIndexK.copy_(dIndexK_f32) _run.score_grad = partial(_score_grad_inplace, index_is_log=score_input_is_log) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py new file mode 100644 index 000000000..cbf1ca15d --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py @@ -0,0 +1,685 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX APIs for sparse and dense DeepSeek indexer backward.""" + +from __future__ import annotations + +import math +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..utils.compiler import compile_options_for_target +from .op import DEFAULT_BLOCK_I, DenseIndexerBackwardOp, IndexerBackwardOp + +_SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) +_SUPPORTED_COMPUTE_CAPABILITY_FAMILIES = (90, 100) + + +def _grad_loss_operand(value: Any) -> Any: + value = jnp.asarray(value, dtype=jnp.float32) + if value.shape == (): + return jnp.reshape(value, (1,)) + if value.shape != (1,): + raise ValueError(f"grad_loss must be scalar or shape (1,), got {value.shape}") + return value + + +def _finite_float(value: float, name: str) -> float: + try: + result = float(value) + except (TypeError, ValueError) as error: + raise TypeError(f"{name} must be a real scalar, got {value!r}") from error + if not math.isfinite(result): + raise ValueError(f"{name} must be finite, got {result}") + return result + + +class _IndexerJaxBase(JaxApiBase): + target_compute_capability: int | None + compute_capability: int | None + + def _resolve_target(self, operation_name: str) -> None: + self.compute_capability = self._resolve_compute_capability( + self.target_compute_capability, + _SUPPORTED_COMPUTE_CAPABILITIES, + operation_name, + ) + + @property + def _architecture_family(self) -> int: + if self.compute_capability is None: + raise RuntimeError("check_support() must resolve the compute capability before lowering") + family = self._compute_capability_family(self.compute_capability, _SUPPORTED_COMPUTE_CAPABILITY_FAMILIES) + if family is None: + raise RuntimeError(f"No indexer-backward kernel for SM{self.compute_capability}") + return family + + @staticmethod + def _output_desc( + sample: Any | None, + *, + source: JaxTensorDesc, + cudnn_dtype: data_type, + shape: tuple[int, ...], + name: str, + ) -> JaxTensorDesc: + if sample is not None: + return JaxApiBase._to_tensor_desc(sample, name) + return source.compact_like(cudnn_dtype=cudnn_dtype, shape=shape, name=name) + + def _last_axis_spec(self, desc: JaxTensorDesc, alignment: int) -> Any: + return self._to_tensor_spec( + desc, + divisibility=(None,) * (desc.ndim - 1) + (alignment,), + ) + + +class IndexerBackward(_IndexerJaxBase): + """Sparse indexer backward for BSHD or packed THD JAX arrays. + + Packed THD uses ``Q=(T_q,H,D)``, ``W=(T_q,H)``, ``K=(T_k,D)``, and + score/index tensors shaped ``(T_q,topk)``. Packed top-K indices must be + global flat K indices because this API intentionally has no host-visible + sequence metadata from which to reconstruct per-batch offsets. + """ + + def __init__( + self, + sample_index_q: Any, + sample_weights: Any, + sample_index_k: Any, + sample_attn_score: Any, + sample_index_score: Any, + sample_topk_indices: Any, + sample_d_index_q: Any | None = None, + sample_d_weights: Any | None = None, + sample_d_index_k: Any | None = None, + sm_scale: float = 1.0, + loss_coeff: float = 1.0, + block_I: int = DEFAULT_BLOCK_I, + topk_indices_global: bool = False, + target_compute_capability: int | None = None, + ) -> None: + self.iq_desc = self._to_tensor_desc(sample_index_q, "sample_index_q") + self.w_desc = self._to_tensor_desc(sample_weights, "sample_weights") + self.ik_desc = self._to_tensor_desc(sample_index_k, "sample_index_k") + self.attn_desc = self._to_tensor_desc(sample_attn_score, "sample_attn_score") + self.idx_score_desc = self._to_tensor_desc(sample_index_score, "sample_index_score") + self.topk_desc = self._to_tensor_desc(sample_topk_indices, "sample_topk_indices") + self.diq_desc = self._output_desc( + sample_d_index_q, + source=self.iq_desc, + cudnn_dtype=data_type.BFLOAT16, + shape=self.iq_desc.shape, + name="sample_d_index_q", + ) + self.dw_desc = self._output_desc( + sample_d_weights, + source=self.w_desc, + cudnn_dtype=data_type.BFLOAT16, + shape=self.w_desc.shape, + name="sample_d_weights", + ) + self.dik_desc = self._output_desc( + sample_d_index_k, + source=self.ik_desc, + cudnn_dtype=data_type.BFLOAT16, + shape=self.ik_desc.shape, + name="sample_d_index_k", + ) + self.grad_loss_desc = self.iq_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(1,), + name="grad_loss", + ) + self._op = IndexerBackwardOp( + index_q=self.iq_desc, + weights=self.w_desc, + index_k=self.ik_desc, + d_index_q=self.diq_desc, + d_weights=self.dw_desc, + d_index_k=self.dik_desc, + attn_score=self.attn_desc, + index_score=self.idx_score_desc, + topk_indices=self.topk_desc, + sm_scale=sm_scale, + block_i=block_I, + topk_indices_global=topk_indices_global, + ) + self.loss_coeff = _finite_float(loss_coeff, "loss_coeff") + self.target_compute_capability = target_compute_capability + self.compute_capability = None + + def check_support(self) -> bool: + self._op.check_support() + self._resolve_target("IndexerBackward") + return True + + def __call__( + self, + index_q: Any, + weights: Any, + index_k: Any, + attn_score: Any, + index_score: Any, + topk_indices: Any, + grad_loss: Any = 1.0, + ) -> TupleDict: + self.check_support() + for value, expected in ( + (index_q, self.iq_desc), + (weights, self.w_desc), + (index_k, self.ik_desc), + (attn_score, self.attn_desc), + (index_score, self.idx_score_desc), + (topk_indices, self.topk_desc), + ): + self._check_tensor_signature(value, expected) + grad_loss = _grad_loss_operand(grad_loss) + self._check_tensor_signature(grad_loss, self.grad_loss_desc) + + dk_accum_desc = self.ik_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=self.ik_desc.shape, + name="d_index_k_accum", + init_value=0.0, + ) + grad_signal_desc = self.attn_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=self.attn_desc.shape, + name="grad_signal", + ) + d_index_q, d_weights, d_index_k_accum = self._call_kernel( + (index_q, weights, index_k, attn_score, index_score, topk_indices, grad_loss), + output_descs=(self.diq_desc, self.dw_desc, dk_accum_desc), + workspace_descs=(grad_signal_desc,), + input_spec=( + self._last_axis_spec(self.iq_desc, self._op.head_dim), + None, + self._last_axis_spec(self.ik_desc, self._op.head_dim), + None, + None, + None, + None, + ), + output_spec=( + self._last_axis_spec(self.diq_desc, self._op.head_dim), + None, + self._last_axis_spec(dk_accum_desc, self._op.head_dim), + ), + compile_options=compile_options_for_target(self.compute_capability, "--opt-level 3"), + ) + return TupleDict( + d_index_q=d_index_q, + d_weights=d_weights, + d_index_k=d_index_k_accum.astype(self.dik_desc.dtype), + ) + + def _launch(self, inputs, outputs, workspaces, stream) -> None: + import cutlass + + index_q, weights, index_k, attn_score, index_score, topk_indices, grad_loss = inputs + d_index_q, d_weights, d_index_k_accum = outputs + (grad_signal,) = workspaces + grad_scale = self.loss_coeff / (self._op.batch * self._op.seqlen_q) + + if self._op.is_thd: + index_q = _prepend_unit_batch(index_q) + weights = _prepend_unit_batch(weights) + index_k = _prepend_unit_batch(index_k) + attn_score = _prepend_unit_batch(attn_score) + index_score = _prepend_unit_batch(index_score) + topk_indices = _prepend_unit_batch(topk_indices) + d_index_q = _prepend_unit_batch(d_index_q) + d_weights = _prepend_unit_batch(d_weights) + d_index_k_accum = _prepend_unit_batch(d_index_k_accum) + grad_signal = _prepend_unit_batch(grad_signal) + + if self._architecture_family == 100: + from .indexer_backward_sm100 import IndexerBackwardSm100, ScoreGradSm100 + + ScoreGradSm100(topk=self._op.topk)( + attn_score, + index_score, + grad_loss, + cutlass.Float32(grad_scale), + stream, + grad_signal, + None, + ) + IndexerBackwardSm100( + head_dim=self._op.head_dim, + heads=self._op.heads, + block_I=self._op.block_i, + topk=self._op.topk, + topk_indices_global=self._op.topk_indices_global, + )( + index_q, + weights, + index_k, + d_index_q, + d_weights, + d_index_k_accum, + grad_signal, + topk_indices, + cutlass.Float32(self._op.sm_scale), + stream, + ) + return + + from .indexer_backward_sm90 import IndexerBackwardSm90, ScoreGradSm90 + + ScoreGradSm90(topk=self._op.topk, index_is_log=False)( + attn_score, + index_score, + grad_loss, + cutlass.Float32(grad_scale), + stream, + grad_signal, + ) + IndexerBackwardSm90( + head_dim=self._op.head_dim, + heads=self._op.heads, + block_I=self._op.block_i, + topk=self._op.topk, + topk_indices_global=self._op.topk_indices_global, + )( + index_q, + weights, + index_k, + d_index_q, + d_weights, + d_index_k_accum, + grad_signal, + topk_indices, + cutlass.Float32(self._op.sm_scale), + stream, + None, + None, + cutlass.Int32(self._op.seqlen_q), + cutlass.Int32(self._op.seqlen_k), + None, + ) + + +def _prepend_unit_batch(tensor: Any) -> Any: + """Expose a compact packed tensor to the sparse kernels as batch one.""" + + import cutlass.cute as cute + + leading_stride = tensor.shape[0] * tensor.stride[0] + return cute.make_tensor( + tensor.iterator, + cute.make_layout( + (1, *tensor.shape), + stride=(leading_stride, *tensor.stride), + ), + ) + + +class DenseIndexerBackward(_IndexerJaxBase): + """Sample-signature-bound dense BSHD or packed-THD backward callable.""" + + def __init__( + self, + sample_index_q: Any, + sample_weights: Any, + sample_index_k: Any, + sample_attn_score: Any, + sample_attn_l1norm: Any, + sample_index_score: Any, + sample_index_lse: Any, + sample_d_index_q: Any | None = None, + sample_d_weights: Any | None = None, + sample_d_index_k: Any | None = None, + sample_cu_seqlens_q: Any | None = None, + sample_cu_seqlens_k: Any | None = None, + sample_q_causal_offsets: Any | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + sm_scale: float = 1.0, + loss_coeff: float = 1.0, + block_I: int = DEFAULT_BLOCK_I, + ratio: int = 1, + target_compute_capability: int | None = None, + ) -> None: + self.iq_desc = self._to_tensor_desc(sample_index_q, "sample_index_q") + self.w_desc = self._to_tensor_desc(sample_weights, "sample_weights") + self.ik_desc = self._to_tensor_desc(sample_index_k, "sample_index_k") + self.attn_desc = self._to_tensor_desc(sample_attn_score, "sample_attn_score") + self.attn_denom_desc = self._to_tensor_desc(sample_attn_l1norm, "sample_attn_l1norm") + self.idx_score_desc = self._to_tensor_desc(sample_index_score, "sample_index_score") + self.idx_lse_desc = self._to_tensor_desc(sample_index_lse, "sample_index_lse") + self.diq_desc = self._output_desc( + sample_d_index_q, + source=self.iq_desc, + cudnn_dtype=data_type.BFLOAT16, + shape=self.iq_desc.shape, + name="sample_d_index_q", + ) + self.dw_desc = self._output_desc( + sample_d_weights, + source=self.w_desc, + cudnn_dtype=data_type.BFLOAT16, + shape=self.w_desc.shape, + name="sample_d_weights", + ) + self.dik_desc = self._output_desc( + sample_d_index_k, + source=self.ik_desc, + cudnn_dtype=data_type.BFLOAT16, + shape=self.ik_desc.shape, + name="sample_d_index_k", + ) + self.grad_loss_desc = self.iq_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(1,), + name="grad_loss", + ) + self.cuq_desc = None if sample_cu_seqlens_q is None else self._to_tensor_desc(sample_cu_seqlens_q, "sample_cu_seqlens_q") + self.cuk_desc = None if sample_cu_seqlens_k is None else self._to_tensor_desc(sample_cu_seqlens_k, "sample_cu_seqlens_k") + self.q_offsets_desc = None if sample_q_causal_offsets is None else self._to_tensor_desc(sample_q_causal_offsets, "sample_q_causal_offsets") + self._op = DenseIndexerBackwardOp( + index_q=self.iq_desc, + weights=self.w_desc, + index_k=self.ik_desc, + d_index_q=self.diq_desc, + d_weights=self.dw_desc, + d_index_k=self.dik_desc, + attn_score=self.attn_desc, + attn_l1norm=self.attn_denom_desc, + index_score=self.idx_score_desc, + index_lse=self.idx_lse_desc, + cu_seqlens_q=self.cuq_desc, + cu_seqlens_k=self.cuk_desc, + q_causal_offsets=self.q_offsets_desc, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + sm_scale=sm_scale, + block_i=block_I, + ratio=ratio, + ) + self.loss_coeff = _finite_float(loss_coeff, "loss_coeff") + self.target_compute_capability = target_compute_capability + self.compute_capability = None + + def check_support(self) -> bool: + self._op.check_support() + self._resolve_target("DenseIndexerBackward") + return True + + def __call__( + self, + index_q: Any, + weights: Any, + index_k: Any, + attn_score: Any, + attn_l1norm: Any, + index_score: Any, + index_lse: Any, + grad_loss: Any = 1.0, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + q_causal_offsets: Any | None = None, + ) -> TupleDict: + self.check_support() + for value, expected in ( + (index_q, self.iq_desc), + (weights, self.w_desc), + (index_k, self.ik_desc), + (attn_score, self.attn_desc), + (attn_l1norm, self.attn_denom_desc), + (index_score, self.idx_score_desc), + (index_lse, self.idx_lse_desc), + ): + self._check_tensor_signature(value, expected) + for value, expected, name in ( + (cu_seqlens_q, self.cuq_desc, "cu_seqlens_q"), + (cu_seqlens_k, self.cuk_desc, "cu_seqlens_k"), + (q_causal_offsets, self.q_offsets_desc, "q_causal_offsets"), + ): + if (value is None) != (expected is None): + raise ValueError(f"{name} presence must match its sample") + if value is not None: + self._check_tensor_signature(value, expected) + grad_loss = _grad_loss_operand(grad_loss) + self._check_tensor_signature(grad_loss, self.grad_loss_desc) + + inputs = (index_q, weights, index_k, attn_score, attn_l1norm, index_score, index_lse, grad_loss) + for optional in (cu_seqlens_q, cu_seqlens_k, q_causal_offsets): + if optional is not None: + inputs += (optional,) + dk_accum_desc = self.ik_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=self.ik_desc.shape, + name="d_index_k_accum", + init_value=0.0, + ) + grad_signal_desc = self.idx_score_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=self.idx_score_desc.shape, + name="grad_signal", + ) + d_index_q, d_weights, d_index_k_accum = self._call_kernel( + inputs, + output_descs=(self.diq_desc, self.dw_desc, dk_accum_desc), + workspace_descs=(grad_signal_desc,), + input_spec=( + self._last_axis_spec(self.iq_desc, self._op.head_dim), + None, + self._last_axis_spec(self.ik_desc, self._op.head_dim), + ) + + (None,) * (len(inputs) - 3), + output_spec=( + self._last_axis_spec(self.diq_desc, self._op.head_dim), + None, + self._last_axis_spec(dk_accum_desc, self._op.head_dim), + ), + compile_options=compile_options_for_target(self.compute_capability, "--opt-level 3"), + ) + return TupleDict( + d_index_q=d_index_q, + d_weights=d_weights, + d_index_k=d_index_k_accum.astype(self.dik_desc.dtype), + ) + + def _launch(self, inputs, outputs, workspaces, stream) -> None: + import cutlass + + index_q, weights, index_k, attn_score, attn_l1norm, index_score, index_lse, grad_loss, *optional = inputs + cursor = 0 + cu_seqlens_q = cu_seqlens_k = q_causal_offsets = None + if self._op.is_thd: + cu_seqlens_q, cu_seqlens_k = optional[:2] + cursor = 2 + if self._op.q_causal_offsets is not None: + q_causal_offsets = optional[cursor] + + d_index_q, d_weights, d_index_k_accum = outputs + (grad_signal,) = workspaces + grad_scale = self.loss_coeff / max(self._op.normalization_tokens, 1) + + if self._architecture_family == 100: + from .dense_indexer_backward_sm100 import DenseIndexerBackward2QGemmSm100, ScoreGradDense + + ScoreGradDense(ratio=self._op.ratio, block_I=self._op.block_i)( + index_score, + attn_score, + index_lse, + attn_l1norm, + cu_seqlens_q, + cu_seqlens_k, + q_causal_offsets, + cutlass.Float32(grad_scale), + cutlass.Int32(self._op.max_seqlen_q), + cutlass.Int32(self._op.max_seqlen_k), + stream, + grad_signal, + grad_loss, + ) + DenseIndexerBackward2QGemmSm100( + head_dim=self._op.head_dim, + heads=self._op.heads, + block_I=self._op.block_i, + ratio=self._op.ratio, + )( + index_q, + weights, + index_k, + d_index_q, + d_weights, + d_index_k_accum, + grad_signal, + cu_seqlens_q, + cu_seqlens_k, + q_causal_offsets, + cutlass.Float32(self._op.sm_scale), + cutlass.Int32(self._op.max_seqlen_q), + cutlass.Int32(self._op.max_seqlen_k), + stream, + ) + return + + from .dense_indexer_backward_sm90 import ScoreGradDenseSm90 + from .indexer_backward_sm90 import IndexerBackwardSm90 + + ScoreGradDenseSm90(ratio=self._op.ratio, block_I=self._op.block_i)( + index_score, + attn_score, + index_lse, + attn_l1norm, + cu_seqlens_q, + cu_seqlens_k, + q_causal_offsets, + cutlass.Float32(grad_scale), + cutlass.Int32(self._op.max_seqlen_q), + cutlass.Int32(self._op.max_seqlen_k), + stream, + grad_signal, + grad_loss, + ) + IndexerBackwardSm90( + head_dim=self._op.head_dim, + heads=self._op.heads, + block_I=self._op.block_i, + topk=self._op.max_seqlen_k, + is_dense=True, + ratio=self._op.ratio, + )( + index_q, + weights, + index_k, + d_index_q, + d_weights, + d_index_k_accum, + grad_signal, + None, + cutlass.Float32(self._op.sm_scale), + stream, + cu_seqlens_q, + cu_seqlens_k, + cutlass.Int32(self._op.max_seqlen_q), + cutlass.Int32(self._op.max_seqlen_k), + q_causal_offsets, + ) + + +@jax.jit(static_argnames=("sm_scale", "loss_coeff", "block_I", "topk_indices_global", "target_compute_capability")) +def indexer_backward_wrapper( + index_q: Any, + weights: Any, + index_k: Any, + attn_score: Any, + index_score: Any, + topk_indices: Any, + sm_scale: float = 1.0, + loss_coeff: float = 1.0, + grad_loss: Any = 1.0, + block_I: int = DEFAULT_BLOCK_I, + topk_indices_global: bool = False, + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute sparse indexer gradients for fixed BSHD or packed THD inputs. + + Packed THD callers must pass ``topk_indices_global=True``. + """ + + samples = tuple(jax.ShapeDtypeStruct(value.shape, value.dtype) for value in (index_q, weights, index_k, attn_score, index_score, topk_indices)) + return IndexerBackward( + *samples, + sm_scale=sm_scale, + loss_coeff=loss_coeff, + block_I=block_I, + topk_indices_global=topk_indices_global, + target_compute_capability=target_compute_capability, + )(index_q, weights, index_k, attn_score, index_score, topk_indices, grad_loss) + + +@jax.jit( + static_argnames=( + "sm_scale", + "loss_coeff", + "block_I", + "ratio", + "max_seqlen_q", + "max_seqlen_k", + "target_compute_capability", + ) +) +def dense_indexer_backward_wrapper( + index_q: Any, + weights: Any, + index_k: Any, + attn_score: Any, + attn_l1norm: Any, + index_score: Any, + index_lse: Any, + sm_scale: float = 1.0, + loss_coeff: float = 1.0, + grad_loss: Any = 1.0, + block_I: int = DEFAULT_BLOCK_I, + ratio: int = 1, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + q_causal_offsets: Any | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + core = (index_q, weights, index_k, attn_score, attn_l1norm, index_score, index_lse) + samples = tuple(jax.ShapeDtypeStruct(value.shape, value.dtype) for value in core) + sample_cuq = None if cu_seqlens_q is None else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype) + sample_cuk = None if cu_seqlens_k is None else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype) + sample_offsets = None if q_causal_offsets is None else jax.ShapeDtypeStruct(q_causal_offsets.shape, q_causal_offsets.dtype) + return DenseIndexerBackward( + *samples, + sample_cu_seqlens_q=sample_cuq, + sample_cu_seqlens_k=sample_cuk, + sample_q_causal_offsets=sample_offsets, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + sm_scale=sm_scale, + loss_coeff=loss_coeff, + block_I=block_I, + ratio=ratio, + target_compute_capability=target_compute_capability, + )( + *core, + grad_loss=grad_loss, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + q_causal_offsets=q_causal_offsets, + ) + + +__all__ = [ + "DenseIndexerBackward", + "IndexerBackward", + "dense_indexer_backward_wrapper", + "indexer_backward_wrapper", +] diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/op.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/op.py new file mode 100644 index 000000000..b19120c71 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/op.py @@ -0,0 +1,438 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral DeepSeek sparse and dense indexer backward operations.""" + +from __future__ import annotations + +import math +from typing import Any, Optional + +from ... import data_type +from ..._op import Op +from ..._tensor_desc import TensorDesc + +DEFAULT_BLOCK_I = 128 +SUPPORTED_HEAD_DIM = 128 +MIN_HEADS = 64 + + +def _require_desc(name: str, desc: TensorDesc[Any] | None, *, optional: bool = False) -> None: + if desc is None and optional: + return + if not isinstance(desc, TensorDesc): + expected = "a TensorDesc or None" if optional else "a TensorDesc" + raise TypeError(f"{name} must be {expected}, got {type(desc).__name__}") + + +def _require_rank(desc: TensorDesc[Any], rank: int, name: str) -> None: + if desc.ndim != rank: + raise ValueError(f"{name} must have rank {rank}, got shape {desc.shape}") + + +def _require_dtype(desc: TensorDesc[Any], expected: data_type | tuple[data_type, ...], name: str) -> None: + expected_values = expected if isinstance(expected, tuple) else (expected,) + if desc.cudnn_dtype not in expected_values: + expected_text = " or ".join(value.name.lower() for value in expected_values) + raise ValueError(f"{name} must have dtype {expected_text}, got {desc.dtype}") + + +def _require_shape(desc: TensorDesc[Any], expected: tuple[int, ...], name: str) -> None: + if desc.shape != expected: + raise ValueError(f"{name} must have shape {expected}, got {desc.shape}") + + +def _require_row_major(desc: TensorDesc[Any]) -> None: + if not desc.is_compact(tuple(reversed(range(desc.ndim)))): + raise ValueError(f"{desc.name or 'tensor'} must be row-major contiguous, got stride {desc.stride}") + + +def _resolve_scale(value: float, name: str) -> float: + try: + resolved = float(value) + except (TypeError, ValueError) as error: + raise TypeError(f"{name} must be a real scalar, got {value!r}") from error + if not math.isfinite(resolved): + raise ValueError(f"{name} must be finite, got {resolved}") + return resolved + + +class IndexerBackwardOp(Op): + """Complete fixed BSHD or packed THD sparse backward signature.""" + + def __init__( + self, + *, + index_q: TensorDesc[Any], + weights: TensorDesc[Any], + index_k: TensorDesc[Any], + d_index_q: TensorDesc[Any], + d_weights: TensorDesc[Any], + d_index_k: TensorDesc[Any], + attn_score: TensorDesc[Any], + index_score: TensorDesc[Any], + topk_indices: TensorDesc[Any], + sm_scale: float = 1.0, + block_i: int = DEFAULT_BLOCK_I, + topk_indices_global: bool = False, + ) -> None: + descriptors = ( + ("index_q", index_q), + ("weights", weights), + ("index_k", index_k), + ("d_index_q", d_index_q), + ("d_weights", d_weights), + ("d_index_k", d_index_k), + ("attn_score", attn_score), + ("index_score", index_score), + ("topk_indices", topk_indices), + ) + for name, desc in descriptors: + _require_desc(name, desc) + + self.index_q = index_q + self.weights = weights + self.index_k = index_k + self.d_index_q = d_index_q + self.d_weights = d_weights + self.d_index_k = d_index_k + self.attn_score = attn_score + self.index_score = index_score + self.topk_indices = topk_indices + self.requested_sm_scale = sm_scale + self.block_i = int(block_i) + self.topk_indices_global = bool(topk_indices_global) + + self.is_thd: Optional[bool] = None + self.batch: Optional[int] = None + self.seqlen_q: Optional[int] = None + self.seqlen_k: Optional[int] = None + self.heads: Optional[int] = None + self.head_dim: Optional[int] = None + self.topk: Optional[int] = None + self.sm_scale: Optional[float] = None + + def check_support(self) -> bool: + self.is_thd = None + self.batch = self.seqlen_q = self.seqlen_k = self.heads = self.head_dim = self.topk = None + self.sm_scale = None + + if self.index_q.ndim not in (3, 4): + raise ValueError(f"index_q must use BSHD rank 4 or packed THD rank 3, got shape {self.index_q.shape}") + is_thd = self.index_q.ndim == 3 + q_rank = 3 if is_thd else 4 + auxiliary_rank = q_rank - 1 + for desc, rank, name in ( + (self.index_q, q_rank, "index_q"), + (self.weights, auxiliary_rank, "weights"), + (self.index_k, auxiliary_rank, "index_k"), + (self.d_index_q, q_rank, "d_index_q"), + (self.d_weights, auxiliary_rank, "d_weights"), + (self.d_index_k, auxiliary_rank, "d_index_k"), + (self.attn_score, auxiliary_rank, "attn_score"), + (self.index_score, auxiliary_rank, "index_score"), + (self.topk_indices, auxiliary_rank, "topk_indices"), + ): + _require_rank(desc, rank, name) + + for desc, name in ((self.index_q, "index_q"), (self.weights, "weights"), (self.index_k, "index_k")): + _require_dtype(desc, data_type.BFLOAT16, name) + _require_dtype(self.d_index_q, data_type.BFLOAT16, "d_index_q") + _require_dtype(self.d_weights, data_type.BFLOAT16, "d_weights") + _require_dtype(self.d_index_k, (data_type.BFLOAT16, data_type.FLOAT), "d_index_k") + _require_dtype(self.attn_score, data_type.FLOAT, "attn_score") + _require_dtype(self.index_score, data_type.FLOAT, "index_score") + _require_dtype(self.topk_indices, data_type.INT32, "topk_indices") + + if is_thd: + seqlen_q, heads, head_dim = self.index_q.shape + seqlen_k, head_dim_k = self.index_k.shape + batch = batch_k = 1 + if not self.topk_indices_global: + raise ValueError("Packed THD IndexerBackward requires topk_indices_global=True") + else: + batch, seqlen_q, heads, head_dim = self.index_q.shape + batch_k, seqlen_k, head_dim_k = self.index_k.shape + topk = self.topk_indices.shape[-1] + dimensions = (batch, seqlen_q, seqlen_k, heads, head_dim, topk) + if any(value <= 0 for value in dimensions): + raise ValueError(f"Indexer-backward dimensions must be positive, got {dimensions}") + if batch_k != batch or head_dim_k != head_dim: + raise ValueError("index_k batch and head dimensions must match index_q") + if heads < MIN_HEADS: + raise ValueError(f"IndexerBackward requires heads >= {MIN_HEADS}, got {heads}") + if head_dim != SUPPORTED_HEAD_DIM: + raise ValueError(f"IndexerBackward requires head_dim={SUPPORTED_HEAD_DIM}, got {head_dim}") + if self.block_i != DEFAULT_BLOCK_I: + raise ValueError(f"block_i must be {DEFAULT_BLOCK_I}, got {self.block_i}") + if topk % self.block_i != 0: + raise ValueError(f"topk ({topk}) must be divisible by block_i ({self.block_i})") + + weights_shape = (seqlen_q, heads) if is_thd else (batch, seqlen_q, heads) + score_shape = (seqlen_q, topk) if is_thd else (batch, seqlen_q, topk) + _require_shape(self.weights, weights_shape, "weights") + _require_shape(self.d_index_q, self.index_q.shape, "d_index_q") + _require_shape(self.d_weights, self.weights.shape, "d_weights") + _require_shape(self.d_index_k, self.index_k.shape, "d_index_k") + _require_shape(self.attn_score, score_shape, "attn_score") + _require_shape(self.index_score, score_shape, "index_score") + _require_shape(self.topk_indices, score_shape, "topk_indices") + for desc in ( + self.index_q, + self.weights, + self.index_k, + self.d_index_q, + self.d_weights, + self.d_index_k, + self.attn_score, + self.index_score, + self.topk_indices, + ): + _require_row_major(desc) + + self.is_thd = is_thd + self.batch = batch + self.seqlen_q = seqlen_q + self.seqlen_k = seqlen_k + self.heads = heads + self.head_dim = head_dim + self.topk = topk + self.sm_scale = _resolve_scale(self.requested_sm_scale, "sm_scale") + return True + + +class DenseIndexerBackwardOp(Op): + """Complete BSHD or packed-THD dense indexer-backward signature.""" + + def __init__( + self, + *, + index_q: TensorDesc[Any], + weights: TensorDesc[Any], + index_k: TensorDesc[Any], + d_index_q: TensorDesc[Any], + d_weights: TensorDesc[Any], + d_index_k: TensorDesc[Any], + attn_score: TensorDesc[Any], + attn_l1norm: TensorDesc[Any], + index_score: TensorDesc[Any], + index_lse: TensorDesc[Any], + cu_seqlens_q: TensorDesc[Any] | None = None, + cu_seqlens_k: TensorDesc[Any] | None = None, + q_causal_offsets: TensorDesc[Any] | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + sm_scale: float = 1.0, + block_i: int = DEFAULT_BLOCK_I, + ratio: int = 1, + ) -> None: + descriptors = ( + ("index_q", index_q), + ("weights", weights), + ("index_k", index_k), + ("d_index_q", d_index_q), + ("d_weights", d_weights), + ("d_index_k", d_index_k), + ("attn_score", attn_score), + ("attn_l1norm", attn_l1norm), + ("index_score", index_score), + ("index_lse", index_lse), + ) + for name, desc in descriptors: + _require_desc(name, desc) + for name, desc in (("cu_seqlens_q", cu_seqlens_q), ("cu_seqlens_k", cu_seqlens_k), ("q_causal_offsets", q_causal_offsets)): + _require_desc(name, desc, optional=True) + + self.index_q = index_q + self.weights = weights + self.index_k = index_k + self.d_index_q = d_index_q + self.d_weights = d_weights + self.d_index_k = d_index_k + self.attn_score = attn_score + self.attn_l1norm = attn_l1norm + self.index_score = index_score + self.index_lse = index_lse + self.cu_seqlens_q = cu_seqlens_q + self.cu_seqlens_k = cu_seqlens_k + self.q_causal_offsets = q_causal_offsets + self.requested_max_seqlen_q = max_seqlen_q + self.requested_max_seqlen_k = max_seqlen_k + self.requested_sm_scale = sm_scale + self.block_i = int(block_i) + self.ratio = int(ratio) + + self.is_thd: Optional[bool] = None + self.batch: Optional[int] = None + self.normalization_tokens: Optional[int] = None + self.total_k: Optional[int] = None + self.heads: Optional[int] = None + self.head_dim: Optional[int] = None + self.max_seqlen_q: Optional[int] = None + self.max_seqlen_k: Optional[int] = None + self.sm_scale: Optional[float] = None + + def check_support(self) -> bool: + self.is_thd = None + self.batch = self.normalization_tokens = self.total_k = self.heads = self.head_dim = None + self.max_seqlen_q = self.max_seqlen_k = None + self.sm_scale = None + + if (self.cu_seqlens_q is None) != (self.cu_seqlens_k is None): + raise ValueError("cu_seqlens_q and cu_seqlens_k must be provided together") + is_thd = self.cu_seqlens_q is not None + if is_thd: + batch, normalization_tokens, total_k, heads, head_dim, max_q, max_k = self._check_thd() + else: + batch, normalization_tokens, total_k, heads, head_dim, max_q, max_k = self._check_bshd() + + for desc, name in ((self.index_q, "index_q"), (self.weights, "weights"), (self.index_k, "index_k")): + _require_dtype(desc, data_type.BFLOAT16, name) + _require_dtype(self.d_index_q, data_type.BFLOAT16, "d_index_q") + _require_dtype(self.d_weights, data_type.BFLOAT16, "d_weights") + _require_dtype(self.d_index_k, (data_type.BFLOAT16, data_type.FLOAT), "d_index_k") + for desc, name in ( + (self.attn_score, "attn_score"), + (self.attn_l1norm, "attn_l1norm"), + (self.index_score, "index_score"), + (self.index_lse, "index_lse"), + ): + _require_dtype(desc, data_type.FLOAT, name) + for desc, name in ((self.cu_seqlens_q, "cu_seqlens_q"), (self.cu_seqlens_k, "cu_seqlens_k"), (self.q_causal_offsets, "q_causal_offsets")): + if desc is not None: + _require_dtype(desc, data_type.INT32, name) + + if heads < MIN_HEADS: + raise ValueError(f"DenseIndexerBackward requires heads >= {MIN_HEADS}, got {heads}") + if head_dim != SUPPORTED_HEAD_DIM: + raise ValueError(f"DenseIndexerBackward requires head_dim={SUPPORTED_HEAD_DIM}, got {head_dim}") + if self.block_i != DEFAULT_BLOCK_I: + raise ValueError(f"block_i must be {DEFAULT_BLOCK_I}, got {self.block_i}") + if self.ratio < 1: + raise ValueError(f"ratio must be >= 1, got {self.ratio}") + + for desc in ( + self.index_q, + self.weights, + self.index_k, + self.d_index_q, + self.d_weights, + self.d_index_k, + self.attn_score, + self.attn_l1norm, + self.index_score, + self.index_lse, + self.cu_seqlens_q, + self.cu_seqlens_k, + self.q_causal_offsets, + ): + if desc is not None: + _require_row_major(desc) + + self.is_thd = is_thd + self.batch = batch + self.normalization_tokens = normalization_tokens + self.total_k = total_k + self.heads = heads + self.head_dim = head_dim + self.max_seqlen_q = max_q + self.max_seqlen_k = max_k + self.sm_scale = _resolve_scale(self.requested_sm_scale, "sm_scale") + return True + + def _check_bshd(self) -> tuple[int, int, int, int, int, int, int]: + for desc, rank, name in ( + (self.index_q, 4, "index_q"), + (self.weights, 3, "weights"), + (self.index_k, 3, "index_k"), + (self.d_index_q, 4, "d_index_q"), + (self.d_weights, 3, "d_weights"), + (self.d_index_k, 3, "d_index_k"), + (self.attn_score, 3, "attn_score"), + (self.attn_l1norm, 2, "attn_l1norm"), + (self.index_score, 3, "index_score"), + (self.index_lse, 2, "index_lse"), + ): + _require_rank(desc, rank, name) + + batch, seqlen_q, heads, head_dim = self.index_q.shape + batch_k, seqlen_k, head_dim_k = self.index_k.shape + if min(batch, seqlen_q, seqlen_k, heads, head_dim) <= 0: + raise ValueError("Dense indexer-backward dimensions must be positive") + if batch_k != batch or head_dim_k != head_dim: + raise ValueError("index_k batch and head dimensions must match index_q") + _require_shape(self.weights, (batch, seqlen_q, heads), "weights") + _require_shape(self.d_index_q, self.index_q.shape, "d_index_q") + _require_shape(self.d_weights, self.weights.shape, "d_weights") + _require_shape(self.d_index_k, self.index_k.shape, "d_index_k") + score_shape = (batch, seqlen_q, seqlen_k) + denom_shape = (batch, seqlen_q) + _require_shape(self.attn_score, score_shape, "attn_score") + _require_shape(self.index_score, score_shape, "index_score") + _require_shape(self.attn_l1norm, denom_shape, "attn_l1norm") + _require_shape(self.index_lse, denom_shape, "index_lse") + if self.q_causal_offsets is not None: + _require_rank(self.q_causal_offsets, 1, "q_causal_offsets") + _require_shape(self.q_causal_offsets, (batch,), "q_causal_offsets") + + max_q = seqlen_q if self.requested_max_seqlen_q is None else int(self.requested_max_seqlen_q) + max_k = seqlen_k if self.requested_max_seqlen_k is None else int(self.requested_max_seqlen_k) + if (max_q, max_k) != (seqlen_q, seqlen_k): + raise ValueError(f"BSHD max_seqlen_q/k must match tensor extents {(seqlen_q, seqlen_k)}, got {(max_q, max_k)}") + return batch, batch * seqlen_q, batch * seqlen_k, heads, head_dim, max_q, max_k + + def _check_thd(self) -> tuple[int, int, int, int, int, int, int]: + for desc, rank, name in ( + (self.index_q, 3, "index_q"), + (self.weights, 2, "weights"), + (self.index_k, 2, "index_k"), + (self.d_index_q, 3, "d_index_q"), + (self.d_weights, 2, "d_weights"), + (self.d_index_k, 2, "d_index_k"), + (self.attn_score, 2, "attn_score"), + (self.attn_l1norm, 1, "attn_l1norm"), + (self.index_score, 2, "index_score"), + (self.index_lse, 1, "index_lse"), + (self.cu_seqlens_q, 1, "cu_seqlens_q"), + (self.cu_seqlens_k, 1, "cu_seqlens_k"), + ): + _require_rank(desc, rank, name) + + total_q, heads, head_dim = self.index_q.shape + total_k, head_dim_k = self.index_k.shape + batch = self.cu_seqlens_q.shape[0] - 1 + if min(total_q, total_k, batch, heads, head_dim) <= 0: + raise ValueError("Dense THD indexer-backward dimensions must be positive") + if self.cu_seqlens_k.shape != (batch + 1,): + raise ValueError("cu_seqlens_q and cu_seqlens_k must describe the same batch") + if head_dim_k != head_dim: + raise ValueError("index_k head dimension must match index_q") + if self.requested_max_seqlen_q is None or self.requested_max_seqlen_k is None: + raise ValueError("THD dense indexer backward requires max_seqlen_q and max_seqlen_k") + max_q = int(self.requested_max_seqlen_q) + max_k = int(self.requested_max_seqlen_k) + if max_q <= 0 or max_k <= 0: + raise ValueError(f"max_seqlen_q/k must be positive, got {(max_q, max_k)}") + + _require_shape(self.weights, (total_q, heads), "weights") + _require_shape(self.d_index_q, self.index_q.shape, "d_index_q") + _require_shape(self.d_weights, self.weights.shape, "d_weights") + _require_shape(self.d_index_k, self.index_k.shape, "d_index_k") + score_shape = (total_q, max_k) + _require_shape(self.attn_score, score_shape, "attn_score") + _require_shape(self.index_score, score_shape, "index_score") + _require_shape(self.attn_l1norm, (total_q,), "attn_l1norm") + _require_shape(self.index_lse, (total_q,), "index_lse") + if self.q_causal_offsets is not None: + _require_rank(self.q_causal_offsets, 1, "q_causal_offsets") + _require_shape(self.q_causal_offsets, (batch,), "q_causal_offsets") + return batch, total_q, total_k, heads, head_dim, max_q, max_k + + +__all__ = [ + "DEFAULT_BLOCK_I", + "DenseIndexerBackwardOp", + "IndexerBackwardOp", + "MIN_HEADS", + "SUPPORTED_HEAD_DIM", +] diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py index 4b8df4c14..012c83fd9 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py @@ -1,3 +1,19 @@ -from .api import IndexerForward, indexer_forward_wrapper +"""Lazy Torch API and framework-neutral indexer-forward exports. -__all__ = ["IndexerForward", "indexer_forward_wrapper"] +JAX APIs are exported through :mod:`cudnn.jax`. +""" + +from ..._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ( + "IndexerForwardOp", + "SUPPORTED_COMPUTE_CAPABILITIES", + "TMA_ALIGN_ELEMENTS", + ), + "api": ("IndexerForward", "indexer_forward_wrapper"), + }, + submodules=("api", "op"), +) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py index d35e0da8b..b66de8af3 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py @@ -26,8 +26,9 @@ from .indexer_fwd_sm100 import IndexerForwardSm100 from ._interface import indexer_fwd as indexer_fwd_sm100 from ._interface_sm90 import indexer_fwd as indexer_fwd_sm90 +from .op import IndexerForwardOp, TMA_ALIGN_ELEMENTS -TMA_ALIGN_ELEMS = 4 # FP32 output => seqlen_k padded to multiples of 4 (16 B) +TMA_ALIGN_ELEMS = TMA_ALIGN_ELEMENTS class IndexerForward(APIBase): @@ -66,6 +67,21 @@ def __init__( self.sm_scale = float(sm_scale) self.qhead_per_kv_head = qhead_per_kv_head + self._op = IndexerForwardOp( + q=self.q_desc, + k=self.k_desc, + weight=self.w_desc, + output=self.o_desc, + ratio=self.ratio, + qhead_per_kv_head=self.qhead_per_kv_head, + m_block_size=self.m_block_size, + n_block_size=self.n_block_size, + q_stage=self.q_stage, + kv_stage=self.kv_stage, + sm_scale=self.sm_scale, + target_compute_capability=100, + ) + self.batch_size = None self.s_q = None self.s_k = None @@ -76,55 +92,7 @@ def __init__( def check_support(self) -> bool: self._logger.debug("Entering check_support") - self._value_error_if( - self.q_desc.ndim != 4, - f"Q must be 4-D (B, S_q, H_q, D), got {self.q_desc.shape}", - ) - self._value_error_if( - self.k_desc.ndim != 4, - f"K must be 4-D (B, S_k, H_kv, D), got {self.k_desc.shape}", - ) - self._value_error_if( - self.w_desc.ndim != 3, - f"W must be 3-D (B, S_q, H_q), got {self.w_desc.shape}", - ) - self._value_error_if( - self.o_desc.ndim != 3, - f"Out must be 3-D (B, S_q, S_k_padded), got {self.o_desc.shape}", - ) - - b, s_q, h_q, d = self.q_desc.shape - b_k, s_k, h_kv, d_k = self.k_desc.shape - b_o, s_q_out, s_k_padded_from_out = self.o_desc.shape - self._value_error_if(b != b_k, f"Batch size mismatch Q={b} vs K={b_k}") - self._value_error_if(b != b_o, f"Batch size mismatch Q={b} vs Out={b_o}") - self._value_error_if(s_q != s_q_out, f"S_q mismatch Q={s_q} vs Out={s_q_out}") - self._value_error_if(d != d_k, f"Head dim mismatch Q={d} vs K={d_k}") - self._value_error_if( - d != 128, - f"IndexerForward is tuned for head_dim=128 only, got {d}", - ) - - qhpkv = self.qhead_per_kv_head if self.qhead_per_kv_head is not None else (h_q // h_kv) - self._value_error_if( - qhpkv * h_kv != h_q, - f"qhead_per_kv_head * h_kv != h_q ({qhpkv} * {h_kv} != {h_q})", - ) - self._value_error_if( - qhpkv not in (32, 64), - f"qhead_per_kv_head must be 32 or 64, got {qhpkv}", - ) - self.qhead_per_kv_head = qhpkv - - self._check_dtype(self.q_desc, torch.bfloat16, name="Q") - self._check_dtype(self.k_desc, torch.bfloat16, name="K") - self._check_dtype(self.w_desc, torch.bfloat16, name="W") - self._check_dtype(self.o_desc, torch.float32, name="Out") - - self._value_error_if( - s_k_padded_from_out % TMA_ALIGN_ELEMS != 0, - f"Out seqlen_k dim must be a multiple of {TMA_ALIGN_ELEMS}, got {s_k_padded_from_out}", - ) + self._op.check_support() major = device_major() self._runtime_error_if( @@ -132,13 +100,14 @@ def check_support(self) -> bool: f"IndexerForward requires SM100+ compute capability, found SM{major}", ) - self.batch_size = b - self.s_q = s_q - self.s_k = s_k - self.s_k_padded = s_k_padded_from_out - self.h_q = h_q - self.h_kv = h_kv - self.head_dim = d + self.batch_size = self._op.batch_size + self.s_q = self._op.s_q + self.s_k = self._op.s_k + self.s_k_padded = self.o_desc.shape[-1] + self.h_q = self._op.h_q + self.h_kv = self._op.h_kv + self.head_dim = self._op.head_dim + self.qhead_per_kv_head = self._op.qhead_per_kv_head self._is_supported = True return True @@ -172,6 +141,7 @@ def compile(self) -> None: cutlass.Float32(self.sm_scale), None, None, + None, fake_stream, options=compile_options(), ) @@ -191,6 +161,7 @@ def tensor_api(q, k, w, out, stream): cutlass.Float32(self.sm_scale), None, None, + None, stream, ) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py new file mode 100644 index 000000000..51edcfe8f --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py @@ -0,0 +1,434 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX API for DeepSeek indexer forward.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +import jax +import jax.numpy as jnp + +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..._jax.layout import mode_from_layout, to_public_axes +from ..utils.compiler import compile_options_for_target +from .op import IndexerForwardOp, SUPPORTED_COMPUTE_CAPABILITIES, TMA_ALIGN_ELEMENTS + + +def _resolve_layout( + name: str, + layout: str | None, + *, + default: str, + kernel_axes: str, + supported: tuple[str, ...], +) -> tuple[str, tuple[int, ...]]: + layout = default if layout is None else layout + mode = mode_from_layout(layout, kernel_axes=kernel_axes) + if layout not in supported: + choices = ", ".join(repr(value) for value in supported) + raise ValueError(f"{name} must be one of ({choices}), got {layout!r}") + return layout, mode + + +class IndexerForward(JaxApiBase): + """JAX callable specialized from fixed BSHD or packed THD metadata. + + Fixed inputs may use batch-major ``BSHD``/``BSH`` or sequence-major + ``SBHD``/``SBH`` public axis orders. The output may similarly use ``BSK`` + or ``SBK``. Packed inputs currently use ``THD``/``TH`` and return ``TK``. + + ``sample_out`` describes the padded physical FP32 kernel buffer. When it + is omitted, the adapter creates that descriptor and returns a logical view + with the unpadded K extent. + """ + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_w: Any, + *, + sample_out: Any | None = None, + sample_cu_seqlens_q: Any | None = None, + sample_cu_seqlens_k: Any | None = None, + sample_q_causal_offsets: Any | None = None, + ratio: int = 4, + qhead_per_kv_head: int | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + m_block_size: int = 128, + n_block_size: int = 128, + q_stage: int = 2, + kv_stage: int = 4, + sm_scale: float = 1.0, + q_layout: str | None = None, + k_layout: str | None = None, + w_layout: str | None = None, + output_layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + if (sample_cu_seqlens_q is None) != (sample_cu_seqlens_k is None): + raise ValueError("THD input requires both sample_cu_seqlens_q and sample_cu_seqlens_k") + + self.is_varlen = sample_cu_seqlens_q is not None + if self.is_varlen: + self.q_layout, self.q_mode = _resolve_layout( + "q_layout", + q_layout, + default="THD", + kernel_axes="THD", + supported=("THD",), + ) + self.k_layout, self.k_mode = _resolve_layout( + "k_layout", + k_layout, + default="THD", + kernel_axes="THD", + supported=("THD",), + ) + self.w_layout, self.w_mode = _resolve_layout( + "w_layout", + w_layout, + default="TH", + kernel_axes="TH", + supported=("TH",), + ) + self.output_layout, self.output_mode = _resolve_layout( + "output_layout", + output_layout, + default="TK", + kernel_axes="TK", + supported=("TK",), + ) + else: + self.q_layout, self.q_mode = _resolve_layout( + "q_layout", + q_layout, + default="BSHD", + kernel_axes="BSHD", + supported=("BSHD", "SBHD"), + ) + self.k_layout, self.k_mode = _resolve_layout( + "k_layout", + k_layout, + default="BSHD", + kernel_axes="BSHD", + supported=("BSHD", "SBHD"), + ) + self.w_layout, self.w_mode = _resolve_layout( + "w_layout", + w_layout, + default="BSH", + kernel_axes="BSH", + supported=("BSH", "SBH"), + ) + self.output_layout, self.output_mode = _resolve_layout( + "output_layout", + output_layout, + default="BSK", + kernel_axes="BSK", + supported=("BSK", "SBK"), + ) + + self.target_compute_capability = self._resolve_compute_capability( + target_compute_capability=target_compute_capability, + supported_compute_capabilities=SUPPORTED_COMPUTE_CAPABILITIES, + operation_name="IndexerForward", + ) + self.q_desc = self._to_tensor_desc(sample_q, "sample_q", mode=self.q_mode) + self.k_desc = self._to_tensor_desc(sample_k, "sample_k", mode=self.k_mode) + self.w_desc = self._to_tensor_desc(sample_w, "sample_w", mode=self.w_mode) + self.cu_seqlens_q_desc = None if sample_cu_seqlens_q is None else self._to_tensor_desc(sample_cu_seqlens_q, "sample_cu_seqlens_q") + self.cu_seqlens_k_desc = None if sample_cu_seqlens_k is None else self._to_tensor_desc(sample_cu_seqlens_k, "sample_cu_seqlens_k") + self.q_causal_offsets_desc = ( + None + if sample_q_causal_offsets is None + else self._to_tensor_desc( + sample_q_causal_offsets, + "sample_q_causal_offsets", + ) + ) + self.max_seqlen_q = max_seqlen_q + self.max_seqlen_k = max_seqlen_k + + if sample_out is None: + self.o_desc = self._default_output_desc(max_seqlen_k) + else: + self.o_desc = self._to_tensor_desc( + sample_out, + "sample_out", + mode=self.output_mode, + init_value=float("-inf"), + ) + + self._op = IndexerForwardOp( + q=self.q_desc, + k=self.k_desc, + weight=self.w_desc, + output=self.o_desc, + cu_seqlens_q=self.cu_seqlens_q_desc, + cu_seqlens_k=self.cu_seqlens_k_desc, + q_causal_offsets=self.q_causal_offsets_desc, + ratio=ratio, + qhead_per_kv_head=qhead_per_kv_head, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + m_block_size=m_block_size, + n_block_size=n_block_size, + q_stage=q_stage, + kv_stage=kv_stage, + sm_scale=sm_scale, + target_compute_capability=self.target_compute_capability, + ) + + def _default_output_desc(self, max_seqlen_k: int | None) -> JaxTensorDesc: + if self.is_varlen: + if self.q_desc.ndim != 3: + raise ValueError(f"THD Q must have rank 3, got {self.q_desc.shape}") + if max_seqlen_k is None: + raise ValueError("THD input requires max_seqlen_k") + leading_shape = (self.q_desc.shape[0],) + logical_seqlen_k = int(max_seqlen_k) + else: + if self.q_desc.ndim != 4 or self.k_desc.ndim != 4: + raise ValueError("Fixed indexer forward requires rank-4 Q and K") + leading_shape = (self.q_desc.shape[0], self.q_desc.shape[1]) + logical_seqlen_k = self.k_desc.shape[1] + + padded_seqlen_k = ((logical_seqlen_k + TMA_ALIGN_ELEMENTS - 1) // TMA_ALIGN_ELEMENTS) * TMA_ALIGN_ELEMENTS + canonical_shape = (*leading_shape, padded_seqlen_k) + return self._to_tensor_desc( + jax.ShapeDtypeStruct( + to_public_axes(canonical_shape, self.output_mode), + jnp.float32, + ), + "sample_out", + mode=self.output_mode, + init_value=float("-inf"), + ) + + def check_support(self) -> bool: + return self._op.check_support() + + def __call__( + self, + q: Any, + k: Any, + w: Any, + *, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + q_causal_offsets: Any | None = None, + ) -> TupleDict: + self.check_support() + self._check_tensor_signature(q, self.q_desc, mode=self.q_mode) + self._check_tensor_signature(k, self.k_desc, mode=self.k_mode) + self._check_tensor_signature(w, self.w_desc, mode=self.w_mode) + self._check_optional_signature(cu_seqlens_q, self.cu_seqlens_q_desc, "cu_seqlens_q") + self._check_optional_signature(cu_seqlens_k, self.cu_seqlens_k_desc, "cu_seqlens_k") + self._check_optional_signature(q_causal_offsets, self.q_causal_offsets_desc, "q_causal_offsets") + + inputs = [q, k, w] + input_bindings = [ + (self.q_desc, self.q_mode), + (self.k_desc, self.k_mode), + (self.w_desc, self.w_mode), + ] + for value, desc in ( + (cu_seqlens_q, self.cu_seqlens_q_desc), + (cu_seqlens_k, self.cu_seqlens_k_desc), + (q_causal_offsets, self.q_causal_offsets_desc), + ): + if desc is not None: + inputs.append(value) + input_bindings.append((desc, None)) + + output_divisibility = (None,) * (self.o_desc.ndim - 1) + (TMA_ALIGN_ELEMENTS,) + (scores_padded,) = self._call_kernel( + tuple(inputs), + output_descs=(self.o_desc,), + input_spec=tuple(self._to_tensor_spec(desc, mode=mode) for desc, mode in input_bindings), + output_spec=( + self._to_tensor_spec( + self.o_desc, + mode=self.output_mode, + divisibility=output_divisibility, + ), + ), + compile_options=compile_options_for_target(self.target_compute_capability), + ) + if self._op.s_k is None: + raise RuntimeError("IndexerForward output shape was not resolved by check_support()") + output_slice = [slice(None)] * self.o_desc.ndim + output_slice[self.output_mode[-1]] = slice(0, self._op.s_k) + return TupleDict(scores=scores_padded[tuple(output_slice)]) + + @staticmethod + def _check_optional_signature(value: Any | None, desc: JaxTensorDesc | None, name: str) -> None: + if (value is None) != (desc is None): + expected = "omitted" if desc is None else "provided" + raise ValueError(f"{name} must be {expected} for this specialized callable") + if desc is not None: + JaxApiBase._check_tensor_signature(value, desc) + + def _launch( + self, + inputs: tuple[Any, ...], + outputs: tuple[Any, ...], + workspaces: tuple[Any, ...], + stream: Any, + ) -> None: + from cutlass import BFloat16, Float32, Int32 + + q, k, w, *optional_inputs = inputs + cu_seqlens_q = optional_inputs.pop(0) if self.cu_seqlens_q_desc is not None else None + cu_seqlens_k = optional_inputs.pop(0) if self.cu_seqlens_k_desc is not None else None + q_causal_offsets = optional_inputs.pop(0) if self.q_causal_offsets_desc is not None else None + if optional_inputs: + raise RuntimeError("Unexpected IndexerForward kernel inputs") + (output,) = outputs + if workspaces: + raise RuntimeError("IndexerForward does not use workspaces") + + resolved = ( + self._op.head_dim, + self._op.qhead_per_kv_head, + self._op.h_kv, + self._op.s_q, + self._op.s_k, + ) + if any(value is None for value in resolved): + raise RuntimeError("IndexerForward launch configuration was not resolved by check_support()") + head_dim, qhead_per_kv_head, h_kv, s_q, s_k = resolved + if self.target_compute_capability < 100: + from .indexer_fwd_sm90 import IndexerForwardSm90 + + kernel = IndexerForwardSm90( + BFloat16, + head_dim=head_dim, + qhead_per_kvhead=qhead_per_kv_head, + ratio=self._op.ratio, + is_varlen=bool(self._op.is_varlen), + ) + else: + from .indexer_fwd_sm100 import IndexerForwardSm100 + + kernel = IndexerForwardSm100( + head_dim=head_dim, + qhead_per_kvhead=qhead_per_kv_head, + ratio=self._op.ratio, + is_varlen=bool(self._op.is_varlen), + m_block_size=self._op.m_block_size, + n_block_size=self._op.n_block_size, + q_stage=self._op.q_stage, + kv_stage=self._op.kv_stage, + ) + + kernel( + q, + k, + w, + output, + Int32(h_kv), + Int32(s_q), + Int32(s_k), + Float32(self._op.sm_scale), + cu_seqlens_q, + cu_seqlens_k, + q_causal_offsets, + stream, + ) + + +@partial( + jax.jit, + static_argnames=( + "ratio", + "qhead_per_kv_head", + "max_seqlen_q", + "max_seqlen_k", + "m_block_size", + "n_block_size", + "q_stage", + "kv_stage", + "sm_scale", + "q_layout", + "k_layout", + "w_layout", + "output_layout", + "target_compute_capability", + ), +) +def indexer_forward_wrapper( + q: Any, + k: Any, + w: Any, + *, + ratio: int = 4, + qhead_per_kv_head: int | None = None, + m_block_size: int = 128, + n_block_size: int = 128, + q_stage: int = 2, + kv_stage: int = 4, + sm_scale: float = 1.0, + q_layout: str | None = None, + k_layout: str | None = None, + w_layout: str | None = None, + output_layout: str | None = None, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + q_causal_offsets: Any | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute fixed batch/sequence-major or packed THD indexer scores.""" + + sample_q = jax.ShapeDtypeStruct(q.shape, q.dtype) + sample_k = jax.ShapeDtypeStruct(k.shape, k.dtype) + sample_w = jax.ShapeDtypeStruct(w.shape, w.dtype) + sample_cu_q = None if cu_seqlens_q is None else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype) + sample_cu_k = None if cu_seqlens_k is None else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype) + sample_offsets = ( + None + if q_causal_offsets is None + else jax.ShapeDtypeStruct( + q_causal_offsets.shape, + q_causal_offsets.dtype, + ) + ) + return IndexerForward( + sample_q, + sample_k, + sample_w, + sample_cu_seqlens_q=sample_cu_q, + sample_cu_seqlens_k=sample_cu_k, + sample_q_causal_offsets=sample_offsets, + ratio=ratio, + qhead_per_kv_head=qhead_per_kv_head, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + m_block_size=m_block_size, + n_block_size=n_block_size, + q_stage=q_stage, + kv_stage=kv_stage, + sm_scale=sm_scale, + q_layout=q_layout, + k_layout=k_layout, + w_layout=w_layout, + output_layout=output_layout, + target_compute_capability=target_compute_capability, + )( + q, + k, + w, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + q_causal_offsets=q_causal_offsets, + ) + + +__all__ = ["IndexerForward", "indexer_forward_wrapper"] diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/op.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/op.py new file mode 100644 index 000000000..5ca7af613 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/op.py @@ -0,0 +1,229 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral validation for DeepSeek indexer forward.""" + +from __future__ import annotations + +from typing import Any + +from ... import data_type +from ..._op import Op +from ..._tensor_desc import TensorDesc + +TMA_ALIGN_ELEMENTS = 4 +SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) + + +def _require_compact( + desc: TensorDesc[Any], + label: str, + contiguous_axis: str, +) -> None: + if not desc.is_compact() or desc.stride[-1] != 1: + raise ValueError( + f"{label} must be compact with its canonical {contiguous_axis} axis " f"contiguous, got stride {desc.stride} and stride order {desc.stride_order}" + ) + + +def _require_desc(name: str, desc: TensorDesc[Any] | None) -> None: + if desc is not None and not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc or None, got {type(desc).__name__}") + + +class IndexerForwardOp(Op): + """Complete tensor signature and static configuration for indexer forward.""" + + def __init__( + self, + *, + q: TensorDesc[Any], + k: TensorDesc[Any], + weight: TensorDesc[Any], + output: TensorDesc[Any], + cu_seqlens_q: TensorDesc[Any] | None = None, + cu_seqlens_k: TensorDesc[Any] | None = None, + q_causal_offsets: TensorDesc[Any] | None = None, + ratio: int = 4, + qhead_per_kv_head: int | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + m_block_size: int = 128, + n_block_size: int = 128, + q_stage: int = 2, + kv_stage: int = 4, + sm_scale: float = 1.0, + target_compute_capability: int = 100, + ) -> None: + for name, desc in (("q", q), ("k", k), ("weight", weight), ("output", output)): + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + for name, desc in ( + ("cu_seqlens_q", cu_seqlens_q), + ("cu_seqlens_k", cu_seqlens_k), + ("q_causal_offsets", q_causal_offsets), + ): + _require_desc(name, desc) + + self.q = q + self.k = k + self.weight = weight + self.output = output + self.cu_seqlens_q = cu_seqlens_q + self.cu_seqlens_k = cu_seqlens_k + self.q_causal_offsets = q_causal_offsets + self.ratio = int(ratio) + self.requested_qhead_per_kv_head = qhead_per_kv_head + self.requested_max_seqlen_q = max_seqlen_q + self.requested_max_seqlen_k = max_seqlen_k + self.m_block_size = int(m_block_size) + self.n_block_size = int(n_block_size) + self.q_stage = int(q_stage) + self.kv_stage = int(kv_stage) + self.sm_scale = float(sm_scale) + self.target_compute_capability = int(target_compute_capability) + + self.is_varlen: bool | None = None + self.batch_size: int | None = None + self.s_q: int | None = None + self.s_k: int | None = None + self.h_q: int | None = None + self.h_kv: int | None = None + self.head_dim: int | None = None + self.qhead_per_kv_head: int | None = None + + def check_support(self) -> bool: + is_varlen = self.cu_seqlens_q is not None or self.cu_seqlens_k is not None + if (self.cu_seqlens_q is None) != (self.cu_seqlens_k is None): + raise ValueError("THD input requires both cu_seqlens_q and cu_seqlens_k") + if self.target_compute_capability not in SUPPORTED_COMPUTE_CAPABILITIES: + raise ValueError("target_compute_capability must be one of " f"{SUPPORTED_COMPUTE_CAPABILITIES}, got {self.target_compute_capability}") + + for desc, label in ((self.q, "Q"), (self.k, "K"), (self.weight, "W")): + if desc.cudnn_dtype != data_type.BFLOAT16: + raise ValueError(f"{label} must have dtype bfloat16, got {desc.dtype}") + if self.output.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"Out must have dtype float32, got {self.output.dtype}") + + if is_varlen: + batch_size, s_q, s_k, h_q, h_kv, head_dim = self._check_varlen_signature() + else: + batch_size, s_q, s_k, h_q, h_kv, head_dim = self._check_fixed_signature() + for desc, label, contiguous_axis in ( + (self.q, "Q", "D"), + (self.k, "K", "D"), + (self.weight, "W", "H"), + (self.output, "Out", "K"), + ): + _require_compact(desc, label, contiguous_axis) + + if any(value <= 0 for value in (batch_size, s_q, s_k, h_q, h_kv, head_dim)): + raise ValueError("IndexerForward dimensions must be positive, got " f"B={batch_size}, S_q={s_q}, S_k={s_k}, H_q={h_q}, H_kv={h_kv}, D={head_dim}") + if head_dim != 128: + raise ValueError(f"IndexerForward requires head_dim=128, got {head_dim}") + + qhead_per_kv_head = self.requested_qhead_per_kv_head + if qhead_per_kv_head is None: + if h_q % h_kv != 0: + raise ValueError(f"H_q ({h_q}) must be divisible by H_kv ({h_kv})") + qhead_per_kv_head = h_q // h_kv + if qhead_per_kv_head * h_kv != h_q: + raise ValueError("qhead_per_kv_head * H_kv must equal H_q, got " f"{qhead_per_kv_head} * {h_kv} != {h_q}") + if qhead_per_kv_head not in (32, 64): + raise ValueError(f"qhead_per_kv_head must be 32 or 64, got {qhead_per_kv_head}") + if self.target_compute_capability < 100 and h_kv != 1: + raise ValueError(f"SM90 IndexerForward requires H_kv=1, got {h_kv}") + if self.ratio < 1: + raise ValueError(f"ratio must be at least 1, got {self.ratio}") + + tuning = (self.m_block_size, self.n_block_size, self.q_stage, self.kv_stage) + if any(value <= 0 for value in tuning): + raise ValueError(f"IndexerForward tuning values must be positive, got {tuning}") + if self.target_compute_capability < 100 and tuning != (128, 128, 2, 4): + raise ValueError("SM90 IndexerForward supports only m_block_size=128, n_block_size=128, " f"q_stage=2, kv_stage=4; got {tuning}") + + if self.q_causal_offsets is not None: + offsets = self.q_causal_offsets + if offsets.ndim != 1 or offsets.shape != (batch_size,): + raise ValueError(f"q_causal_offsets must have shape {(batch_size,)}, got {offsets.shape}") + if offsets.cudnn_dtype != data_type.INT32: + raise ValueError(f"q_causal_offsets must have dtype int32, got {offsets.dtype}") + _require_compact(offsets, "q_causal_offsets", "entry") + + expected_leading_shape = (self.q.shape[0],) if is_varlen else (batch_size, s_q) + output_seqlen_k = self.output.shape[-1] + if self.output.shape[:-1] != expected_leading_shape or output_seqlen_k < s_k or output_seqlen_k % TMA_ALIGN_ELEMENTS: + raise ValueError( + "Out must have leading shape " + f"{expected_leading_shape} and an S_k extent >= {s_k} divisible by " + f"{TMA_ALIGN_ELEMENTS}, got {self.output.shape}" + ) + + self.is_varlen = is_varlen + self.batch_size = batch_size + self.s_q = s_q + self.s_k = s_k + self.h_q = h_q + self.h_kv = h_kv + self.head_dim = head_dim + self.qhead_per_kv_head = qhead_per_kv_head + return True + + def _check_fixed_signature(self) -> tuple[int, int, int, int, int, int]: + if self.q.ndim != 4: + raise ValueError(f"Q must be 4-D (B, S_q, H_q, D), got {self.q.shape}") + if self.k.ndim != 4: + raise ValueError(f"K must be 4-D (B, S_k, H_kv, D), got {self.k.shape}") + if self.weight.ndim != 3: + raise ValueError(f"W must be 3-D (B, S_q, H_q), got {self.weight.shape}") + if self.output.ndim != 3: + raise ValueError(f"Out must be 3-D, got {self.output.shape}") + + batch_size, s_q, h_q, head_dim = self.q.shape + k_batch, s_k, h_kv, k_head_dim = self.k.shape + if k_batch != batch_size: + raise ValueError(f"Q and K batch dimensions must match, got {batch_size} and {k_batch}") + if k_head_dim != head_dim: + raise ValueError(f"Q and K head dimensions must match, got {head_dim} and {k_head_dim}") + if self.weight.shape != (batch_size, s_q, h_q): + raise ValueError(f"W must have shape {(batch_size, s_q, h_q)}, got {self.weight.shape}") + if self.requested_max_seqlen_q not in (None, s_q): + raise ValueError(f"max_seqlen_q must be omitted or equal S_q={s_q}") + if self.requested_max_seqlen_k not in (None, s_k): + raise ValueError(f"max_seqlen_k must be omitted or equal S_k={s_k}") + return batch_size, s_q, s_k, h_q, h_kv, head_dim + + def _check_varlen_signature(self) -> tuple[int, int, int, int, int, int]: + if self.q.ndim != 3: + raise ValueError(f"THD Q must be 3-D (total_q, H_q, D), got {self.q.shape}") + if self.k.ndim != 3: + raise ValueError(f"THD K must be 3-D (total_k, H_kv, D), got {self.k.shape}") + if self.weight.ndim != 2: + raise ValueError(f"THD W must be 2-D (total_q, H_q), got {self.weight.shape}") + if self.output.ndim != 2: + raise ValueError(f"THD Out must be 2-D, got {self.output.shape}") + + if self.cu_seqlens_q is None or self.cu_seqlens_k is None: + raise RuntimeError("THD cumulative sequence descriptors were not configured") + for desc, label in ((self.cu_seqlens_q, "cu_seqlens_q"), (self.cu_seqlens_k, "cu_seqlens_k")): + if desc.ndim != 1 or desc.cudnn_dtype != data_type.INT32: + raise ValueError(f"{label} must be a 1-D int32 tensor, got shape {desc.shape} and dtype {desc.dtype}") + _require_compact(desc, label, "entry") + if self.cu_seqlens_q.shape != self.cu_seqlens_k.shape: + raise ValueError("cu_seqlens_q and cu_seqlens_k must have the same shape") + if self.cu_seqlens_q.shape[0] < 2: + raise ValueError("cu_seqlens_q must contain at least two entries") + if self.requested_max_seqlen_q is None or self.requested_max_seqlen_k is None: + raise ValueError("THD input requires max_seqlen_q and max_seqlen_k") + + total_q, h_q, head_dim = self.q.shape + _, h_kv, k_head_dim = self.k.shape + if k_head_dim != head_dim: + raise ValueError(f"Q and K head dimensions must match, got {head_dim} and {k_head_dim}") + if self.weight.shape != (total_q, h_q): + raise ValueError(f"THD W must have shape {(total_q, h_q)}, got {self.weight.shape}") + batch_size = self.cu_seqlens_q.shape[0] - 1 + return batch_size, int(self.requested_max_seqlen_q), int(self.requested_max_seqlen_k), h_q, h_kv, head_dim + + +__all__ = ["IndexerForwardOp", "SUPPORTED_COMPUTE_CAPABILITIES", "TMA_ALIGN_ELEMENTS"] diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/__init__.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/__init__.py index 481ef694e..909029826 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/__init__.py @@ -1,13 +1,20 @@ -from .api import ( - IndexerTopK, - indexer_top_k_wrapper, - local_to_global_wrapper, - compactify_wrapper, -) +"""Lazy Torch API and framework-neutral indexer top-K exports. + +JAX APIs are exported through :mod:`cudnn.jax`. +""" -__all__ = [ - "IndexerTopK", - "indexer_top_k_wrapper", - "local_to_global_wrapper", - "compactify_wrapper", -] +from ..._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ("IndexerTopKOp", "bucket_num_cols"), + "api": ( + "IndexerTopK", + "compactify_wrapper", + "indexer_top_k_wrapper", + "local_to_global_wrapper", + ), + }, + submodules=("api", "op"), +) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/api.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/api.py index fd5a107a6..b1502bffd 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/api.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/api.py @@ -7,10 +7,13 @@ import torch import cuda.bindings.driver as cuda +from cudnn import data_type from cudnn.api_base import APIBase, TupleDict +from cudnn.deepseek_sparse_attention.utils.runtime import torch_stream_context from .local_to_global_dsl import local_to_global as _local_to_global from .compactify import compactify as _compactify +from .op import IndexerTopKOp _SUPPORTED_DTYPES = (torch.float32, torch.float16, torch.bfloat16) @@ -77,37 +80,52 @@ def __init__( self.return_val = bool(return_val) self.num_copy_bits = int(num_copy_bits) - def check_support(self) -> bool: - self._logger.debug("Entering check_support") + if self.input_desc.ndim != 2: + raise ValueError(f"input_values must be 2-D, got {self.input_desc.shape}") + if self.seq_lens_desc.ndim != 1: + raise ValueError(f"seq_lens must be 1-D, got {self.seq_lens_desc.shape}") self._check_dtype(self.input_desc, list(_SUPPORTED_DTYPES), name="input_values") - self._check_dtype(self.seq_lens_desc, torch.int32, name="seq_lens") - self._value_error_if( - self.input_desc.ndim != 2, - f"input_values must be 2-D (n_rows, num_cols), got {self.input_desc.shape}", + if self.top_k <= 0: + raise ValueError(f"top_k must be positive, got {self.top_k}") + + output_shape = (self.input_desc.shape[0], self.top_k) + self.indices_desc = self.input_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=output_shape, + name="indices", ) - self._value_error_if( - self.seq_lens_desc.ndim != 1, - f"seq_lens must be 1-D, got {self.seq_lens_desc.shape}", + self.values_desc = ( + self.input_desc.compact_like( + cudnn_dtype=self.input_desc.cudnn_dtype, + shape=output_shape, + name="values", + ) + if self.return_val + else None ) - self._value_error_if( - self.top_k <= 0 or self.top_k > 2048, - f"top_k must be in (0, 2048], got {self.top_k}", + buffer_count = 2 if self.input_desc.cudnn_dtype == data_type.FLOAT else 1 + self.workspace_desc = self.input_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=(self.input_desc.shape[0], buffer_count, self.input_desc.shape[1]), + name="extra_buffer", ) - - # Enforce the kernel's n_rows == batch_size * next_n invariant - # up-front so misuse surfaces here rather than as silently-empty - # rows (the stagger formula produces non-positive lengths when - # next_n exceeds n_rows per batch). - n_rows = self.input_desc.shape[0] - batch_size = self.seq_lens_desc.shape[0] - self._value_error_if( - n_rows != batch_size * self.next_n, - f"n_rows ({n_rows}) must equal seq_lens.numel() * next_n " - f"({batch_size} * {self.next_n} = {batch_size * self.next_n}). " - f"For independent top-K over equal-length rows use " - f"next_n=1 and seq_lens of shape (n_rows,).", + self._op = IndexerTopKOp( + input_values=self.input_desc, + seq_lens=self.seq_lens_desc, + output_indices=self.indices_desc, + output_values=self.values_desc, + workspace=self.workspace_desc, + top_k=self.top_k, + next_n=self.next_n, + return_val=self.return_val, + num_copy_bits=self.num_copy_bits, + target_compute_capability=90, ) + def check_support(self) -> bool: + self._logger.debug("Entering check_support") + self._op.check_support() + major, _ = torch.cuda.get_device_capability() self._runtime_error_if( major < 9, @@ -136,14 +154,15 @@ def execute( self.compile() kernel = self._compiled_kernel or _get_cute_dsl_topk_wrapper() - indices, values = kernel( - input_values, - seq_lens, - self.top_k, - self.next_n, - return_val=self.return_val, - num_copy_bits=self.num_copy_bits, - ) + with torch_stream_context(current_stream): + indices, values = kernel( + input_values, + seq_lens, + self.top_k, + self.next_n, + return_val=self.return_val, + num_copy_bits=self.num_copy_bits, + ) return indices, values @@ -170,6 +189,8 @@ def indexer_top_k_wrapper( full details. ``values`` is ``None`` when ``return_val=False``. + Every runtime ``seq_lens[b]`` must satisfy + ``top_k + next_n - 1 <= seq_lens[b] <= input_values.shape[1]``. """ cache_key = ( input_values.dtype, @@ -191,7 +212,7 @@ def indexer_top_k_wrapper( return_val=return_val, num_copy_bits=num_copy_bits, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_IndexerTopKObjects[cache_key] = obj diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/compactify.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/compactify.py index ea1997711..48739b6f4 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/compactify.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/compactify.py @@ -9,27 +9,26 @@ from __future__ import annotations -from typing import Optional, Tuple +from typing import TYPE_CHECKING, Optional, Tuple -import torch import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute from cutlass import Int32, const_expr -from cudnn.deepseek_sparse_attention.utils.compiler import compile_options -from cudnn.deepseek_sparse_attention.utils.runtime import ( - device_major, - resolve_stream as _resolve_stream, -) -from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor as _to_cute +if TYPE_CHECKING: + import torch WARPS_PER_CTA = 4 ROWS_PER_CTA = WARPS_PER_CTA # One warp per row. def is_available() -> bool: + import torch + + from cudnn.deepseek_sparse_attention.utils.runtime import device_major + return torch.cuda.is_available() and device_major() >= 9 @@ -124,6 +123,8 @@ def kernel(self, mIn: cute.Tensor, mOut: cute.Tensor, mLen: cute.Tensor, rows: I def _compile_or_fetch(key, kernel_obj, *args): + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + if key not in _compile_cache: _compile_cache[key] = cute.compile(kernel_obj, *args, options=compile_options("--opt-level 3")) return _compile_cache[key] @@ -139,6 +140,11 @@ def compactify( batch-major to (B*S, K), matching DSA sparse attention's flat top-K layout. Returns (indices, topk_length). """ + import torch + + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream as _resolve_stream + from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor as _to_cute + if not is_available(): raise RuntimeError("compactify requires SM90+ CUDA device") if idxs.dtype != torch.int32: diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_decode_varlen.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_decode_varlen.py index 41ff88173..593b2dd59 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_decode_varlen.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_decode_varlen.py @@ -14,17 +14,12 @@ """SM90+ CuTe DSL indexer top-K decode kernel.""" -import math - import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute import cutlass.utils as utils -import torch from cutlass.utils.distributed import atomicAdd -from cudnn.deepseek_sparse_attention.utils.compiler import compile_options - from .block_scan import block_prefix_sum_kernel from .indexer_top_k_varlen_util import IndexerTopKKernelVarlen @@ -573,13 +568,6 @@ def _next_positive_power_of_2(x: int) -> int: return 1 << (x - 1).bit_length() -_TORCH_TO_CUTLASS_DTYPE = { - torch.float16: cutlass.Float16, - torch.bfloat16: cutlass.BFloat16, - torch.float32: cutlass.Float32, -} - - def _bucket_num_cols(num_cols: int) -> int: """Bucket num_cols to the next power of 2 for compilation caching. @@ -603,8 +591,19 @@ def cute_dsl_topk_wrapper( load_balance=False, num_copy_bits=256, ): + import math + + import torch + + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + + torch_to_cutlass_dtype = { + torch.float16: cutlass.Float16, + torch.bfloat16: cutlass.BFloat16, + torch.float32: cutlass.Float32, + } torch_dtype = input_values.dtype - dtype = _TORCH_TO_CUTLASS_DTYPE[torch_dtype] + dtype = torch_to_cutlass_dtype[torch_dtype] num_rows, num_cols = input_values.shape bucketed_num_cols = _bucket_num_cols(num_cols) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py index 57c7f9f37..ac72aa09c 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py @@ -12,15 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + +from typing import TYPE_CHECKING import cutlass import cutlass.cute as cute -import torch from cutlass._mlir.dialects import llvm from cutlass.utils.distributed import atomicAdd from .block_scan import block_prefix_sum_kernel, fence_acq_rel_cta +if TYPE_CHECKING: + import torch + """ top-k varlen utils. could be used by prefill and decode phase. """ @@ -1089,6 +1094,8 @@ def create_random_logits( Returns: Tensor of shape (num_rows, max_row_length) with random values and -inf padding """ + import torch + torch.manual_seed(seed) torch.cuda.manual_seed(seed) num_rows = row_starts.shape[0] @@ -1145,6 +1152,8 @@ def compare_top_k_results( Returns: True if results match within tolerance, False otherwise """ + import torch + num_rows = cuda_indices.shape[0] # Calculate valid lengths for each row (vectorized) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py new file mode 100644 index 000000000..cfe42d2fd --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py @@ -0,0 +1,404 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX APIs for DeepSeek indexer top-K and index utilities.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._jax import JaxApiBase, TupleDict +from ..utils.compiler import compile_options_for_target +from .op import INT32_MAX, IndexerTopKOp, SUPPORTED_COMPUTE_CAPABILITIES, SUPPORTED_INPUT_DTYPES + + +class IndexerTopK(JaxApiBase): + """JAX callable specialized from top-K input metadata.""" + + def __init__( + self, + sample_input_values: Any, + sample_seq_lens: Any, + top_k: int, + next_n: int = 1, + return_val: bool = True, + num_copy_bits: int = 256, + *, + target_compute_capability: int | None = None, + ) -> None: + self.target_compute_capability = self._resolve_compute_capability( + target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "IndexerTopK", + ) + self.input_desc = self._to_tensor_desc(sample_input_values, "sample_input_values") + self.seq_lens_desc = self._to_tensor_desc(sample_seq_lens, "sample_seq_lens") + self.top_k = int(top_k) + self.next_n = int(next_n) + self.return_val = bool(return_val) + self.num_copy_bits = int(num_copy_bits) + + if self.input_desc.ndim != 2: + raise ValueError(f"input_values must have rank 2, got {self.input_desc.shape}") + if self.seq_lens_desc.ndim != 1: + raise ValueError(f"seq_lens must have rank 1, got {self.seq_lens_desc.shape}") + if self.input_desc.cudnn_dtype not in SUPPORTED_INPUT_DTYPES: + raise ValueError(f"input_values must have dtype float16, bfloat16, or float32, got {self.input_desc.dtype}") + if self.top_k <= 0: + raise ValueError(f"top_k must be positive, got {self.top_k}") + output_shape = (self.input_desc.shape[0], self.top_k) + self.indices_desc = self.input_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=output_shape, + name="indices", + ) + self.values_desc = ( + self.input_desc.compact_like( + cudnn_dtype=self.input_desc.cudnn_dtype, + shape=output_shape, + name="values", + ) + if self.return_val + else None + ) + buffer_count = 2 if self.input_desc.cudnn_dtype == data_type.FLOAT else 1 + self.workspace_desc = self.input_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=(self.input_desc.shape[0], buffer_count, self.input_desc.shape[1]), + name="extra_buffer", + ) + self._op = IndexerTopKOp( + input_values=self.input_desc, + seq_lens=self.seq_lens_desc, + output_indices=self.indices_desc, + output_values=self.values_desc, + workspace=self.workspace_desc, + top_k=self.top_k, + next_n=self.next_n, + return_val=self.return_val, + num_copy_bits=self.num_copy_bits, + target_compute_capability=self.target_compute_capability, + ) + + def check_support(self) -> bool: + self._op.check_support() + workspace_elements = 1 + for extent in self.workspace_desc.shape: + workspace_elements *= extent + if workspace_elements > INT32_MAX: + raise NotImplementedError( + "JAX IndexerTopK does not support the Torch row-chunking fallback when " + f"the workspace exceeds {INT32_MAX} elements (got {workspace_elements})" + ) + return True + + def __call__(self, input_values: Any, seq_lens: Any) -> TupleDict: + self.check_support() + self._check_tensor_signature(input_values, self.input_desc) + self._check_tensor_signature(seq_lens, self.seq_lens_desc) + + output_descs = (self.indices_desc,) if self.values_desc is None else (self.indices_desc, self.values_desc) + results = self._call_kernel( + (input_values, seq_lens), + output_descs=output_descs, + workspace_descs=(self.workspace_desc,), + input_spec=(self._to_tensor_spec(self.input_desc), self._to_tensor_spec(self.seq_lens_desc)), + compile_options=compile_options_for_target(self.target_compute_capability), + ) + indices = results[0] + values = results[1] if self.return_val else None + return TupleDict(indices=indices, values=values) + + def _launch( + self, + inputs: tuple[Any, ...], + outputs: tuple[Any, ...], + workspaces: tuple[Any, ...], + stream: Any, + ) -> None: + import cutlass + + from .indexer_top_k_decode_varlen import IndexerTopKKernelVarlenDecode + + dtype_by_cudnn = { + data_type.HALF: cutlass.Float16, + data_type.BFLOAT16: cutlass.BFloat16, + data_type.FLOAT: cutlass.Float32, + } + input_values, seq_lens = inputs + output_indices = outputs[0] + output_values = outputs[1] if self.return_val else None + (extra_buffer,) = workspaces + resolved = ( + self._op.max_num_cols, + self._op.large_occupancy, + ) + if any(value is None for value in resolved): + raise RuntimeError("IndexerTopK launch configuration was not resolved by check_support()") + max_num_cols, large_occupancy = resolved + kernel = IndexerTopKKernelVarlenDecode( + dtype_by_cudnn[self.input_desc.cudnn_dtype], + max_num_cols, + self.top_k, + self.next_n, + num_copy_bits=self.num_copy_bits, + return_val=self.return_val, + large_occupancy=large_occupancy, + ) + kernel( + input_values, + None, + extra_buffer, + None, + seq_lens, + output_indices, + output_values, + stream, + enable_persistent_dynamic_scheduling=False, + min_blocks_per_mp=1, + ) + + +class _LocalToGlobal(JaxApiBase): + def __init__( + self, + sample_local_indices: Any, + seqlen_k: int, + *, + sample_cu_seqlens_q: Any | None = None, + sample_cu_seqlens_k: Any | None = None, + target_compute_capability: int | None = None, + ) -> None: + if (sample_cu_seqlens_q is None) != (sample_cu_seqlens_k is None): + raise ValueError("Packed local-to-global requires both cumulative sequence tensors") + self.target_compute_capability = self._resolve_compute_capability( + target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "local_to_global", + ) + self.local_desc = self._to_tensor_desc(sample_local_indices, "local_indices") + self.cu_q_desc = None if sample_cu_seqlens_q is None else self._to_tensor_desc(sample_cu_seqlens_q, "cu_seqlens_q") + self.cu_k_desc = None if sample_cu_seqlens_k is None else self._to_tensor_desc(sample_cu_seqlens_k, "cu_seqlens_k") + self.is_varlen = self.cu_q_desc is not None + self.seqlen_k = int(seqlen_k) + self.output_desc = self.local_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=self.local_desc.shape, + name="indices", + ) + + def check_support(self) -> bool: + expected_rank = 2 if self.is_varlen else 3 + if self.local_desc.ndim != expected_rank: + raise ValueError(f"local_indices must have rank {expected_rank}, got {self.local_desc.shape}") + if self.local_desc.cudnn_dtype not in (data_type.INT32, data_type.INT64): + raise ValueError(f"local_indices must have dtype int32 or int64, got {self.local_desc.dtype}") + if any(extent <= 0 for extent in self.local_desc.shape): + raise ValueError(f"local_indices dimensions must be positive, got {self.local_desc.shape}") + if not self.local_desc.is_compact(tuple(reversed(range(expected_rank)))): + raise ValueError("local_indices must be row-major contiguous") + if self.seqlen_k <= 0: + raise ValueError(f"seqlen_k must be positive, got {self.seqlen_k}") + if self.is_varlen: + if self.cu_q_desc is None or self.cu_k_desc is None: + raise RuntimeError("Packed cumulative sequence descriptors were not configured") + for desc, name in ((self.cu_q_desc, "cu_seqlens_q"), (self.cu_k_desc, "cu_seqlens_k")): + if desc.ndim != 1 or desc.cudnn_dtype != data_type.INT32 or not desc.is_compact((0,)): + raise ValueError(f"{name} must be a contiguous rank-1 int32 tensor") + if self.cu_q_desc.shape != self.cu_k_desc.shape: + raise ValueError("cu_seqlens_q and cu_seqlens_k must have the same shape") + if self.cu_q_desc.shape[0] < 2: + raise ValueError("Packed cumulative sequence tensors must contain at least two entries") + return True + + def __call__( + self, + local_indices: Any, + *, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + ) -> TupleDict: + self.check_support() + self._check_tensor_signature(local_indices, self.local_desc) + inputs = [local_indices] + if self.is_varlen: + if cu_seqlens_q is None or cu_seqlens_k is None: + raise ValueError("Packed local-to-global requires both cumulative sequence tensors") + if self.cu_q_desc is None or self.cu_k_desc is None: + raise RuntimeError("Packed cumulative sequence descriptors were not configured") + self._check_tensor_signature(cu_seqlens_q, self.cu_q_desc) + self._check_tensor_signature(cu_seqlens_k, self.cu_k_desc) + inputs.extend((cu_seqlens_q, cu_seqlens_k)) + elif cu_seqlens_q is not None or cu_seqlens_k is not None: + raise ValueError("Fixed local-to-global does not accept cumulative sequence tensors") + + (result,) = self._call_kernel( + tuple(inputs), + output_descs=(self.output_desc,), + compile_options=compile_options_for_target(self.target_compute_capability, "--opt-level 3"), + ) + return TupleDict(indices=result) + + def _launch(self, inputs, outputs, workspaces, stream) -> None: + from cutlass import Int32 + + from .local_to_global_dsl import LocalToGlobalTopK + + local_indices, *optional = inputs + if self.is_varlen: + cu_q, cu_k = optional + else: + cu_q = cu_k = None + (output,) = outputs + if workspaces: + raise RuntimeError("local_to_global does not use workspaces") + LocalToGlobalTopK(is_varlen=self.is_varlen)( + local_indices, + output, + Int32(self.seqlen_k), + cu_q, + cu_k, + stream, + ) + + +class _Compactify(JaxApiBase): + def __init__(self, sample_indices: Any, *, target_compute_capability: int | None = None) -> None: + self.target_compute_capability = self._resolve_compute_capability( + target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "compactify", + ) + self.input_desc = self._to_tensor_desc(sample_indices, "indices") + if self.input_desc.ndim != 2: + raise ValueError(f"Compactify expects flattened rank-2 input, got {self.input_desc.shape}") + rows, cols = self.input_desc.shape + self.rows = rows + self.cols = cols + self.output_desc = self.input_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=(rows, cols), + name="indices", + ) + self.length_desc = self.input_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=(rows,), + name="topk_length", + ) + + def check_support(self) -> bool: + if self.input_desc.cudnn_dtype != data_type.INT32: + raise ValueError(f"indices must have dtype int32, got {self.input_desc.dtype}") + if self.rows <= 0 or self.cols <= 0: + raise ValueError(f"indices dimensions must be positive, got {self.input_desc.shape}") + if not self.input_desc.is_compact((1, 0)): + raise ValueError("indices must be row-major contiguous") + return True + + def __call__(self, indices: Any) -> TupleDict: + self.check_support() + self._check_tensor_signature(indices, self.input_desc) + compact_indices, topk_length = self._call_kernel( + (indices,), + output_descs=(self.output_desc, self.length_desc), + compile_options=compile_options_for_target(self.target_compute_capability, "--opt-level 3"), + ) + return TupleDict(indices=compact_indices, topk_length=topk_length) + + def _launch(self, inputs, outputs, workspaces, stream) -> None: + from cutlass import Int32 + + from .compactify import CompactifyKernel + + (indices,) = inputs + compact_indices, topk_length = outputs + if workspaces: + raise RuntimeError("compactify does not use workspaces") + CompactifyKernel(cols=self.cols)(indices, compact_indices, topk_length, Int32(self.rows), stream) + + +@partial( + jax.jit, + static_argnames=("top_k", "next_n", "return_val", "num_copy_bits", "target_compute_capability"), +) +def indexer_top_k_wrapper( + input_values: Any, + seq_lens: Any, + top_k: int, + next_n: int = 1, + return_val: bool = True, + num_copy_bits: int = 256, + *, + target_compute_capability: int | None = None, +) -> TupleDict: + """Select per-row top-K indices and optional values from JAX. + + Runtime lengths are consumed on device and cannot be inspected while + tracing. Every ``seq_lens[b]`` must satisfy + ``top_k + next_n - 1 <= seq_lens[b] <= input_values.shape[1]``. + """ + + return IndexerTopK( + jax.ShapeDtypeStruct(input_values.shape, input_values.dtype), + jax.ShapeDtypeStruct(seq_lens.shape, seq_lens.dtype), + top_k, + next_n=next_n, + return_val=return_val, + num_copy_bits=num_copy_bits, + target_compute_capability=target_compute_capability, + )(input_values, seq_lens) + + +@partial(jax.jit, static_argnames=("seqlen_k", "target_compute_capability")) +def local_to_global_wrapper( + local_indices: Any, + seqlen_k: int, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + *, + target_compute_capability: int | None = None, +) -> TupleDict: + """Convert local top-K indices to the global flattened KV index space.""" + + sample_cu_q = None if cu_seqlens_q is None else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype) + sample_cu_k = None if cu_seqlens_k is None else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype) + return _LocalToGlobal( + jax.ShapeDtypeStruct(local_indices.shape, local_indices.dtype), + seqlen_k, + sample_cu_seqlens_q=sample_cu_q, + sample_cu_seqlens_k=sample_cu_k, + target_compute_capability=target_compute_capability, + )( + local_indices, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + ) + + +@partial(jax.jit, static_argnames=("target_compute_capability",)) +def compactify_wrapper(indices: Any, *, target_compute_capability: int | None = None) -> TupleDict: + """Pack nonnegative indices to the front of every flattened row.""" + + if indices.ndim not in (2, 3): + raise ValueError(f"indices must have rank 2 or 3, got shape {indices.shape}") + rows = 1 + for extent in indices.shape[:-1]: + rows *= extent + flat_indices = jnp.reshape(indices, (rows, indices.shape[-1])) + return _Compactify( + jax.ShapeDtypeStruct(flat_indices.shape, flat_indices.dtype), + target_compute_capability=target_compute_capability, + )(flat_indices) + + +__all__ = [ + "IndexerTopK", + "compactify_wrapper", + "indexer_top_k_wrapper", + "local_to_global_wrapper", +] diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/local_to_global_dsl.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/local_to_global_dsl.py index bfa61b73f..eaa4ee637 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/local_to_global_dsl.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/local_to_global_dsl.py @@ -9,19 +9,16 @@ from __future__ import annotations -import torch +from typing import TYPE_CHECKING + import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute from cutlass import Int32, Int64, const_expr -from cudnn.deepseek_sparse_attention.utils.compiler import compile_options -from cudnn.deepseek_sparse_attention.utils.runtime import ( - device_major as _get_device_capability, - resolve_stream, -) -from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor as _to_cute_tensor +if TYPE_CHECKING: + import torch class LocalToGlobalTopK: @@ -123,7 +120,11 @@ def is_available() -> bool: # The kernel is arch-agnostic CuTe DSL: thread/block indexing, integer # ALU, and plain global mem load/store — no TMA / TMEM / tcgen05 / wgmma. # SM90+ is enough. - return torch.cuda.is_available() and _get_device_capability() >= 9 + import torch + + from cudnn.deepseek_sparse_attention.utils.runtime import device_major + + return torch.cuda.is_available() and device_major() >= 9 def local_to_global( @@ -133,6 +134,12 @@ def local_to_global( cu_seqlens_k: torch.Tensor | None = None, stream: cuda.CUstream | None = None, ) -> torch.Tensor: + import torch + + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream + from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor as _to_cute_tensor + if not is_available(): raise RuntimeError("CuTe DSL local_to_global requires an SM90+ CUDA device") if not local_indices.is_cuda or not local_indices.is_contiguous(): diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/op.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/op.py new file mode 100644 index 000000000..31c7664e4 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/op.py @@ -0,0 +1,194 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral validation for the DeepSeek indexer top-K operation.""" + +from __future__ import annotations + +from typing import Any + +from ... import data_type +from ..._op import Op +from ..._tensor_desc import TensorDesc + +INT32_MAX = (1 << 31) - 1 +SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) +SUPPORTED_INPUT_DTYPES = (data_type.HALF, data_type.BFLOAT16, data_type.FLOAT) + + +def bucket_num_cols(num_cols: int) -> int: + if num_cols <= 0: + return 1 + return 1 << (num_cols - 1).bit_length() + + +def _dtype_bits(dtype: data_type) -> int: + return 32 if dtype == data_type.FLOAT else 16 + + +def _require_compact(desc: TensorDesc[Any], label: str) -> None: + expected_order = tuple(reversed(range(desc.ndim))) + if not desc.is_compact(expected_order): + raise ValueError(f"{label} must be row-major contiguous, got stride {desc.stride}") + + +def _num_threads_per_cta( + *, + dtype: data_type, + max_num_cols: int, + num_copy_bits: int, + large_occupancy: bool, +) -> int: + if large_occupancy: + return 512 + vec_size = num_copy_bits // _dtype_bits(dtype) + if dtype == data_type.FLOAT: + if max_num_cols >= vec_size * 1024: + return 1024 + if 2048 < max_num_cols < 8192: + return 512 + return 256 + if max_num_cols >= 43008: + return 1024 + if 4096 < max_num_cols < 43008: + return 512 + return 256 + + +class IndexerTopKOp(Op): + """Complete input, output, and workspace signature for indexer top-K.""" + + def __init__( + self, + *, + input_values: TensorDesc[Any], + seq_lens: TensorDesc[Any], + output_indices: TensorDesc[Any], + output_values: TensorDesc[Any] | None, + workspace: TensorDesc[Any], + top_k: int, + next_n: int = 1, + return_val: bool = True, + num_copy_bits: int = 256, + target_compute_capability: int = 90, + ) -> None: + for name, desc in ( + ("input_values", input_values), + ("seq_lens", seq_lens), + ("output_indices", output_indices), + ("workspace", workspace), + ): + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + if output_values is not None and not isinstance(output_values, TensorDesc): + raise TypeError(f"output_values must be a TensorDesc or None, got {type(output_values).__name__}") + + self.input_values = input_values + self.seq_lens = seq_lens + self.output_indices = output_indices + self.output_values = output_values + self.workspace = workspace + self.top_k = int(top_k) + self.next_n = int(next_n) + self.return_val = bool(return_val) + self.num_copy_bits = int(num_copy_bits) + self.target_compute_capability = int(target_compute_capability) + + self.num_rows: int | None = None + self.num_cols: int | None = None + self.batch_size: int | None = None + self.buffer_count: int | None = None + self.dtype_bits: int | None = None + self.max_num_cols: int | None = None + self.large_occupancy: bool | None = None + self.num_threads_per_cta: int | None = None + + def check_support(self) -> bool: + if self.target_compute_capability not in SUPPORTED_COMPUTE_CAPABILITIES: + raise ValueError("target_compute_capability must be one of " f"{SUPPORTED_COMPUTE_CAPABILITIES}, got {self.target_compute_capability}") + if self.input_values.ndim != 2: + raise ValueError(f"input_values must be 2-D, got {self.input_values.shape}") + if self.seq_lens.ndim != 1: + raise ValueError(f"seq_lens must be 1-D, got {self.seq_lens.shape}") + if self.input_values.cudnn_dtype not in SUPPORTED_INPUT_DTYPES: + raise ValueError(f"input_values must have dtype float16, bfloat16, or float32, got {self.input_values.dtype}") + if self.seq_lens.cudnn_dtype != data_type.INT32: + raise ValueError(f"seq_lens must have dtype int32, got {self.seq_lens.dtype}") + _require_compact(self.input_values, "input_values") + _require_compact(self.seq_lens, "seq_lens") + + num_rows, num_cols = self.input_values.shape + (batch_size,) = self.seq_lens.shape + if num_rows <= 0 or num_cols <= 0: + raise ValueError(f"input_values dimensions must be positive, got {(num_rows, num_cols)}") + if batch_size <= 0: + raise ValueError("seq_lens must not be empty") + if self.next_n <= 0: + raise ValueError(f"next_n must be positive, got {self.next_n}") + if num_rows != batch_size * self.next_n: + raise ValueError(f"num_rows ({num_rows}) must equal seq_lens.size * next_n " f"({batch_size} * {self.next_n} = {batch_size * self.next_n})") + if self.top_k <= 0 or self.top_k > min(2048, num_cols): + raise ValueError(f"top_k must be in (0, min(2048, num_cols={num_cols})], got {self.top_k}") + if self.num_copy_bits <= 0 or self.num_copy_bits % 8: + raise ValueError(f"num_copy_bits must be a positive whole-byte width, got {self.num_copy_bits}") + copy_bytes = self.num_copy_bits // 8 + if copy_bytes & (copy_bytes - 1): + raise ValueError(f"num_copy_bits must describe a power-of-two byte alignment, got {self.num_copy_bits}") + + dtype_bits = _dtype_bits(self.input_values.cudnn_dtype) + if self.num_copy_bits % dtype_bits: + raise ValueError(f"num_copy_bits ({self.num_copy_bits}) must be divisible by input dtype width ({dtype_bits})") + + output_shape = (num_rows, self.top_k) + if self.output_indices.shape != output_shape or self.output_indices.cudnn_dtype != data_type.INT32: + raise ValueError(f"output_indices must be int32 with shape {output_shape}, got {self.output_indices.shape}") + _require_compact(self.output_indices, "output_indices") + if self.return_val: + if self.output_values is None: + raise ValueError("output_values is required when return_val=True") + if self.output_values.shape != output_shape or self.output_values.cudnn_dtype != self.input_values.cudnn_dtype: + raise ValueError(f"output_values must have shape {output_shape} and match input_values dtype") + _require_compact(self.output_values, "output_values") + elif self.output_values is not None: + raise ValueError("output_values must be None when return_val=False") + + buffer_count = 2 if self.input_values.cudnn_dtype == data_type.FLOAT else 1 + workspace_shape = (num_rows, buffer_count, num_cols) + if self.workspace.shape != workspace_shape or self.workspace.cudnn_dtype != data_type.INT32: + raise ValueError(f"workspace must be int32 with shape {workspace_shape}, got {self.workspace.shape}") + _require_compact(self.workspace, "workspace") + max_num_cols = bucket_num_cols(num_cols) + large_occupancy = num_rows > 148 + num_threads_per_cta = _num_threads_per_cta( + dtype=self.input_values.cudnn_dtype, + max_num_cols=max_num_cols, + num_copy_bits=self.num_copy_bits, + large_occupancy=large_occupancy, + ) + vector_size = min( + self.top_k, + (self.top_k + num_threads_per_cta - 1) // num_threads_per_cta, + self.num_copy_bits // dtype_bits, + 2, + ) + if self.top_k % vector_size: + raise ValueError(f"top_k ({self.top_k}) must be divisible by the selected output vector width ({vector_size})") + + self.num_rows = num_rows + self.num_cols = num_cols + self.batch_size = batch_size + self.buffer_count = buffer_count + self.dtype_bits = dtype_bits + self.max_num_cols = max_num_cols + self.large_occupancy = large_occupancy + self.num_threads_per_cta = num_threads_per_cta + return True + + +__all__ = [ + "INT32_MAX", + "IndexerTopKOp", + "SUPPORTED_COMPUTE_CAPABILITIES", + "SUPPORTED_INPUT_DTYPES", + "bucket_num_cols", +] diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/__init__.py b/python/cudnn/deepseek_sparse_attention/score_recompute/__init__.py index d5434b6ff..ddaace4f3 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/__init__.py @@ -1,21 +1,38 @@ -from .api import ( - SparseIndexerScoreRecompute, - sparse_indexer_score_recompute_wrapper, - SparseAttnScoreRecompute, - sparse_attn_score_recompute_wrapper, - DenseIndexerScoreRecompute, - dense_indexer_score_recompute_wrapper, - DenseAttnScoreRecompute, - dense_attn_score_recompute_wrapper, -) +"""Lazy Torch API and framework-neutral score-recompute exports. + +JAX APIs are exported through :mod:`cudnn.jax`. +""" -__all__ = [ - "SparseIndexerScoreRecompute", - "sparse_indexer_score_recompute_wrapper", - "SparseAttnScoreRecompute", - "sparse_attn_score_recompute_wrapper", - "DenseIndexerScoreRecompute", - "dense_indexer_score_recompute_wrapper", - "DenseAttnScoreRecompute", - "dense_attn_score_recompute_wrapper", -] +from ..._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "config": ( + "DenseScoreKernelConfig", + "SparseScoreKernelConfig", + "dispatch_sparse_attn_tile_params", + "resolve_dense_score_kernel_config", + "resolve_dense_score_smem_config", + "resolve_sparse_score_kernel_config", + "resolve_sparse_score_smem_config", + ), + "op": ( + "DenseScoreRecomputeOp", + "DenseScoreSm90Config", + "SparseScoreRecomputeOp", + "SparseScoreSm90Config", + ), + "api": ( + "SparseIndexerScoreRecompute", + "sparse_indexer_score_recompute_wrapper", + "SparseAttnScoreRecompute", + "sparse_attn_score_recompute_wrapper", + "DenseIndexerScoreRecompute", + "dense_indexer_score_recompute_wrapper", + "DenseAttnScoreRecompute", + "dense_attn_score_recompute_wrapper", + ), + }, + submodules=("api", "config", "op"), +) diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py b/python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py index 1310cc690..7649f98be 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py @@ -10,7 +10,6 @@ from __future__ import annotations -import math from typing import Optional import torch @@ -21,6 +20,12 @@ from .sparse_score_recompute_sm100 import SparseScoreRecomputeSm100 from .dense_score_recompute_sm100 import DenseScoreRecomputeSm100 +from .config import ( + dispatch_sparse_attn_tile_params as _shared_dispatch_sparse_attn_tile_params, + resolve_dense_score_kernel_config, + resolve_dense_score_smem_config, + resolve_sparse_score_smem_config, +) from cudnn.deepseek_sparse_attention.utils.compiler import compile_options from cudnn.deepseek_sparse_attention.utils.runtime import ( device_major as _get_device_capability, @@ -97,29 +102,14 @@ def _sparse_indexer_score_recompute( with _torch_stream_context(current_stream): out = torch.empty((bs, seqlen_q, topk), dtype=torch.float32, device=device) - # Compute kv_stage and topk_in_smem from SMEM budget (SM100: 228 KB) - SM100_SMEM_BYTES = 228 * 1024 - head_dim_padded = int(math.ceil(head_dim / 16) * 16) - sK_per_stage = n_block_size * head_dim_padded * 2 # BF16 - sQ_size = m_block_size * head_dim_padded * 2 # BF16 - sTopkIdx_bytes = topk * 2 * 4 # double-buffer, INT32 - sPerHead_bytes = m_block_size * 2 * 2 # double-buffer, BF16 - smem_fixed = sPerHead_bytes + 2048 # barriers, sScoreAll, alignment - - topk_in_smem = True - smem_overhead = sTopkIdx_bytes + smem_fixed - kv_stage = min(4, max(1, (SM100_SMEM_BYTES - sQ_size - smem_overhead) // sK_per_stage)) - total_smem_est = sQ_size + sK_per_stage * kv_stage + smem_overhead - if total_smem_est > SM100_SMEM_BYTES: - topk_in_smem = False - smem_overhead = smem_fixed - kv_stage = min(4, max(1, (SM100_SMEM_BYTES - sQ_size - smem_overhead) // sK_per_stage)) - total_smem_est = sQ_size + sK_per_stage * kv_stage + smem_overhead - assert total_smem_est <= SM100_SMEM_BYTES, ( - f"SMEM overflow ({total_smem_est} > {SM100_SMEM_BYTES}) even without sTopkIdx: " - f"topk={topk}, head_dim={head_dim}(padded={head_dim_padded}), " - f"m_block={m_block_size}, n_block={n_block_size}, kv_stage={kv_stage}." - ) + kv_stage, topk_in_smem = resolve_sparse_score_smem_config( + score_type="indexer", + head_dim=head_dim, + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=None, + topk=topk, + ) compute_capability = _get_device_capability() @@ -319,30 +309,14 @@ def _sparse_attn_score_recompute( with _torch_stream_context(current_stream): out = torch.empty((bs, seqlen_q, topk), dtype=torch.float32, device=device) - # Compute kv_stage and topk_in_smem from SMEM budget (SM100: 228 KB) - SM100_SMEM_BYTES = 228 * 1024 - head_dim_padded = int(math.ceil(head_dim / 16) * 16) - k_block_size_eff = k_block_size if k_block_size is not None else head_dim_padded - sK_per_stage = n_block_size * k_block_size_eff * 2 # BF16 - sQ_size = m_block_size * head_dim_padded * 2 # BF16 (Q total unchanged) - sTopkIdx_bytes = topk * 2 * 4 # double-buffer, INT32 - sPerHead_bytes = m_block_size * 2 * 4 # double-buffer, FP32 - smem_fixed = sPerHead_bytes + 2048 # barriers, sScoreAll, alignment - - topk_in_smem = True - smem_overhead = sTopkIdx_bytes + smem_fixed - kv_stage = min(4, max(1, (SM100_SMEM_BYTES - sQ_size - smem_overhead) // sK_per_stage)) - total_smem_est = sQ_size + sK_per_stage * kv_stage + smem_overhead - if total_smem_est > SM100_SMEM_BYTES: - topk_in_smem = False - smem_overhead = smem_fixed - kv_stage = min(4, max(1, (SM100_SMEM_BYTES - sQ_size - smem_overhead) // sK_per_stage)) - total_smem_est = sQ_size + sK_per_stage * kv_stage + smem_overhead - assert total_smem_est <= SM100_SMEM_BYTES, ( - f"SMEM overflow ({total_smem_est} > {SM100_SMEM_BYTES}) even without sTopkIdx: " - f"topk={topk}, head_dim={head_dim}(padded={head_dim_padded}), " - f"m_block={m_block_size}, n_block={n_block_size}, kv_stage={kv_stage}." - ) + kv_stage, topk_in_smem = resolve_sparse_score_smem_config( + score_type="attention", + head_dim=head_dim, + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=k_block_size, + topk=topk, + ) compute_capability = _get_device_capability() @@ -434,37 +408,12 @@ def _dispatch_sparse_attn_tile_params( Rules tuned on B200 via tests/sweep_tile_params.py (--compact / default). """ - m = qhead_per_kv_head - - if compact: - # --- Compact tuned configs --- - # hd=512 - if head_dim == 512 and qhead_per_kv_head == 64: - return m, 64, None # kvs=2, ~789 TFLOPS - if head_dim == 512 and qhead_per_kv_head == 128: - return m, 128, 128 # kvs=2, ~844/1004 TFLOPS - # hd=576 - if head_dim == 576 and qhead_per_kv_head == 64: - return m, 64, None # kvs=2, ~827 TFLOPS - if head_dim == 576 and qhead_per_kv_head == 128: - return m, 64, 192 # kvs=2~3, ~677/1054 TFLOPS - return m, 64, None - - # --- Non-compact tuned configs --- - # hd=512 - if head_dim == 512 and qhead_per_kv_head == 64: - return m, 128, 256 # kvs=2, ~800 TFLOPS - if head_dim == 512 and qhead_per_kv_head == 128 and topk <= 512: - return m, 128, 128 # kvs=2, ~1000 TFLOPS - if head_dim == 512 and qhead_per_kv_head == 128: - return m, 128, 256 # kvs=1, ~910 TFLOPS - # hd=576 - if head_dim == 576 and qhead_per_kv_head == 64: - return m, 128, 192 # kvs=3, ~770 TFLOPS - if head_dim == 576 and qhead_per_kv_head == 128: - return m, 128, 192 # kvs=1, ~830 TFLOPS - - return m, 64, None + return _shared_dispatch_sparse_attn_tile_params( + head_dim, + qhead_per_kv_head, + topk, + compact=compact, + ) def sparse_attn_score_recompute( @@ -550,130 +499,23 @@ def sparse_attn_score_recompute( # Dense backward: full KV via TMA, no topk # ============================================================================= -_SM100_SMEM_BYTES = 228 * 1024 - - -def _select_dense_k_block_size(head_dim_padded, m_block_size, n_block_size, per_head_elem_bytes): - """Auto-select k_block_size so that sQ + sK(1 stage) + overhead fits in SMEM. - - When head_dim_padded is too large (e.g. m_block=128, hd=512 -> sQ=128KB, - sK=128KB > 228KB), we reduce k_block_size to shrink sK per stage. - sQ total = m_block * head_dim_padded * 2 (fixed, since Q has num_k_chunks stages). - - k_block_size must be a multiple of 64 and divide head_dim_padded evenly. - """ - sQ_size = m_block_size * head_dim_padded * 2 - sPerHead_bytes = m_block_size * 2 * per_head_elem_bytes - smem_overhead = sPerHead_bytes + 2048 - - avail_for_sK = _SM100_SMEM_BYTES - sQ_size - smem_overhead - max_kbs = max(64, avail_for_sK // (n_block_size * 2)) - max_kbs = (max_kbs // 64) * 64 - - kbs = head_dim_padded - while kbs > max_kbs: - # Try common divisors of head_dim_padded in descending order - found = False - for candidate in range(kbs - 64, 63, -64): - if candidate > 0 and head_dim_padded % candidate == 0: - kbs = candidate - found = True - break - if not found: - kbs = 64 - break - - assert head_dim_padded % kbs == 0, f"Cannot find valid k_block_size: head_dim_padded={head_dim_padded}, " f"m_block={m_block_size}, n_block={n_block_size}" - return kbs - - -def _compute_dense_kv_stage(head_dim_padded, m_block_size, n_block_size, k_block_size_eff, per_head_elem_bytes): - sK_per_stage = n_block_size * k_block_size_eff * 2 - sQ_size = m_block_size * head_dim_padded * 2 - sPerHead_bytes = m_block_size * 2 * per_head_elem_bytes - smem_overhead = sPerHead_bytes + 2048 - return min(4, max(1, (_SM100_SMEM_BYTES - sQ_size - smem_overhead) // sK_per_stage)) - - -# ---- Dispatch helpers (mirror _dispatch_sparse_attn_tile_params) ------------- - - -def _dense_m_block_with_smem_check(qhead_per_kv_head, head_dim, per_head_elem_bytes): - """Return m=qhpkv*2 if SMEM fits, else fall back to m=qhpkv.""" - head_dim_padded = int(math.ceil(head_dim / 16) * 16) - n = 128 - min_kbs = 64 - smem_overhead = 4096 - - m2 = qhead_per_kv_head * 2 - sQ_m2 = m2 * head_dim_padded * 2 - sK_min = n * min_kbs * 2 - sPerHead_m2 = m2 * 2 * per_head_elem_bytes - if sQ_m2 + sK_min + sPerHead_m2 + smem_overhead <= _SM100_SMEM_BYTES: - return m2 - - return qhead_per_kv_head - def _dispatch_dense_attn_tile_params(head_dim, qhead_per_kv_head): - """Select (m_block_size, n_block_size, k_block_size) for dense attention backward. - - n_block_size is always 128 (required for Ld32x32bOp TMEM path). - m_block_size = qhpkv * 2 (2 q_tokens per tile) when SMEM allows, - falling back to qhpkv when it doesn't (e.g. nh_q=128, hd=512). - - Returns (m_block_size, n_block_size, k_block_size) where k_block_size=None - means no head_dim splitting. - - Rules tuned on B200 via tests/sweep_tile_params.py (dense mode, 2q sweep). - k=64 kvs=4 is consistently best or within 1-2% of best across all sizes. - 2q (m=2*qhpkv) gives 35-46% speedup over 1q for nh_q=64. - """ - m = _dense_m_block_with_smem_check(qhead_per_kv_head, head_dim, per_head_elem_bytes=4) - n = 128 - - # hd=512, nh_q=64: m=128(2q) k=64 kvs=4 ~1600 TFLOPS (16K best, +46% vs 1q) - # hd=512, nh_q=128: m=128(1q) k=64 kvs=4 ~1512 TFLOPS (m=256 doesn't fit SMEM) - # hd=576, nh_q=64: m=128(2q) k=64 kvs=4 ~1601 TFLOPS (16K best, +41% vs 1q) - # hd=576, nh_q=128: m=128(1q) k=64 kvs=4 ~1527 TFLOPS (m=256 doesn't fit SMEM) - if head_dim in (512, 576): - return m, n, 64 - - # Fallback: auto-select from SMEM budget - head_dim_padded = int(math.ceil(head_dim / 16) * 16) - k = _select_dense_k_block_size(head_dim_padded, m, n, per_head_elem_bytes=4) - if k == head_dim_padded: - k = None - return m, n, k + config = resolve_dense_score_kernel_config( + score_type="attention", + head_dim=head_dim, + qhead_per_kv_head=qhead_per_kv_head, + ) + return config.m_block_size, config.n_block_size, config.k_block_size def _dispatch_dense_indexer_tile_params(head_dim, qhead_per_kv_head): - """Select (m_block_size, n_block_size, k_block_size) for dense indexer backward. - - m_block_size = qhpkv * 2 (2 q_tokens per tile) when SMEM allows, - falling back to qhpkv when it doesn't. - - Returns (m_block_size, n_block_size, k_block_size) where k_block_size=None - means no head_dim splitting. - - Rules tuned on B200 via tests/sweep_tile_params.py (dense mode, 2q sweep). - 2q with k=full is consistently best for all tested configs. - nh_q=32: 2q gives ~68% speedup; nh_q=64: 2q gives ~54% speedup. - """ - m = _dense_m_block_with_smem_check(qhead_per_kv_head, head_dim, per_head_elem_bytes=2) - n = 128 - - # hd=128, nh_q=32: m=64(2q) k=full kvs=4 ~541 TFLOPS (16K best, +71% vs 1q) - # hd=128, nh_q=64: m=128(2q) k=full kvs=4 ~871 TFLOPS (16K best, +56% vs 1q) - if head_dim == 128: - return m, n, None - - # Fallback: auto-select from SMEM budget - head_dim_padded = int(math.ceil(head_dim / 16) * 16) - k = _select_dense_k_block_size(head_dim_padded, m, n, per_head_elem_bytes=2) - if k == head_dim_padded: - k = None - return m, n, k + config = resolve_dense_score_kernel_config( + score_type="indexer", + head_dim=head_dim, + qhead_per_kv_head=qhead_per_kv_head, + ) + return config.m_block_size, config.n_block_size, config.k_block_size # ---- Internal dense functions (explicit tile params) ------------------------ @@ -763,10 +605,13 @@ def _dense_indexer_score_recompute( with _torch_stream_context(current_stream): denom_out = torch.empty(denom_shape, dtype=torch.float32, device=device) - head_dim_padded = int(math.ceil(head_dim / 16) * 16) - if k_block_size is None: - k_block_size = _select_dense_k_block_size(head_dim_padded, m_block_size, n_block_size, per_head_elem_bytes=2) - kv_stage = _compute_dense_kv_stage(head_dim_padded, m_block_size, n_block_size, k_block_size, per_head_elem_bytes=2) + k_block_size, kv_stage = resolve_dense_score_smem_config( + score_type="indexer", + head_dim=head_dim, + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=k_block_size, + ) compute_capability = _get_device_capability() if compute_capability < 10: @@ -943,10 +788,13 @@ def _dense_attn_score_recompute( with _torch_stream_context(current_stream): denom_out = torch.empty(denom_shape, dtype=torch.float32, device=device) - head_dim_padded = int(math.ceil(head_dim / 16) * 16) - if k_block_size is None: - k_block_size = _select_dense_k_block_size(head_dim_padded, m_block_size, n_block_size, per_head_elem_bytes=4) - kv_stage = _compute_dense_kv_stage(head_dim_padded, m_block_size, n_block_size, k_block_size, per_head_elem_bytes=4) + k_block_size, kv_stage = resolve_dense_score_smem_config( + score_type="attention", + head_dim=head_dim, + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=k_block_size, + ) compute_capability = _get_device_capability() if compute_capability < 10: diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/api.py b/python/cudnn/deepseek_sparse_attention/score_recompute/api.py index b8d044063..6b31557d4 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/api.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/api.py @@ -17,18 +17,31 @@ from cudnn.api_base import APIBase, TupleDict from . import _interface_sm100 as _iface_sm100 +from .op import DenseScoreRecomputeOp, SparseScoreRecomputeOp # --------------------------------------------------------------------------- # Base helpers # --------------------------------------------------------------------------- -def _check_score_arch(api: APIBase) -> None: - major, _ = torch.cuda.get_device_capability() +def _check_score_arch(api: APIBase) -> int: + major, minor = torch.cuda.get_device_capability() api._runtime_error_if( major != 9 and major < 10, f"{type(api).__name__} requires SM90 or SM100+ compute capability, found SM{major}", ) + return major * 10 + minor + + +def _compact_for_kernel(desc): + if desc is None: + return None + return desc.compact_like( + cudnn_dtype=desc.cudnn_dtype, + shape=desc.shape, + name=desc.name, + init_value=desc.init_value, + ) class _ScoreRecomputeBase(APIBase): @@ -52,36 +65,6 @@ def compile(self) -> None: self._compiled_kernel = True -def _check_sparse_score_shapes( - api: APIBase, - q_desc, - k_desc, - aux_desc, - topk_desc, - out_desc, - topk_length_desc, - aux_name: str, - qhead_per_kv_head: Optional[int], -) -> None: - api._value_error_if(q_desc.ndim != 4, f"Q must be 4-D (B, S_q, H_q, D), got {q_desc.shape}") - api._value_error_if(k_desc.ndim != 3, f"K must be 3-D (B, S_k, D) MQA, got {k_desc.shape}") - api._value_error_if(aux_desc.ndim != 3, f"{aux_name} must be 3-D (B, S_q, H_q), got {aux_desc.shape}") - api._value_error_if(topk_desc.ndim != 3, f"topk_indices must be 3-D (B, S_q, topk), got {topk_desc.shape}") - api._value_error_if(out_desc.ndim != 3, f"out must be 3-D (B, S_q, topk), got {out_desc.shape}") - - b, s_q, h_q, d = q_desc.shape - api._value_error_if(k_desc.shape[0] != b, f"K batch dim {k_desc.shape[0]} must match Q batch dim {b}") - api._value_error_if(k_desc.shape[2] != d, f"K head dim {k_desc.shape[2]} must match Q head dim {d}") - api._value_error_if(aux_desc.shape != (b, s_q, h_q), f"{aux_name} shape must be {(b, s_q, h_q)}, got {aux_desc.shape}") - api._value_error_if(topk_desc.shape[:2] != (b, s_q), f"topk_indices leading dims must be {(b, s_q)}, got {topk_desc.shape[:2]}") - api._value_error_if(out_desc.shape != topk_desc.shape, f"out shape must match topk_indices shape {topk_desc.shape}, got {out_desc.shape}") - if qhead_per_kv_head is not None: - api._value_error_if(qhead_per_kv_head != h_q, f"qhead_per_kv_head must equal H_q ({h_q}) for MQA sparse score, got {qhead_per_kv_head}") - if topk_length_desc is not None: - api._check_dtype(topk_length_desc, torch.int32, name="topk_length") - api._value_error_if(topk_length_desc.shape != (b, s_q), f"topk_length must be shape {(b, s_q)}, got {topk_length_desc.shape}") - - def _check_dense_score_shapes( api: APIBase, q_desc, @@ -154,25 +137,25 @@ def __init__( self.topk_length_desc = self._make_tensor_desc(sample_topk_length, name="sample_topk_length") self.qhead_per_kv_head = qhead_per_kv_head self.topk_indices_global = bool(topk_indices_global) + self._op = None def check_support(self) -> bool: - _check_score_arch(self) - self._check_dtype(self.q_desc, torch.bfloat16, name="Q") - self._check_dtype(self.k_desc, torch.bfloat16, name="K") - self._check_dtype(self.w_desc, torch.bfloat16, name="W") - self._check_dtype(self.topk_desc, torch.int32, name="topk_indices") - self._check_dtype(self.out_desc, torch.float32, name="out") - _check_sparse_score_shapes( - self, - self.q_desc, - self.k_desc, - self.w_desc, - self.topk_desc, - self.out_desc, - self.topk_length_desc, - "W", - self.qhead_per_kv_head, - ) + target_compute_capability = _check_score_arch(self) + if self._op is None: + self._op = SparseScoreRecomputeOp( + q=_compact_for_kernel(self.q_desc), + k=_compact_for_kernel(self.k_desc), + per_head=_compact_for_kernel(self.w_desc), + topk_indices=_compact_for_kernel(self.topk_desc), + output=_compact_for_kernel(self.out_desc), + topk_length=_compact_for_kernel(self.topk_length_desc), + score_type="indexer", + softmax_scale=1.0, + qhead_per_kv_head=self.qhead_per_kv_head, + topk_indices_global=self.topk_indices_global, + target_compute_capability=target_compute_capability, + ) + self._op.check_support() self._is_supported = True return True @@ -269,7 +252,7 @@ def sparse_indexer_score_recompute_wrapper( qhead_per_kv_head=qhead_per_kv_head, topk_indices_global=topk_indices_global, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_SparseIndexerScoreRecomputeObjects[key] = obj @@ -320,25 +303,25 @@ def __init__( self.softmax_scale = float(softmax_scale) self.qhead_per_kv_head = qhead_per_kv_head self.topk_indices_global = bool(topk_indices_global) + self._op = None def check_support(self) -> bool: - _check_score_arch(self) - self._check_dtype(self.q_desc, torch.bfloat16, name="Q") - self._check_dtype(self.k_desc, torch.bfloat16, name="K") - self._check_dtype(self.lse_desc, torch.float32, name="LSE") - self._check_dtype(self.topk_desc, torch.int32, name="topk_indices") - self._check_dtype(self.out_desc, torch.float32, name="out") - _check_sparse_score_shapes( - self, - self.q_desc, - self.k_desc, - self.lse_desc, - self.topk_desc, - self.out_desc, - self.topk_length_desc, - "LSE", - self.qhead_per_kv_head, - ) + target_compute_capability = _check_score_arch(self) + if self._op is None: + self._op = SparseScoreRecomputeOp( + q=_compact_for_kernel(self.q_desc), + k=_compact_for_kernel(self.k_desc), + per_head=_compact_for_kernel(self.lse_desc), + topk_indices=_compact_for_kernel(self.topk_desc), + output=_compact_for_kernel(self.out_desc), + topk_length=_compact_for_kernel(self.topk_length_desc), + score_type="attention", + softmax_scale=self.softmax_scale, + qhead_per_kv_head=self.qhead_per_kv_head, + topk_indices_global=self.topk_indices_global, + target_compute_capability=target_compute_capability, + ) + self._op.check_support() self._is_supported = True return True @@ -442,7 +425,7 @@ def sparse_attn_score_recompute_wrapper( qhead_per_kv_head=qhead_per_kv_head, topk_indices_global=topk_indices_global, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_SparseAttnScoreRecomputeObjects[key] = obj @@ -524,9 +507,28 @@ def __init__( self.sm_scale = float(sm_scale) self.ratio = int(ratio) self.is_thd = bool(is_thd) + self._op = None def check_support(self) -> bool: - _check_score_arch(self) + target_compute_capability = _check_score_arch(self) + if not self.is_thd: + if self._op is None: + self._op = DenseScoreRecomputeOp( + q=_compact_for_kernel(self.q_desc), + k=_compact_for_kernel(self.k_desc), + per_head=_compact_for_kernel(self.w_desc), + output=_compact_for_kernel(self.out_desc), + denominator=_compact_for_kernel(self.denom_desc), + score_type="indexer", + scale=self.sm_scale, + ratio=self.ratio, + qhead_per_kv_head=self.qhead_per_kv_head, + is_thd=False, + target_compute_capability=target_compute_capability, + ) + self._op.check_support() + self._is_supported = True + return True self._check_dtype(self.q_desc, torch.bfloat16, name="Q") self._check_dtype(self.k_desc, torch.bfloat16, name="K") self._check_dtype(self.w_desc, torch.bfloat16, name="W") @@ -683,7 +685,7 @@ def dense_indexer_score_recompute_wrapper( ratio=ratio, is_thd=is_thd, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_DenseIndexerScoreRecomputeObjects[key] = obj @@ -735,9 +737,28 @@ def __init__( self.qhead_per_kv_head = qhead_per_kv_head self.ratio = int(ratio) self.is_thd = bool(is_thd) + self._op = None def check_support(self) -> bool: - _check_score_arch(self) + target_compute_capability = _check_score_arch(self) + if not self.is_thd: + if self._op is None: + self._op = DenseScoreRecomputeOp( + q=_compact_for_kernel(self.q_desc), + k=_compact_for_kernel(self.k_desc), + per_head=_compact_for_kernel(self.lse_desc), + output=_compact_for_kernel(self.out_desc), + denominator=_compact_for_kernel(self.denom_desc), + score_type="attention", + scale=self.softmax_scale, + ratio=self.ratio, + qhead_per_kv_head=self.qhead_per_kv_head, + is_thd=False, + target_compute_capability=target_compute_capability, + ) + self._op.check_support() + self._is_supported = True + return True self._check_dtype(self.q_desc, torch.bfloat16, name="Q") self._check_dtype(self.k_desc, torch.bfloat16, name="K") self._check_dtype(self.lse_desc, torch.float32, name="LSE") @@ -894,7 +915,7 @@ def dense_attn_score_recompute_wrapper( ratio=ratio, is_thd=is_thd, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_DenseAttnScoreRecomputeObjects[key] = obj diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/config.py b/python/cudnn/deepseek_sparse_attention/score_recompute/config.py new file mode 100644 index 000000000..f666747b7 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/config.py @@ -0,0 +1,292 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral SM100 configuration for score-recompute kernels.""" + +from __future__ import annotations + +from dataclasses import dataclass +import math + +_SM100_SMEM_BYTES = 228 * 1024 + + +@dataclass(frozen=True) +class SparseScoreKernelConfig: + """Static launch configuration for ``SparseScoreRecomputeSm100``.""" + + m_block_size: int + n_block_size: int + k_block_size: int | None + kv_stage: int + have_topk_length: bool + topk_in_smem: bool + + +@dataclass(frozen=True) +class DenseScoreKernelConfig: + """Static launch configuration for ``DenseScoreRecomputeSm100``.""" + + m_block_size: int + n_block_size: int + k_block_size: int | None + kv_stage: int + + +def dispatch_sparse_attn_tile_params( + head_dim: int, + qhead_per_kv_head: int, + topk: int, + *, + compact: bool, +) -> tuple[int, int, int | None]: + """Select the tuned SM100 sparse-attention tile.""" + + m = qhead_per_kv_head + if compact: + if head_dim == 512 and qhead_per_kv_head == 64: + return m, 64, None + if head_dim == 512 and qhead_per_kv_head == 128: + return m, 128, 128 + if head_dim == 576 and qhead_per_kv_head == 64: + return m, 64, None + if head_dim == 576 and qhead_per_kv_head == 128: + return m, 64, 192 + return m, 64, None + + if head_dim == 512 and qhead_per_kv_head == 64: + return m, 128, 256 + if head_dim == 512 and qhead_per_kv_head == 128 and topk <= 512: + return m, 128, 128 + if head_dim == 512 and qhead_per_kv_head == 128: + return m, 128, 256 + if head_dim == 576 and qhead_per_kv_head in (64, 128): + return m, 128, 192 + return m, 64, None + + +def resolve_sparse_score_smem_config( + *, + score_type: str, + head_dim: int, + m_block_size: int, + n_block_size: int, + k_block_size: int | None, + topk: int, +) -> tuple[int, bool]: + """Return ``(kv_stage, topk_in_smem)`` for an explicit tile choice.""" + + if score_type not in ("indexer", "attention"): + raise ValueError(f"score_type must be 'indexer' or 'attention', got {score_type!r}") + if head_dim <= 0: + raise ValueError(f"head dimension must be positive, got {head_dim}") + if m_block_size <= 0 or n_block_size <= 0: + raise ValueError(f"m_block_size and n_block_size must be positive, got {m_block_size} and {n_block_size}") + if topk <= 0 or topk % n_block_size: + raise ValueError(f"topk ({topk}) must be positive and a multiple of n_block_size ({n_block_size})") + + head_dim_padded = math.ceil(head_dim / 16) * 16 + k_block_size_eff = head_dim_padded if k_block_size is None else k_block_size + if k_block_size_eff <= 0: + raise ValueError(f"k_block_size must be positive, got {k_block_size_eff}") + if head_dim_padded % k_block_size_eff: + raise ValueError(f"padded head dimension ({head_dim_padded}) must be divisible by k_block_size ({k_block_size_eff})") + if k_block_size_eff % 64: + raise ValueError(f"k_block_size must be a multiple of 64, got {k_block_size_eff}") + + per_head_element_bytes = 2 if score_type == "indexer" else 4 + k_bytes_per_stage = n_block_size * k_block_size_eff * 2 + q_bytes = m_block_size * head_dim_padded * 2 + topk_bytes = topk * 2 * 4 + per_head_bytes = m_block_size * 2 * per_head_element_bytes + fixed_overhead = per_head_bytes + 2048 + + topk_in_smem = True + smem_overhead = topk_bytes + fixed_overhead + kv_stage = min(4, max(1, (_SM100_SMEM_BYTES - q_bytes - smem_overhead) // k_bytes_per_stage)) + total_smem = q_bytes + k_bytes_per_stage * kv_stage + smem_overhead + if total_smem > _SM100_SMEM_BYTES: + topk_in_smem = False + smem_overhead = fixed_overhead + kv_stage = min(4, max(1, (_SM100_SMEM_BYTES - q_bytes - smem_overhead) // k_bytes_per_stage)) + total_smem = q_bytes + k_bytes_per_stage * kv_stage + smem_overhead + if total_smem > _SM100_SMEM_BYTES: + raise ValueError( + "SM100 shared-memory requirement exceeds 228 KiB even with top-K indices in global memory: " + f"head_dim={head_dim}, m_block_size={m_block_size}, " + f"n_block_size={n_block_size}, k_block_size={k_block_size_eff}" + ) + + return kv_stage, topk_in_smem + + +def resolve_sparse_score_kernel_config( + *, + score_type: str, + head_dim: int, + qhead_per_kv_head: int, + topk: int, + have_topk_length: bool, +) -> SparseScoreKernelConfig: + """Resolve the SM100 sparse tile and shared-memory configuration.""" + + if score_type not in ("indexer", "attention"): + raise ValueError(f"score_type must be 'indexer' or 'attention', got {score_type!r}") + if head_dim <= 0 or head_dim % 64: + raise ValueError(f"head dimension must be a positive multiple of 64, got {head_dim}") + if qhead_per_kv_head not in (32, 64, 128): + raise ValueError(f"qhead_per_kv_head must be 32, 64, or 128 for the SM100 sparse score kernel, got {qhead_per_kv_head}") + if topk <= 0: + raise ValueError(f"topk must be positive, got {topk}") + + if score_type == "indexer": + m_block_size, n_block_size, k_block_size = qhead_per_kv_head, 128, None + else: + m_block_size, n_block_size, k_block_size = dispatch_sparse_attn_tile_params( + head_dim, + qhead_per_kv_head, + topk, + compact=have_topk_length, + ) + + if topk % n_block_size: + raise ValueError(f"topk ({topk}) must be a multiple of the selected n_block_size ({n_block_size})") + + kv_stage, topk_in_smem = resolve_sparse_score_smem_config( + score_type=score_type, + head_dim=head_dim, + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=k_block_size, + topk=topk, + ) + return SparseScoreKernelConfig( + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=k_block_size, + kv_stage=kv_stage, + have_topk_length=bool(have_topk_length), + topk_in_smem=topk_in_smem, + ) + + +def _select_dense_k_block_size( + head_dim_padded: int, + m_block_size: int, + n_block_size: int, + per_head_element_bytes: int, +) -> int: + q_bytes = m_block_size * head_dim_padded * 2 + per_head_bytes = m_block_size * 2 * per_head_element_bytes + available_k_bytes = _SM100_SMEM_BYTES - q_bytes - per_head_bytes - 2048 + max_k_block_size = max(64, available_k_bytes // (n_block_size * 2)) + max_k_block_size = (max_k_block_size // 64) * 64 + + for candidate in range(min(head_dim_padded, max_k_block_size), 63, -64): + if head_dim_padded % candidate == 0: + return candidate + raise ValueError( + "No SM100 dense-score K tile divides the padded head dimension: " + f"head_dim_padded={head_dim_padded}, m_block_size={m_block_size}, n_block_size={n_block_size}" + ) + + +def _dense_m_block_size(qhead_per_kv_head: int, head_dim: int, per_head_element_bytes: int) -> int: + head_dim_padded = math.ceil(head_dim / 16) * 16 + n_block_size = 128 + m_block_size = qhead_per_kv_head * 2 + q_bytes = m_block_size * head_dim_padded * 2 + minimum_k_bytes = n_block_size * 64 * 2 + per_head_bytes = m_block_size * 2 * per_head_element_bytes + if q_bytes + minimum_k_bytes + per_head_bytes + 4096 <= _SM100_SMEM_BYTES: + return m_block_size + return qhead_per_kv_head + + +def resolve_dense_score_smem_config( + *, + score_type: str, + head_dim: int, + m_block_size: int, + n_block_size: int, + k_block_size: int | None, +) -> tuple[int, int]: + """Return the effective K tile and stage count for a dense tile choice.""" + + if score_type not in ("indexer", "attention"): + raise ValueError(f"score_type must be 'indexer' or 'attention', got {score_type!r}") + if head_dim <= 0 or m_block_size <= 0 or n_block_size <= 0: + raise ValueError(f"head_dim, m_block_size, and n_block_size must be positive, got {head_dim}, {m_block_size}, and {n_block_size}") + head_dim_padded = math.ceil(head_dim / 16) * 16 + per_head_element_bytes = 2 if score_type == "indexer" else 4 + effective_k_block_size = k_block_size + if effective_k_block_size is None: + effective_k_block_size = _select_dense_k_block_size( + head_dim_padded, + m_block_size, + n_block_size, + per_head_element_bytes, + ) + if effective_k_block_size <= 0 or effective_k_block_size % 64: + raise ValueError(f"k_block_size must be a positive multiple of 64, got {effective_k_block_size}") + if head_dim_padded % effective_k_block_size: + raise ValueError(f"padded head dimension ({head_dim_padded}) must be divisible by k_block_size ({effective_k_block_size})") + + k_bytes_per_stage = n_block_size * effective_k_block_size * 2 + q_bytes = m_block_size * head_dim_padded * 2 + per_head_bytes = m_block_size * 2 * per_head_element_bytes + kv_stage = min(4, max(1, (_SM100_SMEM_BYTES - q_bytes - per_head_bytes - 2048) // k_bytes_per_stage)) + return effective_k_block_size, kv_stage + + +def resolve_dense_score_kernel_config(*, score_type: str, head_dim: int, qhead_per_kv_head: int) -> DenseScoreKernelConfig: + """Resolve the tuned SM100 dense-score tile.""" + + if score_type not in ("indexer", "attention"): + raise ValueError(f"score_type must be 'indexer' or 'attention', got {score_type!r}") + if head_dim <= 0 or head_dim % 64: + raise ValueError(f"head dimension must be a positive multiple of 64, got {head_dim}") + if qhead_per_kv_head not in (32, 64, 128): + raise ValueError(f"qhead_per_kv_head must be 32, 64, or 128 for the SM100 dense score kernel, got {qhead_per_kv_head}") + + per_head_element_bytes = 2 if score_type == "indexer" else 4 + m_block_size = _dense_m_block_size(qhead_per_kv_head, head_dim, per_head_element_bytes) + n_block_size = 128 + head_dim_padded = math.ceil(head_dim / 16) * 16 + + if score_type == "indexer" and head_dim == 128: + k_block_size = head_dim_padded + elif score_type == "attention" and head_dim in (512, 576): + k_block_size = 64 + else: + k_block_size = _select_dense_k_block_size( + head_dim_padded, + m_block_size, + n_block_size, + per_head_element_bytes, + ) + + k_block_size, kv_stage = resolve_dense_score_smem_config( + score_type=score_type, + head_dim=head_dim, + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=k_block_size, + ) + return DenseScoreKernelConfig( + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=None if k_block_size == head_dim_padded else k_block_size, + kv_stage=kv_stage, + ) + + +__all__ = [ + "DenseScoreKernelConfig", + "SparseScoreKernelConfig", + "dispatch_sparse_attn_tile_params", + "resolve_dense_score_kernel_config", + "resolve_dense_score_smem_config", + "resolve_sparse_score_kernel_config", + "resolve_sparse_score_smem_config", +] diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py b/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py new file mode 100644 index 000000000..9c97ec5e8 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py @@ -0,0 +1,714 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX APIs for DSA sparse and dense score recomputation.""" + +from __future__ import annotations + +from functools import partial +from typing import Any, cast + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from .config import DenseScoreKernelConfig, SparseScoreKernelConfig +from .op import ( + DenseScoreRecomputeOp, + DenseScoreSm90Config, + SUPPORTED_COMPUTE_CAPABILITIES, + SparseScoreRecomputeOp, + SparseScoreSm90Config, +) + + +def _compile_options(target_compute_capability: int) -> str: + from ..utils.compiler import compile_options_for_target + + return compile_options_for_target(target_compute_capability) + + +def _optional_desc(api: JaxApiBase, value: Any | None, name: str) -> JaxTensorDesc | None: + return None if value is None else api._to_tensor_desc(value, name) + + +def _check_optional_signature(api: JaxApiBase, value: Any | None, expected: JaxTensorDesc | None, name: str) -> None: + if expected is None: + if value is not None: + raise ValueError(f"{name} was not part of the sample signature") + return + if value is None: + raise ValueError(f"{name} is required by the sample signature") + api._check_tensor_signature(value, expected) + + +class _SparseScoreRecompute(JaxApiBase): + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_per_head: Any, + sample_topk_indices: Any, + *, + score_type: str, + softmax_scale: float, + sample_out: Any | None, + sample_topk_length: Any | None, + qhead_per_kv_head: int | None, + topk_indices_global: bool, + target_compute_capability: int | None, + ) -> None: + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.k_desc = self._to_tensor_desc(sample_k, "sample_k") + self.per_head_desc = self._to_tensor_desc(sample_per_head, "sample_per_head") + self.topk_indices_desc = self._to_tensor_desc(sample_topk_indices, "sample_topk_indices") + self.topk_length_desc = _optional_desc(self, sample_topk_length, "sample_topk_length") + self.out_desc = self._default_output_desc() if sample_out is None else self._to_tensor_desc(sample_out, "sample_out") + self.target_compute_capability = self._resolve_compute_capability( + target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + type(self).__name__, + ) + self._op = SparseScoreRecomputeOp( + q=self.q_desc, + k=self.k_desc, + per_head=self.per_head_desc, + topk_indices=self.topk_indices_desc, + output=self.out_desc, + topk_length=self.topk_length_desc, + score_type=score_type, + softmax_scale=softmax_scale, + qhead_per_kv_head=qhead_per_kv_head, + topk_indices_global=topk_indices_global, + target_compute_capability=self.target_compute_capability, + ) + + def _default_output_desc(self) -> JaxTensorDesc: + if self.topk_indices_desc.ndim != 3: + raise ValueError(f"topk_indices must have rank 3, got shape {self.topk_indices_desc.shape}") + return self.topk_indices_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=self.topk_indices_desc.shape, + name="sample_out", + ) + + def check_support(self) -> bool: + return self._op.check_support() + + def _call(self, q: Any, k: Any, per_head: Any, topk_indices: Any, topk_length: Any | None) -> Any: + self.check_support() + for value, expected in ( + (q, self.q_desc), + (k, self.k_desc), + (per_head, self.per_head_desc), + (topk_indices, self.topk_indices_desc), + ): + self._check_tensor_signature(value, expected) + _check_optional_signature(self, topk_length, self.topk_length_desc, "topk_length") + + if self.target_compute_capability == 90: + kernel_k = jnp.expand_dims(k, axis=2) + kernel_per_head = jnp.transpose(per_head, (0, 2, 1)) + inputs = (q, kernel_k, kernel_per_head, topk_indices) + else: + inputs = (q, k, per_head, topk_indices) + + workspace_descs = () + if topk_length is not None: + inputs += (topk_length,) + elif self.target_compute_capability != 90: + workspace_descs = ( + self.topk_indices_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=(1, 1), + name="topk_length_workspace", + ), + ) + + (output,) = self._call_kernel( + inputs, + output_descs=(self.out_desc,), + workspace_descs=workspace_descs, + compile_options=_compile_options(self.target_compute_capability), + ) + return output + + def _launch( + self, + inputs: tuple[Any, ...], + outputs: tuple[Any, ...], + workspaces: tuple[Any, ...], + stream: Any, + ) -> None: + output = outputs[0] + q, k, per_head, topk_indices, *optional_inputs = inputs + topk_length = optional_inputs[0] if optional_inputs else (None if self.target_compute_capability == 90 else workspaces[0]) + + if self.target_compute_capability == 90: + import cutlass + + from .sparse_score_recompute_sm90 import SparseScoreRecomputeSm90 + + config = cast(SparseScoreSm90Config, self._op.config) + kernel = SparseScoreRecomputeSm90( + cutlass.BFloat16, + head_dim=self.q_desc.shape[-1], + qhead_per_kvhead=cast(int, self._op.qhead_per_kv_head), + tile_m=config.tile_m, + tile_n=config.tile_n, + KV_stage=config.kv_stage, + num_threads=config.num_threads, + swap_AB=True, + topk_max=self.topk_indices_desc.shape[-1], + is_index_scores=self._op.score_type == "indexer", + softmax_scale=self._op.softmax_scale, + has_topk_length=self.topk_length_desc is not None, + num_head_tiles=config.num_head_tiles, + is_sparse=True, + output_log_probs=False, + topk_indices_global=self._op.topk_indices_global, + ) + kernel(q, k, topk_indices, stream, output, per_head, topk_length, None) + return + + import cutlass + + from .sparse_score_recompute_sm100 import SparseScoreRecomputeSm100 + + config = cast(SparseScoreKernelConfig, self._op.config) + kernel = SparseScoreRecomputeSm100( + head_dim=self.q_desc.shape[-1], + qhead_per_kvhead=cast(int, self._op.qhead_per_kv_head), + m_block_size=config.m_block_size, + n_block_size=config.n_block_size, + k_block_size=config.k_block_size, + topk=self.topk_indices_desc.shape[-1], + kv_stage=config.kv_stage, + score_type=self._op.score_type, + have_topk_length=config.have_topk_length, + topk_in_smem=config.topk_in_smem, + topk_indices_global=self._op.topk_indices_global, + ) + kernel(q, k, per_head, topk_indices, output, topk_length, cutlass.Float32(self._op.softmax_scale), stream) + + +class SparseIndexerScoreRecompute(_SparseScoreRecompute): + """JAX callable specialized for sparse indexer score recomputation.""" + + def __init__( + self, + sample_q_indexer: Any, + sample_k_indexer: Any, + sample_weights: Any, + sample_topk_indices: Any, + sample_out: Any | None = None, + sample_topk_length: Any | None = None, + qhead_per_kv_head: int | None = None, + topk_indices_global: bool = False, + target_compute_capability: int | None = None, + ) -> None: + super().__init__( + sample_q_indexer, + sample_k_indexer, + sample_weights, + sample_topk_indices, + score_type="indexer", + softmax_scale=1.0, + sample_out=sample_out, + sample_topk_length=sample_topk_length, + qhead_per_kv_head=qhead_per_kv_head, + topk_indices_global=topk_indices_global, + target_compute_capability=target_compute_capability, + ) + + def __call__( + self, + q_indexer: Any, + k_indexer: Any, + weights: Any, + topk_indices: Any, + topk_length: Any | None = None, + ) -> TupleDict: + return TupleDict(predict=self._call(q_indexer, k_indexer, weights, topk_indices, topk_length)) + + +class SparseAttnScoreRecompute(_SparseScoreRecompute): + """JAX callable specialized for sparse attention score recomputation.""" + + def __init__( + self, + sample_q_attn: Any, + sample_k_attn: Any, + sample_lse: Any, + sample_topk_indices: Any, + softmax_scale: float, + sample_out: Any | None = None, + sample_topk_length: Any | None = None, + qhead_per_kv_head: int | None = None, + topk_indices_global: bool = False, + target_compute_capability: int | None = None, + ) -> None: + super().__init__( + sample_q_attn, + sample_k_attn, + sample_lse, + sample_topk_indices, + score_type="attention", + softmax_scale=softmax_scale, + sample_out=sample_out, + sample_topk_length=sample_topk_length, + qhead_per_kv_head=qhead_per_kv_head, + topk_indices_global=topk_indices_global, + target_compute_capability=target_compute_capability, + ) + + def __call__( + self, + q_attn: Any, + k_attn: Any, + lse: Any, + topk_indices: Any, + topk_length: Any | None = None, + ) -> TupleDict: + return TupleDict(target=self._call(q_attn, k_attn, lse, topk_indices, topk_length)) + + +class _DenseScoreRecompute(JaxApiBase): + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_per_head: Any, + *, + score_type: str, + scale: float, + sample_out: Any | None, + sample_denom_out: Any | None, + qhead_per_kv_head: int | None, + ratio: int, + sample_cu_seqlens_q: Any | None, + sample_cu_seqlens_k: Any | None, + max_seqlen_q: int | None, + max_seqlen_k: int | None, + sample_q_causal_offsets: Any | None, + target_compute_capability: int | None, + ) -> None: + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.k_desc = self._to_tensor_desc(sample_k, "sample_k") + self.per_head_desc = self._to_tensor_desc(sample_per_head, "sample_per_head") + self.cu_seqlens_q_desc = _optional_desc(self, sample_cu_seqlens_q, "sample_cu_seqlens_q") + self.cu_seqlens_k_desc = _optional_desc(self, sample_cu_seqlens_k, "sample_cu_seqlens_k") + self.q_causal_offsets_desc = _optional_desc(self, sample_q_causal_offsets, "sample_q_causal_offsets") + self.is_thd = self.q_desc.ndim == 3 + self.requested_max_seqlen_q = max_seqlen_q + self.requested_max_seqlen_k = max_seqlen_k + + default_out, default_denom = self._default_output_descs(max_seqlen_k) + self.out_desc = default_out if sample_out is None else self._to_tensor_desc(sample_out, "sample_out", init_value=float("-inf")) + self.denom_desc = default_denom if sample_denom_out is None else self._to_tensor_desc(sample_denom_out, "sample_denom_out") + self.target_compute_capability = self._resolve_compute_capability( + target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + type(self).__name__, + ) + self._op = DenseScoreRecomputeOp( + q=self.q_desc, + k=self.k_desc, + per_head=self.per_head_desc, + output=self.out_desc, + denominator=self.denom_desc, + score_type=score_type, + scale=scale, + ratio=ratio, + qhead_per_kv_head=qhead_per_kv_head, + is_thd=self.is_thd, + cu_seqlens_q=self.cu_seqlens_q_desc, + cu_seqlens_k=self.cu_seqlens_k_desc, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + q_causal_offsets=self.q_causal_offsets_desc, + target_compute_capability=self.target_compute_capability, + ) + + def _default_output_descs(self, max_seqlen_k: int | None) -> tuple[JaxTensorDesc, JaxTensorDesc]: + if self.is_thd: + if max_seqlen_k is None: + raise ValueError("max_seqlen_k is required to infer THD dense score outputs") + out_shape = (self.q_desc.shape[0], int(max_seqlen_k)) + denom_shape = (self.q_desc.shape[0],) + elif self.q_desc.ndim == 4 and self.k_desc.ndim == 4: + out_shape = (self.q_desc.shape[0], self.q_desc.shape[1], self.k_desc.shape[1]) + denom_shape = self.q_desc.shape[:2] + else: + raise ValueError(f"Dense score inputs must both use BSHD or THD layout, got Q={self.q_desc.shape}, K={self.k_desc.shape}") + return ( + self.q_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=out_shape, + name="sample_out", + init_value=float("-inf"), + ), + self.q_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=denom_shape, + name="sample_denom_out", + ), + ) + + def check_support(self) -> bool: + return self._op.check_support() + + def _call( + self, + q: Any, + k: Any, + per_head: Any, + cu_seqlens_q: Any | None, + cu_seqlens_k: Any | None, + q_causal_offsets: Any | None, + ) -> TupleDict: + self.check_support() + for value, expected in ((q, self.q_desc), (k, self.k_desc), (per_head, self.per_head_desc)): + self._check_tensor_signature(value, expected) + _check_optional_signature(self, cu_seqlens_q, self.cu_seqlens_q_desc, "cu_seqlens_q") + _check_optional_signature(self, cu_seqlens_k, self.cu_seqlens_k_desc, "cu_seqlens_k") + _check_optional_signature(self, q_causal_offsets, self.q_causal_offsets_desc, "q_causal_offsets") + + if self.target_compute_capability == 90: + kernel_per_head = jnp.transpose(per_head, (0, 2, 1)) + inputs = (q, k, kernel_per_head) + else: + inputs = (q, k, per_head) + if self.is_thd: + inputs += (cu_seqlens_q, cu_seqlens_k) + if q_causal_offsets is not None: + inputs += (q_causal_offsets,) + + output, denominator = self._call_kernel( + inputs, + output_descs=(self.out_desc, self.denom_desc), + compile_options=_compile_options(self.target_compute_capability), + ) + return TupleDict(out=output, denom=denominator) + + def _launch( + self, + inputs: tuple[Any, ...], + outputs: tuple[Any, ...], + workspaces: tuple[Any, ...], + stream: Any, + ) -> None: + del workspaces + q, k, per_head, *optional_inputs = inputs + output, denominator = outputs + + if self.target_compute_capability == 90: + import cutlass + + from .dense_score_recompute_sm90 import DenseScoreRecomputeSm90 + + q_causal_offsets = optional_inputs[0] if optional_inputs else None + config = cast(DenseScoreSm90Config, self._op.config) + kernel = DenseScoreRecomputeSm90( + cutlass.BFloat16, + head_dim=self.q_desc.shape[-1], + qhead_per_kvhead=cast(int, self._op.qhead_per_kv_head), + tile_m=config.tile_m, + tile_n=config.tile_n, + KV_stage=config.kv_stage, + num_threads=config.num_threads, + swap_AB=True, + topk_max=cast(int, self._op.max_seqlen_k), + is_index_scores=self._op.score_type == "indexer", + softmax_scale=self._op.scale, + has_topk_length=False, + num_head_tiles=config.num_head_tiles, + ratio=self._op.ratio, + ) + kernel(q, k, None, stream, output, per_head, None, denominator, q_causal_offsets) + return + + import cutlass + + from .dense_score_recompute_sm100 import DenseScoreRecomputeSm100 + + if self.is_thd: + cu_seqlens_q, cu_seqlens_k, *optional_inputs = optional_inputs + else: + cu_seqlens_q = cu_seqlens_k = None + q_causal_offsets = optional_inputs[0] if optional_inputs else None + config = cast(DenseScoreKernelConfig, self._op.config) + kernel = DenseScoreRecomputeSm100( + head_dim=self.q_desc.shape[-1], + qhead_per_kvhead=cast(int, self._op.qhead_per_kv_head), + m_block_size=config.m_block_size, + n_block_size=config.n_block_size, + k_block_size=config.k_block_size, + kv_stage=config.kv_stage, + score_type=self._op.score_type, + ratio=self._op.ratio, + is_varlen=self.is_thd, + ) + kernel( + q, + k, + per_head, + output, + denominator, + cutlass.Float32(self._op.scale), + cutlass.Int32(cast(int, self._op.max_seqlen_q)), + cutlass.Int32(cast(int, self._op.max_seqlen_k)), + cu_seqlens_q, + cu_seqlens_k, + q_causal_offsets, + stream, + ) + + +class DenseIndexerScoreRecompute(_DenseScoreRecompute): + """JAX callable specialized for dense indexer score recomputation.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_weights: Any, + sample_out: Any | None = None, + sample_denom_out: Any | None = None, + qhead_per_kv_head: int | None = None, + sm_scale: float = 1.0, + ratio: int = 1, + sample_cu_seqlens_q: Any | None = None, + sample_cu_seqlens_k: Any | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + sample_q_causal_offsets: Any | None = None, + target_compute_capability: int | None = None, + ) -> None: + super().__init__( + sample_q, + sample_k, + sample_weights, + score_type="indexer", + scale=sm_scale, + sample_out=sample_out, + sample_denom_out=sample_denom_out, + qhead_per_kv_head=qhead_per_kv_head, + ratio=ratio, + sample_cu_seqlens_q=sample_cu_seqlens_q, + sample_cu_seqlens_k=sample_cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + sample_q_causal_offsets=sample_q_causal_offsets, + target_compute_capability=target_compute_capability, + ) + + def __call__( + self, + q: Any, + k: Any, + weights: Any, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + q_causal_offsets: Any | None = None, + ) -> TupleDict: + return self._call(q, k, weights, cu_seqlens_q, cu_seqlens_k, q_causal_offsets) + + +class DenseAttnScoreRecompute(_DenseScoreRecompute): + """JAX callable specialized for dense attention score recomputation.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_lse: Any, + softmax_scale: float, + sample_out: Any | None = None, + sample_denom_out: Any | None = None, + qhead_per_kv_head: int | None = None, + ratio: int = 1, + sample_cu_seqlens_q: Any | None = None, + sample_cu_seqlens_k: Any | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + sample_q_causal_offsets: Any | None = None, + target_compute_capability: int | None = None, + ) -> None: + super().__init__( + sample_q, + sample_k, + sample_lse, + score_type="attention", + scale=softmax_scale, + sample_out=sample_out, + sample_denom_out=sample_denom_out, + qhead_per_kv_head=qhead_per_kv_head, + ratio=ratio, + sample_cu_seqlens_q=sample_cu_seqlens_q, + sample_cu_seqlens_k=sample_cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + sample_q_causal_offsets=sample_q_causal_offsets, + target_compute_capability=target_compute_capability, + ) + + def __call__( + self, + q: Any, + k: Any, + lse: Any, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + q_causal_offsets: Any | None = None, + ) -> TupleDict: + return self._call(q, k, lse, cu_seqlens_q, cu_seqlens_k, q_causal_offsets) + + +@partial( + jax.jit, + static_argnames=("qhead_per_kv_head", "topk_indices_global", "target_compute_capability"), +) +def sparse_indexer_score_recompute_wrapper( + q_indexer: Any, + k_indexer: Any, + weights: Any, + topk_indices: Any, + qhead_per_kv_head: int | None = None, + topk_length: Any | None = None, + topk_indices_global: bool = False, + target_compute_capability: int | None = None, +) -> TupleDict: + return SparseIndexerScoreRecompute( + jax.ShapeDtypeStruct(q_indexer.shape, q_indexer.dtype), + jax.ShapeDtypeStruct(k_indexer.shape, k_indexer.dtype), + jax.ShapeDtypeStruct(weights.shape, weights.dtype), + jax.ShapeDtypeStruct(topk_indices.shape, topk_indices.dtype), + sample_topk_length=None if topk_length is None else jax.ShapeDtypeStruct(topk_length.shape, topk_length.dtype), + qhead_per_kv_head=qhead_per_kv_head, + topk_indices_global=topk_indices_global, + target_compute_capability=target_compute_capability, + )(q_indexer, k_indexer, weights, topk_indices, topk_length) + + +@partial( + jax.jit, + static_argnames=("softmax_scale", "qhead_per_kv_head", "topk_indices_global", "target_compute_capability"), +) +def sparse_attn_score_recompute_wrapper( + q_attn: Any, + k_attn: Any, + lse: Any, + topk_indices: Any, + softmax_scale: float, + qhead_per_kv_head: int | None = None, + topk_length: Any | None = None, + topk_indices_global: bool = False, + target_compute_capability: int | None = None, +) -> TupleDict: + return SparseAttnScoreRecompute( + jax.ShapeDtypeStruct(q_attn.shape, q_attn.dtype), + jax.ShapeDtypeStruct(k_attn.shape, k_attn.dtype), + jax.ShapeDtypeStruct(lse.shape, lse.dtype), + jax.ShapeDtypeStruct(topk_indices.shape, topk_indices.dtype), + softmax_scale, + sample_topk_length=None if topk_length is None else jax.ShapeDtypeStruct(topk_length.shape, topk_length.dtype), + qhead_per_kv_head=qhead_per_kv_head, + topk_indices_global=topk_indices_global, + target_compute_capability=target_compute_capability, + )(q_attn, k_attn, lse, topk_indices, topk_length) + + +@partial( + jax.jit, + static_argnames=( + "qhead_per_kv_head", + "sm_scale", + "ratio", + "max_seqlen_q", + "max_seqlen_k", + "target_compute_capability", + ), +) +def dense_indexer_score_recompute_wrapper( + q: Any, + k: Any, + weights: Any, + qhead_per_kv_head: int | None = None, + sm_scale: float = 1.0, + ratio: int = 1, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + q_causal_offsets: Any | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + return DenseIndexerScoreRecompute( + jax.ShapeDtypeStruct(q.shape, q.dtype), + jax.ShapeDtypeStruct(k.shape, k.dtype), + jax.ShapeDtypeStruct(weights.shape, weights.dtype), + qhead_per_kv_head=qhead_per_kv_head, + sm_scale=sm_scale, + ratio=ratio, + sample_cu_seqlens_q=None if cu_seqlens_q is None else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype), + sample_cu_seqlens_k=None if cu_seqlens_k is None else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype), + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + sample_q_causal_offsets=None if q_causal_offsets is None else jax.ShapeDtypeStruct(q_causal_offsets.shape, q_causal_offsets.dtype), + target_compute_capability=target_compute_capability, + )(q, k, weights, cu_seqlens_q, cu_seqlens_k, q_causal_offsets) + + +@partial( + jax.jit, + static_argnames=( + "softmax_scale", + "qhead_per_kv_head", + "ratio", + "max_seqlen_q", + "max_seqlen_k", + "target_compute_capability", + ), +) +def dense_attn_score_recompute_wrapper( + q: Any, + k: Any, + lse: Any, + softmax_scale: float, + qhead_per_kv_head: int | None = None, + ratio: int = 1, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + q_causal_offsets: Any | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + return DenseAttnScoreRecompute( + jax.ShapeDtypeStruct(q.shape, q.dtype), + jax.ShapeDtypeStruct(k.shape, k.dtype), + jax.ShapeDtypeStruct(lse.shape, lse.dtype), + softmax_scale, + qhead_per_kv_head=qhead_per_kv_head, + ratio=ratio, + sample_cu_seqlens_q=None if cu_seqlens_q is None else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype), + sample_cu_seqlens_k=None if cu_seqlens_k is None else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype), + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + sample_q_causal_offsets=None if q_causal_offsets is None else jax.ShapeDtypeStruct(q_causal_offsets.shape, q_causal_offsets.dtype), + target_compute_capability=target_compute_capability, + )(q, k, lse, cu_seqlens_q, cu_seqlens_k, q_causal_offsets) + + +__all__ = [ + "DenseAttnScoreRecompute", + "DenseIndexerScoreRecompute", + "SparseAttnScoreRecompute", + "SparseIndexerScoreRecompute", + "dense_attn_score_recompute_wrapper", + "dense_indexer_score_recompute_wrapper", + "sparse_attn_score_recompute_wrapper", + "sparse_indexer_score_recompute_wrapper", +] diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/op.py b/python/cudnn/deepseek_sparse_attention/score_recompute/op.py new file mode 100644 index 000000000..e700f96b5 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/op.py @@ -0,0 +1,427 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral signatures for DSA score-recompute operations.""" + +from __future__ import annotations + +from dataclasses import dataclass +from numbers import Real +from operator import index +from typing import Any + +from ... import data_type +from ..._op import Op +from ..._tensor_desc import TensorDesc +from .config import DenseScoreKernelConfig, SparseScoreKernelConfig, resolve_dense_score_kernel_config, resolve_sparse_score_kernel_config + +SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) + + +@dataclass(frozen=True) +class SparseScoreSm90Config: + tile_m: int + tile_n: int + kv_stage: int + num_threads: int + num_head_tiles: int + + +@dataclass(frozen=True) +class DenseScoreSm90Config: + tile_m: int + tile_n: int + kv_stage: int + num_threads: int + num_head_tiles: int + + +def _require_desc(value: Any, name: str, *, optional: bool = False) -> TensorDesc[Any] | None: + if optional and value is None: + return None + if not isinstance(value, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(value).__name__}") + return value + + +def _require_dtype(desc: TensorDesc[Any], expected: data_type, name: str) -> None: + if desc.cudnn_dtype != expected: + raise ValueError(f"{name} must have dtype {expected}, got {desc.cudnn_dtype}") + + +def _require_rank(desc: TensorDesc[Any], rank: int, name: str) -> None: + if desc.ndim != rank: + raise ValueError(f"{name} must have rank {rank}, got shape {desc.shape}") + + +def _require_compact(desc: TensorDesc[Any], name: str) -> None: + if not desc.is_compact(tuple(reversed(range(desc.ndim)))): + raise ValueError(f"{name} must be row-major compact, got stride {desc.stride}") + + +def _positive_dimensions(named_dimensions: dict[str, int], operation: str) -> None: + invalid = ", ".join(f"{name}={value}" for name, value in named_dimensions.items() if value <= 0) + if invalid: + raise ValueError(f"{operation} dimensions must be positive, got {invalid}") + + +def _normalize_target(target_compute_capability: int) -> int: + if isinstance(target_compute_capability, bool): + raise TypeError("target_compute_capability must be an integer") + try: + target = index(target_compute_capability) + except TypeError as error: + raise TypeError("target_compute_capability must be an integer") from error + if target not in SUPPORTED_COMPUTE_CAPABILITIES: + supported = ", ".join(f"SM{value}" for value in SUPPORTED_COMPUTE_CAPABILITIES) + raise ValueError(f"Unsupported score-recompute target SM{target}; supported targets are {supported}") + return target + + +def _normalize_scale(scale: float, name: str) -> float: + if isinstance(scale, bool) or not isinstance(scale, Real): + raise TypeError(f"{name} must be a real number, got {type(scale).__name__}") + return float(scale) + + +class SparseScoreRecomputeOp(Op): + """Complete sparse score signature and architecture-specific configuration.""" + + def __init__( + self, + *, + q: TensorDesc[Any], + k: TensorDesc[Any], + per_head: TensorDesc[Any], + topk_indices: TensorDesc[Any], + output: TensorDesc[Any], + score_type: str, + softmax_scale: float, + target_compute_capability: int, + topk_length: TensorDesc[Any] | None = None, + qhead_per_kv_head: int | None = None, + topk_indices_global: bool = False, + ) -> None: + self.q = _require_desc(q, "q") + self.k = _require_desc(k, "k") + self.per_head = _require_desc(per_head, "per_head") + self.topk_indices = _require_desc(topk_indices, "topk_indices") + self.output = _require_desc(output, "output") + self.topk_length = _require_desc(topk_length, "topk_length", optional=True) + if score_type not in ("indexer", "attention"): + raise ValueError(f"score_type must be 'indexer' or 'attention', got {score_type!r}") + if not isinstance(topk_indices_global, bool): + raise TypeError(f"topk_indices_global must be a bool, got {type(topk_indices_global).__name__}") + + self.score_type = score_type + self.softmax_scale = _normalize_scale(softmax_scale, "softmax_scale") + self.target_compute_capability = _normalize_target(target_compute_capability) + self.requested_qhead_per_kv_head = qhead_per_kv_head + self.topk_indices_global = topk_indices_global + + self.qhead_per_kv_head: int | None = None + self.config: SparseScoreKernelConfig | SparseScoreSm90Config | None = None + + def check_support(self) -> bool: + self.qhead_per_kv_head = None + self.config = None + + _require_rank(self.q, 4, "Q") + _require_rank(self.k, 3, "K") + _require_rank(self.per_head, 3, "weights" if self.score_type == "indexer" else "LSE") + _require_rank(self.topk_indices, 3, "topk_indices") + _require_rank(self.output, 3, "output") + _require_dtype(self.q, data_type.BFLOAT16, "Q") + _require_dtype(self.k, data_type.BFLOAT16, "K") + _require_dtype( + self.per_head, + data_type.BFLOAT16 if self.score_type == "indexer" else data_type.FLOAT, + "weights" if self.score_type == "indexer" else "LSE", + ) + _require_dtype(self.topk_indices, data_type.INT32, "topk_indices") + _require_dtype(self.output, data_type.FLOAT, "output") + for desc, name in ( + (self.q, "Q"), + (self.k, "K"), + (self.per_head, "weights" if self.score_type == "indexer" else "LSE"), + (self.topk_indices, "topk_indices"), + (self.output, "output"), + ): + _require_compact(desc, name) + + batch, seqlen_q, num_query_heads, head_dim = self.q.shape + k_batch, seqlen_k, k_head_dim = self.k.shape + topk_batch, topk_seqlen_q, topk = self.topk_indices.shape + _positive_dimensions( + { + "B": batch, + "S_q": seqlen_q, + "S_k": seqlen_k, + "H_q": num_query_heads, + "D": head_dim, + "topk": topk, + }, + "Sparse score-recompute", + ) + if k_batch != batch or k_head_dim != head_dim: + raise ValueError(f"K shape must be {(batch, seqlen_k, head_dim)}, got {self.k.shape}") + expected_per_head = (batch, seqlen_q, num_query_heads) + if self.per_head.shape != expected_per_head: + raise ValueError(f"per-head tensor must have shape {expected_per_head}, got {self.per_head.shape}") + if (topk_batch, topk_seqlen_q) != (batch, seqlen_q): + raise ValueError(f"topk_indices leading dimensions must be {(batch, seqlen_q)}, got {self.topk_indices.shape[:2]}") + if self.output.shape != self.topk_indices.shape: + raise ValueError(f"output must have shape {self.topk_indices.shape}, got {self.output.shape}") + + if self.topk_length is not None: + _require_rank(self.topk_length, 2, "topk_length") + _require_dtype(self.topk_length, data_type.INT32, "topk_length") + _require_compact(self.topk_length, "topk_length") + if self.topk_length.shape != (batch, seqlen_q): + raise ValueError(f"topk_length must have shape {(batch, seqlen_q)}, got {self.topk_length.shape}") + + qhead_per_kv_head = num_query_heads if self.requested_qhead_per_kv_head is None else self.requested_qhead_per_kv_head + if isinstance(qhead_per_kv_head, bool): + raise TypeError("qhead_per_kv_head must be an integer") + try: + qhead_per_kv_head = index(qhead_per_kv_head) + except TypeError as error: + raise TypeError("qhead_per_kv_head must be an integer") from error + if qhead_per_kv_head != num_query_heads: + raise ValueError(f"qhead_per_kv_head must equal H_q ({num_query_heads}) for sparse MQA, got {qhead_per_kv_head}") + + if self.target_compute_capability == 90: + if topk % 128: + raise ValueError(f"SM90 sparse score recompute requires topk to be a multiple of 128, got {topk}") + if qhead_per_kv_head <= 1: + raise ValueError("SM90 sparse score recompute requires qhead_per_kv_head > 1") + tile_m = min(qhead_per_kv_head, 64) + if qhead_per_kv_head % tile_m: + raise ValueError(f"qhead_per_kv_head ({qhead_per_kv_head}) must be divisible by tile_m ({tile_m})") + config: SparseScoreKernelConfig | SparseScoreSm90Config = SparseScoreSm90Config( + tile_m=tile_m, + tile_n=64, + kv_stage=2, + num_threads=256, + num_head_tiles=qhead_per_kv_head // tile_m, + ) + else: + config = resolve_sparse_score_kernel_config( + score_type=self.score_type, + head_dim=head_dim, + qhead_per_kv_head=qhead_per_kv_head, + topk=topk, + have_topk_length=self.topk_length is not None, + ) + + self.qhead_per_kv_head = qhead_per_kv_head + self.config = config + return True + + +class DenseScoreRecomputeOp(Op): + """Complete dense score signature and architecture-specific configuration.""" + + def __init__( + self, + *, + q: TensorDesc[Any], + k: TensorDesc[Any], + per_head: TensorDesc[Any], + output: TensorDesc[Any], + denominator: TensorDesc[Any], + score_type: str, + scale: float, + ratio: int, + target_compute_capability: int, + qhead_per_kv_head: int | None = None, + is_thd: bool = False, + cu_seqlens_q: TensorDesc[Any] | None = None, + cu_seqlens_k: TensorDesc[Any] | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + q_causal_offsets: TensorDesc[Any] | None = None, + ) -> None: + self.q = _require_desc(q, "q") + self.k = _require_desc(k, "k") + self.per_head = _require_desc(per_head, "per_head") + self.output = _require_desc(output, "output") + self.denominator = _require_desc(denominator, "denominator") + self.cu_seqlens_q = _require_desc(cu_seqlens_q, "cu_seqlens_q", optional=True) + self.cu_seqlens_k = _require_desc(cu_seqlens_k, "cu_seqlens_k", optional=True) + self.q_causal_offsets = _require_desc(q_causal_offsets, "q_causal_offsets", optional=True) + if score_type not in ("indexer", "attention"): + raise ValueError(f"score_type must be 'indexer' or 'attention', got {score_type!r}") + if not isinstance(is_thd, bool): + raise TypeError(f"is_thd must be a bool, got {type(is_thd).__name__}") + if isinstance(ratio, bool): + raise TypeError("ratio must be an integer") + try: + ratio = index(ratio) + except TypeError as error: + raise TypeError("ratio must be an integer") from error + + self.score_type = score_type + self.scale = _normalize_scale(scale, "scale") + self.ratio = ratio + self.target_compute_capability = _normalize_target(target_compute_capability) + self.requested_qhead_per_kv_head = qhead_per_kv_head + self.is_thd = is_thd + self.requested_max_seqlen_q = max_seqlen_q + self.requested_max_seqlen_k = max_seqlen_k + + self.qhead_per_kv_head: int | None = None + self.max_seqlen_q: int | None = None + self.max_seqlen_k: int | None = None + self.config: DenseScoreKernelConfig | DenseScoreSm90Config | None = None + + @staticmethod + def _positive_static(value: int | None, name: str) -> int: + if value is None or isinstance(value, bool): + raise ValueError(f"{name} must be provided as a positive integer for THD inputs") + try: + value = index(value) + except TypeError as error: + raise TypeError(f"{name} must be an integer") from error + if value <= 0: + raise ValueError(f"{name} must be positive, got {value}") + return value + + def check_support(self) -> bool: + self.qhead_per_kv_head = None + self.max_seqlen_q = None + self.max_seqlen_k = None + self.config = None + if self.ratio < 1: + raise ValueError(f"ratio must be >= 1, got {self.ratio}") + + expected_rank = 3 if self.is_thd else 4 + _require_rank(self.q, expected_rank, "Q") + _require_rank(self.k, expected_rank, "K") + _require_rank(self.per_head, expected_rank - 1, "weights" if self.score_type == "indexer" else "LSE") + _require_rank(self.output, expected_rank - 1, "output") + _require_rank(self.denominator, expected_rank - 2, "denominator") + _require_dtype(self.q, data_type.BFLOAT16, "Q") + _require_dtype(self.k, data_type.BFLOAT16, "K") + _require_dtype( + self.per_head, + data_type.BFLOAT16 if self.score_type == "indexer" else data_type.FLOAT, + "weights" if self.score_type == "indexer" else "LSE", + ) + _require_dtype(self.output, data_type.FLOAT, "output") + _require_dtype(self.denominator, data_type.FLOAT, "denominator") + for desc, name in ( + (self.q, "Q"), + (self.k, "K"), + (self.per_head, "weights" if self.score_type == "indexer" else "LSE"), + (self.output, "output"), + (self.denominator, "denominator"), + ): + _require_compact(desc, name) + + if self.is_thd: + total_q, num_query_heads, head_dim = self.q.shape + total_k, num_kv_heads, k_head_dim = self.k.shape + _positive_dimensions( + {"total_q": total_q, "total_k": total_k, "H_q": num_query_heads, "H_kv": num_kv_heads, "D": head_dim}, + "Dense THD score-recompute", + ) + if self.per_head.shape != (total_q, num_query_heads): + raise ValueError(f"THD per-head tensor must have shape {(total_q, num_query_heads)}, got {self.per_head.shape}") + if self.cu_seqlens_q is None or self.cu_seqlens_k is None: + raise ValueError("THD dense score recompute requires both cu_seqlens_q and cu_seqlens_k") + for desc, name in ((self.cu_seqlens_q, "cu_seqlens_q"), (self.cu_seqlens_k, "cu_seqlens_k")): + _require_rank(desc, 1, name) + _require_dtype(desc, data_type.INT32, name) + _require_compact(desc, name) + if self.cu_seqlens_q.shape != self.cu_seqlens_k.shape: + raise ValueError(f"cu_seqlens_q and cu_seqlens_k shapes must match, got {self.cu_seqlens_q.shape} and {self.cu_seqlens_k.shape}") + batch = self.cu_seqlens_q.shape[0] - 1 + if batch <= 0: + raise ValueError("cumulative sequence-length tensors must contain at least two entries") + max_seqlen_q = self._positive_static(self.requested_max_seqlen_q, "max_seqlen_q") + max_seqlen_k = self._positive_static(self.requested_max_seqlen_k, "max_seqlen_k") + expected_output = (total_q, max_seqlen_k) + expected_denominator = (total_q,) + if self.target_compute_capability == 90: + raise NotImplementedError( + "SM90 THD score recompute uses host-side sequence-length reads and cannot be traced by JAX; use SM100+ or BSHD inputs" + ) + else: + if self.cu_seqlens_q is not None or self.cu_seqlens_k is not None: + raise ValueError("BSHD dense score recompute does not accept cu_seqlens_q or cu_seqlens_k") + batch, seqlen_q, num_query_heads, head_dim = self.q.shape + k_batch, seqlen_k, num_kv_heads, k_head_dim = self.k.shape + _positive_dimensions( + {"B": batch, "S_q": seqlen_q, "S_k": seqlen_k, "H_q": num_query_heads, "H_kv": num_kv_heads, "D": head_dim}, + "Dense BSHD score-recompute", + ) + if k_batch != batch: + raise ValueError(f"K batch dimension must be {batch}, got {k_batch}") + if self.per_head.shape != (batch, seqlen_q, num_query_heads): + raise ValueError(f"BSHD per-head tensor must have shape {(batch, seqlen_q, num_query_heads)}, got {self.per_head.shape}") + max_seqlen_q, max_seqlen_k = seqlen_q, seqlen_k + expected_output = (batch, seqlen_q, seqlen_k) + expected_denominator = (batch, seqlen_q) + + if k_head_dim != head_dim: + raise ValueError(f"K head dimension must be {head_dim}, got {k_head_dim}") + if num_kv_heads != 1: + raise ValueError(f"Dense score recompute currently requires MQA with H_kv=1, got {num_kv_heads}") + if self.output.shape != expected_output: + raise ValueError(f"output must have shape {expected_output}, got {self.output.shape}") + if self.denominator.shape != expected_denominator: + raise ValueError(f"denominator must have shape {expected_denominator}, got {self.denominator.shape}") + + if self.q_causal_offsets is not None: + _require_rank(self.q_causal_offsets, 1, "q_causal_offsets") + _require_dtype(self.q_causal_offsets, data_type.INT32, "q_causal_offsets") + _require_compact(self.q_causal_offsets, "q_causal_offsets") + if self.q_causal_offsets.shape != (batch,): + raise ValueError(f"q_causal_offsets must have shape {(batch,)}, got {self.q_causal_offsets.shape}") + + inferred_qhead_per_kv_head = num_query_heads // num_kv_heads + qhead_per_kv_head = inferred_qhead_per_kv_head if self.requested_qhead_per_kv_head is None else self.requested_qhead_per_kv_head + if isinstance(qhead_per_kv_head, bool): + raise TypeError("qhead_per_kv_head must be an integer") + try: + qhead_per_kv_head = index(qhead_per_kv_head) + except TypeError as error: + raise TypeError("qhead_per_kv_head must be an integer") from error + if qhead_per_kv_head != inferred_qhead_per_kv_head: + raise ValueError(f"qhead_per_kv_head must equal H_q / H_kv ({inferred_qhead_per_kv_head}), got {qhead_per_kv_head}") + + if self.target_compute_capability == 90: + if qhead_per_kv_head <= 1: + raise ValueError("SM90 dense score recompute requires qhead_per_kv_head > 1") + tile_m = min(qhead_per_kv_head, 64) + if qhead_per_kv_head % tile_m: + raise ValueError(f"qhead_per_kv_head ({qhead_per_kv_head}) must be divisible by tile_m ({tile_m})") + config: DenseScoreKernelConfig | DenseScoreSm90Config = DenseScoreSm90Config( + tile_m=tile_m, + tile_n=64, + kv_stage=2, + num_threads=384, + num_head_tiles=qhead_per_kv_head // tile_m, + ) + else: + config = resolve_dense_score_kernel_config( + score_type=self.score_type, + head_dim=head_dim, + qhead_per_kv_head=qhead_per_kv_head, + ) + + self.qhead_per_kv_head = qhead_per_kv_head + self.max_seqlen_q = max_seqlen_q + self.max_seqlen_k = max_seqlen_k + self.config = config + return True + + +__all__ = [ + "DenseScoreRecomputeOp", + "DenseScoreSm90Config", + "SUPPORTED_COMPUTE_CAPABILITIES", + "SparseScoreRecomputeOp", + "SparseScoreSm90Config", +] diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/__init__.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/__init__.py index 039c5be6a..ebd33b73b 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/__init__.py @@ -1,3 +1,20 @@ -from .api import SparseAttentionBackward, sparse_attention_backward_wrapper +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT -__all__ = ["SparseAttentionBackward", "sparse_attention_backward_wrapper"] +"""Lazy Torch API and framework-neutral sparse-attention backward exports.""" + +from ..._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ("SparseAttentionBackwardOp",), + "dsa_bwd_sm90": ("FlashAttentionDSABackwardSm90",), + "dsa_bwd_sm100": ("FlashAttentionDSABackwardSm100",), + "api": ( + "SparseAttentionBackward", + "sparse_attention_backward_wrapper", + ), + }, + submodules=("api", "dsa_bwd_sm90", "dsa_bwd_sm100", "op"), +) diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py index 3030f8198..db3a9603f 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py @@ -59,7 +59,8 @@ def flash_attn_bwd_sm100( head_dim_v = 512 if head_dim == 576 else head_dim device = q.device - assert q.dtype in [torch.float16, torch.bfloat16] + # The SM100 kernel currently fixes element storage to BFloat16. + assert q.dtype == torch.bfloat16 assert q.dtype == kv.dtype == out.dtype == dout.dtype assert lse.dtype == torch.float32 assert attn_sink.dtype == torch.float32 @@ -99,7 +100,7 @@ def flash_attn_bwd_sm100( # Allocate workspace tensors acc_dtype = cutlass.Float32 - ws_lse_odo_shape = FlashAttentionDSABackwardSm100._get_workspace_size_LSE_OdO( + ws_lse_odo_shape = FlashAttentionDSABackwardSm100.get_workspace_size_lse_odo( total_S_q, head_dim, num_head, @@ -112,7 +113,7 @@ def flash_attn_bwd_sm100( device=device, ) - ws_dkv_shape = FlashAttentionDSABackwardSm100._get_workspace_size_dKV( + ws_dkv_shape = FlashAttentionDSABackwardSm100.get_workspace_size_dkv( total_S_kv, head_dim, batch_size, diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py index 7ce340394..7805d71ad 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py @@ -12,9 +12,22 @@ import torch import cuda.bindings.driver as cuda +from cudnn import data_type from cudnn.api_base import APIBase, TupleDict from . import _interface_sm100 as _iface_sm100 +from .op import SparseAttentionBackwardOp + + +def _compact_for_kernel(desc): + """Describe tensors that the Torch backends normalize before lowering.""" + + return desc.compact_like( + cudnn_dtype=desc.cudnn_dtype, + shape=desc.shape, + name=desc.name, + init_value=desc.init_value, + ) class SparseAttentionBackward(APIBase): @@ -42,33 +55,76 @@ def __init__( self.attn_sink_desc = self._make_tensor_desc(sample_attn_sink, name="sample_attn_sink") self.topk_idxs_desc = self._make_tensor_desc(sample_topk_idxs, name="sample_topk_idxs") self.topk_length_desc = self._make_tensor_desc(sample_topk_length, name="sample_topk_length") - self.block_tile = int(block_tile) + self.dq_desc = ( + self._make_tensor_desc(sample_dq, name="sample_dq") + if sample_dq is not None + else self.q_desc.compact_like( + cudnn_dtype=self.q_desc.cudnn_dtype, + shape=self.q_desc.shape, + name="sample_dq", + ) + ) + self.dkv_desc = ( + self._make_tensor_desc(sample_dkv, name="sample_dkv") + if sample_dkv is not None + else self.kv_desc.compact_like( + cudnn_dtype=self.kv_desc.cudnn_dtype, + shape=self.kv_desc.shape, + name="sample_dkv", + ) + ) + self.d_sink_desc = self.attn_sink_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=self.attn_sink_desc.shape, + name="sample_d_sink", + ) + self._op = SparseAttentionBackwardOp( + q=_compact_for_kernel(self.q_desc), + kv=_compact_for_kernel(self.kv_desc), + output=_compact_for_kernel(self.out_desc), + doutput=_compact_for_kernel(self.dout_desc), + lse=_compact_for_kernel(self.lse_desc), + attn_sink=self.attn_sink_desc, + topk_idxs=self.topk_idxs_desc, + topk_length=self.topk_length_desc, + dq=self.dq_desc, + dkv=self.dkv_desc, + d_sink=self.d_sink_desc, + softmax_scale=softmax_scale, + block_tile=block_tile, + ) + self.block_tile = self._op.block_tile self.softmax_scale = softmax_scale def check_support(self) -> bool: - major, _ = torch.cuda.get_device_capability() + major, minor = torch.cuda.get_device_capability(self.q_desc.device) + compute_capability = major * 10 + minor self._runtime_error_if( - major < 9, - f"SparseAttentionBackward requires SM90+, found SM{major}", - ) - self._value_error_if( - self.q_desc.ndim != 3, - f"Q must be 3-D (total_S_q, H, D), got {self.q_desc.shape}", + compute_capability < 90, + f"SparseAttentionBackward requires SM90+, found SM{compute_capability}", ) - self._value_error_if( - self.kv_desc.ndim != 2, - f"KV must be 2-D (total_S_kv, D), got {self.kv_desc.shape}", - ) - self._check_dtype(self.q_desc, [torch.float16, torch.bfloat16], name="Q") - self._check_dtype( - self.kv_desc, - self.q_desc.dtype, - name="KV", - extra_error_msg="KV must have same dtype as Q", - ) - self._check_dtype(self.lse_desc, torch.float32, name="LSE") - self._check_dtype(self.attn_sink_desc, torch.float32, name="attn_sink") - self._check_dtype(self.topk_idxs_desc, torch.int32, name="topk_idxs") + self._op.check_support() + if compute_capability >= 100 and self.q_desc.cudnn_dtype != data_type.BFLOAT16: + raise ValueError("SparseAttentionBackward on SM100+ currently requires bfloat16 inputs") + + devices = { + desc.device + for desc in ( + self.q_desc, + self.kv_desc, + self.out_desc, + self.dout_desc, + self.lse_desc, + self.attn_sink_desc, + self.topk_idxs_desc, + self.topk_length_desc, + self.dq_desc, + self.dkv_desc, + self.d_sink_desc, + ) + if desc is not None + } + self._value_error_if(len(devices) != 1, f"All tensors must be on the same device, got {sorted(map(str, devices))}") self._is_supported = True return True @@ -94,8 +150,8 @@ def execute( softmax_scale: Optional[float] = None, current_stream: Optional[cuda.CUstream] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - major, _ = torch.cuda.get_device_capability() - scale = self.softmax_scale if softmax_scale is None else softmax_scale + major, _ = torch.cuda.get_device_capability(q.device) + scale = self._op.softmax_scale if softmax_scale is None else softmax_scale if major == 9: from . import _interface_sm90 as _iface_sm90 @@ -180,7 +236,7 @@ def sparse_attention_backward_wrapper( softmax_scale=softmax_scale, block_tile=block_tile, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_SparseAttentionBackwardObjects[key] = obj diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py index fd4c69cc6..44dd43119 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py @@ -148,7 +148,7 @@ def _setup_attributes(self): self.compute_tmastore_dQ_stage = 1 @staticmethod - def _get_workspace_size_LSE_OdO(q: int, d: int, h: int, b: int, acc_dtype: Type[cutlass.Numeric]): + def get_workspace_size_lse_odo(q: int, d: int, h: int, b: int, acc_dtype: Type[cutlass.Numeric]): # q is total seqlen, b=1 d = (d + 7) // 8 * 8 # round up to 8 q = (q + 7) // 8 * 8 # round up to 8 @@ -161,13 +161,17 @@ def _get_workspace_size_LSE_OdO(q: int, d: int, h: int, b: int, acc_dtype: Type[ return (b, h, q, workspace_bytes) @staticmethod - def _get_workspace_size_dKV(k: int, d: int, b: int, acc_dtype: Type[cutlass.Numeric]): + def get_workspace_size_dkv(k: int, d: int, b: int, acc_dtype: Type[cutlass.Numeric]): d = (d + 7) // 8 * 8 # round up to 8 k = (k + 7) // 8 * 8 # round up to 8 # FP32 versions of dKV workspace_bytes = d * acc_dtype.width // 8 return (b, 1, k, workspace_bytes) + # Compatibility aliases for the existing Torch interface. + _get_workspace_size_LSE_OdO = get_workspace_size_lse_odo + _get_workspace_size_dKV = get_workspace_size_dkv + def get_workspace_tensor( self, problem_shape: Tuple[Int32, Int32, Int32, Tuple[Int32, Int32]], @@ -563,25 +567,24 @@ class SharedStorage: stream=stream, ) - if cutlass.const_expr(self.same_hdim_kv): - dSink_grid = ( - cute.ceil_div(problem_shape[0], self.dSink_block_q), - problem_shape[3][0], - problem_shape[3][1], - ) - self.sum_dSink( - sum_OdO, - scaled_LSE, - mAttnSink, - mdSink, - problem_shape, - ).launch( - grid=dSink_grid, - block=[self.dSink_num_threads, 1, 1], - cluster=[1, 1, 1], - stream=stream, - min_blocks_per_mp=1, - ) + dSink_grid = ( + cute.ceil_div(problem_shape[0], self.dSink_block_q), + problem_shape[3][0], + problem_shape[3][1], + ) + self.sum_dSink( + sum_OdO, + scaled_LSE, + mAttnSink, + mdSink, + problem_shape, + ).launch( + grid=dSink_grid, + block=[self.dSink_num_threads, 1, 1], + cluster=[1, 1, 1], + stream=stream, + min_blocks_per_mp=1, + ) @cute.kernel def convert( @@ -665,21 +668,18 @@ def sum_OdO( lse_bhq = lse[bidy, (idx_q + offset, bidz)] sum_OdO_bhq = sum_OdO_scale * acc - if cutlass.const_expr(self.same_hdim_kv): - attn_sink_bh = attn_sink[bidy, (0, bidz)] + attn_sink_bh = attn_sink[bidy, (0, bidz)] - log2_e = -lse_scale - lse_log2 = lse_bhq * log2_e - sink_log2 = attn_sink_bh * log2_e - lse_max_log2 = cute.arch.fmax(lse_log2, sink_log2) - sum_exp2 = Float32(cute.math.exp2(lse_log2 - lse_max_log2) + cute.math.exp2(sink_log2 - lse_max_log2)) - lse_with_sink_log2 = lse_max_log2 + cute.math.log2(sum_exp2) - scaled_lse_bhq = -lse_with_sink_log2 + log2_e = -lse_scale + lse_log2 = lse_bhq * log2_e + sink_log2 = attn_sink_bh * log2_e + lse_max_log2 = cute.arch.fmax(lse_log2, sink_log2) + sum_exp2 = Float32(cute.math.exp2(lse_log2 - lse_max_log2) + cute.math.exp2(sink_log2 - lse_max_log2)) + lse_with_sink_log2 = lse_max_log2 + cute.math.log2(sum_exp2) + scaled_lse_bhq = -lse_with_sink_log2 - if lse_bhq == Float32(float("inf")): - scaled_lse_bhq = Float32(float("-inf")) - else: - scaled_lse_bhq = lse_scale * lse_bhq + if lse_bhq == Float32(float("inf")): + scaled_lse_bhq = Float32(float("-inf")) sum_OdO[bidy, (idx_q + offset, bidz)] = sum_OdO_bhq scaled_lse[bidy, (idx_q + offset, bidz)] = scaled_lse_bhq @@ -865,10 +865,14 @@ def bwd( if cutlass.const_expr(mTopkLength is not None): topk = mTopkLength[token_idx] + topk = cutlass.max(Int32(0), cutlass.min(topk, Int32(self.max_topk))) else: topk = mTopkIdxs.shape[0] - tile_count = cute.ceil_div(topk, self.block_tile) + # Keep every warp role on the same balanced one-tile pipeline when a + # query has no selected KV rows. The partial-tile predicates below + # turn that synthetic tile into all zeros. + tile_count = cutlass.max(Int32(1), cute.ceil_div(topk, self.block_tile)) if warp_idx == self.load_warp_id: cute.arch.setmaxregister_decrease(self.num_regs_load) @@ -1296,7 +1300,7 @@ def load_KV( load_mma_K_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.load_mma_K_stage) tile_index = tile_count - 1 - full_tiles = (topk % self.block_tile) == 0 + full_tiles = (topk > 0) and ((topk % self.block_tile) == 0) rows_per_warp = self.block_tile // self.num_load_KV_warps rTopkIdx = cute.make_rmem_tensor((rows_per_warp,), cutlass.Int32) @@ -2127,7 +2131,7 @@ def reduce_dKV( rTopkIdx = cute.make_rmem_tensor((8,), cutlass.Int32) if cutlass.const_expr(not self.same_hdim_kv): rTopkIdx_64 = cute.make_rmem_tensor((8,), cutlass.Int32) - full_tiles = (topk % self.block_tile) == 0 + full_tiles = (topk > 0) and ((topk % self.block_tile) == 0) while tile_index >= 0: # Preload topk indices into rmem (shared across all 4 store_dKV calls) for i in cutlass.range_constexpr(8): diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm90.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm90.py index 0b5aef54d..aad8bb9bc 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm90.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm90.py @@ -1174,9 +1174,12 @@ def mma_wg0( if const_expr(self.have_topk_length): topK = mTopkLength[batch_idx, seq_idx] + topK = cutlass.max(Int32(0), cutlass.min(topK, Int32(self.max_topk))) else: topK = self.max_topk - n_block_max = (topK + self.tile_n - 1) // self.tile_n + # A zero-length row still executes one fully predicated tile so + # both warp groups reach the same barriers and write zero dQ. + n_block_max = cutlass.max(Int32(1), (topK + self.tile_n - 1) // self.tile_n) topk_tail_rows = topK - (n_block_max - 1) * self.tile_n n_block = n_block_max - 1 @@ -1691,9 +1694,10 @@ def mma_wg1( if const_expr(self.have_topk_length): topK = mTopkLength[batch_idx, seq_idx] + topK = cutlass.max(Int32(0), cutlass.min(topK, Int32(self.max_topk))) else: topK = self.max_topk - n_block_max = (topK + self.tile_n - 1) // self.tile_n + n_block_max = cutlass.max(Int32(1), (topK + self.tile_n - 1) // self.tile_n) # scatter AtomicAdd — flat gmem dKVaccum for per-row addressing mdKVaccum_cur = mdKVaccum[None, head_idx_kv, batch_idx] diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py new file mode 100644 index 000000000..e43d3adf8 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py @@ -0,0 +1,457 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for DeepSeek sparse-attention backward on SM90 and SM100+.""" + +from __future__ import annotations + +from typing import Any, Optional + +import jax + +from ... import data_type +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..utils.compiler import compile_options_for_target +from .op import BLOCK_TILE, SparseAttentionBackwardOp + +_SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) +_SUPPORTED_COMPUTE_CAPABILITY_FAMILIES = (90, 100) + + +class SparseAttentionBackward(JaxApiBase): + """JAX callable specialized from sparse-attention tensor metadata.""" + + def __init__( + self, + sample_q: Any, + sample_kv: Any, + sample_out: Any, + sample_dout: Any, + sample_lse: Any, + sample_attn_sink: Any, + sample_topk_idxs: Any, + sample_dq: Any | None = None, + sample_dkv: Any | None = None, + sample_d_sink: Any | None = None, + sample_topk_length: Any | None = None, + softmax_scale: Optional[float] = None, + block_tile: int = BLOCK_TILE, + target_compute_capability: int | None = None, + ) -> None: + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.kv_desc = self._to_tensor_desc(sample_kv, "sample_kv") + self.out_desc = self._to_tensor_desc(sample_out, "sample_out") + self.dout_desc = self._to_tensor_desc(sample_dout, "sample_dout") + self.lse_desc = self._to_tensor_desc(sample_lse, "sample_lse") + self.attn_sink_desc = self._to_tensor_desc(sample_attn_sink, "sample_attn_sink") + self.topk_idxs_desc = self._to_tensor_desc(sample_topk_idxs, "sample_topk_idxs") + self.topk_length_desc = None if sample_topk_length is None else self._to_tensor_desc(sample_topk_length, "sample_topk_length") + + self.dq_desc = self._output_desc( + sample_dq, + source=self.q_desc, + cudnn_dtype=self.q_desc.cudnn_dtype, + shape=self.q_desc.shape, + name="sample_dq", + ) + self.dkv_desc = self._output_desc( + sample_dkv, + source=self.kv_desc, + cudnn_dtype=self.kv_desc.cudnn_dtype, + shape=self.kv_desc.shape, + name="sample_dkv", + ) + self.d_sink_desc = self._output_desc( + sample_d_sink, + source=self.attn_sink_desc, + cudnn_dtype=data_type.FLOAT, + shape=self.attn_sink_desc.shape, + name="sample_d_sink", + init_value=0.0, + ) + self._op = SparseAttentionBackwardOp( + q=self.q_desc, + kv=self.kv_desc, + output=self.out_desc, + doutput=self.dout_desc, + lse=self.lse_desc, + attn_sink=self.attn_sink_desc, + topk_idxs=self.topk_idxs_desc, + topk_length=self.topk_length_desc, + dq=self.dq_desc, + dkv=self.dkv_desc, + d_sink=self.d_sink_desc, + softmax_scale=softmax_scale, + block_tile=block_tile, + ) + self.target_compute_capability = target_compute_capability + self.compute_capability: int | None = None + + @staticmethod + def _output_desc( + sample: Any | None, + *, + source: JaxTensorDesc, + cudnn_dtype: data_type, + shape: tuple[int, ...], + name: str, + init_value: bool | int | float | None = None, + ) -> JaxTensorDesc: + if sample is not None: + desc = JaxApiBase._to_tensor_desc(sample, name, init_value=init_value) + else: + desc = source.compact_like( + cudnn_dtype=cudnn_dtype, + shape=shape, + name=name, + init_value=init_value, + ) + return desc + + def check_support(self) -> bool: + self._op.check_support() + compute_capability = self._resolve_compute_capability( + self.target_compute_capability, + _SUPPORTED_COMPUTE_CAPABILITIES, + "SparseAttentionBackward", + ) + family = self._compute_capability_family(compute_capability, _SUPPORTED_COMPUTE_CAPABILITY_FAMILIES) + if family == 100 and self.q_desc.cudnn_dtype != data_type.BFLOAT16: + raise ValueError("SparseAttentionBackward on SM100+ currently requires bfloat16 inputs") + self.compute_capability = compute_capability + return True + + def __call__( + self, + q: Any, + kv: Any, + out: Any, + dout: Any, + lse: Any, + attn_sink: Any, + topk_idxs: Any, + topk_length: Any | None = None, + ) -> TupleDict: + self.check_support() + for value, expected in ( + (q, self.q_desc), + (kv, self.kv_desc), + (out, self.out_desc), + (dout, self.dout_desc), + (lse, self.lse_desc), + (attn_sink, self.attn_sink_desc), + (topk_idxs, self.topk_idxs_desc), + ): + self._check_tensor_signature(value, expected) + if (topk_length is None) != (self.topk_length_desc is None): + raise ValueError("topk_length presence must match sample_topk_length") + if topk_length is not None: + self._check_tensor_signature(topk_length, self.topk_length_desc) + + inputs = (q, kv, out, dout, lse, attn_sink, topk_idxs) + input_specs = self._input_specs() + if topk_length is not None: + inputs += (topk_length,) + input_specs += (None,) + + dq, dkv, d_sink = self._call_kernel( + inputs, + output_descs=(self.dq_desc, self.dkv_desc, self.d_sink_desc), + workspace_descs=self._workspace_descs(), + input_spec=input_specs, + output_spec=( + self._to_tensor_spec(self.dq_desc, divisibility=(None, None, self._op.head_dim)), + self._to_tensor_spec(self.dkv_desc, divisibility=(None, self._op.head_dim)), + None, + ), + compile_options=compile_options_for_target(self.compute_capability), + ) + return TupleDict(dq=dq, dkv=dkv, d_sink=d_sink) + + def _input_specs(self) -> tuple[Any | None, ...]: + return ( + self._to_tensor_spec(self.q_desc, divisibility=(None, None, self._op.head_dim)), + self._to_tensor_spec(self.kv_desc, divisibility=(None, self._op.head_dim)), + self._to_tensor_spec(self.out_desc, divisibility=(None, None, self._op.head_dim_v)), + self._to_tensor_spec(self.dout_desc, divisibility=(None, None, self._op.head_dim_v)), + None, + None, + None, + ) + + def _workspace_descs(self) -> tuple[JaxTensorDesc, ...]: + if self._architecture_family == 100: + q_rounded = _round_up(self._op.total_seqlen_q, 8) + kv_rounded = _round_up(self._op.total_seqlen_kv, 8) + head_dim_rounded = _round_up(self._op.head_dim, 8) + return ( + self.q_desc.compact_like( + cudnn_dtype=data_type.UINT8, + shape=(1, self._op.num_heads, q_rounded, 8), + name="workspace_lse_odo", + init_value=0, + ), + self.q_desc.compact_like( + cudnn_dtype=data_type.UINT8, + shape=(1, 1, kv_rounded, head_dim_rounded * 4), + name="workspace_dkv", + init_value=0, + ), + ) + + q_rounded = _round_up(self._op.total_seqlen_q, BLOCK_TILE) + kv_rounded = _round_up(self._op.total_seqlen_kv, BLOCK_TILE) + head_dim_rounded = _round_up(self._op.head_dim, 32) + workspaces = ( + self.q_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(1, q_rounded, self._op.num_heads), + name="workspace_dpsum", + ), + self.q_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(1, q_rounded, self._op.num_heads), + name="workspace_lse_log2", + ), + self.q_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(1, 1, kv_rounded * head_dim_rounded), + name="workspace_dkv_accum", + init_value=0.0, + ), + ) + if self.topk_length_desc is None: + workspaces += ( + self.q_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=(1,), + name="workspace_dummy_topk_length", + ), + ) + return workspaces + + @property + def _architecture_family(self) -> int: + if self.compute_capability is None: + raise RuntimeError("check_support() must resolve the compute capability before lowering") + family = self._compute_capability_family(self.compute_capability, _SUPPORTED_COMPUTE_CAPABILITY_FAMILIES) + if family is None: + raise RuntimeError(f"No sparse-attention backward kernel for SM{self.compute_capability}") + return family + + def _launch( + self, + inputs: tuple[Any, ...], + outputs: tuple[Any, ...], + workspaces: tuple[Any, ...], + stream: Any, + ) -> None: + if self._architecture_family == 90: + self._launch_sm90(inputs, outputs, workspaces, stream) + else: + self._launch_sm100(inputs, outputs, workspaces, stream) + + def _launch_sm100( + self, + inputs: tuple[Any, ...], + outputs: tuple[Any, ...], + workspaces: tuple[Any, ...], + stream: Any, + ) -> None: + import cutlass + + from .dsa_bwd_sm100 import FlashAttentionDSABackwardSm100 + + q, kv, out, dout, lse, attn_sink, topk_idxs, *optional = inputs + topk_length = optional[0] if optional else None + dq, dkv, d_sink = outputs + workspace_lse_odo, workspace_dkv = workspaces + kernel = FlashAttentionDSABackwardSm100( + head_dim=self._op.head_dim, + head_dim_v=self._op.head_dim_v, + block_tile=self._op.block_tile, + max_topk=self._op.max_topk, + ) + problem_shape = ( + cutlass.Int32(self._op.total_seqlen_q), + cutlass.Int32(self._op.total_seqlen_kv), + cutlass.Int32(self._op.head_dim), + (cutlass.Int32(self._op.num_heads), cutlass.Int32(1)), + ) + kernel( + problem_shape, + q, + kv, + out, + dout, + lse, + attn_sink, + topk_idxs, + topk_length, + dq, + dkv, + d_sink, + workspace_lse_odo, + workspace_dkv, + cutlass.Float32(self._op.softmax_scale), + stream, + ) + + def _launch_sm90( + self, + inputs: tuple[Any, ...], + outputs: tuple[Any, ...], + workspaces: tuple[Any, ...], + stream: Any, + ) -> None: + import cutlass + + from .dsa_bwd_sm90 import ( + FlashAttentionDSABackwardSm90, + _FlashAttentionDSABackwardPostprocessSm90, + _FlashAttentionDSABackwardPreprocessSm90, + ) + + q, kv, out, dout, lse, attn_sink, topk_idxs, *optional = inputs + dq, dkv, d_sink = outputs + workspace_dpsum, workspace_lse_log2, workspace_dkv_accum, *optional_workspaces = workspaces + topk_length = optional[0] if optional else optional_workspaces[0] + + q4 = _prepend_unit_dim(q) + kv4 = _as_bshd_kv(kv) + out4 = _prepend_unit_dim(out) + dout4 = _prepend_unit_dim(dout) + lse4 = _prepend_unit_dim(lse) + topk4 = _prepend_unit_dim(topk_idxs) + topk_length4 = _prepend_unit_dim(topk_length) + dq4 = _prepend_unit_dim(dq) + dkv4 = _as_bshd_kv(dkv) + + dtype = q.element_type + preprocess = _FlashAttentionDSABackwardPreprocessSm90( + dtype, + self._op.head_dim_v, + 90, + BLOCK_TILE, + num_threads=256, + ) + preprocess( + out4, + dout4, + workspace_dpsum, + lse4, + workspace_lse_log2, + attn_sink, + d_sink, + None, + None, + None, + stream, + ) + + backward = FlashAttentionDSABackwardSm90( + dtype, + self._op.head_dim, + self._op.head_dim_v, + self._op.num_heads, + tile_m=BLOCK_TILE, + tile_n=BLOCK_TILE, + KV_stage=1, + PdS_stage=1, + SdP_swapAB=False, + dKV_swapAB=False, + dQ_swapAB=False, + num_threads=256, + have_topk_length=self.topk_length_desc is not None, + max_topk=self._op.max_topk, + ) + backward( + q4, + kv4, + dout4, + workspace_lse_log2, + workspace_dpsum, + dq4, + workspace_dkv_accum, + topk4, + topk_length4, + cutlass.Float32(self._op.softmax_scale), + stream, + ) + + hdim_chunk = 64 if self._op.head_dim == 576 else min(128, self._op.head_dim) + postprocess = _FlashAttentionDSABackwardPostprocessSm90( + dtype, + hdim_chunk=hdim_chunk, + tile_n=BLOCK_TILE, + head_dim=self._op.head_dim, + num_threads=hdim_chunk, + N_hdim_chunks=self._op.head_dim // hdim_chunk, + ) + postprocess( + workspace_dkv_accum, + dkv4, + cutlass.Int32(self._op.total_seqlen_kv), + stream, + ) + + +def _prepend_unit_dim(tensor: Any) -> Any: + """Return a CuTe view with a leading batch dimension of one.""" + + import cutlass.cute as cute + + leading_stride = tensor.shape[0] * tensor.stride[0] + return cute.make_tensor( + tensor.iterator, + cute.make_layout((1, *tensor.shape), stride=(leading_stride, *tensor.stride)), + ) + + +def _as_bshd_kv(tensor: Any) -> Any: + """Return a CuTe ``(1, sequence, 1, head_dim)`` MQA view.""" + + import cutlass.cute as cute + + leading_stride = tensor.shape[0] * tensor.stride[0] + return cute.make_tensor( + tensor.iterator, + cute.make_layout( + (1, tensor.shape[0], 1, tensor.shape[1]), + stride=(leading_stride, tensor.stride[0], tensor.stride[0], tensor.stride[1]), + ), + ) + + +def _round_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +@jax.jit(static_argnames=("softmax_scale", "block_tile", "target_compute_capability")) +def sparse_attention_backward_wrapper( + q: Any, + kv: Any, + out: Any, + dout: Any, + lse: Any, + attn_sink: Any, + topk_idxs: Any, + softmax_scale: Optional[float] = None, + topk_length: Any | None = None, + block_tile: int = BLOCK_TILE, + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute DeepSeek sparse-attention gradients from JAX arrays.""" + + samples = tuple(jax.ShapeDtypeStruct(value.shape, value.dtype) for value in (q, kv, out, dout, lse, attn_sink, topk_idxs)) + sample_topk_length = None if topk_length is None else jax.ShapeDtypeStruct(topk_length.shape, topk_length.dtype) + return SparseAttentionBackward( + *samples, + sample_topk_length=sample_topk_length, + softmax_scale=softmax_scale, + block_tile=block_tile, + target_compute_capability=target_compute_capability, + )(q, kv, out, dout, lse, attn_sink, topk_idxs, topk_length) + + +__all__ = ["SparseAttentionBackward", "sparse_attention_backward_wrapper"] diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/op.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/op.py new file mode 100644 index 000000000..a28762255 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/op.py @@ -0,0 +1,226 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral DeepSeek sparse-attention backward operation.""" + +from __future__ import annotations + +import math +from typing import Any, Optional + +from ... import data_type +from ..._op import Op +from ..._tensor_desc import TensorDesc + +SUPPORTED_HEAD_DIMS = (512, 576) +BLOCK_TILE = 64 + + +class SparseAttentionBackwardOp(Op): + """Complete logical signature for DeepSeek sparse-attention backward. + + The operation describes the flat MQA interface shared by the Hopper and + Blackwell implementations. Framework adapters own device selection, + allocation, and lowering. + """ + + def __init__( + self, + *, + q: TensorDesc[Any], + kv: TensorDesc[Any], + output: TensorDesc[Any], + doutput: TensorDesc[Any], + lse: TensorDesc[Any], + attn_sink: TensorDesc[Any], + topk_idxs: TensorDesc[Any], + dq: TensorDesc[Any], + dkv: TensorDesc[Any], + d_sink: TensorDesc[Any], + topk_length: Optional[TensorDesc[Any]] = None, + softmax_scale: Optional[float] = None, + block_tile: int = BLOCK_TILE, + ) -> None: + descriptors = ( + ("q", q), + ("kv", kv), + ("output", output), + ("doutput", doutput), + ("lse", lse), + ("attn_sink", attn_sink), + ("topk_idxs", topk_idxs), + ("dq", dq), + ("dkv", dkv), + ("d_sink", d_sink), + ) + for name, desc in descriptors: + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + if topk_length is not None and not isinstance(topk_length, TensorDesc): + raise TypeError(f"topk_length must be a TensorDesc or None, got {type(topk_length).__name__}") + + self.q = q + self.kv = kv + self.output = output + self.doutput = doutput + self.lse = lse + self.attn_sink = attn_sink + self.topk_idxs = topk_idxs + self.topk_length = topk_length + self.dq = dq + self.dkv = dkv + self.d_sink = d_sink + self.requested_softmax_scale = softmax_scale + self.block_tile = int(block_tile) + + self.total_seqlen_q: Optional[int] = None + self.total_seqlen_kv: Optional[int] = None + self.num_heads: Optional[int] = None + self.head_dim: Optional[int] = None + self.head_dim_v: Optional[int] = None + self.max_topk: Optional[int] = None + self.softmax_scale: Optional[float] = None + + def check_support(self) -> bool: + """Validate the complete tensor signature and static configuration.""" + + self.total_seqlen_q = None + self.total_seqlen_kv = None + self.num_heads = None + self.head_dim = None + self.head_dim_v = None + self.max_topk = None + self.softmax_scale = None + + self._check_ranks_and_dtypes() + + total_seqlen_q, num_heads, head_dim = self.q.shape + total_seqlen_kv, kv_head_dim = self.kv.shape + max_topk = self.topk_idxs.shape[1] + head_dim_v = 512 if head_dim == 576 else head_dim + + dimensions = { + "total_seqlen_q": total_seqlen_q, + "total_seqlen_kv": total_seqlen_kv, + "num_heads": num_heads, + "head_dim": head_dim, + "max_topk": max_topk, + } + invalid = ", ".join(f"{name}={value}" for name, value in dimensions.items() if value <= 0) + if invalid: + raise ValueError(f"Sparse-attention dimensions must be positive, got {invalid}") + if kv_head_dim != head_dim: + raise ValueError(f"KV head dimension must match Q ({head_dim}), got {kv_head_dim}") + if head_dim not in SUPPORTED_HEAD_DIMS: + raise ValueError(f"head_dim must be one of {SUPPORTED_HEAD_DIMS}, got {head_dim}") + if num_heads % BLOCK_TILE != 0: + raise ValueError(f"num_heads must be divisible by {BLOCK_TILE}, got {num_heads}") + if self.block_tile != BLOCK_TILE: + raise ValueError(f"block_tile must be {BLOCK_TILE}, got {self.block_tile}") + + self._check_shapes(total_seqlen_q, total_seqlen_kv, num_heads, head_dim, head_dim_v, max_topk) + self._check_compact_layouts() + + if self.requested_softmax_scale is None: + softmax_scale = 1.0 / math.sqrt(head_dim) + else: + try: + softmax_scale = float(self.requested_softmax_scale) + except (TypeError, ValueError) as error: + raise TypeError(f"softmax_scale must be a real scalar or None, got {self.requested_softmax_scale!r}") from error + if not math.isfinite(softmax_scale): + raise ValueError(f"softmax_scale must be finite, got {softmax_scale}") + + self.total_seqlen_q = total_seqlen_q + self.total_seqlen_kv = total_seqlen_kv + self.num_heads = num_heads + self.head_dim = head_dim + self.head_dim_v = head_dim_v + self.max_topk = max_topk + self.softmax_scale = softmax_scale + return True + + def _check_ranks_and_dtypes(self) -> None: + expected_ranks = ( + (self.q, 3, "Q"), + (self.kv, 2, "KV"), + (self.output, 3, "O"), + (self.doutput, 3, "dO"), + (self.lse, 2, "LSE"), + (self.attn_sink, 1, "attn_sink"), + (self.topk_idxs, 2, "topk_idxs"), + (self.dq, 3, "dQ"), + (self.dkv, 2, "dKV"), + (self.d_sink, 1, "d_sink"), + ) + if self.topk_length is not None: + expected_ranks += ((self.topk_length, 1, "topk_length"),) + for desc, rank, name in expected_ranks: + if desc.ndim != rank: + raise ValueError(f"{name} must have rank {rank}, got shape {desc.shape}") + + if self.q.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError(f"Q must have dtype float16 or bfloat16, got {self.q.dtype}") + for desc, name in ( + (self.kv, "KV"), + (self.output, "O"), + (self.doutput, "dO"), + (self.dq, "dQ"), + (self.dkv, "dKV"), + ): + if desc.cudnn_dtype != self.q.cudnn_dtype: + raise ValueError(f"{name} must have the same dtype as Q, got {desc.dtype}") + for desc, name in ((self.lse, "LSE"), (self.attn_sink, "attn_sink"), (self.d_sink, "d_sink")): + if desc.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"{name} must have dtype float32, got {desc.dtype}") + for desc, name in ((self.topk_idxs, "topk_idxs"), (self.topk_length, "topk_length")): + if desc is not None and desc.cudnn_dtype != data_type.INT32: + raise ValueError(f"{name} must have dtype int32, got {desc.dtype}") + + def _check_shapes( + self, + total_seqlen_q: int, + total_seqlen_kv: int, + num_heads: int, + head_dim: int, + head_dim_v: int, + max_topk: int, + ) -> None: + expected_shapes = ( + (self.output, (total_seqlen_q, num_heads, head_dim_v), "O"), + (self.doutput, (total_seqlen_q, num_heads, head_dim_v), "dO"), + (self.lse, (total_seqlen_q, num_heads), "LSE"), + (self.attn_sink, (num_heads,), "attn_sink"), + (self.topk_idxs, (total_seqlen_q, max_topk), "topk_idxs"), + (self.dq, (total_seqlen_q, num_heads, head_dim), "dQ"), + (self.dkv, (total_seqlen_kv, head_dim), "dKV"), + (self.d_sink, (num_heads,), "d_sink"), + ) + if self.topk_length is not None: + expected_shapes += ((self.topk_length, (total_seqlen_q,), "topk_length"),) + for desc, expected, name in expected_shapes: + if desc.shape != expected: + raise ValueError(f"{name} must have shape {expected}, got {desc.shape}") + + def _check_compact_layouts(self) -> None: + descriptors = ( + self.q, + self.kv, + self.output, + self.doutput, + self.lse, + self.attn_sink, + self.topk_idxs, + self.dq, + self.dkv, + self.d_sink, + ) + if self.topk_length is not None: + descriptors += (self.topk_length,) + for desc in descriptors: + if not desc.is_compact(tuple(reversed(range(desc.ndim)))): + name = desc.name or "tensor" + raise ValueError(f"{name} must be row-major contiguous, got stride {desc.stride}") + + +__all__ = ["BLOCK_TILE", "SUPPORTED_HEAD_DIMS", "SparseAttentionBackwardOp"] diff --git a/python/cudnn/deepseek_sparse_attention/utils/compiler.py b/python/cudnn/deepseek_sparse_attention/utils/compiler.py index 65b68e81f..9fdfbe05b 100644 --- a/python/cudnn/deepseek_sparse_attention/utils/compiler.py +++ b/python/cudnn/deepseek_sparse_attention/utils/compiler.py @@ -18,20 +18,41 @@ from functools import lru_cache -import torch - # (compute_capability) → cute DSL --gpu-arch flag value. # H100, B200/B300, and sm_100f require architecture-specific variants because # the kernels use TMA / tcgen05 instructions that are only guaranteed to lower # correctly under the matching SASS gencode. _ARCH_MAP = { - (9, 0): "sm_90a", # Hopper H100 - (10, 0): "sm_100a", # Blackwell B200 - (10, 3): "sm_103a", # Blackwell Ultra B300 - (10, 7): "sm_100f", + 90: "sm_90a", # Hopper H100 + 100: "sm_100a", # Blackwell B200 + 103: "sm_103a", # Blackwell Ultra B300 + 107: "sm_100f", } +def gpu_arch_flag_for_compute_capability(compute_capability: int) -> str: + """Return the CuTe architecture flag for an explicit compilation target.""" + + if isinstance(compute_capability, bool) or not isinstance(compute_capability, int): + raise TypeError(f"compute_capability must be an int, got {type(compute_capability).__name__}") + try: + return _ARCH_MAP[compute_capability] + except KeyError as error: + raise RuntimeError( + f"Unsupported GPU compute capability SM{compute_capability} for DSA CuTe kernels. " + "Add it to deepseek_sparse_attention/utils/compiler.py::_ARCH_MAP." + ) from error + + +def compile_options_for_target(compute_capability: int, extra: str = "") -> str: + """Build CuTe compile options for a framework-supplied GPU target.""" + + parts = ["--enable-tvm-ffi", f"--gpu-arch {gpu_arch_flag_for_compute_capability(compute_capability)}"] + if extra: + parts.append(extra) + return " ".join(parts) + + @lru_cache(maxsize=None) def gpu_arch_flag() -> str: """Return the ``sm_XXX`` value for the current CUDA device. @@ -39,15 +60,12 @@ def gpu_arch_flag() -> str: Cached because torch.cuda.get_device_capability() is cheap but the function gets called inside every cute.compile site. """ + import torch + if not torch.cuda.is_available(): raise RuntimeError("cute.compile requires CUDA; no GPU available") - cap = torch.cuda.get_device_capability() - arch = _ARCH_MAP.get(cap) - if arch is None: - raise RuntimeError( - f"Unsupported GPU compute capability {cap} for DSA CuTe kernels. " f"Add it to deepseek_sparse_attention/utils/compiler.py::_ARCH_MAP." - ) - return arch + major, minor = torch.cuda.get_device_capability() + return gpu_arch_flag_for_compute_capability(major * 10 + minor) def compile_options(extra: str = "") -> str: @@ -59,7 +77,8 @@ def compile_options(extra: str = "") -> str: Example: cute.compile(..., options=compile_options("--opt-level 3")) """ - parts = ["--enable-tvm-ffi", f"--gpu-arch {gpu_arch_flag()}"] + arch = gpu_arch_flag() + parts = ["--enable-tvm-ffi", f"--gpu-arch {arch}"] if extra: parts.append(extra) return " ".join(parts) diff --git a/python/cudnn/jax/__init__.py b/python/cudnn/jax/__init__.py index 453db9d8b..67098c72d 100644 --- a/python/cudnn/jax/__init__.py +++ b/python/cudnn/jax/__init__.py @@ -1,7 +1,15 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Public optional-dependency boundary for JAX operation APIs.""" +"""Public optional-dependency boundary for JAX operation APIs. + +JAX and CUTLASS are validated eagerly at this boundary. Individual operation +adapters remain lazy so importing :mod:`cudnn.jax` does not import every CuTe +kernel or any Torch API module. +""" + +from importlib import import_module +from typing import Any _INSTALL_HINT = "pip install 'nvidia-cudnn-frontend[jax]'" _OPTIONAL_DEPENDENCY_PREFIXES = ("cuda", "cutlass", "jax") @@ -24,16 +32,90 @@ f"Install the optional JAX dependencies with `{_INSTALL_HINT}`." ) -from .._jax import JaxApiBase, JaxTensorDesc, disable_device_compatibility_checks -from ..rmsnorm_rht_amax.jax import ( - RmsNormRhtAmaxSm100, - rmsnorm_rht_amax_sm100, -) +from .._jax import JaxApiBase, JaxTensorDesc, TupleDict, disable_device_compatibility_checks # noqa: E402 + +_OPERATION_EXPORTS = { + "RmsNormRhtAmaxSm100": ("..rmsnorm_rht_amax.jax", "RmsNormRhtAmaxSm100"), + "rmsnorm_rht_amax_sm100": ("..rmsnorm_rht_amax.jax", "rmsnorm_rht_amax_sm100"), + "SparseAttentionBackward": ("..deepseek_sparse_attention.sparse_attention_backward.jax", "SparseAttentionBackward"), + "sparse_attention_backward_wrapper": ( + "..deepseek_sparse_attention.sparse_attention_backward.jax", + "sparse_attention_backward_wrapper", + ), + "IndexerForward": ("..deepseek_sparse_attention.indexer_forward.jax", "IndexerForward"), + "indexer_forward_wrapper": ("..deepseek_sparse_attention.indexer_forward.jax", "indexer_forward_wrapper"), + "IndexerTopK": ("..deepseek_sparse_attention.indexer_top_k.jax", "IndexerTopK"), + "indexer_top_k_wrapper": ("..deepseek_sparse_attention.indexer_top_k.jax", "indexer_top_k_wrapper"), + "local_to_global_wrapper": ("..deepseek_sparse_attention.indexer_top_k.jax", "local_to_global_wrapper"), + "compactify_wrapper": ("..deepseek_sparse_attention.indexer_top_k.jax", "compactify_wrapper"), + "SparseIndexerScoreRecompute": ("..deepseek_sparse_attention.score_recompute.jax", "SparseIndexerScoreRecompute"), + "sparse_indexer_score_recompute_wrapper": ( + "..deepseek_sparse_attention.score_recompute.jax", + "sparse_indexer_score_recompute_wrapper", + ), + "SparseAttnScoreRecompute": ("..deepseek_sparse_attention.score_recompute.jax", "SparseAttnScoreRecompute"), + "sparse_attn_score_recompute_wrapper": ( + "..deepseek_sparse_attention.score_recompute.jax", + "sparse_attn_score_recompute_wrapper", + ), + "DenseIndexerScoreRecompute": ("..deepseek_sparse_attention.score_recompute.jax", "DenseIndexerScoreRecompute"), + "dense_indexer_score_recompute_wrapper": ( + "..deepseek_sparse_attention.score_recompute.jax", + "dense_indexer_score_recompute_wrapper", + ), + "DenseAttnScoreRecompute": ("..deepseek_sparse_attention.score_recompute.jax", "DenseAttnScoreRecompute"), + "dense_attn_score_recompute_wrapper": ( + "..deepseek_sparse_attention.score_recompute.jax", + "dense_attn_score_recompute_wrapper", + ), + "IndexerBackward": ("..deepseek_sparse_attention.indexer_backward.jax", "IndexerBackward"), + "indexer_backward_wrapper": ("..deepseek_sparse_attention.indexer_backward.jax", "indexer_backward_wrapper"), + "DenseIndexerBackward": ("..deepseek_sparse_attention.indexer_backward.jax", "DenseIndexerBackward"), + "dense_indexer_backward_wrapper": ("..deepseek_sparse_attention.indexer_backward.jax", "dense_indexer_backward_wrapper"), +} + +_DSA_EXPORTS = frozenset(name for name in _OPERATION_EXPORTS if name not in {"RmsNormRhtAmaxSm100", "rmsnorm_rht_amax_sm100"}) + + +def _load_operation(name: str) -> Any: + module_name, symbol_name = _OPERATION_EXPORTS[name] + value = getattr(import_module(module_name, __name__), symbol_name) + globals()[name] = value + return value + + +class _DSANamespace: + """Lazy JAX view of the DeepSeek sparse-attention operation family.""" + + def __getattr__(self, name: str) -> Any: + if name not in _DSA_EXPORTS: + raise AttributeError(f"cudnn.jax.DSA has no attribute {name!r}") + value = _load_operation(name) + setattr(self, name, value) + return value + + def __dir__(self) -> list[str]: + return sorted((*vars(self), *_DSA_EXPORTS)) + + +DSA = _DSANamespace() + + +def __getattr__(name: str) -> Any: + if name in _OPERATION_EXPORTS: + return _load_operation(name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__() -> list[str]: + return sorted((*globals(), *_OPERATION_EXPORTS)) + __all__ = [ + "DSA", "JaxApiBase", "JaxTensorDesc", + "TupleDict", "disable_device_compatibility_checks", - "RmsNormRhtAmaxSm100", - "rmsnorm_rht_amax_sm100", + *_OPERATION_EXPORTS, ] diff --git a/test/python/fe_api/dsa/dsa_reference.py b/test/python/fe_api/dsa/dsa_reference.py index fe8dcb334..0a3ceba8c 100644 --- a/test/python/fe_api/dsa/dsa_reference.py +++ b/test/python/fe_api/dsa/dsa_reference.py @@ -66,7 +66,10 @@ def ref_sparse_attention_forward( q_f = q.to(torch.float32) k_f = kv.to(torch.float32) - v_f = kv.to(torch.float32) + # MLA uses a 512-wide value projection for the 576-wide Q/K variant. + # K and V share storage, so V is the leading D_v slice of KV. + head_dim_v = 512 if d == 576 else d + v_f = kv[..., :head_dim_v].to(torch.float32) mask = _make_topk_mask(topk_idxs, topk_length, t_kv) # (T, T_kv) mask = mask.unsqueeze(1).expand(t, h, t_kv) # (T, H, T_kv) @@ -76,8 +79,11 @@ def ref_sparse_attention_forward( scores = scores.masked_fill(~mask, float("-inf")) lse = torch.logsumexp(scores, dim=-1) # (T, H), excludes sink - sink = attn_sink.view(1, h) - lse_with_sink = torch.logaddexp(lse, sink) + sink = attn_sink.view(1, h, 1).expand(t, h, 1) + # Include the finite sink directly in the reduction. Composing + # logaddexp(logsumexp(all -inf), sink) gives NaN score gradients in + # PyTorch even though the sink makes the mathematical denominator finite. + lse_with_sink = torch.logsumexp(torch.cat((scores, sink), dim=-1), dim=-1) weights = torch.exp(scores - lse_with_sink.unsqueeze(-1)) out = torch.einsum("thk,kd->thd", weights, v_f) diff --git a/test/python/fe_api/dsa/dsa_utils.py b/test/python/fe_api/dsa/dsa_utils.py index d62375ca0..84bd7d741 100644 --- a/test/python/fe_api/dsa/dsa_utils.py +++ b/test/python/fe_api/dsa/dsa_utils.py @@ -17,7 +17,7 @@ ] DSA_SPARSE_ATTENTION_BACKWARD_PARAM_MARKS = [ - pytest.mark.parametrize("head_dim", [512]), + pytest.mark.parametrize("head_dim", [512, 576]), pytest.mark.parametrize("head_dim_v", [512]), pytest.mark.parametrize("num_heads", [64]), pytest.mark.parametrize("topk", [512]), diff --git a/test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py b/test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py index c6d10c061..f1ddf70ed 100644 --- a/test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py +++ b/test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py @@ -40,10 +40,44 @@ def _allocate(cfg, has_topk_length: bool): topk_length = None if has_topk_length: topk_length = torch.randint(1, topk_k + 1, (total_s_q,), dtype=torch.int32, device=device) + # Empty compactified rows are valid: attention falls back entirely to + # the sink and all Q/KV gradients for that row are zero. + topk_length[0] = 0 return q, kv, attn_sink, topk_idxs, topk_length +@pytest.mark.L0 +@pytest.mark.parametrize("head_dim", [512, 576]) +def test_sparse_attention_reference_supports_empty_rows(head_dim): + """Keep the D=576 value width and sink-only rows numerically well-defined.""" + + torch.manual_seed(0) + total_q, total_kv, num_heads, topk = 2, 3, 2, 2 + q = torch.randn(total_q, num_heads, head_dim, dtype=torch.float32, requires_grad=True) + kv = torch.randn(total_kv, head_dim, dtype=torch.float32, requires_grad=True) + attn_sink = torch.randn(num_heads, dtype=torch.float32, requires_grad=True) + topk_idxs = torch.tensor([[0, 1], [1, 2]], dtype=torch.int32) + topk_length = torch.tensor([0, topk], dtype=torch.int32) + + out, lse = ref_sparse_attention_forward( + q, + kv, + attn_sink, + topk_idxs, + topk_length=topk_length, + ) + out.sum().backward() + + head_dim_v = 512 if head_dim == 576 else head_dim + assert out.shape == (total_q, num_heads, head_dim_v) + assert torch.isneginf(lse[0]).all() + assert torch.isfinite(q.grad).all() + assert torch.isfinite(kv.grad).all() + assert torch.isfinite(attn_sink.grad).all() + assert torch.count_nonzero(q.grad[0]) == 0 + + @pytest.mark.L0 @torch_fork_set_rng(seed=0) @with_dsa_sparse_attention_backward_params @@ -93,21 +127,18 @@ def test_DSA_sparse_attention_backward_wrapper( ) dout = torch.randn_like(out) - try: - result = DSA.sparse_attention_backward_wrapper( - q, - kv, - out, - dout, - lse, - attn_sink, - topk_idxs, - softmax_scale=softmax_scale, - topk_length=topk_length, - stream=stream, - ) - except (ValueError, NotImplementedError, RuntimeError) as e: - pytest.skip(f"Unsupported testcase: {e}") + result = DSA.sparse_attention_backward_wrapper( + q, + kv, + out, + dout, + lse, + attn_sink, + topk_idxs, + softmax_scale=softmax_scale, + topk_length=topk_length, + stream=stream, + ) dq, dkv, d_sink = result["dq"], result["dkv"], result["d_sink"] diff --git a/test/python/fe_api/test_dsa_compiler.py b/test/python/fe_api/test_dsa_compiler.py new file mode 100644 index 000000000..11ac77981 --- /dev/null +++ b/test/python/fe_api/test_dsa_compiler.py @@ -0,0 +1,52 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for DSA compilation targets.""" + +from importlib.util import module_from_spec, spec_from_file_location +from pathlib import Path +import sys +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_MODULE_PATH = Path(__file__).resolve().parents[3] / "python" / "cudnn" / "deepseek_sparse_attention" / "utils" / "compiler.py" + + +class DsaCompilerTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + spec = spec_from_file_location("dsa_compiler_contract", _MODULE_PATH) + assert spec is not None and spec.loader is not None + cls.module = module_from_spec(spec) + spec.loader.exec_module(cls.module) + + def test_explicit_targets_do_not_import_torch(self): + torch_before = sys.modules.get("torch") + self.assertEqual(self.module.gpu_arch_flag_for_compute_capability(90), "sm_90a") + self.assertEqual(self.module.gpu_arch_flag_for_compute_capability(100), "sm_100a") + self.assertEqual(self.module.gpu_arch_flag_for_compute_capability(103), "sm_103a") + self.assertEqual(self.module.gpu_arch_flag_for_compute_capability(107), "sm_100f") + self.assertIs(sys.modules.get("torch"), torch_before) + + def test_builds_jax_compile_options_from_explicit_target(self): + self.assertEqual( + self.module.compile_options_for_target(103, "--opt-level 3"), + "--enable-tvm-ffi --gpu-arch sm_103a --opt-level 3", + ) + + def test_rejects_unknown_or_non_integer_targets(self): + with self.assertRaisesRegex(RuntimeError, "Unsupported GPU compute capability SM101"): + self.module.compile_options_for_target(101) + with self.assertRaisesRegex(TypeError, "must be an int"): + self.module.compile_options_for_target(True) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_dsa_indexer_backward_op.py b/test/python/fe_api/test_dsa_indexer_backward_op.py new file mode 100644 index 000000000..e861da4c0 --- /dev/null +++ b/test/python/fe_api/test_dsa_indexer_backward_op.py @@ -0,0 +1,231 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for DSA indexer-backward metadata.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "deepseek_sparse_attention" / "indexer_backward" +_PACKAGE = "cudnn_dsa_indexer_backward_op_test" + + +class _DataType(Enum): + NOT_SET = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + + +class IndexerBackwardOpContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + dsa_name = f"{_PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_CUDNN_ROOT / "deepseek_sparse_attention")] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + sys.modules[dsa_name] = dsa + + operation_name = f"{dsa_name}.indexer_backward" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + try: + cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}._op") + cls.op_module = importlib.import_module(f"{operation_name}.op") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def _desc(self, shape, dtype=_DataType.BFLOAT16, *, name=""): + return self.tensor_module.make_compact_tensor_desc( + dtype=dtype, + shape=tuple(shape), + name=name, + ) + + def _sparse(self, **overrides): + batch, seqlen_q, seqlen_k, heads, dim, topk = 2, 64, 256, 64, 128, 128 + arguments = { + "index_q": self._desc((batch, seqlen_q, heads, dim), name="index_q"), + "weights": self._desc((batch, seqlen_q, heads), name="weights"), + "index_k": self._desc((batch, seqlen_k, dim), name="index_k"), + "d_index_q": self._desc((batch, seqlen_q, heads, dim), name="d_index_q"), + "d_weights": self._desc((batch, seqlen_q, heads), name="d_weights"), + "d_index_k": self._desc((batch, seqlen_k, dim), name="d_index_k"), + "attn_score": self._desc((batch, seqlen_q, topk), _DataType.FLOAT, name="attn_score"), + "index_score": self._desc((batch, seqlen_q, topk), _DataType.FLOAT, name="index_score"), + "topk_indices": self._desc((batch, seqlen_q, topk), _DataType.INT32, name="topk_indices"), + } + arguments.update(overrides) + return self.op_module.IndexerBackwardOp(**arguments) + + def _sparse_thd(self, **overrides): + total_q, total_k, heads, dim, topk = 96, 384, 64, 128, 128 + arguments = { + "index_q": self._desc((total_q, heads, dim), name="index_q"), + "weights": self._desc((total_q, heads), name="weights"), + "index_k": self._desc((total_k, dim), name="index_k"), + "d_index_q": self._desc((total_q, heads, dim), name="d_index_q"), + "d_weights": self._desc((total_q, heads), name="d_weights"), + "d_index_k": self._desc((total_k, dim), name="d_index_k"), + "attn_score": self._desc((total_q, topk), _DataType.FLOAT, name="attn_score"), + "index_score": self._desc((total_q, topk), _DataType.FLOAT, name="index_score"), + "topk_indices": self._desc((total_q, topk), _DataType.INT32, name="topk_indices"), + "topk_indices_global": True, + } + arguments.update(overrides) + return self.op_module.IndexerBackwardOp(**arguments) + + def _dense_bshd(self, **overrides): + batch, seqlen_q, seqlen_k, heads, dim = 2, 64, 256, 64, 128 + arguments = { + "index_q": self._desc((batch, seqlen_q, heads, dim), name="index_q"), + "weights": self._desc((batch, seqlen_q, heads), name="weights"), + "index_k": self._desc((batch, seqlen_k, dim), name="index_k"), + "d_index_q": self._desc((batch, seqlen_q, heads, dim), name="d_index_q"), + "d_weights": self._desc((batch, seqlen_q, heads), name="d_weights"), + "d_index_k": self._desc((batch, seqlen_k, dim), name="d_index_k"), + "attn_score": self._desc((batch, seqlen_q, seqlen_k), _DataType.FLOAT, name="attn_score"), + "attn_l1norm": self._desc((batch, seqlen_q), _DataType.FLOAT, name="attn_l1norm"), + "index_score": self._desc((batch, seqlen_q, seqlen_k), _DataType.FLOAT, name="index_score"), + "index_lse": self._desc((batch, seqlen_q), _DataType.FLOAT, name="index_lse"), + } + arguments.update(overrides) + return self.op_module.DenseIndexerBackwardOp(**arguments) + + def _dense_thd(self, **overrides): + batch, total_q, total_k, max_k, heads, dim = 2, 96, 384, 256, 64, 128 + arguments = { + "index_q": self._desc((total_q, heads, dim), name="index_q"), + "weights": self._desc((total_q, heads), name="weights"), + "index_k": self._desc((total_k, dim), name="index_k"), + "d_index_q": self._desc((total_q, heads, dim), name="d_index_q"), + "d_weights": self._desc((total_q, heads), name="d_weights"), + "d_index_k": self._desc((total_k, dim), name="d_index_k"), + "attn_score": self._desc((total_q, max_k), _DataType.FLOAT, name="attn_score"), + "attn_l1norm": self._desc((total_q,), _DataType.FLOAT, name="attn_l1norm"), + "index_score": self._desc((total_q, max_k), _DataType.FLOAT, name="index_score"), + "index_lse": self._desc((total_q,), _DataType.FLOAT, name="index_lse"), + "cu_seqlens_q": self._desc((batch + 1,), _DataType.INT32, name="cu_seqlens_q"), + "cu_seqlens_k": self._desc((batch + 1,), _DataType.INT32, name="cu_seqlens_k"), + "q_causal_offsets": self._desc((batch,), _DataType.INT32, name="q_causal_offsets"), + "max_seqlen_q": 64, + "max_seqlen_k": max_k, + } + arguments.update(overrides) + return self.op_module.DenseIndexerBackwardOp(**arguments) + + def test_sparse_complete_signature(self): + operation = self._sparse() + self.assertIsInstance(operation, self.base_module.Op) + self.assertTrue(operation.check_support()) + self.assertFalse(operation.is_thd) + self.assertEqual( + (operation.batch, operation.seqlen_q, operation.seqlen_k, operation.heads, operation.head_dim, operation.topk), + (2, 64, 256, 64, 128, 128), + ) + + def test_sparse_packed_thd_complete_signature(self): + operation = self._sparse_thd() + self.assertTrue(operation.check_support()) + self.assertTrue(operation.is_thd) + self.assertEqual( + (operation.batch, operation.seqlen_q, operation.seqlen_k, operation.heads, operation.head_dim, operation.topk), + (1, 96, 384, 64, 128, 128), + ) + + def test_sparse_packed_thd_requires_global_indices(self): + with self.assertRaisesRegex(ValueError, "requires topk_indices_global=True"): + self._sparse_thd(topk_indices_global=False).check_support() + + def test_sparse_rejects_invalid_output_score_and_configuration(self): + cases = ( + ({"d_index_q": self._desc((2, 64, 32, 128))}, "d_index_q must have shape"), + ({"index_score": self._desc((2, 64, 64), _DataType.FLOAT)}, "index_score must have shape"), + ({"topk_indices": self._desc((2, 64, 64), _DataType.INT32)}, r"topk \(64\) must be divisible"), + ({"block_i": 64}, "block_i must be 128"), + ({"sm_scale": float("inf")}, "sm_scale must be finite"), + ) + for overrides, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self._sparse(**overrides).check_support() + + def test_scale_preserves_values_supported_by_forward(self): + for scale in (0.0, -0.5): + with self.subTest(scale=scale): + operation = self._sparse(sm_scale=scale) + self.assertTrue(operation.check_support()) + self.assertEqual(operation.sm_scale, scale) + + def test_dense_bshd_complete_signature(self): + operation = self._dense_bshd(q_causal_offsets=self._desc((2,), _DataType.INT32)) + self.assertTrue(operation.check_support()) + self.assertFalse(operation.is_thd) + self.assertEqual( + (operation.batch, operation.normalization_tokens, operation.total_k, operation.max_seqlen_q, operation.max_seqlen_k), + (2, 128, 512, 64, 256), + ) + + def test_dense_thd_complete_signature(self): + operation = self._dense_thd() + self.assertTrue(operation.check_support()) + self.assertTrue(operation.is_thd) + self.assertEqual( + (operation.batch, operation.normalization_tokens, operation.total_k, operation.max_seqlen_q, operation.max_seqlen_k), + (2, 96, 384, 64, 256), + ) + + def test_dense_rejects_partial_or_invalid_thd_signature(self): + cases = ( + ({"cu_seqlens_k": None}, "must be provided together"), + ({"max_seqlen_q": None}, "requires max_seqlen_q and max_seqlen_k"), + ({"q_causal_offsets": self._desc((3,), _DataType.INT32)}, "q_causal_offsets must have shape"), + ({"attn_score": self._desc((96, 128), _DataType.FLOAT)}, "attn_score must have shape"), + ) + for overrides, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self._dense_thd(**overrides).check_support() + + def test_operation_imports_no_framework_or_kernel(self): + source = (_OPERATION_ROOT / "op.py").read_text() + for forbidden in ("import torch", "import jax", "import cutlass", "import cuda"): + self.assertNotIn(forbidden, source) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_dsa_indexer_ops.py b/test/python/fe_api/test_dsa_indexer_ops.py new file mode 100644 index 000000000..11223e905 --- /dev/null +++ b/test/python/fe_api/test_dsa_indexer_ops.py @@ -0,0 +1,534 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for the DSA indexer logical operations.""" + +import ast +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_DSA_ROOT = _CUDNN_ROOT / "deepseek_sparse_attention" +_PACKAGE = "cudnn_dsa_indexer_op_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + INT64 = auto() + + +class _DType: + def __init__(self, name: str, itemsize: int): + self.name = name + self.itemsize = itemsize + + def __repr__(self): + return self.name + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + @property + def ndim(self): + return len(self.shape) + + def __getitem__(self, key): + if isinstance(key, tuple) and key[:-1] == (Ellipsis,) and isinstance(key[-1], slice): + return _Array((*self.shape[:-1], key[-1].stop), self.dtype) + if isinstance(key, tuple) and len(key) == self.ndim and all(isinstance(value, slice) for value in key): + shape = [] + for extent, value in zip(self.shape, key): + start = 0 if value.start is None else value.start + stop = extent if value.stop is None else min(value.stop, extent) + step = 1 if value.step is None else value.step + shape.append(max(0, (stop - start + step - 1) // step)) + return _Array(shape, self.dtype) + raise AssertionError(f"Unexpected array slice {key!r}") + + +class _TensorSpec: + def __init__(self, *, layout=None, mode=None, divisibility=None, **_kwargs): + self.layout = layout + self.mode = mode + self.divisibility = divisibility + + +class DsaIndexerOpContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + dsa_name = f"{_PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_DSA_ROOT)] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + sys.modules[dsa_name] = dsa + + for operation in ("indexer_forward", "indexer_top_k"): + name = f"{dsa_name}.{operation}" + module = types.ModuleType(name) + module.__path__ = [str(_DSA_ROOT / operation)] + module.__package__ = name + module.__spec__ = ModuleSpec(name, loader=None, is_package=True) + sys.modules[name] = module + + cls.tensor = importlib.import_module(f"{_PACKAGE}._tensor_desc") + importlib.import_module(f"{_PACKAGE}._op") + cls.forward = importlib.import_module(f"{dsa_name}.indexer_forward.op") + cls.topk = importlib.import_module(f"{dsa_name}.indexer_top_k.op") + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def desc(self, shape, dtype, name="", stride_order=None): + return self.tensor.make_compact_tensor_desc( + dtype=dtype, + shape=tuple(shape), + stride_order=stride_order, + name=name, + ) + + def test_fixed_forward_complete_signature(self): + operation = self.forward.IndexerForwardOp( + q=self.desc((2, 16, 32, 128), _DataType.BFLOAT16), + k=self.desc((2, 17, 1, 128), _DataType.BFLOAT16), + weight=self.desc((2, 16, 32), _DataType.BFLOAT16), + output=self.desc((2, 16, 20), _DataType.FLOAT), + q_causal_offsets=self.desc((2,), _DataType.INT32), + ratio=4, + target_compute_capability=100, + ) + self.assertTrue(operation.check_support()) + self.assertFalse(operation.is_varlen) + self.assertEqual((operation.s_q, operation.s_k), (16, 17)) + self.assertEqual(operation.qhead_per_kv_head, 32) + + def test_varlen_forward_complete_signature(self): + operation = self.forward.IndexerForwardOp( + q=self.desc((31, 64, 128), _DataType.BFLOAT16), + k=self.desc((47, 1, 128), _DataType.BFLOAT16), + weight=self.desc((31, 64), _DataType.BFLOAT16), + output=self.desc((31, 20), _DataType.FLOAT), + cu_seqlens_q=self.desc((3,), _DataType.INT32), + cu_seqlens_k=self.desc((3,), _DataType.INT32), + q_causal_offsets=self.desc((2,), _DataType.INT32), + max_seqlen_q=16, + max_seqlen_k=17, + target_compute_capability=90, + ) + self.assertTrue(operation.check_support()) + self.assertTrue(operation.is_varlen) + self.assertEqual(operation.batch_size, 2) + + def test_forward_rejects_invalid_output_and_optional_signature(self): + with self.assertRaisesRegex(ValueError, "leading shape"): + self.forward.IndexerForwardOp( + q=self.desc((1, 8, 32, 128), _DataType.BFLOAT16), + k=self.desc((1, 5, 1, 128), _DataType.BFLOAT16), + weight=self.desc((1, 8, 32), _DataType.BFLOAT16), + output=self.desc((1, 8, 5), _DataType.FLOAT), + ).check_support() + + oversized_output = self.forward.IndexerForwardOp( + q=self.desc((1, 8, 32, 128), _DataType.BFLOAT16), + k=self.desc((1, 5, 1, 128), _DataType.BFLOAT16), + weight=self.desc((1, 8, 32), _DataType.BFLOAT16), + output=self.desc((1, 8, 12), _DataType.FLOAT), + ) + self.assertTrue(oversized_output.check_support()) + + with self.assertRaisesRegex(ValueError, "requires both"): + self.forward.IndexerForwardOp( + q=self.desc((8, 32, 128), _DataType.BFLOAT16), + k=self.desc((8, 1, 128), _DataType.BFLOAT16), + weight=self.desc((8, 32), _DataType.BFLOAT16), + output=self.desc((8, 8), _DataType.FLOAT), + cu_seqlens_q=self.desc((2,), _DataType.INT32), + max_seqlen_q=8, + max_seqlen_k=8, + ).check_support() + + def test_forward_requires_the_kernel_vector_axis_to_be_contiguous(self): + operation = self.forward.IndexerForwardOp( + q=self.desc( + (1, 8, 32, 128), + _DataType.BFLOAT16, + stride_order=(2, 3, 1, 0), + ), + k=self.desc((1, 5, 1, 128), _DataType.BFLOAT16), + weight=self.desc((1, 8, 32), _DataType.BFLOAT16), + output=self.desc((1, 8, 8), _DataType.FLOAT), + ratio=1, + ) + with self.assertRaisesRegex(ValueError, "canonical D axis contiguous"): + operation.check_support() + + def make_topk(self, *, return_val=True, **overrides): + input_values = self.desc((4, 64), _DataType.FLOAT) + arguments = { + "input_values": input_values, + "seq_lens": self.desc((4,), _DataType.INT32), + "output_indices": self.desc((4, 8), _DataType.INT32), + "output_values": self.desc((4, 8), _DataType.FLOAT) if return_val else None, + "workspace": self.desc((4, 2, 64), _DataType.INT32), + "top_k": 8, + "return_val": return_val, + } + arguments.update(overrides) + return self.topk.IndexerTopKOp(**arguments) + + def test_topk_complete_signature_with_and_without_values(self): + for return_val in (True, False): + with self.subTest(return_val=return_val): + operation = self.make_topk(return_val=return_val) + self.assertTrue(operation.check_support()) + self.assertEqual(operation.buffer_count, 2) + self.assertEqual(operation.max_num_cols, 64) + + def test_topk_rejects_invalid_rows_workspace_and_copy_width(self): + cases = ( + ({"seq_lens": self.desc((2,), _DataType.INT32)}, "must equal"), + ({"workspace": self.desc((4, 1, 64), _DataType.INT32)}, "workspace must"), + ({"num_copy_bits": 24}, "power-of-two"), + ({"top_k": 65}, "num_cols=64"), + ) + for overrides, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self.make_topk(**overrides).check_support() + + +class DsaIndexerImportContractTest(unittest.TestCase): + @staticmethod + def top_level_imports(path: Path) -> tuple[str, ...]: + tree = ast.parse(path.read_text(), filename=str(path)) + imports = [] + for node in tree.body: + if isinstance(node, ast.Import): + imports.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + imports.append("." * node.level + (node.module or "")) + return tuple(imports) + + def test_operation_packages_and_ops_are_framework_lazy(self): + for operation in ("indexer_forward", "indexer_top_k"): + root = _DSA_ROOT / operation + for filename in ("__init__.py", "op.py"): + imports = self.top_level_imports(root / filename) + self.assertFalse( + any(name == framework or name.startswith(f"{framework}.") for name in imports for framework in ("torch", "jax", "cutlass", "cuda")), + f"{operation}/{filename}: {imports}", + ) + + def test_jax_reachable_kernels_do_not_require_torch(self): + files = ( + _DSA_ROOT / "indexer_forward" / "jax.py", + _DSA_ROOT / "indexer_forward" / "indexer_fwd_sm90.py", + _DSA_ROOT / "indexer_forward" / "indexer_fwd_sm100.py", + _DSA_ROOT / "indexer_top_k" / "jax.py", + _DSA_ROOT / "indexer_top_k" / "indexer_top_k_decode_varlen.py", + _DSA_ROOT / "indexer_top_k" / "indexer_top_k_varlen_util.py", + _DSA_ROOT / "indexer_top_k" / "local_to_global_dsl.py", + _DSA_ROOT / "indexer_top_k" / "compactify.py", + ) + for path in files: + imports = self.top_level_imports(path) + self.assertNotIn("torch", imports, str(path)) + + def test_torch_indexer_forward_passes_all_optional_kernel_arguments(self): + path = _DSA_ROOT / "indexer_forward" / "api.py" + tree = ast.parse(path.read_text(), filename=str(path)) + compile_method = next(node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == "compile" and node.lineno < 220) + compile_call = next( + node for node in ast.walk(compile_method) if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "compile" + ) + self.assertEqual(sum(isinstance(argument, ast.Constant) and argument.value is None for argument in compile_call.args), 3) + + tensor_api = next(node for node in ast.walk(compile_method) if isinstance(node, ast.FunctionDef) and node.name == "tensor_api") + launch_call = next(node for node in ast.walk(tensor_api) if isinstance(node, ast.Call)) + self.assertEqual(sum(isinstance(argument, ast.Constant) and argument.value is None for argument in launch_call.args), 3) + + +class DsaIndexerJaxAdapterContractTest(unittest.TestCase): + _PACKAGE = "cudnn_dsa_indexer_jax_test" + + @classmethod + def setUpClass(cls) -> None: + cls.float16 = _DType("float16", 2) + cls.bfloat16 = _DType("bfloat16", 2) + cls.float32 = _DType("float32", 4) + cls.int32 = _DType("int32", 4) + cls.int64 = _DType("int64", 8) + + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.float16 = cls.float16 + fake_jnp.bfloat16 = cls.bfloat16 + fake_jnp.float32 = cls.float32 + fake_jnp.int32 = cls.int32 + fake_jnp.int64 = cls.int64 + fake_jnp.dtype = lambda value: value.dtype if hasattr(value, "dtype") else value + fake_jnp.full = lambda shape, _value, dtype: _Array(shape, dtype) + fake_jnp.reshape = lambda value, shape: _Array(shape, value.dtype) + + fake_jax = types.ModuleType("jax") + fake_jax.__path__ = [] + fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + fake_jax.numpy = fake_jnp + fake_jax.ShapeDtypeStruct = _Array + + def fake_jit(function=None, **options): + def decorate(resolved): + resolved._test_static_argnames = tuple(options.get("static_argnames", ())) + return resolved + + return decorate if function is None else decorate(function) + + fake_jax.jit = fake_jit + fake_jax.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, + ) + + fake_cutlass = types.ModuleType("cutlass") + fake_cutlass.__path__ = [] + fake_cutlass_jax = types.ModuleType("cutlass.jax") + fake_cutlass_jax.TensorSpec = _TensorSpec + fake_cutlass_jax.cutlass_call = None + fake_cutlass.jax = fake_cutlass_jax + + cls.fake_modules = { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "cutlass": fake_cutlass, + "cutlass.jax": fake_cutlass_jax, + } + cls.module_patch = mock.patch.dict(sys.modules, cls.fake_modules) + cls.module_patch.start() + + root = types.ModuleType(cls._PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = cls._PACKAGE + root.__spec__ = ModuleSpec(cls._PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[cls._PACKAGE] = root + + dsa_name = f"{cls._PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_DSA_ROOT)] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + sys.modules[dsa_name] = dsa + for operation in ("indexer_forward", "indexer_top_k"): + name = f"{dsa_name}.{operation}" + package = types.ModuleType(name) + package.__path__ = [str(_DSA_ROOT / operation)] + package.__package__ = name + package.__spec__ = ModuleSpec(name, loader=None, is_package=True) + sys.modules[name] = package + + cls.jax_base = importlib.import_module(f"{cls._PACKAGE}._jax") + cls.jax_base.disable_device_compatibility_checks(True) + cls.forward = importlib.import_module(f"{dsa_name}.indexer_forward.jax") + cls.topk = importlib.import_module(f"{dsa_name}.indexer_top_k.jax") + + @classmethod + def tearDownClass(cls) -> None: + cls.jax_base.disable_device_compatibility_checks(False) + for name in tuple(sys.modules): + if name == cls._PACKAGE or name.startswith(f"{cls._PACKAGE}."): + sys.modules.pop(name, None) + cls.module_patch.stop() + + @staticmethod + def fake_call(captured): + def call(_inputs, *, output_descs, **options): + captured.update(output_descs=output_descs, **options) + output_specs = options.get("output_spec") + if output_specs is None: + output_specs = (None,) * len(output_descs) + + def public_shape(desc, spec): + if spec is None or spec.mode is None: + return desc.shape + shape = [None] * desc.ndim + for canonical_axis, public_axis in enumerate(spec.mode): + shape[public_axis] = desc.shape[canonical_axis] + return tuple(shape) + + return tuple(_Array(public_shape(desc, spec), desc.dtype) for desc, spec in zip(output_descs, output_specs)) + + return call + + def test_forward_declares_initialized_padded_output_for_fixed_and_thd(self): + fixed = self.forward.IndexerForward( + _Array((1, 8, 32, 128), self.bfloat16), + _Array((1, 5, 1, 128), self.bfloat16), + _Array((1, 8, 32), self.bfloat16), + ratio=1, + target_compute_capability=100, + ) + captured = {} + fixed._call_kernel = self.fake_call(captured) + result = fixed( + _Array((1, 8, 32, 128), self.bfloat16), + _Array((1, 5, 1, 128), self.bfloat16), + _Array((1, 8, 32), self.bfloat16), + ) + self.assertEqual(result["scores"].shape, (1, 8, 5)) + output = captured["output_descs"][0] + self.assertEqual(output.shape, (1, 8, 8)) + self.assertEqual(output.init_value, float("-inf")) + + thd = self.forward.IndexerForward( + _Array((8, 32, 128), self.bfloat16), + _Array((9, 1, 128), self.bfloat16), + _Array((8, 32), self.bfloat16), + sample_cu_seqlens_q=_Array((2,), self.int32), + sample_cu_seqlens_k=_Array((2,), self.int32), + max_seqlen_q=8, + max_seqlen_k=5, + target_compute_capability=90, + ) + captured = {} + thd._call_kernel = self.fake_call(captured) + result = thd( + _Array((8, 32, 128), self.bfloat16), + _Array((9, 1, 128), self.bfloat16), + _Array((8, 32), self.bfloat16), + cu_seqlens_q=_Array((2,), self.int32), + cu_seqlens_k=_Array((2,), self.int32), + ) + self.assertEqual(result["scores"].shape, (8, 5)) + self.assertEqual(captured["output_descs"][0].shape, (8, 8)) + + def test_forward_maps_independent_public_layouts_to_canonical_axes(self): + operation = self.forward.IndexerForward( + _Array((8, 2, 32, 128), self.bfloat16), + _Array((2, 5, 1, 128), self.bfloat16), + _Array((8, 2, 32), self.bfloat16), + ratio=1, + q_layout="SBHD", + k_layout="BSHD", + w_layout="SBH", + output_layout="SBK", + target_compute_capability=100, + ) + captured = {} + operation._call_kernel = self.fake_call(captured) + + result = operation( + _Array((8, 2, 32, 128), self.bfloat16), + _Array((2, 5, 1, 128), self.bfloat16), + _Array((8, 2, 32), self.bfloat16), + ) + + self.assertEqual(operation.q_desc.shape, (2, 8, 32, 128)) + self.assertEqual(operation.q_desc.stride_order, (3, 2, 0, 1)) + self.assertEqual(result["scores"].shape, (8, 2, 5)) + self.assertEqual( + tuple(spec.mode for spec in captured["input_spec"][:3]), + ((1, 0, 2, 3), (0, 1, 2, 3), (1, 0, 2)), + ) + self.assertEqual( + tuple(spec.layout for spec in captured["input_spec"][:3]), + ((3, 2, 1, 0), (3, 2, 1, 0), (2, 1, 0)), + ) + self.assertEqual(captured["output_spec"][0].mode, (1, 0, 2)) + self.assertEqual(captured["output_spec"][0].layout, (2, 1, 0)) + + explicit_output = self.forward.IndexerForward( + _Array((8, 2, 32, 128), self.bfloat16), + _Array((2, 5, 1, 128), self.bfloat16), + _Array((8, 2, 32), self.bfloat16), + sample_out=_Array((8, 2, 8), self.float32), + ratio=1, + q_layout="SBHD", + k_layout="BSHD", + w_layout="SBH", + output_layout="SBK", + target_compute_capability=100, + ) + self.assertEqual(explicit_output.o_desc.shape, (2, 8, 8)) + self.assertEqual(explicit_output.o_desc.stride_order, (2, 0, 1)) + + def test_forward_rejects_unimplemented_feature_axis_permutations(self): + with self.assertRaisesRegex(ValueError, "q_layout must be one of"): + self.forward.IndexerForward( + _Array((1, 32, 8, 128), self.bfloat16), + _Array((1, 5, 1, 128), self.bfloat16), + _Array((1, 8, 32), self.bfloat16), + q_layout="BHSD", + target_compute_capability=100, + ) + + def test_forward_layout_arguments_are_static(self): + static_argnames = set(self.forward.indexer_forward_wrapper._test_static_argnames) + self.assertTrue({"q_layout", "k_layout", "w_layout", "output_layout"}.issubset(static_argnames)) + + def test_topk_declares_outputs_and_hidden_workspace_for_both_modes(self): + for return_val in (True, False): + with self.subTest(return_val=return_val): + api = self.topk.IndexerTopK( + _Array((4, 64), self.float32), + _Array((4,), self.int32), + 8, + return_val=return_val, + target_compute_capability=100, + ) + captured = {} + api._call_kernel = self.fake_call(captured) + result = api(_Array((4, 64), self.float32), _Array((4,), self.int32)) + self.assertEqual(result["indices"].shape, (4, 8)) + self.assertEqual(result["values"] is not None, return_val) + self.assertEqual(captured["workspace_descs"][0].shape, (4, 2, 64)) + self.assertEqual(len(captured["output_descs"]), 2 if return_val else 1) + + def test_topk_constructor_reports_invalid_inputs_before_inference(self): + with self.assertRaisesRegex(ValueError, "input_values must have rank 2"): + self.topk.IndexerTopK( + _Array((64,), self.float32), + _Array((1,), self.int32), + 8, + target_compute_capability=100, + ) + + with self.assertRaisesRegex(ValueError, "input_values must have dtype"): + self.topk.IndexerTopK( + _Array((4, 64), self.int64), + _Array((4,), self.int32), + 8, + target_compute_capability=100, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_dsa_score_recompute_op.py b/test/python/fe_api/test_dsa_score_recompute_op.py new file mode 100644 index 000000000..0a3661275 --- /dev/null +++ b/test/python/fe_api/test_dsa_score_recompute_op.py @@ -0,0 +1,208 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for DSA score-recompute operations.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_DSA_ROOT = _CUDNN_ROOT / "deepseek_sparse_attention" +_SCORE_ROOT = _DSA_ROOT / "score_recompute" +_PACKAGE = "cudnn_dsa_score_op_test" + + +class _DataType(Enum): + NOT_SET = auto() + FLOAT = auto() + BFLOAT16 = auto() + INT32 = auto() + + +class DsaScoreRecomputeOpTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + dsa_name = f"{_PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_DSA_ROOT)] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + sys.modules[dsa_name] = dsa + + score_name = f"{dsa_name}.score_recompute" + score = types.ModuleType(score_name) + score.__path__ = [str(_SCORE_ROOT)] + score.__package__ = score_name + score.__spec__ = ModuleSpec(score_name, loader=None, is_package=True) + sys.modules[score_name] = score + + try: + cls.tensor = importlib.import_module(f"{_PACKAGE}._tensor_desc") + cls.op = importlib.import_module(f"{score_name}.op") + cls.config = importlib.import_module(f"{score_name}.config") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def desc(self, shape, dtype=_DataType.BFLOAT16, *, stride_order=None, name=""): + return self.tensor.make_compact_tensor_desc( + dtype=dtype, + shape=tuple(shape), + stride_order=stride_order, + name=name, + ) + + def sparse(self, *, target=100, score_type="indexer", with_length=False, **overrides): + aux_dtype = _DataType.BFLOAT16 if score_type == "indexer" else _DataType.FLOAT + arguments = { + "q": self.desc((2, 4, 32, 128), name="q"), + "k": self.desc((2, 256, 128), name="k"), + "per_head": self.desc((2, 4, 32), aux_dtype, name="per_head"), + "topk_indices": self.desc((2, 4, 128), _DataType.INT32, name="topk_indices"), + "output": self.desc((2, 4, 128), _DataType.FLOAT, name="output"), + "topk_length": self.desc((2, 4), _DataType.INT32, name="topk_length") if with_length else None, + "score_type": score_type, + "softmax_scale": 0.125 if score_type == "attention" else 1.0, + "target_compute_capability": target, + } + arguments.update(overrides) + return self.op.SparseScoreRecomputeOp(**arguments) + + def dense(self, *, target=100, score_type="indexer", thd=False, **overrides): + aux_dtype = _DataType.BFLOAT16 if score_type == "indexer" else _DataType.FLOAT + if thd: + q_shape, k_shape, aux_shape = (7, 32, 128), (11, 1, 128), (7, 32) + output_shape, denominator_shape = (7, 8), (7,) + cu_q = self.desc((3,), _DataType.INT32, name="cu_q") + cu_k = self.desc((3,), _DataType.INT32, name="cu_k") + max_q, max_k = 4, 8 + else: + q_shape, k_shape, aux_shape = (2, 4, 32, 128), (2, 8, 1, 128), (2, 4, 32) + output_shape, denominator_shape = (2, 4, 8), (2, 4) + cu_q = cu_k = None + max_q = max_k = None + arguments = { + "q": self.desc(q_shape, name="q"), + "k": self.desc(k_shape, name="k"), + "per_head": self.desc(aux_shape, aux_dtype, name="per_head"), + "output": self.desc(output_shape, _DataType.FLOAT, name="output"), + "denominator": self.desc(denominator_shape, _DataType.FLOAT, name="denominator"), + "score_type": score_type, + "scale": 0.125 if score_type == "attention" else 1.0, + "ratio": 2, + "target_compute_capability": target, + "is_thd": thd, + "cu_seqlens_q": cu_q, + "cu_seqlens_k": cu_k, + "max_seqlen_q": max_q, + "max_seqlen_k": max_k, + } + arguments.update(overrides) + return self.op.DenseScoreRecomputeOp(**arguments) + + def test_sparse_resolves_sm100_and_preserves_explicit_lengths(self): + operation = self.sparse(with_length=True) + self.assertTrue(operation.check_support()) + self.assertEqual(operation.qhead_per_kv_head, 32) + self.assertIsInstance(operation.config, self.config.SparseScoreKernelConfig) + self.assertEqual((operation.config.m_block_size, operation.config.n_block_size), (32, 128)) + self.assertTrue(operation.config.have_topk_length) + + def test_sparse_resolves_sm90(self): + operation = self.sparse(target=90) + self.assertTrue(operation.check_support()) + self.assertIsInstance(operation.config, self.op.SparseScoreSm90Config) + self.assertEqual((operation.config.tile_m, operation.config.tile_n, operation.config.num_threads), (32, 64, 256)) + + def test_dense_bshd_and_thd_resolve_sm100(self): + bshd = self.dense() + self.assertTrue(bshd.check_support()) + self.assertEqual((bshd.max_seqlen_q, bshd.max_seqlen_k), (4, 8)) + self.assertIsInstance(bshd.config, self.config.DenseScoreKernelConfig) + + thd = self.dense(thd=True, score_type="attention") + self.assertTrue(thd.check_support()) + self.assertEqual((thd.max_seqlen_q, thd.max_seqlen_k), (4, 8)) + + def test_dense_sm90_thd_is_explicitly_unsupported(self): + with self.assertRaisesRegex(NotImplementedError, "host-side sequence-length reads"): + self.dense(target=90, thd=True).check_support() + + def test_complete_signatures_are_cross_checked(self): + with self.assertRaisesRegex(ValueError, "output must have shape"): + self.sparse(output=self.desc((2, 4, 64), _DataType.FLOAT)).check_support() + with self.assertRaisesRegex(ValueError, "per-head tensor must have shape"): + self.dense(per_head=self.desc((2, 4, 16))).check_support() + with self.assertRaisesRegex(ValueError, "qhead_per_kv_head must equal"): + self.dense(qhead_per_kv_head=16).check_support() + with self.assertRaisesRegex(ValueError, "ratio must be >= 1"): + self.dense(ratio=0).check_support() + + def test_native_abi_descriptors_must_be_compact(self): + noncompact = self.tensor.TensorDesc( + dtype=_DataType.BFLOAT16, + shape=(2, 4, 32, 128), + stride=(16384, 128, 512, 1), + stride_order=(3, 1, 2, 0), + ) + with self.assertRaisesRegex(ValueError, "Q must be row-major compact"): + self.sparse(q=noncompact).check_support() + + def test_sm100_configuration_rejects_incompatible_topk(self): + with self.assertRaisesRegex(ValueError, "multiple of the selected n_block_size"): + self.sparse( + topk_indices=self.desc((2, 4, 96), _DataType.INT32), + output=self.desc((2, 4, 96), _DataType.FLOAT), + ).check_support() + + def test_rejects_signatures_the_native_kernels_cannot_cover(self): + with self.assertRaisesRegex(ValueError, "SM90.*multiple of 128"): + self.sparse( + target=90, + topk_indices=self.desc((2, 4, 96), _DataType.INT32), + output=self.desc((2, 4, 96), _DataType.FLOAT), + ).check_support() + + with self.assertRaisesRegex(ValueError, "requires MQA with H_kv=1"): + self.dense( + q=self.desc((2, 4, 64, 128)), + k=self.desc((2, 8, 2, 128)), + per_head=self.desc((2, 4, 64)), + ).check_support() + + def test_common_modules_do_not_load_frameworks_or_kernels(self): + self.sparse().check_support() + prefix = f"{_PACKAGE}.deepseek_sparse_attention.score_recompute" + self.assertNotIn(f"{prefix}.jax", sys.modules) + self.assertNotIn(f"{prefix}.api", sys.modules) + self.assertNotIn(f"{prefix}.sparse_score_recompute_sm100", sys.modules) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_dsa_sparse_attention_backward_op.py b/test/python/fe_api/test_dsa_sparse_attention_backward_op.py new file mode 100644 index 000000000..da72104a3 --- /dev/null +++ b/test/python/fe_api/test_dsa_sparse_attention_backward_op.py @@ -0,0 +1,193 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for DSA sparse-attention backward metadata.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "deepseek_sparse_attention" / "sparse_attention_backward" +_PACKAGE = "cudnn_dsa_sparse_attention_backward_op_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + UINT8 = auto() + + +class SparseAttentionBackwardOpContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + dsa_name = f"{_PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_CUDNN_ROOT / "deepseek_sparse_attention")] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + sys.modules[dsa_name] = dsa + + operation_name = f"{dsa_name}.sparse_attention_backward" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + try: + cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}._op") + cls.op_module = importlib.import_module(f"{operation_name}.op") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def _desc(self, shape, dtype=_DataType.BFLOAT16, *, stride=None, name=""): + shape = tuple(shape) + if stride is None: + stride_values = [0] * len(shape) + running = 1 + for dimension in reversed(range(len(shape))): + stride_values[dimension] = running + running *= max(shape[dimension], 1) + stride = tuple(stride_values) + else: + stride = tuple(stride) + stride_order = tuple(index for index, _ in sorted(enumerate(stride), key=lambda item: (item[1], shape[item[0]]))) + return self.tensor_module.TensorDesc( + dtype=dtype, + shape=shape, + stride=stride, + stride_order=stride_order, + name=name, + ) + + def _op(self, *, head_dim=512, dtype=_DataType.BFLOAT16, with_length=True, **overrides): + m, n, h, topk = 65, 130, 64, 32 + head_dim_v = 512 if head_dim == 576 else head_dim + arguments = { + "q": self._desc((m, h, head_dim), dtype, name="q"), + "kv": self._desc((n, head_dim), dtype, name="kv"), + "output": self._desc((m, h, head_dim_v), dtype, name="output"), + "doutput": self._desc((m, h, head_dim_v), dtype, name="doutput"), + "lse": self._desc((m, h), _DataType.FLOAT, name="lse"), + "attn_sink": self._desc((h,), _DataType.FLOAT, name="attn_sink"), + "topk_idxs": self._desc((m, topk), _DataType.INT32, name="topk_idxs"), + "topk_length": self._desc((m,), _DataType.INT32, name="topk_length") if with_length else None, + "dq": self._desc((m, h, head_dim), dtype, name="dq"), + "dkv": self._desc((n, head_dim), dtype, name="dkv"), + "d_sink": self._desc((h,), _DataType.FLOAT, name="d_sink"), + } + arguments.update(overrides) + return self.op_module.SparseAttentionBackwardOp(**arguments) + + def test_validates_both_supported_head_dimensions(self): + for head_dim, expected_v in ((512, 512), (576, 512)): + with self.subTest(head_dim=head_dim): + operation = self._op(head_dim=head_dim) + self.assertIsInstance(operation, self.base_module.Op) + self.assertTrue(operation.check_support()) + self.assertEqual(operation.head_dim_v, expected_v) + self.assertEqual(operation.softmax_scale, 1.0 / head_dim**0.5) + self.assertEqual( + (operation.total_seqlen_q, operation.total_seqlen_kv, operation.num_heads, operation.max_topk), + (65, 130, 64, 32), + ) + + def test_supports_fp16_and_optional_topk_length(self): + operation = self._op(dtype=_DataType.HALF, with_length=False, softmax_scale=0.125) + self.assertTrue(operation.check_support()) + self.assertIsNone(operation.topk_length) + self.assertEqual(operation.softmax_scale, 0.125) + + def test_rejects_incomplete_signatures(self): + cases = ( + ({"q": self._desc((65, 512))}, "Q must have rank 3"), + ({"kv": self._desc((130, 256))}, "KV head dimension must match Q"), + ({"output": self._desc((65, 64, 576))}, "O must have shape"), + ({"doutput": self._desc((65, 64, 576))}, "dO must have shape"), + ({"lse": self._desc((65, 32), _DataType.FLOAT)}, "LSE must have shape"), + ({"attn_sink": self._desc((32,), _DataType.FLOAT)}, "attn_sink must have shape"), + ({"topk_idxs": self._desc((64, 32), _DataType.INT32)}, "topk_idxs must have shape"), + ({"topk_length": self._desc((64,), _DataType.INT32)}, "topk_length must have shape"), + ({"dq": self._desc((64, 64, 512))}, "dQ must have shape"), + ({"dkv": self._desc((130, 256))}, "dKV must have shape"), + ({"d_sink": self._desc((32,), _DataType.FLOAT)}, "d_sink must have shape"), + ({"q": self._desc((65, 64, 512), _DataType.FLOAT)}, "Q must have dtype"), + ({"topk_idxs": self._desc((65, 32), _DataType.FLOAT)}, "topk_idxs must have dtype int32"), + ( + {"q": self._desc((65, 64, 256)), "kv": self._desc((130, 256)), "dq": self._desc((65, 64, 256)), "dkv": self._desc((130, 256))}, + "head_dim must be one of", + ), + ( + { + "q": self._desc((65, 32, 512)), + "output": self._desc((65, 32, 512)), + "doutput": self._desc((65, 32, 512)), + "lse": self._desc((65, 32), _DataType.FLOAT), + "attn_sink": self._desc((32,), _DataType.FLOAT), + "dq": self._desc((65, 32, 512)), + "d_sink": self._desc((32,), _DataType.FLOAT), + }, + "num_heads must be divisible", + ), + ({"q": self._desc((65, 64, 512), stride=(1, 65 * 512, 65))}, "must be row-major contiguous"), + ({"block_tile": 32}, "block_tile must be 64"), + ({"softmax_scale": float("inf")}, "softmax_scale must be finite"), + ) + for overrides, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self._op(**overrides).check_support() + + def test_explicit_scale_preserves_finite_values(self): + for scale in (0.0, -0.125): + with self.subTest(scale=scale): + operation = self._op(softmax_scale=scale) + self.assertTrue(operation.check_support()) + self.assertEqual(operation.softmax_scale, scale) + + def test_operation_imports_no_framework_or_kernel(self): + source = (_OPERATION_ROOT / "op.py").read_text() + for forbidden in ("import torch", "import jax", "import cutlass", "import cuda"): + self.assertNotIn(forbidden, source) + operation = self._op() + operation.check_support() + prefix = f"{_PACKAGE}.deepseek_sparse_attention.sparse_attention_backward" + self.assertNotIn(f"{prefix}.api", sys.modules) + self.assertNotIn(f"{prefix}.jax", sys.modules) + self.assertNotIn(f"{prefix}.dsa_bwd_sm90", sys.modules) + self.assertNotIn(f"{prefix}.dsa_bwd_sm100", sys.modules) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_api_base.py b/test/python/fe_api/test_jax_api_base.py index 63d29bff2..a67b2f636 100644 --- a/test/python/fe_api/test_jax_api_base.py +++ b/test/python/fe_api/test_jax_api_base.py @@ -197,6 +197,61 @@ def test_device_compatibility_checks_can_be_disabled_and_reenabled(self): with self.assertRaisesRegex(TypeError, "disabled must be a bool"): self.module.disable_device_compatibility_checks(1) + def test_resolves_exact_local_target_for_multi_arch_operation(self): + jax = types.ModuleType("jax") + jax.local_devices = lambda *, backend: (_Device(0, "10.3"),) + + with mock.patch.dict(sys.modules, {"jax": jax}): + target = self.module.JaxApiBase._resolve_compute_capability( + None, + (90, 100, 103, 107), + "TestOp", + ) + + self.assertEqual(target, 103) + + def test_rejects_heterogeneous_implicit_targets(self): + jax = types.ModuleType("jax") + jax.local_devices = lambda *, backend: (_Device(0, "10.0"), _Device(1, "10.3")) + + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaisesRegex(RuntimeError, r"heterogeneous targets.*target_compute_capability"): + self.module.JaxApiBase._resolve_compute_capability(None, (90, 100, 103, 107), "TestOp") + + def test_explicit_cross_compile_target_requires_opt_out(self): + jax = types.ModuleType("jax") + jax.local_devices = lambda *, backend: (_Device(0, "10.0"),) + + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaisesRegex(RuntimeError, r"targets SM90.*found CudaDevice\(id=0\) \(SM100\)"): + self.module.JaxApiBase._resolve_compute_capability(90, (90, 100, 103, 107), "TestOp") + + self.module.disable_device_compatibility_checks(True) + with mock.patch.dict(sys.modules, {"jax": None}): + self.assertEqual( + self.module.JaxApiBase._resolve_compute_capability(90, (90, 100), "TestOp"), + 90, + ) + + def test_rejects_unknown_explicit_target_before_lowering(self): + self.module.disable_device_compatibility_checks(True) + + with self.assertRaisesRegex(ValueError, r"no kernel for SM101.*supported targets"): + self.module.JaxApiBase._resolve_compute_capability( + 101, + (90, 100, 103, 107), + "TestOp", + ) + + def test_implicit_target_still_needs_a_device_when_checks_are_disabled(self): + self.module.disable_device_compatibility_checks(True) + jax = types.ModuleType("jax") + jax.local_devices = lambda *, backend: () + + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaisesRegex(RuntimeError, r"no local JAX GPU.*target_compute_capability"): + self.module.JaxApiBase._resolve_compute_capability(None, (90, 100), "TestOp") + def test_converts_array_metadata_to_shared_tensor_desc(self): desc = self.module.JaxApiBase._to_tensor_desc( _ArrayMetadata((2, 3, 4), "bfloat16"), diff --git a/test/python/fe_api/test_jax_dsa_backward.py b/test/python/fe_api/test_jax_dsa_backward.py new file mode 100644 index 000000000..1aa29a7f2 --- /dev/null +++ b/test/python/fe_api/test_jax_dsa_backward.py @@ -0,0 +1,347 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX integration coverage for DSA backward adapters.""" + +import pytest + + +def _jax_runtime(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + pytest.importorskip("cudnn") + return jax, jnp + + +def _supported_gpu(jax, minimum_compute_capability): + supported_targets = {90, 100, 103, 107} + for device in jax.local_devices(): + if device.platform != "gpu": + continue + reported = getattr(device, "compute_capability", None) + if isinstance(reported, (tuple, list)): + major, minor = int(reported[0]), int(reported[1]) + else: + text = str(reported) + if "." in text: + major_text, minor_text = text.split(".", 1) + major, minor = int(major_text), int(minor_text) + else: + try: + major, minor = divmod(int(text), 10) + except ValueError: + continue + capability = major * 10 + minor + if capability >= minimum_compute_capability and capability in supported_targets: + return device, capability + pytest.skip(f"A supported JAX SM{minimum_compute_capability}+ GPU is not available") + + +@pytest.mark.L0 +def test_jax_dsa_backward_abstract_wrappers(monkeypatch): + """Trace every backward signature on a CPU-only JAX installation.""" + + jax, jnp = _jax_runtime() + from cudnn._jax import JaxApiBase, disable_device_compatibility_checks + from cudnn.deepseek_sparse_attention.indexer_backward.jax import ( + dense_indexer_backward_wrapper, + indexer_backward_wrapper, + ) + from cudnn.deepseek_sparse_attention.sparse_attention_backward.jax import sparse_attention_backward_wrapper + + def abstract_call(_self, _inputs, *, output_descs, **_options): + return tuple(jnp.empty(desc.shape, dtype=desc.dtype) for desc in output_descs) + + monkeypatch.setattr(JaxApiBase, "_call_kernel", abstract_call) + disable_device_compatibility_checks(True) + try: + q = jax.ShapeDtypeStruct((1, 2, 64, 128), jnp.bfloat16) + weights = jax.ShapeDtypeStruct((1, 2, 64), jnp.bfloat16) + k = jax.ShapeDtypeStruct((1, 128, 128), jnp.bfloat16) + sparse_score = jax.ShapeDtypeStruct((1, 2, 128), jnp.float32) + topk = jax.ShapeDtypeStruct((1, 2, 128), jnp.int32) + sparse = jax.eval_shape( + lambda q, w, k, target, predict, topk: indexer_backward_wrapper( + q, + w, + k, + target, + predict, + topk, + target_compute_capability=100, + ), + q, + weights, + k, + sparse_score, + sparse_score, + topk, + ) + assert sparse["d_index_q"].shape == q.shape + assert sparse["d_weights"].shape == weights.shape + assert sparse["d_index_k"].shape == k.shape + + q_sparse_thd = jax.ShapeDtypeStruct((2, 64, 128), jnp.bfloat16) + weights_sparse_thd = jax.ShapeDtypeStruct((2, 64), jnp.bfloat16) + k_sparse_thd = jax.ShapeDtypeStruct((128, 128), jnp.bfloat16) + score_sparse_thd = jax.ShapeDtypeStruct((2, 128), jnp.float32) + topk_sparse_thd = jax.ShapeDtypeStruct((2, 128), jnp.int32) + sparse_thd = jax.eval_shape( + lambda q, w, k, target, predict, topk: indexer_backward_wrapper( + q, + w, + k, + target, + predict, + topk, + topk_indices_global=True, + target_compute_capability=100, + ), + q_sparse_thd, + weights_sparse_thd, + k_sparse_thd, + score_sparse_thd, + score_sparse_thd, + topk_sparse_thd, + ) + assert sparse_thd["d_index_q"].shape == q_sparse_thd.shape + assert sparse_thd["d_weights"].shape == weights_sparse_thd.shape + assert sparse_thd["d_index_k"].shape == k_sparse_thd.shape + + dense_score = jax.ShapeDtypeStruct((1, 2, 128), jnp.float32) + dense_denom = jax.ShapeDtypeStruct((1, 2), jnp.float32) + dense = jax.eval_shape( + lambda q, w, k, target, denom, predict, lse: dense_indexer_backward_wrapper( + q, + w, + k, + target, + denom, + predict, + lse, + target_compute_capability=100, + ), + q, + weights, + k, + dense_score, + dense_denom, + dense_score, + dense_denom, + ) + assert dense["d_index_q"].shape == q.shape + assert dense["d_weights"].shape == weights.shape + assert dense["d_index_k"].shape == k.shape + + q_thd = jax.ShapeDtypeStruct((2, 64, 128), jnp.bfloat16) + weights_thd = jax.ShapeDtypeStruct((2, 64), jnp.bfloat16) + k_thd = jax.ShapeDtypeStruct((128, 128), jnp.bfloat16) + score_thd = jax.ShapeDtypeStruct((2, 128), jnp.float32) + denom_thd = jax.ShapeDtypeStruct((2,), jnp.float32) + cu = jax.ShapeDtypeStruct((2,), jnp.int32) + dense_thd = jax.eval_shape( + lambda q, w, k, target, denom, predict, lse, cu_q, cu_k: dense_indexer_backward_wrapper( + q, + w, + k, + target, + denom, + predict, + lse, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=2, + max_seqlen_k=128, + target_compute_capability=90, + ), + q_thd, + weights_thd, + k_thd, + score_thd, + denom_thd, + score_thd, + denom_thd, + cu, + cu, + ) + assert dense_thd["d_index_q"].shape == q_thd.shape + assert dense_thd["d_index_k"].shape == k_thd.shape + + attn_q = jax.ShapeDtypeStruct((2, 64, 512), jnp.bfloat16) + attn_kv = jax.ShapeDtypeStruct((128, 512), jnp.bfloat16) + attn_out = jax.ShapeDtypeStruct((2, 64, 512), jnp.bfloat16) + lse = jax.ShapeDtypeStruct((2, 64), jnp.float32) + sink = jax.ShapeDtypeStruct((64,), jnp.float32) + attn_topk = jax.ShapeDtypeStruct((2, 32), jnp.int32) + attention = jax.eval_shape( + lambda q, kv, out, dout, lse, sink, topk: sparse_attention_backward_wrapper( + q, + kv, + out, + dout, + lse, + sink, + topk, + target_compute_capability=100, + ), + attn_q, + attn_kv, + attn_out, + attn_out, + lse, + sink, + attn_topk, + ) + assert attention["dq"].shape == attn_q.shape + assert attention["dkv"].shape == attn_kv.shape + assert attention["d_sink"].shape == sink.shape + finally: + disable_device_compatibility_checks(False) + + +@pytest.mark.L0 +def test_jax_dense_indexer_backward_sm100_jit_and_numerics(): + jax, jnp = _jax_runtime() + cutlass_jax = pytest.importorskip("cutlass.jax") + if not cutlass_jax.is_available(): + pytest.skip("Installed JAX version is unsupported by CUTLASS JAX") + device, compute_capability = _supported_gpu(jax, 100) + + from cudnn.deepseek_sparse_attention.indexer_backward.jax import dense_indexer_backward_wrapper + + batch, seqlen_q, seqlen_k = 1, 128, 128 + heads, head_dim = 64, 128 + sm_scale = 0.5 + index_q = jax.device_put( + jax.random.normal( + jax.random.key(40), + (batch, seqlen_q, heads, head_dim), + dtype=jnp.bfloat16, + ), + device, + ) + weights = jax.device_put( + jax.random.normal( + jax.random.key(41), + (batch, seqlen_q, heads), + dtype=jnp.bfloat16, + ) + * jnp.bfloat16(1.0 / heads), + device, + ) + index_k = jax.device_put( + jax.random.normal( + jax.random.key(42), + (batch, seqlen_k, head_dim), + dtype=jnp.bfloat16, + ), + device, + ) + valid = jnp.arange(seqlen_k)[None, :] < jnp.arange(seqlen_q)[:, None] + 1 + attn_score = jax.device_put( + jax.random.uniform( + jax.random.key(43), + (batch, seqlen_q, seqlen_k), + dtype=jnp.float32, + minval=0.1, + maxval=1.0, + ), + device, + ) + attn_score = jnp.where(valid[None, :, :], attn_score, 0.0) + attn_l1norm = jnp.sum(attn_score, axis=-1) + scores_per_head = jnp.einsum( + "bqhd,bkd->bqhk", + index_q.astype(jnp.float32), + index_k.astype(jnp.float32), + ) + index_score = jnp.sum( + jnp.maximum(scores_per_head, 0.0) * (weights.astype(jnp.float32) * sm_scale)[..., None], + axis=2, + ) + index_lse = jax.nn.logsumexp( + jnp.where(valid[None, :, :], index_score, -jnp.inf), + axis=-1, + ) + + @jax.jit + def run(q, w, k, target, target_denom, predict, predict_lse, grad_loss): + return dense_indexer_backward_wrapper( + q, + w, + k, + target, + target_denom, + predict, + predict_lse, + grad_loss=grad_loss, + sm_scale=sm_scale, + target_compute_capability=compute_capability, + ) + + grad_loss = jax.device_put(jnp.ones((1,), dtype=jnp.float32), device) + arguments = ( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + grad_loss, + ) + lowered = run.lower(*arguments) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + compiled = lowered.compile() + result = compiled(*arguments) + result["d_index_k"].block_until_ready() + assert result["d_index_q"].shape == index_q.shape + assert result["d_weights"].shape == weights.shape + assert result["d_index_k"].shape == index_k.shape + assert all(jnp.all(jnp.isfinite(value)) for value in result.values()) + + target = attn_score / attn_l1norm[..., None] + predict = jax.nn.softmax( + jnp.where(valid[None, :, :], index_score, -jnp.inf), + axis=-1, + ) + grad_scale = 1.0 / (batch * seqlen_q) + g = jnp.where( + valid[None, :, :], + -jnp.maximum(target, jnp.exp(-100.0)) * grad_scale, + 0.0, + ) + grad_signal = jnp.where( + valid[None, :, :], + g - predict * jnp.sum(g, axis=-1, keepdims=True), + 0.0, + ) + + def raw_scores(q, w, k): + per_head = jnp.einsum("bqhd,bkd->bqhk", q, k) + return jnp.sum( + jnp.maximum(per_head, 0.0) * (w * sm_scale)[..., None], + axis=2, + ) + + _, pullback = jax.vjp( + raw_scores, + index_q.astype(jnp.float32), + weights.astype(jnp.float32), + index_k.astype(jnp.float32), + ) + reference = pullback(grad_signal) + for name, expected in zip(("d_index_q", "d_weights", "d_index_k"), reference): + actual_f32 = result[name].astype(jnp.float32).ravel() + expected_f32 = expected.astype(jnp.float32).ravel() + cosine = jnp.vdot(actual_f32, expected_f32) / (jnp.linalg.norm(actual_f32) * jnp.linalg.norm(expected_f32)) + relative_rms = jnp.sqrt(jnp.mean(jnp.square(actual_f32 - expected_f32))) / jnp.sqrt(jnp.mean(jnp.square(expected_f32))) + assert cosine >= 0.97, f"{name} cosine similarity: {cosine}" + assert relative_rms <= 0.55, f"{name} relative RMS error: {relative_rms}" + + zero_result = compiled(*arguments[:-1], jnp.zeros_like(grad_loss)) + zero_result["d_index_k"].block_until_ready() + assert all(jnp.array_equal(value, jnp.zeros_like(value)) for value in zero_result.values()) diff --git a/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py b/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py new file mode 100644 index 000000000..7d900a59e --- /dev/null +++ b/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py @@ -0,0 +1,551 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free JAX contracts for DSA indexer backward.""" + +from __future__ import annotations + +import ast +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_DSA_ROOT = _CUDNN_ROOT / "deepseek_sparse_attention" +_OPERATION_ROOT = _DSA_ROOT / "indexer_backward" +_PACKAGE = "cudnn_jax_dsa_indexer_backward_test" + + +class _DataType(Enum): + NOT_SET = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + + +class _DType: + def __init__(self, name: str): + self.name = name + + def __repr__(self): + return self.name + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + def astype(self, dtype): + return _Array(self.shape, dtype) + + +class _TupleDict(dict): + def __iter__(self): + return iter(self.values()) + + +class _TensorSpec: + def __init__(self, *, mode=None, divisibility=None): + self.mode = mode + self.divisibility = divisibility + + +class _Scalar: + def __init__(self, value): + self.value = value + + +class _CuteTensor: + def __init__(self, shape, stride): + self.shape = tuple(shape) + self.stride = tuple(stride) + self.iterator = self + + +class _RecordedKernel: + events = [] + + def __init__(self, label, *configuration, **keywords): + self.label = label + self.configuration = (configuration, keywords) + + def __call__(self, *arguments): + self.events.append((self.label, self.configuration, arguments)) + + +class JaxIndexerBackwardContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.bfloat16 = _DType("bfloat16") + cls.float32 = _DType("float32") + cls.int32 = _DType("int32") + cls.dtype_to_cudnn = { + cls.bfloat16: _DataType.BFLOAT16, + cls.float32: _DataType.FLOAT, + cls.int32: _DataType.INT32, + } + cls.cudnn_to_dtype = {value: key for key, value in cls.dtype_to_cudnn.items()} + + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + dsa_name = f"{_PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_DSA_ROOT)] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + sys.modules[dsa_name] = dsa + + operation_name = f"{dsa_name}.indexer_backward" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + + class JaxTensorDesc(tensor_module.TensorDesc): + @property + def cudnn_dtype(self): + return cls.dtype_to_cudnn.get(self.dtype, _DataType.NOT_SET) + + def compact_like(self, *, cudnn_dtype, shape, stride_order=None, name="", init_value=None): + canonical = tensor_module.make_compact_tensor_desc( + dtype=cudnn_dtype, + shape=shape, + stride_order=stride_order, + name=name, + init_value=init_value, + ) + return JaxTensorDesc( + dtype=cls.cudnn_to_dtype[cudnn_dtype], + shape=canonical.shape, + stride=canonical.stride, + stride_order=canonical.stride_order, + name=name, + init_value=init_value, + ) + + class JaxApiBase: + captured_call = None + + @staticmethod + def _to_tensor_desc(value, name, *, mode=None, init_value=None): + del mode + shape = tuple(value.shape) + stride = [0] * len(shape) + running = 1 + for dimension in reversed(range(len(shape))): + stride[dimension] = running + running *= max(shape[dimension], 1) + return JaxTensorDesc( + dtype=value.dtype, + shape=shape, + stride=tuple(stride), + stride_order=tuple(reversed(range(len(shape)))), + name=name, + init_value=init_value, + ) + + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del supported, operation_name + if target is None: + raise RuntimeError("target required by dependency-free test") + return target + + @staticmethod + def _compute_capability_family(target, supported): + return max((value for value in supported if value <= target), default=None) + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + del mode + actual_dtype = cls.dtype_to_cudnn.get(value.dtype, _DataType.NOT_SET) + if tuple(value.shape) != expected.shape or actual_dtype != expected.cudnn_dtype: + raise ValueError(f"{expected.name} tensor signature mismatch") + + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + del desc + return _TensorSpec(mode=mode, divisibility=divisibility) + + def _call_kernel(self, inputs, **options): + JaxApiBase.captured_call = (inputs, options) + return tuple(_Array(desc.shape, cls.cudnn_to_dtype[desc.cudnn_dtype]) for desc in options["output_descs"]) + + internal = types.ModuleType(f"{_PACKAGE}._jax") + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = _TupleDict + sys.modules[internal.__name__] = internal + + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.bfloat16 = cls.bfloat16 + fake_jnp.float32 = cls.float32 + fake_jnp.int32 = cls.int32 + fake_jnp.dtype = lambda value: value.dtype if hasattr(value, "dtype") else value + fake_jnp.asarray = lambda value, dtype: _Array(getattr(value, "shape", ()), dtype) + fake_jnp.reshape = lambda value, shape: _Array(shape, value.dtype) + + fake_jax = types.ModuleType("jax") + fake_jax.__path__ = [] + fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + fake_jax.numpy = fake_jnp + fake_jax.ShapeDtypeStruct = _Array + fake_jax.jit = lambda function=None, **_kwargs: (lambda fn: fn) if function is None else function + + try: + with mock.patch.dict(sys.modules, {"jax": fake_jax, "jax.numpy": fake_jnp}): + cls.module = importlib.import_module(f"{operation_name}.jax") + except Exception: + cls.tearDownClass() + raise + cls.JaxApiBase = JaxApiBase + cls.fake_jax_modules = {"jax": fake_jax, "jax.numpy": fake_jnp} + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def setUp(self): + self.JaxApiBase.captured_call = None + _RecordedKernel.events.clear() + + @classmethod + def _sparse_inputs(cls, *, topk=128): + score_shape = (2, 64, topk) + return ( + _Array((2, 64, 64, 128), cls.bfloat16), + _Array((2, 64, 64), cls.bfloat16), + _Array((2, 256, 128), cls.bfloat16), + _Array(score_shape, cls.float32), + _Array(score_shape, cls.float32), + _Array(score_shape, cls.int32), + ) + + @classmethod + def _sparse_thd_inputs(cls, *, topk=128): + score_shape = (96, topk) + return ( + _Array((96, 64, 128), cls.bfloat16), + _Array((96, 64), cls.bfloat16), + _Array((384, 128), cls.bfloat16), + _Array(score_shape, cls.float32), + _Array(score_shape, cls.float32), + _Array(score_shape, cls.int32), + ) + + @classmethod + def _dense_bshd_inputs(cls): + return ( + _Array((2, 64, 64, 128), cls.bfloat16), + _Array((2, 64, 64), cls.bfloat16), + _Array((2, 256, 128), cls.bfloat16), + _Array((2, 64, 256), cls.float32), + _Array((2, 64), cls.float32), + _Array((2, 64, 256), cls.float32), + _Array((2, 64), cls.float32), + ) + + @classmethod + def _dense_thd_inputs(cls): + return ( + _Array((96, 64, 128), cls.bfloat16), + _Array((96, 64), cls.bfloat16), + _Array((384, 128), cls.bfloat16), + _Array((96, 256), cls.float32), + _Array((96,), cls.float32), + _Array((96, 256), cls.float32), + _Array((96,), cls.float32), + ) + + def test_sparse_wrapper_declares_functional_outputs_and_zeroed_dk(self): + inputs = self._sparse_inputs() + with mock.patch.dict(sys.modules, self.fake_jax_modules): + result = self.module.indexer_backward_wrapper(*inputs, target_compute_capability=103) + + self.assertEqual(tuple(result), (result["d_index_q"], result["d_weights"], result["d_index_k"])) + self.assertEqual(result["d_index_q"].shape, inputs[0].shape) + self.assertEqual(result["d_weights"].shape, inputs[1].shape) + self.assertEqual(result["d_index_k"].shape, inputs[2].shape) + kernel_inputs, options = self.JaxApiBase.captured_call + self.assertEqual(kernel_inputs[:-1], inputs) + self.assertEqual((kernel_inputs[-1].shape, kernel_inputs[-1].dtype), ((1,), self.float32)) + self.assertEqual(options["output_descs"][2].init_value, 0.0) + self.assertEqual(options["workspace_descs"][0].shape, inputs[3].shape) + self.assertIn("--gpu-arch sm_103a", options["compile_options"]) + self.assertEqual(options["input_spec"][0].divisibility, (None, None, None, 128)) + self.assertEqual(options["input_spec"][2].divisibility, (None, None, 128)) + + def test_sparse_packed_thd_signature_requires_global_indices(self): + inputs = self._sparse_thd_inputs() + with mock.patch.dict(sys.modules, self.fake_jax_modules): + with self.assertRaisesRegex(ValueError, "requires topk_indices_global=True"): + self.module.indexer_backward_wrapper( + *inputs, + target_compute_capability=100, + ) + result = self.module.indexer_backward_wrapper( + *inputs, + topk_indices_global=True, + target_compute_capability=100, + ) + + self.assertEqual(result["d_index_q"].shape, inputs[0].shape) + self.assertEqual(result["d_weights"].shape, inputs[1].shape) + self.assertEqual(result["d_index_k"].shape, inputs[2].shape) + _, options = self.JaxApiBase.captured_call + self.assertEqual(options["input_spec"][0].divisibility, (None, None, 128)) + self.assertEqual(options["input_spec"][2].divisibility, (None, 128)) + + def test_grad_loss_uses_a_fixed_runtime_abi_descriptor(self): + inputs = self._sparse_inputs() + with mock.patch.dict(sys.modules, self.fake_jax_modules): + api = self.module.IndexerBackward( + *inputs, + target_compute_capability=100, + ) + self.assertEqual(api.grad_loss_desc.shape, (1,)) + self.assertIs(api.grad_loss_desc.dtype, self.float32) + + def test_dense_bshd_declares_functional_outputs_and_runtime_grad_loss(self): + inputs = self._dense_bshd_inputs() + grad_loss = _Array((1,), self.float32) + with mock.patch.dict(sys.modules, self.fake_jax_modules): + result = self.module.dense_indexer_backward_wrapper( + *inputs, + grad_loss=grad_loss, + target_compute_capability=100, + ) + + self.assertEqual(result["d_index_q"].shape, inputs[0].shape) + self.assertEqual(result["d_weights"].shape, inputs[1].shape) + self.assertEqual(result["d_index_k"].shape, inputs[2].shape) + kernel_inputs, options = self.JaxApiBase.captured_call + self.assertEqual((kernel_inputs[7].shape, kernel_inputs[7].dtype), ((1,), self.float32)) + self.assertEqual(options["output_descs"][2].init_value, 0.0) + self.assertEqual(options["workspace_descs"][0].shape, inputs[3].shape) + self.assertIn("--gpu-arch sm_100a", options["compile_options"]) + + def test_dense_thd_signature_and_optional_runtime_operands(self): + inputs = self._dense_thd_inputs() + cu_q = _Array((3,), self.int32) + cu_k = _Array((3,), self.int32) + offsets = _Array((2,), self.int32) + with mock.patch.dict(sys.modules, self.fake_jax_modules): + api = self.module.DenseIndexerBackward( + *inputs, + sample_cu_seqlens_q=cu_q, + sample_cu_seqlens_k=cu_k, + sample_q_causal_offsets=offsets, + max_seqlen_q=64, + max_seqlen_k=256, + target_compute_capability=90, + ) + result = api( + *inputs, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + q_causal_offsets=offsets, + ) + + self.assertTrue(api._op.is_thd) + self.assertEqual(result["d_index_q"].shape, inputs[0].shape) + kernel_inputs, options = self.JaxApiBase.captured_call + self.assertEqual(kernel_inputs[-3:], (cu_q, cu_k, offsets)) + self.assertEqual(options["input_spec"][0].divisibility, (None, None, 128)) + self.assertEqual(options["input_spec"][2].divisibility, (None, 128)) + self.assertIn("--gpu-arch sm_90a", options["compile_options"]) + + def _kernel_modules(self): + package = self.module.__package__ + + def kernel(label): + return lambda *args, **kwargs: _RecordedKernel(label, *args, **kwargs) + + sparse_100 = types.ModuleType(f"{package}.indexer_backward_sm100") + sparse_100.ScoreGradSm100 = kernel("sparse_score_100") + sparse_100.IndexerBackwardSm100 = kernel("sparse_gemm_100") + sparse_90 = types.ModuleType(f"{package}.indexer_backward_sm90") + sparse_90.ScoreGradSm90 = kernel("sparse_score_90") + sparse_90.IndexerBackwardSm90 = kernel("sparse_gemm_90") + dense_100 = types.ModuleType(f"{package}.dense_indexer_backward_sm100") + dense_100.ScoreGradDense = kernel("dense_score_100") + dense_100.DenseIndexerBackward2QGemmSm100 = kernel("dense_gemm_100") + dense_90 = types.ModuleType(f"{package}.dense_indexer_backward_sm90") + dense_90.ScoreGradDenseSm90 = kernel("dense_score_90") + return {module.__name__: module for module in (sparse_100, sparse_90, dense_100, dense_90)} + + @staticmethod + def _fake_cutlass(): + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass.Float32 = _Scalar + cutlass.Int32 = _Scalar + return cutlass + + @staticmethod + def _fake_cute(): + cute = types.ModuleType("cutlass.cute") + cute.make_layout = lambda shape, stride: types.SimpleNamespace( + shape=tuple(shape), + stride=tuple(stride), + ) + cute.make_tensor = lambda _iterator, layout: _CuteTensor(layout.shape, layout.stride) + return cute + + @staticmethod + def _compact_cute(shape): + stride = [0] * len(shape) + running = 1 + for dimension in reversed(range(len(shape))): + stride[dimension] = running + running *= shape[dimension] + return _CuteTensor(shape, stride) + + def test_sparse_sm100_and_sm90_launch_functionally(self): + inputs = self._sparse_inputs() + kernel_inputs = tuple(object() for _ in range(7)) + outputs = tuple(object() for _ in range(3)) + workspace = (object(),) + modules = {**self.fake_jax_modules, "cutlass": self._fake_cutlass(), **self._kernel_modules()} + + for target, labels in ((100, ("sparse_score_100", "sparse_gemm_100")), (90, ("sparse_score_90", "sparse_gemm_90"))): + with self.subTest(target=target), mock.patch.dict(sys.modules, modules): + _RecordedKernel.events.clear() + api = self.module.IndexerBackward(*inputs, target_compute_capability=target) + api.check_support() + api._launch(kernel_inputs, outputs, workspace, object()) + self.assertEqual(tuple(event[0] for event in _RecordedKernel.events), labels) + score_args = _RecordedKernel.events[0][2] + self.assertIs(score_args[-2 if target == 100 else -1], workspace[0]) + gemm_args = _RecordedKernel.events[1][2] + self.assertIs(gemm_args[6], workspace[0]) + + def test_sparse_packed_thd_launch_adds_a_synthetic_batch_view(self): + inputs = self._sparse_thd_inputs() + cute = self._fake_cute() + cutlass = self._fake_cutlass() + cutlass.cute = cute + modules = { + **self.fake_jax_modules, + "cutlass": cutlass, + "cutlass.cute": cute, + **self._kernel_modules(), + } + kernel_inputs = tuple(self._compact_cute(value.shape) for value in inputs) + (self._compact_cute((1,)),) + outputs = ( + self._compact_cute(inputs[0].shape), + self._compact_cute(inputs[1].shape), + self._compact_cute(inputs[2].shape), + ) + workspace = (self._compact_cute(inputs[3].shape),) + + for target, labels in ((100, ("sparse_score_100", "sparse_gemm_100")), (90, ("sparse_score_90", "sparse_gemm_90"))): + with self.subTest(target=target), mock.patch.dict(sys.modules, modules): + _RecordedKernel.events.clear() + api = self.module.IndexerBackward( + *inputs, + topk_indices_global=True, + target_compute_capability=target, + ) + api.check_support() + api._launch(kernel_inputs, outputs, workspace, object()) + + self.assertEqual(tuple(event[0] for event in _RecordedKernel.events), labels) + score_args = _RecordedKernel.events[0][2] + self.assertEqual(score_args[0].shape, (1, 96, 128)) + self.assertEqual(score_args[1].shape, (1, 96, 128)) + self.assertEqual(score_args[5].shape, (1, 96, 128)) + gemm_args = _RecordedKernel.events[1][2] + self.assertEqual(gemm_args[0].shape, (1, 96, 64, 128)) + self.assertEqual(gemm_args[1].shape, (1, 96, 64)) + self.assertEqual(gemm_args[2].shape, (1, 384, 128)) + self.assertEqual(gemm_args[7].shape, (1, 96, 128)) + + def test_dense_sm100_bshd_and_sm90_thd_launch_functionally(self): + modules = {**self.fake_jax_modules, "cutlass": self._fake_cutlass(), **self._kernel_modules()} + outputs = tuple(object() for _ in range(3)) + workspace = (object(),) + + bshd = self._dense_bshd_inputs() + with mock.patch.dict(sys.modules, modules): + api = self.module.DenseIndexerBackward(*bshd, target_compute_capability=100) + api.check_support() + kernel_inputs = tuple(object() for _ in range(8)) + api._launch(kernel_inputs, outputs, workspace, object()) + self.assertEqual(tuple(event[0] for event in _RecordedKernel.events), ("dense_score_100", "dense_gemm_100")) + score_args = _RecordedKernel.events[0][2] + self.assertIs(score_args[11], workspace[0]) + self.assertIs(score_args[12], kernel_inputs[7]) + + _RecordedKernel.events.clear() + thd = self._dense_thd_inputs() + cu_q = _Array((3,), self.int32) + cu_k = _Array((3,), self.int32) + offsets = _Array((2,), self.int32) + with mock.patch.dict(sys.modules, modules): + api = self.module.DenseIndexerBackward( + *thd, + sample_cu_seqlens_q=cu_q, + sample_cu_seqlens_k=cu_k, + sample_q_causal_offsets=offsets, + max_seqlen_q=64, + max_seqlen_k=256, + target_compute_capability=90, + ) + api.check_support() + kernel_inputs = tuple(object() for _ in range(11)) + api._launch(kernel_inputs, outputs, workspace, object()) + self.assertEqual(tuple(event[0] for event in _RecordedKernel.events), ("dense_score_90", "sparse_gemm_90")) + score_args = _RecordedKernel.events[0][2] + self.assertEqual(score_args[4:7], kernel_inputs[8:11]) + self.assertIs(score_args[11], workspace[0]) + self.assertIs(score_args[12], kernel_inputs[7]) + gemm_args = _RecordedKernel.events[1][2] + self.assertEqual(gemm_args[10:12], kernel_inputs[8:10]) + self.assertIs(gemm_args[14], kernel_inputs[10]) + + def test_jax_reachable_modules_do_not_import_torch_at_module_scope(self): + for filename in ( + "jax.py", + "indexer_backward_sm90.py", + "indexer_backward_sm100.py", + "dense_indexer_backward_sm90.py", + "dense_indexer_backward_sm100.py", + ): + path = _OPERATION_ROOT / filename + tree = ast.parse(path.read_text(), filename=str(path)) + imports = [] + for node in tree.body: + if isinstance(node, ast.Import): + imports.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + imports.append(node.module or "") + self.assertNotIn("torch", imports, filename) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_dsa_indexers.py b/test/python/fe_api/test_jax_dsa_indexers.py new file mode 100644 index 000000000..36ef3034e --- /dev/null +++ b/test/python/fe_api/test_jax_dsa_indexers.py @@ -0,0 +1,278 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX integration coverage for DSA indexer forward and top-K.""" + +import pytest + + +def _jax_runtime(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + cutlass_jax = pytest.importorskip("cutlass.jax") + if not cutlass_jax.is_available(): + pytest.skip("Installed JAX version is unsupported by CUTLASS JAX") + return jax, jnp + + +def _supported_gpu(jax, minimum=90): + supported_targets = {90, 100, 103, 107} + for device in jax.local_devices(): + if device.platform != "gpu": + continue + reported = getattr(device, "compute_capability", None) + if isinstance(reported, (tuple, list)): + major, minor = int(reported[0]), int(reported[1]) + else: + text = str(reported) + if "." not in text: + continue + major_text, minor_text = text.split(".", 1) + major, minor = int(major_text), int(minor_text) + capability = major * 10 + minor + if capability >= minimum and capability in supported_targets: + return device, capability + pytest.skip(f"A supported JAX SM{minimum}+ GPU is not available") + + +@pytest.fixture +def disabled_device_checks(): + _jax_runtime() + from cudnn.jax import disable_device_compatibility_checks + + disable_device_compatibility_checks(True) + try: + yield + finally: + disable_device_compatibility_checks(False) + + +@pytest.mark.L0 +def test_jax_dsa_indexer_abstract_shapes(disabled_device_checks): + jax, jnp = _jax_runtime() + from cudnn.jax import ( + compactify_wrapper, + indexer_forward_wrapper, + indexer_top_k_wrapper, + local_to_global_wrapper, + ) + + q = jax.ShapeDtypeStruct((1, 4, 32, 128), jnp.bfloat16) + k = jax.ShapeDtypeStruct((1, 5, 1, 128), jnp.bfloat16) + w = jax.ShapeDtypeStruct((1, 4, 32), jnp.bfloat16) + fixed = jax.eval_shape( + lambda q, k, w: indexer_forward_wrapper(q, k, w, ratio=1, target_compute_capability=100), + q, + k, + w, + )["scores"] + assert fixed.shape == (1, 4, 5) + assert fixed.dtype == jnp.float32 + + q_sequence_major = jax.ShapeDtypeStruct((4, 2, 32, 128), jnp.bfloat16) + k_sequence_major = jax.ShapeDtypeStruct((5, 2, 1, 128), jnp.bfloat16) + w_sequence_major = jax.ShapeDtypeStruct((4, 2, 32), jnp.bfloat16) + sequence_major = jax.eval_shape( + lambda q, k, w: indexer_forward_wrapper( + q, + k, + w, + ratio=1, + q_layout="SBHD", + k_layout="SBHD", + w_layout="SBH", + output_layout="SBK", + target_compute_capability=100, + ), + q_sequence_major, + k_sequence_major, + w_sequence_major, + )["scores"] + assert sequence_major.shape == (4, 2, 5) + assert sequence_major.dtype == jnp.float32 + + q_thd = jax.ShapeDtypeStruct((4, 32, 128), jnp.bfloat16) + k_thd = jax.ShapeDtypeStruct((5, 1, 128), jnp.bfloat16) + w_thd = jax.ShapeDtypeStruct((4, 32), jnp.bfloat16) + cu = jax.ShapeDtypeStruct((2,), jnp.int32) + packed = jax.eval_shape( + lambda q, k, w, cu_q, cu_k: indexer_forward_wrapper( + q, + k, + w, + ratio=1, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=4, + max_seqlen_k=5, + target_compute_capability=100, + ), + q_thd, + k_thd, + w_thd, + cu, + cu, + )["scores"] + assert packed.shape == (4, 5) + + values = jax.ShapeDtypeStruct((2, 64), jnp.float32) + lengths = jax.ShapeDtypeStruct((2,), jnp.int32) + with_values = jax.eval_shape( + lambda values, lengths: indexer_top_k_wrapper( + values, + lengths, + 8, + target_compute_capability=100, + ), + values, + lengths, + ) + assert with_values["indices"].shape == (2, 8) + assert with_values["indices"].dtype == jnp.int32 + assert with_values["values"].shape == (2, 8) + without_values = jax.eval_shape( + lambda values, lengths: indexer_top_k_wrapper( + values, + lengths, + 8, + return_val=False, + target_compute_capability=100, + ), + values, + lengths, + ) + assert without_values["indices"].shape == (2, 8) + assert without_values["values"] is None + + local = jax.ShapeDtypeStruct((1, 2, 8), jnp.int64) + global_indices = jax.eval_shape( + lambda indices: local_to_global_wrapper(indices, 64, target_compute_capability=100), + local, + )["indices"] + assert global_indices.shape == local.shape + assert global_indices.dtype == jnp.int32 + packed_global = jax.eval_shape( + lambda indices, cu_q, cu_k: local_to_global_wrapper( + indices, + 64, + cu_q, + cu_k, + target_compute_capability=100, + ), + jax.ShapeDtypeStruct((2, 8), jnp.int32), + jax.ShapeDtypeStruct((2,), jnp.int32), + jax.ShapeDtypeStruct((2,), jnp.int32), + )["indices"] + assert packed_global.shape == (2, 8) + + compact = jax.eval_shape( + lambda indices: compactify_wrapper(indices, target_compute_capability=100), + jax.ShapeDtypeStruct((1, 2, 8), jnp.int32), + ) + assert compact["indices"].shape == (2, 8) + assert compact["topk_length"].shape == (2,) + + +@pytest.mark.L0 +def test_jax_indexer_forward_jit_and_numerics(disabled_device_checks): + jax, jnp = _jax_runtime() + device, capability = _supported_gpu(jax, 90) + from cudnn.jax import indexer_forward_wrapper + + batch, seqlen_q, seqlen_k = 2, 4, 5 + num_query_heads, head_dim = 32, 128 + q = jax.device_put( + jax.random.normal(jax.random.key(0), (batch, seqlen_q, num_query_heads, head_dim), dtype=jnp.bfloat16), + device, + ) + k = jax.device_put( + jax.random.normal(jax.random.key(1), (batch, seqlen_k, 1, head_dim), dtype=jnp.bfloat16), + device, + ) + w = jax.device_put( + jax.random.normal(jax.random.key(2), (batch, seqlen_q, num_query_heads), dtype=jnp.bfloat16), + device, + ) + + lowered = indexer_forward_wrapper.lower( + q, + k, + w, + ratio=1, + target_compute_capability=capability, + ) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + scores = lowered.compile()(q, k, w)["scores"] + scores.block_until_ready() + + expanded_k = jnp.repeat(k.astype(jnp.float32), num_query_heads, axis=2) + per_head = jnp.einsum("bqhd,bkhd->bqhk", q.astype(jnp.float32), expanded_k) + expected = jnp.sum(jnp.maximum(per_head, 0.0) * w.astype(jnp.float32)[..., None], axis=2) + valid = jnp.arange(seqlen_k)[None, :] < (jnp.arange(seqlen_q)[:, None] + 1) + expected = jnp.where(valid[None, :, :], expected, -jnp.inf) + assert jnp.array_equal(jnp.isneginf(scores), jnp.isneginf(expected)) + finite = jnp.isfinite(expected) + assert jnp.allclose(scores[finite], expected[finite], atol=5e-2, rtol=2e-2) + + q_sequence_major = jnp.transpose(q, (1, 0, 2, 3)) + k_sequence_major = jnp.transpose(k, (1, 0, 2, 3)) + w_sequence_major = jnp.transpose(w, (1, 0, 2)) + sequence_major_lowered = indexer_forward_wrapper.lower( + q_sequence_major, + k_sequence_major, + w_sequence_major, + ratio=1, + q_layout="SBHD", + k_layout="SBHD", + w_layout="SBH", + output_layout="SBK", + target_compute_capability=capability, + ) + sequence_major_scores = sequence_major_lowered.compile()( + q_sequence_major, + k_sequence_major, + w_sequence_major, + )["scores"] + sequence_major_scores.block_until_ready() + expected_sequence_major = jnp.transpose(expected, (1, 0, 2)) + assert jnp.array_equal( + jnp.isneginf(sequence_major_scores), + jnp.isneginf(expected_sequence_major), + ) + finite = jnp.isfinite(expected_sequence_major) + assert jnp.allclose( + sequence_major_scores[finite], + expected_sequence_major[finite], + atol=5e-2, + rtol=2e-2, + ) + + +@pytest.mark.L0 +def test_jax_indexer_top_k_jit_and_numerics(disabled_device_checks): + jax, jnp = _jax_runtime() + device, capability = _supported_gpu(jax, 90) + from cudnn.jax import indexer_top_k_wrapper + + num_rows, num_cols, top_k = 4, 64, 8 + row = jnp.arange(num_cols, dtype=jnp.float32) + input_values = jax.device_put(jnp.broadcast_to(row, (num_rows, num_cols)), device) + seq_lens = jax.device_put(jnp.asarray((64, 63, 62, 61), dtype=jnp.int32), device) + lowered = indexer_top_k_wrapper.lower( + input_values, + seq_lens, + top_k, + target_compute_capability=capability, + ) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + result = lowered.compile()(input_values, seq_lens) + result["indices"].block_until_ready() + + valid = jnp.arange(num_cols)[None, :] < seq_lens[:, None] + reference_values, reference_indices = jax.lax.top_k(jnp.where(valid, input_values, -jnp.inf), top_k) + assert jnp.array_equal(jnp.sort(result["indices"], axis=-1), jnp.sort(reference_indices, axis=-1)) + assert jnp.array_equal(jnp.sort(result["values"], axis=-1), jnp.sort(reference_values, axis=-1)) diff --git a/test/python/fe_api/test_jax_dsa_score_recompute.py b/test/python/fe_api/test_jax_dsa_score_recompute.py new file mode 100644 index 000000000..c1db33783 --- /dev/null +++ b/test/python/fe_api/test_jax_dsa_score_recompute.py @@ -0,0 +1,360 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX integration tests for DSA score-recompute kernels.""" + +import pytest + + +def _jax_runtime(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + cutlass_jax = pytest.importorskip("cutlass.jax") + if not cutlass_jax.is_available(): + pytest.skip("Installed JAX version is unsupported by CUTLASS JAX") + return jax, jnp + + +def _sm100_device(jax): + supported = {100, 103, 107} + try: + devices = tuple(jax.local_devices(backend="gpu")) + except RuntimeError as error: + pytest.skip(f"A JAX GPU is not available: {error}") + if not devices: + pytest.skip("A JAX GPU is not available") + capabilities = [] + for device in devices: + reported = getattr(device, "compute_capability", None) + try: + if isinstance(reported, (tuple, list)): + major, minor = (int(value) for value in reported[:2]) + capability = major * 10 + minor + elif "." in str(reported): + major, minor = (int(value) for value in str(reported).split(".", 1)) + capability = major * 10 + minor + else: + capability = int(reported) + if capability < 10: + capability *= 10 + except (TypeError, ValueError): + pytest.skip(f"JAX reported an unsupported compute capability {reported!r}") + capabilities.append(capability) + if len(set(capabilities)) != 1 or capabilities[0] not in supported: + pytest.skip(f"Score recompute requires homogeneous SM100/SM103/SM107 devices, found {capabilities}") + return devices[0], capabilities[0] + + +@pytest.fixture +def disabled_device_checks(): + _jax_runtime() + from cudnn.jax import disable_device_compatibility_checks + + disable_device_compatibility_checks(True) + try: + yield + finally: + disable_device_compatibility_checks(False) + + +@pytest.mark.L0 +def test_jax_score_recompute_eval_shapes(disabled_device_checks): + jax, jnp = _jax_runtime() + from cudnn.jax import ( + dense_attn_score_recompute_wrapper, + dense_indexer_score_recompute_wrapper, + sparse_attn_score_recompute_wrapper, + sparse_indexer_score_recompute_wrapper, + ) + + sparse_q = jax.ShapeDtypeStruct((1, 4, 32, 128), jnp.bfloat16) + sparse_k = jax.ShapeDtypeStruct((1, 128, 128), jnp.bfloat16) + sparse_weights = jax.ShapeDtypeStruct((1, 4, 32), jnp.bfloat16) + sparse_lse = jax.ShapeDtypeStruct((1, 4, 32), jnp.float32) + sparse_indices = jax.ShapeDtypeStruct((1, 4, 128), jnp.int32) + predict = jax.eval_shape( + lambda q, k, weights, indices: sparse_indexer_score_recompute_wrapper( + q, + k, + weights, + indices, + target_compute_capability=100, + ), + sparse_q, + sparse_k, + sparse_weights, + sparse_indices, + )["predict"] + target = jax.eval_shape( + lambda q, k, lse, indices: sparse_attn_score_recompute_wrapper( + q, + k, + lse, + indices, + softmax_scale=0.125, + target_compute_capability=100, + ), + sparse_q, + sparse_k, + sparse_lse, + sparse_indices, + )["target"] + assert (predict.shape, predict.dtype) == ((1, 4, 128), jnp.float32) + assert (target.shape, target.dtype) == ((1, 4, 128), jnp.float32) + + dense_q = jax.ShapeDtypeStruct((1, 4, 32, 128), jnp.bfloat16) + dense_k = jax.ShapeDtypeStruct((1, 8, 1, 128), jnp.bfloat16) + dense_weights = jax.ShapeDtypeStruct((1, 4, 32), jnp.bfloat16) + dense_lse = jax.ShapeDtypeStruct((1, 4, 32), jnp.float32) + dense_indexer = jax.eval_shape( + lambda q, k, weights: dense_indexer_score_recompute_wrapper(q, k, weights, target_compute_capability=100), + dense_q, + dense_k, + dense_weights, + ) + dense_attention = jax.eval_shape( + lambda q, k, lse: dense_attn_score_recompute_wrapper( + q, + k, + lse, + softmax_scale=0.125, + target_compute_capability=100, + ), + dense_q, + dense_k, + dense_lse, + ) + for result in (dense_indexer, dense_attention): + assert (result["out"].shape, result["out"].dtype) == ((1, 4, 8), jnp.float32) + assert (result["denom"].shape, result["denom"].dtype) == ((1, 4), jnp.float32) + + +@pytest.mark.L0 +def test_jax_dense_score_recompute_jit_and_numerics(): + jax, jnp = _jax_runtime() + device, target_compute_capability = _sm100_device(jax) + from cudnn.jax import dense_attn_score_recompute_wrapper, dense_indexer_score_recompute_wrapper + + batch, seqlen_q, seqlen_k = 1, 2, 128 + num_query_heads, head_dim = 32, 128 + ratio = 1 + softmax_scale = head_dim**-0.5 + q = jax.device_put(jax.random.normal(jax.random.key(40), (batch, seqlen_q, num_query_heads, head_dim), dtype=jnp.bfloat16), device) + k = jax.device_put(jax.random.normal(jax.random.key(41), (batch, seqlen_k, 1, head_dim), dtype=jnp.bfloat16), device) + weights = jax.device_put( + jax.random.normal(jax.random.key(42), (batch, seqlen_q, num_query_heads), dtype=jnp.bfloat16) * jnp.bfloat16(1.0 / num_query_heads), + device, + ) + q_causal_offsets = jax.device_put(jnp.array([63], dtype=jnp.int32), device) + + qk = jnp.einsum("bqhd,bkd->bqhk", q.astype(jnp.float32), k[:, :, 0, :].astype(jnp.float32)) + scaled_qk = qk * softmax_scale + row_max = jnp.max(scaled_qk, axis=-1) + lse = row_max + jnp.log(jnp.sum(jnp.exp(scaled_qk - row_max[..., None]), axis=-1)) + + @jax.jit + def run_indexer(q, k, weights, offsets): + return dense_indexer_score_recompute_wrapper( + q, + k, + weights, + ratio=ratio, + q_causal_offsets=offsets, + target_compute_capability=target_compute_capability, + ) + + @jax.jit + def run_attention(q, k, lse, offsets): + return dense_attn_score_recompute_wrapper( + q, + k, + lse, + softmax_scale, + ratio=ratio, + q_causal_offsets=offsets, + target_compute_capability=target_compute_capability, + ) + + indexer_lowered = run_indexer.lower(q, k, weights, q_causal_offsets) + attention_lowered = run_attention.lower(q, k, lse, q_causal_offsets) + for lowered in (indexer_lowered, attention_lowered): + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + indexer = indexer_lowered.compile()(q, k, weights, q_causal_offsets) + attention = attention_lowered.compile()(q, k, lse, q_causal_offsets) + attention["denom"].block_until_ready() + + q_positions = jnp.arange(seqlen_q, dtype=jnp.int32)[None, :, None] + col_limits = (q_causal_offsets[:, None, None] + q_positions + 1) // ratio + valid = jnp.arange(seqlen_k, dtype=jnp.int32)[None, None, :] < col_limits + + indexer_scores = jnp.sum(jnp.maximum(qk, 0.0) * weights.astype(jnp.float32)[..., None], axis=2) + expected_indexer_out = jnp.where(valid, indexer_scores, -jnp.inf) + indexer_max = jnp.max(expected_indexer_out, axis=-1) + expected_indexer_denom = indexer_max + jnp.log(jnp.sum(jnp.exp(expected_indexer_out - indexer_max[..., None]), axis=-1)) + + attention_scores = jnp.sum(jnp.exp(scaled_qk - lse[..., None]), axis=2) + expected_attention_out = jnp.where(valid, attention_scores, -jnp.inf) + expected_attention_denom = jnp.sum(jnp.where(valid, attention_scores, 0.0), axis=-1) + + assert jnp.allclose(indexer["out"], expected_indexer_out, atol=6e-3, rtol=2e-2) + assert jnp.allclose(indexer["denom"], expected_indexer_denom, atol=6e-3, rtol=2e-2) + assert jnp.allclose(attention["out"], expected_attention_out, atol=6e-3, rtol=2e-2) + assert jnp.allclose(attention["denom"], expected_attention_denom, atol=6e-3, rtol=2e-2) + assert jnp.all(jnp.isneginf(indexer["out"][~valid])) + assert jnp.all(jnp.isneginf(attention["out"][~valid])) + + +@pytest.mark.L0 +def test_jax_sparse_indexer_score_recompute_jit(): + jax, jnp = _jax_runtime() + device, target_compute_capability = _sm100_device(jax) + from cudnn.jax import disable_device_compatibility_checks, sparse_indexer_score_recompute_wrapper + + disable_device_compatibility_checks(False) + batch, seqlen_q, seqlen_k = 1, 2, 128 + num_query_heads, head_dim, topk = 32, 128, 128 + q = jax.device_put(jax.random.normal(jax.random.key(10), (batch, seqlen_q, num_query_heads, head_dim), dtype=jnp.bfloat16), device) + k = jax.device_put(jax.random.normal(jax.random.key(11), (batch, seqlen_k, head_dim), dtype=jnp.bfloat16), device) + weights = jax.device_put( + jax.random.normal(jax.random.key(12), (batch, seqlen_q, num_query_heads), dtype=jnp.bfloat16) * jnp.bfloat16(1.0 / num_query_heads), + device, + ) + indices = jax.device_put( + jnp.broadcast_to(jnp.arange(seqlen_k - 1, -1, -1, dtype=jnp.int32), (batch, seqlen_q, topk)), + device, + ) + + lowered = sparse_indexer_score_recompute_wrapper.lower( + q, + k, + weights, + indices, + target_compute_capability=target_compute_capability, + ) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + predict = lowered.compile()(q, k, weights, indices)["predict"] + predict.block_until_ready() + + k_gather = jax.vmap(lambda batch_k, batch_indices: batch_k[batch_indices])(k.astype(jnp.float32), indices) + scores_per_head = jnp.einsum("bqhd,bqtd->bqht", q.astype(jnp.float32), k_gather) + scores = jnp.sum(jnp.maximum(scores_per_head, 0.0) * weights.astype(jnp.float32)[..., None], axis=2) + expected = jax.nn.softmax(scores, axis=-1) + assert jnp.allclose(predict, expected, atol=3e-3, rtol=2e-2) + + +@pytest.mark.L0 +def test_jax_sparse_attn_score_recompute_jit(): + jax, jnp = _jax_runtime() + device, target_compute_capability = _sm100_device(jax) + from cudnn.jax import disable_device_compatibility_checks, sparse_attn_score_recompute_wrapper + + disable_device_compatibility_checks(False) + batch, seqlen_q, seqlen_k = 1, 2, 128 + num_query_heads, head_dim, topk = 32, 128, 128 + softmax_scale = head_dim**-0.5 + q = jax.device_put(jax.random.normal(jax.random.key(20), (batch, seqlen_q, num_query_heads, head_dim), dtype=jnp.bfloat16), device) + k = jax.device_put(jax.random.normal(jax.random.key(21), (batch, seqlen_k, head_dim), dtype=jnp.bfloat16), device) + indices = jax.device_put( + jnp.broadcast_to(jnp.arange(seqlen_k - 1, -1, -1, dtype=jnp.int32), (batch, seqlen_q, topk)), + device, + ) + full_scores = jnp.einsum("bqhd,bkd->bqhk", q.astype(jnp.float32), k.astype(jnp.float32)) * softmax_scale + row_max = jnp.max(full_scores, axis=-1) + lse = row_max + jnp.log(jnp.sum(jnp.exp(full_scores - row_max[..., None]), axis=-1)) + + lowered = sparse_attn_score_recompute_wrapper.lower( + q, + k, + lse, + indices, + softmax_scale=softmax_scale, + target_compute_capability=target_compute_capability, + ) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + target = lowered.compile()(q, k, lse, indices)["target"] + target.block_until_ready() + + k_gather = jax.vmap(lambda batch_k, batch_indices: batch_k[batch_indices])(k.astype(jnp.float32), indices) + sparse_scores = jnp.einsum("bqhd,bqtd->bqht", q.astype(jnp.float32), k_gather) * softmax_scale + expected = jnp.sum(jnp.exp(sparse_scores - lse[..., None]), axis=2) + expected /= jnp.sum(expected, axis=-1, keepdims=True) + assert jnp.allclose(target, expected, atol=3e-3, rtol=2e-2) + + +@pytest.mark.L0 +def test_jax_sparse_score_recompute_global_indices_and_partial_lengths(): + jax, jnp = _jax_runtime() + device, target_compute_capability = _sm100_device(jax) + from cudnn.jax import sparse_attn_score_recompute_wrapper, sparse_indexer_score_recompute_wrapper + + batch, seqlen_q, seqlen_k = 2, 1, 128 + num_query_heads, head_dim, topk = 32, 128, 128 + softmax_scale = head_dim**-0.5 + q = jax.device_put(jax.random.normal(jax.random.key(30), (batch, seqlen_q, num_query_heads, head_dim), dtype=jnp.bfloat16), device) + k = jax.device_put(jax.random.normal(jax.random.key(31), (batch, seqlen_k, head_dim), dtype=jnp.bfloat16), device) + weights = jax.device_put( + jax.random.normal(jax.random.key(32), (batch, seqlen_q, num_query_heads), dtype=jnp.bfloat16) * jnp.bfloat16(1.0 / num_query_heads), + device, + ) + local_indices = jnp.broadcast_to(jnp.arange(topk, dtype=jnp.int32), (batch, seqlen_q, topk)) + global_indices = jax.device_put(local_indices + jnp.arange(batch, dtype=jnp.int32)[:, None, None] * seqlen_k, device) + topk_length = jax.device_put(jnp.array([[0], [73]], dtype=jnp.int32), device) + + full_scores = jnp.einsum("bqhd,bkd->bqhk", q.astype(jnp.float32), k.astype(jnp.float32)) * softmax_scale + row_max = jnp.max(full_scores, axis=-1) + lse = row_max + jnp.log(jnp.sum(jnp.exp(full_scores - row_max[..., None]), axis=-1)) + + @jax.jit + def run(q, k, weights, lse, indices, lengths): + predict = sparse_indexer_score_recompute_wrapper( + q, + k, + weights, + indices, + topk_length=lengths, + topk_indices_global=True, + target_compute_capability=target_compute_capability, + )["predict"] + target = sparse_attn_score_recompute_wrapper( + q, + k, + lse, + indices, + softmax_scale, + topk_length=lengths, + topk_indices_global=True, + target_compute_capability=target_compute_capability, + )["target"] + return predict, target + + lowered = run.lower(q, k, weights, lse, global_indices, topk_length) + assert lowered.as_text("stablehlo").count("stablehlo.custom_call") == 2 + predict, target = lowered.compile()(q, k, weights, lse, global_indices, topk_length) + target.block_until_ready() + + k_gather = jax.vmap(lambda batch_k, batch_indices: batch_k[batch_indices])(k.astype(jnp.float32), local_indices) + qk = jnp.einsum("bqhd,bqtd->bqht", q.astype(jnp.float32), k_gather) + valid = jnp.arange(topk)[None, None, :] < topk_length[..., None] + + indexer_scores = jnp.sum(jnp.maximum(qk, 0.0) * weights.astype(jnp.float32)[..., None], axis=2) + expected_predict = jax.nn.softmax(jnp.where(valid, indexer_scores, -jnp.inf), axis=-1) + expected_predict = jnp.where(topk_length[..., None] > 0, expected_predict, 0.0) + + attention_scores = jnp.sum(jnp.exp(qk * softmax_scale - lse[..., None]), axis=2) + attention_scores = jnp.where(valid, attention_scores, 0.0) + attention_denom = jnp.sum(attention_scores, axis=-1, keepdims=True) + expected_target = jnp.where(attention_denom > 0, attention_scores / jnp.maximum(attention_denom, 1e-12), 0.0) + + assert jnp.allclose(predict, expected_predict, atol=3e-3, rtol=2e-2) + assert jnp.allclose(target, expected_target, atol=3e-3, rtol=2e-2) + assert jnp.array_equal(predict[0], jnp.zeros_like(predict[0])) + assert jnp.array_equal(target[0], jnp.zeros_like(target[0])) + assert jnp.array_equal(predict[1, :, 73:], jnp.zeros_like(predict[1, :, 73:])) + assert jnp.array_equal(target[1, :, 73:], jnp.zeros_like(target[1, :, 73:])) diff --git a/test/python/fe_api/test_jax_dsa_score_recompute_contract.py b/test/python/fe_api/test_jax_dsa_score_recompute_contract.py new file mode 100644 index 000000000..9b0c652f1 --- /dev/null +++ b/test/python/fe_api/test_jax_dsa_score_recompute_contract.py @@ -0,0 +1,475 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for JAX DSA score-recompute adapters.""" + +from enum import Enum, auto +import ast +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_DSA_ROOT = _CUDNN_ROOT / "deepseek_sparse_attention" +_SCORE_ROOT = _DSA_ROOT / "score_recompute" +_PACKAGE = "cudnn_jax_dsa_score_contract_test" + + +class _DataType(Enum): + NOT_SET = auto() + FLOAT = auto() + BFLOAT16 = auto() + INT32 = auto() + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + @property + def ndim(self): + return len(self.shape) + + +class _TupleDict(dict): + pass + + +def _identity_jit(fn=None, **_kwargs): + return (lambda decorated: decorated) if fn is None else fn + + +class JaxDsaScoreRecomputeContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + cls.tensor = importlib.import_module(f"{_PACKAGE}._tensor_desc") + + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.bfloat16 = _DataType.BFLOAT16 + fake_jnp.float32 = _DataType.FLOAT + fake_jnp.int32 = _DataType.INT32 + fake_jnp.expand_dims = lambda value, axis: _Array(value.shape[:axis] + (1,) + value.shape[axis:], value.dtype) + fake_jnp.transpose = lambda value, axes: _Array(tuple(value.shape[axis] for axis in axes), value.dtype) + + fake_jax = types.ModuleType("jax") + fake_jax.__path__ = [] + fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + fake_jax.numpy = fake_jnp + fake_jax.jit = _identity_jit + fake_jax.ShapeDtypeStruct = _Array + + tensor_module = cls.tensor + + class FakeJaxApiBase: + @staticmethod + def _to_tensor_desc(value, name, *, mode=None, init_value=None): + del mode + return tensor_module.make_compact_tensor_desc( + dtype=value.dtype, + shape=value.shape, + name=name, + init_value=init_value, + ) + + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del operation_name + resolved = 100 if target is None else target + if resolved not in supported: + raise ValueError(f"unsupported target {resolved}") + return resolved + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + del mode + if tuple(value.shape) != expected.shape: + raise ValueError(f"{expected.name} tensor shape mismatch") + if value.dtype != expected.cudnn_dtype: + raise ValueError(f"{expected.name} tensor dtype mismatch") + + def _call_kernel(self, inputs, *, output_descs, workspace_descs=(), **options): + self.captured_call = { + "inputs": tuple(inputs), + "outputs": tuple(output_descs), + "workspaces": tuple(workspace_descs), + **options, + } + return tuple(_Array(desc.shape, desc.cudnn_dtype) for desc in output_descs) + + fake_internal_jax = types.ModuleType(f"{_PACKAGE}._jax") + fake_internal_jax.JaxApiBase = FakeJaxApiBase + fake_internal_jax.JaxTensorDesc = tensor_module.TensorDesc + fake_internal_jax.TupleDict = _TupleDict + + dsa_name = f"{_PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_DSA_ROOT)] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + + score_name = f"{dsa_name}.score_recompute" + score = types.ModuleType(score_name) + score.__path__ = [str(_SCORE_ROOT)] + score.__package__ = score_name + score.__spec__ = ModuleSpec(score_name, loader=None, is_package=True) + + utils_name = f"{dsa_name}.utils" + utils = types.ModuleType(utils_name) + utils.__path__ = [] + compiler = types.ModuleType(f"{utils_name}.compiler") + compiler.compile_options_for_target = lambda target: f"--gpu-arch sm_{target}" + + cls.optional_modules = { + "jax": fake_jax, + "jax.numpy": fake_jnp, + f"{_PACKAGE}._jax": fake_internal_jax, + dsa_name: dsa, + score_name: score, + utils_name: utils, + compiler.__name__: compiler, + } + sys.modules.update(cls.optional_modules) + try: + cls.adapter = importlib.import_module(f"{score_name}.jax") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + for name in ("jax.numpy", "jax"): + module = sys.modules.get(name) + if module is cls.optional_modules.get(name): + sys.modules.pop(name, None) + + def array(self, shape, dtype=_DataType.BFLOAT16): + return _Array(shape, dtype) + + def sparse_samples(self): + return ( + self.array((2, 4, 32, 128)), + self.array((2, 256, 128)), + self.array((2, 4, 32)), + self.array((2, 4, 128), _DataType.INT32), + ) + + def dense_samples(self, *, attention=False): + return ( + self.array((2, 4, 32, 128)), + self.array((2, 8, 1, 128)), + self.array((2, 4, 32), _DataType.FLOAT if attention else _DataType.BFLOAT16), + ) + + def test_sm100_sparse_infers_output_and_hidden_length_workspace(self): + q, k, weights, indices = self.sparse_samples() + api = self.adapter.SparseIndexerScoreRecompute(q, k, weights, indices, target_compute_capability=100) + result = api(q, k, weights, indices) + + self.assertIsInstance(result, _TupleDict) + self.assertEqual((result["predict"].shape, result["predict"].dtype), ((2, 4, 128), _DataType.FLOAT)) + self.assertEqual(len(api.captured_call["workspaces"]), 1) + workspace = api.captured_call["workspaces"][0] + self.assertEqual((workspace.shape, workspace.cudnn_dtype), ((1, 1), _DataType.INT32)) + self.assertEqual(api.captured_call["compile_options"], "--gpu-arch sm_100") + + def test_explicit_topk_length_remains_a_real_input(self): + q, k, _, indices = self.sparse_samples() + lse = self.array((2, 4, 32), _DataType.FLOAT) + lengths = self.array((2, 4), _DataType.INT32) + api = self.adapter.SparseAttnScoreRecompute( + q, + k, + lse, + indices, + 0.125, + sample_topk_length=lengths, + target_compute_capability=100, + ) + result = api(q, k, lse, indices, lengths) + + self.assertEqual(result["target"].shape, (2, 4, 128)) + self.assertEqual(api.captured_call["workspaces"], ()) + self.assertIs(api.captured_call["inputs"][-1], lengths) + self.assertTrue(api._op.config.have_topk_length) + + def test_sm90_sparse_adapts_k_and_per_head_layouts(self): + q, k, weights, indices = self.sparse_samples() + api = self.adapter.SparseIndexerScoreRecompute(q, k, weights, indices, target_compute_capability=90) + api(q, k, weights, indices) + + _, kernel_k, kernel_weights, _ = api.captured_call["inputs"] + self.assertEqual(kernel_k.shape, (2, 256, 1, 128)) + self.assertEqual(kernel_weights.shape, (2, 32, 4)) + self.assertEqual(api.captured_call["workspaces"], ()) + + def test_dense_bshd_outputs_are_functional_and_initialized(self): + q, k, weights = self.dense_samples() + offsets = self.array((2,), _DataType.INT32) + api = self.adapter.DenseIndexerScoreRecompute( + q, + k, + weights, + ratio=2, + sample_q_causal_offsets=offsets, + target_compute_capability=100, + ) + result = api(q, k, weights, q_causal_offsets=offsets) + + self.assertEqual(result["out"].shape, (2, 4, 8)) + self.assertEqual(result["denom"].shape, (2, 4)) + self.assertEqual(api.out_desc.init_value, float("-inf")) + self.assertIsNone(api.denom_desc.init_value) + self.assertIs(api.captured_call["inputs"][-1], offsets) + + explicit = self.adapter.DenseIndexerScoreRecompute( + q, + k, + weights, + sample_out=self.array((2, 4, 8), _DataType.FLOAT), + sample_denom_out=self.array((2, 4), _DataType.FLOAT), + target_compute_capability=100, + ) + self.assertEqual(explicit.out_desc.init_value, float("-inf")) + self.assertIsNone(explicit.denom_desc.init_value) + + def test_dense_thd_sm100_and_sm90_rejection(self): + q = self.array((7, 32, 128)) + k = self.array((11, 1, 128)) + lse = self.array((7, 32), _DataType.FLOAT) + cu_q = self.array((3,), _DataType.INT32) + cu_k = self.array((3,), _DataType.INT32) + api = self.adapter.DenseAttnScoreRecompute( + q, + k, + lse, + 0.125, + sample_cu_seqlens_q=cu_q, + sample_cu_seqlens_k=cu_k, + max_seqlen_q=4, + max_seqlen_k=8, + target_compute_capability=100, + ) + result = api(q, k, lse, cu_q, cu_k) + self.assertEqual((result["out"].shape, result["denom"].shape), ((7, 8), (7,))) + self.assertEqual(api.captured_call["inputs"][-2:], (cu_q, cu_k)) + + sm90 = self.adapter.DenseAttnScoreRecompute( + q, + k, + lse, + 0.125, + sample_cu_seqlens_q=cu_q, + sample_cu_seqlens_k=cu_k, + max_seqlen_q=4, + max_seqlen_k=8, + target_compute_capability=90, + ) + with self.assertRaisesRegex(NotImplementedError, "cannot be traced by JAX"): + sm90(q, k, lse, cu_q, cu_k) + + def test_explicit_output_samples_are_cross_checked(self): + q, k, weights = self.dense_samples() + bad_output = self.array((2, 4, 4), _DataType.FLOAT) + denominator = self.array((2, 4), _DataType.FLOAT) + api = self.adapter.DenseIndexerScoreRecompute( + q, + k, + weights, + sample_out=bad_output, + sample_denom_out=denominator, + target_compute_capability=100, + ) + with self.assertRaisesRegex(ValueError, "output must have shape"): + api(q, k, weights) + + def test_high_level_wrapper_infers_shapes(self): + q, k, weights, indices = self.sparse_samples() + with mock.patch.dict(sys.modules, self.optional_modules): + result = self.adapter.sparse_indexer_score_recompute_wrapper( + q, + k, + weights, + indices, + target_compute_capability=100, + ) + self.assertEqual(result["predict"].shape, indices.shape) + + def test_sparse_launchers_preserve_native_argument_order(self): + calls = [] + + class FakeSm90Kernel: + def __init__(self, dtype, **options): + calls.append(("sm90_init", dtype, options)) + + def __call__(self, *args): + calls.append(("sm90_call", args)) + + class FakeSm100Kernel: + def __init__(self, **options): + calls.append(("sm100_init", options)) + + def __call__(self, *args): + calls.append(("sm100_call", args)) + + cutlass = types.ModuleType("cutlass") + cutlass.BFloat16 = "BFloat16" + cutlass.Float32 = lambda value: ("Float32", value) + package = self.adapter.__package__ + sm90_module = types.ModuleType(f"{package}.sparse_score_recompute_sm90") + sm90_module.SparseScoreRecomputeSm90 = FakeSm90Kernel + sm100_module = types.ModuleType(f"{package}.sparse_score_recompute_sm100") + sm100_module.SparseScoreRecomputeSm100 = FakeSm100Kernel + + q, k, weights, indices = self.sparse_samples() + sm90 = self.adapter.SparseIndexerScoreRecompute(q, k, weights, indices, target_compute_capability=90) + sm90.check_support() + sm100 = self.adapter.SparseIndexerScoreRecompute(q, k, weights, indices, target_compute_capability=100) + sm100.check_support() + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + sm90_module.__name__: sm90_module, + sm100_module.__name__: sm100_module, + }, + ): + sm90._launch(("q90", "k90", "aux90", "indices90"), ("out90",), (), "stream90") + sm100._launch(("q100", "k100", "aux100", "indices100"), ("out100",), ("dummy",), "stream100") + + self.assertEqual( + [entry for entry in calls if entry[0].endswith("call")], + [ + ("sm90_call", ("q90", "k90", "indices90", "stream90", "out90", "aux90", None, None)), + ( + "sm100_call", + ("q100", "k100", "aux100", "indices100", "out100", "dummy", ("Float32", 1.0), "stream100"), + ), + ], + ) + + def test_dense_launchers_preserve_native_argument_order(self): + calls = [] + + class FakeSm90Kernel: + def __init__(self, dtype, **options): + calls.append(("sm90_init", dtype, options)) + + def __call__(self, *args): + calls.append(("sm90_call", args)) + + class FakeSm100Kernel: + def __init__(self, **options): + calls.append(("sm100_init", options)) + + def __call__(self, *args): + calls.append(("sm100_call", args)) + + cutlass = types.ModuleType("cutlass") + cutlass.BFloat16 = "BFloat16" + cutlass.Float32 = lambda value: ("Float32", value) + cutlass.Int32 = lambda value: ("Int32", value) + package = self.adapter.__package__ + sm90_module = types.ModuleType(f"{package}.dense_score_recompute_sm90") + sm90_module.DenseScoreRecomputeSm90 = FakeSm90Kernel + sm100_module = types.ModuleType(f"{package}.dense_score_recompute_sm100") + sm100_module.DenseScoreRecomputeSm100 = FakeSm100Kernel + + q, k, weights = self.dense_samples() + offsets = self.array((2,), _DataType.INT32) + sm90 = self.adapter.DenseIndexerScoreRecompute( + q, + k, + weights, + sample_q_causal_offsets=offsets, + target_compute_capability=90, + ) + sm90.check_support() + sm100 = self.adapter.DenseIndexerScoreRecompute( + q, + k, + weights, + sample_q_causal_offsets=offsets, + target_compute_capability=100, + ) + sm100.check_support() + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + sm90_module.__name__: sm90_module, + sm100_module.__name__: sm100_module, + }, + ): + sm90._launch(("q90", "k90", "aux90", "offset90"), ("out90", "denom90"), (), "stream90") + sm100._launch(("q100", "k100", "aux100", "offset100"), ("out100", "denom100"), (), "stream100") + + self.assertEqual( + [entry for entry in calls if entry[0].endswith("call")], + [ + ("sm90_call", ("q90", "k90", None, "stream90", "out90", "aux90", None, "denom90", "offset90")), + ( + "sm100_call", + ( + "q100", + "k100", + "aux100", + "out100", + "denom100", + ("Float32", 1.0), + ("Int32", 4), + ("Int32", 8), + None, + None, + "offset100", + "stream100", + ), + ), + ], + ) + + def test_common_and_jax_modules_keep_framework_import_boundaries(self): + def imports_framework(path, framework): + tree = ast.parse(path.read_text(), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import) and any(alias.name == framework or alias.name.startswith(f"{framework}.") for alias in node.names): + return True + if isinstance(node, ast.ImportFrom) and node.level == 0 and (node.module == framework or (node.module or "").startswith(f"{framework}.")): + return True + return False + + for filename in ("config.py", "op.py"): + path = _SCORE_ROOT / filename + for framework in ("torch", "jax", "cutlass", "cuda"): + self.assertFalse(imports_framework(path, framework), f"{filename} imports {framework}") + self.assertFalse(imports_framework(_SCORE_ROOT / "jax.py", "torch")) + + init_tree = ast.parse((_SCORE_ROOT / "__init__.py").read_text()) + imported_modules = {node.module for node in ast.walk(init_tree) if isinstance(node, ast.ImportFrom)} + self.assertNotIn("api", imported_modules) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py b/test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py new file mode 100644 index 000000000..8a7634616 --- /dev/null +++ b/test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py @@ -0,0 +1,385 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free JAX contracts for DSA sparse-attention backward.""" + +from __future__ import annotations + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_DSA_ROOT = _CUDNN_ROOT / "deepseek_sparse_attention" +_OPERATION_ROOT = _DSA_ROOT / "sparse_attention_backward" +_PACKAGE = "cudnn_jax_dsa_sparse_attention_backward_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + UINT8 = auto() + + +_DTYPE_TO_CUDNN = { + "float16": _DataType.HALF, + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, + "int32": _DataType.INT32, + "uint8": _DataType.UINT8, +} +_CUDNN_TO_DTYPE = {value: key for key, value in _DTYPE_TO_CUDNN.items()} + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + +class _TupleDict(dict): + def __iter__(self): + return iter(self.values()) + + +class _TensorSpec: + def __init__(self, *, mode=None, divisibility=None): + self.mode = mode + self.divisibility = divisibility + + +class JaxSparseAttentionBackwardContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + dsa_name = f"{_PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_DSA_ROOT)] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + sys.modules[dsa_name] = dsa + + operation_name = f"{dsa_name}.sparse_attention_backward" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + + class JaxTensorDesc(tensor_module.TensorDesc): + @property + def cudnn_dtype(self): + return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) + + def compact_like(self, *, cudnn_dtype, shape, stride_order=None, name="", init_value=None): + canonical = tensor_module.make_compact_tensor_desc( + dtype=cudnn_dtype, + shape=shape, + stride_order=stride_order, + name=name, + init_value=init_value, + ) + return JaxTensorDesc( + dtype=_CUDNN_TO_DTYPE[cudnn_dtype], + shape=canonical.shape, + stride=canonical.stride, + stride_order=canonical.stride_order, + name=name, + init_value=init_value, + ) + + class JaxApiBase: + captured_call = None + + @staticmethod + def _to_tensor_desc(value, name, *, mode=None, init_value=None): + del mode + shape = tuple(value.shape) + stride = [] + running = 1 + for size in reversed(shape): + stride.append(running) + running *= max(size, 1) + stride = tuple(reversed(stride)) + return JaxTensorDesc( + dtype=value.dtype, + shape=shape, + stride=stride, + stride_order=tuple(reversed(range(len(shape)))), + name=name, + init_value=init_value, + ) + + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del supported, operation_name + if target is None: + raise RuntimeError("target required by dependency-free test") + return target + + @staticmethod + def _compute_capability_family(target, supported): + return max((value for value in supported if value <= target), default=None) + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + del mode + if tuple(value.shape) != expected.shape or _DTYPE_TO_CUDNN[value.dtype] != expected.cudnn_dtype: + raise ValueError(f"{expected.name} tensor signature mismatch") + + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + del desc + return _TensorSpec(mode=mode, divisibility=divisibility) + + def _call_kernel(self, inputs, **options): + JaxApiBase.captured_call = (inputs, options) + return tuple(_Array(desc.shape, _CUDNN_TO_DTYPE[desc.cudnn_dtype]) for desc in options["output_descs"]) + + internal = types.ModuleType(f"{_PACKAGE}._jax") + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = _TupleDict + sys.modules[internal.__name__] = internal + + fake_jax = types.ModuleType("jax") + fake_jax.jit = lambda function=None, **_kwargs: (lambda fn: fn) if function is None else function + fake_jax.ShapeDtypeStruct = _Array + + try: + with mock.patch.dict(sys.modules, {"jax": fake_jax}): + cls.module = importlib.import_module(f"{operation_name}.jax") + except Exception: + cls.tearDownClass() + raise + cls.JaxApiBase = JaxApiBase + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + @staticmethod + def _inputs(*, head_dim=512, dtype="bfloat16", with_length=False): + m, n, h, topk = 65, 130, 64, 32 + head_dim_v = 512 if head_dim == 576 else head_dim + values = ( + _Array((m, h, head_dim), dtype), + _Array((n, head_dim), dtype), + _Array((m, h, head_dim_v), dtype), + _Array((m, h, head_dim_v), dtype), + _Array((m, h), "float32"), + _Array((h,), "float32"), + _Array((m, topk), "int32"), + ) + return (*values, _Array((m,), "int32") if with_length else None) + + def test_sm100_declares_functional_outputs_and_zeroed_workspaces(self): + inputs = self._inputs(with_length=True) + api = self.module.SparseAttentionBackward( + *inputs[:-1], + sample_topk_length=inputs[-1], + target_compute_capability=103, + ) + result = api(*inputs[:-1], inputs[-1]) + self.assertEqual(tuple(result), (result["dq"], result["dkv"], result["d_sink"])) + self.assertEqual(result["dq"].shape, inputs[0].shape) + self.assertEqual(result["dkv"].shape, inputs[1].shape) + self.assertEqual(result["d_sink"].shape, inputs[5].shape) + + _, options = self.JaxApiBase.captured_call + outputs = options["output_descs"] + self.assertIsNone(outputs[0].init_value) + self.assertIsNone(outputs[1].init_value) + self.assertEqual(outputs[2].init_value, 0.0) + workspaces = options["workspace_descs"] + self.assertEqual( + [(desc.shape, desc.cudnn_dtype, desc.init_value) for desc in workspaces], + [ + ((1, 64, 72, 8), _DataType.UINT8, 0), + ((1, 1, 136, 2048), _DataType.UINT8, 0), + ], + ) + self.assertIn("--gpu-arch sm_103a", options["compile_options"]) + + def test_sm90_declares_pipeline_workspaces_and_optional_length_dummy(self): + inputs = self._inputs(dtype="float16") + api = self.module.SparseAttentionBackward(*inputs[:-1], target_compute_capability=90) + api(*inputs[:-1]) + _, options = self.JaxApiBase.captured_call + self.assertEqual( + [(desc.shape, desc.cudnn_dtype, desc.init_value) for desc in options["workspace_descs"]], + [ + ((1, 128, 64), _DataType.FLOAT, None), + ((1, 128, 64), _DataType.FLOAT, None), + ((1, 1, 192 * 512), _DataType.FLOAT, 0.0), + ((1,), _DataType.INT32, None), + ], + ) + self.assertIn("--gpu-arch sm_90a", options["compile_options"]) + + def test_target_specific_dtype_and_576_shape_validation(self): + fp16 = self._inputs(dtype="float16") + with self.assertRaisesRegex(ValueError, "SM100.*bfloat16"): + self.module.SparseAttentionBackward(*fp16[:-1], target_compute_capability=100)(*fp16[:-1]) + + inputs_576 = self._inputs(head_dim=576) + result = self.module.SparseAttentionBackward(*inputs_576[:-1], target_compute_capability=100)(*inputs_576[:-1]) + self.assertEqual(result["dq"].shape, (65, 64, 576)) + self.assertEqual(result["dkv"].shape, (130, 576)) + + def test_sm100_launcher_passes_max_topk_to_kernel(self): + inputs = self._inputs() + api = self.module.SparseAttentionBackward(*inputs[:-1], target_compute_capability=100) + api.check_support() + seen = {} + + class Kernel: + def __init__(self, **configuration): + seen["configuration"] = configuration + + def __call__(self, *arguments): + seen["arguments"] = arguments + + kernel_module = types.ModuleType(f"{self.module.__package__}.dsa_bwd_sm100") + kernel_module.FlashAttentionDSABackwardSm100 = Kernel + cutlass = types.ModuleType("cutlass") + cutlass.Int32 = int + cutlass.Float32 = float + placeholders = tuple(object() for _ in range(12)) + with mock.patch.dict(sys.modules, {"cutlass": cutlass, kernel_module.__name__: kernel_module}): + api._launch_sm100(placeholders[:7], placeholders[7:10], placeholders[10:12], object()) + + self.assertEqual(seen["configuration"]["max_topk"], 32) + self.assertEqual(seen["configuration"]["head_dim"], 512) + self.assertIsNone(seen["arguments"][8]) + + def test_sm90_launcher_builds_flat_mqa_views_and_orders_pipeline(self): + inputs = self._inputs() + api = self.module.SparseAttentionBackward(*inputs[:-1], target_compute_capability=90) + api.check_support() + seen = {} + + class CuteTensor: + def __init__(self, shape, stride, *, element_type="BFloat16", label=""): + self.shape = tuple(shape) + self.stride = tuple(stride) + self.element_type = element_type + self.iterator = self + self.label = label + + def compact(shape, *, element_type="BFloat16", label=""): + stride = [] + running = 1 + for size in reversed(shape): + stride.append(running) + running *= max(size, 1) + return CuteTensor(shape, tuple(reversed(stride)), element_type=element_type, label=label) + + class Stage: + def __init__(self, label, *configuration, **keywords): + seen[f"{label}_configuration"] = (configuration, keywords) + self.label = label + + def __call__(self, *arguments): + seen[f"{self.label}_arguments"] = arguments + + kernel_module = types.ModuleType(f"{self.module.__package__}.dsa_bwd_sm90") + kernel_module._FlashAttentionDSABackwardPreprocessSm90 = lambda *args, **kwargs: Stage("pre", *args, **kwargs) + kernel_module.FlashAttentionDSABackwardSm90 = lambda *args, **kwargs: Stage("main", *args, **kwargs) + kernel_module._FlashAttentionDSABackwardPostprocessSm90 = lambda *args, **kwargs: Stage("post", *args, **kwargs) + + cute = types.ModuleType("cutlass.cute") + cute.make_layout = lambda shape, stride: types.SimpleNamespace(shape=tuple(shape), stride=tuple(stride)) + cute.make_tensor = lambda iterator, layout: CuteTensor( + layout.shape, + layout.stride, + element_type=iterator.element_type, + label=iterator.label, + ) + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass.Int32 = int + cutlass.Float32 = float + cutlass.cute = cute + + kernel_inputs = ( + compact((65, 64, 512), label="q"), + compact((130, 512), label="kv"), + compact((65, 64, 512), label="out"), + compact((65, 64, 512), label="dout"), + compact((65, 64), element_type="Float32", label="lse"), + compact((64,), element_type="Float32", label="sink"), + compact((65, 32), element_type="Int32", label="topk"), + ) + kernel_outputs = ( + compact((65, 64, 512), label="dq"), + compact((130, 512), label="dkv"), + compact((64,), element_type="Float32", label="dsink"), + ) + kernel_workspaces = ( + compact((1, 128, 64), element_type="Float32", label="dpsum"), + compact((1, 128, 64), element_type="Float32", label="lse_log2"), + compact((1, 1, 192 * 512), element_type="Float32", label="dkv_accum"), + compact((1,), element_type="Int32", label="dummy_length"), + ) + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + "cutlass.cute": cute, + kernel_module.__name__: kernel_module, + }, + ): + api._launch_sm90(kernel_inputs, kernel_outputs, kernel_workspaces, object()) + + pre_args = seen["pre_arguments"] + self.assertEqual(pre_args[0].shape, (1, 65, 64, 512)) + self.assertEqual(pre_args[1].shape, (1, 65, 64, 512)) + main_args = seen["main_arguments"] + self.assertEqual(main_args[0].shape, (1, 65, 64, 512)) + self.assertEqual(main_args[1].shape, (1, 130, 1, 512)) + self.assertEqual(main_args[7].shape, (1, 65, 32)) + self.assertEqual(main_args[8].shape, (1, 1)) + main_config, main_keywords = seen["main_configuration"] + self.assertEqual(main_config[:4], ("BFloat16", 512, 512, 64)) + self.assertEqual(main_keywords["max_topk"], 32) + self.assertFalse(main_keywords["have_topk_length"]) + post_args = seen["post_arguments"] + self.assertEqual(post_args[1].shape, (1, 130, 1, 512)) + self.assertEqual(post_args[2], 130) + + def test_jax_module_has_no_torch_import(self): + source = (_OPERATION_ROOT / "jax.py").read_text() + self.assertNotIn("import torch", source) + for kernel_name in ("dsa_bwd_sm90.py", "dsa_bwd_sm100.py"): + self.assertNotIn("import torch", (_OPERATION_ROOT / kernel_name).read_text()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_import_contract.py b/test/python/fe_api/test_jax_import_contract.py new file mode 100644 index 000000000..bbfdacd41 --- /dev/null +++ b/test/python/fe_api/test_jax_import_contract.py @@ -0,0 +1,225 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Automatically discovered import-boundary contracts for JAX APIs.""" + +from __future__ import annotations + +import ast +from collections import deque +import importlib.util +from pathlib import Path +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CUDNN_ROOT = _REPO_ROOT / "python" / "cudnn" + + +def _module_name(path): + relative = path.relative_to(_CUDNN_ROOT) + parts = relative.parts[:-1] if path.name == "__init__.py" else relative.with_suffix("").parts + return ".".join(("cudnn", *parts)) + + +def _static_condition(node): + if isinstance(node, ast.Constant) and isinstance(node.value, bool): + return node.value + if isinstance(node, ast.Name) and node.id == "TYPE_CHECKING": + return False + if isinstance(node, ast.Attribute) and node.attr == "TYPE_CHECKING": + return False + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not): + value = _static_condition(node.operand) + return None if value is None else not value + if ( + isinstance(node, ast.Compare) + and isinstance(node.left, ast.Name) + and node.left.id == "__name__" + and len(node.ops) == len(node.comparators) == 1 + and isinstance(node.comparators[0], ast.Constant) + and node.comparators[0].value == "__main__" + ): + if isinstance(node.ops[0], ast.Eq): + return False + if isinstance(node.ops[0], ast.NotEq): + return True + return None + + +def _catches_import_error(node): + if isinstance(node, ast.Name): + return node.id in {"ImportError", "ModuleNotFoundError"} + if isinstance(node, ast.Attribute): + return node.attr in {"ImportError", "ModuleNotFoundError"} + if isinstance(node, ast.Tuple): + return any(_catches_import_error(item) for item in node.elts) + return False + + +def _imports(path, *, all_scopes): + """Collect imports and whether a direct Torch import is optional.""" + + imports = [] + + def visit(node, torch_optional=False): + if isinstance(node, (ast.Import, ast.ImportFrom)): + imports.append((node, torch_optional)) + return + if isinstance(node, ast.If): + condition = _static_condition(node.test) + if condition is True: + branches = node.body + elif condition is False: + branches = node.orelse + else: + branches = (*node.body, *node.orelse) + for child in branches: + visit(child, torch_optional) + return + if isinstance(node, ast.Try): + catches_import = any(handler.type is not None and _catches_import_error(handler.type) for handler in node.handlers) + for child in node.body: + visit(child, torch_optional or catches_import) + for handler in node.handlers: + for child in handler.body: + visit(child, torch_optional) + for child in (*node.orelse, *node.finalbody): + visit(child, torch_optional) + return + if not all_scopes and isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)): + return + for child in ast.iter_child_nodes(node): + visit(child, torch_optional) + + tree = ast.parse(path.read_text(), filename=str(path)) + for node in tree.body: + visit(node) + return imports + + +def _import_targets(path, node, modules): + if isinstance(node, ast.Import): + return [alias.name for alias in node.names] + + if node.level: + current_module = _module_name(path) + package = current_module if path.name == "__init__.py" else current_module.rpartition(".")[0] + base = importlib.util.resolve_name("." * node.level + (node.module or ""), package) + else: + base = node.module or "" + + targets = [base] if base else [] + for alias in node.names: + candidate = f"{base}.{alias.name}" if base else alias.name + if alias.name != "*" and candidate in modules: + targets.append(candidate) + return targets + + +def _package_initializers(path): + initializers = [_CUDNN_ROOT / "__init__.py"] + current = _CUDNN_ROOT + for part in path.parent.relative_to(_CUDNN_ROOT).parts: + current /= part + initializer = current / "__init__.py" + if initializer != path and initializer.is_file(): + initializers.append(initializer) + return initializers + + +def _jax_import_violations(root, modules): + # Scan every scope in JAX roots because kernel imports are deferred until + # tracing. In reached modules, scan only import-time code so shared kernel + # files may retain function-local Torch wrappers. + queue = deque([(root, True, (root,))]) + queue.extend((path, False, (root, path)) for path in _package_initializers(root)) + scanned_all_scopes = set() + scanned_module_scope = set() + visited = set() + violations = [] + + while queue: + path, all_scopes, chain = queue.popleft() + if path in scanned_all_scopes or (not all_scopes and path in scanned_module_scope): + continue + (scanned_all_scopes if all_scopes else scanned_module_scope).add(path) + visited.add(path) + + for node, torch_optional in _imports(path, all_scopes=all_scopes): + for target in _import_targets(path, node, modules): + chain_text = " -> ".join(str(item.relative_to(_REPO_ROOT)) for item in chain) + if target == "torch" or target.startswith("torch."): + if not torch_optional: + violations.append(f"{chain_text}:{node.lineno} imports {target}") + continue + + dependency = modules.get(target) + if dependency is None: + continue + if dependency.name == "api.py": + violations.append(f"{chain_text}:{node.lineno} imports Torch API module {dependency.relative_to(_REPO_ROOT)}") + continue + queue.extend((initializer, False, (*chain, initializer)) for initializer in _package_initializers(dependency)) + queue.append((dependency, False, (*chain, dependency))) + + return visited, violations + + +class JaxImportContractTest(unittest.TestCase): + def test_jax_dsa_facade_matches_the_torch_operation_surface(self): + def literal_assignment(path, name): + tree = ast.parse(path.read_text(), filename=str(path)) + assignment = next( + node + for node in tree.body + if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == name for target in node.targets) + ) + return ast.literal_eval(assignment.value) + + torch_exports = literal_assignment( + _CUDNN_ROOT / "deepseek_sparse_attention" / "__init__.py", + "_SYMBOLS", + ) + jax_exports = literal_assignment(_CUDNN_ROOT / "jax" / "__init__.py", "_OPERATION_EXPORTS") + jax_dsa_exports = {name for name in jax_exports if name not in {"RmsNormRhtAmaxSm100", "rmsnorm_rht_amax_sm100"}} + + self.assertEqual(jax_dsa_exports, set(torch_exports)) + + def test_jax_import_graph_has_no_required_torch_dependency(self): + modules = {_module_name(path): path for path in _CUDNN_ROOT.rglob("*.py")} + roots = set(_CUDNN_ROOT.rglob("jax.py")) + roots.update((_CUDNN_ROOT / "jax").rglob("*.py")) + self.assertTrue(roots, "No JAX API modules were discovered") + + visited = set() + violations = [] + for root in sorted(roots): + root_visited, root_violations = _jax_import_violations(root, modules) + visited.update(root_visited) + violations.extend(root_violations) + + self.assertTrue( + any(path.name == "__init__.py" and path not in roots for path in visited), + "JAX import graph did not discover package initializers", + ) + self.assertTrue( + any(path.name != "__init__.py" and path not in roots for path in visited), + "JAX import graph did not discover deferred kernel dependencies", + ) + self.assertEqual( + violations, + [], + "JAX import paths must not require Torch:\n" + "\n".join(violations), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_layout.py b/test/python/fe_api/test_jax_layout.py index 57ad2429b..9873ea034 100644 --- a/test/python/fe_api/test_jax_layout.py +++ b/test/python/fe_api/test_jax_layout.py @@ -43,6 +43,45 @@ def test_non_self_inverse_stride_order_maps_to_cutlass_axis_ranks(self): (2, 0, 1), ) + def test_named_layout_maps_kernel_axes_to_public_axes(self): + cases = ( + ("LMK", "MKL", (1, 2, 0)), + ("LKM", "MKL", (2, 1, 0)), + ("LNK", "NKL", (1, 2, 0)), + ("LKN", "NKL", (2, 1, 0)), + ("LMN", "MNL", (1, 2, 0)), + ("LNM", "MNL", (2, 1, 0)), + ) + for public_layout, kernel_axes, expected in cases: + with self.subTest(public_layout=public_layout, kernel_axes=kernel_axes): + self.assertEqual( + self.layout.mode_from_layout( + public_layout, + kernel_axes=kernel_axes, + ), + expected, + ) + + def test_named_layout_requires_strings(self): + with self.assertRaisesRegex(TypeError, "layout must be a string"): + self.layout.mode_from_layout(("L", "M", "K"), kernel_axes="MKL") + with self.assertRaisesRegex(TypeError, "kernel_axes must be a string"): + self.layout.mode_from_layout("LMK", kernel_axes=("M", "K", "L")) + + def test_named_layout_rejects_rank_mismatch(self): + with self.assertRaisesRegex(ValueError, "layout rank must match"): + self.layout.mode_from_layout("MK", kernel_axes="MKL") + + def test_named_layout_rejects_duplicate_axes(self): + with self.assertRaisesRegex(ValueError, "layout axes must be unique"): + self.layout.mode_from_layout("MML", kernel_axes="MKL") + with self.assertRaisesRegex(ValueError, "kernel_axes must be unique"): + self.layout.mode_from_layout("MKL", kernel_axes="MML") + + def test_named_layout_requires_the_exact_axis_set(self): + with self.assertRaisesRegex(ValueError, "layout must contain exactly"): + self.layout.mode_from_layout("MKN", kernel_axes="MKL") + def test_noncompact_stride_is_rejected(self): with self.assertRaisesRegex(ValueError, "cannot represent non-compact stride"): self.layout.to_cutlass_layout( diff --git a/test/python/fe_api/test_result.py b/test/python/fe_api/test_result.py new file mode 100644 index 000000000..5a6918037 --- /dev/null +++ b/test/python/fe_api/test_result.py @@ -0,0 +1,60 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for framework-neutral frontend results.""" + +import importlib.util +from pathlib import Path +import sys +import types + +import pytest + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_COMMON_SPEC = importlib.util.spec_from_file_location("_cudnn_result_contract", _CUDNN_ROOT / "_result.py") +assert _COMMON_SPEC is not None and _COMMON_SPEC.loader is not None +_COMMON_MODULE = importlib.util.module_from_spec(_COMMON_SPEC) +_COMMON_SPEC.loader.exec_module(_COMMON_MODULE) +TupleDict = _COMMON_MODULE.TupleDict + + +pytestmark = pytest.mark.L0 + + +def test_tuple_dict_preserves_named_and_positional_access(): + result = TupleDict(output=1, amax=2) + + assert tuple(result) == (1, 2) + assert result[0] == result["output"] == 1 + assert result[1] == result["amax"] == 2 + with pytest.raises(IndexError, match="index -1 out of range"): + _ = result[-1] + + +def test_tuple_dict_is_a_jax_pytree(monkeypatch): + jax = pytest.importorskip("jax") + + package_name = "_cudnn_result_jax_contract" + package = types.ModuleType(package_name) + package.__path__ = [str(_CUDNN_ROOT)] + jax_package = types.ModuleType(f"{package_name}._jax") + jax_package.__path__ = [str(_CUDNN_ROOT / "_jax")] + monkeypatch.setitem(sys.modules, package_name, package) + monkeypatch.setitem(sys.modules, f"{package_name}._jax", jax_package) + monkeypatch.setitem(sys.modules, f"{package_name}._result", _COMMON_MODULE) + + spec = importlib.util.spec_from_file_location( + f"{package_name}._jax.result", + _CUDNN_ROOT / "_jax" / "result.py", + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + JaxTupleDict = module.TupleDict + + result = JaxTupleDict(output=1, amax=2) + leaves, treedef = jax.tree_util.tree_flatten(result) + + assert leaves == [1, 2] + assert jax.tree_util.tree_unflatten(treedef, leaves) == result diff --git a/test/python/fe_api/test_shared_tensor_desc.py b/test/python/fe_api/test_shared_tensor_desc.py index 45d8b3bc3..218ad44b4 100644 --- a/test/python/fe_api/test_shared_tensor_desc.py +++ b/test/python/fe_api/test_shared_tensor_desc.py @@ -132,6 +132,27 @@ def test_compact_descriptor_supports_scalars(self): self.assertEqual(desc.stride, ()) self.assertEqual(desc.stride_order, ()) self.assertEqual(desc.ndim, 0) + self.assertTrue(desc.is_compact()) + + def test_is_compact_uses_canonical_or_explicit_dimension_order(self): + row_major = self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2, 3, 4), + ) + self.assertTrue(row_major.is_compact()) + self.assertTrue(row_major.is_compact((2, 1, 0))) + self.assertFalse(row_major.is_compact((0, 1, 2))) + + size_one = self.tensor_desc.TensorDesc( + dtype=_DataType.FLOAT, + shape=(2, 1, 4), + stride=(4, 17, 1), + stride_order=(2, 0, 1), + ) + self.assertTrue(size_one.is_compact((2, 1, 0))) + + with self.assertRaisesRegex(ValueError, "must be a permutation"): + row_major.is_compact((2, 2, 0)) def test_compact_descriptor_preserves_order_for_zero_extents(self): desc = self.tensor_desc.make_compact_tensor_desc( From 829688f8af6a7bb423f34fd2072c22089447ca94 Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Tue, 7 Jul 2026 21:50:57 -0500 Subject: [PATCH 05/13] Add JAX support for dense GEMM SwiGLU --- docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md | 85 +++- python/cudnn/_cute_compiler.py | 40 ++ python/cudnn/_dense_gemm.py | 215 +++++++++ python/cudnn/_jax/api_base.py | 30 +- python/cudnn/_jax/gemm.py | 62 +++ .../utils/compiler.py | 34 +- python/cudnn/gemm_swiglu/__init__.py | 25 +- python/cudnn/gemm_swiglu/api.py | 54 ++- .../dense_gemm_persistent_swiglu.py | 4 +- python/cudnn/gemm_swiglu/jax.py | 278 +++++++++++ python/cudnn/gemm_swiglu/op.py | 165 +++++++ python/cudnn/jax/__init__.py | 4 +- test/python/fe_api/test_dense_gemm_helpers.py | 198 ++++++++ test/python/fe_api/test_dsa_compiler.py | 6 +- .../python/fe_api/test_gemm_swiglu_imports.py | 190 ++++++++ test/python/fe_api/test_gemm_swiglu_op.py | 241 ++++++++++ test/python/fe_api/test_jax_api_base.py | 42 ++ test/python/fe_api/test_jax_gemm_layout.py | 127 +++++ test/python/fe_api/test_jax_gemm_swiglu.py | 201 ++++++++ .../fe_api/test_jax_gemm_swiglu_contract.py | 438 ++++++++++++++++++ .../python/fe_api/test_jax_import_contract.py | 2 +- 21 files changed, 2377 insertions(+), 64 deletions(-) create mode 100644 python/cudnn/_cute_compiler.py create mode 100644 python/cudnn/_dense_gemm.py create mode 100644 python/cudnn/_jax/gemm.py create mode 100644 python/cudnn/gemm_swiglu/jax.py create mode 100644 python/cudnn/gemm_swiglu/op.py create mode 100644 test/python/fe_api/test_dense_gemm_helpers.py create mode 100644 test/python/fe_api/test_gemm_swiglu_imports.py create mode 100644 test/python/fe_api/test_gemm_swiglu_op.py create mode 100644 test/python/fe_api/test_jax_gemm_layout.py create mode 100644 test/python/fe_api/test_jax_gemm_swiglu.py create mode 100644 test/python/fe_api/test_jax_gemm_swiglu_contract.py diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md index e03c041e4..1e9816f4e 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md @@ -10,6 +10,9 @@ This API supports two modes: 1. **Standard mode**: High-precision GEMM with SwiGLU epilogue 2. **Quantized mode** (block-scaled): Low-precision GEMM using block scaling supporting FP4 and FP8 data types +The Torch API supports both modes. The initial JAX API supports standard mode; +the block-scaled JAX surface will be added with the shared scale-factor ABI. + ### Shapes - Inputs: @@ -178,12 +181,72 @@ gemm.execute( ) ``` +### JAX API (Standard Mode) + +The JAX wrapper is already decorated with `jax.jit`. Public arrays are compact +row-major arrays whose axis order is described explicitly: + +```python +import jax +import jax.numpy as jnp + +from cudnn.jax import gemm_swiglu_wrapper_sm100 + +# A: (L, M, K), B: (L, N, K) +result = gemm_swiglu_wrapper_sm100( + a, + b, + alpha=1.0, + a_layout="LMK", + b_layout="LNK", + c_layout="LMN", + ab12_dtype=jnp.float32, + c_dtype=jnp.bfloat16, + acc_dtype=jnp.float32, +) + +# AB12: (L, M, N), C: (L, M, N/2) +ab12 = result["ab12_tensor"] +c = result["c_tensor"] +``` + +The advanced class API specializes from shape/dtype metadata and may optionally +accept explicit output exemplars: + +```python +from cudnn.jax import GemmSwigluSm100 + +operation = GemmSwigluSm100( + jax.ShapeDtypeStruct(a.shape, a.dtype), + jax.ShapeDtypeStruct(b.shape, b.dtype), + a_layout="LMK", + b_layout="LNK", + c_layout="LMN", +) +result = operation(a, b) +``` + +Supported public layouts are: + +- A: `LMK` (K-major) or `LKM` (M-major) +- B: `LNK` (K-major) or `LKN` (N-major) +- AB12/C: `LMN` (N-major) or `LNM` (M-major) + +An explicit `target_compute_capability` may be used after calling +`cudnn.jax.disable_device_compatibility_checks(True)` to bypass the local-device +compatibility check. CuTe lowering still queries CUDA occupancy and therefore +requires access to a CUDA device. + --- ## Parameters ### Input/Output tensors +The shapes and strides below use the kernel's canonical axis order, which is +also the Torch API order. JAX arrays use the public axis orders selected by +`a_layout`, `b_layout`, and `c_layout` above. + - Input tensor **A**: `a_tensor` (wrapper) or `sample_a`, `a_tensor` (class) - Shape: `(M, K, L)` - Stride: `(1, M, M·K)` for `m`-major or `(K, 1, M·K)` for `k`-major @@ -234,20 +297,20 @@ gemm.execute( - `alpha: float` - Scalar multiplier applied to the GEMM result before SwiGLU. - Default: `1.0` -- `acc_dtype: torch.dtype` +- `acc_dtype`: `torch.dtype` for Torch or a JAX dtype-like value for JAX - Accumulator dtype. - - Standard mode: `{float32, float16}`. Default: `torch.float32` + - Standard mode: `{float32, float16}`. Default: `float32` - Quantized mode: Must be `float32` - `mma_tiler_mn: Tuple[int, int]` - Kernel tile size `(TILE_M, TILE_N)`. Default: `(128, 128)` - `TILE_M ∈ {128, 256}` - - Standard mode: `TILE_N ∈ {32, 64, ..., 224, 256}` + - Standard mode: `TILE_N ∈ {64, 128, 192, 256}` - Quantized mode: `TILE_N ∈ {64, 128, 192, 256}` - `cluster_shape_mn: Tuple[int, int] | None` - Thread Block cluster shape `(CLUSTER_M, CLUSTER_N)` - Constraints: positive powers of 2, `CLUSTER_M*CLUSTER_N ≤ 16`. - Default: `(1,1)` if `mma_tiler_mn[0] != 256` else `(2,2)`. -- CUDA stream (`current_stream` in class API, `stream` in wrapper) +- CUDA stream (`current_stream` in the Torch class API, `stream` in the Torch wrapper). JAX supplies its stream during lowering. - **Quantization-specific parameters** - `sf_vec_size: int` - Scale factor vector size (number of elements per scale factor) @@ -265,11 +328,9 @@ gemm.execute( ### Wrapper-specific parameters: `gemm_swiglu_wrapper_sm100` - `a_tensor`, `b_tensor`: see Input/Output tensors -- `c_major: str`: see Input/Output tensors. Default: `"n"` -- `ab12_dtype: torch.dtype`: see Input/Output tensors. Default: `torch.float32` -- `c_dtype: torch.dtype`: see Input/Output tensors. Default: `torch.float16` -- `sfa_tensor`, `sfb_tensor`, `norm_const_tensor`: see Quantization-specific tensors -- `sf_vec_size`, `vector_f32`, `ab12_stages`: see Quantization-specific parameters +- Torch: `c_major` selects `"m"` or `"n"`; `ab12_dtype` defaults to `torch.float32`; `c_dtype` defaults to `torch.float16`. +- JAX: `a_layout`, `b_layout`, and `c_layout` select public axis order; `ab12_dtype` defaults to `jnp.float32`; `c_dtype` defaults to `jnp.float16`. +- Torch quantized mode: `sfa_tensor`, `sfb_tensor`, `norm_const_tensor`, `sf_vec_size`, `vector_f32`, and `ab12_stages` use the quantization rules above. ### Wrapper return values @@ -290,8 +351,8 @@ Tuple unpacking order is always: #### `GemmSwigluSm100` (constructor) -- `sample_a`, `sample_b`, `sample_ab12`, `sample_c` – see Input/Output tensors -- `sample_sfa`, `sample_sfb`, `sample_sfc`, `sample_amax`, `sample_norm_const` – see Scale factor tensors (quantized mode) +- `sample_a`, `sample_b`, `sample_ab12`, `sample_c` – see Input/Output tensors. JAX requires only the input samples and accepts the output samples optionally. +- `sample_sfa`, `sample_sfb`, `sample_sfc`, `sample_amax`, `sample_norm_const` – Torch quantized mode only. #### `GemmSwigluSm100.execute` @@ -305,7 +366,7 @@ Tuple unpacking order is always: ### Layouts and strides - `AB12` and `C` must have the same major order. -- `A`, `B`, `AB12` must be 16-byte aligned along the contiguous dimension. +- `A`, `B`, `AB12`, and `C` must be 16-byte aligned along the contiguous dimension. - For FP4 inputs (quantized mode): `A` and `B` must be `k`-major, `AB12` must be `n`-major. ### Dtypes diff --git a/python/cudnn/_cute_compiler.py b/python/cudnn/_cute_compiler.py new file mode 100644 index 000000000..7e6c8cb21 --- /dev/null +++ b/python/cudnn/_cute_compiler.py @@ -0,0 +1,40 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free CuTe compilation options for framework adapters.""" + +from __future__ import annotations + +_ARCH_MAP = { + 90: "sm_90a", + 100: "sm_100a", + 103: "sm_103a", + 107: "sm_100f", +} + + +def gpu_arch_flag_for_compute_capability(compute_capability: int) -> str: + """Return the explicit CuTe architecture flag for a compilation target.""" + + if isinstance(compute_capability, bool) or not isinstance(compute_capability, int): + raise TypeError(f"compute_capability must be an int, got {type(compute_capability).__name__}") + try: + return _ARCH_MAP[compute_capability] + except KeyError as error: + supported = ", ".join(f"SM{value}" for value in sorted(_ARCH_MAP)) + raise RuntimeError(f"Unsupported GPU compute capability SM{compute_capability} for CuTe kernels; supported targets are {supported}") from error + + +def compile_options_for_target(compute_capability: int, extra: str = "") -> str: + """Build CuTe compile options for an explicit framework-supplied target.""" + + parts = [ + "--enable-tvm-ffi", + f"--gpu-arch {gpu_arch_flag_for_compute_capability(compute_capability)}", + ] + if extra: + parts.append(extra) + return " ".join(parts) + + +__all__ = ["compile_options_for_target", "gpu_arch_flag_for_compute_capability"] diff --git a/python/cudnn/_dense_gemm.py b/python/cudnn/_dense_gemm.py new file mode 100644 index 000000000..2287ef1a6 --- /dev/null +++ b/python/cudnn/_dense_gemm.py @@ -0,0 +1,215 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral validation helpers for dense GEMM operations.""" + +from __future__ import annotations + +from operator import index +from typing import Iterable + +from . import data_type +from ._tensor_desc import TensorDesc + +_STANDARD_MMA_M = (128, 256) +_STANDARD_MMA_N = tuple(range(32, 257, 32)) + +_DATA_TYPE_BITS_BY_NAME = { + "BOOLEAN": 8, + "INT4": 4, + "INT8": 8, + "UINT8": 8, + "HALF": 16, + "BFLOAT16": 16, + "FLOAT": 32, + "DOUBLE": 64, + "INT32": 32, + "INT64": 64, + "FP4_E2M1": 4, + "FP8_E4M3": 8, + "FP8_E5M2": 8, + "FP8_E8M0": 8, + "FAST_FLOAT_FOR_FP8": 32, +} + + +def _require_tensor_desc(tensor: TensorDesc, label: str) -> None: + if not isinstance(tensor, TensorDesc): + raise TypeError(f"{label} must be a TensorDesc, got {type(tensor).__name__}") + + +def _integer_pair(value: Iterable[int], label: str) -> tuple[int, int]: + try: + values = tuple(value) + except TypeError as error: + raise TypeError(f"{label} must contain two integers, got {value!r}") from error + + if len(values) != 2: + raise ValueError(f"{label} must contain two integers, got {values}") + + normalized = [] + for entry in values: + if isinstance(entry, bool): + raise TypeError(f"{label} must contain integers, got {values}") + try: + normalized.append(index(entry)) + except TypeError as error: + raise TypeError(f"{label} must contain integers, got {values}") from error + return normalized[0], normalized[1] + + +def _format_values(values: Iterable[int]) -> str: + return "{" + ", ".join(str(value) for value in values) + "}" + + +def require_gemm_inputs(a: TensorDesc, b: TensorDesc) -> tuple[int, int, int, int]: + """Validate canonical ``A[M,K,L]`` and ``B[N,K,L]`` descriptors. + + Returns ``(M, N, K, L)``. Layout and dtype validation are deliberately + separate because their supported combinations differ between kernels. + """ + + _require_tensor_desc(a, "A") + _require_tensor_desc(b, "B") + if a.ndim != 3: + raise ValueError(f"A must have rank 3, got shape {a.shape}") + if b.ndim != 3: + raise ValueError(f"B must have rank 3, got shape {b.shape}") + + m, k, l = a.shape + n, b_k, b_l = b.shape + for axis, extent in (("M", m), ("N", n), ("K", k), ("L", l)): + if extent <= 0: + raise ValueError(f"{axis} must be positive, got {extent}") + + if (b_k, b_l) != (k, l): + raise ValueError(f"B shape mismatch: expected (N, {k}, {l}), got {b.shape}") + return m, n, k, l + + +def require_tensor_shape( + tensor: TensorDesc, + expected: tuple[int, ...], + *, + label: str | None = None, +) -> None: + """Require a descriptor to have an exact logical shape.""" + + tensor_label = label or getattr(tensor, "name", "") or "Tensor" + _require_tensor_desc(tensor, tensor_label) + expected = tuple(expected) + if tensor.shape != expected: + raise ValueError(f"{tensor_label} must have shape {expected}, got {tensor.shape}") + + +def require_compact_major( + tensor: TensorDesc, + mode0_label: str, + mode1_label: str, +) -> str: + """Resolve the compact major mode of a canonical rank-three GEMM tensor. + + Canonical GEMM tensors may make mode 0 or mode 1 contiguous. The batch + mode (mode 2) must remain outermost. The returned label is supplied by the + caller, for example ``("m", "k")`` for ``A[M,K,L]``. + """ + + tensor_label = getattr(tensor, "name", "") or "Tensor" + _require_tensor_desc(tensor, tensor_label) + if tensor.ndim != 3: + raise ValueError(f"{tensor_label} must have rank 3, got shape {tensor.shape}") + + labels_by_order = { + (0, 1, 2): mode0_label, + (1, 0, 2): mode1_label, + } + major = labels_by_order.get(tensor.stride_order) + if major is None or not tensor.is_compact(tensor.stride_order): + raise ValueError(f"{tensor_label} must be compact {mode0_label}-major or " f"{mode1_label}-major, got shape {tensor.shape} and stride {tensor.stride}") + return major + + +def data_type_bits(dtype: data_type) -> int: + """Return the storage width of a canonical cuDNN data type.""" + + if not isinstance(dtype, data_type): + raise TypeError(f"dtype must be a cudnn.data_type, got {type(dtype).__name__}") + + for name, bits in _DATA_TYPE_BITS_BY_NAME.items(): + member = getattr(data_type, name, None) + if member is not None and dtype == member: + return bits + raise ValueError(f"Unsupported cuDNN data type {dtype}") + + +def require_16_byte_alignment(tensor: TensorDesc) -> None: + """Require the tensor's contiguous extent to cover whole 16-byte units. + + This is the logical-extent requirement used by dense GEMM TMA operands; + framework adapters remain responsible for the base-pointer alignment. + ``require_compact_major`` should be called first for GEMM matrices. + """ + + tensor_label = getattr(tensor, "name", "") or "Tensor" + _require_tensor_desc(tensor, tensor_label) + if tensor.ndim == 0: + raise ValueError(f"{tensor_label} must have at least one dimension") + + contiguous_dimension = tensor.stride_order[0] + bits = data_type_bits(tensor.cudnn_dtype) + extent = tensor.shape[contiguous_dimension] + if extent * bits % 128 != 0: + required_multiple = 128 // bits + raise ValueError(f"{tensor_label} contiguous extent must be a multiple of " f"{required_multiple} elements for 16-byte alignment, got {extent}") + + +def require_mma_tiler( + mma_tiler_mn: Iterable[int], + *, + allowed_m: Iterable[int] = _STANDARD_MMA_M, + allowed_n: Iterable[int] = _STANDARD_MMA_N, +) -> tuple[int, int]: + """Validate and normalize an ``(M, N)`` MMA tile.""" + + mma_m, mma_n = _integer_pair(mma_tiler_mn, "mma_tiler_mn") + allowed_m = tuple(allowed_m) + allowed_n = tuple(allowed_n) + if mma_m not in allowed_m: + raise ValueError(f"mma_tiler_mn[0] must be in {_format_values(allowed_m)}, got {mma_m}") + if mma_n not in allowed_n: + raise ValueError(f"mma_tiler_mn[1] must be in {_format_values(allowed_n)}, got {mma_n}") + return mma_m, mma_n + + +def require_cluster_shape( + cluster_shape_mn: Iterable[int], + *, + cta_group_size: int = 1, + max_cluster_size: int = 16, +) -> tuple[int, int]: + """Validate and normalize an ``(M, N)`` Blackwell cluster shape.""" + + cluster_m, cluster_n = _integer_pair(cluster_shape_mn, "cluster_shape_mn") + if isinstance(cta_group_size, bool) or not isinstance(cta_group_size, int) or cta_group_size <= 0: + raise ValueError(f"cta_group_size must be a positive integer, got {cta_group_size!r}") + if isinstance(max_cluster_size, bool) or not isinstance(max_cluster_size, int) or max_cluster_size <= 0: + raise ValueError(f"max_cluster_size must be a positive integer, got {max_cluster_size!r}") + + if cluster_m <= 0 or cluster_n <= 0 or cluster_m & (cluster_m - 1) or cluster_n & (cluster_n - 1): + raise ValueError("cluster_shape_mn entries must be positive powers of two, " f"got {(cluster_m, cluster_n)}") + if cluster_m * cluster_n > max_cluster_size: + raise ValueError(f"cluster_shape_mn product must not exceed {max_cluster_size}, " f"got {(cluster_m, cluster_n)}") + if cluster_m % cta_group_size != 0: + raise ValueError(f"cluster_shape_mn[0] must be divisible by CTA group size " f"{cta_group_size}, got {cluster_m}") + return cluster_m, cluster_n + + +__all__ = [ + "data_type_bits", + "require_16_byte_alignment", + "require_cluster_shape", + "require_compact_major", + "require_gemm_inputs", + "require_mma_tiler", + "require_tensor_shape", +] diff --git a/python/cudnn/_jax/api_base.py b/python/cudnn/_jax/api_base.py index 3180b1f34..1a45184da 100644 --- a/python/cudnn/_jax/api_base.py +++ b/python/cudnn/_jax/api_base.py @@ -7,6 +7,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable +from operator import index from typing import Any from .. import data_type @@ -14,7 +15,7 @@ from .layout import compact_stride, normalize_mode, to_canonical_axes, to_cutlass_layout, to_public_axes _DEVICE_COMPATIBILITY_CHECKS_DISABLED = False -_DEVICE_COMPATIBILITY_CHECK_DISABLE_HINT = "Call cudnn.jax.disable_device_compatibility_checks(True) before tracing to allow AOT or remote compilation." +_DEVICE_COMPATIBILITY_CHECK_DISABLE_HINT = "Call cudnn.jax.disable_device_compatibility_checks(True) before tracing for a different or unavailable local GPU." _TARGET_COMPUTE_CAPABILITY_HINT = "Pass target_compute_capability explicitly when tracing for AOT or on a host without the target GPU." @@ -78,6 +79,27 @@ def _require_array_metadata(value: Any, name: str) -> tuple[Any, ...]: return tuple(value.shape) +def _normalize_stride_order(rank: int, stride_order: tuple[int, ...] | None) -> tuple[int, ...]: + if stride_order is None: + return tuple(reversed(range(rank))) + + normalized = [] + for dimension in stride_order: + if isinstance(dimension, bool): + raise TypeError(f"public_stride_order entries must be integers, got {dimension!r}") + try: + normalized.append(index(dimension)) + except TypeError as error: + raise TypeError(f"public_stride_order entries must be integers, got {dimension!r}") from error + + normalized_order = tuple(normalized) + if len(normalized_order) != rank: + raise ValueError(f"public_stride_order rank mismatch: expected {rank}, got {len(normalized_order)}") + if tuple(sorted(normalized_order)) != tuple(range(rank)): + raise ValueError(f"public_stride_order must be a permutation of [0, {rank - 1}], got {normalized_order}") + return normalized_order + + class JaxApiBase(ABC): """Common tensor metadata, validation, and CuTe binding for JAX adapters.""" @@ -223,18 +245,20 @@ def _to_tensor_desc( name: str, *, mode: tuple[int, ...] | None = None, + public_stride_order: tuple[int, ...] | None = None, init_value: bool | int | float | None = None, ) -> JaxTensorDesc: """Describe a public JAX array in canonical kernel-axis order. ``mode[kernel_axis]`` selects the corresponding public array axis. - Public JAX arrays are modeled as compact row-major buffers. + ``public_stride_order`` lists the public array dimensions from fastest + to slowest. It defaults to compact row-major storage. """ public_shape = _require_array_metadata(value, name) rank = len(public_shape) mode = normalize_mode(rank, mode) - public_stride_order = tuple(reversed(range(rank))) + public_stride_order = _normalize_stride_order(rank, public_stride_order) public_stride = compact_stride(public_shape, public_stride_order) canonical_axis_by_public_axis = to_public_axes(tuple(range(rank)), mode) diff --git a/python/cudnn/_jax/gemm.py b/python/cudnn/_jax/gemm.py new file mode 100644 index 000000000..7a82f8d40 --- /dev/null +++ b/python/cudnn/_jax/gemm.py @@ -0,0 +1,62 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free public-axis bindings for dense JAX GEMMs.""" + +from __future__ import annotations + +from .layout import mode_from_layout + +GEMM_A_LAYOUTS = ("LMK", "LKM") +GEMM_B_LAYOUTS = ("LNK", "LKN") +GEMM_OUTPUT_LAYOUTS = ("LMN", "LNM") + +ROW_MAJOR_STRIDE_ORDER_3D = (2, 1, 0) +BLOCK_SCALE_STRIDE_ORDER = (3, 1, 0, 4, 2, 5) +PROBABILITY_STRIDE_ORDER = (0, 1, 2) + + +def require_layout(name: str, layout: str, supported: tuple[str, ...]) -> str: + """Validate an explicit, case-sensitive public axis-order string.""" + + if not isinstance(layout, str): + raise TypeError(f"{name} must be a string, got {type(layout).__name__}") + if layout not in supported: + choices = ", ".join(repr(value) for value in supported) + raise ValueError(f"{name} must be one of ({choices}), got {layout!r}") + return layout + + +def gemm_a_mode(layout: str) -> tuple[int, ...]: + """Map public A layout ``LMK`` or ``LKM`` to canonical ``MKL`` axes.""" + + layout = require_layout("a_layout", layout, GEMM_A_LAYOUTS) + return mode_from_layout(layout, kernel_axes="MKL") + + +def gemm_b_mode(layout: str) -> tuple[int, ...]: + """Map public B layout ``LNK`` or ``LKN`` to canonical ``NKL`` axes.""" + + layout = require_layout("b_layout", layout, GEMM_B_LAYOUTS) + return mode_from_layout(layout, kernel_axes="NKL") + + +def gemm_output_mode(layout: str, *, name: str = "c_layout") -> tuple[int, ...]: + """Map public output layout ``LMN`` or ``LNM`` to canonical ``MNL`` axes.""" + + layout = require_layout(name, layout, GEMM_OUTPUT_LAYOUTS) + return mode_from_layout(layout, kernel_axes="MNL") + + +__all__ = [ + "BLOCK_SCALE_STRIDE_ORDER", + "GEMM_A_LAYOUTS", + "GEMM_B_LAYOUTS", + "GEMM_OUTPUT_LAYOUTS", + "PROBABILITY_STRIDE_ORDER", + "ROW_MAJOR_STRIDE_ORDER_3D", + "gemm_a_mode", + "gemm_b_mode", + "gemm_output_mode", + "require_layout", +] diff --git a/python/cudnn/deepseek_sparse_attention/utils/compiler.py b/python/cudnn/deepseek_sparse_attention/utils/compiler.py index 9fdfbe05b..9ce687138 100644 --- a/python/cudnn/deepseek_sparse_attention/utils/compiler.py +++ b/python/cudnn/deepseek_sparse_attention/utils/compiler.py @@ -18,39 +18,7 @@ from functools import lru_cache -# (compute_capability) → cute DSL --gpu-arch flag value. -# H100, B200/B300, and sm_100f require architecture-specific variants because -# the kernels use TMA / tcgen05 instructions that are only guaranteed to lower -# correctly under the matching SASS gencode. -_ARCH_MAP = { - 90: "sm_90a", # Hopper H100 - 100: "sm_100a", # Blackwell B200 - 103: "sm_103a", # Blackwell Ultra B300 - 107: "sm_100f", -} - - -def gpu_arch_flag_for_compute_capability(compute_capability: int) -> str: - """Return the CuTe architecture flag for an explicit compilation target.""" - - if isinstance(compute_capability, bool) or not isinstance(compute_capability, int): - raise TypeError(f"compute_capability must be an int, got {type(compute_capability).__name__}") - try: - return _ARCH_MAP[compute_capability] - except KeyError as error: - raise RuntimeError( - f"Unsupported GPU compute capability SM{compute_capability} for DSA CuTe kernels. " - "Add it to deepseek_sparse_attention/utils/compiler.py::_ARCH_MAP." - ) from error - - -def compile_options_for_target(compute_capability: int, extra: str = "") -> str: - """Build CuTe compile options for a framework-supplied GPU target.""" - - parts = ["--enable-tvm-ffi", f"--gpu-arch {gpu_arch_flag_for_compute_capability(compute_capability)}"] - if extra: - parts.append(extra) - return " ".join(parts) +from ..._cute_compiler import compile_options_for_target, gpu_arch_flag_for_compute_capability @lru_cache(maxsize=None) diff --git a/python/cudnn/gemm_swiglu/__init__.py b/python/cudnn/gemm_swiglu/__init__.py index cf2170c32..6f11dd595 100644 --- a/python/cudnn/gemm_swiglu/__init__.py +++ b/python/cudnn/gemm_swiglu/__init__.py @@ -1,9 +1,18 @@ -from .api import ( - GemmSwigluSm100, - gemm_swiglu_wrapper_sm100, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy Torch API and framework-neutral GEMM + SwiGLU operation exports.""" -__all__ = [ - "GemmSwigluSm100", - "gemm_swiglu_wrapper_sm100", -] +from .._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ("GemmSwigluSm100Op",), + "api": ( + "GemmSwigluSm100", + "gemm_swiglu_wrapper_sm100", + ), + }, + submodules=("api", "op"), +) diff --git a/python/cudnn/gemm_swiglu/api.py b/python/cudnn/gemm_swiglu/api.py index 2daeaa740..5a49b9034 100644 --- a/python/cudnn/gemm_swiglu/api.py +++ b/python/cudnn/gemm_swiglu/api.py @@ -41,8 +41,10 @@ import cutlass.cute as cute from cutlass.cute.runtime import make_fake_stream -from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn import data_type +from cudnn.datatypes import _convert_to_cutlass_data_type, _torch_to_cudnn_data_type from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from .op import GemmSwigluSm100Op import os @@ -104,6 +106,24 @@ def __init__( self._logger.debug("Quantization arguments provided, using quantized GEMM swiglu kernel") self._kernel = Sm100BlockScaledPersistentDenseGemmKernel + self._op = None + if self._kernel is PersistentDenseGemmKernel: + canonical_acc_dtype = _torch_to_cudnn_data_type(self.acc_dtype) + if canonical_acc_dtype is None: + canonical_acc_dtype = data_type.NOT_SET + self._op = GemmSwigluSm100Op( + a=self.a_desc, + b=self.b_desc, + ab12=self.ab12_desc, + c=self.c_desc, + alpha=self.alpha, + acc_dtype=canonical_acc_dtype, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + ) + self.mma_tiler_mn = self._op.mma_tiler_mn + self.cluster_shape_mn = self._op.cluster_shape_mn + self._logger.debug( f"__init__ completed with args: sample_a {self.a_desc.shape}, sample_b {self.b_desc.shape}, sample_ab12 {self.ab12_desc.shape}, sample_c {self.c_desc.shape}, alpha {alpha}, acc_dtype {acc_dtype}, mma_tiler_mn {mma_tiler_mn}, cluster_shape_mn {cluster_shape_mn}, sample_sfa {self.sfa_desc.shape if self.sfa_desc is not None else None}, sample_sfb {self.sfb_desc.shape if self.sfb_desc is not None else None}, sample_amax {self.amax_desc.shape if self.amax_desc is not None else None}, sample_sfc {self.sfc_desc.shape if self.sfc_desc is not None else None}, sample_norm_const {self.norm_const_desc.shape if self.norm_const_desc is not None else None}, sf_vec_size {sf_vec_size}, vector_f32 {vector_f32}, ab12_stages {ab12_stages}" ) @@ -111,6 +131,38 @@ def __init__( self._interpret_uint8_as_fp4x2 = True def check_support(self) -> bool: + """Validate the shared operation contract and Torch device support.""" + + self._is_supported = False + if self._kernel is Sm100BlockScaledPersistentDenseGemmKernel: + return self._check_quantized_support() + + if self._op is None: + raise RuntimeError("Standard GEMM + SwiGLU operation was not initialized") + self._op.check_support() + self.mma_tiler_mn = self._op.mma_tiler_mn + self.cluster_shape_mn = self._op.cluster_shape_mn + + self.ab_dtype = self.a_desc.dtype + self.ab12_dtype = self.ab12_desc.dtype + self.c_dtype = self.c_desc.dtype + + self._logger.debug("Checking environment") + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available") + device = torch.cuda.current_device() + major, minor = torch.cuda.get_device_capability(device) + compute_capability = major * 10 + minor + if compute_capability < 100: + raise RuntimeError(f"GemmSwiglu requires SM100+ compute capability, but found SM{compute_capability} on device {device}") + + self._is_supported = True + self._logger.debug("check_support completed successfully") + return True + + def _check_quantized_support(self) -> bool: + """Preserve the existing Torch-only block-scaled validation path.""" + self._logger.debug("Entering check_support") self._logger.debug("Checking tensor shapes, strides, and dtypes") diff --git a/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py b/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py index 28f6aff86..5b6769395 100644 --- a/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py +++ b/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py @@ -78,7 +78,7 @@ see detailed valid dtype combinations in below PersistentDenseGemmKernel class documentation * A/B tensor must have the same data type * Mma tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) -* Mma tiler N must be 32-256, step 32 +* Mma tiler N must be 64/128/192/256 so each tile contains complete SwiGLU block pairs * Cluster shape M/N must be positive and power of 2, total cluster size <= 16 * Cluster shape M must be multiple of 2 if use_2cta_instrs=True * The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, @@ -125,7 +125,7 @@ class PersistentDenseGemmKernel: :note: Constraints: - MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) - - MMA tiler N must be 32-256, step 32 + - MMA tiler N must be 64/128/192/256 so each tile contains complete SwiGLU block pairs - Cluster shape M must be multiple of 2 if use_2cta_instrs=True - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 diff --git a/python/cudnn/gemm_swiglu/jax.py b/python/cudnn/gemm_swiglu/jax.py new file mode 100644 index 000000000..91151045b --- /dev/null +++ b/python/cudnn/gemm_swiglu/jax.py @@ -0,0 +1,278 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX API for the standard dense GEMM + SwiGLU kernel.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any + +import jax +import jax.numpy as jnp + +from .. import data_type +from .._cute_compiler import compile_options_for_target +from .._dense_gemm import require_gemm_inputs +from .._jax import JaxApiBase, JaxTensorDesc, TupleDict +from .._jax.datatypes import jax_to_cudnn_dtype +from .._jax.gemm import gemm_a_mode, gemm_b_mode, gemm_output_mode +from .._jax.layout import to_public_axes +from .op import GemmSwigluSm100Op + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) + + +def _normalize_dtype(value: Any | None, default: Any, name: str) -> Any: + try: + return jnp.dtype(default if value is None else value) + except TypeError as error: + raise TypeError(f"{name} must be a JAX dtype, got {value!r}") from error + + +class GemmSwigluSm100(JaxApiBase): + """JAX callable specialized from dense GEMM input metadata. + + ``a_layout``, ``b_layout``, and ``c_layout`` describe the public JAX axis + order. Descriptors stored by the operation remain in the kernel's + canonical ``MKL``, ``NKL``, and ``MNL`` orders. + + This first JAX surface intentionally selects only the standard dense + kernel. Block-scaled inputs and quantized outputs are added separately so + their packed scale-factor ABI does not complicate the base contract. + """ + + def __init__( + self, + sample_a: Any, + sample_b: Any, + *, + sample_ab12: Any | None = None, + sample_c: Any | None = None, + alpha: float = 1.0, + c_layout: str = "LMN", + ab12_dtype: Any | None = None, + c_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: tuple[int, int] | None = None, + a_layout: str = "LMK", + b_layout: str = "LNK", + target_compute_capability: int | None = None, + ) -> None: + self.a_layout = a_layout + self.b_layout = b_layout + self.c_layout = c_layout + self.a_mode = gemm_a_mode(a_layout) + self.b_mode = gemm_b_mode(b_layout) + self.output_mode = gemm_output_mode(c_layout) + + self.target_compute_capability = self._resolve_compute_capability( + target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "GemmSwigluSm100", + ) + self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=self.a_mode) + self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=self.b_mode) + self.acc_dtype = _normalize_dtype(acc_dtype, jnp.float32, "acc_dtype") + + if (sample_ab12 is None) != (sample_c is None): + raise ValueError("sample_ab12 and sample_c must be provided together") + if sample_ab12 is None: + resolved_ab12_dtype = _normalize_dtype(ab12_dtype, jnp.float32, "ab12_dtype") + resolved_c_dtype = _normalize_dtype(c_dtype, jnp.float16, "c_dtype") + self.ab12_desc, self.c_desc = self._default_output_descs( + resolved_ab12_dtype, + resolved_c_dtype, + ) + else: + self.ab12_desc = self._to_tensor_desc( + sample_ab12, + "sample_ab12", + mode=self.output_mode, + ) + self.c_desc = self._to_tensor_desc( + sample_c, + "sample_c", + mode=self.output_mode, + ) + self._check_requested_output_dtype(ab12_dtype, self.ab12_desc, "ab12_dtype") + self._check_requested_output_dtype(c_dtype, self.c_desc, "c_dtype") + + acc_cudnn_dtype = jax_to_cudnn_dtype(self.acc_dtype) + if acc_cudnn_dtype == data_type.NOT_SET: + raise ValueError(f"Unsupported JAX accumulator dtype {self.acc_dtype}") + + self._op = GemmSwigluSm100Op( + a=self.a_desc, + b=self.b_desc, + ab12=self.ab12_desc, + c=self.c_desc, + alpha=alpha, + acc_dtype=acc_cudnn_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + ) + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + + def _default_output_descs( + self, + ab12_dtype: Any, + c_dtype: Any, + ) -> tuple[JaxTensorDesc, JaxTensorDesc]: + m, n, _, batch = require_gemm_inputs(self.a_desc, self.b_desc) + if n % 2: + raise ValueError(f"SwiGLU requires an even N dimension, got {n}") + + return ( + self._to_tensor_desc( + jax.ShapeDtypeStruct( + to_public_axes((m, n, batch), self.output_mode), + ab12_dtype, + ), + "sample_ab12", + mode=self.output_mode, + ), + self._to_tensor_desc( + jax.ShapeDtypeStruct( + to_public_axes((m, n // 2, batch), self.output_mode), + c_dtype, + ), + "sample_c", + mode=self.output_mode, + ), + ) + + @staticmethod + def _check_requested_output_dtype( + requested: Any | None, + desc: JaxTensorDesc, + name: str, + ) -> None: + if requested is None: + return + requested_dtype = _normalize_dtype(requested, requested, name) + actual_dtype = jnp.dtype(desc.dtype) + if requested_dtype != actual_dtype: + raise ValueError(f"{name}={requested_dtype} does not match the explicit sample dtype {actual_dtype}") + + def check_support(self) -> bool: + return self._op.check_support() + + def __call__(self, a_tensor: Any, b_tensor: Any) -> TupleDict: + self.check_support() + self._check_tensor_signature(a_tensor, self.a_desc, mode=self.a_mode) + self._check_tensor_signature(b_tensor, self.b_desc, mode=self.b_mode) + + ab12_tensor, c_tensor = self._call_kernel( + (a_tensor, b_tensor), + output_descs=(self.ab12_desc, self.c_desc), + input_spec=( + self._to_tensor_spec(self.a_desc, mode=self.a_mode), + self._to_tensor_spec(self.b_desc, mode=self.b_mode), + ), + output_spec=( + self._to_tensor_spec(self.ab12_desc, mode=self.output_mode), + self._to_tensor_spec(self.c_desc, mode=self.output_mode), + ), + compile_options=compile_options_for_target(self.target_compute_capability), + ) + return TupleDict( + ab12_tensor=ab12_tensor, + c_tensor=c_tensor, + sfc_tensor=None, + amax_tensor=None, + ) + + def _launch( + self, + inputs: tuple[Any, ...], + outputs: tuple[Any, ...], + workspaces: tuple[Any, ...], + stream: Any, + ) -> None: + if workspaces: + raise RuntimeError(f"GemmSwigluSm100 received unexpected workspaces: {len(workspaces)}") + a, b = inputs + ab12, c = outputs + + import cutlass + from cutlass.jax import jax_to_cutlass_dtype + + from .dense_gemm_persistent_swiglu import PersistentDenseGemmKernel + + kernel = PersistentDenseGemmKernel( + acc_dtype=jax_to_cutlass_dtype(self.acc_dtype), + use_2cta_instrs=self._op.mma_tiler_mn[0] == 256, + mma_tiler_mn=self._op.mma_tiler_mn, + cluster_shape_mn=self._op.cluster_shape_mn, + ) + max_active_clusters = cutlass.utils.HardwareInfo().get_max_active_clusters(self._op.cluster_shape_mn[0] * self._op.cluster_shape_mn[1]) + max_active_clusters -= self.num_cluster_overlap_margin + if max_active_clusters <= 0: + raise ValueError("max_active_clusters must be positive after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") + kernel( + a, + b, + ab12, + c, + cutlass.Float32(self._op.alpha), + max_active_clusters, + stream, + ) + + +@partial( + jax.jit, + static_argnames=( + "alpha", + "c_layout", + "ab12_dtype", + "c_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "a_layout", + "b_layout", + "target_compute_capability", + ), +) +def gemm_swiglu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + alpha: float = 1.0, + c_layout: str = "LMN", + ab12_dtype: Any | None = None, + c_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: tuple[int, int] | None = None, + *, + a_layout: str = "LMK", + b_layout: str = "LNK", + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute the standard dense batched GEMM and fused SwiGLU projection.""" + + return GemmSwigluSm100( + jax.ShapeDtypeStruct(a_tensor.shape, a_tensor.dtype), + jax.ShapeDtypeStruct(b_tensor.shape, b_tensor.dtype), + alpha=alpha, + c_layout=c_layout, + ab12_dtype=ab12_dtype, + c_dtype=c_dtype, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + a_layout=a_layout, + b_layout=b_layout, + target_compute_capability=target_compute_capability, + )(a_tensor, b_tensor) + + +__all__ = [ + "GemmSwigluSm100", + "SUPPORTED_COMPUTE_CAPABILITIES", + "gemm_swiglu_wrapper_sm100", +] diff --git a/python/cudnn/gemm_swiglu/op.py b/python/cudnn/gemm_swiglu/op.py new file mode 100644 index 000000000..787ef4624 --- /dev/null +++ b/python/cudnn/gemm_swiglu/op.py @@ -0,0 +1,165 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral dense GEMM + SwiGLU operation.""" + +from __future__ import annotations + +from typing import Any, Optional + +from .. import data_type +from .._dense_gemm import ( + require_16_byte_alignment, + require_cluster_shape, + require_compact_major, + require_gemm_inputs, + require_mma_tiler, + require_tensor_shape, +) +from .._op import Op +from .._tensor_desc import TensorDesc + +_STANDARD_AB_DTYPES = frozenset( + { + data_type.HALF, + data_type.BFLOAT16, + data_type.FLOAT, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + } +) +_STANDARD_OUTPUT_DTYPES = frozenset( + { + data_type.HALF, + data_type.BFLOAT16, + } +) +_STANDARD_AB12_DTYPES = _STANDARD_OUTPUT_DTYPES | { + data_type.FLOAT, + data_type.FP8_E4M3, + data_type.FP8_E5M2, +} +_FP8_DTYPES = frozenset({data_type.FP8_E4M3, data_type.FP8_E5M2}) +_SWIGLU_BLOCK_COLUMNS = 32 +_SWIGLU_COLUMNS_PER_PAIR = 2 * _SWIGLU_BLOCK_COLUMNS +_SWIGLU_MMA_N = tuple(range(_SWIGLU_COLUMNS_PER_PAIR, 257, _SWIGLU_COLUMNS_PER_PAIR)) + + +class GemmSwigluSm100Op(Op): + """Logical signature and launch configuration for standard GEMM + SwiGLU.""" + + def __init__( + self, + *, + a: TensorDesc[Any], + b: TensorDesc[Any], + ab12: TensorDesc[Any], + c: TensorDesc[Any], + alpha: float = 1.0, + acc_dtype: data_type = data_type.FLOAT, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: Optional[tuple[int, int]] = None, + ) -> None: + for name, desc in (("a", a), ("b", b), ("ab12", ab12), ("c", c)): + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + if not isinstance(acc_dtype, data_type): + raise TypeError(f"acc_dtype must be a cudnn.data_type, got {type(acc_dtype).__name__}") + + self.a = a + self.b = b + self.ab12 = ab12 + self.c = c + + self.alpha = alpha + self.acc_dtype = acc_dtype + self.mma_tiler_mn = tuple(mma_tiler_mn) + if cluster_shape_mn is None: + use_2cta_default = len(self.mma_tiler_mn) == 2 and self.mma_tiler_mn[0] == 256 + self.cluster_shape_mn = (2, 2) if use_2cta_default else (1, 1) + else: + self.cluster_shape_mn = tuple(cluster_shape_mn) + + self.m: Optional[int] = None + self.n: Optional[int] = None + self.k: Optional[int] = None + self.l: Optional[int] = None + self.output_n: Optional[int] = None + self.a_major: Optional[str] = None + self.b_major: Optional[str] = None + self.output_major: Optional[str] = None + + def check_support(self) -> bool: + """Validate the logical signature and resolve its canonical modes.""" + + self.m = self.n = self.k = self.l = self.output_n = None + self.a_major = self.b_major = self.output_major = None + + m, n, k, l = require_gemm_inputs(self.a, self.b) + if n % 2 != 0: + raise ValueError(f"N must be even for SwiGLU input/gate pairs, got {n}") + if n % _SWIGLU_COLUMNS_PER_PAIR != 0: + raise ValueError( + f"N must be divisible by {_SWIGLU_COLUMNS_PER_PAIR} because SwiGLU pairs consecutive {_SWIGLU_BLOCK_COLUMNS}-column blocks, got {n}" + ) + output_n = n // 2 + + require_tensor_shape(self.ab12, (m, n, l), label="AB12") + require_tensor_shape(self.c, (m, output_n, l), label="C") + + a_major = require_compact_major(self.a, "m", "k") + b_major = require_compact_major(self.b, "n", "k") + ab12_major = require_compact_major(self.ab12, "m", "n") + c_major = require_compact_major(self.c, "m", "n") + if ab12_major != c_major: + raise ValueError(f"AB12 and C must use the same major mode, got {ab12_major}-major and {c_major}-major") + + self.m, self.n, self.k, self.l = m, n, k, l + self.output_n = output_n + self.a_major, self.b_major, self.output_major = a_major, b_major, ab12_major + + self._check_standard_dtypes() + self.mma_tiler_mn = require_mma_tiler( + self.mma_tiler_mn, + allowed_n=_SWIGLU_MMA_N, + ) + cta_group_size = 2 if self.mma_tiler_mn[0] == 256 else 1 + if cta_group_size == 2 and m % self.mma_tiler_mn[0] != 0: + raise ValueError(f"M must be divisible by {self.mma_tiler_mn[0]} for 2-CTA MMA, got {m}") + self.cluster_shape_mn = require_cluster_shape( + self.cluster_shape_mn, + cta_group_size=cta_group_size, + ) + if cta_group_size == 1 and self.cluster_shape_mn != (1, 1): + raise ValueError("cluster_shape_mn must be (1, 1) with a 128-wide M tile") + + for tensor in (self.a, self.b, self.ab12, self.c): + require_16_byte_alignment(tensor) + return True + + def _check_standard_dtypes(self) -> None: + ab_dtype = self.a.cudnn_dtype + if ab_dtype not in _STANDARD_AB_DTYPES: + raise ValueError(f"A dtype must be one of the supported dense GEMM dtypes, got {self.a.dtype}") + if self.b.cudnn_dtype != ab_dtype: + raise ValueError(f"A and B must have the same dtype, got {self.a.dtype} and {self.b.dtype}") + + ab12_dtype = self.ab12.cudnn_dtype + if self.acc_dtype == data_type.FLOAT: + if ab12_dtype not in _STANDARD_AB12_DTYPES: + raise ValueError(f"AB12 has unsupported dtype {self.ab12.dtype} for float32 accumulation") + if ab12_dtype in _FP8_DTYPES: + raise NotImplementedError("FP8 AB12 output is currently disabled") + elif self.acc_dtype == data_type.HALF: + if ab12_dtype not in _STANDARD_OUTPUT_DTYPES: + raise ValueError(f"AB12 has unsupported dtype {self.ab12.dtype} for float16 accumulation") + if ab_dtype not in _FP8_DTYPES | {data_type.HALF}: + raise ValueError(f"A and B dtype {self.a.dtype} is unsupported for float16 accumulation") + else: + raise ValueError(f"Accumulator dtype must be float32 or float16, got {self.acc_dtype}") + + if self.c.cudnn_dtype not in _STANDARD_OUTPUT_DTYPES: + raise ValueError(f"C dtype must be float16 or bfloat16, got {self.c.dtype}") + + +__all__ = ["GemmSwigluSm100Op"] diff --git a/python/cudnn/jax/__init__.py b/python/cudnn/jax/__init__.py index 67098c72d..92ac00453 100644 --- a/python/cudnn/jax/__init__.py +++ b/python/cudnn/jax/__init__.py @@ -37,6 +37,8 @@ _OPERATION_EXPORTS = { "RmsNormRhtAmaxSm100": ("..rmsnorm_rht_amax.jax", "RmsNormRhtAmaxSm100"), "rmsnorm_rht_amax_sm100": ("..rmsnorm_rht_amax.jax", "rmsnorm_rht_amax_sm100"), + "GemmSwigluSm100": ("..gemm_swiglu.jax", "GemmSwigluSm100"), + "gemm_swiglu_wrapper_sm100": ("..gemm_swiglu.jax", "gemm_swiglu_wrapper_sm100"), "SparseAttentionBackward": ("..deepseek_sparse_attention.sparse_attention_backward.jax", "SparseAttentionBackward"), "sparse_attention_backward_wrapper": ( "..deepseek_sparse_attention.sparse_attention_backward.jax", @@ -74,7 +76,7 @@ "dense_indexer_backward_wrapper": ("..deepseek_sparse_attention.indexer_backward.jax", "dense_indexer_backward_wrapper"), } -_DSA_EXPORTS = frozenset(name for name in _OPERATION_EXPORTS if name not in {"RmsNormRhtAmaxSm100", "rmsnorm_rht_amax_sm100"}) +_DSA_EXPORTS = frozenset(name for name, (module_name, _) in _OPERATION_EXPORTS.items() if module_name.startswith("..deepseek_sparse_attention.")) def _load_operation(name: str) -> Any: diff --git a/test/python/fe_api/test_dense_gemm_helpers.py b/test/python/fe_api/test_dense_gemm_helpers.py new file mode 100644 index 000000000..2db698500 --- /dev/null +++ b/test/python/fe_api/test_dense_gemm_helpers.py @@ -0,0 +1,198 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for framework-neutral dense GEMM helpers.""" + +import ast +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_PACKAGE = "cudnn_dense_gemm_test" + + +class _DataType(Enum): + NOT_SET = auto() + FLOAT = auto() + HALF = auto() + BFLOAT16 = auto() + FP8_E4M3 = auto() + FP8_E5M2 = auto() + FP8_E8M0 = auto() + FP4_E2M1 = auto() + + +class DenseGemmHelperTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + package = types.ModuleType(_PACKAGE) + package.__path__ = [str(_CUDNN_ROOT)] + package.__package__ = _PACKAGE + package.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + package.data_type = _DataType + sys.modules[_PACKAGE] = package + + cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + cls.helpers = importlib.import_module(f"{_PACKAGE}._dense_gemm") + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def desc(self, shape, dtype=_DataType.BFLOAT16, *, stride_order=None, name=""): + return self.tensor_module.make_compact_tensor_desc( + dtype=dtype, + shape=tuple(shape), + stride_order=stride_order, + name=name, + ) + + def test_gemm_inputs_resolve_mnkl(self): + a = self.desc((17, 31, 3), name="A") + b = self.desc((19, 31, 3), name="B") + + self.assertEqual(self.helpers.require_gemm_inputs(a, b), (17, 19, 31, 3)) + + def test_gemm_inputs_reject_invalid_rank_extents_and_matching(self): + cases = ( + (self.desc((17, 31)), self.desc((19, 31, 3)), "A must have rank 3"), + (self.desc((17, 31, 3)), self.desc((19, 31)), "B must have rank 3"), + (self.desc((0, 31, 3)), self.desc((19, 31, 3)), "M must be positive"), + (self.desc((17, 31, 3)), self.desc((0, 31, 3)), "N must be positive"), + (self.desc((17, 0, 3)), self.desc((19, 0, 3)), "K must be positive"), + (self.desc((17, 31, 0)), self.desc((19, 31, 0)), "L must be positive"), + (self.desc((17, 31, 3)), self.desc((19, 29, 3)), "B shape mismatch"), + (self.desc((17, 31, 3)), self.desc((19, 31, 4)), "B shape mismatch"), + ) + + for a, b, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self.helpers.require_gemm_inputs(a, b) + + def test_exact_output_shape_uses_logical_descriptor_shape(self): + output = self.desc((17, 19, 3), name="AB12") + self.helpers.require_tensor_shape(output, (17, 19, 3)) + + with self.assertRaisesRegex(ValueError, "AB12 must have shape"): + self.helpers.require_tensor_shape(output, (17, 18, 3)) + + def test_compact_major_recognizes_mode_zero_and_mode_one(self): + m_major = self.desc((8, 16, 2), stride_order=(0, 1, 2), name="A") + k_major = self.desc((8, 16, 2), stride_order=(1, 0, 2), name="A") + + self.assertEqual(self.helpers.require_compact_major(m_major, "m", "k"), "m") + self.assertEqual(self.helpers.require_compact_major(k_major, "m", "k"), "k") + + def test_compact_major_rejects_batch_major_and_noncompact_layouts(self): + batch_major = self.desc((8, 16, 2), stride_order=(2, 1, 0), name="A") + noncompact = self.tensor_module.TensorDesc( + dtype=_DataType.BFLOAT16, + shape=(8, 16, 2), + stride=(1, 16, 256), + stride_order=(0, 1, 2), + name="A", + ) + + for tensor in (batch_major, noncompact): + with self.subTest(stride=tensor.stride): + with self.assertRaisesRegex(ValueError, "compact m-major or k-major"): + self.helpers.require_compact_major(tensor, "m", "k") + + def test_data_type_widths_and_16_byte_alignment(self): + cases = ( + (_DataType.FLOAT, 32, 4), + (_DataType.HALF, 16, 8), + (_DataType.BFLOAT16, 16, 8), + (_DataType.FP8_E4M3, 8, 16), + (_DataType.FP8_E5M2, 8, 16), + (_DataType.FP4_E2M1, 4, 32), + ) + + for dtype, bits, extent in cases: + with self.subTest(dtype=dtype): + tensor = self.desc((extent, 7, 2), dtype, stride_order=(0, 1, 2), name="A") + self.assertEqual(self.helpers.data_type_bits(dtype), bits) + self.helpers.require_16_byte_alignment(tensor) + + misaligned = self.desc((extent - 1, 7, 2), dtype, stride_order=(0, 1, 2), name="A") + with self.assertRaisesRegex(ValueError, "16-byte alignment"): + self.helpers.require_16_byte_alignment(misaligned) + + with self.assertRaisesRegex(ValueError, "Unsupported cuDNN data type"): + self.helpers.data_type_bits(_DataType.NOT_SET) + + def test_16_byte_alignment_uses_the_selected_major_extent(self): + tensor = self.desc((7, 8, 2), stride_order=(1, 0, 2), name="A") + self.helpers.require_16_byte_alignment(tensor) + + def test_mma_tiler_validation_is_configurable(self): + self.assertEqual(self.helpers.require_mma_tiler((128, 32)), (128, 32)) + self.assertEqual(self.helpers.require_mma_tiler((256, 256)), (256, 256)) + self.assertEqual( + self.helpers.require_mma_tiler((64, 96), allowed_m=(64,), allowed_n=(96,)), + (64, 96), + ) + + cases = ( + ((64, 128), "mma_tiler_mn\\[0\\]"), + ((128, 48), "mma_tiler_mn\\[1\\]"), + ((128,), "must contain two integers"), + ) + for tiler, message in cases: + with self.subTest(tiler=tiler): + with self.assertRaisesRegex(ValueError, message): + self.helpers.require_mma_tiler(tiler) + + with self.assertRaisesRegex(TypeError, "must contain integers"): + self.helpers.require_mma_tiler((True, 128)) + + def test_cluster_shape_validation_handles_cta_groups(self): + self.assertEqual(self.helpers.require_cluster_shape((1, 1)), (1, 1)) + self.assertEqual( + self.helpers.require_cluster_shape((2, 4), cta_group_size=2), + (2, 4), + ) + + cases = ( + ((3, 1), {}, "positive powers of two"), + ((4, 8), {}, "must not exceed 16"), + ((1, 2), {"cta_group_size": 2}, "divisible by CTA group size"), + ) + for shape, kwargs, message in cases: + with self.subTest(shape=shape): + with self.assertRaisesRegex(ValueError, message): + self.helpers.require_cluster_shape(shape, **kwargs) + + def test_module_has_no_framework_or_kernel_dependencies(self): + tree = ast.parse((_CUDNN_ROOT / "_dense_gemm.py").read_text()) + imports = [] + for node in tree.body: + if isinstance(node, ast.Import): + imports.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + imports.append("." * node.level + (node.module or "")) + + self.assertFalse( + any(imported == dependency or imported.startswith(f"{dependency}.") for imported in imports for dependency in ("torch", "jax", "cutlass", "cuda")), + imports, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_dsa_compiler.py b/test/python/fe_api/test_dsa_compiler.py index 11ac77981..54ad206f7 100644 --- a/test/python/fe_api/test_dsa_compiler.py +++ b/test/python/fe_api/test_dsa_compiler.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for DSA compilation targets.""" +"""Dependency-free contracts for JAX CuTe compilation targets.""" from importlib.util import module_from_spec, spec_from_file_location from pathlib import Path @@ -16,10 +16,10 @@ pytestmark = pytest.mark.L0 -_MODULE_PATH = Path(__file__).resolve().parents[3] / "python" / "cudnn" / "deepseek_sparse_attention" / "utils" / "compiler.py" +_MODULE_PATH = Path(__file__).resolve().parents[3] / "python" / "cudnn" / "_cute_compiler.py" -class DsaCompilerTest(unittest.TestCase): +class JaxCompilerTest(unittest.TestCase): @classmethod def setUpClass(cls) -> None: spec = spec_from_file_location("dsa_compiler_contract", _MODULE_PATH) diff --git a/test/python/fe_api/test_gemm_swiglu_imports.py b/test/python/fe_api/test_gemm_swiglu_imports.py new file mode 100644 index 000000000..b40c9c162 --- /dev/null +++ b/test/python/fe_api/test_gemm_swiglu_imports.py @@ -0,0 +1,190 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Import and Torch-adapter contracts for dense GEMM + SwiGLU.""" + +import ast +import importlib.util +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "gemm_swiglu" + + +def _load_operation_package(name: str): + parent_name, separator, _ = name.rpartition(".") + if not separator: + raise ValueError("Synthetic operation package must have a parent") + parent = types.ModuleType(parent_name) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = parent_name + sys.modules[parent_name] = parent + + spec = importlib.util.spec_from_file_location( + name, + _OPERATION_ROOT / "__init__.py", + submodule_search_locations=[str(_OPERATION_ROOT)], + ) + if spec is None or spec.loader is None: + raise RuntimeError("Unable to load GEMM + SwiGLU package") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _remove_package(name: str) -> None: + root_name = name.split(".", 1)[0] + for module_name in tuple(sys.modules): + if module_name == root_name or module_name.startswith(f"{root_name}."): + sys.modules.pop(module_name, None) + + +class GemmSwigluImportContractTest(unittest.TestCase): + def test_package_import_loads_no_framework_adapter(self): + name = "cudnn_gemm_swiglu_import_test.gemm_swiglu" + blocked = {module_name: None for module_name in ("torch", "jax", "cutlass", "cuda")} + try: + with mock.patch.dict(sys.modules, blocked): + package = _load_operation_package(name) + self.assertNotIn(f"{name}.api", sys.modules) + self.assertNotIn(f"{name}.op", sys.modules) + self.assertTrue( + { + "GemmSwigluSm100Op", + "GemmSwigluSm100", + "gemm_swiglu_wrapper_sm100", + "api", + "op", + }.issubset(dir(package)) + ) + finally: + _remove_package(name) + + def test_exports_route_without_loading_the_other_layer(self): + name = "cudnn_gemm_swiglu_route_test.gemm_swiglu" + api_sentinel = object() + op_sentinel = object() + api = types.ModuleType(f"{name}.api") + api.GemmSwigluSm100 = api_sentinel + api.gemm_swiglu_wrapper_sm100 = object() + op = types.ModuleType(f"{name}.op") + op.GemmSwigluSm100Op = op_sentinel + try: + sys.modules[api.__name__] = api + sys.modules[op.__name__] = op + package = _load_operation_package(name) + self.assertIs(package.GemmSwigluSm100Op, op_sentinel) + self.assertIs(package.GemmSwigluSm100, api_sentinel) + finally: + _remove_package(name) + + def test_static_adapter_boundaries_and_torch_constructor(self): + def imports_framework(node, framework: str): + if isinstance(node, ast.Import): + return any(alias.name == framework or alias.name.startswith(f"{framework}.") for alias in node.names) + return node.level == 0 and (node.module == framework or (node.module or "").startswith(f"{framework}.")) + + op_tree = ast.parse((_OPERATION_ROOT / "op.py").read_text(), filename="op.py") + api_tree = ast.parse((_OPERATION_ROOT / "api.py").read_text(), filename="api.py") + op_imports = [node for node in ast.walk(op_tree) if isinstance(node, (ast.Import, ast.ImportFrom))] + api_imports = [node for node in ast.walk(api_tree) if isinstance(node, (ast.Import, ast.ImportFrom))] + self.assertFalse(any(imports_framework(node, framework) for node in op_imports for framework in ("torch", "jax", "cutlass", "cuda"))) + self.assertFalse(any(imports_framework(node, "jax") for node in api_imports)) + + adapter = next(node for node in api_tree.body if isinstance(node, ast.ClassDef) and node.name == "GemmSwigluSm100") + self.assertEqual( + [base.id for base in adapter.bases if isinstance(base, ast.Name)], + ["APIBase"], + ) + constructor = next(node for node in adapter.body if isinstance(node, ast.FunctionDef) and node.name == "__init__") + self.assertEqual( + [argument.arg for argument in constructor.args.args], + [ + "self", + "sample_a", + "sample_b", + "sample_ab12", + "sample_c", + "alpha", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sample_sfa", + "sample_sfb", + "sample_amax", + "sample_sfc", + "sample_norm_const", + "sf_vec_size", + "vector_f32", + "ab12_stages", + ], + ) + required_count = len(constructor.args.args) - len(constructor.args.defaults) + self.assertEqual( + [argument.arg for argument in constructor.args.args[:required_count]], + ["self", "sample_a", "sample_b", "sample_ab12", "sample_c"], + ) + op_constructors = [ + node for node in ast.walk(constructor) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "GemmSwigluSm100Op" + ] + self.assertEqual(len(op_constructors), 1) + self.assertTrue({"a", "b", "ab12", "c"}.issubset({keyword.arg for keyword in op_constructors[0].keywords})) + + execute = next(node for node in adapter.body if isinstance(node, ast.FunctionDef) and node.name == "execute") + self.assertEqual( + [argument.arg for argument in execute.args.args], + [ + "self", + "a_tensor", + "b_tensor", + "ab12_tensor", + "c_tensor", + "sfa_tensor", + "sfb_tensor", + "amax_tensor", + "sfc_tensor", + "norm_const_tensor", + "alpha", + "current_stream", + ], + ) + + wrapper = next(node for node in api_tree.body if isinstance(node, ast.FunctionDef) and node.name == "gemm_swiglu_wrapper_sm100") + self.assertEqual( + [argument.arg for argument in wrapper.args.args], + [ + "a_tensor", + "b_tensor", + "alpha", + "c_major", + "ab12_dtype", + "c_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sfa_tensor", + "sfb_tensor", + "norm_const_tensor", + "sf_vec_size", + "vector_f32", + "ab12_stages", + "stream", + ], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_gemm_swiglu_op.py b/test/python/fe_api/test_gemm_swiglu_op.py new file mode 100644 index 000000000..a549251aa --- /dev/null +++ b/test/python/fe_api/test_gemm_swiglu_op.py @@ -0,0 +1,241 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for the dense GEMM + SwiGLU operation.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "gemm_swiglu" +_PACKAGE = "cudnn_gemm_swiglu_op_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + FP8_E4M3 = auto() + FP8_E5M2 = auto() + FP8_E8M0 = auto() + FP4_E2M1 = auto() + + +class GemmSwigluOpContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + operation_name = f"{_PACKAGE}.gemm_swiglu" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + try: + cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}._op") + cls.op_module = importlib.import_module(f"{operation_name}.op") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def _desc(self, shape, dtype, order, name): + shape = tuple(shape) + stride = [0] * len(shape) + running = 1 + for dimension in order: + stride[dimension] = running + running *= max(shape[dimension], 1) + return self.tensor_module.TensorDesc( + dtype=dtype, + shape=shape, + stride=tuple(stride), + stride_order=tuple(order), + name=name, + ) + + def _op(self, **overrides): + m, n, k, batch = 128, 128, 128, 2 + arguments = { + "a": self._desc((m, k, batch), _DataType.BFLOAT16, (1, 0, 2), "A"), + "b": self._desc((n, k, batch), _DataType.BFLOAT16, (1, 0, 2), "B"), + "ab12": self._desc((m, n, batch), _DataType.FLOAT, (1, 0, 2), "AB12"), + "c": self._desc((m, n // 2, batch), _DataType.HALF, (1, 0, 2), "C"), + } + arguments.update(overrides) + return self.op_module.GemmSwigluSm100Op(**arguments) + + def test_validates_complete_signature_and_resolves_configuration(self): + operation = self._op(alpha=0.5) + + self.assertIsInstance(operation, self.base_module.Op) + self.assertTrue(operation.check_support()) + self.assertEqual((operation.m, operation.n, operation.k, operation.l), (128, 128, 128, 2)) + self.assertEqual(operation.output_n, 64) + self.assertEqual( + (operation.a_major, operation.b_major, operation.output_major), + ("k", "k", "n"), + ) + self.assertEqual(operation.mma_tiler_mn, (128, 128)) + self.assertEqual(operation.cluster_shape_mn, (1, 1)) + self.assertEqual(operation.alpha, 0.5) + self.assertNotIn(f"{_PACKAGE}.gemm_swiglu.api", sys.modules) + self.assertNotIn(f"{_PACKAGE}.gemm_swiglu.dense_gemm_persistent_swiglu", sys.modules) + + def test_accepts_each_supported_compact_major_mode(self): + m, n, k, batch = 128, 128, 128, 2 + for a_order, b_order, output_order, expected in ( + ((0, 1, 2), (0, 1, 2), (0, 1, 2), ("m", "n", "m")), + ((1, 0, 2), (1, 0, 2), (1, 0, 2), ("k", "k", "n")), + ): + with self.subTest(expected=expected): + operation = self._op( + a=self._desc((m, k, batch), _DataType.BFLOAT16, a_order, "A"), + b=self._desc((n, k, batch), _DataType.BFLOAT16, b_order, "B"), + ab12=self._desc((m, n, batch), _DataType.FLOAT, output_order, "AB12"), + c=self._desc((m, n // 2, batch), _DataType.HALF, output_order, "C"), + ) + self.assertTrue(operation.check_support()) + self.assertEqual( + (operation.a_major, operation.b_major, operation.output_major), + expected, + ) + + def test_rejects_invalid_shapes_and_layouts(self): + cases = ( + ( + {"a": self._desc((128, 128), _DataType.BFLOAT16, (1, 0), "A")}, + "A must have rank 3", + ), + ( + {"b": self._desc((128, 64, 2), _DataType.BFLOAT16, (1, 0, 2), "B")}, + "B shape mismatch", + ), + ( + { + "b": self._desc((127, 128, 2), _DataType.BFLOAT16, (1, 0, 2), "B"), + "ab12": self._desc((128, 127, 2), _DataType.FLOAT, (1, 0, 2), "AB12"), + "c": self._desc((128, 63, 2), _DataType.HALF, (1, 0, 2), "C"), + }, + "N must be even", + ), + ( + { + "b": self._desc((96, 128, 2), _DataType.BFLOAT16, (1, 0, 2), "B"), + "ab12": self._desc((128, 96, 2), _DataType.FLOAT, (1, 0, 2), "AB12"), + "c": self._desc((128, 48, 2), _DataType.HALF, (1, 0, 2), "C"), + }, + "N must be divisible by 64", + ), + ( + {"c": self._desc((128, 128, 2), _DataType.HALF, (1, 0, 2), "C")}, + "C must have shape", + ), + ( + {"c": self._desc((128, 64, 2), _DataType.HALF, (0, 1, 2), "C")}, + "AB12 and C must use the same major mode", + ), + ) + for overrides, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self._op(**overrides).check_support() + + def test_rejects_invalid_standard_dtypes(self): + cases = ( + ( + {"b": self._desc((128, 128, 2), _DataType.HALF, (1, 0, 2), "B")}, + ValueError, + "A and B must have the same dtype", + ), + ( + {"ab12": self._desc((128, 128, 2), _DataType.FP8_E4M3, (1, 0, 2), "AB12")}, + NotImplementedError, + "FP8 AB12 output is currently disabled", + ), + ( + { + "acc_dtype": _DataType.HALF, + "ab12": self._desc((128, 128, 2), _DataType.HALF, (1, 0, 2), "AB12"), + }, + ValueError, + "unsupported for float16 accumulation", + ), + ( + {"c": self._desc((128, 64, 2), _DataType.FLOAT, (1, 0, 2), "C")}, + ValueError, + "C dtype must be float16 or bfloat16", + ), + ) + for overrides, error_type, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(error_type, message): + self._op(**overrides).check_support() + + def test_validates_alignment_tiler_and_cluster(self): + with self.assertRaisesRegex(ValueError, "A contiguous extent must be a multiple"): + self._op( + a=self._desc((128, 124, 2), _DataType.BFLOAT16, (1, 0, 2), "A"), + b=self._desc((128, 124, 2), _DataType.BFLOAT16, (1, 0, 2), "B"), + ).check_support() + + with self.assertRaisesRegex(ValueError, r"mma_tiler_mn\[1\]"): + self._op(mma_tiler_mn=(128, 32)).check_support() + + with self.assertRaisesRegex(ValueError, "must contain two integers"): + self._op(mma_tiler_mn=()).check_support() + + with self.assertRaisesRegex(ValueError, r"must be \(1, 1\)"): + self._op(cluster_shape_mn=(2, 1)).check_support() + + with self.assertRaisesRegex(ValueError, "M must be divisible by 256"): + self._op(mma_tiler_mn=(256, 128)).check_support() + + operation = self._op( + a=self._desc((256, 128, 2), _DataType.BFLOAT16, (1, 0, 2), "A"), + ab12=self._desc((256, 128, 2), _DataType.FLOAT, (1, 0, 2), "AB12"), + c=self._desc((256, 64, 2), _DataType.HALF, (1, 0, 2), "C"), + mma_tiler_mn=(256, 128), + ) + self.assertTrue(operation.check_support()) + self.assertEqual(operation.cluster_shape_mn, (2, 2)) + + def test_constructor_requires_canonical_descriptors_and_dtype(self): + for name in ("a", "b", "ab12", "c"): + with self.subTest(name=name): + with self.assertRaisesRegex(TypeError, f"{name} must be a TensorDesc"): + self._op(**{name: types.SimpleNamespace(shape=(1,))}) + + with self.assertRaisesRegex(TypeError, "acc_dtype must be a cudnn.data_type"): + self._op(acc_dtype="float32") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_api_base.py b/test/python/fe_api/test_jax_api_base.py index a67b2f636..b9e184a51 100644 --- a/test/python/fe_api/test_jax_api_base.py +++ b/test/python/fe_api/test_jax_api_base.py @@ -278,6 +278,28 @@ def test_converts_public_array_axes_to_canonical_descriptor_axes(self): self.assertEqual(desc.stride, (1, 12, 4)) self.assertEqual(desc.stride_order, (0, 2, 1)) + def test_converts_public_physical_layout_to_canonical_descriptor_axes(self): + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3, 4), "bfloat16"), + "sample", + mode=(1, 2, 0), + public_stride_order=(0, 2, 1), + ) + + self.assertEqual(desc.shape, (3, 4, 2)) + self.assertEqual(desc.stride, (8, 2, 1)) + self.assertEqual(desc.stride_order, (2, 1, 0)) + + def test_rejects_invalid_public_stride_orders(self): + value = _ArrayMetadata((2, 3, 4), "bfloat16") + + with self.assertRaisesRegex(ValueError, "public_stride_order rank mismatch"): + self.module.JaxApiBase._to_tensor_desc(value, "sample", public_stride_order=(1, 0)) + with self.assertRaisesRegex(ValueError, "public_stride_order must be a permutation"): + self.module.JaxApiBase._to_tensor_desc(value, "sample", public_stride_order=(2, 2, 0)) + with self.assertRaisesRegex(TypeError, "public_stride_order entries must be integers"): + self.module.JaxApiBase._to_tensor_desc(value, "sample", public_stride_order=(2, True, 0)) + def test_jax_descriptor_derives_compact_output_metadata(self): source = self.module.JaxApiBase._to_tensor_desc( _ArrayMetadata((2, 3), "bfloat16"), @@ -354,6 +376,26 @@ def test_builds_public_tensor_spec_from_canonical_metadata(self): self.assertEqual(spec.mode, mode) self.assertEqual(spec.divisibility, (2, 4, 8)) + def test_tensor_spec_preserves_explicit_public_physical_layout(self): + mode = (1, 2, 0) + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3, 4), "bfloat16"), + "sample", + mode=mode, + public_stride_order=(0, 2, 1), + ) + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass_jax = types.ModuleType("cutlass.jax") + cutlass_jax.TensorSpec = _TensorSpec + cutlass.jax = cutlass_jax + + with mock.patch.dict(sys.modules, {"cutlass": cutlass, "cutlass.jax": cutlass_jax}): + spec = self.module.JaxApiBase._to_tensor_spec(desc, mode=mode) + + self.assertEqual(spec.layout, (0, 2, 1)) + self.assertEqual(spec.mode, mode) + def test_tensor_spec_defaults_to_identity_axis_binding(self): desc = self.module.JaxApiBase._to_tensor_desc( _ArrayMetadata((2, 3), "bfloat16"), diff --git a/test/python/fe_api/test_jax_gemm_layout.py b/test/python/fe_api/test_jax_gemm_layout.py new file mode 100644 index 000000000..e6c0f2972 --- /dev/null +++ b/test/python/fe_api/test_jax_gemm_layout.py @@ -0,0 +1,127 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for dense JAX GEMM axis bindings.""" + +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_PACKAGE = "cudnn_jax_gemm_layout_test" + + +class JaxGemmLayoutTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + sys.modules[_PACKAGE] = root + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + try: + with mock.patch.dict( + sys.modules, + { + "jax": None, + "jax.numpy": None, + "cutlass": None, + "cutlass.jax": None, + "torch": None, + }, + ): + cls.module = importlib.import_module(f"{internal_name}.gemm") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls): + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def test_maps_supported_public_layouts_to_canonical_modes(self): + cases = ( + (self.module.gemm_a_mode, "LMK", (1, 2, 0)), + (self.module.gemm_a_mode, "LKM", (2, 1, 0)), + (self.module.gemm_b_mode, "LNK", (1, 2, 0)), + (self.module.gemm_b_mode, "LKN", (2, 1, 0)), + (self.module.gemm_output_mode, "LMN", (1, 2, 0)), + (self.module.gemm_output_mode, "LNM", (2, 1, 0)), + ) + + for helper, layout, expected in cases: + with self.subTest(layout=layout): + self.assertEqual(helper(layout), expected) + + def test_rejects_unsupported_axis_orders(self): + with self.assertRaisesRegex(ValueError, r"a_layout must be one of \('LMK', 'LKM'\)"): + self.module.gemm_a_mode("MKL") + with self.assertRaisesRegex(ValueError, r"b_layout must be one of \('LNK', 'LKN'\)"): + self.module.gemm_b_mode("NKL") + with self.assertRaisesRegex(ValueError, r"c_layout must be one of \('LMN', 'LNM'\)"): + self.module.gemm_output_mode("MNL") + + with self.assertRaisesRegex(ValueError, r"d_layout must be one of \('LMN', 'LNM'\)"): + self.module.gemm_output_mode("MNL", name="d_layout") + + def test_layout_validation_is_case_sensitive(self): + with self.assertRaisesRegex(ValueError, "got 'lmk'"): + self.module.gemm_a_mode("lmk") + + def test_layout_validation_requires_a_string(self): + with self.assertRaisesRegex(TypeError, "a_layout must be a string"): + self.module.gemm_a_mode(("L", "M", "K")) + + def test_fixed_auxiliary_stride_orders_match_kernel_abis(self): + self.assertEqual(self.module.ROW_MAJOR_STRIDE_ORDER_3D, (2, 1, 0)) + self.assertEqual(self.module.BLOCK_SCALE_STRIDE_ORDER, (3, 1, 0, 4, 2, 5)) + self.assertEqual(self.module.PROBABILITY_STRIDE_ORDER, (0, 1, 2)) + + layout = importlib.import_module(f"{self.module.__package__}.layout") + block_scale_shape = (32, 4, 2, 4, 3, 5) + block_scale_stride = layout.compact_stride(block_scale_shape, self.module.BLOCK_SCALE_STRIDE_ORDER) + self.assertEqual( + layout.to_cutlass_layout( + block_scale_shape, + block_scale_stride, + self.module.BLOCK_SCALE_STRIDE_ORDER, + ), + (2, 1, 4, 0, 3, 5), + ) + + probability_shape = (7, 1, 5) + probability_stride = layout.compact_stride(probability_shape, self.module.PROBABILITY_STRIDE_ORDER) + self.assertEqual( + layout.to_cutlass_layout( + probability_shape, + probability_stride, + self.module.PROBABILITY_STRIDE_ORDER, + ), + (0, 1, 2), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_gemm_swiglu.py b/test/python/fe_api/test_jax_gemm_swiglu.py new file mode 100644 index 000000000..77c5ea6f4 --- /dev/null +++ b/test/python/fe_api/test_jax_gemm_swiglu.py @@ -0,0 +1,201 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX integration tests for the standard dense GEMM + SwiGLU adapter.""" + +import pytest + + +def _jax_runtime(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + cutlass_jax = pytest.importorskip("cutlass.jax") + if not cutlass_jax.is_available(): + pytest.skip("Installed JAX version is unsupported by CUTLASS JAX") + return jax, jnp + + +def _compute_capability(device): + reported = getattr(device, "compute_capability", None) + if isinstance(reported, (tuple, list)): + major, minor = int(reported[0]), int(reported[1]) + else: + text = str(reported) + if "." in text: + major_text, minor_text = text.split(".", 1) + major, minor = int(major_text), int(minor_text) + else: + major, minor = divmod(int(text), 10) + return major * 10 + minor + + +def _supported_gpu(jax): + devices = tuple(device for device in jax.local_devices() if device.platform == "gpu") + try: + capabilities = {_compute_capability(device) for device in devices} + except (TypeError, ValueError): + pytest.skip("JAX GPU compute capability is unavailable") + if not devices or len(capabilities) != 1: + pytest.skip("A homogeneous JAX SM100-family GPU configuration is required") + capability = capabilities.pop() + if capability not in {100, 103, 107}: + pytest.skip("A supported JAX SM100-family GPU is not available") + return devices[0], capability + + +@pytest.fixture +def device_compatibility_checks_disabled(): + _jax_runtime() + from cudnn.jax import disable_device_compatibility_checks + + disable_device_compatibility_checks(True) + try: + yield + finally: + disable_device_compatibility_checks(False) + + +@pytest.mark.L0 +def test_jax_gemm_swiglu_abstract_contract( + monkeypatch, + device_compatibility_checks_disabled, +): + jax, jnp = _jax_runtime() + + from cudnn._jax import JaxApiBase + from cudnn.jax import GemmSwigluSm100, gemm_swiglu_wrapper_sm100 + + def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): + return tuple( + jnp.empty( + self._materialize_tensor_desc(desc, mode=spec.mode).shape, + dtype=desc.dtype, + ) + for desc, spec in zip(output_descs, output_spec) + ) + + monkeypatch.setattr(JaxApiBase, "_call_kernel", abstract_call) + + sample_a = jax.ShapeDtypeStruct((3, 128, 64), jnp.bfloat16) + sample_b = jax.ShapeDtypeStruct((3, 192, 64), jnp.bfloat16) + api = GemmSwigluSm100( + sample_a, + sample_b, + target_compute_capability=100, + ) + result = jax.eval_shape(api, sample_a, sample_b) + assert tuple(result.keys()) == ( + "ab12_tensor", + "c_tensor", + "sfc_tensor", + "amax_tensor", + ) + assert result["ab12_tensor"].shape == (3, 128, 192) + assert result["ab12_tensor"].dtype == jnp.float32 + assert result["c_tensor"].shape == (3, 128, 96) + assert result["c_tensor"].dtype == jnp.float16 + assert result["sfc_tensor"] is None + assert result["amax_tensor"] is None + + alternate_a = jax.ShapeDtypeStruct((3, 64, 128), jnp.bfloat16) + alternate_b = jax.ShapeDtypeStruct((3, 64, 192), jnp.bfloat16) + alternate = jax.eval_shape( + lambda a, b: gemm_swiglu_wrapper_sm100( + a, + b, + a_layout="LKM", + b_layout="LKN", + c_layout="LNM", + c_dtype=jnp.bfloat16, + target_compute_capability=103, + ), + alternate_a, + alternate_b, + ) + assert alternate["ab12_tensor"].shape == (3, 192, 128) + assert alternate["c_tensor"].shape == (3, 96, 128) + assert alternate["c_tensor"].dtype == jnp.bfloat16 + + explicit = GemmSwigluSm100( + sample_a, + sample_b, + sample_ab12=jax.ShapeDtypeStruct((3, 128, 192), jnp.bfloat16), + sample_c=jax.ShapeDtypeStruct((3, 128, 96), jnp.bfloat16), + target_compute_capability=100, + ) + explicit_result = jax.eval_shape(explicit, sample_a, sample_b) + assert explicit_result["ab12_tensor"].dtype == jnp.bfloat16 + assert explicit_result["c_tensor"].dtype == jnp.bfloat16 + + +@pytest.mark.L0 +def test_jax_gemm_swiglu_sm100_jit_and_numerics(): + jax, jnp = _jax_runtime() + device, capability = _supported_gpu(jax) + + from cudnn.jax import gemm_swiglu_wrapper_sm100 + + batch, m, n, k = 1, 128, 128, 128 + alpha = 0.5 + a = jax.device_put( + jax.random.normal( + jax.random.key(0), + (batch, k, m), + dtype=jnp.bfloat16, + ) + * jnp.bfloat16(0.1), + device, + ) + b = jax.device_put( + jax.random.normal( + jax.random.key(1), + (batch, k, n), + dtype=jnp.bfloat16, + ) + * jnp.bfloat16(0.1), + device, + ) + + lowered = gemm_swiglu_wrapper_sm100.lower( + a, + b, + alpha=alpha, + a_layout="LKM", + b_layout="LKN", + c_layout="LNM", + target_compute_capability=capability, + ) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + compiled = lowered.compile() + + def reference(a_value): + ab12_lmn = alpha * jnp.einsum( + "lmk,lnk->lmn", + a_value.transpose(0, 2, 1).astype(jnp.float32), + b.transpose(0, 2, 1).astype(jnp.float32), + ) + blocks = ab12_lmn.reshape(batch, m, n // 32, 32) + input_blocks = blocks[:, :, 0::2].reshape(batch, m, n // 2) + gate_blocks = blocks[:, :, 1::2].reshape(batch, m, n // 2) + c_lmn = input_blocks * jax.nn.silu(gate_blocks) + return ab12_lmn.transpose(0, 2, 1), c_lmn.transpose(0, 2, 1) + + for a_value in (a, -a): + result = compiled(a_value, b) + result["c_tensor"].block_until_ready() + reference_ab12, reference_c = reference(a_value) + assert jnp.allclose( + result["ab12_tensor"], + reference_ab12, + atol=5e-2, + rtol=2e-2, + ) + assert jnp.allclose( + result["c_tensor"].astype(jnp.float32), + reference_c, + atol=8e-2, + rtol=3e-2, + ) diff --git a/test/python/fe_api/test_jax_gemm_swiglu_contract.py b/test/python/fe_api/test_jax_gemm_swiglu_contract.py new file mode 100644 index 000000000..5fb446a4a --- /dev/null +++ b/test/python/fe_api/test_jax_gemm_swiglu_contract.py @@ -0,0 +1,438 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for the standard JAX GEMM + SwiGLU adapter.""" + +from __future__ import annotations + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +import os +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "gemm_swiglu" +_PACKAGE = "cudnn_jax_gemm_swiglu_contract_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + FP8_E4M3 = auto() + FP8_E5M2 = auto() + + +_DTYPE_TO_CUDNN = { + "float16": _DataType.HALF, + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, +} +_CUDNN_TO_DTYPE = {value: key for key, value in _DTYPE_TO_CUDNN.items()} + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + +class _TensorSpec: + def __init__(self, *, layout, mode, divisibility=None): + self.layout = layout + self.mode = mode + self.divisibility = divisibility + + +class JaxGemmSwigluContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + operation_name = f"{_PACKAGE}.gemm_swiglu" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + layout_module = importlib.import_module(f"{internal_name}.layout") + result_module = importlib.import_module(f"{_PACKAGE}._result") + + class JaxTensorDesc(tensor_module.TensorDesc): + @property + def cudnn_dtype(self): + return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) + + class JaxApiBase: + @staticmethod + def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): + public_shape = tuple(value.shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + public_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride(public_shape, public_order) + canonical_axis_by_public_axis = layout_module.to_public_axes(tuple(range(len(public_shape))), mode) + return JaxTensorDesc( + dtype=value.dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple(canonical_axis_by_public_axis[axis] for axis in public_order), + name=name, + init_value=init_value, + ) + + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del operation_name + resolved = 100 if target is None else target + if resolved not in supported: + raise ValueError(f"unsupported target {resolved}") + return resolved + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + actual_shape = layout_module.to_canonical_axes(tuple(value.shape), mode) + if actual_shape != expected.shape: + raise ValueError(f"{expected.name} tensor shape mismatch: " f"expected {expected.shape}, got {actual_shape}") + actual_dtype = _DTYPE_TO_CUDNN.get(value.dtype, _DataType.NOT_SET) + if actual_dtype != expected.cudnn_dtype: + raise ValueError(f"{expected.name} tensor dtype mismatch: " f"expected {expected.cudnn_dtype}, got {actual_dtype}") + + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + return _TensorSpec( + layout=layout_module.to_cutlass_layout( + desc.shape, + desc.stride, + desc.stride_order, + mode=mode, + name=desc.name, + ), + mode=mode, + divisibility=divisibility, + ) + + def _call_kernel(self, inputs, **options): + self.captured_call = (tuple(inputs), options) + return tuple( + _Array( + layout_module.to_public_axes(desc.shape, spec.mode), + _CUDNN_TO_DTYPE[desc.cudnn_dtype], + ) + for desc, spec in zip(options["output_descs"], options["output_spec"]) + ) + + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = result_module.TupleDict + + datatypes = types.ModuleType(f"{internal_name}.datatypes") + datatypes.jax_to_cudnn_dtype = lambda dtype: _DTYPE_TO_CUDNN.get(dtype, _DataType.NOT_SET) + datatypes.cudnn_to_jax_dtype = lambda dtype: _CUDNN_TO_DTYPE[dtype] + sys.modules[datatypes.__name__] = datatypes + + cls.jit_static_argnames = {} + fake_jax = types.ModuleType("jax") + + def jit(function=None, *, static_argnames=()): + def decorate(target): + cls.jit_static_argnames[target.__name__] = tuple(static_argnames) + return target + + return decorate if function is None else decorate(function) + + fake_jax.jit = jit + fake_jax.ShapeDtypeStruct = _Array + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.dtype = lambda dtype: dtype + fake_jnp.float16 = "float16" + fake_jnp.bfloat16 = "bfloat16" + fake_jnp.float32 = "float32" + fake_jax.numpy = fake_jnp + + try: + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "torch": None, + "cutlass": None, + }, + ): + cls.module = importlib.import_module(f"{operation_name}.jax") + except Exception: + cls.tearDownClass() + raise + + cls.JaxApiBase = JaxApiBase + cls.operation_name = operation_name + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def setUp(self) -> None: + self.JaxApiBase.captured_call = None + + @staticmethod + def _samples(a_layout="LMK", b_layout="LNK"): + m, n, k, batch = 128, 192, 64, 3 + a_shape = (batch, m, k) if a_layout == "LMK" else (batch, k, m) + b_shape = (batch, n, k) if b_layout == "LNK" else (batch, k, n) + return _Array(a_shape, "bfloat16"), _Array(b_shape, "bfloat16") + + def test_default_layouts_infer_outputs_and_declare_cutlass_call(self): + sample_a, sample_b = self._samples() + api = self.module.GemmSwigluSm100( + sample_a, + sample_b, + target_compute_capability=100, + ) + + self.assertEqual(api.a_mode, (1, 2, 0)) + self.assertEqual(api.b_mode, (1, 2, 0)) + self.assertEqual(api.output_mode, (1, 2, 0)) + self.assertEqual(api.a_desc.shape, (128, 64, 3)) + self.assertEqual(api.b_desc.shape, (192, 64, 3)) + self.assertEqual(api.ab12_desc.shape, (128, 192, 3)) + self.assertEqual(api.c_desc.shape, (128, 96, 3)) + self.assertIs(api._op.a, api.a_desc) + self.assertIs(api._op.b, api.b_desc) + self.assertIs(api._op.ab12, api.ab12_desc) + self.assertIs(api._op.c, api.c_desc) + + result = api(sample_a, sample_b) + self.assertEqual( + tuple(result.keys()), + ("ab12_tensor", "c_tensor", "sfc_tensor", "amax_tensor"), + ) + self.assertEqual(result["ab12_tensor"].shape, (3, 128, 192)) + self.assertEqual(result["ab12_tensor"].dtype, "float32") + self.assertEqual(result["c_tensor"].shape, (3, 128, 96)) + self.assertEqual(result["c_tensor"].dtype, "float16") + self.assertIsNone(result["sfc_tensor"]) + self.assertIsNone(result["amax_tensor"]) + + inputs, options = api.captured_call + self.assertEqual(inputs, (sample_a, sample_b)) + self.assertEqual(options["output_descs"], (api.ab12_desc, api.c_desc)) + self.assertEqual( + tuple(spec.mode for spec in options["input_spec"]), + (api.a_mode, api.b_mode), + ) + self.assertEqual( + tuple(spec.mode for spec in options["output_spec"]), + (api.output_mode, api.output_mode), + ) + self.assertEqual( + tuple(spec.layout for spec in options["input_spec"]), + ((2, 1, 0), (2, 1, 0)), + ) + self.assertEqual( + tuple(spec.layout for spec in options["output_spec"]), + ((2, 1, 0), (2, 1, 0)), + ) + self.assertIn("--gpu-arch sm_100a", options["compile_options"]) + + def test_alternate_layouts_map_public_shapes_to_the_same_canonical_op(self): + sample_a, sample_b = self._samples(a_layout="LKM", b_layout="LKN") + api = self.module.GemmSwigluSm100( + sample_a, + sample_b, + a_layout="LKM", + b_layout="LKN", + c_layout="LNM", + target_compute_capability=103, + ) + result = api(sample_a, sample_b) + + self.assertEqual(api.a_mode, (2, 1, 0)) + self.assertEqual(api.b_mode, (2, 1, 0)) + self.assertEqual(api.output_mode, (2, 1, 0)) + self.assertEqual(api.a_desc.shape, (128, 64, 3)) + self.assertEqual(api.b_desc.shape, (192, 64, 3)) + self.assertEqual(api.a_desc.stride_order, (0, 1, 2)) + self.assertEqual(api.b_desc.stride_order, (0, 1, 2)) + self.assertEqual(api.ab12_desc.stride_order, (0, 1, 2)) + self.assertEqual(api.c_desc.stride_order, (0, 1, 2)) + self.assertEqual(result["ab12_tensor"].shape, (3, 192, 128)) + self.assertEqual(result["c_tensor"].shape, (3, 96, 128)) + self.assertIn("--gpu-arch sm_103a", api.captured_call[1]["compile_options"]) + + def test_explicit_output_exemplars_are_checked_and_preserved(self): + sample_a, sample_b = self._samples() + sample_ab12 = _Array((3, 128, 192), "bfloat16") + sample_c = _Array((3, 128, 96), "bfloat16") + api = self.module.GemmSwigluSm100( + sample_a, + sample_b, + sample_ab12=sample_ab12, + sample_c=sample_c, + ab12_dtype="bfloat16", + c_dtype="bfloat16", + target_compute_capability=100, + ) + result = api(sample_a, sample_b) + self.assertEqual(result["ab12_tensor"].dtype, "bfloat16") + self.assertEqual(result["c_tensor"].dtype, "bfloat16") + + with self.assertRaisesRegex(ValueError, "sample_ab12 and sample_c must be provided together"): + self.module.GemmSwigluSm100( + sample_a, + sample_b, + sample_ab12=sample_ab12, + target_compute_capability=100, + ) + with self.assertRaisesRegex(ValueError, "c_dtype=float16 does not match the explicit sample dtype"): + self.module.GemmSwigluSm100( + sample_a, + sample_b, + sample_ab12=sample_ab12, + sample_c=sample_c, + c_dtype="float16", + target_compute_capability=100, + ) + + def test_runtime_inputs_must_match_the_specialized_signature(self): + sample_a, sample_b = self._samples() + api = self.module.GemmSwigluSm100( + sample_a, + sample_b, + target_compute_capability=100, + ) + with self.assertRaisesRegex(ValueError, "sample_a tensor shape mismatch"): + api(_Array((3, 64, 64), "bfloat16"), sample_b) + with self.assertRaisesRegex(ValueError, "sample_b tensor dtype mismatch"): + api(sample_a, _Array(sample_b.shape, "float16")) + + def test_wrapper_marks_configuration_as_static_and_remains_functional(self): + static = self.jit_static_argnames["gemm_swiglu_wrapper_sm100"] + self.assertEqual( + static, + ( + "alpha", + "c_layout", + "ab12_dtype", + "c_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "a_layout", + "b_layout", + "target_compute_capability", + ), + ) + + sample_a, sample_b = self._samples(a_layout="LKM", b_layout="LKN") + result = self.module.gemm_swiglu_wrapper_sm100( + sample_a, + sample_b, + c_layout="LNM", + a_layout="LKM", + b_layout="LKN", + target_compute_capability=107, + ) + self.assertEqual(result["ab12_tensor"].shape, (3, 192, 128)) + self.assertEqual(result["c_tensor"].shape, (3, 96, 128)) + + def test_launch_constructs_the_kernel_internally_and_preserves_abi_order(self): + _, sample_b = self._samples() + sample_a = _Array((3, 256, 64), "bfloat16") + with mock.patch.dict(os.environ, {"CUDNNFE_CLUSTER_OVERLAP_MARGIN": "2"}): + api = self.module.GemmSwigluSm100( + sample_a, + sample_b, + alpha=0.25, + mma_tiler_mn=(256, 128), + target_compute_capability=100, + ) + api.check_support() + seen = {} + + class HardwareInfo: + def get_max_active_clusters(self, cluster_size): + seen["cluster_size"] = cluster_size + return 12 + + class Kernel: + def __init__(self, **configuration): + seen["configuration"] = configuration + + def __call__(self, *arguments): + seen["arguments"] = arguments + + cutlass = types.ModuleType("cutlass") + cutlass.Float32 = lambda value: ("Float32", value) + cutlass.utils = types.SimpleNamespace(HardwareInfo=HardwareInfo) + cutlass_jax = types.ModuleType("cutlass.jax") + cutlass_jax.jax_to_cutlass_dtype = lambda dtype: ("cutlass_dtype", dtype) + cutlass.jax = cutlass_jax + kernel_module = types.ModuleType(f"{self.operation_name}.dense_gemm_persistent_swiglu") + kernel_module.PersistentDenseGemmKernel = Kernel + stream = object() + + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + "cutlass.jax": cutlass_jax, + kernel_module.__name__: kernel_module, + }, + ): + api._launch(("A", "B"), ("AB12", "C"), (), stream) + + self.assertEqual( + seen["configuration"], + { + "acc_dtype": ("cutlass_dtype", "float32"), + "use_2cta_instrs": True, + "mma_tiler_mn": (256, 128), + "cluster_shape_mn": (2, 2), + }, + ) + self.assertEqual(seen["cluster_size"], 4) + self.assertEqual( + seen["arguments"], + ("A", "B", "AB12", "C", ("Float32", 0.25), 10, stream), + ) + + with self.assertRaisesRegex(RuntimeError, "unexpected workspaces"): + api._launch(("A", "B"), ("AB12", "C"), (object(),), stream) + + def test_adapter_import_does_not_load_torch_or_the_cute_kernel(self): + self.assertNotIn(f"{self.operation_name}.api", sys.modules) + self.assertNotIn(f"{self.operation_name}.dense_gemm_persistent_swiglu", sys.modules) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_import_contract.py b/test/python/fe_api/test_jax_import_contract.py index bbfdacd41..7bfb03769 100644 --- a/test/python/fe_api/test_jax_import_contract.py +++ b/test/python/fe_api/test_jax_import_contract.py @@ -189,7 +189,7 @@ def literal_assignment(path, name): "_SYMBOLS", ) jax_exports = literal_assignment(_CUDNN_ROOT / "jax" / "__init__.py", "_OPERATION_EXPORTS") - jax_dsa_exports = {name for name in jax_exports if name not in {"RmsNormRhtAmaxSm100", "rmsnorm_rht_amax_sm100"}} + jax_dsa_exports = {name for name, (module_name, _) in jax_exports.items() if module_name.startswith("..deepseek_sparse_attention.")} self.assertEqual(jax_dsa_exports, set(torch_exports)) From 6e19a911156db2309ef0650fd93d9a6187b48f01 Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Tue, 7 Jul 2026 21:58:17 -0500 Subject: [PATCH 06/13] Make JAX GEMM SwiGLU wrapper target-generic --- docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md | 9 ++++++--- python/cudnn/gemm_swiglu/jax.py | 6 +++--- python/cudnn/jax/__init__.py | 2 +- test/python/fe_api/test_jax_gemm_swiglu.py | 8 ++++---- test/python/fe_api/test_jax_gemm_swiglu_contract.py | 6 ++++-- 5 files changed, 18 insertions(+), 13 deletions(-) diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md index 1e9816f4e..e826e820f 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md @@ -190,10 +190,10 @@ row-major arrays whose axis order is described explicitly: import jax import jax.numpy as jnp -from cudnn.jax import gemm_swiglu_wrapper_sm100 +from cudnn.jax import gemm_swiglu_wrapper # A: (L, M, K), B: (L, N, K) -result = gemm_swiglu_wrapper_sm100( +result = gemm_swiglu_wrapper( a, b, alpha=1.0, @@ -325,7 +325,10 @@ also the Torch API order. JAX arrays use the public axis orders selected by - Number of pipeline stages for AB12 output - Default: `4` -### Wrapper-specific parameters: `gemm_swiglu_wrapper_sm100` +### Wrapper-specific parameters + +The Torch wrapper is `gemm_swiglu_wrapper_sm100`. The JAX wrapper is +`gemm_swiglu_wrapper`; it resolves the concrete kernel target while tracing. - `a_tensor`, `b_tensor`: see Input/Output tensors - Torch: `c_major` selects `"m"` or `"n"`; `ab12_dtype` defaults to `torch.float32`; `c_dtype` defaults to `torch.float16`. diff --git a/python/cudnn/gemm_swiglu/jax.py b/python/cudnn/gemm_swiglu/jax.py index 91151045b..bc2a2033d 100644 --- a/python/cudnn/gemm_swiglu/jax.py +++ b/python/cudnn/gemm_swiglu/jax.py @@ -238,7 +238,7 @@ def _launch( "target_compute_capability", ), ) -def gemm_swiglu_wrapper_sm100( +def gemm_swiglu_wrapper( a_tensor: Any, b_tensor: Any, alpha: float = 1.0, @@ -253,7 +253,7 @@ def gemm_swiglu_wrapper_sm100( b_layout: str = "LNK", target_compute_capability: int | None = None, ) -> TupleDict: - """Compute the standard dense batched GEMM and fused SwiGLU projection.""" + """Resolve a supported GPU target and compute the standard GEMM + SwiGLU projection.""" return GemmSwigluSm100( jax.ShapeDtypeStruct(a_tensor.shape, a_tensor.dtype), @@ -274,5 +274,5 @@ def gemm_swiglu_wrapper_sm100( __all__ = [ "GemmSwigluSm100", "SUPPORTED_COMPUTE_CAPABILITIES", - "gemm_swiglu_wrapper_sm100", + "gemm_swiglu_wrapper", ] diff --git a/python/cudnn/jax/__init__.py b/python/cudnn/jax/__init__.py index 92ac00453..0972b4e42 100644 --- a/python/cudnn/jax/__init__.py +++ b/python/cudnn/jax/__init__.py @@ -38,7 +38,7 @@ "RmsNormRhtAmaxSm100": ("..rmsnorm_rht_amax.jax", "RmsNormRhtAmaxSm100"), "rmsnorm_rht_amax_sm100": ("..rmsnorm_rht_amax.jax", "rmsnorm_rht_amax_sm100"), "GemmSwigluSm100": ("..gemm_swiglu.jax", "GemmSwigluSm100"), - "gemm_swiglu_wrapper_sm100": ("..gemm_swiglu.jax", "gemm_swiglu_wrapper_sm100"), + "gemm_swiglu_wrapper": ("..gemm_swiglu.jax", "gemm_swiglu_wrapper"), "SparseAttentionBackward": ("..deepseek_sparse_attention.sparse_attention_backward.jax", "SparseAttentionBackward"), "sparse_attention_backward_wrapper": ( "..deepseek_sparse_attention.sparse_attention_backward.jax", diff --git a/test/python/fe_api/test_jax_gemm_swiglu.py b/test/python/fe_api/test_jax_gemm_swiglu.py index 77c5ea6f4..188d66921 100644 --- a/test/python/fe_api/test_jax_gemm_swiglu.py +++ b/test/python/fe_api/test_jax_gemm_swiglu.py @@ -63,7 +63,7 @@ def test_jax_gemm_swiglu_abstract_contract( jax, jnp = _jax_runtime() from cudnn._jax import JaxApiBase - from cudnn.jax import GemmSwigluSm100, gemm_swiglu_wrapper_sm100 + from cudnn.jax import GemmSwigluSm100, gemm_swiglu_wrapper def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): return tuple( @@ -100,7 +100,7 @@ def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): alternate_a = jax.ShapeDtypeStruct((3, 64, 128), jnp.bfloat16) alternate_b = jax.ShapeDtypeStruct((3, 64, 192), jnp.bfloat16) alternate = jax.eval_shape( - lambda a, b: gemm_swiglu_wrapper_sm100( + lambda a, b: gemm_swiglu_wrapper( a, b, a_layout="LKM", @@ -133,7 +133,7 @@ def test_jax_gemm_swiglu_sm100_jit_and_numerics(): jax, jnp = _jax_runtime() device, capability = _supported_gpu(jax) - from cudnn.jax import gemm_swiglu_wrapper_sm100 + from cudnn.jax import gemm_swiglu_wrapper batch, m, n, k = 1, 128, 128, 128 alpha = 0.5 @@ -156,7 +156,7 @@ def test_jax_gemm_swiglu_sm100_jit_and_numerics(): device, ) - lowered = gemm_swiglu_wrapper_sm100.lower( + lowered = gemm_swiglu_wrapper.lower( a, b, alpha=alpha, diff --git a/test/python/fe_api/test_jax_gemm_swiglu_contract.py b/test/python/fe_api/test_jax_gemm_swiglu_contract.py index 5fb446a4a..dbaa8c6ee 100644 --- a/test/python/fe_api/test_jax_gemm_swiglu_contract.py +++ b/test/python/fe_api/test_jax_gemm_swiglu_contract.py @@ -336,7 +336,9 @@ def test_runtime_inputs_must_match_the_specialized_signature(self): api(sample_a, _Array(sample_b.shape, "float16")) def test_wrapper_marks_configuration_as_static_and_remains_functional(self): - static = self.jit_static_argnames["gemm_swiglu_wrapper_sm100"] + self.assertIn("gemm_swiglu_wrapper", self.module.__all__) + self.assertNotIn("gemm_swiglu_wrapper_sm100", self.module.__all__) + static = self.jit_static_argnames["gemm_swiglu_wrapper"] self.assertEqual( static, ( @@ -354,7 +356,7 @@ def test_wrapper_marks_configuration_as_static_and_remains_functional(self): ) sample_a, sample_b = self._samples(a_layout="LKM", b_layout="LKN") - result = self.module.gemm_swiglu_wrapper_sm100( + result = self.module.gemm_swiglu_wrapper( sample_a, sample_b, c_layout="LNM", From 86792fd73a51c853c137a152962596a940308d59 Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Tue, 7 Jul 2026 22:06:37 -0500 Subject: [PATCH 07/13] Require a local GPU for JAX GEMM SwiGLU --- docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md | 16 ++++----- python/cudnn/_jax/api_base.py | 14 +++++--- python/cudnn/gemm_swiglu/jax.py | 17 ++++------ python/cudnn/jax/__init__.py | 2 +- test/python/fe_api/test_jax_api_base.py | 14 ++++++++ test/python/fe_api/test_jax_gemm_swiglu.py | 34 ++++++------------- .../fe_api/test_jax_gemm_swiglu_contract.py | 22 ++++-------- 7 files changed, 57 insertions(+), 62 deletions(-) diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md index e826e820f..2d51e3684 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md @@ -190,10 +190,10 @@ row-major arrays whose axis order is described explicitly: import jax import jax.numpy as jnp -from cudnn.jax import gemm_swiglu_wrapper +from cudnn.jax import gemm_swiglu_wrapper_sm100 # A: (L, M, K), B: (L, N, K) -result = gemm_swiglu_wrapper( +result = gemm_swiglu_wrapper_sm100( a, b, alpha=1.0, @@ -232,10 +232,10 @@ Supported public layouts are: - B: `LNK` (K-major) or `LKN` (N-major) - AB12/C: `LMN` (N-major) or `LNM` (M-major) -An explicit `target_compute_capability` may be used after calling -`cudnn.jax.disable_device_compatibility_checks(True)` to bypass the local-device -compatibility check. CuTe lowering still queries CUDA occupancy and therefore -requires access to a CUDA device. +The JAX API requires a homogeneous, supported local SM100-family GPU. It +resolves the exact local compute capability for CuTe compilation and queries +that GPU's occupancy during lowering. Disabling device compatibility checks +does not enable deviceless or remote AOT compilation for this operation. --- @@ -327,8 +327,8 @@ also the Torch API order. JAX arrays use the public axis orders selected by ### Wrapper-specific parameters -The Torch wrapper is `gemm_swiglu_wrapper_sm100`. The JAX wrapper is -`gemm_swiglu_wrapper`; it resolves the concrete kernel target while tracing. +Torch and JAX both expose `gemm_swiglu_wrapper_sm100`. The JAX wrapper resolves +the exact local SM100-family compilation target while tracing. - `a_tensor`, `b_tensor`: see Input/Output tensors - Torch: `c_major` selects `"m"` or `"n"`; `ab12_dtype` defaults to `torch.float32`; `c_dtype` defaults to `torch.float16`. diff --git a/python/cudnn/_jax/api_base.py b/python/cudnn/_jax/api_base.py index 1a45184da..d227525d7 100644 --- a/python/cudnn/_jax/api_base.py +++ b/python/cudnn/_jax/api_base.py @@ -150,6 +150,8 @@ def _resolve_compute_capability( target_compute_capability: int | None, supported_compute_capabilities: tuple[int, ...], operation_name: str, + *, + require_local: bool = False, ) -> int: """Resolve an exact JAX compilation target for a multi-arch adapter. @@ -157,7 +159,8 @@ def _resolve_compute_capability( targets implemented by the adapter. The returned value remains exact (for example ``103``) so the compiler selects the matching target. With compatibility checks enabled, an explicit target must match every - local GPU exactly. Disable checks for cross-compilation. + local GPU exactly. Disable checks for cross-compilation. ``require_local`` + keeps local-device resolution mandatory even when checks are disabled. """ supported = tuple(sorted(set(supported_compute_capabilities))) @@ -172,7 +175,7 @@ def _resolve_compute_capability( if target_compute_capability not in supported: supported_text = ", ".join(f"SM{value}" for value in supported) raise ValueError(f"{operation_name} has no kernel for SM{target_compute_capability}; supported targets are {supported_text}") - if _DEVICE_COMPATIBILITY_CHECKS_DISABLED: + if _DEVICE_COMPATIBILITY_CHECKS_DISABLED and not require_local: return target_compute_capability local = JaxApiBase._local_gpu_capabilities(operation_name) @@ -186,17 +189,18 @@ def _resolve_compute_capability( raise RuntimeError(f"{operation_name} targets SM{target_compute_capability}, but found {found}. " f"{_DEVICE_COMPATIBILITY_CHECK_DISABLE_HINT}") return target_compute_capability + local_target_hint = "This operation requires a homogeneous supported local JAX GPU." if require_local else _TARGET_COMPUTE_CAPABILITY_HINT try: local = JaxApiBase._local_gpu_capabilities(operation_name) except RuntimeError as error: - raise RuntimeError(f"{error}. {_TARGET_COMPUTE_CAPABILITY_HINT}") from error + raise RuntimeError(f"{error}. {local_target_hint}") from error if not local: - raise RuntimeError(f"{operation_name}: no local JAX GPU is available. {_TARGET_COMPUTE_CAPABILITY_HINT}") + raise RuntimeError(f"{operation_name}: no local JAX GPU is available. {local_target_hint}") capabilities = tuple(sorted({capability for _, capability in local})) if len(capabilities) != 1: found = ", ".join(f"SM{capability}" for capability in capabilities) - raise RuntimeError(f"{operation_name}: local JAX GPUs have heterogeneous targets ({found}). {_TARGET_COMPUTE_CAPABILITY_HINT}") + raise RuntimeError(f"{operation_name}: local JAX GPUs have heterogeneous targets ({found}). {local_target_hint}") resolved = capabilities[0] if resolved not in supported: diff --git a/python/cudnn/gemm_swiglu/jax.py b/python/cudnn/gemm_swiglu/jax.py index bc2a2033d..33844eb68 100644 --- a/python/cudnn/gemm_swiglu/jax.py +++ b/python/cudnn/gemm_swiglu/jax.py @@ -59,7 +59,6 @@ def __init__( cluster_shape_mn: tuple[int, int] | None = None, a_layout: str = "LMK", b_layout: str = "LNK", - target_compute_capability: int | None = None, ) -> None: self.a_layout = a_layout self.b_layout = b_layout @@ -68,10 +67,11 @@ def __init__( self.b_mode = gemm_b_mode(b_layout) self.output_mode = gemm_output_mode(c_layout) - self.target_compute_capability = self._resolve_compute_capability( - target_compute_capability, + self.compute_capability = self._resolve_compute_capability( + None, SUPPORTED_COMPUTE_CAPABILITIES, "GemmSwigluSm100", + require_local=True, ) self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=self.a_mode) self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=self.b_mode) @@ -176,7 +176,7 @@ def __call__(self, a_tensor: Any, b_tensor: Any) -> TupleDict: self._to_tensor_spec(self.ab12_desc, mode=self.output_mode), self._to_tensor_spec(self.c_desc, mode=self.output_mode), ), - compile_options=compile_options_for_target(self.target_compute_capability), + compile_options=compile_options_for_target(self.compute_capability), ) return TupleDict( ab12_tensor=ab12_tensor, @@ -235,10 +235,9 @@ def _launch( "cluster_shape_mn", "a_layout", "b_layout", - "target_compute_capability", ), ) -def gemm_swiglu_wrapper( +def gemm_swiglu_wrapper_sm100( a_tensor: Any, b_tensor: Any, alpha: float = 1.0, @@ -251,9 +250,8 @@ def gemm_swiglu_wrapper( *, a_layout: str = "LMK", b_layout: str = "LNK", - target_compute_capability: int | None = None, ) -> TupleDict: - """Resolve a supported GPU target and compute the standard GEMM + SwiGLU projection.""" + """Compute the standard GEMM + SwiGLU projection on a local SM100-family GPU.""" return GemmSwigluSm100( jax.ShapeDtypeStruct(a_tensor.shape, a_tensor.dtype), @@ -267,12 +265,11 @@ def gemm_swiglu_wrapper( cluster_shape_mn=cluster_shape_mn, a_layout=a_layout, b_layout=b_layout, - target_compute_capability=target_compute_capability, )(a_tensor, b_tensor) __all__ = [ "GemmSwigluSm100", "SUPPORTED_COMPUTE_CAPABILITIES", - "gemm_swiglu_wrapper", + "gemm_swiglu_wrapper_sm100", ] diff --git a/python/cudnn/jax/__init__.py b/python/cudnn/jax/__init__.py index 0972b4e42..92ac00453 100644 --- a/python/cudnn/jax/__init__.py +++ b/python/cudnn/jax/__init__.py @@ -38,7 +38,7 @@ "RmsNormRhtAmaxSm100": ("..rmsnorm_rht_amax.jax", "RmsNormRhtAmaxSm100"), "rmsnorm_rht_amax_sm100": ("..rmsnorm_rht_amax.jax", "rmsnorm_rht_amax_sm100"), "GemmSwigluSm100": ("..gemm_swiglu.jax", "GemmSwigluSm100"), - "gemm_swiglu_wrapper": ("..gemm_swiglu.jax", "gemm_swiglu_wrapper"), + "gemm_swiglu_wrapper_sm100": ("..gemm_swiglu.jax", "gemm_swiglu_wrapper_sm100"), "SparseAttentionBackward": ("..deepseek_sparse_attention.sparse_attention_backward.jax", "SparseAttentionBackward"), "sparse_attention_backward_wrapper": ( "..deepseek_sparse_attention.sparse_attention_backward.jax", diff --git a/test/python/fe_api/test_jax_api_base.py b/test/python/fe_api/test_jax_api_base.py index b9e184a51..9998b199d 100644 --- a/test/python/fe_api/test_jax_api_base.py +++ b/test/python/fe_api/test_jax_api_base.py @@ -252,6 +252,20 @@ def test_implicit_target_still_needs_a_device_when_checks_are_disabled(self): with self.assertRaisesRegex(RuntimeError, r"no local JAX GPU.*target_compute_capability"): self.module.JaxApiBase._resolve_compute_capability(None, (90, 100), "TestOp") + def test_require_local_keeps_device_resolution_strict(self): + self.module.disable_device_compatibility_checks(True) + jax = types.ModuleType("jax") + jax.local_devices = lambda *, backend: () + + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaisesRegex(RuntimeError, r"no local JAX GPU.*requires a homogeneous supported local JAX GPU"): + self.module.JaxApiBase._resolve_compute_capability( + None, + (100, 103, 107), + "TestOp", + require_local=True, + ) + def test_converts_array_metadata_to_shared_tensor_desc(self): desc = self.module.JaxApiBase._to_tensor_desc( _ArrayMetadata((2, 3, 4), "bfloat16"), diff --git a/test/python/fe_api/test_jax_gemm_swiglu.py b/test/python/fe_api/test_jax_gemm_swiglu.py index 188d66921..75a8911f3 100644 --- a/test/python/fe_api/test_jax_gemm_swiglu.py +++ b/test/python/fe_api/test_jax_gemm_swiglu.py @@ -40,30 +40,17 @@ def _supported_gpu(jax): capability = capabilities.pop() if capability not in {100, 103, 107}: pytest.skip("A supported JAX SM100-family GPU is not available") - return devices[0], capability - - -@pytest.fixture -def device_compatibility_checks_disabled(): - _jax_runtime() - from cudnn.jax import disable_device_compatibility_checks - - disable_device_compatibility_checks(True) - try: - yield - finally: - disable_device_compatibility_checks(False) + return devices[0] @pytest.mark.L0 def test_jax_gemm_swiglu_abstract_contract( monkeypatch, - device_compatibility_checks_disabled, ): jax, jnp = _jax_runtime() from cudnn._jax import JaxApiBase - from cudnn.jax import GemmSwigluSm100, gemm_swiglu_wrapper + from cudnn.jax import GemmSwigluSm100, gemm_swiglu_wrapper_sm100 def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): return tuple( @@ -75,13 +62,17 @@ def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): ) monkeypatch.setattr(JaxApiBase, "_call_kernel", abstract_call) + monkeypatch.setattr( + JaxApiBase, + "_resolve_compute_capability", + staticmethod(lambda _target, _supported, _operation, **_options: 100), + ) sample_a = jax.ShapeDtypeStruct((3, 128, 64), jnp.bfloat16) sample_b = jax.ShapeDtypeStruct((3, 192, 64), jnp.bfloat16) api = GemmSwigluSm100( sample_a, sample_b, - target_compute_capability=100, ) result = jax.eval_shape(api, sample_a, sample_b) assert tuple(result.keys()) == ( @@ -100,14 +91,13 @@ def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): alternate_a = jax.ShapeDtypeStruct((3, 64, 128), jnp.bfloat16) alternate_b = jax.ShapeDtypeStruct((3, 64, 192), jnp.bfloat16) alternate = jax.eval_shape( - lambda a, b: gemm_swiglu_wrapper( + lambda a, b: gemm_swiglu_wrapper_sm100( a, b, a_layout="LKM", b_layout="LKN", c_layout="LNM", c_dtype=jnp.bfloat16, - target_compute_capability=103, ), alternate_a, alternate_b, @@ -121,7 +111,6 @@ def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): sample_b, sample_ab12=jax.ShapeDtypeStruct((3, 128, 192), jnp.bfloat16), sample_c=jax.ShapeDtypeStruct((3, 128, 96), jnp.bfloat16), - target_compute_capability=100, ) explicit_result = jax.eval_shape(explicit, sample_a, sample_b) assert explicit_result["ab12_tensor"].dtype == jnp.bfloat16 @@ -131,9 +120,9 @@ def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): @pytest.mark.L0 def test_jax_gemm_swiglu_sm100_jit_and_numerics(): jax, jnp = _jax_runtime() - device, capability = _supported_gpu(jax) + device = _supported_gpu(jax) - from cudnn.jax import gemm_swiglu_wrapper + from cudnn.jax import gemm_swiglu_wrapper_sm100 batch, m, n, k = 1, 128, 128, 128 alpha = 0.5 @@ -156,14 +145,13 @@ def test_jax_gemm_swiglu_sm100_jit_and_numerics(): device, ) - lowered = gemm_swiglu_wrapper.lower( + lowered = gemm_swiglu_wrapper_sm100.lower( a, b, alpha=alpha, a_layout="LKM", b_layout="LKN", c_layout="LNM", - target_compute_capability=capability, ) stablehlo = lowered.as_text("stablehlo") assert stablehlo.count("stablehlo.custom_call") == 1 diff --git a/test/python/fe_api/test_jax_gemm_swiglu_contract.py b/test/python/fe_api/test_jax_gemm_swiglu_contract.py index dbaa8c6ee..d4be00289 100644 --- a/test/python/fe_api/test_jax_gemm_swiglu_contract.py +++ b/test/python/fe_api/test_jax_gemm_swiglu_contract.py @@ -109,8 +109,9 @@ def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): ) @staticmethod - def _resolve_compute_capability(target, supported, operation_name): + def _resolve_compute_capability(target, supported, operation_name, *, require_local=False): del operation_name + assert require_local resolved = 100 if target is None else target if resolved not in supported: raise ValueError(f"unsupported target {resolved}") @@ -216,7 +217,6 @@ def test_default_layouts_infer_outputs_and_declare_cutlass_call(self): api = self.module.GemmSwigluSm100( sample_a, sample_b, - target_compute_capability=100, ) self.assertEqual(api.a_mode, (1, 2, 0)) @@ -272,7 +272,6 @@ def test_alternate_layouts_map_public_shapes_to_the_same_canonical_op(self): a_layout="LKM", b_layout="LKN", c_layout="LNM", - target_compute_capability=103, ) result = api(sample_a, sample_b) @@ -287,7 +286,7 @@ def test_alternate_layouts_map_public_shapes_to_the_same_canonical_op(self): self.assertEqual(api.c_desc.stride_order, (0, 1, 2)) self.assertEqual(result["ab12_tensor"].shape, (3, 192, 128)) self.assertEqual(result["c_tensor"].shape, (3, 96, 128)) - self.assertIn("--gpu-arch sm_103a", api.captured_call[1]["compile_options"]) + self.assertIn("--gpu-arch sm_100a", api.captured_call[1]["compile_options"]) def test_explicit_output_exemplars_are_checked_and_preserved(self): sample_a, sample_b = self._samples() @@ -300,7 +299,6 @@ def test_explicit_output_exemplars_are_checked_and_preserved(self): sample_c=sample_c, ab12_dtype="bfloat16", c_dtype="bfloat16", - target_compute_capability=100, ) result = api(sample_a, sample_b) self.assertEqual(result["ab12_tensor"].dtype, "bfloat16") @@ -311,7 +309,6 @@ def test_explicit_output_exemplars_are_checked_and_preserved(self): sample_a, sample_b, sample_ab12=sample_ab12, - target_compute_capability=100, ) with self.assertRaisesRegex(ValueError, "c_dtype=float16 does not match the explicit sample dtype"): self.module.GemmSwigluSm100( @@ -320,7 +317,6 @@ def test_explicit_output_exemplars_are_checked_and_preserved(self): sample_ab12=sample_ab12, sample_c=sample_c, c_dtype="float16", - target_compute_capability=100, ) def test_runtime_inputs_must_match_the_specialized_signature(self): @@ -328,7 +324,6 @@ def test_runtime_inputs_must_match_the_specialized_signature(self): api = self.module.GemmSwigluSm100( sample_a, sample_b, - target_compute_capability=100, ) with self.assertRaisesRegex(ValueError, "sample_a tensor shape mismatch"): api(_Array((3, 64, 64), "bfloat16"), sample_b) @@ -336,9 +331,9 @@ def test_runtime_inputs_must_match_the_specialized_signature(self): api(sample_a, _Array(sample_b.shape, "float16")) def test_wrapper_marks_configuration_as_static_and_remains_functional(self): - self.assertIn("gemm_swiglu_wrapper", self.module.__all__) - self.assertNotIn("gemm_swiglu_wrapper_sm100", self.module.__all__) - static = self.jit_static_argnames["gemm_swiglu_wrapper"] + self.assertIn("gemm_swiglu_wrapper_sm100", self.module.__all__) + self.assertNotIn("gemm_swiglu_wrapper", self.module.__all__) + static = self.jit_static_argnames["gemm_swiglu_wrapper_sm100"] self.assertEqual( static, ( @@ -351,18 +346,16 @@ def test_wrapper_marks_configuration_as_static_and_remains_functional(self): "cluster_shape_mn", "a_layout", "b_layout", - "target_compute_capability", ), ) sample_a, sample_b = self._samples(a_layout="LKM", b_layout="LKN") - result = self.module.gemm_swiglu_wrapper( + result = self.module.gemm_swiglu_wrapper_sm100( sample_a, sample_b, c_layout="LNM", a_layout="LKM", b_layout="LKN", - target_compute_capability=107, ) self.assertEqual(result["ab12_tensor"].shape, (3, 192, 128)) self.assertEqual(result["c_tensor"].shape, (3, 96, 128)) @@ -376,7 +369,6 @@ def test_launch_constructs_the_kernel_internally_and_preserves_abi_order(self): sample_b, alpha=0.25, mma_tiler_mn=(256, 128), - target_compute_capability=100, ) api.check_support() seen = {} From 705dc1bdff46a7ed3584c47bcc11da7f98085de4 Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Tue, 7 Jul 2026 22:14:54 -0500 Subject: [PATCH 08/13] Remove JAX device compatibility bypass --- docs/fe-oss-apis/dsa.md | 22 +----- docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md | 4 +- docs/fe-oss-apis/rmsnorm_rht_amax.md | 16 +---- python/cudnn/_jax/__init__.py | 4 +- python/cudnn/_jax/api_base.py | 50 +++----------- python/cudnn/gemm_swiglu/jax.py | 1 - python/cudnn/jax/__init__.py | 3 +- test/python/fe_api/test_dsa_indexer_ops.py | 21 +++++- test/python/fe_api/test_jax_api_base.py | 67 +++---------------- test/python/fe_api/test_jax_dsa_backward.py | 23 +++++-- test/python/fe_api/test_jax_dsa_indexers.py | 59 ++++++++-------- .../fe_api/test_jax_dsa_score_recompute.py | 23 ++----- .../fe_api/test_jax_gemm_swiglu_contract.py | 3 +- .../fe_api/test_jax_rmsnorm_rht_amax.py | 20 ++---- .../fe_api/test_rmsnorm_rht_amax_imports.py | 4 -- 15 files changed, 105 insertions(+), 215 deletions(-) diff --git a/docs/fe-oss-apis/dsa.md b/docs/fe-oss-apis/dsa.md index c5db25664..820a20805 100644 --- a/docs/fe-oss-apis/dsa.md +++ b/docs/fe-oss-apis/dsa.md @@ -185,26 +185,8 @@ and `d_index_k`. `target_compute_capability` is the exact compilation target, such as `90`, `100`, `103`, or `107`. When it is omitted, the adapter infers it from a -homogeneous set of local JAX GPUs. With compatibility checks enabled, an -explicit target must match every local GPU. - -For AOT or remote compilation without the target GPU, disable the local-device -check and provide the target before tracing: - -```python -cudnn_jax.disable_device_compatibility_checks(True) -lowered = cudnn_jax.indexer_forward_wrapper.lower( - q_spec, - k_spec, - w_spec, - target_compute_capability=100, -) -``` - -Pass `False` to re-enable checks. This process-wide setting only disables the -local-device validation; it does not enable an unsupported kernel target. -Changing it does not invalidate existing JAX compilation-cache entries, so -configure it before the first `jax.jit` or `lower` call for a signature. +homogeneous set of local JAX GPUs. An explicit target must match every local +GPU. Device-free AOT and remote compilation are not currently supported. ### JAX support matrix diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md index 2d51e3684..ba0e1891b 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md @@ -234,8 +234,8 @@ Supported public layouts are: The JAX API requires a homogeneous, supported local SM100-family GPU. It resolves the exact local compute capability for CuTe compilation and queries -that GPU's occupancy during lowering. Disabling device compatibility checks -does not enable deviceless or remote AOT compilation for this operation. +that GPU's occupancy during lowering. Device-free AOT and remote compilation +are not currently supported. --- diff --git a/docs/fe-oss-apis/rmsnorm_rht_amax.md b/docs/fe-oss-apis/rmsnorm_rht_amax.md index 398e17122..69493cd5c 100644 --- a/docs/fe-oss-apis/rmsnorm_rht_amax.md +++ b/docs/fe-oss-apis/rmsnorm_rht_amax.md @@ -143,20 +143,8 @@ results. Callers may instead provide both `sample_o` and `sample_amax` as explicit output placeholders. In either case, the common operation validates the complete signature through `check_support()`. -By default, JAX support validation requires every local GPU to meet the -operation's SM100 minimum. AOT or remote compilation without a compatible -local target device requires an explicit process-wide opt-out before tracing: - -```python -import cudnn.jax - -cudnn.jax.disable_device_compatibility_checks(True) -lowered = rmsnorm_rht_amax_sm100.lower(x_spec, w_spec) -``` - -Pass `False` to re-enable the checks. Changing this setting does not -invalidate existing JAX compilation-cache entries, so configure it before the -first `jax.jit` or `lower` call for a signature. +JAX support validation requires every local GPU to meet the operation's SM100 +minimum. Device-free AOT and remote compilation are not currently supported. The current JAX adapter models its buffers as row-major with the public and kernel axes in the same order. Additional JAX-specific axis mappings can be diff --git a/python/cudnn/_jax/__init__.py b/python/cudnn/_jax/__init__.py index 4048c3437..25f9cc8a7 100644 --- a/python/cudnn/_jax/__init__.py +++ b/python/cudnn/_jax/__init__.py @@ -3,7 +3,7 @@ """Internal implementation helpers for optional JAX operation adapters.""" -from .api_base import JaxApiBase, JaxTensorDesc, disable_device_compatibility_checks +from .api_base import JaxApiBase, JaxTensorDesc from .result import TupleDict -__all__ = ["JaxApiBase", "JaxTensorDesc", "TupleDict", "disable_device_compatibility_checks"] +__all__ = ["JaxApiBase", "JaxTensorDesc", "TupleDict"] diff --git a/python/cudnn/_jax/api_base.py b/python/cudnn/_jax/api_base.py index d227525d7..6c4d370da 100644 --- a/python/cudnn/_jax/api_base.py +++ b/python/cudnn/_jax/api_base.py @@ -14,24 +14,6 @@ from .._tensor_desc import TensorDesc from .layout import compact_stride, normalize_mode, to_canonical_axes, to_cutlass_layout, to_public_axes -_DEVICE_COMPATIBILITY_CHECKS_DISABLED = False -_DEVICE_COMPATIBILITY_CHECK_DISABLE_HINT = "Call cudnn.jax.disable_device_compatibility_checks(True) before tracing for a different or unavailable local GPU." -_TARGET_COMPUTE_CAPABILITY_HINT = "Pass target_compute_capability explicitly when tracing for AOT or on a host without the target GPU." - - -def disable_device_compatibility_checks(disabled: bool) -> None: - """Enable or disable local-device compatibility checks for future traces. - - This process-global setting does not invalidate existing JAX compilation - cache entries. Configure it before calling ``jax.jit`` or ``lower``. - """ - - if not isinstance(disabled, bool): - raise TypeError(f"disabled must be a bool, got {type(disabled).__name__}") - - global _DEVICE_COMPATIBILITY_CHECKS_DISABLED - _DEVICE_COMPATIBILITY_CHECKS_DISABLED = disabled - class JaxTensorDesc(TensorDesc[Any]): """Framework-neutral tensor metadata backed by a JAX dtype.""" @@ -150,17 +132,13 @@ def _resolve_compute_capability( target_compute_capability: int | None, supported_compute_capabilities: tuple[int, ...], operation_name: str, - *, - require_local: bool = False, ) -> int: """Resolve an exact JAX compilation target for a multi-arch adapter. ``supported_compute_capabilities`` lists the exact CuTe compilation targets implemented by the adapter. The returned value remains exact (for example ``103``) so the compiler selects the matching target. - With compatibility checks enabled, an explicit target must match every - local GPU exactly. Disable checks for cross-compilation. ``require_local`` - keeps local-device resolution mandatory even when checks are disabled. + An explicit target must match every local GPU exactly. """ supported = tuple(sorted(set(supported_compute_capabilities))) @@ -175,32 +153,23 @@ def _resolve_compute_capability( if target_compute_capability not in supported: supported_text = ", ".join(f"SM{value}" for value in supported) raise ValueError(f"{operation_name} has no kernel for SM{target_compute_capability}; supported targets are {supported_text}") - if _DEVICE_COMPATIBILITY_CHECKS_DISABLED and not require_local: - return target_compute_capability - local = JaxApiBase._local_gpu_capabilities(operation_name) if not local: - raise RuntimeError( - f"{operation_name} targets SM{target_compute_capability}, but no local JAX GPU is available. " f"{_DEVICE_COMPATIBILITY_CHECK_DISABLE_HINT}" - ) + raise RuntimeError(f"{operation_name} targets SM{target_compute_capability}, but no local JAX GPU is available") mismatched = tuple((device, capability) for device, capability in local if capability != target_compute_capability) if mismatched: found = ", ".join(f"{device} (SM{capability})" for device, capability in mismatched) - raise RuntimeError(f"{operation_name} targets SM{target_compute_capability}, but found {found}. " f"{_DEVICE_COMPATIBILITY_CHECK_DISABLE_HINT}") + raise RuntimeError(f"{operation_name} targets SM{target_compute_capability}, but found {found}") return target_compute_capability - local_target_hint = "This operation requires a homogeneous supported local JAX GPU." if require_local else _TARGET_COMPUTE_CAPABILITY_HINT - try: - local = JaxApiBase._local_gpu_capabilities(operation_name) - except RuntimeError as error: - raise RuntimeError(f"{error}. {local_target_hint}") from error + local = JaxApiBase._local_gpu_capabilities(operation_name) if not local: - raise RuntimeError(f"{operation_name}: no local JAX GPU is available. {local_target_hint}") + raise RuntimeError(f"{operation_name}: no local JAX GPU is available") capabilities = tuple(sorted({capability for _, capability in local})) if len(capabilities) != 1: found = ", ".join(f"SM{capability}" for capability in capabilities) - raise RuntimeError(f"{operation_name}: local JAX GPUs have heterogeneous targets ({found}). {local_target_hint}") + raise RuntimeError(f"{operation_name}: local JAX GPUs have heterogeneous targets ({found})") resolved = capabilities[0] if resolved not in supported: @@ -216,11 +185,8 @@ def _check_device_compatibility( ) -> None: """Require every local JAX GPU to satisfy an operation's minimum SM.""" - if _DEVICE_COMPATIBILITY_CHECKS_DISABLED: - return - def compatibility_error(reason: str) -> RuntimeError: - return RuntimeError(f"{operation_name} requires SM{minimum_compute_capability}+, {reason}. " f"{_DEVICE_COMPATIBILITY_CHECK_DISABLE_HINT}") + return RuntimeError(f"{operation_name} requires SM{minimum_compute_capability}+, {reason}") try: devices = JaxApiBase._local_gpu_capabilities(operation_name) @@ -499,4 +465,4 @@ def get_jax_callable(self) -> Callable[..., Any]: return self -__all__ = ["JaxApiBase", "JaxTensorDesc", "disable_device_compatibility_checks"] +__all__ = ["JaxApiBase", "JaxTensorDesc"] diff --git a/python/cudnn/gemm_swiglu/jax.py b/python/cudnn/gemm_swiglu/jax.py index 33844eb68..62236f268 100644 --- a/python/cudnn/gemm_swiglu/jax.py +++ b/python/cudnn/gemm_swiglu/jax.py @@ -71,7 +71,6 @@ def __init__( None, SUPPORTED_COMPUTE_CAPABILITIES, "GemmSwigluSm100", - require_local=True, ) self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=self.a_mode) self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=self.b_mode) diff --git a/python/cudnn/jax/__init__.py b/python/cudnn/jax/__init__.py index 92ac00453..4b61bf0b0 100644 --- a/python/cudnn/jax/__init__.py +++ b/python/cudnn/jax/__init__.py @@ -32,7 +32,7 @@ f"Install the optional JAX dependencies with `{_INSTALL_HINT}`." ) -from .._jax import JaxApiBase, JaxTensorDesc, TupleDict, disable_device_compatibility_checks # noqa: E402 +from .._jax import JaxApiBase, JaxTensorDesc, TupleDict # noqa: E402 _OPERATION_EXPORTS = { "RmsNormRhtAmaxSm100": ("..rmsnorm_rht_amax.jax", "RmsNormRhtAmaxSm100"), @@ -118,6 +118,5 @@ def __dir__() -> list[str]: "JaxApiBase", "JaxTensorDesc", "TupleDict", - "disable_device_compatibility_checks", *_OPERATION_EXPORTS, ] diff --git a/test/python/fe_api/test_dsa_indexer_ops.py b/test/python/fe_api/test_dsa_indexer_ops.py index 11223e905..736aad8cb 100644 --- a/test/python/fe_api/test_dsa_indexer_ops.py +++ b/test/python/fe_api/test_dsa_indexer_ops.py @@ -356,13 +356,30 @@ def decorate(resolved): sys.modules[name] = package cls.jax_base = importlib.import_module(f"{cls._PACKAGE}._jax") - cls.jax_base.disable_device_compatibility_checks(True) + + def resolve_compute_capability( + target_compute_capability, + supported_compute_capabilities, + operation_name, + ): + del operation_name + resolved = 100 if target_compute_capability is None else target_compute_capability + if resolved not in supported_compute_capabilities: + raise ValueError(f"unsupported synthetic target SM{resolved}") + return resolved + + cls.target_resolution_patch = mock.patch.object( + cls.jax_base.JaxApiBase, + "_resolve_compute_capability", + staticmethod(resolve_compute_capability), + ) + cls.target_resolution_patch.start() cls.forward = importlib.import_module(f"{dsa_name}.indexer_forward.jax") cls.topk = importlib.import_module(f"{dsa_name}.indexer_top_k.jax") @classmethod def tearDownClass(cls) -> None: - cls.jax_base.disable_device_compatibility_checks(False) + cls.target_resolution_patch.stop() for name in tuple(sys.modules): if name == cls._PACKAGE or name.startswith(f"{cls._PACKAGE}."): sys.modules.pop(name, None) diff --git a/test/python/fe_api/test_jax_api_base.py b/test/python/fe_api/test_jax_api_base.py index 9998b199d..328d83d93 100644 --- a/test/python/fe_api/test_jax_api_base.py +++ b/test/python/fe_api/test_jax_api_base.py @@ -112,12 +112,6 @@ def tearDownClass(cls) -> None: if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): sys.modules.pop(name, None) - def setUp(self) -> None: - self.module.disable_device_compatibility_checks(False) - - def tearDown(self) -> None: - self.module.disable_device_compatibility_checks(False) - def test_checks_all_local_jax_gpu_compute_capabilities(self): jax = types.ModuleType("jax") @@ -143,12 +137,12 @@ def test_rejects_an_incompatible_local_jax_gpu(self): operation_name="TestOp", ) - def test_requires_an_explicit_opt_out_without_a_local_jax_gpu(self): + def test_requires_a_local_jax_gpu(self): jax = types.ModuleType("jax") jax.local_devices = lambda *, backend: () with mock.patch.dict(sys.modules, {"jax": jax}): - with self.assertRaisesRegex(RuntimeError, r"no local JAX GPU.*disable_device_compatibility_checks\(True\)"): + with self.assertRaisesRegex(RuntimeError, r"no local JAX GPU"): self.module.JaxApiBase._check_device_compatibility( minimum_compute_capability=100, operation_name="TestOp", @@ -162,7 +156,7 @@ def fail_discovery(*, backend): jax.local_devices = fail_discovery with mock.patch.dict(sys.modules, {"jax": jax}): - with self.assertRaisesRegex(RuntimeError, r"could not discover a local GPU.*disable_device_compatibility_checks\(True\)"): + with self.assertRaisesRegex(RuntimeError, r"could not discover a local GPU"): self.module.JaxApiBase._check_device_compatibility( minimum_compute_capability=100, operation_name="TestOp", @@ -170,33 +164,12 @@ def fail_discovery(*, backend): jax.local_devices = lambda *, backend: (_Device(0, "unknown"),) with mock.patch.dict(sys.modules, {"jax": jax}): - with self.assertRaisesRegex(RuntimeError, r"invalid compute capability 'unknown'.*disable_device_compatibility_checks\(True\)"): - self.module.JaxApiBase._check_device_compatibility( - minimum_compute_capability=100, - operation_name="TestOp", - ) - - def test_device_compatibility_checks_can_be_disabled_and_reenabled(self): - self.module.disable_device_compatibility_checks(True) - with mock.patch.dict(sys.modules, {"jax": None}): - self.module.JaxApiBase._check_device_compatibility( - minimum_compute_capability=100, - operation_name="TestOp", - ) - - self.module.disable_device_compatibility_checks(False) - jax = types.ModuleType("jax") - jax.local_devices = lambda *, backend: () - with mock.patch.dict(sys.modules, {"jax": jax}): - with self.assertRaises(RuntimeError): + with self.assertRaisesRegex(RuntimeError, r"invalid compute capability 'unknown'"): self.module.JaxApiBase._check_device_compatibility( minimum_compute_capability=100, operation_name="TestOp", ) - with self.assertRaisesRegex(TypeError, "disabled must be a bool"): - self.module.disable_device_compatibility_checks(1) - def test_resolves_exact_local_target_for_multi_arch_operation(self): jax = types.ModuleType("jax") jax.local_devices = lambda *, backend: (_Device(0, "10.3"),) @@ -215,27 +188,22 @@ def test_rejects_heterogeneous_implicit_targets(self): jax.local_devices = lambda *, backend: (_Device(0, "10.0"), _Device(1, "10.3")) with mock.patch.dict(sys.modules, {"jax": jax}): - with self.assertRaisesRegex(RuntimeError, r"heterogeneous targets.*target_compute_capability"): + with self.assertRaisesRegex(RuntimeError, r"heterogeneous targets"): self.module.JaxApiBase._resolve_compute_capability(None, (90, 100, 103, 107), "TestOp") - def test_explicit_cross_compile_target_requires_opt_out(self): + def test_explicit_target_must_match_local_devices(self): jax = types.ModuleType("jax") jax.local_devices = lambda *, backend: (_Device(0, "10.0"),) with mock.patch.dict(sys.modules, {"jax": jax}): with self.assertRaisesRegex(RuntimeError, r"targets SM90.*found CudaDevice\(id=0\) \(SM100\)"): self.module.JaxApiBase._resolve_compute_capability(90, (90, 100, 103, 107), "TestOp") - - self.module.disable_device_compatibility_checks(True) - with mock.patch.dict(sys.modules, {"jax": None}): self.assertEqual( - self.module.JaxApiBase._resolve_compute_capability(90, (90, 100), "TestOp"), - 90, + self.module.JaxApiBase._resolve_compute_capability(100, (90, 100, 103, 107), "TestOp"), + 100, ) def test_rejects_unknown_explicit_target_before_lowering(self): - self.module.disable_device_compatibility_checks(True) - with self.assertRaisesRegex(ValueError, r"no kernel for SM101.*supported targets"): self.module.JaxApiBase._resolve_compute_capability( 101, @@ -243,29 +211,14 @@ def test_rejects_unknown_explicit_target_before_lowering(self): "TestOp", ) - def test_implicit_target_still_needs_a_device_when_checks_are_disabled(self): - self.module.disable_device_compatibility_checks(True) + def test_implicit_target_requires_a_local_device(self): jax = types.ModuleType("jax") jax.local_devices = lambda *, backend: () with mock.patch.dict(sys.modules, {"jax": jax}): - with self.assertRaisesRegex(RuntimeError, r"no local JAX GPU.*target_compute_capability"): + with self.assertRaisesRegex(RuntimeError, r"no local JAX GPU"): self.module.JaxApiBase._resolve_compute_capability(None, (90, 100), "TestOp") - def test_require_local_keeps_device_resolution_strict(self): - self.module.disable_device_compatibility_checks(True) - jax = types.ModuleType("jax") - jax.local_devices = lambda *, backend: () - - with mock.patch.dict(sys.modules, {"jax": jax}): - with self.assertRaisesRegex(RuntimeError, r"no local JAX GPU.*requires a homogeneous supported local JAX GPU"): - self.module.JaxApiBase._resolve_compute_capability( - None, - (100, 103, 107), - "TestOp", - require_local=True, - ) - def test_converts_array_metadata_to_shared_tensor_desc(self): desc = self.module.JaxApiBase._to_tensor_desc( _ArrayMetadata((2, 3, 4), "bfloat16"), diff --git a/test/python/fe_api/test_jax_dsa_backward.py b/test/python/fe_api/test_jax_dsa_backward.py index 1aa29a7f2..1947026a3 100644 --- a/test/python/fe_api/test_jax_dsa_backward.py +++ b/test/python/fe_api/test_jax_dsa_backward.py @@ -42,7 +42,7 @@ def test_jax_dsa_backward_abstract_wrappers(monkeypatch): """Trace every backward signature on a CPU-only JAX installation.""" jax, jnp = _jax_runtime() - from cudnn._jax import JaxApiBase, disable_device_compatibility_checks + from cudnn._jax import JaxApiBase from cudnn.deepseek_sparse_attention.indexer_backward.jax import ( dense_indexer_backward_wrapper, indexer_backward_wrapper, @@ -53,8 +53,23 @@ def abstract_call(_self, _inputs, *, output_descs, **_options): return tuple(jnp.empty(desc.shape, dtype=desc.dtype) for desc in output_descs) monkeypatch.setattr(JaxApiBase, "_call_kernel", abstract_call) - disable_device_compatibility_checks(True) - try: + + def resolve_compute_capability( + target_compute_capability, + supported_compute_capabilities, + operation_name, + ): + del operation_name + resolved = 100 if target_compute_capability is None else target_compute_capability + assert resolved in supported_compute_capabilities + return resolved + + with monkeypatch.context() as target_resolution: + target_resolution.setattr( + JaxApiBase, + "_resolve_compute_capability", + staticmethod(resolve_compute_capability), + ) q = jax.ShapeDtypeStruct((1, 2, 64, 128), jnp.bfloat16) weights = jax.ShapeDtypeStruct((1, 2, 64), jnp.bfloat16) k = jax.ShapeDtypeStruct((1, 128, 128), jnp.bfloat16) @@ -195,8 +210,6 @@ def abstract_call(_self, _inputs, *, output_descs, **_options): assert attention["dq"].shape == attn_q.shape assert attention["dkv"].shape == attn_kv.shape assert attention["d_sink"].shape == sink.shape - finally: - disable_device_compatibility_checks(False) @pytest.mark.L0 diff --git a/test/python/fe_api/test_jax_dsa_indexers.py b/test/python/fe_api/test_jax_dsa_indexers.py index 36ef3034e..b18b7774a 100644 --- a/test/python/fe_api/test_jax_dsa_indexers.py +++ b/test/python/fe_api/test_jax_dsa_indexers.py @@ -17,39 +17,38 @@ def _jax_runtime(): def _supported_gpu(jax, minimum=90): supported_targets = {90, 100, 103, 107} - for device in jax.local_devices(): - if device.platform != "gpu": - continue - reported = getattr(device, "compute_capability", None) - if isinstance(reported, (tuple, list)): - major, minor = int(reported[0]), int(reported[1]) - else: - text = str(reported) - if "." not in text: - continue - major_text, minor_text = text.split(".", 1) - major, minor = int(major_text), int(minor_text) - capability = major * 10 + minor - if capability >= minimum and capability in supported_targets: - return device, capability - pytest.skip(f"A supported JAX SM{minimum}+ GPU is not available") - + try: + devices = tuple(jax.local_devices(backend="gpu")) + except RuntimeError as error: + pytest.skip(f"A JAX GPU is not available: {error}") + if not devices: + pytest.skip("A JAX GPU is not available") -@pytest.fixture -def disabled_device_checks(): - _jax_runtime() - from cudnn.jax import disable_device_compatibility_checks + capabilities = [] + for device in devices: + reported = getattr(device, "compute_capability", None) + try: + if isinstance(reported, (tuple, list)): + major, minor = int(reported[0]), int(reported[1]) + else: + major_text, minor_text = str(reported).split(".", 1) + major, minor = int(major_text), int(minor_text) + except (TypeError, ValueError): + pytest.skip(f"JAX reported an unsupported compute capability {reported!r}") + capabilities.append(major * 10 + minor) - disable_device_compatibility_checks(True) - try: - yield - finally: - disable_device_compatibility_checks(False) + if len(set(capabilities)) != 1: + pytest.skip(f"DSA indexers require homogeneous local GPU targets, found {capabilities}") + capability = capabilities[0] + if capability < minimum or capability not in supported_targets: + pytest.skip(f"A supported JAX SM{minimum}+ GPU is not available") + return devices[0], capability @pytest.mark.L0 -def test_jax_dsa_indexer_abstract_shapes(disabled_device_checks): +def test_jax_dsa_indexer_abstract_shapes(monkeypatch): jax, jnp = _jax_runtime() + from cudnn._jax import JaxApiBase from cudnn.jax import ( compactify_wrapper, indexer_forward_wrapper, @@ -57,6 +56,8 @@ def test_jax_dsa_indexer_abstract_shapes(disabled_device_checks): local_to_global_wrapper, ) + monkeypatch.setattr(JaxApiBase, "_local_gpu_capabilities", staticmethod(lambda _operation_name: ((object(), 100),))) + q = jax.ShapeDtypeStruct((1, 4, 32, 128), jnp.bfloat16) k = jax.ShapeDtypeStruct((1, 5, 1, 128), jnp.bfloat16) w = jax.ShapeDtypeStruct((1, 4, 32), jnp.bfloat16) @@ -174,7 +175,7 @@ def test_jax_dsa_indexer_abstract_shapes(disabled_device_checks): @pytest.mark.L0 -def test_jax_indexer_forward_jit_and_numerics(disabled_device_checks): +def test_jax_indexer_forward_jit_and_numerics(): jax, jnp = _jax_runtime() device, capability = _supported_gpu(jax, 90) from cudnn.jax import indexer_forward_wrapper @@ -251,7 +252,7 @@ def test_jax_indexer_forward_jit_and_numerics(disabled_device_checks): @pytest.mark.L0 -def test_jax_indexer_top_k_jit_and_numerics(disabled_device_checks): +def test_jax_indexer_top_k_jit_and_numerics(): jax, jnp = _jax_runtime() device, capability = _supported_gpu(jax, 90) from cudnn.jax import indexer_top_k_wrapper diff --git a/test/python/fe_api/test_jax_dsa_score_recompute.py b/test/python/fe_api/test_jax_dsa_score_recompute.py index c1db33783..3b9ee967f 100644 --- a/test/python/fe_api/test_jax_dsa_score_recompute.py +++ b/test/python/fe_api/test_jax_dsa_score_recompute.py @@ -45,21 +45,10 @@ def _sm100_device(jax): return devices[0], capabilities[0] -@pytest.fixture -def disabled_device_checks(): - _jax_runtime() - from cudnn.jax import disable_device_compatibility_checks - - disable_device_compatibility_checks(True) - try: - yield - finally: - disable_device_compatibility_checks(False) - - @pytest.mark.L0 -def test_jax_score_recompute_eval_shapes(disabled_device_checks): +def test_jax_score_recompute_eval_shapes(monkeypatch): jax, jnp = _jax_runtime() + from cudnn._jax import JaxApiBase from cudnn.jax import ( dense_attn_score_recompute_wrapper, dense_indexer_score_recompute_wrapper, @@ -67,6 +56,8 @@ def test_jax_score_recompute_eval_shapes(disabled_device_checks): sparse_indexer_score_recompute_wrapper, ) + monkeypatch.setattr(JaxApiBase, "_local_gpu_capabilities", staticmethod(lambda _operation_name: ((object(), 100),))) + sparse_q = jax.ShapeDtypeStruct((1, 4, 32, 128), jnp.bfloat16) sparse_k = jax.ShapeDtypeStruct((1, 128, 128), jnp.bfloat16) sparse_weights = jax.ShapeDtypeStruct((1, 4, 32), jnp.bfloat16) @@ -211,9 +202,8 @@ def run_attention(q, k, lse, offsets): def test_jax_sparse_indexer_score_recompute_jit(): jax, jnp = _jax_runtime() device, target_compute_capability = _sm100_device(jax) - from cudnn.jax import disable_device_compatibility_checks, sparse_indexer_score_recompute_wrapper + from cudnn.jax import sparse_indexer_score_recompute_wrapper - disable_device_compatibility_checks(False) batch, seqlen_q, seqlen_k = 1, 2, 128 num_query_heads, head_dim, topk = 32, 128, 128 q = jax.device_put(jax.random.normal(jax.random.key(10), (batch, seqlen_q, num_query_heads, head_dim), dtype=jnp.bfloat16), device) @@ -251,9 +241,8 @@ def test_jax_sparse_indexer_score_recompute_jit(): def test_jax_sparse_attn_score_recompute_jit(): jax, jnp = _jax_runtime() device, target_compute_capability = _sm100_device(jax) - from cudnn.jax import disable_device_compatibility_checks, sparse_attn_score_recompute_wrapper + from cudnn.jax import sparse_attn_score_recompute_wrapper - disable_device_compatibility_checks(False) batch, seqlen_q, seqlen_k = 1, 2, 128 num_query_heads, head_dim, topk = 32, 128, 128 softmax_scale = head_dim**-0.5 diff --git a/test/python/fe_api/test_jax_gemm_swiglu_contract.py b/test/python/fe_api/test_jax_gemm_swiglu_contract.py index d4be00289..bcff25aaf 100644 --- a/test/python/fe_api/test_jax_gemm_swiglu_contract.py +++ b/test/python/fe_api/test_jax_gemm_swiglu_contract.py @@ -109,9 +109,8 @@ def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): ) @staticmethod - def _resolve_compute_capability(target, supported, operation_name, *, require_local=False): + def _resolve_compute_capability(target, supported, operation_name): del operation_name - assert require_local resolved = 100 if target is None else target if resolved not in supported: raise ValueError(f"unsupported target {resolved}") diff --git a/test/python/fe_api/test_jax_rmsnorm_rht_amax.py b/test/python/fe_api/test_jax_rmsnorm_rht_amax.py index b5d23f7af..0a1e1ba29 100644 --- a/test/python/fe_api/test_jax_rmsnorm_rht_amax.py +++ b/test/python/fe_api/test_jax_rmsnorm_rht_amax.py @@ -22,26 +22,16 @@ def _hadamard_16(jnp): return jnp.asarray(values, dtype=jnp.float32) / math.sqrt(16) -@pytest.fixture -def device_compatibility_checks_disabled(): - _jax_runtime() - from cudnn.jax import disable_device_compatibility_checks - - disable_device_compatibility_checks(True) - try: - yield - finally: - disable_device_compatibility_checks(False) - - @pytest.mark.L0 -def test_jax_rmsnorm_rht_amax_abstract_contract(device_compatibility_checks_disabled): +def test_jax_rmsnorm_rht_amax_abstract_contract(monkeypatch): jax, jnp = _jax_runtime() from cudnn import Op, TensorDesc from cudnn.jax import JaxApiBase, RmsNormRhtAmaxSm100, rmsnorm_rht_amax_sm100 from cudnn.rmsnorm_rht_amax.op import RmsNormRhtAmaxSm100Op + monkeypatch.setattr(JaxApiBase, "_local_gpu_capabilities", staticmethod(lambda _operation_name: ((object(), 100),))) + sample_x = jax.ShapeDtypeStruct((256, 2048), jnp.bfloat16) sample_weight = jax.ShapeDtypeStruct((2048,), jnp.bfloat16) api = RmsNormRhtAmaxSm100(sample_x, sample_weight) @@ -134,9 +124,7 @@ def test_jax_rmsnorm_rht_amax_jit(): if not all_local_devices_are_supported: pytest.skip("RMSNorm + RHT requires every local JAX GPU to be SM100+") - from cudnn.jax import disable_device_compatibility_checks, rmsnorm_rht_amax_sm100 - - disable_device_compatibility_checks(False) + from cudnn.jax import rmsnorm_rht_amax_sm100 device = gpu_devices[0] m, n = 256, 2048 diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_imports.py b/test/python/fe_api/test_rmsnorm_rht_amax_imports.py index ff93d42f2..736677b05 100644 --- a/test/python/fe_api/test_rmsnorm_rht_amax_imports.py +++ b/test/python/fe_api/test_rmsnorm_rht_amax_imports.py @@ -342,10 +342,6 @@ def imports_framework(node, framework: str): any(isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "cutlass_call" for node in ast.walk(jax_adapter)) ) self.assertIn("_INSTALL_HINT", {node.id for node in ast.walk(jax_facade_tree) if isinstance(node, ast.Name)}) - self.assertIn( - "disable_device_compatibility_checks", - {alias.name for node in ast.walk(jax_facade_tree) if isinstance(node, ast.ImportFrom) for alias in node.names}, - ) self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "is_available" for node in ast.walk(jax_facade_tree))) self.assertTrue(any(isinstance(node, ast.Try) for node in jax_facade_tree.body)) self.assertTrue(any(isinstance(node, ast.Import) and any(alias.name == "jax" for alias in node.names) for node in jax_tree.body)) From c35b46db1d938643ed327d93177aca78881cdd60 Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Tue, 7 Jul 2026 22:27:29 -0500 Subject: [PATCH 09/13] Cache JAX cluster occupancy before lowering --- docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md | 4 +- python/cudnn/_jax/api_base.py | 27 ++++++++++++ python/cudnn/gemm_swiglu/jax.py | 10 ++--- test/python/fe_api/test_jax_api_base.py | 43 +++++++++++++++++++ .../fe_api/test_jax_gemm_swiglu_contract.py | 16 +++++-- 5 files changed, 89 insertions(+), 11 deletions(-) diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md index ba0e1891b..068c6c10f 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md @@ -234,8 +234,8 @@ Supported public layouts are: The JAX API requires a homogeneous, supported local SM100-family GPU. It resolves the exact local compute capability for CuTe compilation and queries -that GPU's occupancy during lowering. Device-free AOT and remote compilation -are not currently supported. +that GPU's occupancy once per specialized adapter before CuTe lowering. +Device-free AOT and remote compilation are not currently supported. --- diff --git a/python/cudnn/_jax/api_base.py b/python/cudnn/_jax/api_base.py index 6c4d370da..ff50cf4fc 100644 --- a/python/cudnn/_jax/api_base.py +++ b/python/cudnn/_jax/api_base.py @@ -209,6 +209,33 @@ def compatibility_error(reason: str) -> RuntimeError: found = ", ".join(f"{device} (SM{compute_capability})" for device, compute_capability in incompatible) raise compatibility_error(f"found {found}") + def _get_max_active_clusters( + self, + cluster_size: int, + *, + overlap_margin: int = 0, + ) -> int: + """Return cached occupancy for this adapter and cluster size. + + The hardware query performs its own CuTe compilation, so adapters must + call this before entering ``cutlass.jax.cutlass_call`` lowering. + """ + + cache = getattr(self, "_max_active_clusters_cache", None) + if cache is None: + cache = {} + self._max_active_clusters_cache = cache + + if cluster_size not in cache: + import cutlass + + cache[cluster_size] = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_size) + + max_active_clusters = cache[cluster_size] - overlap_margin + if max_active_clusters <= 0: + raise ValueError("max_active_clusters must be positive after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") + return max_active_clusters + @staticmethod def _to_tensor_desc( value: Any, diff --git a/python/cudnn/gemm_swiglu/jax.py b/python/cudnn/gemm_swiglu/jax.py index 62236f268..f5560032f 100644 --- a/python/cudnn/gemm_swiglu/jax.py +++ b/python/cudnn/gemm_swiglu/jax.py @@ -163,6 +163,10 @@ def __call__(self, a_tensor: Any, b_tensor: Any) -> TupleDict: self.check_support() self._check_tensor_signature(a_tensor, self.a_desc, mode=self.a_mode) self._check_tensor_signature(b_tensor, self.b_desc, mode=self.b_mode) + self._max_active_clusters = self._get_max_active_clusters( + self._op.cluster_shape_mn[0] * self._op.cluster_shape_mn[1], + overlap_margin=self.num_cluster_overlap_margin, + ) ab12_tensor, c_tensor = self._call_kernel( (a_tensor, b_tensor), @@ -207,17 +211,13 @@ def _launch( mma_tiler_mn=self._op.mma_tiler_mn, cluster_shape_mn=self._op.cluster_shape_mn, ) - max_active_clusters = cutlass.utils.HardwareInfo().get_max_active_clusters(self._op.cluster_shape_mn[0] * self._op.cluster_shape_mn[1]) - max_active_clusters -= self.num_cluster_overlap_margin - if max_active_clusters <= 0: - raise ValueError("max_active_clusters must be positive after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") kernel( a, b, ab12, c, cutlass.Float32(self._op.alpha), - max_active_clusters, + self._max_active_clusters, stream, ) diff --git a/test/python/fe_api/test_jax_api_base.py b/test/python/fe_api/test_jax_api_base.py index 328d83d93..6b2de6cb0 100644 --- a/test/python/fe_api/test_jax_api_base.py +++ b/test/python/fe_api/test_jax_api_base.py @@ -219,6 +219,49 @@ def test_implicit_target_requires_a_local_device(self): with self.assertRaisesRegex(RuntimeError, r"no local JAX GPU"): self.module.JaxApiBase._resolve_compute_capability(None, (90, 100), "TestOp") + def test_caches_raw_max_active_clusters_per_instance_and_cluster_size(self): + query_calls = [] + + class HardwareInfo: + def get_max_active_clusters(self, cluster_size): + query_calls.append(cluster_size) + return {2: 8, 4: 12}[cluster_size] + + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass_utils = types.ModuleType("cutlass.utils") + cutlass_utils.HardwareInfo = HardwareInfo + cutlass.utils = cutlass_utils + + class Adapter(self.module.JaxApiBase): + def check_support(self): + return True + + def __call__(self): + return self.check_support() + + def _launch(self, inputs, outputs, workspaces, stream): + raise AssertionError("not called") + + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + "cutlass.utils": cutlass_utils, + }, + ): + api = Adapter() + self.assertEqual(api._get_max_active_clusters(4, overlap_margin=2), 10) + self.assertEqual(api._get_max_active_clusters(4, overlap_margin=3), 9) + self.assertEqual(api._get_max_active_clusters(2, overlap_margin=1), 7) + with self.assertRaisesRegex(ValueError, "max_active_clusters must be positive"): + api._get_max_active_clusters(4, overlap_margin=12) + + other_api = Adapter() + self.assertEqual(other_api._get_max_active_clusters(4), 12) + + self.assertEqual(query_calls, [4, 2, 4]) + def test_converts_array_metadata_to_shared_tensor_desc(self): desc = self.module.JaxApiBase._to_tensor_desc( _ArrayMetadata((2, 3, 4), "bfloat16"), diff --git a/test/python/fe_api/test_jax_gemm_swiglu_contract.py b/test/python/fe_api/test_jax_gemm_swiglu_contract.py index bcff25aaf..d934729c7 100644 --- a/test/python/fe_api/test_jax_gemm_swiglu_contract.py +++ b/test/python/fe_api/test_jax_gemm_swiglu_contract.py @@ -140,6 +140,7 @@ def _to_tensor_spec(desc, *, mode=None, divisibility=None): ) def _call_kernel(self, inputs, **options): + self.captured_max_active_clusters_at_call = getattr(self, "_max_active_clusters", None) self.captured_call = (tuple(inputs), options) return tuple( _Array( @@ -149,6 +150,12 @@ def _call_kernel(self, inputs, **options): for desc, spec in zip(options["output_descs"], options["output_spec"]) ) + def _get_max_active_clusters(self, cluster_size, *, overlap_margin=0): + if not hasattr(self, "captured_active_cluster_queries"): + self.captured_active_cluster_queries = [] + self.captured_active_cluster_queries.append((cluster_size, overlap_margin)) + return 12 - overlap_margin + internal.JaxApiBase = JaxApiBase internal.JaxTensorDesc = JaxTensorDesc internal.TupleDict = result_module.TupleDict @@ -369,13 +376,15 @@ def test_launch_constructs_the_kernel_internally_and_preserves_abi_order(self): alpha=0.25, mma_tiler_mn=(256, 128), ) - api.check_support() + self.assertFalse(hasattr(api, "captured_active_cluster_queries")) + api(sample_a, sample_b) + self.assertEqual(api.captured_active_cluster_queries, [(4, 2)]) + self.assertEqual(api.captured_max_active_clusters_at_call, 10) seen = {} class HardwareInfo: def get_max_active_clusters(self, cluster_size): - seen["cluster_size"] = cluster_size - return 12 + raise AssertionError(f"HardwareInfo queried from _launch for cluster size {cluster_size}") class Kernel: def __init__(self, **configuration): @@ -413,7 +422,6 @@ def __call__(self, *arguments): "cluster_shape_mn": (2, 2), }, ) - self.assertEqual(seen["cluster_size"], 4) self.assertEqual( seen["arguments"], ("A", "B", "AB12", "C", ("Float32", 0.25), 10, stream), From f5817e1cc68404870df0f9c3076d63fbbc66ef34 Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Tue, 7 Jul 2026 22:29:15 -0500 Subject: [PATCH 10/13] Share JAX dtype normalization --- python/cudnn/_jax/datatypes.py | 11 ++++++++++- python/cudnn/gemm_swiglu/jax.py | 17 +++++------------ test/python/fe_api/test_jax_datatypes.py | 8 ++++++++ .../fe_api/test_jax_gemm_swiglu_contract.py | 1 + 4 files changed, 24 insertions(+), 13 deletions(-) diff --git a/python/cudnn/_jax/datatypes.py b/python/cudnn/_jax/datatypes.py index eac294c5e..005a23b19 100644 --- a/python/cudnn/_jax/datatypes.py +++ b/python/cudnn/_jax/datatypes.py @@ -46,6 +46,15 @@ def _make_mappings(): _JAX_TO_CUDNN_DATA_TYPE, _CUDNN_TO_JAX_DATA_TYPE = _make_mappings() +def normalize_jax_dtype(value: Any | None, default: Any, name: str) -> Any: + """Normalize a JAX dtype-like value and identify invalid arguments by name.""" + + try: + return jnp.dtype(default if value is None else value) + except TypeError as error: + raise TypeError(f"{name} must be a JAX dtype, got {value!r}") from error + + def jax_to_cudnn_dtype(dtype: Any) -> data_type: """Return the canonical cuDNN type for a JAX dtype-like value.""" @@ -61,4 +70,4 @@ def cudnn_to_jax_dtype(dtype: data_type) -> Any: raise ValueError(f"Unsupported JAX data type {dtype}") from error -__all__ = ["cudnn_to_jax_dtype", "jax_to_cudnn_dtype"] +__all__ = ["cudnn_to_jax_dtype", "jax_to_cudnn_dtype", "normalize_jax_dtype"] diff --git a/python/cudnn/gemm_swiglu/jax.py b/python/cudnn/gemm_swiglu/jax.py index f5560032f..f47c11dfe 100644 --- a/python/cudnn/gemm_swiglu/jax.py +++ b/python/cudnn/gemm_swiglu/jax.py @@ -16,7 +16,7 @@ from .._cute_compiler import compile_options_for_target from .._dense_gemm import require_gemm_inputs from .._jax import JaxApiBase, JaxTensorDesc, TupleDict -from .._jax.datatypes import jax_to_cudnn_dtype +from .._jax.datatypes import jax_to_cudnn_dtype, normalize_jax_dtype from .._jax.gemm import gemm_a_mode, gemm_b_mode, gemm_output_mode from .._jax.layout import to_public_axes from .op import GemmSwigluSm100Op @@ -24,13 +24,6 @@ SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) -def _normalize_dtype(value: Any | None, default: Any, name: str) -> Any: - try: - return jnp.dtype(default if value is None else value) - except TypeError as error: - raise TypeError(f"{name} must be a JAX dtype, got {value!r}") from error - - class GemmSwigluSm100(JaxApiBase): """JAX callable specialized from dense GEMM input metadata. @@ -74,13 +67,13 @@ def __init__( ) self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=self.a_mode) self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=self.b_mode) - self.acc_dtype = _normalize_dtype(acc_dtype, jnp.float32, "acc_dtype") + self.acc_dtype = normalize_jax_dtype(acc_dtype, jnp.float32, "acc_dtype") if (sample_ab12 is None) != (sample_c is None): raise ValueError("sample_ab12 and sample_c must be provided together") if sample_ab12 is None: - resolved_ab12_dtype = _normalize_dtype(ab12_dtype, jnp.float32, "ab12_dtype") - resolved_c_dtype = _normalize_dtype(c_dtype, jnp.float16, "c_dtype") + resolved_ab12_dtype = normalize_jax_dtype(ab12_dtype, jnp.float32, "ab12_dtype") + resolved_c_dtype = normalize_jax_dtype(c_dtype, jnp.float16, "c_dtype") self.ab12_desc, self.c_desc = self._default_output_descs( resolved_ab12_dtype, resolved_c_dtype, @@ -151,7 +144,7 @@ def _check_requested_output_dtype( ) -> None: if requested is None: return - requested_dtype = _normalize_dtype(requested, requested, name) + requested_dtype = normalize_jax_dtype(requested, requested, name) actual_dtype = jnp.dtype(desc.dtype) if requested_dtype != actual_dtype: raise ValueError(f"{name}={requested_dtype} does not match the explicit sample dtype {actual_dtype}") diff --git a/test/python/fe_api/test_jax_datatypes.py b/test/python/fe_api/test_jax_datatypes.py index 54d39d705..40d820dea 100644 --- a/test/python/fe_api/test_jax_datatypes.py +++ b/test/python/fe_api/test_jax_datatypes.py @@ -74,6 +74,14 @@ def test_supported_types_round_trip(self): self.assertEqual(self.module.cudnn_to_jax_dtype(_DataType.BFLOAT16), "bfloat16") self.assertEqual(self.module.cudnn_to_jax_dtype(_DataType.FLOAT), "float32") + def test_normalizes_dtype_values_and_defaults(self): + self.assertEqual(self.module.normalize_jax_dtype("bfloat16", "float32", "output_dtype"), "bfloat16") + self.assertEqual(self.module.normalize_jax_dtype(None, "float32", "output_dtype"), "float32") + + with mock.patch.object(self.module.jnp, "dtype", side_effect=TypeError("invalid dtype")): + with self.assertRaisesRegex(TypeError, "output_dtype must be a JAX dtype, got 'invalid'"): + self.module.normalize_jax_dtype("invalid", "float32", "output_dtype") + def test_unsupported_types_have_explicit_behavior(self): self.assertEqual(self.module.jax_to_cudnn_dtype("unsupported"), _DataType.NOT_SET) with self.assertRaisesRegex(ValueError, "Unsupported JAX data type"): diff --git a/test/python/fe_api/test_jax_gemm_swiglu_contract.py b/test/python/fe_api/test_jax_gemm_swiglu_contract.py index d934729c7..ec123a664 100644 --- a/test/python/fe_api/test_jax_gemm_swiglu_contract.py +++ b/test/python/fe_api/test_jax_gemm_swiglu_contract.py @@ -162,6 +162,7 @@ def _get_max_active_clusters(self, cluster_size, *, overlap_margin=0): datatypes = types.ModuleType(f"{internal_name}.datatypes") datatypes.jax_to_cudnn_dtype = lambda dtype: _DTYPE_TO_CUDNN.get(dtype, _DataType.NOT_SET) + datatypes.normalize_jax_dtype = lambda value, default, name: default if value is None else value datatypes.cudnn_to_jax_dtype = lambda dtype: _CUDNN_TO_DTYPE[dtype] sys.modules[datatypes.__name__] = datatypes From 5f2238ade87fb9706b192f183f1c825bad77dcc6 Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Tue, 7 Jul 2026 23:15:34 -0500 Subject: [PATCH 11/13] Add JAX support for dense GEMM fusions --- docs/fe-oss-apis/gemm_fusions/gemm_amax.md | 67 +- docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md | 69 +- docs/fe-oss-apis/gemm_fusions/gemm_srelu.md | 68 +- docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md | 68 +- python/cudnn/_dense_gemm.py | 38 ++ python/cudnn/_gemm_relu.py | 220 +++++++ python/cudnn/_jax/api_base.py | 22 +- python/cudnn/_jax/gemm.py | 6 + .../indexer_backward/jax.py | 26 +- .../indexer_forward/jax.py | 12 +- .../indexer_top_k/jax.py | 27 +- .../score_recompute/jax.py | 26 +- .../sparse_attention_backward/jax.py | 12 +- python/cudnn/gemm_amax/__init__.py | 25 +- python/cudnn/gemm_amax/api.py | 317 +++------- python/cudnn/gemm_amax/jax.py | 282 +++++++++ python/cudnn/gemm_amax/op.py | 236 +++++++ python/cudnn/gemm_dsrelu/__init__.py | 25 +- python/cudnn/gemm_dsrelu/api.py | 7 +- ...lockscaled_gemm_persistent_dsrelu_quant.py | 2 - python/cudnn/gemm_dsrelu/jax.py | 325 ++++++++++ python/cudnn/gemm_dsrelu/op.py | 21 + python/cudnn/gemm_srelu/__init__.py | 25 +- python/cudnn/gemm_srelu/api.py | 7 +- ...blockscaled_gemm_persistent_srelu_quant.py | 15 +- python/cudnn/gemm_srelu/jax.py | 319 ++++++++++ python/cudnn/gemm_srelu/op.py | 18 + python/cudnn/gemm_swiglu/__init__.py | 5 +- python/cudnn/gemm_swiglu/jax.py | 304 +++++++-- python/cudnn/gemm_swiglu/op.py | 208 +++++- python/cudnn/jax/__init__.py | 6 + python/cudnn/rmsnorm_rht_amax/jax.py | 34 +- test/python/fe_api/test_gemm_amax_imports.py | 148 +++++ test/python/fe_api/test_gemm_amax_op.py | 295 +++++++++ test/python/fe_api/test_gemm_relu_imports.py | 121 ++++ .../python/fe_api/test_gemm_swiglu_imports.py | 2 + test/python/fe_api/test_gemm_swiglu_op.py | 91 +++ test/python/fe_api/test_jax_api_base.py | 36 +- .../test_jax_dsa_indexer_backward_contract.py | 11 +- .../test_jax_dsa_score_recompute_contract.py | 11 +- ..._dsa_sparse_attention_backward_contract.py | 3 +- test/python/fe_api/test_jax_gemm_amax.py | 194 ++++++ .../fe_api/test_jax_gemm_amax_contract.py | 414 ++++++++++++ test/python/fe_api/test_jax_gemm_layout.py | 2 + test/python/fe_api/test_jax_gemm_relu.py | 275 ++++++++ .../fe_api/test_jax_gemm_relu_contract.py | 597 ++++++++++++++++++ test/python/fe_api/test_jax_gemm_swiglu.py | 51 ++ .../fe_api/test_jax_gemm_swiglu_contract.py | 167 ++++- .../fe_api/test_rmsnorm_rht_amax_imports.py | 9 +- 49 files changed, 4796 insertions(+), 473 deletions(-) create mode 100644 python/cudnn/_gemm_relu.py create mode 100644 python/cudnn/gemm_amax/jax.py create mode 100644 python/cudnn/gemm_amax/op.py create mode 100644 python/cudnn/gemm_dsrelu/jax.py create mode 100644 python/cudnn/gemm_dsrelu/op.py create mode 100644 python/cudnn/gemm_srelu/jax.py create mode 100644 python/cudnn/gemm_srelu/op.py create mode 100644 test/python/fe_api/test_gemm_amax_imports.py create mode 100644 test/python/fe_api/test_gemm_amax_op.py create mode 100644 test/python/fe_api/test_gemm_relu_imports.py create mode 100644 test/python/fe_api/test_jax_gemm_amax.py create mode 100644 test/python/fe_api/test_jax_gemm_amax_contract.py create mode 100644 test/python/fe_api/test_jax_gemm_relu.py create mode 100644 test/python/fe_api/test_jax_gemm_relu_contract.py diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_amax.md b/docs/fe-oss-apis/gemm_fusions/gemm_amax.md index 12f14f11e..2022d9037 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_amax.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_amax.md @@ -103,6 +103,66 @@ op.compile() op.execute(a, b, sfa, sfb, c, amax, current_stream=None) ``` +### JAX API + +The JAX wrapper is decorated with `jax.jit`. GEMM arrays are compact row-major +arrays whose public axis order is selected explicitly: + +```python +import jax +import jax.numpy as jnp + +from cudnn.jax import gemm_amax_wrapper_sm100 + +# A: (L, M, K), B: (L, N, K) +# SFA/SFB: (L, ceil(rows/128), K_tiles, 32, 4, 4) +result = gemm_amax_wrapper_sm100( + a, + b, + sfa, + sfb, + a_layout="LMK", + b_layout="LNK", + c_layout="LMN", + c_dtype=jnp.float32, + acc_dtype=jnp.float32, + sf_vec_size=32, +) + +c = result["c_tensor"] # (L, M, N) +amax = result["amax_tensor"] # (1, 1, 1) +``` + +The scale arrays use ordinary row-major public storage with shape +`(L, ceil(rows/128), ceil(ceil(K/sf_vec_size)/4), 32, 4, 4)`. The adapter +maps that public representation to the kernel's canonical packed descriptor +`(32, 4, row_tiles, 4, k_tiles, L)`. + +The advanced class specializes from shape/dtype metadata and optionally accepts +explicit output exemplars: + +```python +from cudnn.jax import GemmAmaxSm100 + +operation = GemmAmaxSm100( + jax.ShapeDtypeStruct(a.shape, a.dtype), + jax.ShapeDtypeStruct(b.shape, b.dtype), + jax.ShapeDtypeStruct(sfa.shape, sfa.dtype), + jax.ShapeDtypeStruct(sfb.shape, sfb.dtype), +) +result = operation(a, b, sfa, sfb) +``` + +Supported public layouts are: + +- A: `LMK` (K-major) or `LKM` (M-major) +- B: `LNK` (K-major) or `LKN` (N-major) +- C: `LMN` (N-major) or `LNM` (M-major) + +JAX uses native FP4/FP8 dtypes. The Torch-only `uint8` packed-FP4 and `int8` +E8M0 storage aliases are not accepted by the JAX adapter. The JAX API requires +a homogeneous local GPU with an exact supported target in the SM100 family. + --- ## Parameters @@ -131,7 +191,7 @@ op.execute(a, b, sfa, sfb, c, amax, current_stream=None) - Dtype: `float32` ### Common parameters -- `acc_dtype: torch.dtype` +- `acc_dtype`: a Torch dtype or JAX dtype-like value - Accumulator dtype. Default: `torch.float32` (only supported value) - `mma_tiler_mn: Tuple[int, int]` - Kernel tile size `(TILE_M, TILE_N)`. Default: `(128, 128)` @@ -143,12 +203,17 @@ op.execute(a, b, sfa, sfb, c, amax, current_stream=None) - `sf_vec_size: int` - Size of K-group per scale factor: `{16, 32}`. Default: `32` - CUDA stream (`current_stream` in class API, `stream` in wrapper) + JAX receives its stream from lowering and does not expose a stream argument. ### Wrapper-specific parameters: `gemm_amax_wrapper_sm100` - `a_tensor`, `b_tensor`, `sfa_tensor`, `sfb_tensor`: see Input/Output tensors - `c_major: str`: see Input/Output tensors. Default: `"n"` - `c_dtype: torch.dtype`: see Input/Output tensors. Default: `torch.float32` +For JAX, `a_layout`, `b_layout`, and `c_layout` select the public axis orders +listed above. `c_dtype` defaults to `jnp.float32`; output allocation and amax +initialization are owned by the JAX wrapper. + ### Wrapper return values Returns a `TupleDict` with keys: diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md b/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md index a18656744..915506cdf 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md @@ -7,7 +7,7 @@ **Block-scaled GEMM + dsReLU backward fusion**: A persistent, batched dense GEMM on NVIDIA Blackwell GPUs (SM100+) that supports block-scaled FP4 and FP8 inputs and produces both the backward output `D` and the probability gradient `dprob` in a single kernel launch. - **Inputs**: quantized `A` and `B`, the forward/intermediate tensor `C`, scale-factor tensors `SFA` and `SFB`, and a per-row probability tensor `prob` -- **Outputs**: backward output `D`, probability gradient `dprob`, and optional output scale factors `SFD` / `Amax` +- **Outputs**: backward output `D`, probability gradient `dprob`, and optional `Amax` ### Shapes @@ -22,7 +22,6 @@ - **Outputs** - `D`: shape `(M, N, L)` - `dprob`: shape `(M, 1, L)` - - `SFD`: shape `(32, 4, ceil_div(M, 128), 4, ceil_div(ceil_div(N, sf_vec_size), 4), L)` when `D` is FP8 - `Amax`: shape `(1,)` when FP4 input is written to fp16/bf16/fp32 output `L` is the batch dimension. @@ -43,7 +42,9 @@ $$ \mathrm{dprob}[m, 0, l] = \sum_n C[m, n, l] \cdot \mathrm{relu}(G[m, n, l])^2 $$ -As with the forward `srelu` kernel, FP8 `D` also emits `SFD` using the provided `norm_const_tensor`, and FP4 input written to fp16/bf16/fp32 `D` can emit `Amax`. +As with the forward `srelu` kernel, FP4 input written to fp16/bf16/fp32 `D` +can emit `Amax`. The FP8 `D`/SFD epilogue is reserved by the API but is not +implemented by the current kernel. ### Diagram @@ -63,16 +64,18 @@ A (MxKxL), SFA B (NxKxL), SFB | +--> dprob (Mx1xL) | - +-----------+-----------+ - | | - v v - SFD Amax + v + Amax ``` --- ## API Usage +The existing APIs in this section use Torch tensors. The JAX API is described +separately below because its public arrays are row-major and name their logical +axis order explicitly. + ### High-level wrapper ```python @@ -142,6 +145,47 @@ op.execute( ) ``` +### JAX API + +```python +import jax.numpy as jnp +from cudnn.jax import gemm_dsrelu_wrapper_sm100 + +result = gemm_dsrelu_wrapper_sm100( + a, # (L, M, K), float4_e2m1fn or FP8 + b, # (L, N, K), same dtype as A + c, # (L, M, N) + sfa, # (L, ceil(M/128), rest_k, 32, 4, 4) + sfb, # (L, ceil(N/128), rest_k, 32, 4, 4) + prob, # (L, 1, M), float32 + d_dtype=jnp.bfloat16, + sf_vec_size=16, # use 32 for MXFP8 + E8M0 scales +) +``` + +`GemmDsreluSm100` provides the corresponding class API and accepts optional +`sample_d` and `sample_dprob` output exemplars. The high-level wrapper is JIT +compiled and infers its outputs. + +JAX arrays use compact row-major storage: + +- `a_layout="LMK"` and `"LKM"` map public A shapes `(L,M,K)` and `(L,K,M)` + to kernel axes `(M,K,L)`. +- `b_layout="LNK"` and `"LKN"` map public B shapes `(L,N,K)` and `(L,K,N)` + to kernel axes `(N,K,L)`. +- `d_layout="LMN"` consumes C and produces D as `(L,M,N)`; `"LNM"` uses + `(L,N,M)`. +- SFA/SFB use public shape `(L, tiles, rest, 32, 4, 4)`. The adapter maps + this row-major representation to the packed six-dimensional kernel ABI. +- `prob` and the inferred `dprob_tensor` use public shape `(L,1,M)`. + +Here `rest_k = ceil(ceil(K / sf_vec_size) / 4)`. + +Native JAX `float4_e2m1fn` inputs return an initialized `amax_tensor` with +shape `(1,)`. JAX currently rejects FP8 `D`, `norm_const_tensor`, and SFD +generation because the kernel's SFD epilogue is not implemented. `D` must use +`float16`, `bfloat16`, or `float32`. + --- ## Parameters @@ -169,14 +213,14 @@ op.execute( - Dtype: `float32` - Output tensor **D**: `result["d_tensor"]` (wrapper) or `sample_d` / `d_tensor` (class) - Shape: `(M, N, L)` - - Dtype: `{float16, bfloat16, float32, float8_e4m3fn, float8_e5m2}` + - Functional dtypes: `{float16, bfloat16, float32}` - Output tensor **dprob**: `result["dprob_tensor"]` (wrapper) or `sample_dprob` / `dprob_tensor` (class) - Shape: `(M, 1, L)` - Dtype: `float32` - Output tensor **SFD**: `result["sfd_tensor"]` (wrapper) or `sample_sfd` / `sfd_tensor` (class) - Shape: `(32, 4, ceil_div(M, 128), 4, ceil_div(ceil_div(N, sf_vec_size), 4), L)` - Dtype: Must match `SFA` - - Required when `D` is FP8 + - Reserved for FP8 `D`; generation is not implemented - Output tensor **Amax**: `result["amax_tensor"]` (wrapper) or `sample_amax` / `amax_tensor` (class) - Shape: `(1,)` - Dtype: `float32` @@ -184,7 +228,7 @@ op.execute( - Input tensor **Norm Const**: `norm_const_tensor` (wrapper) or `sample_norm_const` / `norm_const_tensor` (class) - Shape: `(1,)` - Dtype: `float32` - - Required when `D` is FP8 + - Reserved for FP8 `D`; currently unsupported ### Common parameters @@ -237,11 +281,13 @@ Tuple unpacking order is: `(d_tensor, dprob_tensor, amax_tensor, sfd_tensor)`. - `sf_vec_size == 32` is unsupported with `sf_dtype == float8_e4m3fn` - FP8 input requires `sf_vec_size == 32` - FP4 input with FP8 `D` is unsupported -- FP8 `D` requires both `SFD` and `norm_const_tensor` +- FP8 `D` and SFD generation are not implemented ### Environment - Requires CUDA with SM100+ compute capability +- The JAX API requires a homogeneous local SM100-family GPU and resolves + occupancy before CuTe lowering. Device-free compilation is not supported. --- @@ -251,3 +297,4 @@ For end-to-end usage and regression coverage, see: - `test/python/fe_api/test_gemm_dsrelu.py` - `test/python/fe_api/test_gemm_dsrelu_utils.py` +- `test/python/fe_api/test_jax_gemm_relu.py` diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md b/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md index 963165e3f..242f10117 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md @@ -7,7 +7,7 @@ **Block-scaled GEMM + sReLU fusion**: A persistent, batched dense GEMM on NVIDIA Blackwell GPUs (SM100+) that supports block-scaled FP4 and FP8 inputs and produces both the full GEMM result `C` and a probability-gated squared-ReLU output `D` in a single kernel launch. - **Inputs**: quantized `A` and `B`, scale-factor tensors `SFA` and `SFB`, and a per-row probability tensor `prob` -- **Outputs**: full GEMM result `C`, squared-ReLU output `D`, and optional output scale factors `SFD` / `Amax` +- **Outputs**: full GEMM result `C`, squared-ReLU output `D`, and optional `Amax` ### Shapes @@ -21,7 +21,6 @@ - **Outputs** - `C`: shape `(M, N, L)` - `D`: shape `(M, N, L)` - - `SFD`: shape `(32, 4, ceil_div(M, 128), 4, ceil_div(ceil_div(N, sf_vec_size), 4), L)` when `D` is FP8 - `Amax`: shape `(1,)` when FP4 input is written to fp16/bf16/fp32 output `L` is the batch dimension. @@ -38,7 +37,9 @@ $$ D[m, n, l] = \mathrm{prob}[m, 0, l] \cdot \mathrm{relu}(C[m, n, l])^2 $$ -When `D` is FP8, the kernel also emits output scale factors `SFD` using the provided `norm_const_tensor`. When FP4 input is written to a higher-precision `D`, the kernel can also emit `Amax`. +When FP4 input is written to a higher-precision `D`, the kernel can also emit +`Amax`. The FP8 `D`/SFD epilogue is reserved by the API but is not implemented +by the current kernel. ### Diagram @@ -56,16 +57,18 @@ A (MxKxL), SFA B (NxKxL), SFB v D (MxNxL) | - +-----------+-----------+ - | | - v v - SFD Amax + v + Amax ``` --- ## API Usage +The existing APIs in this section use Torch tensors. The JAX API is described +separately below because its public arrays are row-major and name their logical +axis order explicitly. + ### High-level wrapper ```python @@ -133,6 +136,46 @@ op.execute( ) ``` +### JAX API + +```python +import jax.numpy as jnp +from cudnn.jax import gemm_srelu_wrapper_sm100 + +result = gemm_srelu_wrapper_sm100( + a, # (L, M, K), float4_e2m1fn or FP8 + b, # (L, N, K), same dtype as A + sfa, # (L, ceil(M/128), rest_k, 32, 4, 4) + sfb, # (L, ceil(N/128), rest_k, 32, 4, 4) + prob, # (L, 1, M), float32 + c_dtype=jnp.bfloat16, + d_dtype=jnp.bfloat16, + sf_vec_size=16, # use 32 for MXFP8 + E8M0 scales +) +``` + +`GemmSreluSm100` provides the corresponding class API and accepts optional +`sample_c` and `sample_d` output exemplars. The high-level wrapper is JIT +compiled and infers its outputs. + +JAX arrays use compact row-major storage: + +- `a_layout="LMK"` and `"LKM"` map public A shapes `(L,M,K)` and `(L,K,M)` + to kernel axes `(M,K,L)`. +- `b_layout="LNK"` and `"LKN"` map public B shapes `(L,N,K)` and `(L,K,N)` + to kernel axes `(N,K,L)`. +- `c_layout="LMN"` produces `(L,M,N)`; `"LNM"` produces `(L,N,M)`. +- SFA/SFB use public shape `(L, tiles, rest, 32, 4, 4)`. The adapter maps + this row-major representation to the packed six-dimensional kernel ABI. +- `prob` uses public shape `(L,1,M)`. + +Here `rest_k = ceil(ceil(K / sf_vec_size) / 4)`. + +Native JAX `float4_e2m1fn` inputs return an initialized `amax_tensor` with +shape `(1,)`. JAX currently rejects FP8 `D`, `norm_const_tensor`, and SFD +generation because the kernel's SFD epilogue is not implemented. `D` must use +`float16`, `bfloat16`, or `float32`. + --- ## Parameters @@ -160,11 +203,11 @@ op.execute( - Dtype: `{float16, bfloat16, float32, float8_e4m3fn, float8_e5m2}` - Output tensor **D**: `result["d_tensor"]` (wrapper) or `sample_d` / `d_tensor` (class) - Shape: `(M, N, L)` - - Dtype: `{float16, bfloat16, float32, float8_e4m3fn, float8_e5m2}` + - Functional dtypes: `{float16, bfloat16, float32}` - Output tensor **SFD**: `result["sfd_tensor"]` (wrapper) or `sample_sfd` / `sfd_tensor` (class) - Shape: `(32, 4, ceil_div(M, 128), 4, ceil_div(ceil_div(N, sf_vec_size), 4), L)` - Dtype: Must match `SFA` - - Required when `D` is FP8 + - Reserved for FP8 `D`; generation is not implemented - Output tensor **Amax**: `result["amax_tensor"]` (wrapper) or `sample_amax` / `amax_tensor` (class) - Shape: `(1,)` - Dtype: `float32` @@ -172,7 +215,7 @@ op.execute( - Input tensor **Norm Const**: `norm_const_tensor` (wrapper) or `sample_norm_const` / `norm_const_tensor` (class) - Shape: `(1,)` - Dtype: `float32` - - Required when `D` is FP8 + - Reserved for FP8 `D`; currently unsupported ### Common parameters @@ -222,11 +265,13 @@ Tuple unpacking order is: `(c_tensor, d_tensor, amax_tensor, sfd_tensor)`. - `sf_vec_size == 32` is unsupported with `sf_dtype == float8_e4m3fn` - FP8 input requires `sf_vec_size == 32` - FP4 input with FP8 `D` is unsupported -- FP8 `D` requires both `SFD` and `norm_const_tensor` +- FP8 `D` and SFD generation are not implemented ### Environment - Requires CUDA with SM100+ compute capability +- The JAX API requires a homogeneous local SM100-family GPU and resolves + occupancy before CuTe lowering. Device-free compilation is not supported. --- @@ -236,3 +281,4 @@ For end-to-end usage and regression coverage, see: - `test/python/fe_api/test_gemm_srelu.py` - `test/python/fe_api/test_gemm_srelu_utils.py` +- `test/python/fe_api/test_jax_gemm_relu.py` diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md index 068c6c10f..6c5976232 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md @@ -10,8 +10,9 @@ This API supports two modes: 1. **Standard mode**: High-precision GEMM with SwiGLU epilogue 2. **Quantized mode** (block-scaled): Low-precision GEMM using block scaling supporting FP4 and FP8 data types -The Torch API supports both modes. The initial JAX API supports standard mode; -the block-scaled JAX surface will be added with the shared scale-factor ABI. +The Torch and JAX APIs support both modes. JAX scale-factor arrays use an +ordinary row-major public shape and are mapped to the kernel's packed ABI by +the adapter. ### Shapes @@ -226,6 +227,42 @@ operation = GemmSwigluSm100( result = operation(a, b) ``` +### JAX API (Block-Scaled Mode) + +Providing both `sfa_tensor` and `sfb_tensor` selects the block-scaled kernel. +The output dtypes remain explicit because the wrapper defaults are unchanged +from standard mode; MXFP8 callers must request a supported AB12 dtype. + +```python +from cudnn.jax import gemm_swiglu_wrapper_sm100 + +# A: (L, M, K), B: (L, N, K) +# SFA/SFB: (L, row_tiles, k_tiles, 32, 4, 4) +result = gemm_swiglu_wrapper_sm100( + a_fp8, + b_fp8, + sfa_tensor=sfa_e8m0, + sfb_tensor=sfb_e8m0, + sf_vec_size=32, + ab12_dtype=jnp.bfloat16, + c_dtype=jnp.bfloat16, + a_layout="LMK", + b_layout="LNK", + c_layout="LMN", +) +``` + +For a scale tensor associated with `rows` and reduction extent `K`, the public +JAX shape is +`(L, ceil(rows/128), ceil(ceil(K/sf_vec_size)/4), 32, 4, 4)`. +Its fixed mode `(3, 4, 1, 5, 2, 0)` binds that row-major array to the canonical +kernel shape `(32, 4, ceil(rows/128), 4, k_tiles, L)`. + +The current JAX surface supports `C` in `float16`, `bfloat16`, or `float32`. +FP8 `C`, SFC generation, and `norm_const_tensor` are disabled because none of +the currently supported input configurations can use the kernel's FP8-C path. +`result["sfc_tensor"]` is therefore `None`. + Supported public layouts are: - A: `LMK` (K-major) or `LKM` (M-major) @@ -265,7 +302,7 @@ also the Torch API order. JAX arrays use the public axis orders selected by - Quantized mode: Must be `n`-major for FP4 outputs - Dtype (`ab12_dtype`, provided as `ab12_dtype` argument for wrapper): - Standard mode: `{float32, float16, bfloat16}` if `acc_dtype == float32`, `{float16, bfloat16}` if `acc_dtype == float16` - - Quantized mode: `{float32, float16, bfloat16, float8_e4m3fn, float8_e5m2}` + - Quantized mode: `{float32, float16, bfloat16}` - Output tensor **C**: `result["c_tensor"]` (wrapper) or `sample_c`, `c_tensor` (class) - Shape: `(M, N/2, L)` - Stride: `(1, M, M·N/2)` for `m`-major or `(N/2, 1, M·N/2)` for `n`-major. Must match with `AB12` @@ -279,10 +316,12 @@ also the Torch API order. JAX arrays use the public axis orders selected by - Input tensor **SFB** (B scale factor): `sfb_tensor` (wrapper) or `sample_sfb`, `sfb_tensor` (class) - Shape: `(32, 4, ceil(N/128), 4, ceil(ceil(K/sf_vec_size)/4), L)` - Dtype: Must match `SFA` + - JAX public shape for SFA/SFB is + `(L, row_tiles, k_tiles, 32, 4, 4)` as described above. - Output tensor **SFC** (C scale factor, **Optional**): `result["sfc_tensor"]` (wrapper) or `sample_sfc`, `sfc_tensor` (class) - Shape: `(32, 4, ceil(M/128), 4, ceil(ceil((N/2)/sf_vec_size)/4), L)` - Dtype: Must match `SFA` - - **Required when**: `c_dtype ∈ {float8_e4m3fn, float8_e5m2}` + - Reserved for FP8 `C`; current supported input combinations reject FP8 `C` - Output tensor **AMAX** (**Optional**): `result["amax_tensor"]` (wrapper) or `sample_amax`, `amax_tensor` (class) - Shape: `(1,)` - Dtype: `float32` @@ -290,7 +329,7 @@ also the Torch API order. JAX arrays use the public axis orders selected by - Input tensor **Norm Const** (**Optional**): `norm_const_tensor` (wrapper) or `sample_norm_const`, `norm_const_tensor` (class) - Shape: `(1,)` - Dtype: `float32` - - **Required when**: `c_dtype ∈ {float8_e4m3fn, float8_e5m2}` + - Reserved for FP8 `C`; currently unsupported ### Common parameters @@ -333,7 +372,12 @@ the exact local SM100-family compilation target while tracing. - `a_tensor`, `b_tensor`: see Input/Output tensors - Torch: `c_major` selects `"m"` or `"n"`; `ab12_dtype` defaults to `torch.float32`; `c_dtype` defaults to `torch.float16`. - JAX: `a_layout`, `b_layout`, and `c_layout` select public axis order; `ab12_dtype` defaults to `jnp.float32`; `c_dtype` defaults to `jnp.float16`. -- Torch quantized mode: `sfa_tensor`, `sfb_tensor`, `norm_const_tensor`, `sf_vec_size`, `vector_f32`, and `ab12_stages` use the quantization rules above. +- Torch quantized mode: `sfa_tensor`, `sfb_tensor`, `norm_const_tensor`, + `sf_vec_size`, `vector_f32`, and `ab12_stages` use the quantization rules + above. +- JAX block-scaled mode: `sfa_tensor`, `sfb_tensor`, `sf_vec_size`, + `vector_f32`, and `ab12_stages` are supported. `norm_const_tensor` is + rejected while FP8 `C` remains disabled. ### Wrapper return values @@ -355,7 +399,10 @@ Tuple unpacking order is always: #### `GemmSwigluSm100` (constructor) - `sample_a`, `sample_b`, `sample_ab12`, `sample_c` – see Input/Output tensors. JAX requires only the input samples and accepts the output samples optionally. -- `sample_sfa`, `sample_sfb`, `sample_sfc`, `sample_amax`, `sample_norm_const` – Torch quantized mode only. +- `sample_sfa`, `sample_sfb`, `sample_sfc`, `sample_amax`, + `sample_norm_const` – quantized-mode samples. JAX supports SFA/SFB and can + infer AMAX; it rejects SFC and norm-constant samples while FP8 `C` remains + disabled. #### `GemmSwigluSm100.execute` @@ -391,11 +438,14 @@ The quantized kernel supports the following configurations: | **MXFP4** | `float4_e2m1fn_x2` or `uint8` | `float8_e4m3fn` | 16 | NVF4 variant | | **MXFP8** | `float8_e4m3fn` or `float8_e5m2` | `float8_e8m0fnu` | 32 | Standard MX FP8 | +JAX uses its native logical `float4_e2m1fn` dtype for FP4 when that dtype is +available in the installed JAX runtime. The raw `uint8` packed-FP4 alias is a +Torch compatibility path and is not reinterpreted by the JAX adapter. + Additional constraints: - `acc_dtype` must be `float32` - Not compatible with FP8 c_dtype. BF16 `c_dtype` is expected. -- For MXFP8 inputs, ab12_dtype` should be float16 or bfloat16. -- When `c_dtype ∈ {float8_e4m3fn, float8_e5m2}`: `sfc_tensor` and `norm_const_tensor` are required +- For MXFP8 inputs, `ab12_dtype` should be float16 or bfloat16. - When `ab_dtype` is FP4 and `c_dtype == bfloat16`: `amax_tensor` is required - `c_dtype` and `ab12_dtype` cannot both be `float32` diff --git a/python/cudnn/_dense_gemm.py b/python/cudnn/_dense_gemm.py index 2287ef1a6..3326d1b01 100644 --- a/python/cudnn/_dense_gemm.py +++ b/python/cudnn/_dense_gemm.py @@ -13,6 +13,7 @@ _STANDARD_MMA_M = (128, 256) _STANDARD_MMA_N = tuple(range(32, 257, 32)) +BLOCK_SCALE_STRIDE_ORDER = (3, 1, 0, 4, 2, 5) _DATA_TYPE_BITS_BY_NAME = { "BOOLEAN": 8, @@ -87,6 +88,40 @@ def require_gemm_inputs(a: TensorDesc, b: TensorDesc) -> tuple[int, int, int, in return m, n, k, l +def block_scale_shape( + rows: int, + k: int, + batch: int, + sf_vec_size: int, +) -> tuple[int, int, int, int, int, int]: + """Return the canonical packed scale-factor shape for a dense GEMM.""" + + if isinstance(sf_vec_size, bool): + raise TypeError(f"sf_vec_size must be an integer, got {sf_vec_size!r}") + try: + sf_vec_size = index(sf_vec_size) + except TypeError as error: + raise TypeError(f"sf_vec_size must be an integer, got {sf_vec_size!r}") from error + if sf_vec_size <= 0: + raise ValueError(f"sf_vec_size must be positive, got {sf_vec_size}") + row_tiles = (rows + 127) // 128 + scale_k = (k + sf_vec_size - 1) // sf_vec_size + k_tiles = (scale_k + 3) // 4 + return (32, 4, row_tiles, 4, k_tiles, batch) + + +def require_block_scale_layout(tensor: TensorDesc, label: str) -> None: + """Require the packed scale-factor layout used by dense SM100 GEMMs.""" + + _require_tensor_desc(tensor, label) + if not tensor.is_compact(BLOCK_SCALE_STRIDE_ORDER): + raise ValueError( + f"{label} must use the packed block-scale layout with stride order " + f"{BLOCK_SCALE_STRIDE_ORDER}, got stride {tensor.stride} and " + f"stride order {tensor.stride_order}" + ) + + def require_tensor_shape( tensor: TensorDesc, expected: tuple[int, ...], @@ -205,8 +240,11 @@ def require_cluster_shape( __all__ = [ + "BLOCK_SCALE_STRIDE_ORDER", + "block_scale_shape", "data_type_bits", "require_16_byte_alignment", + "require_block_scale_layout", "require_cluster_shape", "require_compact_major", "require_gemm_inputs", diff --git a/python/cudnn/_gemm_relu.py b/python/cudnn/_gemm_relu.py new file mode 100644 index 000000000..2720521a2 --- /dev/null +++ b/python/cudnn/_gemm_relu.py @@ -0,0 +1,220 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral validation shared by dense sReLU fusions.""" + +from __future__ import annotations + +from typing import Any + +from . import data_type +from ._dense_gemm import ( + block_scale_shape, + require_16_byte_alignment, + require_block_scale_layout, + require_cluster_shape, + require_compact_major, + require_gemm_inputs, + require_mma_tiler, + require_tensor_shape, +) +from ._op import Op +from ._tensor_desc import TensorDesc + +_AB_DTYPES = frozenset( + { + data_type.FP4_E2M1, + data_type.UINT8, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + } +) +_FP4_DTYPES = frozenset({data_type.FP4_E2M1, data_type.UINT8}) +_FP8_DTYPES = frozenset({data_type.FP8_E4M3, data_type.FP8_E5M2}) +_WIDE_OUTPUT_DTYPES = frozenset({data_type.HALF, data_type.BFLOAT16, data_type.FLOAT}) +_C_DTYPES = _WIDE_OUTPUT_DTYPES | _FP8_DTYPES + + +class BlockScaledGemmReluSm100Op(Op): + """Logical tensor signature shared by sReLU forward and backward.""" + + def __init__( + self, + *, + a: TensorDesc[Any], + b: TensorDesc[Any], + c: TensorDesc[Any], + d: TensorDesc[Any], + sfa: TensorDesc[Any], + sfb: TensorDesc[Any], + prob: TensorDesc[Any], + dprob: TensorDesc[Any] | None = None, + sfd: TensorDesc[Any] | None = None, + amax: TensorDesc[Any] | None = None, + norm_const: TensorDesc[Any] | None = None, + alpha: float = 1.0, + acc_dtype: data_type = data_type.FLOAT, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 16, + ) -> None: + for name, desc in ( + ("a", a), + ("b", b), + ("c", c), + ("d", d), + ("sfa", sfa), + ("sfb", sfb), + ("prob", prob), + ): + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + for name, desc in ( + ("dprob", dprob), + ("sfd", sfd), + ("amax", amax), + ("norm_const", norm_const), + ): + if desc is not None and not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc or None, got {type(desc).__name__}") + if not isinstance(acc_dtype, data_type): + raise TypeError(f"acc_dtype must be a cudnn.data_type, got {type(acc_dtype).__name__}") + if isinstance(sf_vec_size, bool) or not isinstance(sf_vec_size, int): + raise TypeError(f"sf_vec_size must be an int, got {type(sf_vec_size).__name__}") + + self.a = a + self.b = b + self.c = c + self.d = d + self.sfa = sfa + self.sfb = sfb + self.prob = prob + self.dprob = dprob + self.sfd = sfd + self.amax = amax + self.norm_const = norm_const + self.alpha = alpha + self.acc_dtype = acc_dtype + self.mma_tiler_mn = tuple(mma_tiler_mn) + self.cluster_shape_mn = None if cluster_shape_mn is None else tuple(cluster_shape_mn) + self.sf_vec_size = sf_vec_size + + self.m: int | None = None + self.n: int | None = None + self.k: int | None = None + self.l: int | None = None + self.a_major: str | None = None + self.b_major: str | None = None + self.output_major: str | None = None + + def check_support(self) -> bool: + """Validate shapes, dtypes, compact layouts, and static tiling.""" + + m, n, k, batch = require_gemm_inputs(self.a, self.b) + require_tensor_shape(self.c, (m, n, batch), label="C") + require_tensor_shape(self.d, (m, n, batch), label="D") + require_tensor_shape(self.prob, (m, 1, batch), label="prob") + if self.dprob is not None: + require_tensor_shape(self.dprob, (m, 1, batch), label="dprob") + + if self.sf_vec_size not in (16, 32): + raise ValueError(f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}") + require_tensor_shape( + self.sfa, + block_scale_shape(m, k, batch, self.sf_vec_size), + label="SFA", + ) + require_tensor_shape( + self.sfb, + block_scale_shape(n, k, batch, self.sf_vec_size), + label="SFB", + ) + for label, desc in (("SFA", self.sfa), ("SFB", self.sfb)): + require_block_scale_layout(desc, label) + + if self.sfd is not None: + raise NotImplementedError("SFD generation is not implemented by the current sReLU kernels") + if self.amax is not None: + require_tensor_shape(self.amax, (1,), label="amax") + if self.norm_const is not None: + raise ValueError("norm_const is only used with FP8 D output, which is not implemented") + + a_storage_dtype = self.a.cudnn_dtype + b_storage_dtype = self.b.cudnn_dtype + if b_storage_dtype != a_storage_dtype: + raise ValueError(f"A and B must have the same dtype, got {self.a.dtype} and {self.b.dtype}") + ab_dtype = data_type.FP4_E2M1 if a_storage_dtype == data_type.UINT8 else a_storage_dtype + if ab_dtype not in _AB_DTYPES: + raise ValueError(f"A has unsupported block-scaled GEMM dtype {self.a.dtype}") + + sf_dtype = self.sfa.cudnn_dtype + if self.sfb.cudnn_dtype != sf_dtype: + raise ValueError(f"SFA and SFB must have the same dtype, got {self.sfa.dtype} and {self.sfb.dtype}") + if sf_dtype not in {data_type.FP8_E8M0, data_type.FP8_E4M3}: + raise ValueError(f"SFA has unsupported scale-factor dtype {self.sfa.dtype}") + if ab_dtype in _FP8_DTYPES and (sf_dtype, self.sf_vec_size) != (data_type.FP8_E8M0, 32): + raise ValueError("FP8 inputs require FP8_E8M0 scales with sf_vec_size=32") + if ab_dtype in _FP4_DTYPES and self.sf_vec_size == 32 and sf_dtype != data_type.FP8_E8M0: + raise ValueError("FP4 inputs with sf_vec_size=32 require FP8_E8M0 scales") + + if self.c.cudnn_dtype not in _C_DTYPES: + raise ValueError(f"C has unsupported dtype {self.c.dtype}") + if self.d.cudnn_dtype in _FP8_DTYPES: + raise NotImplementedError("FP8 D output is unavailable because SFD generation is not implemented") + if self.d.cudnn_dtype not in _WIDE_OUTPUT_DTYPES: + raise ValueError(f"D has unsupported dtype {self.d.dtype}") + if self.prob.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"prob must have float32 dtype, got {self.prob.dtype}") + if self.dprob is not None and self.dprob.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"dprob must have float32 dtype, got {self.dprob.dtype}") + if self.acc_dtype != data_type.FLOAT: + raise ValueError(f"Accumulator dtype must be float32, got {self.acc_dtype}") + if self.amax is not None and self.amax.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"amax must have float32 dtype, got {self.amax.dtype}") + + a_major = require_compact_major(self.a, "m", "k") + b_major = require_compact_major(self.b, "n", "k") + c_major = require_compact_major(self.c, "m", "n") + d_major = require_compact_major(self.d, "m", "n") + if c_major != d_major: + raise ValueError(f"C and D must use the same major mode, got {c_major}-major and {d_major}-major") + if ab_dtype in _FP4_DTYPES and (a_major != "k" or b_major != "k"): + raise ValueError("FP4 A and B must use K-major layouts") + + self.mma_tiler_mn = require_mma_tiler( + self.mma_tiler_mn, + allowed_n=(64, 128, 192, 256), + ) + cta_group_size = 2 if self.mma_tiler_mn[0] == 256 else 1 + cta_tile_m = self.mma_tiler_mn[0] // cta_group_size + if m % cta_tile_m: + raise ValueError(f"M must be divisible by CTA_TILE_M={cta_tile_m} because the probability load is not predicated, got {m}") + if self.cluster_shape_mn is None: + self.cluster_shape_mn = (2, 1) if cta_group_size == 2 else (1, 1) + self.cluster_shape_mn = require_cluster_shape( + self.cluster_shape_mn, + cta_group_size=cta_group_size, + ) + if self.cluster_shape_mn[0] > 4 or self.cluster_shape_mn[1] > 4: + raise ValueError("cluster_shape_mn entries must be at most 4 for scale-factor multicast, " f"got {self.cluster_shape_mn}") + + for tensor in (self.c, self.d): + require_16_byte_alignment(tensor) + for tensor in (self.a, self.b): + if tensor.cudnn_dtype == data_type.UINT8: + contiguous_extent = tensor.shape[tensor.stride_order[0]] + if contiguous_extent * 4 % 128 != 0: + raise ValueError( + f"{tensor.name or 'Tensor'} contiguous extent must be a multiple of 32 logical FP4 elements " + f"for 16-byte alignment, got {contiguous_extent}" + ) + else: + require_16_byte_alignment(tensor) + + self.m, self.n, self.k, self.l = m, n, k, batch + self.ab_dtype = ab_dtype + self.a_major, self.b_major, self.output_major = a_major, b_major, c_major + return True + + +__all__ = ["BlockScaledGemmReluSm100Op"] diff --git a/python/cudnn/_jax/api_base.py b/python/cudnn/_jax/api_base.py index ff50cf4fc..fa468d14d 100644 --- a/python/cudnn/_jax/api_base.py +++ b/python/cudnn/_jax/api_base.py @@ -355,6 +355,7 @@ def _call_kernel( self, inputs: tuple[Any, ...], *, + launch: Callable[..., None], output_descs: tuple[TensorDesc[Any], ...], workspace_descs: tuple[TensorDesc[Any], ...] = (), input_spec: tuple[Any | None, ...] | None = None, @@ -364,10 +365,10 @@ def _call_kernel( compile_options: Any = None, use_static_tensors: bool = True, ) -> tuple[Any, ...]: - """Bind this adapter's launch hook to JAX and return public outputs. + """Bind an explicit kernel launcher to JAX and return public outputs. - CUTLASS JAX supplies the stream first; the launcher adapts that call to - the adapter's ``inputs, outputs, workspaces, stream`` launch hook. Workspaces + CUTLASS JAX supplies the stream first; the launcher receives + ``stream, *inputs, *outputs, *workspaces`` in that order. Workspaces are declared as custom-call results so XLA owns their lifetime, then omitted from this method's return value. Descriptors with a non-``None`` ``init_value`` are materialized as JAX inputs and aliased to their @@ -378,6 +379,9 @@ def _call_kernel( import cutlass.jax as cutlass_jax import jax.numpy as jnp + if not callable(launch): + raise TypeError(f"launch must be callable, got {type(launch).__name__}") + inputs = tuple(inputs) for index, value in enumerate(inputs): _require_array_metadata(value, f"input #{index}") @@ -451,7 +455,7 @@ def launcher(stream, *args): ordered_buffers[index] = value output_buffers = tuple(ordered_buffers[: len(outputs)]) workspace_buffers = tuple(ordered_buffers[len(outputs) :]) - self._launch(tuple(kernel_inputs), output_buffers, workspace_buffers, stream) + launch(stream, *kernel_inputs, *output_buffers, *workspace_buffers) call = cutlass_jax.cutlass_call( launcher, @@ -468,16 +472,6 @@ def launcher(stream, *args): results = (results,) return tuple(results[: len(outputs)]) - @abstractmethod - def _launch( - self, - inputs: tuple[Any, ...], - outputs: tuple[Any, ...], - workspaces: tuple[Any, ...], - stream: Any, - ) -> None: - """Launch the concrete CuTe kernel for one traced custom call.""" - @abstractmethod def check_support(self) -> bool: """Validate the operation signature and static configuration.""" diff --git a/python/cudnn/_jax/gemm.py b/python/cudnn/_jax/gemm.py index 7a82f8d40..4734e89f8 100644 --- a/python/cudnn/_jax/gemm.py +++ b/python/cudnn/_jax/gemm.py @@ -12,7 +12,11 @@ GEMM_OUTPUT_LAYOUTS = ("LMN", "LNM") ROW_MAJOR_STRIDE_ORDER_3D = (2, 1, 0) +# Public row-major shape: (L, tiles_M_or_N, rest_K_or_N, 32, 4, 4). +BLOCK_SCALE_MODE = (3, 4, 1, 5, 2, 0) BLOCK_SCALE_STRIDE_ORDER = (3, 1, 0, 4, 2, 5) +# Public row-major shape: (L, 1, M). +PROBABILITY_MODE = (2, 1, 0) PROBABILITY_STRIDE_ORDER = (0, 1, 2) @@ -49,10 +53,12 @@ def gemm_output_mode(layout: str, *, name: str = "c_layout") -> tuple[int, ...]: __all__ = [ + "BLOCK_SCALE_MODE", "BLOCK_SCALE_STRIDE_ORDER", "GEMM_A_LAYOUTS", "GEMM_B_LAYOUTS", "GEMM_OUTPUT_LAYOUTS", + "PROBABILITY_MODE", "PROBABILITY_STRIDE_ORDER", "ROW_MAJOR_STRIDE_ORDER_3D", "gemm_a_mode", diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py index cbf1ca15d..dee253739 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py @@ -196,6 +196,7 @@ def __call__( ) d_index_q, d_weights, d_index_k_accum = self._call_kernel( (index_q, weights, index_k, attn_score, index_score, topk_indices, grad_loss), + launch=self._launch_kernel, output_descs=(self.diq_desc, self.dw_desc, dk_accum_desc), workspace_descs=(grad_signal_desc,), input_spec=( @@ -220,12 +221,23 @@ def __call__( d_index_k=d_index_k_accum.astype(self.dik_desc.dtype), ) - def _launch(self, inputs, outputs, workspaces, stream) -> None: + def _launch_kernel( + self, + stream, + index_q, + weights, + index_k, + attn_score, + index_score, + topk_indices, + grad_loss, + d_index_q, + d_weights, + d_index_k_accum, + grad_signal, + ) -> None: import cutlass - index_q, weights, index_k, attn_score, index_score, topk_indices, grad_loss = inputs - d_index_q, d_weights, d_index_k_accum = outputs - (grad_signal,) = workspaces grad_scale = self.loss_coeff / (self._op.batch * self._op.seqlen_q) if self._op.is_thd: @@ -467,6 +479,7 @@ def __call__( ) d_index_q, d_weights, d_index_k_accum = self._call_kernel( inputs, + launch=self._launch_kernel, output_descs=(self.diq_desc, self.dw_desc, dk_accum_desc), workspace_descs=(grad_signal_desc,), input_spec=( @@ -488,9 +501,10 @@ def __call__( d_index_k=d_index_k_accum.astype(self.dik_desc.dtype), ) - def _launch(self, inputs, outputs, workspaces, stream) -> None: + def _launch_kernel(self, stream, *arguments) -> None: import cutlass + *inputs, d_index_q, d_weights, d_index_k_accum, grad_signal = arguments index_q, weights, index_k, attn_score, attn_l1norm, index_score, index_lse, grad_loss, *optional = inputs cursor = 0 cu_seqlens_q = cu_seqlens_k = q_causal_offsets = None @@ -500,8 +514,6 @@ def _launch(self, inputs, outputs, workspaces, stream) -> None: if self._op.q_causal_offsets is not None: q_causal_offsets = optional[cursor] - d_index_q, d_weights, d_index_k_accum = outputs - (grad_signal,) = workspaces grad_scale = self.loss_coeff / max(self._op.normalization_tokens, 1) if self._architecture_family == 100: diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py index 51edcfe8f..db19225f7 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py @@ -249,6 +249,7 @@ def __call__( output_divisibility = (None,) * (self.o_desc.ndim - 1) + (TMA_ALIGN_ELEMENTS,) (scores_padded,) = self._call_kernel( tuple(inputs), + launch=self._launch_kernel, output_descs=(self.o_desc,), input_spec=tuple(self._to_tensor_spec(desc, mode=mode) for desc, mode in input_bindings), output_spec=( @@ -274,25 +275,20 @@ def _check_optional_signature(value: Any | None, desc: JaxTensorDesc | None, nam if desc is not None: JaxApiBase._check_tensor_signature(value, desc) - def _launch( + def _launch_kernel( self, - inputs: tuple[Any, ...], - outputs: tuple[Any, ...], - workspaces: tuple[Any, ...], stream: Any, + *arguments: Any, ) -> None: from cutlass import BFloat16, Float32, Int32 + *inputs, output = arguments q, k, w, *optional_inputs = inputs cu_seqlens_q = optional_inputs.pop(0) if self.cu_seqlens_q_desc is not None else None cu_seqlens_k = optional_inputs.pop(0) if self.cu_seqlens_k_desc is not None else None q_causal_offsets = optional_inputs.pop(0) if self.q_causal_offsets_desc is not None else None if optional_inputs: raise RuntimeError("Unexpected IndexerForward kernel inputs") - (output,) = outputs - if workspaces: - raise RuntimeError("IndexerForward does not use workspaces") - resolved = ( self._op.head_dim, self._op.qhead_per_kv_head, diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py index cfe42d2fd..6a424e165 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py @@ -105,6 +105,7 @@ def __call__(self, input_values: Any, seq_lens: Any) -> TupleDict: output_descs = (self.indices_desc,) if self.values_desc is None else (self.indices_desc, self.values_desc) results = self._call_kernel( (input_values, seq_lens), + launch=self._launch_kernel, output_descs=output_descs, workspace_descs=(self.workspace_desc,), input_spec=(self._to_tensor_spec(self.input_desc), self._to_tensor_spec(self.seq_lens_desc)), @@ -114,12 +115,12 @@ def __call__(self, input_values: Any, seq_lens: Any) -> TupleDict: values = results[1] if self.return_val else None return TupleDict(indices=indices, values=values) - def _launch( + def _launch_kernel( self, - inputs: tuple[Any, ...], - outputs: tuple[Any, ...], - workspaces: tuple[Any, ...], stream: Any, + input_values: Any, + seq_lens: Any, + *buffers: Any, ) -> None: import cutlass @@ -130,10 +131,10 @@ def _launch( data_type.BFLOAT16: cutlass.BFloat16, data_type.FLOAT: cutlass.Float32, } - input_values, seq_lens = inputs + outputs = buffers[:-1] + extra_buffer = buffers[-1] output_indices = outputs[0] output_values = outputs[1] if self.return_val else None - (extra_buffer,) = workspaces resolved = ( self._op.max_num_cols, self._op.large_occupancy, @@ -239,24 +240,23 @@ def __call__( (result,) = self._call_kernel( tuple(inputs), + launch=self._launch_kernel, output_descs=(self.output_desc,), compile_options=compile_options_for_target(self.target_compute_capability, "--opt-level 3"), ) return TupleDict(indices=result) - def _launch(self, inputs, outputs, workspaces, stream) -> None: + def _launch_kernel(self, stream, *arguments) -> None: from cutlass import Int32 from .local_to_global_dsl import LocalToGlobalTopK + *inputs, output = arguments local_indices, *optional = inputs if self.is_varlen: cu_q, cu_k = optional else: cu_q = cu_k = None - (output,) = outputs - if workspaces: - raise RuntimeError("local_to_global does not use workspaces") LocalToGlobalTopK(is_varlen=self.is_varlen)( local_indices, output, @@ -305,20 +305,17 @@ def __call__(self, indices: Any) -> TupleDict: self._check_tensor_signature(indices, self.input_desc) compact_indices, topk_length = self._call_kernel( (indices,), + launch=self._launch_kernel, output_descs=(self.output_desc, self.length_desc), compile_options=compile_options_for_target(self.target_compute_capability, "--opt-level 3"), ) return TupleDict(indices=compact_indices, topk_length=topk_length) - def _launch(self, inputs, outputs, workspaces, stream) -> None: + def _launch_kernel(self, stream, indices, compact_indices, topk_length) -> None: from cutlass import Int32 from .compactify import CompactifyKernel - (indices,) = inputs - compact_indices, topk_length = outputs - if workspaces: - raise RuntimeError("compactify does not use workspaces") CompactifyKernel(cols=self.cols)(indices, compact_indices, topk_length, Int32(self.rows), stream) diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py b/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py index 9c97ec5e8..5c689a6fa 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py @@ -128,22 +128,26 @@ def _call(self, q: Any, k: Any, per_head: Any, topk_indices: Any, topk_length: A (output,) = self._call_kernel( inputs, + launch=self._launch_kernel, output_descs=(self.out_desc,), workspace_descs=workspace_descs, compile_options=_compile_options(self.target_compute_capability), ) return output - def _launch( + def _launch_kernel( self, - inputs: tuple[Any, ...], - outputs: tuple[Any, ...], - workspaces: tuple[Any, ...], stream: Any, + *arguments: Any, ) -> None: - output = outputs[0] + uses_workspace = self.topk_length_desc is None and self.target_compute_capability != 90 + if uses_workspace: + *inputs, output, topk_length_workspace = arguments + else: + *inputs, output = arguments + topk_length_workspace = None q, k, per_head, topk_indices, *optional_inputs = inputs - topk_length = optional_inputs[0] if optional_inputs else (None if self.target_compute_capability == 90 else workspaces[0]) + topk_length = optional_inputs[0] if optional_inputs else topk_length_workspace if self.target_compute_capability == 90: import cutlass @@ -387,21 +391,19 @@ def _call( output, denominator = self._call_kernel( inputs, + launch=self._launch_kernel, output_descs=(self.out_desc, self.denom_desc), compile_options=_compile_options(self.target_compute_capability), ) return TupleDict(out=output, denom=denominator) - def _launch( + def _launch_kernel( self, - inputs: tuple[Any, ...], - outputs: tuple[Any, ...], - workspaces: tuple[Any, ...], stream: Any, + *arguments: Any, ) -> None: - del workspaces + *inputs, output, denominator = arguments q, k, per_head, *optional_inputs = inputs - output, denominator = outputs if self.target_compute_capability == 90: import cutlass diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py index e43d3adf8..f769b9e04 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py @@ -156,6 +156,7 @@ def __call__( dq, dkv, d_sink = self._call_kernel( inputs, + launch=self._launch_kernel, output_descs=(self.dq_desc, self.dkv_desc, self.d_sink_desc), workspace_descs=self._workspace_descs(), input_spec=input_specs, @@ -239,13 +240,16 @@ def _architecture_family(self) -> int: raise RuntimeError(f"No sparse-attention backward kernel for SM{self.compute_capability}") return family - def _launch( + def _launch_kernel( self, - inputs: tuple[Any, ...], - outputs: tuple[Any, ...], - workspaces: tuple[Any, ...], stream: Any, + *arguments: Any, ) -> None: + workspace_count = len(self._workspace_descs()) + input_count = len(arguments) - 3 - workspace_count + inputs = arguments[:input_count] + outputs = arguments[input_count : input_count + 3] + workspaces = arguments[input_count + 3 :] if self._architecture_family == 90: self._launch_sm90(inputs, outputs, workspaces, stream) else: diff --git a/python/cudnn/gemm_amax/__init__.py b/python/cudnn/gemm_amax/__init__.py index b65cf491e..33844177a 100644 --- a/python/cudnn/gemm_amax/__init__.py +++ b/python/cudnn/gemm_amax/__init__.py @@ -1,9 +1,18 @@ -from .api import ( - GemmAmaxSm100, - gemm_amax_wrapper_sm100, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy Torch API and framework-neutral GEMM + amax operation exports.""" -__all__ = [ - "GemmAmaxSm100", - "gemm_amax_wrapper_sm100", -] +from .._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ("GemmAmaxSm100Op",), + "api": ( + "GemmAmaxSm100", + "gemm_amax_wrapper_sm100", + ), + }, + submodules=("api", "op"), +) diff --git a/python/cudnn/gemm_amax/api.py b/python/cudnn/gemm_amax/api.py index 62fc475bb..a43eaabd9 100644 --- a/python/cudnn/gemm_amax/api.py +++ b/python/cudnn/gemm_amax/api.py @@ -1,18 +1,26 @@ -from .dense_blockscaled_gemm_persistent_amax import ( - Sm100BlockScaledPersistentDenseGemmKernel, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT -from cuda.bindings import driver as cuda +"""Torch API for block-scaled dense GEMM + amax on SM100.""" + +import logging import os -import torch -from typing import Tuple, Optional +from typing import Optional, Tuple +from cuda.bindings import driver as cuda import cutlass import cutlass.cute as cute from cutlass.cute.runtime import make_fake_stream +import torch -from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.api_base import APIBase, TensorDesc, TupleDict, is_power_of_2, ceil_div +from cudnn import data_type +from cudnn.api_base import APIBase, TupleDict +from cudnn.datatypes import _torch_to_cudnn_data_type + +from .dense_blockscaled_gemm_persistent_amax import ( + Sm100BlockScaledPersistentDenseGemmKernel, +) +from .op import GemmAmaxSm100Op class GemmAmaxSm100(APIBase): @@ -39,216 +47,75 @@ def __init__( self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa") self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb") self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") - self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax") + self.amax_desc = self._pad_tensor_to_ndim( + self._make_tensor_desc(sample_amax, name="sample_amax"), + 3, + "sample_amax", + ) self.acc_dtype = acc_dtype self.mma_tiler_mn = mma_tiler_mn self.cluster_shape_mn = cluster_shape_mn self.sf_vec_size = sf_vec_size - self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) - self._logger.debug(f"setting num_cluster_overlap_margin: {self.num_cluster_overlap_margin}") - - # used to reshape sfa/sfb tensors to atom layout - self.atom_m = (32, 4) - self.atom_k = 4 - - self._interpret_uint8_as_fp4x2 = True self._logger.debug( - f"__init__ completed with args: sample_a {self.a_desc.shape}, sample_b {self.b_desc.shape}, sample_sfa {self.sfa_desc.shape}, sample_sfb {self.sfb_desc.shape}, sample_c {self.c_desc.shape}, sample_amax {self.amax_desc.shape}, acc_dtype {acc_dtype}, mma_tiler_mn {mma_tiler_mn}, cluster_shape_mn {cluster_shape_mn}, sf_vec_size {sf_vec_size}" + "setting num_cluster_overlap_margin: %s", + self.num_cluster_overlap_margin, ) - def check_support(self) -> bool: - self._logger.debug("Entering check_support") - - self._logger.debug("Checking dtypes and sf_vec_size") - ab_dtype = self._check_dtype( - self.a_desc, - dtype=[torch.float4_e2m1fn_x2, torch.uint8, torch.float8_e5m2, torch.float8_e4m3fn], - name="A", - ) - self._check_dtype( - self.b_desc, - dtype=ab_dtype, - name="B", - extra_error_msg="A and B tensor dtypes must match", + # Torch uint8 is the legacy packed-FP4 storage alias for this API. + self._interpret_uint8_as_fp4x2 = True + canonical_acc_dtype = _torch_to_cudnn_data_type(self.acc_dtype) + if canonical_acc_dtype is None: + canonical_acc_dtype = data_type.NOT_SET + self._op = GemmAmaxSm100Op( + a=self.a_desc, + b=self.b_desc, + sfa=self.sfa_desc, + sfb=self.sfb_desc, + c=self.c_desc, + amax=self.amax_desc, + acc_dtype=canonical_acc_dtype, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=self.cluster_shape_mn, + sf_vec_size=self.sf_vec_size, ) - if ab_dtype == torch.uint8: - self._logger.warning("Uint8 ab_dtype will be interpreted as packed fp4, not as native uint8") + self._kernel = Sm100BlockScaledPersistentDenseGemmKernel - self._value_error_if( - self.sf_vec_size not in {16, 32}, - f"Unsupported sf_vec_size: received {self.sf_vec_size}, expected {{16, 32}}", + self._logger.debug( + "__init__ completed with args: sample_a %s, sample_b %s, " + "sample_sfa %s, sample_sfb %s, sample_c %s, sample_amax %s, " + "acc_dtype %s, mma_tiler_mn %s, cluster_shape_mn %s, sf_vec_size %s", + self.a_desc.shape, + self.b_desc.shape, + self.sfa_desc.shape, + self.sfb_desc.shape, + self.c_desc.shape, + self.amax_desc.shape, + acc_dtype, + mma_tiler_mn, + cluster_shape_mn, + sf_vec_size, ) - sf_dtype = self._check_dtype( - self.sfa_desc, - dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn, torch.int8], - name="sfa", - ) - self._check_dtype( - self.sfb_desc, - dtype=sf_dtype, - name="sfb", - extra_error_msg="sfa and sfb tensor dtypes must match", - ) - if sf_dtype == torch.int8: + def check_support(self) -> bool: + """Validate the shared operation contract and Torch device support.""" + + self._is_supported = False + self._op.check_support() + self.mma_tiler_mn = self._op.mma_tiler_mn + self.cluster_shape_mn = self._op.cluster_shape_mn + self.a_major = self._op.a_major + self.b_major = self._op.b_major + self.c_major = self._op.c_major + self.ab_dtype = self.a_desc.dtype + self.c_dtype = self.c_desc.dtype + + if self.a_desc.dtype == torch.uint8: + self._logger.warning("Uint8 ab_dtype will be interpreted as packed fp4, not as native uint8") + if self.sfa_desc.dtype == torch.int8: self._logger.warning("Int8 sf_dtype will be interpreted as float8_e8m0fnu, not as native int8") - self._value_error_if( - sf_dtype == torch.float8_e4m3fn and self.sf_vec_size == 32, - "Unsupported sf_dtype and sf_vec_size combination: float8_e4m3fn and 32 is not supported", - ) - self._value_error_if( - ab_dtype in {torch.float8_e5m2, torch.float8_e4m3fn} and self.sf_vec_size == 16, - f"Unsupported ab_dtype and sf_vec_size combination: {{float8_e5m2, float8_e4m3fn}} and 16 is not supported", - ) - - c_dtype = self._check_dtype( - self.c_desc, - dtype=[torch.float32, torch.float16, torch.bfloat16, torch.float8_e5m2, torch.float8_e4m3fn, torch.float4_e2m1fn_x2, torch.uint8], - name="C", - ) - self._value_error_if( - self._is_fp4x2(c_dtype) and not self._is_fp4x2(ab_dtype), - f"Unsupported c_dtype and ab_dtype combination: fp4 c_dtype requires fp4 ab_dtype, got {ab_dtype}", - ) - self._not_implemented_error_if( - self._is_fp8(c_dtype) and self._is_fp8(ab_dtype), - "Unsupported c_dtype and ab_dtype combination: fp8 ab_dtype and fp8 c_dtype (fails to launch)", - ) - self._check_dtype( - self.acc_dtype, - dtype=torch.float32, - name="Accumulator", - extra_error_msg="Accumulator must be float32", - ) - - self.ab_dtype = ab_dtype - self.c_dtype = c_dtype - - self._logger.debug("Checking tensor layout") - m, k, l = self.a_desc.shape - n, _, _ = self.b_desc.shape - _, _, m_div_atom_m0_m1, _, sf_k_div_atom_k, _ = self.sfa_desc.shape - _, _, n_div_atom_m0_m1, _, sf_k_div_atom_k, _ = self.sfb_desc.shape - - self._check_tensor_shape(self.a_desc, (m, k, l), "A") - self._check_tensor_shape(self.b_desc, (n, k, l), "B") - self._check_tensor_shape(self.c_desc, (m, n, l), "C") - self._check_tensor_shape( - self.sfa_desc, - (self.atom_m[0], self.atom_m[1], m_div_atom_m0_m1, self.atom_k, sf_k_div_atom_k, l), - "sfa", - ) - self._check_tensor_shape( - self.sfb_desc, - (self.atom_m[0], self.atom_m[1], n_div_atom_m0_m1, self.atom_k, sf_k_div_atom_k, l), - "sfb", - ) - self.amax_desc = self._pad_tensor_to_ndim(self.amax_desc, 3, "amax") - self._check_tensor_shape(self.amax_desc, (1, 1, 1), "amax") - - expected_m_div_atom = ceil_div(m, self.atom_m[0] * self.atom_m[1]) - expected_n_div_atom = ceil_div(n, self.atom_m[0] * self.atom_m[1]) - self._value_error_if( - m_div_atom_m0_m1 != expected_m_div_atom, - f"Input/Output shape mismatch: expected m_div_atom_m0_m1 (sfa.shape[2]) = {expected_m_div_atom}, got {m_div_atom_m0_m1}", - ) - self._value_error_if( - n_div_atom_m0_m1 != expected_n_div_atom, - f"Input/Output shape mismatch: expected n_div_atom_m0_m1 (sfb.shape[2]) = {expected_n_div_atom}, got {n_div_atom_m0_m1}", - ) - - # Check tensor strides - _ = self._check_tensor_stride( - self.a_desc, - stride=[(1, m, m * k), (k, 1, m * k)], - name="A", - ) - _ = self._check_tensor_stride( - self.b_desc, - stride=[(1, n, n * k), (k, 1, n * k)], - name="B", - ) - _ = self._check_tensor_stride( - self.c_desc, - stride=[(1, m, m * n), (n, 1, m * n)], - name="C", - ) - - # Derive major mode from stride order - self.a_major = "m" if self.a_desc.stride_order == (0, 1, 2) else "k" - self.b_major = "n" if self.b_desc.stride_order == (0, 1, 2) else "k" - self.c_major = "m" if self.c_desc.stride_order == (0, 1, 2) else "n" - - self._value_error_if( - self._is_fp4x2(ab_dtype) and not (self.a_major == "k" and self.b_major == "k"), - f"Unsupported A or B tensor stride: Float4 tensors require k-major layout for hardware efficiency, got {self.a_major} and {self.b_major}", - ) - self._value_error_if( - self._is_fp4x2(c_dtype) and self.c_major == "m", - f"Unsupported C tensor stride: Float4 tensors require n-major layout for hardware efficiency, got {self.c_major}", - ) - - self._logger.debug("Checking mma tiler and cluster shape") - self._value_error_if( - self.mma_tiler_mn[0] not in [128, 256], - f"Unsupported mma_tiler_mn[0]: expected {{128, 256}}, got {self.mma_tiler_mn[0]}", - ) - self._value_error_if( - self.mma_tiler_mn[1] not in [128, 256], - f"Unsupported mma_tiler_mn[1]: expected {{128, 256}}, got {self.mma_tiler_mn[1]}", - ) - self._not_implemented_error_if( - self.mma_tiler_mn[0] == 256, - "mma_tiler_mn[0] == 256 currently hangs", - ) - self._value_error_if( - self._is_fp4x2(self.ab_dtype) and self.mma_tiler_mn[1] == 256 and k <= 128, - f"mma_tiler_mn (X, 256) requires k > 128 (packed x2), got {k}", - ) - self._value_error_if( - not (self.cluster_shape_mn[0] % (2 if self.mma_tiler_mn[0] == 256 else 1) == 0), - "Illegal cluster shape", - ) - self._not_implemented_error_if( - self.mma_tiler_mn == (128, 256) and self.sf_vec_size == 16 and c_dtype in {torch.float32, torch.float16, torch.bfloat16}, - "mma_tiler_mn (128, 256), sf_vec_size 16, c_dtype {torch.float32, torch.float16, torch.bfloat16} fails to launch", - ) - - # Special cluster shape check for scale factor multicasts. - # Due to limited size of scale factors, we can't multicast among more than 4 CTAs. - self._value_error_if( - not ( - self.cluster_shape_mn[0] <= 4 - and self.cluster_shape_mn[1] <= 4 - and self.cluster_shape_mn[0] > 0 - and self.cluster_shape_mn[1] > 0 - and is_power_of_2(self.cluster_shape_mn[0]) - and is_power_of_2(self.cluster_shape_mn[1]) - ), - f"Invalid cluster shape: expected cluster_shape_mn values in {{1, 2, 4}}, got {self.cluster_shape_mn}", - ) - - self._logger.debug("Checking tensor alignment") - - def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): - major_mode_idx = 0 if is_mode0_major else 1 - num_major_elements = tensor_shape[major_mode_idx] - num_contiguous_elements = 16 * 8 // (_convert_to_cutlass_data_type(dtype).width) - return num_major_elements % num_contiguous_elements == 0 - - self._value_error_if( - not ( - check_contigous_16B_alignment(ab_dtype, self.a_major == "m", (m, k, l)) - and check_contigous_16B_alignment(ab_dtype, self.b_major == "n", (n, k, l)) - and check_contigous_16B_alignment(c_dtype, self.c_major == "m", (m, n, l)) - ), - "Unsupported tensor alignment: tensors must be 16B aligned", - ) - - self._logger.debug("Checking environment") self._runtime_error_if(not torch.cuda.is_available(), "CUDA is not available") device = torch.cuda.current_device() major, minor = torch.cuda.get_device_capability(device) @@ -258,8 +125,6 @@ def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): f"GemmAmax requires SM100+ compute capability, but found SM{compute_capability} on device {device}", ) - self._kernel = Sm100BlockScaledPersistentDenseGemmKernel - self._is_supported = True self._logger.debug("check_support completed successfully") return True @@ -281,7 +146,7 @@ def compile(self) -> None: max_active_clusters -= self.num_cluster_overlap_margin self._value_error_if( max_active_clusters <= 0, - "max_active_clusters must be > 0 after applying overlap margin; reduce CUDNNFE_CLUSTER_OVERLAP_MARGIN", + "max_active_clusters must be > 0 after applying overlap margin; " "reduce CUDNNFE_CLUSTER_OVERLAP_MARGIN", ) self._logger.debug("Compiling gemm_amax") @@ -293,7 +158,7 @@ def compile(self) -> None: amax_cute = self._make_fake_cute_tensor_from_desc(self.amax_desc, assumed_align=16) fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) - _compiled_kernel = cute.compile( + compiled_kernel = cute.compile( gemm_amax, a_tensor=a_cute, b_tensor=b_cute, @@ -316,8 +181,7 @@ def tensor_api( stream: cuda.CUstream, ): amax_tensor = self._pad_tensor_to_ndim(amax_tensor, 3, "amax") - - return _compiled_kernel( + return compiled_kernel( a_tensor, b_tensor, sfa_tensor, @@ -328,11 +192,6 @@ def tensor_api( ) self._compiled_kernel = tensor_api - - for attr_name in tuple(vars(self)): - if attr_name.startswith("sample_"): - setattr(self, attr_name, None) - self._logger.debug("Kernel compiled successfully") def execute( @@ -353,7 +212,6 @@ def execute( "GemmAmaxSm100 kernel not compiled; call compile() first", ) self._logger.debug("Executing with compiled kernel") - self._compiled_kernel( a_tensor=a_tensor, b_tensor=b_tensor, @@ -366,8 +224,6 @@ def execute( self._logger.debug("Executed with compiled kernel successfully") -import logging - _logger = logging.getLogger(__name__) _cache_of_GemmAmaxSm100Objects = {} @@ -385,19 +241,32 @@ def gemm_amax_wrapper_sm100( sf_vec_size: int = 32, stream: Optional[cuda.CUstream] = None, ) -> TupleDict: - _logger.debug("gemm_amax_wrapper_sm100: Creating empty output tensors c and amax") m, _, l = a_tensor.shape n, _, l = b_tensor.shape - c_tensor = None if c_major == "m": - c_tensor = torch.empty_strided((m, n, l), (1, m, m * n), dtype=c_dtype, device=a_tensor.device) + c_tensor = torch.empty_strided( + (m, n, l), + (1, m, m * n), + dtype=c_dtype, + device=a_tensor.device, + ) elif c_major == "n": - c_tensor = torch.empty_strided((m, n, l), (n, 1, m * n), dtype=c_dtype, device=a_tensor.device) + c_tensor = torch.empty_strided( + (m, n, l), + (n, 1, m * n), + dtype=c_dtype, + device=a_tensor.device, + ) else: raise ValueError(f"c_major must be either 'm' or 'n', got {c_major}") - amax_tensor = torch.full((1, 1, 1), -float("inf"), device=a_tensor.device, dtype=torch.float32) + amax_tensor = torch.full( + (1, 1, 1), + -float("inf"), + device=a_tensor.device, + dtype=torch.float32, + ) cache_key = ( a_tensor.shape, @@ -432,7 +301,7 @@ def gemm_amax_wrapper_sm100( current_stream=stream, ) else: - _logger.debug("gemm_amax_wrapper_sm100: No previously cached GemmAmaxSm100 object found, creating new GemmAmaxSm100 object") + _logger.debug("gemm_amax_wrapper_sm100: No previously cached GemmAmaxSm100 object " "found, creating new GemmAmaxSm100 object") gemm_amax = GemmAmaxSm100( sample_a=a_tensor, sample_b=b_tensor, @@ -445,7 +314,7 @@ def gemm_amax_wrapper_sm100( cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, ) - assert gemm_amax.check_support() + gemm_amax.check_support() gemm_amax.compile() gemm_amax.execute( a_tensor=a_tensor, diff --git a/python/cudnn/gemm_amax/jax.py b/python/cudnn/gemm_amax/jax.py new file mode 100644 index 000000000..71436d36d --- /dev/null +++ b/python/cudnn/gemm_amax/jax.py @@ -0,0 +1,282 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX API for block-scaled dense GEMM + amax on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any + +import jax +import jax.numpy as jnp + +from .. import data_type +from .._cute_compiler import compile_options_for_target +from .._dense_gemm import require_gemm_inputs +from .._jax import JaxApiBase, JaxTensorDesc, TupleDict +from .._jax.datatypes import jax_to_cudnn_dtype, normalize_jax_dtype +from .._jax.gemm import BLOCK_SCALE_MODE, gemm_a_mode, gemm_b_mode, gemm_output_mode +from .._jax.layout import to_public_axes +from .op import GemmAmaxSm100Op + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) + + +class GemmAmaxSm100(JaxApiBase): + """JAX callable specialized from block-scaled GEMM input metadata.""" + + def __init__( + self, + sample_a: Any, + sample_b: Any, + sample_sfa: Any, + sample_sfb: Any, + *, + sample_c: Any | None = None, + sample_amax: Any | None = None, + c_layout: str = "LMN", + c_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: tuple[int, int] = (1, 1), + sf_vec_size: int = 32, + a_layout: str = "LMK", + b_layout: str = "LNK", + ) -> None: + self.a_layout = a_layout + self.b_layout = b_layout + self.c_layout = c_layout + self.a_mode = gemm_a_mode(a_layout) + self.b_mode = gemm_b_mode(b_layout) + self.output_mode = gemm_output_mode(c_layout) + self.scale_mode = BLOCK_SCALE_MODE + + self.compute_capability = self._resolve_compute_capability( + None, + SUPPORTED_COMPUTE_CAPABILITIES, + "GemmAmaxSm100", + ) + self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=self.a_mode) + self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=self.b_mode) + self.sfa_desc = self._to_tensor_desc( + sample_sfa, + "sample_sfa", + mode=self.scale_mode, + ) + self.sfb_desc = self._to_tensor_desc( + sample_sfb, + "sample_sfb", + mode=self.scale_mode, + ) + self.acc_dtype = normalize_jax_dtype(acc_dtype, jnp.float32, "acc_dtype") + + if (sample_c is None) != (sample_amax is None): + raise ValueError("sample_c and sample_amax must be provided together") + if sample_c is None: + resolved_c_dtype = normalize_jax_dtype(c_dtype, jnp.float32, "c_dtype") + self.c_desc, self.amax_desc = self._default_output_descs(resolved_c_dtype) + else: + self.c_desc = self._to_tensor_desc( + sample_c, + "sample_c", + mode=self.output_mode, + ) + self.amax_desc = self._to_tensor_desc( + sample_amax, + "sample_amax", + init_value=float("-inf"), + ) + self._check_requested_c_dtype(c_dtype) + + acc_cudnn_dtype = jax_to_cudnn_dtype(self.acc_dtype) + if acc_cudnn_dtype == data_type.NOT_SET: + raise ValueError(f"Unsupported JAX accumulator dtype {self.acc_dtype}") + self._op = GemmAmaxSm100Op( + a=self.a_desc, + b=self.b_desc, + sfa=self.sfa_desc, + sfb=self.sfb_desc, + c=self.c_desc, + amax=self.amax_desc, + acc_dtype=acc_cudnn_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + ) + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + + def _default_output_descs( + self, + c_dtype: Any, + ) -> tuple[JaxTensorDesc, JaxTensorDesc]: + m, n, _, batch = require_gemm_inputs(self.a_desc, self.b_desc) + return ( + self._to_tensor_desc( + jax.ShapeDtypeStruct( + to_public_axes((m, n, batch), self.output_mode), + c_dtype, + ), + "sample_c", + mode=self.output_mode, + ), + self._to_tensor_desc( + jax.ShapeDtypeStruct((1, 1, 1), jnp.float32), + "sample_amax", + init_value=float("-inf"), + ), + ) + + def _check_requested_c_dtype(self, requested: Any | None) -> None: + if requested is None: + return + requested_dtype = normalize_jax_dtype(requested, requested, "c_dtype") + actual_dtype = jnp.dtype(self.c_desc.dtype) + if requested_dtype != actual_dtype: + raise ValueError(f"c_dtype={requested_dtype} does not match the explicit sample dtype " f"{actual_dtype}") + + def check_support(self) -> bool: + """Validate the JAX dtype surface and common operation contract.""" + + if self.a_desc.cudnn_dtype not in { + data_type.FP4_E2M1, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + }: + raise ValueError(f"The JAX GEMM + amax API requires FP4 or FP8 A and B, got " f"{self.a_desc.dtype}") + if self.b_desc.cudnn_dtype not in { + data_type.FP4_E2M1, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + }: + raise ValueError(f"The JAX GEMM + amax API requires FP4 or FP8 A and B, got " f"{self.b_desc.dtype}") + if self.sfa_desc.cudnn_dtype not in { + data_type.FP8_E8M0, + data_type.FP8_E4M3, + }: + raise ValueError(f"The JAX GEMM + amax API requires E8M0 or E4M3 scale factors, " f"got {self.sfa_desc.dtype}") + if self.sfb_desc.cudnn_dtype not in { + data_type.FP8_E8M0, + data_type.FP8_E4M3, + }: + raise ValueError(f"The JAX GEMM + amax API requires E8M0 or E4M3 scale factors, " f"got {self.sfb_desc.dtype}") + if self.c_desc.cudnn_dtype not in { + data_type.FLOAT, + data_type.HALF, + data_type.BFLOAT16, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + data_type.FP4_E2M1, + }: + raise ValueError(f"The JAX GEMM + amax API received unsupported C dtype " f"{self.c_desc.dtype}") + return self._op.check_support() + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + ) -> TupleDict: + self.check_support() + self._check_tensor_signature(a_tensor, self.a_desc, mode=self.a_mode) + self._check_tensor_signature(b_tensor, self.b_desc, mode=self.b_mode) + self._check_tensor_signature(sfa_tensor, self.sfa_desc, mode=self.scale_mode) + self._check_tensor_signature(sfb_tensor, self.sfb_desc, mode=self.scale_mode) + max_active_clusters = self._get_max_active_clusters( + self._op.cluster_shape_mn[0] * self._op.cluster_shape_mn[1], + overlap_margin=self.num_cluster_overlap_margin, + ) + + def launch(stream, a, b, sfa, sfb, c, amax): + from .dense_blockscaled_gemm_persistent_amax import ( + Sm100BlockScaledPersistentDenseGemmKernel, + ) + + kernel = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size=self._op.sf_vec_size, + mma_tiler_mn=self._op.mma_tiler_mn, + cluster_shape_mn=self._op.cluster_shape_mn, + ) + kernel( + a, + b, + sfa, + sfb, + c, + amax, + max_active_clusters, + stream, + ) + + c_tensor, amax_tensor = self._call_kernel( + (a_tensor, b_tensor, sfa_tensor, sfb_tensor), + launch=launch, + output_descs=(self.c_desc, self.amax_desc), + input_spec=( + self._to_tensor_spec(self.a_desc, mode=self.a_mode), + self._to_tensor_spec(self.b_desc, mode=self.b_mode), + self._to_tensor_spec(self.sfa_desc, mode=self.scale_mode), + self._to_tensor_spec(self.sfb_desc, mode=self.scale_mode), + ), + output_spec=( + self._to_tensor_spec(self.c_desc, mode=self.output_mode), + self._to_tensor_spec(self.amax_desc), + ), + compile_options=compile_options_for_target(self.compute_capability), + ) + return TupleDict(c_tensor=c_tensor, amax_tensor=amax_tensor) + + +@partial( + jax.jit, + static_argnames=( + "c_layout", + "c_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "a_layout", + "b_layout", + ), +) +def gemm_amax_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + c_layout: str = "LMN", + c_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: tuple[int, int] = (1, 1), + sf_vec_size: int = 32, + *, + a_layout: str = "LMK", + b_layout: str = "LNK", +) -> TupleDict: + """Compute a block-scaled GEMM and global max-absolute reduction.""" + + return GemmAmaxSm100( + jax.ShapeDtypeStruct(a_tensor.shape, a_tensor.dtype), + jax.ShapeDtypeStruct(b_tensor.shape, b_tensor.dtype), + jax.ShapeDtypeStruct(sfa_tensor.shape, sfa_tensor.dtype), + jax.ShapeDtypeStruct(sfb_tensor.shape, sfb_tensor.dtype), + c_layout=c_layout, + c_dtype=c_dtype, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + a_layout=a_layout, + b_layout=b_layout, + )(a_tensor, b_tensor, sfa_tensor, sfb_tensor) + + +__all__ = [ + "GemmAmaxSm100", + "SUPPORTED_COMPUTE_CAPABILITIES", + "gemm_amax_wrapper_sm100", +] diff --git a/python/cudnn/gemm_amax/op.py b/python/cudnn/gemm_amax/op.py new file mode 100644 index 000000000..f05b09f89 --- /dev/null +++ b/python/cudnn/gemm_amax/op.py @@ -0,0 +1,236 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral block-scaled dense GEMM + amax operation.""" + +from __future__ import annotations + +from typing import Any, Iterable, Optional + +from .. import data_type +from .._dense_gemm import ( + block_scale_shape, + data_type_bits, + require_block_scale_layout, + require_cluster_shape, + require_compact_major, + require_gemm_inputs, + require_mma_tiler, + require_tensor_shape, +) +from .._op import Op +from .._tensor_desc import TensorDesc + +_AB_DTYPES = frozenset( + { + data_type.FP4_E2M1, + data_type.UINT8, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + } +) +_SCALE_DTYPES = frozenset( + { + data_type.FP8_E4M3, + data_type.FP8_E8M0, + data_type.INT8, + } +) +_C_DTYPES = frozenset( + { + data_type.FLOAT, + data_type.HALF, + data_type.BFLOAT16, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + data_type.FP4_E2M1, + data_type.UINT8, + } +) +_WIDE_C_DTYPES = frozenset( + { + data_type.FLOAT, + data_type.HALF, + data_type.BFLOAT16, + } +) +_FP8_DTYPES = frozenset({data_type.FP8_E4M3, data_type.FP8_E5M2}) +_FP4_STORAGE_DTYPES = frozenset({data_type.FP4_E2M1, data_type.UINT8}) +_E8M0_STORAGE_DTYPES = frozenset({data_type.FP8_E8M0, data_type.INT8}) + + +def _require_supported_dtype( + label: str, + dtype: data_type, + supported: Iterable[data_type], +) -> None: + if dtype not in supported: + raise ValueError(f"{label} has unsupported dtype {dtype}") + + +def _require_logical_alignment( + tensor: TensorDesc[Any], + logical_dtype: data_type, +) -> None: + bits = 4 if logical_dtype == data_type.FP4_E2M1 else data_type_bits(logical_dtype) + extent = tensor.shape[tensor.stride_order[0]] + if extent * bits % 128: + required_multiple = 128 // bits + raise ValueError( + f"{tensor.name or 'Tensor'} contiguous extent must be a multiple of " f"{required_multiple} elements for 16-byte alignment, got {extent}" + ) + + +class GemmAmaxSm100Op(Op): + """Complete logical signature and static configuration for GEMM + amax.""" + + def __init__( + self, + *, + a: TensorDesc[Any], + b: TensorDesc[Any], + sfa: TensorDesc[Any], + sfb: TensorDesc[Any], + c: TensorDesc[Any], + amax: TensorDesc[Any], + acc_dtype: data_type = data_type.FLOAT, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: tuple[int, int] = (1, 1), + sf_vec_size: int = 32, + ) -> None: + for name, desc in ( + ("a", a), + ("b", b), + ("sfa", sfa), + ("sfb", sfb), + ("c", c), + ("amax", amax), + ): + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + if not isinstance(acc_dtype, data_type): + raise TypeError(f"acc_dtype must be a cudnn.data_type, got {type(acc_dtype).__name__}") + + self.a = a + self.b = b + self.sfa = sfa + self.sfb = sfb + self.c = c + self.amax = amax + self.acc_dtype = acc_dtype + self.requested_mma_tiler_mn = mma_tiler_mn + self.requested_cluster_shape_mn = cluster_shape_mn + self.sf_vec_size = sf_vec_size + + self.m: Optional[int] = None + self.n: Optional[int] = None + self.k: Optional[int] = None + self.l: Optional[int] = None + self.a_major: Optional[str] = None + self.b_major: Optional[str] = None + self.c_major: Optional[str] = None + self.ab_dtype: Optional[data_type] = None + self.scale_dtype: Optional[data_type] = None + self.c_dtype: Optional[data_type] = None + self.mma_tiler_mn: Optional[tuple[int, int]] = None + self.cluster_shape_mn: Optional[tuple[int, int]] = None + + def check_support(self) -> bool: + """Validate the complete tensor signature and static configuration.""" + + self.m = self.n = self.k = self.l = None + self.a_major = self.b_major = self.c_major = None + self.ab_dtype = self.scale_dtype = self.c_dtype = None + self.mma_tiler_mn = self.cluster_shape_mn = None + + m, n, k, batch = require_gemm_inputs(self.a, self.b) + require_tensor_shape(self.c, (m, n, batch), label="C") + require_tensor_shape(self.amax, (1, 1, 1), label="Amax") + + if self.a.cudnn_dtype != self.b.cudnn_dtype: + raise ValueError(f"A and B must have the same dtype, got {self.a.dtype} and {self.b.dtype}") + ab_storage_dtype = self.a.cudnn_dtype + _require_supported_dtype("A and B", ab_storage_dtype, _AB_DTYPES) + ab_dtype = data_type.FP4_E2M1 if ab_storage_dtype in _FP4_STORAGE_DTYPES else ab_storage_dtype + + if self.sfa.cudnn_dtype != self.sfb.cudnn_dtype: + raise ValueError(f"SFA and SFB must have the same dtype, got {self.sfa.dtype} and {self.sfb.dtype}") + scale_storage_dtype = self.sfa.cudnn_dtype + _require_supported_dtype("SFA and SFB", scale_storage_dtype, _SCALE_DTYPES) + scale_dtype = data_type.FP8_E8M0 if scale_storage_dtype in _E8M0_STORAGE_DTYPES else scale_storage_dtype + + if self.sf_vec_size not in (16, 32): + raise ValueError(f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}") + if scale_dtype == data_type.FP8_E4M3 and self.sf_vec_size == 32: + raise ValueError("FP8 E4M3 scale factors do not support sf_vec_size=32") + if ab_dtype in _FP8_DTYPES and self.sf_vec_size == 16: + raise ValueError("FP8 A and B do not support sf_vec_size=16") + require_tensor_shape( + self.sfa, + block_scale_shape(m, k, batch, self.sf_vec_size), + label="SFA", + ) + require_tensor_shape( + self.sfb, + block_scale_shape(n, k, batch, self.sf_vec_size), + label="SFB", + ) + require_block_scale_layout(self.sfa, "SFA") + require_block_scale_layout(self.sfb, "SFB") + + c_storage_dtype = self.c.cudnn_dtype + _require_supported_dtype("C", c_storage_dtype, _C_DTYPES) + c_dtype = data_type.FP4_E2M1 if c_storage_dtype in _FP4_STORAGE_DTYPES else c_storage_dtype + if self.amax.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"Amax must have dtype float32, got {self.amax.dtype}") + if self.acc_dtype != data_type.FLOAT: + raise ValueError(f"Accumulator dtype must be float32, got {self.acc_dtype}") + + a_major = require_compact_major(self.a, "m", "k") + b_major = require_compact_major(self.b, "n", "k") + c_major = require_compact_major(self.c, "m", "n") + + if ab_dtype == data_type.FP4_E2M1 and (a_major, b_major) != ("k", "k"): + raise ValueError("FP4 A and B require k-major layouts, got " f"{a_major}-major and {b_major}-major") + if c_dtype == data_type.FP4_E2M1 and c_major != "n": + raise ValueError(f"FP4 C requires n-major layout, got {c_major}-major") + if c_dtype == data_type.FP4_E2M1 and ab_dtype != data_type.FP4_E2M1: + raise ValueError("FP4 C requires FP4 A and B") + if ab_dtype in _FP8_DTYPES and c_dtype in _FP8_DTYPES: + raise NotImplementedError("FP8 A and B with FP8 C are unsupported because the kernel fails to launch") + + mma_tiler_mn = require_mma_tiler( + self.requested_mma_tiler_mn, + allowed_m=(128, 256), + allowed_n=(128, 256), + ) + if mma_tiler_mn[0] == 256: + raise NotImplementedError("mma_tiler_mn[0] == 256 currently hangs") + cluster_shape_mn = require_cluster_shape( + self.requested_cluster_shape_mn, + cta_group_size=2 if mma_tiler_mn[0] == 256 else 1, + max_cluster_size=16, + ) + if cluster_shape_mn[0] > 4 or cluster_shape_mn[1] > 4: + raise ValueError("cluster_shape_mn entries must be at most 4 for scale-factor multicast, " f"got {cluster_shape_mn}") + + if ab_dtype == data_type.FP4_E2M1 and mma_tiler_mn[1] == 256 and k <= 128: + raise ValueError(f"mma_tiler_mn (X, 256) requires K > 128 for FP4, got {k}") + if mma_tiler_mn == (128, 256) and self.sf_vec_size == 16 and c_dtype in _WIDE_C_DTYPES: + raise NotImplementedError("mma_tiler_mn (128, 256), sf_vec_size=16, and a 16/32-bit C " "dtype are unsupported because the kernel fails to launch") + + _require_logical_alignment(self.a, ab_dtype) + _require_logical_alignment(self.b, ab_dtype) + _require_logical_alignment(self.c, c_dtype) + + self.m, self.n, self.k, self.l = m, n, k, batch + self.a_major, self.b_major, self.c_major = a_major, b_major, c_major + self.ab_dtype = ab_dtype + self.scale_dtype = scale_dtype + self.c_dtype = c_dtype + self.mma_tiler_mn = mma_tiler_mn + self.cluster_shape_mn = cluster_shape_mn + return True + + +__all__ = ["GemmAmaxSm100Op", "block_scale_shape"] diff --git a/python/cudnn/gemm_dsrelu/__init__.py b/python/cudnn/gemm_dsrelu/__init__.py index ade25458b..5fb97a9eb 100644 --- a/python/cudnn/gemm_dsrelu/__init__.py +++ b/python/cudnn/gemm_dsrelu/__init__.py @@ -1,9 +1,18 @@ -from .api import ( - GemmDsreluSm100, - gemm_dsrelu_wrapper_sm100, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy Torch API and framework-neutral GEMM + dsReLU exports.""" -__all__ = [ - "GemmDsreluSm100", - "gemm_dsrelu_wrapper_sm100", -] +from .._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ("GemmDsreluSm100Op",), + "api": ( + "GemmDsreluSm100", + "gemm_dsrelu_wrapper_sm100", + ), + }, + submodules=("api", "op"), +) diff --git a/python/cudnn/gemm_dsrelu/api.py b/python/cudnn/gemm_dsrelu/api.py index fb02fa027..d2e71490c 100644 --- a/python/cudnn/gemm_dsrelu/api.py +++ b/python/cudnn/gemm_dsrelu/api.py @@ -118,6 +118,10 @@ def check_support(self) -> bool: dtype=[torch.float16, torch.bfloat16, torch.float32, torch.float8_e4m3fn, torch.float8_e5m2], name="D", ) + self._not_implemented_error_if( + self._is_fp8(self.d_dtype), + "FP8 D output is unavailable because the current dsReLU kernel does not implement SFD generation", + ) self._check_dtype(self.prob_desc, dtype=torch.float32, name="prob") self._check_dtype(self.dprob_desc, dtype=torch.float32, name="dprob") @@ -128,9 +132,6 @@ def check_support(self) -> bool: self._check_dtype(self.acc_dtype, dtype=torch.float32, name="Accumulator") self._value_error_if(self.sf_vec_size not in {16, 32}, f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}") - self._value_error_if( - self._is_fp8(self.d_desc) and (self.sfd_desc is None or self.norm_const_desc is None), "sfd and norm_const are required when D is FP8" - ) self._value_error_if( self._is_fp4x2(self.ab_dtype) and self.d_dtype in {torch.float8_e4m3fn, torch.float8_e5m2}, "FP4 input with FP8 output is not supported" ) diff --git a/python/cudnn/gemm_dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.py b/python/cudnn/gemm_dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.py index 66f7d6ddb..9712c71bd 100644 --- a/python/cudnn/gemm_dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.py +++ b/python/cudnn/gemm_dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.py @@ -29,8 +29,6 @@ from typing import Type, Tuple, Union, Optional import cuda.bindings.driver as cuda -import torch - import cutlass import cutlass.cute as cute from cutlass.cute.nvgpu import cpasync, tcgen05 diff --git a/python/cudnn/gemm_dsrelu/jax.py b/python/cudnn/gemm_dsrelu/jax.py new file mode 100644 index 000000000..b809958d9 --- /dev/null +++ b/python/cudnn/gemm_dsrelu/jax.py @@ -0,0 +1,325 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX API for block-scaled dense GEMM + squared-ReLU backward.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any + +import jax +import jax.numpy as jnp + +from .. import data_type +from .._cute_compiler import compile_options_for_target +from .._dense_gemm import require_gemm_inputs +from .._jax import JaxApiBase, JaxTensorDesc, TupleDict +from .._jax.datatypes import jax_to_cudnn_dtype, normalize_jax_dtype +from .._jax.gemm import BLOCK_SCALE_MODE, PROBABILITY_MODE, gemm_a_mode, gemm_b_mode, gemm_output_mode +from .._jax.layout import to_public_axes +from .op import GemmDsreluSm100Op + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) +_JAX_INPUT_DTYPES = frozenset({data_type.FP4_E2M1, data_type.FP8_E4M3, data_type.FP8_E5M2}) +_WIDE_OUTPUT_DTYPES = frozenset({data_type.HALF, data_type.BFLOAT16, data_type.FLOAT}) + + +class GemmDsreluSm100(JaxApiBase): + """JAX callable specialized from a block-scaled squared-ReLU backward signature.""" + + def __init__( + self, + sample_a: Any, + sample_b: Any, + sample_c: Any, + sample_sfa: Any, + sample_sfb: Any, + sample_prob: Any, + *, + sample_d: Any | None = None, + sample_dprob: Any | None = None, + sample_norm_const: Any | None = None, + alpha: float = 1.0, + d_layout: str = "LMN", + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + a_layout: str = "LMK", + b_layout: str = "LNK", + ) -> None: + self.a_mode = gemm_a_mode(a_layout) + self.b_mode = gemm_b_mode(b_layout) + self.output_mode = gemm_output_mode(d_layout, name="d_layout") + self.scale_mode = BLOCK_SCALE_MODE + self.probability_mode = PROBABILITY_MODE + + self.compute_capability = self._resolve_compute_capability( + None, + SUPPORTED_COMPUTE_CAPABILITIES, + "GemmDsreluSm100", + ) + self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=self.a_mode) + self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=self.b_mode) + self.c_desc = self._to_tensor_desc(sample_c, "sample_c", mode=self.output_mode) + self.sfa_desc = self._to_tensor_desc(sample_sfa, "sample_sfa", mode=self.scale_mode) + self.sfb_desc = self._to_tensor_desc(sample_sfb, "sample_sfb", mode=self.scale_mode) + self.prob_desc = self._to_tensor_desc(sample_prob, "sample_prob", mode=self.probability_mode) + + self.acc_dtype = normalize_jax_dtype(acc_dtype, jnp.float32, "acc_dtype") + resolved_d_dtype = normalize_jax_dtype(d_dtype, jnp.bfloat16, "d_dtype") + if (sample_d is None) != (sample_dprob is None): + raise ValueError("sample_d and sample_dprob must be provided together") + if sample_d is None: + self.d_desc, self.dprob_desc = self._default_output_descs(resolved_d_dtype) + else: + self.d_desc = self._to_tensor_desc(sample_d, "sample_d", mode=self.output_mode) + self.dprob_desc = self._to_tensor_desc( + sample_dprob, + "sample_dprob", + mode=self.probability_mode, + init_value=0.0, + ) + self._check_requested_output_dtype(d_dtype, self.d_desc, "d_dtype") + + self.amax_desc = None + if self.a_desc.cudnn_dtype == data_type.FP4_E2M1 and self.d_desc.cudnn_dtype in _WIDE_OUTPUT_DTYPES: + self.amax_desc = self.d_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(1,), + name="amax_tensor", + init_value=float("-inf"), + ) + self.norm_const_desc = None + self.sfd_desc = None + if self.d_desc.cudnn_dtype in {data_type.FP8_E4M3, data_type.FP8_E5M2}: + raise NotImplementedError("FP8 D output is unavailable because the current dsReLU kernel does not implement SFD generation") + if sample_norm_const is not None: + raise ValueError("sample_norm_const is only used with FP8 D output, which is not implemented") + + acc_cudnn_dtype = jax_to_cudnn_dtype(self.acc_dtype) + if acc_cudnn_dtype == data_type.NOT_SET: + raise ValueError(f"Unsupported JAX accumulator dtype {self.acc_dtype}") + if not isinstance(vector_f32, bool): + raise TypeError(f"vector_f32 must be a bool, got {type(vector_f32).__name__}") + + self._op = GemmDsreluSm100Op( + a=self.a_desc, + b=self.b_desc, + c=self.c_desc, + d=self.d_desc, + sfa=self.sfa_desc, + sfb=self.sfb_desc, + prob=self.prob_desc, + dprob=self.dprob_desc, + sfd=self.sfd_desc, + amax=self.amax_desc, + norm_const=self.norm_const_desc, + alpha=alpha, + acc_dtype=acc_cudnn_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + ) + self.vector_f32 = vector_f32 + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + + def _default_output_descs(self, d_dtype: Any) -> tuple[JaxTensorDesc, JaxTensorDesc]: + m, n, _, batch = require_gemm_inputs(self.a_desc, self.b_desc) + d_shape = to_public_axes((m, n, batch), self.output_mode) + dprob_shape = to_public_axes((m, 1, batch), self.probability_mode) + return ( + self._to_tensor_desc( + jax.ShapeDtypeStruct(d_shape, d_dtype), + "sample_d", + mode=self.output_mode, + ), + self._to_tensor_desc( + jax.ShapeDtypeStruct(dprob_shape, jnp.float32), + "sample_dprob", + mode=self.probability_mode, + init_value=0.0, + ), + ) + + @staticmethod + def _check_requested_output_dtype(requested: Any | None, desc: JaxTensorDesc, name: str) -> None: + if requested is None: + return + requested_dtype = normalize_jax_dtype(requested, requested, name) + actual_dtype = jnp.dtype(desc.dtype) + if requested_dtype != actual_dtype: + raise ValueError(f"{name}={requested_dtype} does not match the explicit sample dtype {actual_dtype}") + + def check_support(self) -> bool: + self._op.check_support() + if self.a_desc.cudnn_dtype not in _JAX_INPUT_DTYPES: + raise NotImplementedError("The JAX GEMM + dsReLU API supports native FP4 and FP8 inputs") + if self.c_desc.cudnn_dtype not in _WIDE_OUTPUT_DTYPES or self.d_desc.cudnn_dtype not in _WIDE_OUTPUT_DTYPES: + raise NotImplementedError("The JAX GEMM + dsReLU API received an unsupported output dtype") + return True + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + prob_tensor: Any, + norm_const_tensor: Any | None = None, + ) -> TupleDict: + self.check_support() + for value, desc, mode in ( + (a_tensor, self.a_desc, self.a_mode), + (b_tensor, self.b_desc, self.b_mode), + (c_tensor, self.c_desc, self.output_mode), + (sfa_tensor, self.sfa_desc, self.scale_mode), + (sfb_tensor, self.sfb_desc, self.scale_mode), + (prob_tensor, self.prob_desc, self.probability_mode), + ): + self._check_tensor_signature(value, desc, mode=mode) + if norm_const_tensor is not None: + raise ValueError("norm_const_tensor is only used with FP8 D output, which is not implemented") + + max_active_clusters = self._get_max_active_clusters( + self._op.cluster_shape_mn[0] * self._op.cluster_shape_mn[1], + overlap_margin=self.num_cluster_overlap_margin, + ) + + def launch(stream, a, b, c, sfa, sfb, prob, d, dprob, *optional_outputs): + expected_optional_outputs = int(self.amax_desc is not None) + if len(optional_outputs) != expected_optional_outputs: + raise RuntimeError(f"GemmDsreluSm100 received {len(optional_outputs)} optional output buffers; expected {expected_optional_outputs}") + amax = optional_outputs[0] if self.amax_desc is not None else None + + import cutlass + import cutlass.cute as cute + + from .dense_blockscaled_gemm_persistent_dsrelu_quant import Sm100BlockScaledPersistentDenseGemmKernel + + kernel = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size=self._op.sf_vec_size, + mma_tiler_mn=self._op.mma_tiler_mn, + cluster_shape_mn=self._op.cluster_shape_mn, + vector_f32=self.vector_f32, + ) + + def squared_relu_backward(x, upstream): + return cute.where(x > 0, x, cute.full_like(x, 0)) * 2 * upstream + + kernel( + a, + b, + sfa, + sfb, + c, + d, + prob, + dprob, + amax, + None, + None, + cutlass.Float32(self._op.alpha), + max_active_clusters, + stream, + epilogue_op=squared_relu_backward, + ) + + optional_output_descs = () if self.amax_desc is None else (self.amax_desc,) + output_descs = (self.d_desc, self.dprob_desc) + optional_output_descs + output_specs = ( + self._to_tensor_spec(self.d_desc, mode=self.output_mode), + self._to_tensor_spec(self.dprob_desc, mode=self.probability_mode), + ) + tuple(self._to_tensor_spec(desc) for desc in optional_output_descs) + results = self._call_kernel( + (a_tensor, b_tensor, c_tensor, sfa_tensor, sfb_tensor, prob_tensor), + launch=launch, + output_descs=output_descs, + input_spec=( + self._to_tensor_spec(self.a_desc, mode=self.a_mode), + self._to_tensor_spec(self.b_desc, mode=self.b_mode), + self._to_tensor_spec(self.c_desc, mode=self.output_mode), + self._to_tensor_spec(self.sfa_desc, mode=self.scale_mode), + self._to_tensor_spec(self.sfb_desc, mode=self.scale_mode), + self._to_tensor_spec(self.prob_desc, mode=self.probability_mode), + ), + output_spec=output_specs, + compile_options=compile_options_for_target(self.compute_capability), + ) + d_tensor, dprob_tensor, *optional_results = results + return TupleDict( + d_tensor=d_tensor, + dprob_tensor=dprob_tensor, + amax_tensor=optional_results[0] if self.amax_desc is not None else None, + sfd_tensor=None, + ) + + +@partial( + jax.jit, + static_argnames=( + "alpha", + "d_layout", + "d_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "a_layout", + "b_layout", + ), +) +def gemm_dsrelu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + prob_tensor: Any, + alpha: float = 1.0, + d_layout: str = "LMN", + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + norm_const_tensor: Any | None = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + *, + a_layout: str = "LMK", + b_layout: str = "LNK", +) -> TupleDict: + """Compute the block-scaled squared-ReLU backward fusion.""" + + return GemmDsreluSm100( + jax.ShapeDtypeStruct(a_tensor.shape, a_tensor.dtype), + jax.ShapeDtypeStruct(b_tensor.shape, b_tensor.dtype), + jax.ShapeDtypeStruct(c_tensor.shape, c_tensor.dtype), + jax.ShapeDtypeStruct(sfa_tensor.shape, sfa_tensor.dtype), + jax.ShapeDtypeStruct(sfb_tensor.shape, sfb_tensor.dtype), + jax.ShapeDtypeStruct(prob_tensor.shape, prob_tensor.dtype), + sample_norm_const=(None if norm_const_tensor is None else jax.ShapeDtypeStruct(norm_const_tensor.shape, norm_const_tensor.dtype)), + alpha=alpha, + d_layout=d_layout, + d_dtype=d_dtype, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + a_layout=a_layout, + b_layout=b_layout, + )(a_tensor, b_tensor, c_tensor, sfa_tensor, sfb_tensor, prob_tensor, norm_const_tensor) + + +__all__ = [ + "GemmDsreluSm100", + "SUPPORTED_COMPUTE_CAPABILITIES", + "gemm_dsrelu_wrapper_sm100", +] diff --git a/python/cudnn/gemm_dsrelu/op.py b/python/cudnn/gemm_dsrelu/op.py new file mode 100644 index 000000000..6046628ab --- /dev/null +++ b/python/cudnn/gemm_dsrelu/op.py @@ -0,0 +1,21 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral dense GEMM + squared-ReLU backward operation.""" + +from .. import data_type +from .._gemm_relu import BlockScaledGemmReluSm100Op + + +class GemmDsreluSm100Op(BlockScaledGemmReluSm100Op): + """Logical signature for block-scaled GEMM + squared-ReLU backward.""" + + def check_support(self) -> bool: + if self.dprob is None: + raise ValueError("dprob is required by the dsReLU backward signature") + if self.c.cudnn_dtype not in {data_type.HALF, data_type.BFLOAT16, data_type.FLOAT}: + raise ValueError(f"dsReLU C must use float16, bfloat16, or float32, got {self.c.dtype}") + return super().check_support() + + +__all__ = ["GemmDsreluSm100Op"] diff --git a/python/cudnn/gemm_srelu/__init__.py b/python/cudnn/gemm_srelu/__init__.py index 29324f369..11913e2a5 100644 --- a/python/cudnn/gemm_srelu/__init__.py +++ b/python/cudnn/gemm_srelu/__init__.py @@ -1,9 +1,18 @@ -from .api import ( - GemmSreluSm100, - gemm_srelu_wrapper_sm100, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy Torch API and framework-neutral GEMM + sReLU exports.""" -__all__ = [ - "GemmSreluSm100", - "gemm_srelu_wrapper_sm100", -] +from .._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ("GemmSreluSm100Op",), + "api": ( + "GemmSreluSm100", + "gemm_srelu_wrapper_sm100", + ), + }, + submodules=("api", "op"), +) diff --git a/python/cudnn/gemm_srelu/api.py b/python/cudnn/gemm_srelu/api.py index 3bf21f95f..90ed43f44 100644 --- a/python/cudnn/gemm_srelu/api.py +++ b/python/cudnn/gemm_srelu/api.py @@ -117,6 +117,10 @@ def check_support(self) -> bool: dtype=[torch.float16, torch.bfloat16, torch.float32, torch.float8_e4m3fn, torch.float8_e5m2], name="D", ) + self._not_implemented_error_if( + self._is_fp8(self.d_dtype), + "FP8 D output is unavailable because the current sReLU kernel does not implement SFD generation", + ) self._check_dtype(self.prob_desc, dtype=torch.float32, name="prob") self.sf_dtype = self._check_dtype(self.sfa_desc, dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn], name="SFA") @@ -126,9 +130,6 @@ def check_support(self) -> bool: self._check_dtype(self.acc_dtype, dtype=torch.float32, name="Accumulator") self._value_error_if(self.sf_vec_size not in {16, 32}, f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}") - self._value_error_if( - self._is_fp8(self.d_desc) and (self.sfd_desc is None or self.norm_const_desc is None), "sfd and norm_const are required when D is FP8" - ) self._value_error_if( self._is_fp4x2(self.ab_dtype) and self.d_dtype in {torch.float8_e4m3fn, torch.float8_e5m2}, "FP4 input with FP8 output is not supported" ) diff --git a/python/cudnn/gemm_srelu/dense_blockscaled_gemm_persistent_srelu_quant.py b/python/cudnn/gemm_srelu/dense_blockscaled_gemm_persistent_srelu_quant.py index 1b338ba70..1ef5eb36a 100644 --- a/python/cudnn/gemm_srelu/dense_blockscaled_gemm_persistent_srelu_quant.py +++ b/python/cudnn/gemm_srelu/dense_blockscaled_gemm_persistent_srelu_quant.py @@ -29,8 +29,6 @@ from typing import Type, Tuple, Union, Optional import cuda.bindings.driver as cuda -import torch - import cutlass import cutlass.cute as cute from cutlass.cute.nvgpu import cpasync, tcgen05 @@ -1375,15 +1373,18 @@ def kernel( self.epilog_sync_barrier.arrive_and_wait() + # + # Apply squared ReLU and the row probability to D. + # + acc_values = acc_vec.load() + acc_values = epilogue_op(acc_values) + acc_values = acc_values * mProb + acc_vec.store(acc_values) + # # Generate amax # if cutlass.const_expr(self.generate_amax): - acc_values = acc_vec.load() - acc_values = epilogue_op(acc_values) - acc_values = acc_values * mProb - acc_vec.store(acc_values) - # Apply element-wise absolute value using math.absf (supports vectors) abs_acc_values_ir = cutlass._mlir.dialects.math.absf(acc_values.ir_value()) # operand (positional) abs_acc_values = type(acc_values)(abs_acc_values_ir, acc_values.shape, acc_values.dtype) diff --git a/python/cudnn/gemm_srelu/jax.py b/python/cudnn/gemm_srelu/jax.py new file mode 100644 index 000000000..75bfd8b06 --- /dev/null +++ b/python/cudnn/gemm_srelu/jax.py @@ -0,0 +1,319 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX API for block-scaled dense GEMM + squared ReLU.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any + +import jax +import jax.numpy as jnp + +from .. import data_type +from .._cute_compiler import compile_options_for_target +from .._dense_gemm import require_gemm_inputs +from .._jax import JaxApiBase, JaxTensorDesc, TupleDict +from .._jax.datatypes import jax_to_cudnn_dtype, normalize_jax_dtype +from .._jax.gemm import BLOCK_SCALE_MODE, PROBABILITY_MODE, gemm_a_mode, gemm_b_mode, gemm_output_mode +from .._jax.layout import to_public_axes +from .op import GemmSreluSm100Op + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) +_JAX_INPUT_DTYPES = frozenset({data_type.FP4_E2M1, data_type.FP8_E4M3, data_type.FP8_E5M2}) +_WIDE_OUTPUT_DTYPES = frozenset({data_type.HALF, data_type.BFLOAT16, data_type.FLOAT}) +_JAX_C_DTYPES = _WIDE_OUTPUT_DTYPES | {data_type.FP8_E4M3, data_type.FP8_E5M2} + + +class GemmSreluSm100(JaxApiBase): + """JAX callable specialized from a block-scaled GEMM + sReLU signature.""" + + def __init__( + self, + sample_a: Any, + sample_b: Any, + sample_sfa: Any, + sample_sfb: Any, + sample_prob: Any, + *, + sample_c: Any | None = None, + sample_d: Any | None = None, + sample_norm_const: Any | None = None, + alpha: float = 1.0, + c_layout: str = "LMN", + c_dtype: Any | None = None, + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + a_layout: str = "LMK", + b_layout: str = "LNK", + ) -> None: + self.a_mode = gemm_a_mode(a_layout) + self.b_mode = gemm_b_mode(b_layout) + self.output_mode = gemm_output_mode(c_layout) + self.scale_mode = BLOCK_SCALE_MODE + self.probability_mode = PROBABILITY_MODE + + self.compute_capability = self._resolve_compute_capability( + None, + SUPPORTED_COMPUTE_CAPABILITIES, + "GemmSreluSm100", + ) + self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=self.a_mode) + self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=self.b_mode) + self.sfa_desc = self._to_tensor_desc(sample_sfa, "sample_sfa", mode=self.scale_mode) + self.sfb_desc = self._to_tensor_desc(sample_sfb, "sample_sfb", mode=self.scale_mode) + self.prob_desc = self._to_tensor_desc(sample_prob, "sample_prob", mode=self.probability_mode) + + self.acc_dtype = normalize_jax_dtype(acc_dtype, jnp.float32, "acc_dtype") + resolved_c_dtype = normalize_jax_dtype(c_dtype, jnp.bfloat16, "c_dtype") + resolved_d_dtype = normalize_jax_dtype(d_dtype, jnp.bfloat16, "d_dtype") + if (sample_c is None) != (sample_d is None): + raise ValueError("sample_c and sample_d must be provided together") + if sample_c is None: + self.c_desc, self.d_desc = self._default_output_descs( + resolved_c_dtype, + resolved_d_dtype, + ) + else: + self.c_desc = self._to_tensor_desc(sample_c, "sample_c", mode=self.output_mode) + self.d_desc = self._to_tensor_desc(sample_d, "sample_d", mode=self.output_mode) + self._check_requested_output_dtype(c_dtype, self.c_desc, "c_dtype") + self._check_requested_output_dtype(d_dtype, self.d_desc, "d_dtype") + + self.amax_desc = None + if self.a_desc.cudnn_dtype == data_type.FP4_E2M1 and self.d_desc.cudnn_dtype in _WIDE_OUTPUT_DTYPES: + self.amax_desc = self.d_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(1,), + name="amax_tensor", + init_value=float("-inf"), + ) + self.norm_const_desc = None + self.sfd_desc = None + if self.d_desc.cudnn_dtype in {data_type.FP8_E4M3, data_type.FP8_E5M2}: + raise NotImplementedError("FP8 D output is unavailable because the current sReLU kernel does not implement SFD generation") + if sample_norm_const is not None: + raise ValueError("sample_norm_const is only used with FP8 D output, which is not implemented") + + acc_cudnn_dtype = jax_to_cudnn_dtype(self.acc_dtype) + if acc_cudnn_dtype == data_type.NOT_SET: + raise ValueError(f"Unsupported JAX accumulator dtype {self.acc_dtype}") + if not isinstance(vector_f32, bool): + raise TypeError(f"vector_f32 must be a bool, got {type(vector_f32).__name__}") + + self._op = GemmSreluSm100Op( + a=self.a_desc, + b=self.b_desc, + c=self.c_desc, + d=self.d_desc, + sfa=self.sfa_desc, + sfb=self.sfb_desc, + prob=self.prob_desc, + sfd=self.sfd_desc, + amax=self.amax_desc, + norm_const=self.norm_const_desc, + alpha=alpha, + acc_dtype=acc_cudnn_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + ) + self.vector_f32 = vector_f32 + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + + def _default_output_descs(self, c_dtype: Any, d_dtype: Any) -> tuple[JaxTensorDesc, JaxTensorDesc]: + m, n, _, batch = require_gemm_inputs(self.a_desc, self.b_desc) + public_shape = to_public_axes((m, n, batch), self.output_mode) + return ( + self._to_tensor_desc( + jax.ShapeDtypeStruct(public_shape, c_dtype), + "sample_c", + mode=self.output_mode, + ), + self._to_tensor_desc( + jax.ShapeDtypeStruct(public_shape, d_dtype), + "sample_d", + mode=self.output_mode, + ), + ) + + @staticmethod + def _check_requested_output_dtype(requested: Any | None, desc: JaxTensorDesc, name: str) -> None: + if requested is None: + return + requested_dtype = normalize_jax_dtype(requested, requested, name) + actual_dtype = jnp.dtype(desc.dtype) + if requested_dtype != actual_dtype: + raise ValueError(f"{name}={requested_dtype} does not match the explicit sample dtype {actual_dtype}") + + def check_support(self) -> bool: + self._op.check_support() + if self.a_desc.cudnn_dtype not in _JAX_INPUT_DTYPES: + raise NotImplementedError("The JAX GEMM + sReLU API supports native FP4 and FP8 inputs") + if self.c_desc.cudnn_dtype not in _JAX_C_DTYPES or self.d_desc.cudnn_dtype not in _WIDE_OUTPUT_DTYPES: + raise NotImplementedError("The JAX GEMM + sReLU API received an unsupported output dtype") + return True + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + prob_tensor: Any, + norm_const_tensor: Any | None = None, + ) -> TupleDict: + self.check_support() + for value, desc, mode in ( + (a_tensor, self.a_desc, self.a_mode), + (b_tensor, self.b_desc, self.b_mode), + (sfa_tensor, self.sfa_desc, self.scale_mode), + (sfb_tensor, self.sfb_desc, self.scale_mode), + (prob_tensor, self.prob_desc, self.probability_mode), + ): + self._check_tensor_signature(value, desc, mode=mode) + if norm_const_tensor is not None: + raise ValueError("norm_const_tensor is only used with FP8 D output, which is not implemented") + + max_active_clusters = self._get_max_active_clusters( + self._op.cluster_shape_mn[0] * self._op.cluster_shape_mn[1], + overlap_margin=self.num_cluster_overlap_margin, + ) + + def launch(stream, a, b, sfa, sfb, prob, c, d, *optional_outputs): + expected_optional_outputs = int(self.amax_desc is not None) + if len(optional_outputs) != expected_optional_outputs: + raise RuntimeError(f"GemmSreluSm100 received {len(optional_outputs)} optional output buffers; expected {expected_optional_outputs}") + amax = optional_outputs[0] if self.amax_desc is not None else None + + import cutlass + import cutlass.cute as cute + + from .dense_blockscaled_gemm_persistent_srelu_quant import Sm100BlockScaledPersistentDenseGemmKernel + + kernel = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size=self._op.sf_vec_size, + mma_tiler_mn=self._op.mma_tiler_mn, + cluster_shape_mn=self._op.cluster_shape_mn, + vector_f32=self.vector_f32, + ) + + def squared_relu(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) ** 2 + + kernel( + a, + b, + sfa, + sfb, + c, + d, + prob, + amax, + None, + None, + cutlass.Float32(self._op.alpha), + max_active_clusters, + stream, + epilogue_op=squared_relu, + ) + + optional_output_descs = () if self.amax_desc is None else (self.amax_desc,) + output_descs = (self.c_desc, self.d_desc) + optional_output_descs + output_specs = ( + self._to_tensor_spec(self.c_desc, mode=self.output_mode), + self._to_tensor_spec(self.d_desc, mode=self.output_mode), + ) + tuple(self._to_tensor_spec(desc) for desc in optional_output_descs) + results = self._call_kernel( + (a_tensor, b_tensor, sfa_tensor, sfb_tensor, prob_tensor), + launch=launch, + output_descs=output_descs, + input_spec=( + self._to_tensor_spec(self.a_desc, mode=self.a_mode), + self._to_tensor_spec(self.b_desc, mode=self.b_mode), + self._to_tensor_spec(self.sfa_desc, mode=self.scale_mode), + self._to_tensor_spec(self.sfb_desc, mode=self.scale_mode), + self._to_tensor_spec(self.prob_desc, mode=self.probability_mode), + ), + output_spec=output_specs, + compile_options=compile_options_for_target(self.compute_capability), + ) + c_tensor, d_tensor, *optional_results = results + return TupleDict( + c_tensor=c_tensor, + d_tensor=d_tensor, + amax_tensor=optional_results[0] if self.amax_desc is not None else None, + sfd_tensor=None, + ) + + +@partial( + jax.jit, + static_argnames=( + "alpha", + "c_layout", + "c_dtype", + "d_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "a_layout", + "b_layout", + ), +) +def gemm_srelu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + prob_tensor: Any, + alpha: float = 1.0, + c_layout: str = "LMN", + c_dtype: Any | None = None, + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + norm_const_tensor: Any | None = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + *, + a_layout: str = "LMK", + b_layout: str = "LNK", +) -> TupleDict: + """Compute block-scaled GEMM and ``D = prob * relu(C) ** 2``.""" + + return GemmSreluSm100( + jax.ShapeDtypeStruct(a_tensor.shape, a_tensor.dtype), + jax.ShapeDtypeStruct(b_tensor.shape, b_tensor.dtype), + jax.ShapeDtypeStruct(sfa_tensor.shape, sfa_tensor.dtype), + jax.ShapeDtypeStruct(sfb_tensor.shape, sfb_tensor.dtype), + jax.ShapeDtypeStruct(prob_tensor.shape, prob_tensor.dtype), + sample_norm_const=(None if norm_const_tensor is None else jax.ShapeDtypeStruct(norm_const_tensor.shape, norm_const_tensor.dtype)), + alpha=alpha, + c_layout=c_layout, + c_dtype=c_dtype, + d_dtype=d_dtype, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + a_layout=a_layout, + b_layout=b_layout, + )(a_tensor, b_tensor, sfa_tensor, sfb_tensor, prob_tensor, norm_const_tensor) + + +__all__ = [ + "GemmSreluSm100", + "SUPPORTED_COMPUTE_CAPABILITIES", + "gemm_srelu_wrapper_sm100", +] diff --git a/python/cudnn/gemm_srelu/op.py b/python/cudnn/gemm_srelu/op.py new file mode 100644 index 000000000..538aee587 --- /dev/null +++ b/python/cudnn/gemm_srelu/op.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral dense GEMM + squared-ReLU operation.""" + +from .._gemm_relu import BlockScaledGemmReluSm100Op + + +class GemmSreluSm100Op(BlockScaledGemmReluSm100Op): + """Logical signature for block-scaled GEMM + squared ReLU.""" + + def check_support(self) -> bool: + if self.dprob is not None: + raise ValueError("dprob is only part of the dsReLU backward signature") + return super().check_support() + + +__all__ = ["GemmSreluSm100Op"] diff --git a/python/cudnn/gemm_swiglu/__init__.py b/python/cudnn/gemm_swiglu/__init__.py index 6f11dd595..917c2fa19 100644 --- a/python/cudnn/gemm_swiglu/__init__.py +++ b/python/cudnn/gemm_swiglu/__init__.py @@ -8,7 +8,10 @@ __all__, __getattr__, __dir__ = make_operation_api( globals(), exports={ - "op": ("GemmSwigluSm100Op",), + "op": ( + "BlockScaledGemmSwigluSm100Op", + "GemmSwigluSm100Op", + ), "api": ( "GemmSwigluSm100", "gemm_swiglu_wrapper_sm100", diff --git a/python/cudnn/gemm_swiglu/jax.py b/python/cudnn/gemm_swiglu/jax.py index f47c11dfe..1b138c7f1 100644 --- a/python/cudnn/gemm_swiglu/jax.py +++ b/python/cudnn/gemm_swiglu/jax.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Optional JAX API for the standard dense GEMM + SwiGLU kernel.""" +"""Optional JAX API for standard and block-scaled dense GEMM + SwiGLU.""" from __future__ import annotations @@ -17,9 +17,9 @@ from .._dense_gemm import require_gemm_inputs from .._jax import JaxApiBase, JaxTensorDesc, TupleDict from .._jax.datatypes import jax_to_cudnn_dtype, normalize_jax_dtype -from .._jax.gemm import gemm_a_mode, gemm_b_mode, gemm_output_mode +from .._jax.gemm import BLOCK_SCALE_MODE, gemm_a_mode, gemm_b_mode, gemm_output_mode from .._jax.layout import to_public_axes -from .op import GemmSwigluSm100Op +from .op import BlockScaledGemmSwigluSm100Op, GemmSwigluSm100Op SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) @@ -31,9 +31,9 @@ class GemmSwigluSm100(JaxApiBase): order. Descriptors stored by the operation remain in the kernel's canonical ``MKL``, ``NKL``, and ``MNL`` orders. - This first JAX surface intentionally selects only the standard dense - kernel. Block-scaled inputs and quantized outputs are added separately so - their packed scale-factor ABI does not complicate the base contract. + Providing SFA and SFB samples selects the block-scaled kernel. Public scale + arrays are compact row-major ``[L, row_tiles, k_tiles, 32, 4, 4]`` arrays; + the adapter maps them to the kernel's packed scale-factor layout. """ def __init__( @@ -43,6 +43,11 @@ def __init__( *, sample_ab12: Any | None = None, sample_c: Any | None = None, + sample_sfa: Any | None = None, + sample_sfb: Any | None = None, + sample_sfc: Any | None = None, + sample_amax: Any | None = None, + sample_norm_const: Any | None = None, alpha: float = 1.0, c_layout: str = "LMN", ab12_dtype: Any | None = None, @@ -50,6 +55,9 @@ def __init__( acc_dtype: Any | None = None, mma_tiler_mn: tuple[int, int] = (128, 128), cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + ab12_stages: int = 4, a_layout: str = "LMK", b_layout: str = "LNK", ) -> None: @@ -59,6 +67,8 @@ def __init__( self.a_mode = gemm_a_mode(a_layout) self.b_mode = gemm_b_mode(b_layout) self.output_mode = gemm_output_mode(c_layout) + self.scale_mode = BLOCK_SCALE_MODE + self.sf_vec_size = sf_vec_size self.compute_capability = self._resolve_compute_capability( None, @@ -68,6 +78,29 @@ def __init__( self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=self.a_mode) self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=self.b_mode) self.acc_dtype = normalize_jax_dtype(acc_dtype, jnp.float32, "acc_dtype") + if sample_sfc is not None or sample_norm_const is not None: + raise NotImplementedError("The JAX GEMM + SwiGLU API does not expose FP8 C or SFC because the current supported input combinations reject FP8 C") + + self.is_block_scaled = any( + sample is not None + for sample in ( + sample_sfa, + sample_sfb, + sample_amax, + ) + ) + if self.is_block_scaled and (sample_sfa is None or sample_sfb is None): + raise ValueError("sample_sfa and sample_sfb are required for block-scaled GEMM + SwiGLU") + if not self.is_block_scaled: + if sf_vec_size != 16: + raise ValueError("sf_vec_size only applies to block-scaled GEMM + SwiGLU") + if vector_f32: + raise ValueError("vector_f32 only applies to block-scaled GEMM + SwiGLU") + if ab12_stages != 4: + raise ValueError("ab12_stages only applies to block-scaled GEMM + SwiGLU") + self.sfa_desc = None if sample_sfa is None else self._to_tensor_desc(sample_sfa, "sample_sfa", mode=self.scale_mode) + self.sfb_desc = None if sample_sfb is None else self._to_tensor_desc(sample_sfb, "sample_sfb", mode=self.scale_mode) + self.norm_const_desc = None if (sample_ab12 is None) != (sample_c is None): raise ValueError("sample_ab12 and sample_c must be provided together") @@ -92,20 +125,51 @@ def __init__( self._check_requested_output_dtype(ab12_dtype, self.ab12_desc, "ab12_dtype") self._check_requested_output_dtype(c_dtype, self.c_desc, "c_dtype") + self.sfc_desc = None + self.amax_desc = None + if self.c_desc.cudnn_dtype in {data_type.FP8_E4M3, data_type.FP8_E5M2}: + raise NotImplementedError("The JAX GEMM + SwiGLU API does not support FP8 C output") + if self.is_block_scaled: + needs_amax = self.a_desc.cudnn_dtype == data_type.FP4_E2M1 and self.c_desc.cudnn_dtype == data_type.BFLOAT16 + if sample_amax is not None: + self.amax_desc = self._to_tensor_desc(sample_amax, "sample_amax", init_value=float("-inf")) + elif needs_amax: + self.amax_desc = self._default_amax_desc() + acc_cudnn_dtype = jax_to_cudnn_dtype(self.acc_dtype) if acc_cudnn_dtype == data_type.NOT_SET: raise ValueError(f"Unsupported JAX accumulator dtype {self.acc_dtype}") - self._op = GemmSwigluSm100Op( - a=self.a_desc, - b=self.b_desc, - ab12=self.ab12_desc, - c=self.c_desc, - alpha=alpha, - acc_dtype=acc_cudnn_dtype, - mma_tiler_mn=mma_tiler_mn, - cluster_shape_mn=cluster_shape_mn, - ) + if self.is_block_scaled: + self._op = BlockScaledGemmSwigluSm100Op( + a=self.a_desc, + b=self.b_desc, + sfa=self.sfa_desc, + sfb=self.sfb_desc, + ab12=self.ab12_desc, + c=self.c_desc, + sfc=self.sfc_desc, + amax=self.amax_desc, + norm_const=self.norm_const_desc, + alpha=alpha, + acc_dtype=acc_cudnn_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + ab12_stages=ab12_stages, + ) + else: + self._op = GemmSwigluSm100Op( + a=self.a_desc, + b=self.b_desc, + ab12=self.ab12_desc, + c=self.c_desc, + alpha=alpha, + acc_dtype=acc_cudnn_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + ) self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) def _default_output_descs( @@ -136,6 +200,13 @@ def _default_output_descs( ), ) + def _default_amax_desc(self) -> JaxTensorDesc: + return self._to_tensor_desc( + jax.ShapeDtypeStruct((1,), jnp.float32), + "sample_amax", + init_value=float("-inf"), + ) + @staticmethod def _check_requested_output_dtype( requested: Any | None, @@ -152,66 +223,156 @@ def _check_requested_output_dtype( def check_support(self) -> bool: return self._op.check_support() - def __call__(self, a_tensor: Any, b_tensor: Any) -> TupleDict: + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any | None = None, + sfb_tensor: Any | None = None, + norm_const_tensor: Any | None = None, + ) -> TupleDict: self.check_support() self._check_tensor_signature(a_tensor, self.a_desc, mode=self.a_mode) self._check_tensor_signature(b_tensor, self.b_desc, mode=self.b_mode) - self._max_active_clusters = self._get_max_active_clusters( + if self.is_block_scaled: + if sfa_tensor is None or sfb_tensor is None: + raise ValueError("sfa_tensor and sfb_tensor are required for this block-scaled callable") + self._check_tensor_signature(sfa_tensor, self.sfa_desc, mode=self.scale_mode) + self._check_tensor_signature(sfb_tensor, self.sfb_desc, mode=self.scale_mode) + if norm_const_tensor is not None: + raise NotImplementedError("norm_const_tensor is unavailable because the JAX API does not support FP8 C output") + elif sfa_tensor is not None or sfb_tensor is not None or norm_const_tensor is not None: + raise ValueError("Scale-factor tensors are not part of this standard GEMM + SwiGLU callable") + + max_active_clusters = self._get_max_active_clusters( self._op.cluster_shape_mn[0] * self._op.cluster_shape_mn[1], overlap_margin=self.num_cluster_overlap_margin, ) - ab12_tensor, c_tensor = self._call_kernel( - (a_tensor, b_tensor), - output_descs=(self.ab12_desc, self.c_desc), - input_spec=( - self._to_tensor_spec(self.a_desc, mode=self.a_mode), - self._to_tensor_spec(self.b_desc, mode=self.b_mode), - ), - output_spec=( - self._to_tensor_spec(self.ab12_desc, mode=self.output_mode), - self._to_tensor_spec(self.c_desc, mode=self.output_mode), - ), - compile_options=compile_options_for_target(self.compute_capability), + if not self.is_block_scaled: + + def launch( + stream: Any, + a: Any, + b: Any, + ab12: Any, + c: Any, + ) -> None: + import cutlass + from cutlass.jax import jax_to_cutlass_dtype + + from .dense_gemm_persistent_swiglu import PersistentDenseGemmKernel + + kernel = PersistentDenseGemmKernel( + acc_dtype=jax_to_cutlass_dtype(self.acc_dtype), + use_2cta_instrs=self._op.mma_tiler_mn[0] == 256, + mma_tiler_mn=self._op.mma_tiler_mn, + cluster_shape_mn=self._op.cluster_shape_mn, + ) + kernel( + a, + b, + ab12, + c, + cutlass.Float32(self._op.alpha), + max_active_clusters, + stream, + ) + + ab12_tensor, c_tensor = self._call_kernel( + (a_tensor, b_tensor), + launch=launch, + output_descs=(self.ab12_desc, self.c_desc), + input_spec=( + self._to_tensor_spec(self.a_desc, mode=self.a_mode), + self._to_tensor_spec(self.b_desc, mode=self.b_mode), + ), + output_spec=( + self._to_tensor_spec(self.ab12_desc, mode=self.output_mode), + self._to_tensor_spec(self.c_desc, mode=self.output_mode), + ), + compile_options=compile_options_for_target(self.compute_capability), + ) + return TupleDict( + ab12_tensor=ab12_tensor, + c_tensor=c_tensor, + sfc_tensor=None, + amax_tensor=None, + ) + + kernel_inputs = (a_tensor, b_tensor, sfa_tensor, sfb_tensor) + input_specs = ( + self._to_tensor_spec(self.a_desc, mode=self.a_mode), + self._to_tensor_spec(self.b_desc, mode=self.b_mode), + self._to_tensor_spec(self.sfa_desc, mode=self.scale_mode), + self._to_tensor_spec(self.sfb_desc, mode=self.scale_mode), ) - return TupleDict( - ab12_tensor=ab12_tensor, - c_tensor=c_tensor, - sfc_tensor=None, - amax_tensor=None, + + output_descs = (self.ab12_desc, self.c_desc) + output_specs = ( + self._to_tensor_spec(self.ab12_desc, mode=self.output_mode), + self._to_tensor_spec(self.c_desc, mode=self.output_mode), ) + if self.amax_desc is not None: + output_descs += (self.amax_desc,) + output_specs += (self._to_tensor_spec(self.amax_desc),) - def _launch( - self, - inputs: tuple[Any, ...], - outputs: tuple[Any, ...], - workspaces: tuple[Any, ...], - stream: Any, - ) -> None: - if workspaces: - raise RuntimeError(f"GemmSwigluSm100 received unexpected workspaces: {len(workspaces)}") - a, b = inputs - ab12, c = outputs + def launch_block_scaled( + stream: Any, + a: Any, + b: Any, + sfa: Any, + sfb: Any, + ab12: Any, + c: Any, + *optional_outputs: Any, + ) -> None: + expected_optional_outputs = int(self.amax_desc is not None) + if len(optional_outputs) != expected_optional_outputs: + raise RuntimeError(f"GemmSwigluSm100 received {len(optional_outputs)} optional output buffers; expected {expected_optional_outputs}") + amax = optional_outputs[0] if self.amax_desc is not None else None - import cutlass - from cutlass.jax import jax_to_cutlass_dtype + import cutlass - from .dense_gemm_persistent_swiglu import PersistentDenseGemmKernel + from .dense_blockscaled_gemm_persistent_swiglu_interleaved_quant import Sm100BlockScaledPersistentDenseGemmKernel - kernel = PersistentDenseGemmKernel( - acc_dtype=jax_to_cutlass_dtype(self.acc_dtype), - use_2cta_instrs=self._op.mma_tiler_mn[0] == 256, - mma_tiler_mn=self._op.mma_tiler_mn, - cluster_shape_mn=self._op.cluster_shape_mn, + kernel = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size=self._op.sf_vec_size, + mma_tiler_mn=self._op.mma_tiler_mn, + cluster_shape_mn=self._op.cluster_shape_mn, + vector_f32=self._op.vector_f32, + ab12_stages=self._op.ab12_stages, + ) + kernel( + a, + b, + sfa, + sfb, + c, + ab12, + amax, + None, + None, + cutlass.Float32(self._op.alpha), + max_active_clusters, + stream, + ) + + results = self._call_kernel( + kernel_inputs, + launch=launch_block_scaled, + output_descs=output_descs, + input_spec=input_specs, + output_spec=output_specs, + compile_options=compile_options_for_target(self.compute_capability), ) - kernel( - a, - b, - ab12, - c, - cutlass.Float32(self._op.alpha), - self._max_active_clusters, - stream, + ab12_tensor, c_tensor, *optional_results = results + amax_tensor = optional_results[0] if self.amax_desc is not None else None + return TupleDict( + ab12_tensor=ab12_tensor, + c_tensor=c_tensor, + sfc_tensor=None, + amax_tensor=amax_tensor, ) @@ -225,6 +386,9 @@ def _launch( "acc_dtype", "mma_tiler_mn", "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "ab12_stages", "a_layout", "b_layout", ), @@ -239,15 +403,24 @@ def gemm_swiglu_wrapper_sm100( acc_dtype: Any | None = None, mma_tiler_mn: tuple[int, int] = (128, 128), cluster_shape_mn: tuple[int, int] | None = None, + sfa_tensor: Any | None = None, + sfb_tensor: Any | None = None, + norm_const_tensor: Any | None = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + ab12_stages: int = 4, *, a_layout: str = "LMK", b_layout: str = "LNK", ) -> TupleDict: - """Compute the standard GEMM + SwiGLU projection on a local SM100-family GPU.""" + """Compute standard or block-scaled GEMM + SwiGLU on a local SM100-family GPU.""" return GemmSwigluSm100( jax.ShapeDtypeStruct(a_tensor.shape, a_tensor.dtype), jax.ShapeDtypeStruct(b_tensor.shape, b_tensor.dtype), + sample_sfa=None if sfa_tensor is None else jax.ShapeDtypeStruct(sfa_tensor.shape, sfa_tensor.dtype), + sample_sfb=None if sfb_tensor is None else jax.ShapeDtypeStruct(sfb_tensor.shape, sfb_tensor.dtype), + sample_norm_const=None if norm_const_tensor is None else jax.ShapeDtypeStruct(norm_const_tensor.shape, norm_const_tensor.dtype), alpha=alpha, c_layout=c_layout, ab12_dtype=ab12_dtype, @@ -255,9 +428,12 @@ def gemm_swiglu_wrapper_sm100( acc_dtype=acc_dtype, mma_tiler_mn=mma_tiler_mn, cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + ab12_stages=ab12_stages, a_layout=a_layout, b_layout=b_layout, - )(a_tensor, b_tensor) + )(a_tensor, b_tensor, sfa_tensor, sfb_tensor, norm_const_tensor) __all__ = [ diff --git a/python/cudnn/gemm_swiglu/op.py b/python/cudnn/gemm_swiglu/op.py index 787ef4624..81f5cc423 100644 --- a/python/cudnn/gemm_swiglu/op.py +++ b/python/cudnn/gemm_swiglu/op.py @@ -9,7 +9,9 @@ from .. import data_type from .._dense_gemm import ( + block_scale_shape, require_16_byte_alignment, + require_block_scale_layout, require_cluster_shape, require_compact_major, require_gemm_inputs, @@ -40,6 +42,16 @@ data_type.FP8_E5M2, } _FP8_DTYPES = frozenset({data_type.FP8_E4M3, data_type.FP8_E5M2}) +_BLOCK_SCALED_AB_DTYPES = frozenset({data_type.FP4_E2M1, *_FP8_DTYPES}) +_BLOCK_SCALED_OUTPUT_DTYPES = frozenset( + { + data_type.HALF, + data_type.BFLOAT16, + data_type.FLOAT, + *_FP8_DTYPES, + } +) +_BLOCK_SCALE_DTYPES = frozenset({data_type.FP8_E4M3, data_type.FP8_E8M0}) _SWIGLU_BLOCK_COLUMNS = 32 _SWIGLU_COLUMNS_PER_PAIR = 2 * _SWIGLU_BLOCK_COLUMNS _SWIGLU_MMA_N = tuple(range(_SWIGLU_COLUMNS_PER_PAIR, 257, _SWIGLU_COLUMNS_PER_PAIR)) @@ -162,4 +174,198 @@ def _check_standard_dtypes(self) -> None: raise ValueError(f"C dtype must be float16 or bfloat16, got {self.c.dtype}") -__all__ = ["GemmSwigluSm100Op"] +class BlockScaledGemmSwigluSm100Op(Op): + """Logical signature for block-scaled GEMM + SwiGLU and quantized outputs.""" + + def __init__( + self, + *, + a: TensorDesc[Any], + b: TensorDesc[Any], + sfa: TensorDesc[Any], + sfb: TensorDesc[Any], + ab12: TensorDesc[Any], + c: TensorDesc[Any], + sfc: TensorDesc[Any] | None = None, + amax: TensorDesc[Any] | None = None, + norm_const: TensorDesc[Any] | None = None, + alpha: float = 1.0, + acc_dtype: data_type = data_type.FLOAT, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + ab12_stages: int = 4, + ) -> None: + for name, desc in ( + ("a", a), + ("b", b), + ("sfa", sfa), + ("sfb", sfb), + ("ab12", ab12), + ("c", c), + ): + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + for name, desc in (("sfc", sfc), ("amax", amax), ("norm_const", norm_const)): + if desc is not None and not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc or None, got {type(desc).__name__}") + if not isinstance(acc_dtype, data_type): + raise TypeError(f"acc_dtype must be a cudnn.data_type, got {type(acc_dtype).__name__}") + if isinstance(sf_vec_size, bool) or not isinstance(sf_vec_size, int): + raise TypeError(f"sf_vec_size must be an int, got {type(sf_vec_size).__name__}") + if not isinstance(vector_f32, bool): + raise TypeError(f"vector_f32 must be a bool, got {type(vector_f32).__name__}") + if isinstance(ab12_stages, bool) or not isinstance(ab12_stages, int): + raise TypeError(f"ab12_stages must be an int, got {type(ab12_stages).__name__}") + + self.a = a + self.b = b + self.sfa = sfa + self.sfb = sfb + self.ab12 = ab12 + self.c = c + self.sfc = sfc + self.amax = amax + self.norm_const = norm_const + self.alpha = alpha + self.acc_dtype = acc_dtype + self.requested_mma_tiler_mn = mma_tiler_mn + self.requested_cluster_shape_mn = cluster_shape_mn + self.sf_vec_size = sf_vec_size + self.vector_f32 = vector_f32 + self.ab12_stages = ab12_stages + + self.m: Optional[int] = None + self.n: Optional[int] = None + self.k: Optional[int] = None + self.l: Optional[int] = None + self.output_n: Optional[int] = None + self.a_major: Optional[str] = None + self.b_major: Optional[str] = None + self.output_major: Optional[str] = None + self.ab_dtype: Optional[data_type] = None + self.sf_dtype: Optional[data_type] = None + self.ab12_dtype: Optional[data_type] = None + self.c_dtype: Optional[data_type] = None + self.mma_tiler_mn: Optional[tuple[int, int]] = None + self.cluster_shape_mn: Optional[tuple[int, int]] = None + + def check_support(self) -> bool: + """Validate shapes, packed scale layouts, dtypes, and static tiling.""" + + self.m = self.n = self.k = self.l = self.output_n = None + self.a_major = self.b_major = self.output_major = None + self.ab_dtype = self.sf_dtype = self.ab12_dtype = self.c_dtype = None + self.mma_tiler_mn = self.cluster_shape_mn = None + + m, n, k, batch = require_gemm_inputs(self.a, self.b) + if n % _SWIGLU_COLUMNS_PER_PAIR: + raise ValueError(f"N must be divisible by {_SWIGLU_COLUMNS_PER_PAIR} for SwiGLU block pairs, got {n}") + output_n = n // 2 + require_tensor_shape(self.ab12, (m, n, batch), label="AB12") + require_tensor_shape(self.c, (m, output_n, batch), label="C") + + if self.sf_vec_size not in (16, 32): + raise ValueError(f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}") + require_tensor_shape(self.sfa, block_scale_shape(m, k, batch, self.sf_vec_size), label="SFA") + require_tensor_shape(self.sfb, block_scale_shape(n, k, batch, self.sf_vec_size), label="SFB") + require_block_scale_layout(self.sfa, "SFA") + require_block_scale_layout(self.sfb, "SFB") + + ab_dtype = self.a.cudnn_dtype + if ab_dtype not in _BLOCK_SCALED_AB_DTYPES: + raise ValueError(f"A has unsupported block-scaled dtype {self.a.dtype}") + if self.b.cudnn_dtype != ab_dtype: + raise ValueError(f"A and B must have the same dtype, got {self.a.dtype} and {self.b.dtype}") + + sf_dtype = self.sfa.cudnn_dtype + if sf_dtype not in _BLOCK_SCALE_DTYPES: + raise ValueError(f"SFA has unsupported scale-factor dtype {self.sfa.dtype}") + if self.sfb.cudnn_dtype != sf_dtype: + raise ValueError(f"SFA and SFB must have the same dtype, got {self.sfa.dtype} and {self.sfb.dtype}") + if ab_dtype in _FP8_DTYPES and (sf_dtype, self.sf_vec_size) != (data_type.FP8_E8M0, 32): + raise ValueError("FP8 A and B require FP8_E8M0 scales with sf_vec_size=32") + if ab_dtype == data_type.FP4_E2M1 and (sf_dtype, self.sf_vec_size) == (data_type.FP8_E4M3, 32): + raise ValueError("FP4 A and B do not support FP8_E4M3 scales with sf_vec_size=32") + + ab12_dtype = self.ab12.cudnn_dtype + c_dtype = self.c.cudnn_dtype + if ab12_dtype not in _BLOCK_SCALED_OUTPUT_DTYPES: + raise ValueError(f"AB12 has unsupported dtype {self.ab12.dtype}") + if c_dtype not in _BLOCK_SCALED_OUTPUT_DTYPES: + raise ValueError(f"C has unsupported dtype {self.c.dtype}") + if self.acc_dtype != data_type.FLOAT: + raise ValueError(f"Accumulator dtype must be float32, got {self.acc_dtype}") + if ab_dtype == data_type.FP4_E2M1 and c_dtype in _FP8_DTYPES: + raise ValueError("FP4 A and B are not compatible with FP8 C") + if c_dtype == data_type.FLOAT and ab12_dtype == data_type.FLOAT: + raise NotImplementedError("float32 C and float32 AB12 are disabled because the kernel fails to launch") + if ab_dtype in _FP8_DTYPES and (c_dtype in _FP8_DTYPES or ab12_dtype in _FP8_DTYPES or ab12_dtype == data_type.FLOAT): + raise NotImplementedError("MXFP8 inputs require float16 or bfloat16 AB12 and non-FP8 C") + + if c_dtype in _FP8_DTYPES: + if self.sfc is None or self.norm_const is None: + raise ValueError("sfc and norm_const are required when C is FP8") + if self.sfc is not None: + require_tensor_shape(self.sfc, block_scale_shape(m, output_n, batch, self.sf_vec_size), label="SFC") + require_block_scale_layout(self.sfc, "SFC") + if self.sfc.cudnn_dtype != sf_dtype: + raise ValueError(f"SFC must have the same dtype as SFA, got {self.sfc.dtype} and {self.sfa.dtype}") + if self.norm_const is not None: + require_tensor_shape(self.norm_const, (1,), label="norm_const") + if self.norm_const.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"norm_const must have float32 dtype, got {self.norm_const.dtype}") + if ab_dtype == data_type.FP4_E2M1 and c_dtype == data_type.BFLOAT16 and self.amax is None: + raise ValueError("amax is required when A and B are FP4 and C is bfloat16") + if self.amax is not None: + require_tensor_shape(self.amax, (1,), label="amax") + if self.amax.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"amax must have float32 dtype, got {self.amax.dtype}") + + a_major = require_compact_major(self.a, "m", "k") + b_major = require_compact_major(self.b, "n", "k") + ab12_major = require_compact_major(self.ab12, "m", "n") + c_major = require_compact_major(self.c, "m", "n") + if ab12_major != c_major: + raise ValueError(f"AB12 and C must use the same major mode, got {ab12_major}-major and {c_major}-major") + if ab_dtype == data_type.FP4_E2M1 and (a_major != "k" or b_major != "k"): + raise ValueError("FP4 A and B must use K-major layouts") + if ab_dtype == data_type.FP4_E2M1 and ab12_major != "n": + raise ValueError("FP4 AB12 and C must use N-major layouts") + + mma_tiler_mn = require_mma_tiler( + self.requested_mma_tiler_mn, + allowed_n=_SWIGLU_MMA_N, + ) + cta_group_size = 2 if mma_tiler_mn[0] == 256 else 1 + default_cluster = (2, 2) if cta_group_size == 2 else (1, 1) + cluster_shape_mn = require_cluster_shape( + default_cluster if self.requested_cluster_shape_mn is None else self.requested_cluster_shape_mn, + cta_group_size=cta_group_size, + ) + if cluster_shape_mn[0] > 4 or cluster_shape_mn[1] > 4: + raise ValueError(f"cluster_shape_mn entries must be at most 4 for scale-factor multicast, got {cluster_shape_mn}") + if m % mma_tiler_mn[0] or n % mma_tiler_mn[1]: + raise ValueError(f"M and N must be divisible by mma_tiler_mn, got M={m}, N={n}, tile={mma_tiler_mn}") + if self.ab12_stages <= 0: + raise ValueError(f"ab12_stages must be positive, got {self.ab12_stages}") + + for tensor in (self.a, self.b, self.ab12, self.c): + require_16_byte_alignment(tensor) + + self.m, self.n, self.k, self.l = m, n, k, batch + self.output_n = output_n + self.a_major, self.b_major, self.output_major = a_major, b_major, ab12_major + self.ab_dtype, self.sf_dtype = ab_dtype, sf_dtype + self.ab12_dtype, self.c_dtype = ab12_dtype, c_dtype + self.mma_tiler_mn = mma_tiler_mn + self.cluster_shape_mn = cluster_shape_mn + return True + + +__all__ = [ + "BlockScaledGemmSwigluSm100Op", + "GemmSwigluSm100Op", + "block_scale_shape", +] diff --git a/python/cudnn/jax/__init__.py b/python/cudnn/jax/__init__.py index 4b61bf0b0..ae7f5aef6 100644 --- a/python/cudnn/jax/__init__.py +++ b/python/cudnn/jax/__init__.py @@ -39,6 +39,12 @@ "rmsnorm_rht_amax_sm100": ("..rmsnorm_rht_amax.jax", "rmsnorm_rht_amax_sm100"), "GemmSwigluSm100": ("..gemm_swiglu.jax", "GemmSwigluSm100"), "gemm_swiglu_wrapper_sm100": ("..gemm_swiglu.jax", "gemm_swiglu_wrapper_sm100"), + "GemmSreluSm100": ("..gemm_srelu.jax", "GemmSreluSm100"), + "gemm_srelu_wrapper_sm100": ("..gemm_srelu.jax", "gemm_srelu_wrapper_sm100"), + "GemmDsreluSm100": ("..gemm_dsrelu.jax", "GemmDsreluSm100"), + "gemm_dsrelu_wrapper_sm100": ("..gemm_dsrelu.jax", "gemm_dsrelu_wrapper_sm100"), + "GemmAmaxSm100": ("..gemm_amax.jax", "GemmAmaxSm100"), + "gemm_amax_wrapper_sm100": ("..gemm_amax.jax", "gemm_amax_wrapper_sm100"), "SparseAttentionBackward": ("..deepseek_sparse_attention.sparse_attention_backward.jax", "SparseAttentionBackward"), "sparse_attention_backward_wrapper": ( "..deepseek_sparse_attention.sparse_attention_backward.jax", diff --git a/python/cudnn/rmsnorm_rht_amax/jax.py b/python/cudnn/rmsnorm_rht_amax/jax.py index a64fe1da0..046b973ef 100644 --- a/python/cudnn/rmsnorm_rht_amax/jax.py +++ b/python/cudnn/rmsnorm_rht_amax/jax.py @@ -93,30 +93,30 @@ def __call__(self, x: Any, weight: Any) -> tuple[Any, Any]: self.o_desc, divisibility=(self._op.rows_per_cta, 16), ) + + def launch( + stream: Any, + x: Any, + weight: Any, + output: Any, + amax: Any, + ) -> None: + kernel = RMSNormRHTAmaxKernel( + n=self._op.n, + num_threads=self._op.num_threads, + eps=self._op.eps, + rows_per_cta=self._op.rows_per_cta, + ) + kernel(x, weight, output, amax, stream) + return self._call_kernel( (x, weight), + launch=launch, output_descs=(self.o_desc, self.amax_desc), input_spec=(x_spec, weight_spec), output_spec=(output_spec, None), ) - def _launch( - self, - inputs: tuple[Any, ...], - outputs: tuple[Any, ...], - workspaces: tuple[Any, ...], - stream: Any, - ) -> None: - x, weight = inputs - output, amax = outputs - kernel = RMSNormRHTAmaxKernel( - n=self._op.n, - num_threads=self._op.num_threads, - eps=self._op.eps, - rows_per_cta=self._op.rows_per_cta, - ) - kernel(x, weight, output, amax, stream) - @jax.jit(static_argnames=("eps", "num_threads", "rows_per_cta")) def rmsnorm_rht_amax_sm100( diff --git a/test/python/fe_api/test_gemm_amax_imports.py b/test/python/fe_api/test_gemm_amax_imports.py new file mode 100644 index 000000000..7b19102cc --- /dev/null +++ b/test/python/fe_api/test_gemm_amax_imports.py @@ -0,0 +1,148 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Import and Torch-adapter contracts for dense GEMM + amax.""" + +import ast +import importlib.util +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "gemm_amax" + + +def _load_operation_package(name: str): + parent_name, separator, _ = name.rpartition(".") + if not separator: + raise ValueError("Synthetic operation package must have a parent") + parent = types.ModuleType(parent_name) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = parent_name + sys.modules[parent_name] = parent + + spec = importlib.util.spec_from_file_location( + name, + _OPERATION_ROOT / "__init__.py", + submodule_search_locations=[str(_OPERATION_ROOT)], + ) + if spec is None or spec.loader is None: + raise RuntimeError("Unable to load GEMM + amax package") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _remove_package(name: str) -> None: + root_name = name.split(".", 1)[0] + for module_name in tuple(sys.modules): + if module_name == root_name or module_name.startswith(f"{root_name}."): + sys.modules.pop(module_name, None) + + +class GemmAmaxImportContractTest(unittest.TestCase): + def test_package_import_loads_no_framework_adapter(self): + name = "cudnn_gemm_amax_import_test.gemm_amax" + blocked = {module_name: None for module_name in ("torch", "jax", "cutlass", "cuda")} + try: + with mock.patch.dict(sys.modules, blocked): + package = _load_operation_package(name) + self.assertNotIn(f"{name}.api", sys.modules) + self.assertNotIn(f"{name}.op", sys.modules) + self.assertTrue( + { + "GemmAmaxSm100Op", + "GemmAmaxSm100", + "gemm_amax_wrapper_sm100", + "api", + "op", + }.issubset(dir(package)) + ) + finally: + _remove_package(name) + + def test_static_adapter_boundaries_and_public_torch_signatures(self): + def imports_framework(node, framework: str): + if isinstance(node, ast.Import): + return any(alias.name == framework or alias.name.startswith(f"{framework}.") for alias in node.names) + return node.level == 0 and (node.module == framework or (node.module or "").startswith(f"{framework}.")) + + op_tree = ast.parse((_OPERATION_ROOT / "op.py").read_text(), filename="op.py") + api_tree = ast.parse((_OPERATION_ROOT / "api.py").read_text(), filename="api.py") + op_imports = [node for node in ast.walk(op_tree) if isinstance(node, (ast.Import, ast.ImportFrom))] + api_imports = [node for node in ast.walk(api_tree) if isinstance(node, (ast.Import, ast.ImportFrom))] + self.assertFalse(any(imports_framework(node, framework) for node in op_imports for framework in ("torch", "jax", "cutlass", "cuda"))) + self.assertFalse(any(imports_framework(node, "jax") for node in api_imports)) + + adapter = next(node for node in api_tree.body if isinstance(node, ast.ClassDef) and node.name == "GemmAmaxSm100") + constructor = next(node for node in adapter.body if isinstance(node, ast.FunctionDef) and node.name == "__init__") + self.assertEqual( + [argument.arg for argument in constructor.args.args], + [ + "self", + "sample_a", + "sample_b", + "sample_sfa", + "sample_sfb", + "sample_c", + "sample_amax", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + ], + ) + op_constructors = [ + node for node in ast.walk(constructor) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "GemmAmaxSm100Op" + ] + self.assertEqual(len(op_constructors), 1) + + execute = next(node for node in adapter.body if isinstance(node, ast.FunctionDef) and node.name == "execute") + self.assertEqual( + [argument.arg for argument in execute.args.args], + [ + "self", + "a_tensor", + "b_tensor", + "sfa_tensor", + "sfb_tensor", + "c_tensor", + "amax_tensor", + "current_stream", + ], + ) + + wrapper = next(node for node in api_tree.body if isinstance(node, ast.FunctionDef) and node.name == "gemm_amax_wrapper_sm100") + self.assertEqual( + [argument.arg for argument in wrapper.args.args], + [ + "a_tensor", + "b_tensor", + "sfa_tensor", + "sfb_tensor", + "c_major", + "c_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "stream", + ], + ) + self.assertFalse(any(isinstance(node, ast.FunctionDef) and node.name.startswith("_check_torch_support") for node in adapter.body)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_gemm_amax_op.py b/test/python/fe_api/test_gemm_amax_op.py new file mode 100644 index 000000000..b61e6ab57 --- /dev/null +++ b/test/python/fe_api/test_gemm_amax_op.py @@ -0,0 +1,295 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for the dense GEMM + amax operation.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "gemm_amax" +_PACKAGE = "cudnn_gemm_amax_op_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT8 = auto() + UINT8 = auto() + FP4_E2M1 = auto() + FP8_E4M3 = auto() + FP8_E5M2 = auto() + FP8_E8M0 = auto() + + +class GemmAmaxOpContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + operation_name = f"{_PACKAGE}.gemm_amax" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + try: + cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}._op") + cls.op_module = importlib.import_module(f"{operation_name}.op") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def _desc(self, shape, dtype, order, name): + shape = tuple(shape) + stride = [0] * len(shape) + running = 1 + for dimension in order: + stride[dimension] = running + running *= max(shape[dimension], 1) + return self.tensor_module.TensorDesc( + dtype=dtype, + shape=shape, + stride=tuple(stride), + stride_order=tuple(order), + name=name, + ) + + def _op(self, **overrides): + m, n, k, batch = 128, 256, 128, 2 + scale_order = (3, 1, 0, 4, 2, 5) + arguments = { + "a": self._desc((m, k, batch), _DataType.FP8_E4M3, (1, 0, 2), "A"), + "b": self._desc((n, k, batch), _DataType.FP8_E4M3, (1, 0, 2), "B"), + "sfa": self._desc((32, 4, 1, 4, 1, batch), _DataType.FP8_E8M0, scale_order, "SFA"), + "sfb": self._desc((32, 4, 2, 4, 1, batch), _DataType.FP8_E8M0, scale_order, "SFB"), + "c": self._desc((m, n, batch), _DataType.FLOAT, (1, 0, 2), "C"), + "amax": self._desc((1, 1, 1), _DataType.FLOAT, (2, 1, 0), "Amax"), + } + arguments.update(overrides) + return self.op_module.GemmAmaxSm100Op(**arguments) + + def test_validates_complete_signature_and_resolves_configuration(self): + operation = self._op() + + self.assertIsInstance(operation, self.base_module.Op) + self.assertTrue(operation.check_support()) + self.assertEqual((operation.m, operation.n, operation.k, operation.l), (128, 256, 128, 2)) + self.assertEqual((operation.a_major, operation.b_major, operation.c_major), ("k", "k", "n")) + self.assertEqual(operation.ab_dtype, _DataType.FP8_E4M3) + self.assertEqual(operation.scale_dtype, _DataType.FP8_E8M0) + self.assertEqual(operation.c_dtype, _DataType.FLOAT) + self.assertEqual(operation.mma_tiler_mn, (128, 128)) + self.assertEqual(operation.cluster_shape_mn, (1, 1)) + self.assertNotIn(f"{_PACKAGE}.gemm_amax.api", sys.modules) + self.assertNotIn( + f"{_PACKAGE}.gemm_amax.dense_blockscaled_gemm_persistent_amax", + sys.modules, + ) + + def test_accepts_supported_major_modes_and_storage_aliases(self): + scale_order = (3, 1, 0, 4, 2, 5) + operation = self._op( + a=self._desc((128, 128, 2), _DataType.UINT8, (1, 0, 2), "A"), + b=self._desc((256, 128, 2), _DataType.UINT8, (1, 0, 2), "B"), + sfa=self._desc((32, 4, 1, 4, 2, 2), _DataType.INT8, scale_order, "SFA"), + sfb=self._desc((32, 4, 2, 4, 2, 2), _DataType.INT8, scale_order, "SFB"), + c=self._desc((128, 256, 2), _DataType.UINT8, (1, 0, 2), "C"), + sf_vec_size=16, + ) + self.assertTrue(operation.check_support()) + self.assertEqual(operation.ab_dtype, _DataType.FP4_E2M1) + self.assertEqual(operation.scale_dtype, _DataType.FP8_E8M0) + self.assertEqual(operation.c_dtype, _DataType.FP4_E2M1) + + alternate = self._op( + a=self._desc((128, 128, 2), _DataType.FP8_E4M3, (0, 1, 2), "A"), + b=self._desc((256, 128, 2), _DataType.FP8_E4M3, (0, 1, 2), "B"), + c=self._desc((128, 256, 2), _DataType.BFLOAT16, (0, 1, 2), "C"), + ) + self.assertTrue(alternate.check_support()) + self.assertEqual((alternate.a_major, alternate.b_major, alternate.c_major), ("m", "n", "m")) + + def test_scale_layout_accepts_size_one_stride_order_ties(self): + scale = self.tensor_module.TensorDesc( + dtype=_DataType.FP8_E8M0, + shape=(32, 4, 1, 4, 1, 1), + stride=(16, 4, 512, 1, 512, 512), + stride_order=(3, 1, 0, 2, 4, 5), + name="scale", + ) + operation = self._op( + a=self._desc((128, 128, 1), _DataType.FP8_E4M3, (1, 0, 2), "A"), + b=self._desc((128, 128, 1), _DataType.FP8_E4M3, (1, 0, 2), "B"), + sfa=scale, + sfb=scale, + c=self._desc((128, 128, 1), _DataType.FLOAT, (1, 0, 2), "C"), + ) + self.assertTrue(operation.check_support()) + + def test_existing_torch_configuration_families_remain_accepted(self): + scale_order = (3, 1, 0, 4, 2, 5) + cases = ( + { + "a": self._desc((128, 256, 2), _DataType.FP8_E5M2, (0, 1, 2), "A"), + "b": self._desc((256, 256, 2), _DataType.FP8_E5M2, (0, 1, 2), "B"), + "sfa": self._desc((32, 4, 1, 4, 2, 2), _DataType.FP8_E8M0, scale_order, "SFA"), + "sfb": self._desc((32, 4, 2, 4, 2, 2), _DataType.FP8_E8M0, scale_order, "SFB"), + "c": self._desc((128, 256, 2), _DataType.BFLOAT16, (0, 1, 2), "C"), + "mma_tiler_mn": (128, 256), + "cluster_shape_mn": (2, 2), + }, + { + "a": self._desc((128, 128, 2), _DataType.FP4_E2M1, (1, 0, 2), "A"), + "b": self._desc((256, 128, 2), _DataType.FP4_E2M1, (1, 0, 2), "B"), + "sfa": self._desc((32, 4, 1, 4, 2, 2), _DataType.FP8_E4M3, scale_order, "SFA"), + "sfb": self._desc((32, 4, 2, 4, 2, 2), _DataType.FP8_E4M3, scale_order, "SFB"), + "c": self._desc((128, 256, 2), _DataType.FP8_E4M3, (1, 0, 2), "C"), + "sf_vec_size": 16, + "cluster_shape_mn": (2, 2), + }, + { + "a": self._desc((128, 128, 2), _DataType.FP4_E2M1, (1, 0, 2), "A"), + "b": self._desc((256, 128, 2), _DataType.FP4_E2M1, (1, 0, 2), "B"), + "c": self._desc((128, 256, 2), _DataType.FLOAT, (0, 1, 2), "C"), + "cluster_shape_mn": (4, 4), + }, + ) + for overrides in cases: + with self.subTest(overrides=tuple(sorted(overrides))): + self.assertTrue(self._op(**overrides).check_support()) + + def test_rejects_inconsistent_shapes_and_scale_layout(self): + scale_order = (3, 1, 0, 4, 2, 5) + cases = ( + ( + {"b": self._desc((256, 64, 2), _DataType.FP8_E4M3, (1, 0, 2), "B")}, + "B shape mismatch", + ), + ( + {"c": self._desc((128, 128, 2), _DataType.FLOAT, (1, 0, 2), "C")}, + "C must have shape", + ), + ( + { + "sfa": self._desc( + (32, 4, 1, 4, 2, 2), + _DataType.FP8_E8M0, + scale_order, + "SFA", + ) + }, + "SFA must have shape", + ), + ( + { + "sfa": self._desc( + (32, 4, 1, 4, 1, 2), + _DataType.FP8_E8M0, + (5, 4, 3, 2, 1, 0), + "SFA", + ) + }, + "packed block-scale layout", + ), + ( + {"amax": self._desc((1,), _DataType.FLOAT, (0,), "Amax")}, + "Amax must have shape", + ), + ) + for overrides, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self._op(**overrides).check_support() + + def test_rejects_invalid_dtype_combinations(self): + cases = ( + ( + {"b": self._desc((256, 128, 2), _DataType.FP8_E5M2, (1, 0, 2), "B")}, + ValueError, + "A and B must have the same dtype", + ), + ( + {"amax": self._desc((1, 1, 1), _DataType.HALF, (2, 1, 0), "Amax")}, + ValueError, + "Amax must have dtype float32", + ), + ( + {"acc_dtype": _DataType.HALF}, + ValueError, + "Accumulator dtype must be float32", + ), + ( + {"c": self._desc((128, 256, 2), _DataType.FP8_E5M2, (1, 0, 2), "C")}, + NotImplementedError, + "FP8 A and B with FP8 C", + ), + ( + {"sf_vec_size": 16}, + ValueError, + "FP8 A and B do not support sf_vec_size=16", + ), + ) + for overrides, error_type, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(error_type, message): + self._op(**overrides).check_support() + + def test_validates_tiler_cluster_and_alignment(self): + with self.assertRaisesRegex(ValueError, r"mma_tiler_mn\[1\]"): + self._op(mma_tiler_mn=(128, 64)).check_support() + with self.assertRaisesRegex(NotImplementedError, "currently hangs"): + self._op(mma_tiler_mn=(256, 128)).check_support() + with self.assertRaisesRegex(ValueError, "entries must be at most 4"): + self._op(cluster_shape_mn=(8, 1)).check_support() + with self.assertRaisesRegex(TypeError, "mma_tiler_mn must contain two integers"): + self._op(mma_tiler_mn=None).check_support() + + scale_order = (3, 1, 0, 4, 2, 5) + with self.assertRaisesRegex(ValueError, "16-byte alignment"): + self._op( + a=self._desc((128, 120, 2), _DataType.FP8_E4M3, (1, 0, 2), "A"), + b=self._desc((256, 120, 2), _DataType.FP8_E4M3, (1, 0, 2), "B"), + sfa=self._desc((32, 4, 1, 4, 1, 2), _DataType.FP8_E8M0, scale_order, "SFA"), + sfb=self._desc((32, 4, 2, 4, 1, 2), _DataType.FP8_E8M0, scale_order, "SFB"), + ).check_support() + + def test_constructor_requires_descriptors_and_canonical_accumulator(self): + for name in ("a", "b", "sfa", "sfb", "c", "amax"): + with self.subTest(name=name): + with self.assertRaisesRegex(TypeError, f"{name} must be a TensorDesc"): + self._op(**{name: types.SimpleNamespace(shape=(1,))}) + with self.assertRaisesRegex(TypeError, "acc_dtype must be a cudnn.data_type"): + self._op(acc_dtype="float32") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_gemm_relu_imports.py b/test/python/fe_api/test_gemm_relu_imports.py new file mode 100644 index 000000000..d336471d6 --- /dev/null +++ b/test/python/fe_api/test_gemm_relu_imports.py @@ -0,0 +1,121 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Import boundaries for dense GEMM + sReLU and dsReLU.""" + +import ast +import importlib.util +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" + + +def _load_operation_package(name: str, operation: str): + parent_name = name.rpartition(".")[0] + parent = types.ModuleType(parent_name) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = parent_name + sys.modules[parent_name] = parent + + operation_root = _CUDNN_ROOT / operation + spec = importlib.util.spec_from_file_location( + name, + operation_root / "__init__.py", + submodule_search_locations=[str(operation_root)], + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"Unable to load {operation} package") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _remove_package(name: str) -> None: + root_name = name.split(".", 1)[0] + for module_name in tuple(sys.modules): + if module_name == root_name or module_name.startswith(f"{root_name}."): + sys.modules.pop(module_name, None) + + +class GemmReluImportContractTest(unittest.TestCase): + def test_public_jax_facade_routes_both_operations(self): + tree = ast.parse((_CUDNN_ROOT / "jax" / "__init__.py").read_text(), filename="cudnn/jax/__init__.py") + assignment = next( + node + for node in tree.body + if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == "_OPERATION_EXPORTS" for target in node.targets) + ) + exports = ast.literal_eval(assignment.value) + self.assertEqual(exports["GemmSreluSm100"], ("..gemm_srelu.jax", "GemmSreluSm100")) + self.assertEqual(exports["gemm_srelu_wrapper_sm100"], ("..gemm_srelu.jax", "gemm_srelu_wrapper_sm100")) + self.assertEqual(exports["GemmDsreluSm100"], ("..gemm_dsrelu.jax", "GemmDsreluSm100")) + self.assertEqual(exports["gemm_dsrelu_wrapper_sm100"], ("..gemm_dsrelu.jax", "gemm_dsrelu_wrapper_sm100")) + + def test_package_imports_are_framework_free_and_lazy(self): + blocked = {module_name: None for module_name in ("torch", "jax", "cutlass", "cuda")} + for operation, op_class, api_class, wrapper in ( + ("gemm_srelu", "GemmSreluSm100Op", "GemmSreluSm100", "gemm_srelu_wrapper_sm100"), + ("gemm_dsrelu", "GemmDsreluSm100Op", "GemmDsreluSm100", "gemm_dsrelu_wrapper_sm100"), + ): + name = f"cudnn_{operation}_import_test.{operation}" + try: + with self.subTest(operation=operation), mock.patch.dict(sys.modules, blocked): + package = _load_operation_package(name, operation) + self.assertNotIn(f"{name}.api", sys.modules) + self.assertNotIn(f"{name}.op", sys.modules) + self.assertTrue({op_class, api_class, wrapper, "api", "op"}.issubset(dir(package))) + finally: + _remove_package(name) + + def test_exports_route_to_op_and_torch_modules_independently(self): + for operation, op_class, api_class, wrapper in ( + ("gemm_srelu", "GemmSreluSm100Op", "GemmSreluSm100", "gemm_srelu_wrapper_sm100"), + ("gemm_dsrelu", "GemmDsreluSm100Op", "GemmDsreluSm100", "gemm_dsrelu_wrapper_sm100"), + ): + name = f"cudnn_{operation}_route_test.{operation}" + op_sentinel = object() + api_sentinel = object() + op_module = types.ModuleType(f"{name}.op") + setattr(op_module, op_class, op_sentinel) + api_module = types.ModuleType(f"{name}.api") + setattr(api_module, api_class, api_sentinel) + setattr(api_module, wrapper, object()) + try: + with self.subTest(operation=operation): + sys.modules[op_module.__name__] = op_module + sys.modules[api_module.__name__] = api_module + package = _load_operation_package(name, operation) + self.assertIs(getattr(package, op_class), op_sentinel) + self.assertIs(getattr(package, api_class), api_sentinel) + finally: + _remove_package(name) + + def test_framework_neutral_op_modules_have_no_framework_imports(self): + paths = [ + _CUDNN_ROOT / "_gemm_relu.py", + _CUDNN_ROOT / "gemm_srelu" / "op.py", + _CUDNN_ROOT / "gemm_dsrelu" / "op.py", + ] + for path in paths: + imports = [node for node in ast.walk(ast.parse(path.read_text(), filename=str(path))) if isinstance(node, (ast.Import, ast.ImportFrom))] + imported = {alias.name.split(".", 1)[0] for node in imports if isinstance(node, ast.Import) for alias in node.names} + imported.update((node.module or "").split(".", 1)[0] for node in imports if isinstance(node, ast.ImportFrom) and node.level == 0) + with self.subTest(path=path.name): + self.assertTrue(imported.isdisjoint({"torch", "jax", "cutlass", "cuda"})) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_gemm_swiglu_imports.py b/test/python/fe_api/test_gemm_swiglu_imports.py index b40c9c162..db44c06aa 100644 --- a/test/python/fe_api/test_gemm_swiglu_imports.py +++ b/test/python/fe_api/test_gemm_swiglu_imports.py @@ -63,6 +63,7 @@ def test_package_import_loads_no_framework_adapter(self): self.assertNotIn(f"{name}.op", sys.modules) self.assertTrue( { + "BlockScaledGemmSwigluSm100Op", "GemmSwigluSm100Op", "GemmSwigluSm100", "gemm_swiglu_wrapper_sm100", @@ -81,6 +82,7 @@ def test_exports_route_without_loading_the_other_layer(self): api.GemmSwigluSm100 = api_sentinel api.gemm_swiglu_wrapper_sm100 = object() op = types.ModuleType(f"{name}.op") + op.BlockScaledGemmSwigluSm100Op = object() op.GemmSwigluSm100Op = op_sentinel try: sys.modules[api.__name__] = api diff --git a/test/python/fe_api/test_gemm_swiglu_op.py b/test/python/fe_api/test_gemm_swiglu_op.py index a549251aa..fba11954f 100644 --- a/test/python/fe_api/test_gemm_swiglu_op.py +++ b/test/python/fe_api/test_gemm_swiglu_op.py @@ -92,6 +92,21 @@ def _op(self, **overrides): arguments.update(overrides) return self.op_module.GemmSwigluSm100Op(**arguments) + def _block_op(self, **overrides): + m, n, k, batch = 128, 128, 128, 2 + scale_order = (3, 1, 0, 4, 2, 5) + arguments = { + "a": self._desc((m, k, batch), _DataType.FP8_E4M3, (1, 0, 2), "A"), + "b": self._desc((n, k, batch), _DataType.FP8_E4M3, (1, 0, 2), "B"), + "sfa": self._desc((32, 4, 1, 4, 1, batch), _DataType.FP8_E8M0, scale_order, "SFA"), + "sfb": self._desc((32, 4, 1, 4, 1, batch), _DataType.FP8_E8M0, scale_order, "SFB"), + "ab12": self._desc((m, n, batch), _DataType.BFLOAT16, (1, 0, 2), "AB12"), + "c": self._desc((m, n // 2, batch), _DataType.BFLOAT16, (1, 0, 2), "C"), + "sf_vec_size": 32, + } + arguments.update(overrides) + return self.op_module.BlockScaledGemmSwigluSm100Op(**arguments) + def test_validates_complete_signature_and_resolves_configuration(self): operation = self._op(alpha=0.5) @@ -236,6 +251,82 @@ def test_constructor_requires_canonical_descriptors_and_dtype(self): with self.assertRaisesRegex(TypeError, "acc_dtype must be a cudnn.data_type"): self._op(acc_dtype="float32") + def test_validates_block_scaled_mxfp8_signature(self): + operation = self._block_op(alpha=0.25, vector_f32=True, ab12_stages=3) + + self.assertTrue(operation.check_support()) + self.assertEqual((operation.m, operation.n, operation.k, operation.l), (128, 128, 128, 2)) + self.assertEqual(operation.output_n, 64) + self.assertEqual((operation.ab_dtype, operation.sf_dtype), (_DataType.FP8_E4M3, _DataType.FP8_E8M0)) + self.assertEqual((operation.ab12_dtype, operation.c_dtype), (_DataType.BFLOAT16, _DataType.BFLOAT16)) + self.assertEqual(operation.mma_tiler_mn, (128, 128)) + self.assertEqual(operation.cluster_shape_mn, (1, 1)) + self.assertTrue(operation.vector_f32) + self.assertEqual(operation.ab12_stages, 3) + + def test_validates_block_scaled_fp4_amax_signature(self): + m, n, k, batch = 256, 256, 128, 1 + scale_order = (3, 1, 0, 4, 2, 5) + operation = self._block_op( + a=self._desc((m, k, batch), _DataType.FP4_E2M1, (1, 0, 2), "A"), + b=self._desc((n, k, batch), _DataType.FP4_E2M1, (1, 0, 2), "B"), + sfa=self._desc((32, 4, 2, 4, 2, batch), _DataType.FP8_E4M3, scale_order, "SFA"), + sfb=self._desc((32, 4, 2, 4, 2, batch), _DataType.FP8_E4M3, scale_order, "SFB"), + ab12=self._desc((m, n, batch), _DataType.HALF, (1, 0, 2), "AB12"), + c=self._desc((m, n // 2, batch), _DataType.BFLOAT16, (1, 0, 2), "C"), + amax=self._desc((1,), _DataType.FLOAT, (0,), "amax"), + sf_vec_size=16, + mma_tiler_mn=(256, 256), + ) + + self.assertTrue(operation.check_support()) + self.assertEqual(operation.cluster_shape_mn, (2, 2)) + self.assertEqual(operation.output_major, "n") + + def test_rejects_invalid_block_scaled_signatures(self): + scale_order = (3, 1, 0, 4, 2, 5) + cases = ( + ( + {"sfa": self._desc((32, 4, 1, 4, 2, 2), _DataType.FP8_E8M0, scale_order, "SFA")}, + ValueError, + "SFA must have shape", + ), + ( + {"sfa": self._desc((32, 4, 1, 4, 1, 2), _DataType.FP8_E8M0, (5, 4, 3, 2, 1, 0), "SFA")}, + ValueError, + "packed block-scale layout", + ), + ( + { + "sfa": self._desc((32, 4, 1, 4, 1, 2), _DataType.FP8_E4M3, scale_order, "SFA"), + "sfb": self._desc((32, 4, 1, 4, 1, 2), _DataType.FP8_E4M3, scale_order, "SFB"), + }, + ValueError, + "FP8 A and B require FP8_E8M0 scales with sf_vec_size=32", + ), + ( + {"ab12": self._desc((128, 128, 2), _DataType.FLOAT, (1, 0, 2), "AB12")}, + NotImplementedError, + "MXFP8 inputs require", + ), + ( + {"mma_tiler_mn": (128, 192)}, + ValueError, + "M and N must be divisible by mma_tiler_mn", + ), + ) + for overrides, error_type, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(error_type, message): + self._block_op(**overrides).check_support() + + fp4 = _DataType.FP4_E2M1 + with self.assertRaisesRegex(ValueError, "amax is required"): + self._block_op( + a=self._desc((128, 128, 2), fp4, (1, 0, 2), "A"), + b=self._desc((128, 128, 2), fp4, (1, 0, 2), "B"), + ).check_support() + if __name__ == "__main__": unittest.main() diff --git a/test/python/fe_api/test_jax_api_base.py b/test/python/fe_api/test_jax_api_base.py index 6b2de6cb0..e3aaef2eb 100644 --- a/test/python/fe_api/test_jax_api_base.py +++ b/test/python/fe_api/test_jax_api_base.py @@ -240,9 +240,6 @@ def check_support(self): def __call__(self): return self.check_support() - def _launch(self, inputs, outputs, workspaces, stream): - raise AssertionError("not called") - with mock.patch.dict( sys.modules, { @@ -452,7 +449,7 @@ def test_materializes_canonical_output_in_public_axis_order(self): self.assertEqual(output.shape, (2, 3, 4)) self.assertEqual(output.dtype, "bfloat16") - def test_calls_launch_hook_with_initialized_outputs_and_hidden_workspaces(self): + def test_calls_explicit_launcher_with_initialized_outputs_and_hidden_workspaces(self): tensor_module = self.tensor_module seen = {} @@ -483,6 +480,9 @@ def test_calls_launch_hook_with_initialized_outputs_and_hidden_workspaces(self): ), ) + def launch(stream, *buffers): + seen["launch_args"] = (stream, *buffers) + class Adapter(self.module.JaxApiBase): def check_support(self): return True @@ -490,13 +490,11 @@ def check_support(self): def __call__(self, value): return self._call_kernel( (value,), + launch=launch, output_descs=output_descs, workspace_descs=workspace_descs, ) - def _launch(self, inputs, outputs, workspaces, stream): - seen["launch_args"] = (inputs, outputs, workspaces, stream) - full_calls = [] jax = types.ModuleType("jax") jax.__path__ = [] @@ -557,6 +555,7 @@ def invoke(*args): ): results = api._call_kernel( (input_value,), + launch=launch, output_descs=output_descs, workspace_descs=workspace_descs, input_spec=(input_tensor_spec,), @@ -565,16 +564,24 @@ def invoke(*args): with self.assertRaisesRegex(TypeError, "input #0 must have shape and dtype metadata"): api._call_kernel( ([input_value],), + launch=launch, output_descs=output_descs, workspace_descs=workspace_descs, ) with self.assertRaisesRegex(TypeError, "output_spec must contain only TensorSpec or None"): api._call_kernel( (input_value,), + launch=launch, output_descs=output_descs, workspace_descs=workspace_descs, output_spec=((0,), initialized_output_spec), ) + with self.assertRaisesRegex(TypeError, "launch must be callable"): + api._call_kernel( + (input_value,), + launch=None, + output_descs=output_descs, + ) self.assertEqual( [(shape, value, dtype) for shape, value, dtype, _ in full_calls], @@ -588,13 +595,13 @@ def invoke(*args): self.assertIs(seen["call_options"]["input_spec"][1], initialized_output_spec) self.assertIs(seen["call_options"]["input_spec"][2], seen["call_options"]["output_spec"][2]) - launch_inputs, launch_outputs, launch_workspaces, stream = seen["launch_args"] - self.assertEqual(launch_inputs, (input_value,)) - self.assertEqual(launch_outputs[0].label, "allocated(0)") - self.assertIs(launch_outputs[1], full_calls[0][3]) - self.assertIs(launch_workspaces[0], full_calls[1][3]) - self.assertEqual(launch_workspaces[1].label, "allocated(3)") + stream, launch_input, output, amax, counter, scratch = seen["launch_args"] self.assertEqual(stream, "stream") + self.assertIs(launch_input, input_value) + self.assertEqual(output.label, "allocated(0)") + self.assertIs(amax, full_calls[0][3]) + self.assertIs(counter, full_calls[1][3]) + self.assertEqual(scratch.label, "allocated(3)") self.assertEqual([result.label for result in results], ["allocated(0)", "full(-inf)"]) self.assertNotIn(full_calls[0][3], vars(api).values()) self.assertNotIn(full_calls[1][3], vars(api).values()) @@ -609,9 +616,6 @@ def check_support(self): def __call__(self): return self.check_support() - def _launch(self, inputs, outputs, workspaces, stream): - raise AssertionError("not called") - api = Adapter() api.option = 1 self.assertIs(api.get_jax_callable(), api) diff --git a/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py b/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py index 7d900a59e..2dec243da 100644 --- a/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py +++ b/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py @@ -189,7 +189,8 @@ def _to_tensor_spec(desc, *, mode=None, divisibility=None): del desc return _TensorSpec(mode=mode, divisibility=divisibility) - def _call_kernel(self, inputs, **options): + def _call_kernel(self, inputs, *, launch, **options): + options["launch"] = launch JaxApiBase.captured_call = (inputs, options) return tuple(_Array(desc.shape, cls.cudnn_to_dtype[desc.cudnn_dtype]) for desc in options["output_descs"]) @@ -437,7 +438,7 @@ def test_sparse_sm100_and_sm90_launch_functionally(self): _RecordedKernel.events.clear() api = self.module.IndexerBackward(*inputs, target_compute_capability=target) api.check_support() - api._launch(kernel_inputs, outputs, workspace, object()) + api._launch_kernel(object(), *kernel_inputs, *outputs, *workspace) self.assertEqual(tuple(event[0] for event in _RecordedKernel.events), labels) score_args = _RecordedKernel.events[0][2] self.assertIs(score_args[-2 if target == 100 else -1], workspace[0]) @@ -472,7 +473,7 @@ def test_sparse_packed_thd_launch_adds_a_synthetic_batch_view(self): target_compute_capability=target, ) api.check_support() - api._launch(kernel_inputs, outputs, workspace, object()) + api._launch_kernel(object(), *kernel_inputs, *outputs, *workspace) self.assertEqual(tuple(event[0] for event in _RecordedKernel.events), labels) score_args = _RecordedKernel.events[0][2] @@ -495,7 +496,7 @@ def test_dense_sm100_bshd_and_sm90_thd_launch_functionally(self): api = self.module.DenseIndexerBackward(*bshd, target_compute_capability=100) api.check_support() kernel_inputs = tuple(object() for _ in range(8)) - api._launch(kernel_inputs, outputs, workspace, object()) + api._launch_kernel(object(), *kernel_inputs, *outputs, *workspace) self.assertEqual(tuple(event[0] for event in _RecordedKernel.events), ("dense_score_100", "dense_gemm_100")) score_args = _RecordedKernel.events[0][2] self.assertIs(score_args[11], workspace[0]) @@ -518,7 +519,7 @@ def test_dense_sm100_bshd_and_sm90_thd_launch_functionally(self): ) api.check_support() kernel_inputs = tuple(object() for _ in range(11)) - api._launch(kernel_inputs, outputs, workspace, object()) + api._launch_kernel(object(), *kernel_inputs, *outputs, *workspace) self.assertEqual(tuple(event[0] for event in _RecordedKernel.events), ("dense_score_90", "sparse_gemm_90")) score_args = _RecordedKernel.events[0][2] self.assertEqual(score_args[4:7], kernel_inputs[8:11]) diff --git a/test/python/fe_api/test_jax_dsa_score_recompute_contract.py b/test/python/fe_api/test_jax_dsa_score_recompute_contract.py index 9b0c652f1..cee84f8e1 100644 --- a/test/python/fe_api/test_jax_dsa_score_recompute_contract.py +++ b/test/python/fe_api/test_jax_dsa_score_recompute_contract.py @@ -107,11 +107,12 @@ def _check_tensor_signature(value, expected, *, mode=None): if value.dtype != expected.cudnn_dtype: raise ValueError(f"{expected.name} tensor dtype mismatch") - def _call_kernel(self, inputs, *, output_descs, workspace_descs=(), **options): + def _call_kernel(self, inputs, *, launch, output_descs, workspace_descs=(), **options): self.captured_call = { "inputs": tuple(inputs), "outputs": tuple(output_descs), "workspaces": tuple(workspace_descs), + "launch": launch, **options, } return tuple(_Array(desc.shape, desc.cudnn_dtype) for desc in output_descs) @@ -356,8 +357,8 @@ def __call__(self, *args): sm100_module.__name__: sm100_module, }, ): - sm90._launch(("q90", "k90", "aux90", "indices90"), ("out90",), (), "stream90") - sm100._launch(("q100", "k100", "aux100", "indices100"), ("out100",), ("dummy",), "stream100") + sm90._launch_kernel("stream90", "q90", "k90", "aux90", "indices90", "out90") + sm100._launch_kernel("stream100", "q100", "k100", "aux100", "indices100", "out100", "dummy") self.assertEqual( [entry for entry in calls if entry[0].endswith("call")], @@ -423,8 +424,8 @@ def __call__(self, *args): sm100_module.__name__: sm100_module, }, ): - sm90._launch(("q90", "k90", "aux90", "offset90"), ("out90", "denom90"), (), "stream90") - sm100._launch(("q100", "k100", "aux100", "offset100"), ("out100", "denom100"), (), "stream100") + sm90._launch_kernel("stream90", "q90", "k90", "aux90", "offset90", "out90", "denom90") + sm100._launch_kernel("stream100", "q100", "k100", "aux100", "offset100", "out100", "denom100") self.assertEqual( [entry for entry in calls if entry[0].endswith("call")], diff --git a/test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py b/test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py index 8a7634616..26ee1d383 100644 --- a/test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py +++ b/test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py @@ -156,7 +156,8 @@ def _to_tensor_spec(desc, *, mode=None, divisibility=None): del desc return _TensorSpec(mode=mode, divisibility=divisibility) - def _call_kernel(self, inputs, **options): + def _call_kernel(self, inputs, *, launch, **options): + options["launch"] = launch JaxApiBase.captured_call = (inputs, options) return tuple(_Array(desc.shape, _CUDNN_TO_DTYPE[desc.cudnn_dtype]) for desc in options["output_descs"]) diff --git a/test/python/fe_api/test_jax_gemm_amax.py b/test/python/fe_api/test_jax_gemm_amax.py new file mode 100644 index 000000000..09a563b51 --- /dev/null +++ b/test/python/fe_api/test_jax_gemm_amax.py @@ -0,0 +1,194 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX integration tests for block-scaled dense GEMM + amax.""" + +import pytest + + +def _jax_runtime(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + cutlass_jax = pytest.importorskip("cutlass.jax") + if not cutlass_jax.is_available(): + pytest.skip("Installed JAX version is unsupported by CUTLASS JAX") + if not hasattr(jnp, "float8_e8m0fnu"): + pytest.skip("JAX does not provide the E8M0 scale-factor dtype") + return jax, jnp + + +def _compute_capability(device): + reported = getattr(device, "compute_capability", None) + if isinstance(reported, (tuple, list)): + major, minor = int(reported[0]), int(reported[1]) + else: + text = str(reported) + if "." in text: + major_text, minor_text = text.split(".", 1) + major, minor = int(major_text), int(minor_text) + else: + major, minor = divmod(int(text), 10) + return major * 10 + minor + + +def _supported_gpu(jax): + try: + devices = tuple(jax.local_devices(backend="gpu")) + capabilities = {_compute_capability(device) for device in devices} + except RuntimeError: + pytest.skip("A JAX GPU backend is not available") + except (TypeError, ValueError): + pytest.skip("JAX GPU compute capability is unavailable") + if not devices or len(capabilities) != 1: + pytest.skip("A homogeneous JAX SM100-family GPU configuration is required") + capability = capabilities.pop() + if capability not in {100, 103, 107}: + pytest.skip("A supported JAX SM100-family GPU is not available") + return devices[0] + + +def _scale_shape(batch, rows, k, sf_vec_size): + return ( + batch, + (rows + 127) // 128, + ((k + sf_vec_size - 1) // sf_vec_size + 3) // 4, + 32, + 4, + 4, + ) + + +@pytest.mark.L0 +def test_jax_gemm_amax_abstract_contract(monkeypatch): + jax, jnp = _jax_runtime() + + from cudnn._jax import JaxApiBase + from cudnn.jax import GemmAmaxSm100, gemm_amax_wrapper_sm100 + + def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): + values = [] + for desc, spec in zip(output_descs, output_spec): + shape = self._materialize_tensor_desc(desc, mode=spec.mode).shape + fill = 0 if desc.init_value is None else desc.init_value + values.append(jnp.full(shape, fill, dtype=desc.dtype)) + return tuple(values) + + monkeypatch.setattr(JaxApiBase, "_call_kernel", abstract_call) + monkeypatch.setattr( + JaxApiBase, + "_resolve_compute_capability", + staticmethod(lambda _target, _supported, _operation, **_options: 100), + ) + monkeypatch.setattr( + JaxApiBase, + "_get_max_active_clusters", + lambda self, _cluster_size, **_options: 8, + ) + + batch, m, n, k = 2, 128, 256, 128 + sample_a = jax.ShapeDtypeStruct((batch, m, k), jnp.float8_e4m3fn) + sample_b = jax.ShapeDtypeStruct((batch, n, k), jnp.float8_e4m3fn) + sample_sfa = jax.ShapeDtypeStruct( + _scale_shape(batch, m, k, 32), + jnp.float8_e8m0fnu, + ) + sample_sfb = jax.ShapeDtypeStruct( + _scale_shape(batch, n, k, 32), + jnp.float8_e8m0fnu, + ) + operation = GemmAmaxSm100(sample_a, sample_b, sample_sfa, sample_sfb) + result = jax.eval_shape( + operation, + sample_a, + sample_b, + sample_sfa, + sample_sfb, + ) + assert tuple(result.keys()) == ("c_tensor", "amax_tensor") + assert result["c_tensor"].shape == (batch, m, n) + assert result["c_tensor"].dtype == jnp.float32 + assert result["amax_tensor"].shape == (1, 1, 1) + assert result["amax_tensor"].dtype == jnp.float32 + + alternate_a = jax.ShapeDtypeStruct((batch, k, m), jnp.float8_e4m3fn) + alternate_b = jax.ShapeDtypeStruct((batch, k, n), jnp.float8_e4m3fn) + alternate = jax.eval_shape( + lambda a, b, sfa, sfb: gemm_amax_wrapper_sm100( + a, + b, + sfa, + sfb, + a_layout="LKM", + b_layout="LKN", + c_layout="LNM", + c_dtype=jnp.bfloat16, + ), + alternate_a, + alternate_b, + sample_sfa, + sample_sfb, + ) + assert alternate["c_tensor"].shape == (batch, n, m) + assert alternate["c_tensor"].dtype == jnp.bfloat16 + + +@pytest.mark.L0 +def test_jax_gemm_amax_sm100_jit_numerics_and_amax_reset(): + jax, jnp = _jax_runtime() + device = _supported_gpu(jax) + + from cudnn.jax import gemm_amax_wrapper_sm100 + + batch, m, n, k = 1, 128, 128, 128 + a_base = jax.random.uniform( + jax.random.key(0), + (batch, m, k), + minval=-0.5, + maxval=0.5, + dtype=jnp.float32, + ).astype(jnp.float8_e4m3fn) + b = jax.random.uniform( + jax.random.key(1), + (batch, n, k), + minval=-0.5, + maxval=0.5, + dtype=jnp.float32, + ).astype(jnp.float8_e4m3fn) + a_large = (a_base.astype(jnp.float32) * 2).astype(jnp.float8_e4m3fn) + sfa = jnp.ones( + _scale_shape(batch, m, k, 32), + dtype=jnp.float8_e8m0fnu, + ) + sfb = jnp.ones( + _scale_shape(batch, n, k, 32), + dtype=jnp.float8_e8m0fnu, + ) + a_base, a_large, b, sfa, sfb = (jax.device_put(value, device) for value in (a_base, a_large, b, sfa, sfb)) + + lowered = gemm_amax_wrapper_sm100.lower(a_large, b, sfa, sfb) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + compiled = lowered.compile() + + def reference(a_value): + c = jnp.einsum( + "lmk,lnk->lmn", + a_value.astype(jnp.float32), + b.astype(jnp.float32), + ) + return c, jnp.max(jnp.abs(c)).reshape(1, 1, 1) + + # Run the larger input first. The smaller second result proves that the + # aliased amax seed is reset to -inf for each invocation. + for a_value in (a_large, a_base): + result = compiled(a_value, b, sfa, sfb) + result["amax_tensor"].block_until_ready() + reference_c, reference_amax = reference(a_value) + assert jnp.allclose(result["c_tensor"], reference_c, atol=2e-1, rtol=5e-2) + assert jnp.allclose( + result["amax_tensor"], + reference_amax, + atol=2e-1, + rtol=5e-2, + ) diff --git a/test/python/fe_api/test_jax_gemm_amax_contract.py b/test/python/fe_api/test_jax_gemm_amax_contract.py new file mode 100644 index 000000000..e510c6d5c --- /dev/null +++ b/test/python/fe_api/test_jax_gemm_amax_contract.py @@ -0,0 +1,414 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for the JAX GEMM + amax adapter.""" + +from __future__ import annotations + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +import os +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "gemm_amax" +_PACKAGE = "cudnn_jax_gemm_amax_contract_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT8 = auto() + UINT8 = auto() + FP4_E2M1 = auto() + FP8_E4M3 = auto() + FP8_E5M2 = auto() + FP8_E8M0 = auto() + + +_DTYPE_TO_CUDNN = { + "float16": _DataType.HALF, + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, + "float4_e2m1fn": _DataType.FP4_E2M1, + "float8_e4m3fn": _DataType.FP8_E4M3, + "float8_e5m2": _DataType.FP8_E5M2, + "float8_e8m0fnu": _DataType.FP8_E8M0, +} +_CUDNN_TO_DTYPE = {value: key for key, value in _DTYPE_TO_CUDNN.items()} + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + +class _TensorSpec: + def __init__(self, *, layout, mode, divisibility=None): + self.layout = layout + self.mode = mode + self.divisibility = divisibility + + +class JaxGemmAmaxContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + operation_name = f"{_PACKAGE}.gemm_amax" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + layout_module = importlib.import_module(f"{internal_name}.layout") + result_module = importlib.import_module(f"{_PACKAGE}._result") + + class JaxTensorDesc(tensor_module.TensorDesc): + @property + def cudnn_dtype(self): + return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) + + class JaxApiBase: + @staticmethod + def _to_tensor_desc( + value, + name, + *, + mode=None, + public_stride_order=None, + init_value=None, + ): + public_shape = tuple(value.shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + if public_stride_order is None: + public_stride_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride( + public_shape, + public_stride_order, + ) + canonical_axis_by_public_axis = layout_module.to_public_axes( + tuple(range(len(public_shape))), + mode, + ) + return JaxTensorDesc( + dtype=value.dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple(canonical_axis_by_public_axis[axis] for axis in public_stride_order), + name=name, + init_value=init_value, + ) + + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del operation_name + resolved = 100 if target is None else target + if resolved not in supported: + raise ValueError(f"unsupported target {resolved}") + return resolved + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + actual_shape = layout_module.to_canonical_axes(tuple(value.shape), mode) + if actual_shape != expected.shape: + raise ValueError(f"{expected.name} tensor shape mismatch: expected " f"{expected.shape}, got {actual_shape}") + actual_dtype = _DTYPE_TO_CUDNN.get(value.dtype, _DataType.NOT_SET) + if actual_dtype != expected.cudnn_dtype: + raise ValueError(f"{expected.name} tensor dtype mismatch: expected " f"{expected.cudnn_dtype}, got {actual_dtype}") + + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + mode = layout_module.normalize_mode(desc.ndim, mode) + return _TensorSpec( + layout=layout_module.to_cutlass_layout( + desc.shape, + desc.stride, + desc.stride_order, + mode=mode, + name=desc.name, + ), + mode=mode, + divisibility=divisibility, + ) + + def _call_kernel(self, inputs, **options): + self.captured_call = (tuple(inputs), options) + return tuple( + _Array( + layout_module.to_public_axes(desc.shape, spec.mode), + _CUDNN_TO_DTYPE[desc.cudnn_dtype], + ) + for desc, spec in zip( + options["output_descs"], + options["output_spec"], + ) + ) + + def _get_max_active_clusters(self, cluster_size, *, overlap_margin=0): + self.captured_active_cluster_queries = getattr( + self, + "captured_active_cluster_queries", + [], + ) + self.captured_active_cluster_queries.append((cluster_size, overlap_margin)) + return 12 - overlap_margin + + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = result_module.TupleDict + + datatypes = types.ModuleType(f"{internal_name}.datatypes") + datatypes.jax_to_cudnn_dtype = lambda dtype: _DTYPE_TO_CUDNN.get(dtype, _DataType.NOT_SET) + datatypes.normalize_jax_dtype = lambda value, default, name: (default if value is None else value) + datatypes.cudnn_to_jax_dtype = lambda dtype: _CUDNN_TO_DTYPE[dtype] + sys.modules[datatypes.__name__] = datatypes + + cls.jit_static_argnames = {} + fake_jax = types.ModuleType("jax") + + def jit(function=None, *, static_argnames=()): + def decorate(target): + cls.jit_static_argnames[target.__name__] = tuple(static_argnames) + return target + + return decorate if function is None else decorate(function) + + fake_jax.jit = jit + fake_jax.ShapeDtypeStruct = _Array + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.dtype = lambda dtype: dtype + for dtype_name in _DTYPE_TO_CUDNN: + setattr(fake_jnp, dtype_name, dtype_name) + fake_jax.numpy = fake_jnp + + try: + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "torch": None, + "cutlass": None, + }, + ): + cls.module = importlib.import_module(f"{operation_name}.jax") + except Exception: + cls.tearDownClass() + raise + + cls.operation_name = operation_name + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + @staticmethod + def _samples( + *, + a_layout="LMK", + b_layout="LNK", + ab_dtype="float8_e4m3fn", + sf_dtype="float8_e8m0fnu", + sf_vec_size=32, + ): + m, n, k, batch = 128, 256, 128, 2 + a_shape = (batch, m, k) if a_layout == "LMK" else (batch, k, m) + b_shape = (batch, n, k) if b_layout == "LNK" else (batch, k, n) + k_tiles = ((k + sf_vec_size - 1) // sf_vec_size + 3) // 4 + sfa_shape = (batch, 1, k_tiles, 32, 4, 4) + sfb_shape = (batch, 2, k_tiles, 32, 4, 4) + return ( + _Array(a_shape, ab_dtype), + _Array(b_shape, ab_dtype), + _Array(sfa_shape, sf_dtype), + _Array(sfb_shape, sf_dtype), + ) + + def test_default_layouts_bind_row_major_scales_and_infer_outputs(self): + samples = self._samples() + api = self.module.GemmAmaxSm100(*samples) + + self.assertEqual(api.a_mode, (1, 2, 0)) + self.assertEqual(api.b_mode, (1, 2, 0)) + self.assertEqual(api.output_mode, (1, 2, 0)) + self.assertEqual(api.scale_mode, (3, 4, 1, 5, 2, 0)) + self.assertEqual(api.sfa_desc.shape, (32, 4, 1, 4, 1, 2)) + self.assertEqual(api.sfa_desc.stride_order, (3, 1, 0, 4, 2, 5)) + self.assertEqual(api.amax_desc.init_value, float("-inf")) + + result = api(*samples) + self.assertEqual(tuple(result.keys()), ("c_tensor", "amax_tensor")) + self.assertEqual(result["c_tensor"].shape, (2, 128, 256)) + self.assertEqual(result["c_tensor"].dtype, "float32") + self.assertEqual(result["amax_tensor"].shape, (1, 1, 1)) + + inputs, options = api.captured_call + self.assertEqual(inputs, samples) + self.assertTrue(callable(options["launch"])) + self.assertEqual( + tuple(spec.mode for spec in options["input_spec"]), + (api.a_mode, api.b_mode, api.scale_mode, api.scale_mode), + ) + self.assertEqual(options["input_spec"][2].layout, (5, 4, 3, 2, 1, 0)) + self.assertEqual(options["input_spec"][3].layout, (5, 4, 3, 2, 1, 0)) + self.assertEqual(options["output_descs"], (api.c_desc, api.amax_desc)) + self.assertIn("--gpu-arch sm_100a", options["compile_options"]) + + def test_alternate_gemm_layouts_and_native_fp4_are_supported(self): + alternate_samples = self._samples( + a_layout="LKM", + b_layout="LKN", + ) + alternate = self.module.GemmAmaxSm100( + *alternate_samples, + a_layout="LKM", + b_layout="LKN", + c_layout="LNM", + c_dtype="bfloat16", + ) + alternate_result = alternate(*alternate_samples) + self.assertEqual(alternate.a_desc.stride_order, (0, 1, 2)) + self.assertEqual(alternate.b_desc.stride_order, (0, 1, 2)) + self.assertEqual(alternate_result["c_tensor"].shape, (2, 256, 128)) + + fp4_samples = self._samples( + ab_dtype="float4_e2m1fn", + sf_dtype="float8_e4m3fn", + sf_vec_size=16, + ) + fp4 = self.module.GemmAmaxSm100( + *fp4_samples, + c_dtype="bfloat16", + sf_vec_size=16, + ) + fp4_result = fp4(*fp4_samples) + self.assertEqual(fp4_result["c_tensor"].shape, (2, 128, 256)) + self.assertEqual(fp4_result["c_tensor"].dtype, "bfloat16") + + for c_dtype in ("float8_e4m3fn", "float8_e5m2", "float4_e2m1fn"): + with self.subTest(c_dtype=c_dtype): + quantized_output = self.module.GemmAmaxSm100( + *fp4_samples, + c_dtype=c_dtype, + sf_vec_size=16, + )(*fp4_samples) + self.assertEqual(quantized_output["c_tensor"].dtype, c_dtype) + + def test_explicit_outputs_and_runtime_signatures_are_checked(self): + samples = self._samples() + api = self.module.GemmAmaxSm100( + *samples, + sample_c=_Array((2, 128, 256), "bfloat16"), + sample_amax=_Array((1, 1, 1), "float32"), + c_dtype="bfloat16", + ) + result = api(*samples) + self.assertEqual(result["c_tensor"].dtype, "bfloat16") + + with self.assertRaisesRegex( + ValueError, + "sample_c and sample_amax must be provided together", + ): + self.module.GemmAmaxSm100( + *samples, + sample_c=_Array((2, 128, 256), "float32"), + ) + with self.assertRaisesRegex(ValueError, "sample_a tensor shape mismatch"): + api(_Array((2, 64, 128), samples[0].dtype), *samples[1:]) + + def test_wrapper_configuration_is_static(self): + self.assertEqual( + self.jit_static_argnames["gemm_amax_wrapper_sm100"], + ( + "c_layout", + "c_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "a_layout", + "b_layout", + ), + ) + samples = self._samples() + result = self.module.gemm_amax_wrapper_sm100(*samples) + self.assertEqual(result["c_tensor"].shape, (2, 128, 256)) + + def test_inline_launch_constructs_kernel_and_preserves_abi_order(self): + samples = self._samples() + with mock.patch.dict(os.environ, {"CUDNNFE_CLUSTER_OVERLAP_MARGIN": "2"}): + api = self.module.GemmAmaxSm100(*samples, cluster_shape_mn=(2, 2)) + api(*samples) + self.assertEqual(api.captured_active_cluster_queries, [(4, 2)]) + _, options = api.captured_call + seen = {} + + class Kernel: + def __init__(self, **configuration): + seen["configuration"] = configuration + + def __call__(self, *arguments): + seen["arguments"] = arguments + + kernel_module = types.ModuleType(f"{self.operation_name}.dense_blockscaled_gemm_persistent_amax") + kernel_module.Sm100BlockScaledPersistentDenseGemmKernel = Kernel + stream = object() + with mock.patch.dict(sys.modules, {kernel_module.__name__: kernel_module}): + options["launch"](stream, "A", "B", "SFA", "SFB", "C", "Amax") + + self.assertEqual( + seen["configuration"], + { + "sf_vec_size": 32, + "mma_tiler_mn": (128, 128), + "cluster_shape_mn": (2, 2), + }, + ) + self.assertEqual( + seen["arguments"], + ("A", "B", "SFA", "SFB", "C", "Amax", 10, stream), + ) + + def test_adapter_import_does_not_load_torch_or_kernel(self): + self.assertNotIn(f"{self.operation_name}.api", sys.modules) + self.assertNotIn( + f"{self.operation_name}.dense_blockscaled_gemm_persistent_amax", + sys.modules, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_gemm_layout.py b/test/python/fe_api/test_jax_gemm_layout.py index e6c0f2972..5ce45a971 100644 --- a/test/python/fe_api/test_jax_gemm_layout.py +++ b/test/python/fe_api/test_jax_gemm_layout.py @@ -96,7 +96,9 @@ def test_layout_validation_requires_a_string(self): def test_fixed_auxiliary_stride_orders_match_kernel_abis(self): self.assertEqual(self.module.ROW_MAJOR_STRIDE_ORDER_3D, (2, 1, 0)) + self.assertEqual(self.module.BLOCK_SCALE_MODE, (3, 4, 1, 5, 2, 0)) self.assertEqual(self.module.BLOCK_SCALE_STRIDE_ORDER, (3, 1, 0, 4, 2, 5)) + self.assertEqual(self.module.PROBABILITY_MODE, (2, 1, 0)) self.assertEqual(self.module.PROBABILITY_STRIDE_ORDER, (0, 1, 2)) layout = importlib.import_module(f"{self.module.__package__}.layout") diff --git a/test/python/fe_api/test_jax_gemm_relu.py b/test/python/fe_api/test_jax_gemm_relu.py new file mode 100644 index 000000000..ef12440ba --- /dev/null +++ b/test/python/fe_api/test_jax_gemm_relu.py @@ -0,0 +1,275 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX integration tests for dense block-scaled sReLU fusions.""" + +import pytest + + +def _jax_runtime(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + cutlass_jax = pytest.importorskip("cutlass.jax") + if not cutlass_jax.is_available(): + pytest.skip("Installed JAX version is unsupported by CUTLASS JAX") + return jax, jnp + + +def _compute_capability(device): + reported = getattr(device, "compute_capability", None) + if isinstance(reported, (tuple, list)): + major, minor = int(reported[0]), int(reported[1]) + else: + text = str(reported) + if "." in text: + major_text, minor_text = text.split(".", 1) + major, minor = int(major_text), int(minor_text) + else: + major, minor = divmod(int(text), 10) + return major * 10 + minor + + +def _supported_gpu(jax): + devices = tuple(device for device in jax.local_devices() if device.platform == "gpu") + try: + capabilities = {_compute_capability(device) for device in devices} + except (TypeError, ValueError): + pytest.skip("JAX GPU compute capability is unavailable") + if not devices or len(capabilities) != 1: + pytest.skip("A homogeneous JAX SM100-family GPU configuration is required") + capability = capabilities.pop() + if capability not in {100, 103, 107}: + pytest.skip("A supported JAX SM100-family GPU is not available") + return devices[0] + + +@pytest.mark.L0 +def test_jax_gemm_relu_abstract_contract(monkeypatch): + jax, jnp = _jax_runtime() + + from cudnn._jax import JaxApiBase + from cudnn.jax import GemmDsreluSm100, GemmSreluSm100, gemm_dsrelu_wrapper_sm100, gemm_srelu_wrapper_sm100 + + def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): + return tuple( + jnp.empty( + self._materialize_tensor_desc(desc, mode=spec.mode).shape, + dtype=desc.dtype, + ) + for desc, spec in zip(output_descs, output_spec) + ) + + monkeypatch.setattr(JaxApiBase, "_call_kernel", abstract_call) + monkeypatch.setattr( + JaxApiBase, + "_resolve_compute_capability", + staticmethod(lambda _target, _supported, _operation: 100), + ) + monkeypatch.setattr(JaxApiBase, "_get_max_active_clusters", lambda self, _size, **_options: 8) + + batch, m, n, k = 2, 256, 256, 512 + a = jax.ShapeDtypeStruct((batch, m, k), jnp.float8_e4m3fn) + b = jax.ShapeDtypeStruct((batch, n, k), jnp.float8_e4m3fn) + scales = jax.ShapeDtypeStruct((batch, m // 128, k // 128, 32, 4, 4), jnp.float8_e8m0fnu) + prob = jax.ShapeDtypeStruct((batch, 1, m), jnp.float32) + c = jax.ShapeDtypeStruct((batch, m, n), jnp.bfloat16) + + forward = jax.eval_shape( + GemmSreluSm100(a, b, scales, scales, prob, sf_vec_size=32), + a, + b, + scales, + scales, + prob, + ) + backward = jax.eval_shape( + GemmDsreluSm100(a, b, c, scales, scales, prob, sf_vec_size=32), + a, + b, + c, + scales, + scales, + prob, + ) + assert tuple(forward.keys()) == ("c_tensor", "d_tensor", "amax_tensor", "sfd_tensor") + assert forward["c_tensor"].shape == (batch, m, n) + assert forward["d_tensor"].shape == (batch, m, n) + assert forward["amax_tensor"] is None + assert backward["d_tensor"].shape == (batch, m, n) + assert backward["dprob_tensor"].shape == (batch, 1, m) + + alternate = jax.eval_shape( + lambda a_value, b_value, scale_a, scale_b, probability: gemm_srelu_wrapper_sm100( + a_value, + b_value, + scale_a, + scale_b, + probability, + c_layout="LNM", + sf_vec_size=32, + ), + jax.ShapeDtypeStruct((batch, k, m), jnp.float8_e4m3fn), + jax.ShapeDtypeStruct((batch, k, n), jnp.float8_e4m3fn), + scales, + scales, + prob, + ) + assert alternate["c_tensor"].shape == (batch, n, m) + + norm_const = jax.ShapeDtypeStruct((1,), jnp.float32) + with pytest.raises(NotImplementedError, match="does not implement SFD generation"): + jax.eval_shape( + lambda *args: gemm_dsrelu_wrapper_sm100( + *args[:-1], + d_dtype=jnp.float8_e4m3fn, + norm_const_tensor=args[-1], + sf_vec_size=32, + ), + a, + b, + c, + scales, + scales, + prob, + norm_const, + ) + + fp4_dtype = getattr(jnp, "float4_e2m1fn", None) + if fp4_dtype is not None: + fp4_a = jax.ShapeDtypeStruct((1, 128, 128), fp4_dtype) + fp4_b = jax.ShapeDtypeStruct((1, 128, 128), fp4_dtype) + fp4_scales = jax.ShapeDtypeStruct((1, 1, 2, 32, 4, 4), jnp.float8_e8m0fnu) + fp4_prob = jax.ShapeDtypeStruct((1, 1, 128), jnp.float32) + fp4_result = jax.eval_shape( + lambda *args: gemm_srelu_wrapper_sm100( + *args, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ), + fp4_a, + fp4_b, + fp4_scales, + fp4_scales, + fp4_prob, + ) + assert fp4_result["amax_tensor"].shape == (1,) + + +@pytest.mark.L0 +def test_jax_gemm_relu_sm100_jit_and_numerics(): + jax, jnp = _jax_runtime() + device = _supported_gpu(jax) + + from cudnn.jax import gemm_dsrelu_wrapper_sm100, gemm_srelu_wrapper_sm100 + + batch = 1 + m = n = k = 128 + alpha = 0.5 + fp8_dtype = jnp.float8_e4m3fn + a = jax.device_put( + jax.random.uniform(jax.random.key(0), (batch, m, k), minval=-0.5, maxval=0.5).astype(fp8_dtype), + device, + ) + b = jax.device_put( + jax.random.uniform(jax.random.key(1), (batch, n, k), minval=-0.5, maxval=0.5).astype(fp8_dtype), + device, + ) + c = jax.device_put( + jax.random.uniform(jax.random.key(2), (batch, m, n), minval=-0.5, maxval=0.5).astype(jnp.bfloat16), + device, + ) + prob = jax.device_put( + jax.random.uniform(jax.random.key(3), (batch, 1, m), minval=0.25, maxval=1.0), + device, + ) + scales = jax.device_put( + jnp.ones((batch, 1, 1, 32, 4, 4), dtype=jnp.float32).astype(jnp.float8_e8m0fnu), + device, + ) + + forward_lowered = gemm_srelu_wrapper_sm100.lower( + a, + b, + scales, + scales, + prob, + alpha=alpha, + sf_vec_size=32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + backward_lowered = gemm_dsrelu_wrapper_sm100.lower( + a, + b, + c, + scales, + scales, + prob, + alpha=alpha, + sf_vec_size=32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + for lowered in (forward_lowered, backward_lowered): + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + gemm = alpha * jnp.einsum("lmk,lnk->lmn", a.astype(jnp.float32), b.astype(jnp.float32)) + relu_gemm = jnp.maximum(gemm, 0.0) + probability = prob.transpose(0, 2, 1) + + forward_result = forward_lowered.compile()(a, b, scales, scales, prob) + forward_result["d_tensor"].block_until_ready() + assert jnp.allclose(forward_result["c_tensor"].astype(jnp.float32), gemm, atol=2e-1, rtol=3e-2) + assert jnp.allclose( + forward_result["d_tensor"].astype(jnp.float32), + jnp.square(relu_gemm) * probability, + atol=4e-1, + rtol=5e-2, + ) + + backward_compiled = backward_lowered.compile() + backward_result = backward_compiled(a, b, c, scales, scales, prob) + backward_result["dprob_tensor"].block_until_ready() + expected_d = c.astype(jnp.float32) * probability * 2 * relu_gemm + expected_dprob = jnp.sum(c.astype(jnp.float32) * jnp.square(relu_gemm), axis=2, keepdims=True).transpose(0, 2, 1) + assert jnp.allclose(backward_result["d_tensor"].astype(jnp.float32), expected_d, atol=4e-1, rtol=5e-2) + assert jnp.allclose(backward_result["dprob_tensor"], expected_dprob, atol=2.0, rtol=5e-2) + + zero_result = backward_compiled(a, b, jnp.zeros_like(c), scales, scales, prob) + zero_result["dprob_tensor"].block_until_ready() + assert jnp.array_equal(zero_result["dprob_tensor"], jnp.zeros_like(zero_result["dprob_tensor"])) + + +@pytest.mark.L0 +def test_jax_gemm_srelu_native_fp4_amax_smoke(): + jax, jnp = _jax_runtime() + device = _supported_gpu(jax) + fp4_dtype = getattr(jnp, "float4_e2m1fn", None) + if fp4_dtype is None: + pytest.skip("JAX does not provide native float4_e2m1fn") + + from cudnn.jax import gemm_srelu_wrapper_sm100 + + m = n = k = 128 + a = jax.device_put(jnp.ones((1, m, k), dtype=jnp.float32).astype(fp4_dtype), device) + b = jax.device_put(jnp.ones((1, n, k), dtype=jnp.float32).astype(fp4_dtype), device) + scales = jax.device_put( + jnp.ones((1, 1, 2, 32, 4, 4), dtype=jnp.float32).astype(jnp.float8_e8m0fnu), + device, + ) + prob = jax.device_put(jnp.ones((1, 1, m), dtype=jnp.float32), device) + lowered = gemm_srelu_wrapper_sm100.lower( + a, + b, + scales, + scales, + prob, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + result = lowered.compile()(a, b, scales, scales, prob) + result["amax_tensor"].block_until_ready() + assert result["amax_tensor"].shape == (1,) + assert jnp.all(jnp.isfinite(result["amax_tensor"])) diff --git a/test/python/fe_api/test_jax_gemm_relu_contract.py b/test/python/fe_api/test_jax_gemm_relu_contract.py new file mode 100644 index 000000000..6f6f3feff --- /dev/null +++ b/test/python/fe_api/test_jax_gemm_relu_contract.py @@ -0,0 +1,597 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for the JAX sReLU and dsReLU adapters.""" + +from __future__ import annotations + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +import os +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_PACKAGE = "cudnn_jax_gemm_relu_contract_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + UINT8 = auto() + FP4_E2M1 = auto() + FP8_E4M3 = auto() + FP8_E5M2 = auto() + FP8_E8M0 = auto() + + +_DTYPE_TO_CUDNN = { + "float4_e2m1fn": _DataType.FP4_E2M1, + "float8_e4m3fn": _DataType.FP8_E4M3, + "float8_e5m2": _DataType.FP8_E5M2, + "float8_e8m0fnu": _DataType.FP8_E8M0, + "float16": _DataType.HALF, + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, +} +_CUDNN_TO_DTYPE = {value: key for key, value in _DTYPE_TO_CUDNN.items()} + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + +class _TensorSpec: + def __init__(self, *, layout, mode, divisibility=None): + self.layout = layout + self.mode = mode + self.divisibility = divisibility + + +class JaxGemmReluContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + for operation in ("gemm_srelu", "gemm_dsrelu"): + operation_name = f"{_PACKAGE}.{operation}" + operation_module = types.ModuleType(operation_name) + operation_module.__path__ = [str(_CUDNN_ROOT / operation)] + operation_module.__package__ = operation_name + operation_module.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation_module + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + layout_module = importlib.import_module(f"{internal_name}.layout") + result_module = importlib.import_module(f"{_PACKAGE}._result") + + class JaxTensorDesc(tensor_module.TensorDesc): + @property + def cudnn_dtype(self): + return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) + + def compact_like(self, *, cudnn_dtype, shape, stride_order=None, name="", init_value=None): + canonical = tensor_module.make_compact_tensor_desc( + dtype=cudnn_dtype, + shape=shape, + stride_order=stride_order, + name=name, + init_value=init_value, + ) + return JaxTensorDesc( + dtype=_CUDNN_TO_DTYPE[cudnn_dtype], + shape=canonical.shape, + stride=canonical.stride, + stride_order=canonical.stride_order, + name=name, + init_value=init_value, + ) + + class JaxApiBase: + @staticmethod + def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): + public_shape = tuple(value.shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + public_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride(public_shape, public_order) + canonical_axis_by_public_axis = layout_module.to_public_axes(tuple(range(len(public_shape))), mode) + return JaxTensorDesc( + dtype=value.dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple(canonical_axis_by_public_axis[axis] for axis in public_order), + name=name, + init_value=init_value, + ) + + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del operation_name + resolved = 100 if target is None else target + if resolved not in supported: + raise ValueError(f"unsupported target {resolved}") + return resolved + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + actual_shape = layout_module.to_canonical_axes(tuple(value.shape), mode) + if actual_shape != expected.shape: + raise ValueError(f"{expected.name} tensor shape mismatch: expected {expected.shape}, got {actual_shape}") + actual_dtype = _DTYPE_TO_CUDNN.get(value.dtype, _DataType.NOT_SET) + if actual_dtype != expected.cudnn_dtype: + raise ValueError(f"{expected.name} tensor dtype mismatch: expected {expected.cudnn_dtype}, got {actual_dtype}") + + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + return _TensorSpec( + layout=layout_module.to_cutlass_layout( + desc.shape, + desc.stride, + desc.stride_order, + mode=mode, + name=desc.name, + ), + mode=mode, + divisibility=divisibility, + ) + + def _call_kernel(self, inputs, **options): + self.captured_call = (tuple(inputs), options) + return tuple( + _Array( + layout_module.to_public_axes(desc.shape, spec.mode), + _CUDNN_TO_DTYPE[desc.cudnn_dtype], + ) + for desc, spec in zip(options["output_descs"], options["output_spec"]) + ) + + def _get_max_active_clusters(self, cluster_size, *, overlap_margin=0): + self.captured_active_cluster_query = (cluster_size, overlap_margin) + return 12 - overlap_margin + + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = result_module.TupleDict + + datatypes = types.ModuleType(f"{internal_name}.datatypes") + datatypes.jax_to_cudnn_dtype = lambda dtype: _DTYPE_TO_CUDNN.get(dtype, _DataType.NOT_SET) + datatypes.normalize_jax_dtype = lambda value, default, name: default if value is None else value + datatypes.cudnn_to_jax_dtype = lambda dtype: _CUDNN_TO_DTYPE[dtype] + sys.modules[datatypes.__name__] = datatypes + + cls.jit_static_argnames = {} + fake_jax = types.ModuleType("jax") + + def jit(function=None, *, static_argnames=()): + def decorate(target): + cls.jit_static_argnames[target.__name__] = tuple(static_argnames) + return target + + return decorate if function is None else decorate(function) + + fake_jax.jit = jit + fake_jax.ShapeDtypeStruct = _Array + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.dtype = lambda dtype: dtype + for dtype_name in _DTYPE_TO_CUDNN: + setattr(fake_jnp, dtype_name, dtype_name) + fake_jax.numpy = fake_jnp + + try: + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "torch": None, + "cutlass": None, + }, + ): + cls.srelu = importlib.import_module(f"{_PACKAGE}.gemm_srelu.jax") + cls.dsrelu = importlib.import_module(f"{_PACKAGE}.gemm_dsrelu.jax") + except Exception: + cls.tearDownClass() + raise + + cls.operation_names = { + cls.srelu: f"{_PACKAGE}.gemm_srelu", + cls.dsrelu: f"{_PACKAGE}.gemm_dsrelu", + } + cls.tensor_module = tensor_module + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + @staticmethod + def _samples(*, dtype="float8_e4m3fn", sf_vec_size=32, batch=2, m=256, n=256, k=512): + rest_k = ((k + sf_vec_size - 1) // sf_vec_size + 3) // 4 + a = _Array((batch, m, k), dtype) + b = _Array((batch, n, k), dtype) + sfa = _Array((batch, (m + 127) // 128, rest_k, 32, 4, 4), "float8_e8m0fnu") + sfb = _Array((batch, (n + 127) // 128, rest_k, 32, 4, 4), "float8_e8m0fnu") + prob = _Array((batch, 1, m), "float32") + c = _Array((batch, m, n), "bfloat16") + return a, b, c, sfa, sfb, prob + + def test_forward_and_backward_use_row_major_public_auxiliary_layouts(self): + a, b, c, sfa, sfb, prob = self._samples() + forward = self.srelu.GemmSreluSm100(a, b, sfa, sfb, prob, sf_vec_size=32) + backward = self.dsrelu.GemmDsreluSm100(a, b, c, sfa, sfb, prob, sf_vec_size=32) + + forward_result = forward(a, b, sfa, sfb, prob) + backward_result = backward(a, b, c, sfa, sfb, prob) + + self.assertEqual(forward.sfa_desc.shape, (32, 4, 2, 4, 4, 2)) + self.assertEqual(forward.sfa_desc.stride_order, (3, 1, 0, 4, 2, 5)) + self.assertEqual(forward.prob_desc.shape, (256, 1, 2)) + self.assertEqual(forward.prob_desc.stride_order, (0, 1, 2)) + self.assertIs(forward._op.sfa, forward.sfa_desc) + self.assertIs(backward._op.dprob, backward.dprob_desc) + + self.assertEqual(forward_result["c_tensor"].shape, (2, 256, 256)) + self.assertEqual(forward_result["d_tensor"].shape, (2, 256, 256)) + self.assertIsNone(forward_result["amax_tensor"]) + self.assertEqual(backward_result["d_tensor"].shape, (2, 256, 256)) + self.assertEqual(backward_result["dprob_tensor"].shape, (2, 1, 256)) + self.assertEqual(backward.dprob_desc.init_value, 0.0) + + forward_inputs, forward_options = forward.captured_call + self.assertEqual(forward_inputs, (a, b, sfa, sfb, prob)) + self.assertTrue(callable(forward_options["launch"])) + self.assertEqual( + tuple(spec.mode for spec in forward_options["input_spec"]), + (forward.a_mode, forward.b_mode, forward.scale_mode, forward.scale_mode, forward.probability_mode), + ) + self.assertEqual(forward_options["input_spec"][2].layout, (5, 4, 3, 2, 1, 0)) + self.assertEqual(forward_options["input_spec"][4].layout, (2, 1, 0)) + self.assertIn("--gpu-arch sm_100a", forward_options["compile_options"]) + + def test_native_fp4_returns_initialized_amax(self): + a, b, c, sfa, sfb, prob = self._samples(dtype="float4_e2m1fn", sf_vec_size=16) + forward = self.srelu.GemmSreluSm100(a, b, sfa, sfb, prob, sf_vec_size=16) + backward = self.dsrelu.GemmDsreluSm100(a, b, c, sfa, sfb, prob, sf_vec_size=16) + + forward_result = forward(a, b, sfa, sfb, prob) + backward_result = backward(a, b, c, sfa, sfb, prob) + + for api, result in ((forward, forward_result), (backward, backward_result)): + self.assertIsNotNone(api.amax_desc) + self.assertEqual(api.amax_desc.shape, (1,)) + self.assertEqual(api.amax_desc.init_value, float("-inf")) + self.assertEqual(result["amax_tensor"].shape, (1,)) + self.assertIs(api.captured_call[1]["output_descs"][-1], api.amax_desc) + + def test_fp8_d_output_is_rejected_until_sfd_generation_is_implemented(self): + a, b, c, sfa, sfb, prob = self._samples() + norm_const = _Array((1,), "float32") + constructors = ( + lambda: self.srelu.GemmSreluSm100( + a, + b, + sfa, + sfb, + prob, + sample_norm_const=norm_const, + d_dtype="float8_e4m3fn", + sf_vec_size=32, + ), + lambda: self.dsrelu.GemmDsreluSm100( + a, + b, + c, + sfa, + sfb, + prob, + sample_norm_const=norm_const, + d_dtype="float8_e4m3fn", + sf_vec_size=32, + ), + ) + for construct in constructors: + with self.subTest(construct=construct), self.assertRaisesRegex(NotImplementedError, "does not implement SFD generation"): + construct() + + def test_unit_scale_tiles_and_batch_preserve_packed_layout(self): + a, b, c, sfa, sfb, prob = self._samples(batch=1, m=128, n=128, k=128) + forward = self.srelu.GemmSreluSm100( + a, + b, + sfa, + sfb, + prob, + sf_vec_size=32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + backward = self.dsrelu.GemmDsreluSm100( + a, + b, + c, + sfa, + sfb, + prob, + sf_vec_size=32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + + self.assertTrue(forward.check_support()) + self.assertTrue(backward.check_support()) + self.assertEqual(forward.sfa_desc.shape, (32, 4, 1, 4, 1, 1)) + self.assertTrue(forward.sfa_desc.is_compact((3, 1, 0, 4, 2, 5))) + + def test_launch_callbacks_preserve_native_kernel_abis(self): + a, b, c, sfa, sfb, prob = self._samples(dtype="float4_e2m1fn", sf_vec_size=16) + with mock.patch.dict(os.environ, {"CUDNNFE_CLUSTER_OVERLAP_MARGIN": "2"}): + forward = self.srelu.GemmSreluSm100(a, b, sfa, sfb, prob, alpha=0.25) + backward = self.dsrelu.GemmDsreluSm100(a, b, c, sfa, sfb, prob, alpha=0.5) + forward(a, b, sfa, sfb, prob) + backward(a, b, c, sfa, sfb, prob) + self.assertEqual(forward.captured_active_cluster_query, (2, 2)) + self.assertEqual(backward.captured_active_cluster_query, (2, 2)) + + calls = {} + + def kernel_type(label): + class Kernel: + def __init__(self, **configuration): + calls[f"{label}_configuration"] = configuration + + def __call__(self, *arguments, **keywords): + calls[f"{label}_arguments"] = arguments + calls[f"{label}_keywords"] = keywords + + return Kernel + + cutlass = types.ModuleType("cutlass") + cutlass.Float32 = lambda value: ("Float32", value) + cute = types.ModuleType("cutlass.cute") + cute.where = lambda condition, yes, no: yes if condition else no + cute.full_like = lambda value, fill: fill + cutlass.cute = cute + + forward_kernel_name = f"{self.operation_names[self.srelu]}.dense_blockscaled_gemm_persistent_srelu_quant" + forward_kernel = types.ModuleType(forward_kernel_name) + forward_kernel.Sm100BlockScaledPersistentDenseGemmKernel = kernel_type("forward") + backward_kernel_name = f"{self.operation_names[self.dsrelu]}.dense_blockscaled_gemm_persistent_dsrelu_quant" + backward_kernel = types.ModuleType(backward_kernel_name) + backward_kernel.Sm100BlockScaledPersistentDenseGemmKernel = kernel_type("backward") + + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + "cutlass.cute": cute, + forward_kernel_name: forward_kernel, + backward_kernel_name: backward_kernel, + }, + ): + forward.captured_call[1]["launch"]("stream", "A", "B", "SFA", "SFB", "PROB", "C", "D", "AMAX") + backward.captured_call[1]["launch"]( + "stream", + "A", + "B", + "C", + "SFA", + "SFB", + "PROB", + "D", + "DPROB", + "AMAX", + ) + + self.assertEqual( + calls["forward_arguments"], + ("A", "B", "SFA", "SFB", "C", "D", "PROB", "AMAX", None, None, ("Float32", 0.25), 10, "stream"), + ) + self.assertEqual( + calls["backward_arguments"], + ("A", "B", "SFA", "SFB", "C", "D", "PROB", "DPROB", "AMAX", None, None, ("Float32", 0.5), 10, "stream"), + ) + self.assertEqual(set(calls["forward_keywords"]), {"epilogue_op"}) + self.assertEqual(set(calls["backward_keywords"]), {"epilogue_op"}) + + def test_wrappers_are_jitted_with_static_configuration(self): + for wrapper_name in ("gemm_srelu_wrapper_sm100", "gemm_dsrelu_wrapper_sm100"): + static = self.jit_static_argnames[wrapper_name] + self.assertIn("mma_tiler_mn", static) + self.assertIn("cluster_shape_mn", static) + self.assertIn("sf_vec_size", static) + self.assertIn("a_layout", static) + self.assertIn("b_layout", static) + + def test_validation_rejects_bad_scale_shape_and_partial_mma_rows(self): + a, b, _, sfa, sfb, prob = self._samples() + bad_sfa = _Array((2, 2, 3, 32, 4, 4), "float8_e8m0fnu") + with self.assertRaisesRegex(ValueError, "SFA must have shape"): + self.srelu.GemmSreluSm100(a, b, bad_sfa, sfb, prob, sf_vec_size=32)(a, b, bad_sfa, sfb, prob) + + partial_a, partial_b, _, partial_sfa, partial_sfb, partial_prob = self._samples(m=129) + with self.assertRaisesRegex(ValueError, "CTA_TILE_M=128"): + self.srelu.GemmSreluSm100( + partial_a, + partial_b, + partial_sfa, + partial_sfb, + partial_prob, + sf_vec_size=32, + )(partial_a, partial_b, partial_sfa, partial_sfb, partial_prob) + + def test_shared_op_normalizes_legacy_uint8_storage_to_logical_fp4(self): + make_desc = self.tensor_module.make_compact_tensor_desc + a = make_desc(dtype=_DataType.UINT8, shape=(128, 128, 1), stride_order=(1, 0, 2), name="A") + b = make_desc(dtype=_DataType.UINT8, shape=(128, 128, 1), stride_order=(1, 0, 2), name="B") + c = make_desc(dtype=_DataType.BFLOAT16, shape=(128, 128, 1), stride_order=(1, 0, 2), name="C") + d = make_desc(dtype=_DataType.BFLOAT16, shape=(128, 128, 1), stride_order=(1, 0, 2), name="D") + sfa = make_desc(dtype=_DataType.FP8_E8M0, shape=(32, 4, 1, 4, 2, 1), stride_order=(3, 1, 0, 4, 2, 5), name="SFA") + sfb = make_desc(dtype=_DataType.FP8_E8M0, shape=(32, 4, 1, 4, 2, 1), stride_order=(3, 1, 0, 4, 2, 5), name="SFB") + prob = make_desc(dtype=_DataType.FLOAT, shape=(128, 1, 1), stride_order=(0, 1, 2), name="prob") + op = self.srelu.GemmSreluSm100Op( + a=a, + b=b, + c=c, + d=d, + sfa=sfa, + sfb=sfb, + prob=prob, + sf_vec_size=16, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + + self.assertTrue(op.check_support()) + self.assertEqual(op.ab_dtype, _DataType.FP4_E2M1) + + with self.assertRaisesRegex(ValueError, "at most 4"): + self.srelu.GemmSreluSm100Op( + a=a, + b=b, + c=c, + d=d, + sfa=sfa, + sfb=sfb, + prob=prob, + sf_vec_size=16, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(8, 1), + ).check_support() + + def test_relu_ops_enforce_forward_and_backward_specific_signatures(self): + make_desc = self.tensor_module.make_compact_tensor_desc + a = make_desc(dtype=_DataType.FP4_E2M1, shape=(128, 128, 1), stride_order=(1, 0, 2), name="A") + b = make_desc(dtype=_DataType.FP4_E2M1, shape=(128, 128, 1), stride_order=(1, 0, 2), name="B") + c = make_desc(dtype=_DataType.BFLOAT16, shape=(128, 128, 1), stride_order=(1, 0, 2), name="C") + d = make_desc(dtype=_DataType.BFLOAT16, shape=(128, 128, 1), stride_order=(1, 0, 2), name="D") + dprob = make_desc(dtype=_DataType.FLOAT, shape=(128, 1, 1), stride_order=(0, 1, 2), name="dprob") + sfa = make_desc(dtype=_DataType.FP8_E8M0, shape=(32, 4, 1, 4, 2, 1), stride_order=(3, 1, 0, 4, 2, 5), name="SFA") + sfb = make_desc(dtype=_DataType.FP8_E8M0, shape=(32, 4, 1, 4, 2, 1), stride_order=(3, 1, 0, 4, 2, 5), name="SFB") + prob = make_desc(dtype=_DataType.FLOAT, shape=(128, 1, 1), stride_order=(0, 1, 2), name="prob") + arguments = dict( + a=a, + b=b, + c=c, + d=d, + sfa=sfa, + sfb=sfb, + prob=prob, + sf_vec_size=16, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + + with self.assertRaisesRegex(ValueError, "dprob is only part of"): + self.srelu.GemmSreluSm100Op(**arguments, dprob=dprob).check_support() + with self.assertRaisesRegex(ValueError, "dprob is required"): + self.dsrelu.GemmDsreluSm100Op(**arguments).check_support() + self.assertTrue(self.dsrelu.GemmDsreluSm100Op(**arguments, dprob=dprob).check_support()) + + fp8_c = make_desc(dtype=_DataType.FP8_E4M3, shape=(128, 128, 1), stride_order=(1, 0, 2), name="C") + with self.assertRaisesRegex(ValueError, "dsReLU C must use"): + self.dsrelu.GemmDsreluSm100Op(**{**arguments, "c": fp8_c}, dprob=dprob).check_support() + + def test_shared_op_rejects_unimplemented_fp8_output(self): + make_desc = self.tensor_module.make_compact_tensor_desc + a = make_desc(dtype=_DataType.FP8_E4M3, shape=(128, 128, 1), stride_order=(1, 0, 2), name="A") + b = make_desc(dtype=_DataType.FP8_E4M3, shape=(128, 128, 1), stride_order=(1, 0, 2), name="B") + c = make_desc(dtype=_DataType.BFLOAT16, shape=(128, 128, 1), stride_order=(1, 0, 2), name="C") + d = make_desc(dtype=_DataType.FP8_E4M3, shape=(128, 128, 1), stride_order=(1, 0, 2), name="D") + sfa = make_desc(dtype=_DataType.FP8_E8M0, shape=(32, 4, 1, 4, 1, 1), stride_order=(3, 1, 0, 4, 2, 5), name="SFA") + sfb = make_desc(dtype=_DataType.FP8_E8M0, shape=(32, 4, 1, 4, 1, 1), stride_order=(3, 1, 0, 4, 2, 5), name="SFB") + prob = make_desc(dtype=_DataType.FLOAT, shape=(128, 1, 1), stride_order=(0, 1, 2), name="prob") + arguments = dict( + a=a, + b=b, + c=c, + d=d, + sfa=sfa, + sfb=sfb, + prob=prob, + sf_vec_size=32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + with self.assertRaisesRegex(NotImplementedError, "SFD generation is not implemented"): + self.srelu.GemmSreluSm100Op(**arguments).check_support() + + sfd = make_desc(dtype=_DataType.FP8_E8M0, shape=(32, 4, 1, 4, 1, 1), stride_order=(3, 1, 0, 4, 2, 5), name="SFD") + norm_const = make_desc(dtype=_DataType.FLOAT, shape=(1,), name="norm_const") + with self.assertRaisesRegex(NotImplementedError, "SFD generation is not implemented"): + self.srelu.GemmSreluSm100Op( + **arguments, + sfd=sfd, + norm_const=norm_const, + ).check_support() + + def test_shared_op_rejects_scale_formats_that_do_not_match_the_mma(self): + make_desc = self.tensor_module.make_compact_tensor_desc + + def descriptor(dtype, shape, stride_order, name): + return make_desc(dtype=dtype, shape=shape, stride_order=stride_order, name=name) + + c = descriptor(_DataType.BFLOAT16, (128, 128, 1), (1, 0, 2), "C") + d = descriptor(_DataType.BFLOAT16, (128, 128, 1), (1, 0, 2), "D") + prob = descriptor(_DataType.FLOAT, (128, 1, 1), (0, 1, 2), "prob") + + cases = ( + (_DataType.FP8_E4M3, _DataType.FP8_E4M3, 16, 2, "FP8 inputs require FP8_E8M0"), + (_DataType.FP4_E2M1, _DataType.FP8_E4M3, 32, 1, "FP4 inputs with sf_vec_size=32 require FP8_E8M0"), + ) + for ab_dtype, sf_dtype, sf_vec_size, rest_k, message in cases: + a = descriptor(ab_dtype, (128, 128, 1), (1, 0, 2), "A") + b = descriptor(ab_dtype, (128, 128, 1), (1, 0, 2), "B") + sfa = descriptor(sf_dtype, (32, 4, 1, 4, rest_k, 1), (3, 1, 0, 4, 2, 5), "SFA") + sfb = descriptor(sf_dtype, (32, 4, 1, 4, rest_k, 1), (3, 1, 0, 4, 2, 5), "SFB") + with self.subTest(ab_dtype=ab_dtype, sf_dtype=sf_dtype, sf_vec_size=sf_vec_size), self.assertRaisesRegex(ValueError, message): + self.srelu.GemmSreluSm100Op( + a=a, + b=b, + c=c, + d=d, + sfa=sfa, + sfb=sfb, + prob=prob, + sf_vec_size=sf_vec_size, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ).check_support() + + def test_jax_modules_do_not_load_torch_or_kernels(self): + self.assertNotIn(f"{_PACKAGE}.gemm_srelu.api", sys.modules) + self.assertNotIn(f"{_PACKAGE}.gemm_dsrelu.api", sys.modules) + self.assertNotIn(f"{_PACKAGE}.gemm_srelu.dense_blockscaled_gemm_persistent_srelu_quant", sys.modules) + self.assertNotIn(f"{_PACKAGE}.gemm_dsrelu.dense_blockscaled_gemm_persistent_dsrelu_quant", sys.modules) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_gemm_swiglu.py b/test/python/fe_api/test_jax_gemm_swiglu.py index 75a8911f3..6f6737eb2 100644 --- a/test/python/fe_api/test_jax_gemm_swiglu.py +++ b/test/python/fe_api/test_jax_gemm_swiglu.py @@ -117,6 +117,57 @@ def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): assert explicit_result["c_tensor"].dtype == jnp.bfloat16 +@pytest.mark.L0 +def test_jax_gemm_swiglu_block_scaled_sm100_jit_and_numerics(): + jax, jnp = _jax_runtime() + device = _supported_gpu(jax) + if not hasattr(jnp, "float8_e4m3fn") or not hasattr(jnp, "float8_e8m0fnu"): + pytest.skip("The installed JAX runtime does not expose MXFP8 dtypes") + + from cudnn.jax import gemm_swiglu_wrapper_sm100 + + batch, m, n, k = 1, 128, 128, 128 + alpha = 0.5 + a = jax.device_put( + (jax.random.normal(jax.random.key(10), (batch, m, k), dtype=jnp.float32) * 0.1).astype(jnp.float8_e4m3fn), + device, + ) + b = jax.device_put( + (jax.random.normal(jax.random.key(11), (batch, n, k), dtype=jnp.float32) * 0.1).astype(jnp.float8_e4m3fn), + device, + ) + scale_shape = (batch, 1, 1, 32, 4, 4) + sfa = jax.device_put(jnp.ones(scale_shape, dtype=jnp.float32).astype(jnp.float8_e8m0fnu), device) + sfb = jax.device_put(jnp.ones(scale_shape, dtype=jnp.float32).astype(jnp.float8_e8m0fnu), device) + + lowered = gemm_swiglu_wrapper_sm100.lower( + a, + b, + alpha=alpha, + ab12_dtype=jnp.bfloat16, + c_dtype=jnp.bfloat16, + sfa_tensor=sfa, + sfb_tensor=sfb, + sf_vec_size=32, + ) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + result = lowered.compile()(a, b, sfa_tensor=sfa, sfb_tensor=sfb) + result["c_tensor"].block_until_ready() + reference_ab12 = alpha * jnp.einsum("lmk,lnk->lmn", a.astype(jnp.float32), b.astype(jnp.float32)) + blocks = reference_ab12.reshape(batch, m, n // 32, 32) + input_blocks = blocks[:, :, 0::2].reshape(batch, m, n // 2) + gate_blocks = blocks[:, :, 1::2].reshape(batch, m, n // 2) + reference_c = input_blocks * jax.nn.silu(gate_blocks) + + assert jnp.allclose(result["ab12_tensor"].astype(jnp.float32), reference_ab12, atol=8e-2, rtol=4e-2) + assert jnp.allclose(result["c_tensor"].astype(jnp.float32), reference_c, atol=1e-1, rtol=5e-2) + assert result["sfc_tensor"] is None + assert result["amax_tensor"] is None + + @pytest.mark.L0 def test_jax_gemm_swiglu_sm100_jit_and_numerics(): jax, jnp = _jax_runtime() diff --git a/test/python/fe_api/test_jax_gemm_swiglu_contract.py b/test/python/fe_api/test_jax_gemm_swiglu_contract.py index ec123a664..32b368528 100644 --- a/test/python/fe_api/test_jax_gemm_swiglu_contract.py +++ b/test/python/fe_api/test_jax_gemm_swiglu_contract.py @@ -35,12 +35,18 @@ class _DataType(Enum): FLOAT = auto() FP8_E4M3 = auto() FP8_E5M2 = auto() + FP8_E8M0 = auto() + FP4_E2M1 = auto() _DTYPE_TO_CUDNN = { "float16": _DataType.HALF, "bfloat16": _DataType.BFLOAT16, "float32": _DataType.FLOAT, + "float4_e2m1fn": _DataType.FP4_E2M1, + "float8_e4m3fn": _DataType.FP8_E4M3, + "float8_e5m2": _DataType.FP8_E5M2, + "float8_e8m0fnu": _DataType.FP8_E8M0, } _CUDNN_TO_DTYPE = {value: key for key, value in _DTYPE_TO_CUDNN.items()} @@ -139,8 +145,8 @@ def _to_tensor_spec(desc, *, mode=None, divisibility=None): divisibility=divisibility, ) - def _call_kernel(self, inputs, **options): - self.captured_max_active_clusters_at_call = getattr(self, "_max_active_clusters", None) + def _call_kernel(self, inputs, *, launch, **options): + options["launch"] = launch self.captured_call = (tuple(inputs), options) return tuple( _Array( @@ -183,6 +189,10 @@ def decorate(target): fake_jnp.float16 = "float16" fake_jnp.bfloat16 = "bfloat16" fake_jnp.float32 = "float32" + fake_jnp.float4_e2m1fn = "float4_e2m1fn" + fake_jnp.float8_e4m3fn = "float8_e4m3fn" + fake_jnp.float8_e5m2 = "float8_e5m2" + fake_jnp.float8_e8m0fnu = "float8_e8m0fnu" fake_jax.numpy = fake_jnp try: @@ -252,6 +262,7 @@ def test_default_layouts_infer_outputs_and_declare_cutlass_call(self): inputs, options = api.captured_call self.assertEqual(inputs, (sample_a, sample_b)) + self.assertTrue(callable(options["launch"])) self.assertEqual(options["output_descs"], (api.ab12_desc, api.c_desc)) self.assertEqual( tuple(spec.mode for spec in options["input_spec"]), @@ -337,6 +348,29 @@ def test_runtime_inputs_must_match_the_specialized_signature(self): with self.assertRaisesRegex(ValueError, "sample_b tensor dtype mismatch"): api(sample_a, _Array(sample_b.shape, "float16")) + def test_rejects_unreachable_fp8_c_and_block_only_standard_options(self): + sample_a, sample_b = self._samples() + sample_ab12 = _Array((3, 128, 192), "bfloat16") + sample_c_fp8 = _Array((3, 128, 96), "float8_e4m3fn") + with self.assertRaisesRegex(NotImplementedError, "does not support FP8 C"): + self.module.GemmSwigluSm100( + sample_a, + sample_b, + sample_ab12=sample_ab12, + sample_c=sample_c_fp8, + ) + + with self.assertRaisesRegex(NotImplementedError, "does not expose FP8 C or SFC"): + self.module.GemmSwigluSm100( + sample_a, + sample_b, + sample_sfc=_Array((1,), "float8_e8m0fnu"), + ) + + for option, value in (("sf_vec_size", 32), ("vector_f32", True), ("ab12_stages", 3)): + with self.subTest(option=option), self.assertRaisesRegex(ValueError, "only applies to block-scaled"): + self.module.GemmSwigluSm100(sample_a, sample_b, **{option: value}) + def test_wrapper_marks_configuration_as_static_and_remains_functional(self): self.assertIn("gemm_swiglu_wrapper_sm100", self.module.__all__) self.assertNotIn("gemm_swiglu_wrapper", self.module.__all__) @@ -351,6 +385,9 @@ def test_wrapper_marks_configuration_as_static_and_remains_functional(self): "acc_dtype", "mma_tiler_mn", "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "ab12_stages", "a_layout", "b_layout", ), @@ -367,6 +404,122 @@ def test_wrapper_marks_configuration_as_static_and_remains_functional(self): self.assertEqual(result["ab12_tensor"].shape, (3, 192, 128)) self.assertEqual(result["c_tensor"].shape, (3, 96, 128)) + def test_block_scaled_mxfp8_uses_packed_scale_modes_and_kernel_abi(self): + batch, m, n, k = 2, 128, 128, 128 + sample_a = _Array((batch, m, k), "float8_e4m3fn") + sample_b = _Array((batch, n, k), "float8_e4m3fn") + scale_shape = (batch, 1, 1, 32, 4, 4) + sample_sfa = _Array(scale_shape, "float8_e8m0fnu") + sample_sfb = _Array(scale_shape, "float8_e8m0fnu") + api = self.module.GemmSwigluSm100( + sample_a, + sample_b, + sample_sfa=sample_sfa, + sample_sfb=sample_sfb, + alpha=0.5, + ab12_dtype="bfloat16", + c_dtype="bfloat16", + sf_vec_size=32, + vector_f32=True, + ab12_stages=3, + ) + + self.assertTrue(api.is_block_scaled) + self.assertEqual(api.scale_mode, (3, 4, 1, 5, 2, 0)) + self.assertEqual(api.sfa_desc.shape, (32, 4, 1, 4, 1, batch)) + self.assertEqual(api.sfa_desc.stride_order, (3, 1, 0, 4, 2, 5)) + result = api(sample_a, sample_b, sample_sfa, sample_sfb) + self.assertEqual(result["ab12_tensor"].shape, (batch, m, n)) + self.assertEqual(result["ab12_tensor"].dtype, "bfloat16") + self.assertEqual(result["c_tensor"].shape, (batch, m, n // 2)) + self.assertEqual(result["c_tensor"].dtype, "bfloat16") + self.assertIsNone(result["sfc_tensor"]) + self.assertIsNone(result["amax_tensor"]) + + inputs, options = api.captured_call + self.assertEqual(inputs, (sample_a, sample_b, sample_sfa, sample_sfb)) + self.assertEqual( + tuple(spec.mode for spec in options["input_spec"]), + (api.a_mode, api.b_mode, api.scale_mode, api.scale_mode), + ) + self.assertEqual(tuple(spec.mode for spec in options["output_spec"]), (api.output_mode, api.output_mode)) + self.assertEqual(api.captured_active_cluster_queries, [(1, 0)]) + + seen = {} + + class Kernel: + def __init__(self, **configuration): + seen["configuration"] = configuration + + def __call__(self, *arguments): + seen["arguments"] = arguments + + cutlass = types.ModuleType("cutlass") + cutlass.Float32 = lambda value: ("Float32", value) + kernel_module = types.ModuleType(f"{self.operation_name}.dense_blockscaled_gemm_persistent_swiglu_interleaved_quant") + kernel_module.Sm100BlockScaledPersistentDenseGemmKernel = Kernel + stream = object() + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + kernel_module.__name__: kernel_module, + }, + ): + options["launch"](stream, "A", "B", "SFA", "SFB", "AB12", "C") + + self.assertEqual( + seen["configuration"], + { + "sf_vec_size": 32, + "mma_tiler_mn": (128, 128), + "cluster_shape_mn": (1, 1), + "vector_f32": True, + "ab12_stages": 3, + }, + ) + self.assertEqual( + seen["arguments"], + ("A", "B", "SFA", "SFB", "C", "AB12", None, None, None, ("Float32", 0.5), 12, stream), + ) + + wrapped = self.module.gemm_swiglu_wrapper_sm100( + sample_a, + sample_b, + alpha=0.5, + ab12_dtype="bfloat16", + c_dtype="bfloat16", + sfa_tensor=sample_sfa, + sfb_tensor=sample_sfb, + sf_vec_size=32, + ) + self.assertEqual(wrapped["ab12_tensor"].shape, (batch, m, n)) + self.assertEqual(wrapped["c_tensor"].shape, (batch, m, n // 2)) + + def test_block_scaled_fp4_infers_initialized_amax(self): + sample_a = _Array((1, 128, 128), "float4_e2m1fn") + sample_b = _Array((1, 128, 128), "float4_e2m1fn") + scale_shape = (1, 1, 2, 32, 4, 4) + sample_sfa = _Array(scale_shape, "float8_e8m0fnu") + sample_sfb = _Array(scale_shape, "float8_e8m0fnu") + api = self.module.GemmSwigluSm100( + sample_a, + sample_b, + sample_sfa=sample_sfa, + sample_sfb=sample_sfb, + c_dtype="bfloat16", + sf_vec_size=16, + ) + + result = api(sample_a, sample_b, sample_sfa, sample_sfb) + self.assertEqual(result["amax_tensor"].shape, (1,)) + self.assertEqual(result["amax_tensor"].dtype, "float32") + self.assertEqual(api.amax_desc.init_value, float("-inf")) + self.assertEqual(api.captured_call[1]["output_descs"], (api.ab12_desc, api.c_desc, api.amax_desc)) + + with self.assertRaisesRegex(ValueError, "sample_sfa and sample_sfb are required"): + self.module.GemmSwigluSm100(sample_a, sample_b, sample_sfa=sample_sfa) + def test_launch_constructs_the_kernel_internally_and_preserves_abi_order(self): _, sample_b = self._samples() sample_a = _Array((3, 256, 64), "bfloat16") @@ -380,12 +533,12 @@ def test_launch_constructs_the_kernel_internally_and_preserves_abi_order(self): self.assertFalse(hasattr(api, "captured_active_cluster_queries")) api(sample_a, sample_b) self.assertEqual(api.captured_active_cluster_queries, [(4, 2)]) - self.assertEqual(api.captured_max_active_clusters_at_call, 10) + launch = api.captured_call[1]["launch"] seen = {} class HardwareInfo: def get_max_active_clusters(self, cluster_size): - raise AssertionError(f"HardwareInfo queried from _launch for cluster size {cluster_size}") + raise AssertionError(f"HardwareInfo queried from launch callback for cluster size {cluster_size}") class Kernel: def __init__(self, **configuration): @@ -412,7 +565,7 @@ def __call__(self, *arguments): kernel_module.__name__: kernel_module, }, ): - api._launch(("A", "B"), ("AB12", "C"), (), stream) + launch(stream, "A", "B", "AB12", "C") self.assertEqual( seen["configuration"], @@ -428,8 +581,8 @@ def __call__(self, *arguments): ("A", "B", "AB12", "C", ("Float32", 0.25), 10, stream), ) - with self.assertRaisesRegex(RuntimeError, "unexpected workspaces"): - api._launch(("A", "B"), ("AB12", "C"), (object(),), stream) + with self.assertRaises(TypeError): + launch(stream, "A", "B", "AB12", "C", object()) def test_adapter_import_does_not_load_torch_or_the_cute_kernel(self): self.assertNotIn(f"{self.operation_name}.api", sys.modules) diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_imports.py b/test/python/fe_api/test_rmsnorm_rht_amax_imports.py index 736677b05..0f576f44e 100644 --- a/test/python/fe_api/test_rmsnorm_rht_amax_imports.py +++ b/test/python/fe_api/test_rmsnorm_rht_amax_imports.py @@ -258,10 +258,17 @@ def imports_framework(node, framework: str): ) jax_check_support = next(node for node in jax_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "check_support") self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "_check_device_compatibility" for node in ast.walk(jax_check_support))) - jax_launch = next(node for node in jax_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "_launch") + jax_call = next(node for node in jax_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "__call__") + jax_launch = next(node for node in ast.walk(jax_call) if isinstance(node, ast.FunctionDef) and node.name == "launch") self.assertTrue( any(isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RMSNormRHTAmaxKernel" for node in ast.walk(jax_launch)) ) + jax_kernel_call = next( + node for node in ast.walk(jax_call) if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "_call_kernel" + ) + launch_keyword = next(keyword.value for keyword in jax_kernel_call.keywords if keyword.arg == "launch") + self.assertIsInstance(launch_keyword, ast.Name) + self.assertEqual(launch_keyword.id, "launch") jax_wrapper = next(node for node in jax_tree.body if isinstance(node, ast.FunctionDef) and node.name == "rmsnorm_rht_amax_sm100") jit_decorator = next( From f14710e382dbdf72ba3a4f20578dee5c4a9b0a6e Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Wed, 8 Jul 2026 08:22:24 -0500 Subject: [PATCH 12/13] Add JAX support for remaining fused operators Add grouped and discrete GEMMs, native/block sparse attention, and specialized SDPA adapters. Extend sequence layout handling, functional output allocation, native FP4 coverage, and contract tests. --- docs/fe-oss-apis/dsa.md | 14 +- docs/fe-oss-apis/nsa.md | 40 + python/cudnn/_cute_compiler.py | 2 + python/cudnn/_jax/api_base.py | 92 +- python/cudnn/_jax/layout.py | 44 +- .../cudnn/block_sparse_attention/__init__.py | 10 +- .../csrc/bwd/bucketed_k2q_csr.py | 13 +- .../csrc/bwd/sm100_blk128/bsa_bwd_sm100.py | 20 +- .../csrc/utils/cute_dsl_utils.py | 36 +- python/cudnn/block_sparse_attention/jax.py | 1428 +++++++++++++++++ python/cudnn/block_sparse_attention/op.py | 481 ++++++ .../indexer_backward/jax.py | 421 ++++- .../indexer_backward/op.py | 158 +- .../score_recompute/jax.py | 546 ++++++- .../score_recompute/op.py | 243 ++- .../cudnn/discrete_grouped_gemm/__init__.py | 39 +- .../discrete_grouped_gemm/_jax_common.py | 321 ++++ .../discrete_grouped_gemm_dswiglu/__init__.py | 23 +- ...e_B_blockscaled_grouped_gemm_dglu_dbias.py | 34 +- .../discrete_grouped_gemm_dswiglu/jax.py | 601 +++++++ .../discrete_grouped_gemm_swiglu/__init__.py | 23 +- ...ete_B_blockscaled_grouped_gemm_glu_bias.py | 38 +- .../discrete_grouped_gemm_swiglu/jax.py | 574 +++++++ .../discrete_kernel_utils.py | 17 +- python/cudnn/grouped_gemm/__init__.py | 92 +- python/cudnn/grouped_gemm/_jax_api.py | 721 +++++++++ .../grouped_gemm_dglu/__init__.py | 16 +- .../grouped_gemm/grouped_gemm_dglu/jax.py | 752 +++++++++ .../grouped_gemm_dsrelu/__init__.py | 16 +- .../grouped_gemm/grouped_gemm_dsrelu/jax.py | 670 ++++++++ .../grouped_gemm_dswiglu/__init__.py | 16 +- .../grouped_gemm/grouped_gemm_dswiglu/jax.py | 625 ++++++++ .../grouped_gemm/grouped_gemm_glu/__init__.py | 16 +- .../grouped_gemm/grouped_gemm_glu/jax.py | 687 ++++++++ .../grouped_gemm_glu_hadamard/__init__.py | 19 +- .../hadamard_utils.py | 3 +- .../grouped_gemm_glu_hadamard/jax.py | 648 ++++++++ .../grouped_gemm_quant/__init__.py | 22 +- .../grouped_gemm/grouped_gemm_quant/jax.py | 626 ++++++++ .../grouped_gemm_srelu/__init__.py | 21 +- .../grouped_gemm/grouped_gemm_srelu/jax.py | 627 ++++++++ .../grouped_gemm_swiglu/__init__.py | 21 +- .../grouped_gemm/grouped_gemm_swiglu/jax.py | 564 +++++++ .../grouped_gemm_wgrad/__init__.py | 16 +- .../grouped_gemm/grouped_gemm_wgrad/jax.py | 577 +++++++ .../cudnn/grouped_gemm/moe_kernel_helpers.py | 13 +- python/cudnn/grouped_gemm/utils.py | 11 +- python/cudnn/jax/__init__.py | 256 ++- .../cudnn/native_sparse_attention/__init__.py | 47 +- .../compression/__init__.py | 13 +- .../compression/fmha.py | 3 + .../compression/fmha_helpers.py | 16 +- .../compression/jax.py | 559 +++++++ .../native_sparse_attention/jax_utils.py | 207 +++ .../selection/__init__.py | 13 +- .../native_sparse_attention/selection/jax.py | 343 ++++ .../sliding_window_attention/__init__.py | 13 +- .../sliding_window_attention/jax.py | 369 +++++ .../native_sparse_attention/top_k/__init__.py | 13 +- .../native_sparse_attention/top_k/jax.py | 630 ++++++++ python/cudnn/sdpa/__init__.py | 18 +- python/cudnn/sdpa/bwd/__init__.py | 13 +- python/cudnn/sdpa/bwd/jax.py | 723 +++++++++ python/cudnn/sdpa/fwd/__init__.py | 13 +- python/cudnn/sdpa/fwd/jax.py | 548 +++++++ python/cudnn/sdpa/jax_utils.py | 187 +++ .../fe_api/test_dsa_score_recompute_op.py | 47 +- test/python/fe_api/test_jax_api_base.py | 48 +- test/python/fe_api/test_jax_bsa_contract.py | 389 +++++ ...test_jax_discrete_grouped_gemm_contract.py | 458 ++++++ .../test_jax_dsa_indexer_backward_contract.py | 246 ++- .../test_jax_dsa_score_recompute_contract.py | 250 ++- .../fe_api/test_jax_grouped_gemm_contract.py | 567 +++++++ ...jax_grouped_wgrad_accumulation_contract.py | 84 + .../python/fe_api/test_jax_import_contract.py | 141 +- test/python/fe_api/test_jax_layout.py | 31 +- .../test_jax_nsa_sdpa_source_contract.py | 355 ++++ .../fe_api/test_jax_nsa_topk_thd_contract.py | 400 +++++ .../fe_api/test_jax_sdpa_thd_contract.py | 184 +++ .../test_jax_sliding_packed_contract.py | 64 + 80 files changed, 18564 insertions(+), 681 deletions(-) create mode 100644 python/cudnn/block_sparse_attention/jax.py create mode 100644 python/cudnn/block_sparse_attention/op.py create mode 100644 python/cudnn/discrete_grouped_gemm/_jax_common.py create mode 100644 python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/jax.py create mode 100644 python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/jax.py create mode 100644 python/cudnn/grouped_gemm/_jax_api.py create mode 100644 python/cudnn/grouped_gemm/grouped_gemm_dglu/jax.py create mode 100644 python/cudnn/grouped_gemm/grouped_gemm_dsrelu/jax.py create mode 100644 python/cudnn/grouped_gemm/grouped_gemm_dswiglu/jax.py create mode 100644 python/cudnn/grouped_gemm/grouped_gemm_glu/jax.py create mode 100644 python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/jax.py create mode 100644 python/cudnn/grouped_gemm/grouped_gemm_quant/jax.py create mode 100644 python/cudnn/grouped_gemm/grouped_gemm_srelu/jax.py create mode 100644 python/cudnn/grouped_gemm/grouped_gemm_swiglu/jax.py create mode 100644 python/cudnn/grouped_gemm/grouped_gemm_wgrad/jax.py create mode 100644 python/cudnn/native_sparse_attention/compression/jax.py create mode 100644 python/cudnn/native_sparse_attention/jax_utils.py create mode 100644 python/cudnn/native_sparse_attention/selection/jax.py create mode 100644 python/cudnn/native_sparse_attention/sliding_window_attention/jax.py create mode 100644 python/cudnn/native_sparse_attention/top_k/jax.py create mode 100644 python/cudnn/sdpa/bwd/jax.py create mode 100644 python/cudnn/sdpa/fwd/jax.py create mode 100644 python/cudnn/sdpa/jax_utils.py create mode 100644 test/python/fe_api/test_jax_bsa_contract.py create mode 100644 test/python/fe_api/test_jax_discrete_grouped_gemm_contract.py create mode 100644 test/python/fe_api/test_jax_grouped_gemm_contract.py create mode 100644 test/python/fe_api/test_jax_grouped_wgrad_accumulation_contract.py create mode 100644 test/python/fe_api/test_jax_nsa_sdpa_source_contract.py create mode 100644 test/python/fe_api/test_jax_nsa_topk_thd_contract.py create mode 100644 test/python/fe_api/test_jax_sdpa_thd_contract.py create mode 100644 test/python/fe_api/test_jax_sliding_packed_contract.py diff --git a/docs/fe-oss-apis/dsa.md b/docs/fe-oss-apis/dsa.md index 820a20805..7e543d52a 100644 --- a/docs/fe-oss-apis/dsa.md +++ b/docs/fe-oss-apis/dsa.md @@ -151,8 +151,10 @@ scores_sbk = result["scores"] ``` The defaults are `q_layout="BSHD"`, `k_layout="BSHD"`, -`w_layout="BSH"`, and `output_layout="BSK"`. Packed inputs use `THD` and -`TH`, with a `TK` output. The adapter maps these public axes to the kernel's +`w_layout="BSH"`, and `output_layout="BSK"`. Backward and score-recompute +adapters use the same convention and additionally expose layout selectors for +their score, denominator, and gradient tensors. Packed inputs use `THD` and +`TH`, with a `TK` output. The adapters map these public axes to the kernel's canonical axes; callers do not need to transpose arrays into canonical order. Use the class API when the caller needs to choose the enclosing JAX @@ -198,11 +200,11 @@ and SM107. | Sparse Attention Backward | Flat MQA tensors | Yes | Yes; inputs must be BF16 | | Indexer Forward | Fixed BSHD/SBHD or packed THD | Yes; `H_kv=1` and default tuning | Yes | | Indexer Top-K and index utilities | Flattened rows; fixed or packed index conversion | Yes | Yes | -| Sparse score recompute | Fixed BSHD/MQA | Yes | Yes | -| Dense score recompute | Fixed BSHD/MQA | Yes | Yes | +| Sparse score recompute | Fixed BSHD/SBHD outer layouts, MQA | Yes | Yes | +| Dense score recompute | Fixed BSHD/SBHD outer layouts, MQA | Yes | Yes | | Dense score recompute | Packed THD/MQA | No | Yes | -| Indexer Backward | Fixed BSHD; JAX-only packed THD with global indices | Yes | Yes | -| Dense Indexer Backward | Fixed BSHD or packed THD | Yes | Yes | +| Indexer Backward | Fixed BSHD/SBHD; JAX-only packed THD with global indices | Yes | Yes | +| Dense Indexer Backward | Fixed BSHD/SBHD or packed THD | Yes | Yes | SM90 packed THD dense score recompute is rejected because that backend needs host-side cumulative-length reads, which cannot be represented during JAX diff --git a/docs/fe-oss-apis/nsa.md b/docs/fe-oss-apis/nsa.md index 10a2458a6..559c52318 100644 --- a/docs/fe-oss-apis/nsa.md +++ b/docs/fe-oss-apis/nsa.md @@ -74,6 +74,46 @@ NSA.TopKReduction NSA.topk_reduction_wrapper ``` +### JAX namespace + +Install the optional integration with `nvidia-cudnn-frontend[jax]`, then use +the lazy `cudnn.jax.NSA` namespace (the symbols are also available directly +from `cudnn.jax`): + +```python +import cudnn.jax as cudnn_jax + +result = cudnn_jax.NSA.compression_attention_wrapper( + q, + k, + v, + layout="BSHD", + enable_lse=True, +) +output = result["o_tensor"] +lse = result["lse_tensor"] +``` + +Functional JAX wrappers are decorated with `jax.jit`. Their layout, tuning, +maximum-sequence-length, and target arguments are static; cumulative sequence +lengths remain dynamic array operands. Classes accept array-like exemplars and +can be placed under a caller-owned JAX transformation. + +| Component | JAX public layouts | +| --- | --- | +| Selection | Packed `THD` | +| Compression | Fixed `BHSD`/`BSHD` or packed `THD` | +| Sliding Window | Fixed `BHSD`/`BSHD` or fully packed `THD` | +| Top-K | Fixed `BHSD`/`BSHD` or packed `THD` | + +Packed calls require cumulative Q/K sequence lengths and static `max_s_q` and +`max_s_k` bounds. The arrays must start at zero, be monotonic, end at the +corresponding packed token count, and respect those bounds. Sliding Window +uses those offsets to form padded cuDNN attention operands in the lowered +computation; it does not accept arbitrary ragged byte-offset tensors. Fixed +data outputs follow the selected layout, while LSE/statistics remain in +`(B, H, S[, 1])` order. + --- ## Components diff --git a/python/cudnn/_cute_compiler.py b/python/cudnn/_cute_compiler.py index 7e6c8cb21..23549f2d6 100644 --- a/python/cudnn/_cute_compiler.py +++ b/python/cudnn/_cute_compiler.py @@ -10,6 +10,8 @@ 100: "sm_100a", 103: "sm_103a", 107: "sm_100f", + 110: "sm_110a", + 120: "sm_120a", } diff --git a/python/cudnn/_jax/api_base.py b/python/cudnn/_jax/api_base.py index fa468d14d..6c4fa13ba 100644 --- a/python/cudnn/_jax/api_base.py +++ b/python/cudnn/_jax/api_base.py @@ -236,6 +236,19 @@ def _get_max_active_clusters( raise ValueError("max_active_clusters must be positive after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") return max_active_clusters + def _get_device_multiprocessor_count(self) -> int: + """Return a cached local SM count before entering custom-call lowering.""" + + count = getattr(self, "_device_multiprocessor_count", None) + if count is None: + import cutlass + + count = int(cutlass.utils.HardwareInfo().get_device_multiprocessor_count()) + if count <= 0: + raise RuntimeError(f"CUDA reported an invalid multiprocessor count {count}") + self._device_multiprocessor_count = count + return count + @staticmethod def _to_tensor_desc( value: Any, @@ -357,6 +370,7 @@ def _call_kernel( *, launch: Callable[..., None], output_descs: tuple[TensorDesc[Any], ...], + output_seeds: tuple[Any | None, ...] | None = None, workspace_descs: tuple[TensorDesc[Any], ...] = (), input_spec: tuple[Any | None, ...] | None = None, output_spec: tuple[Any | None, ...] | None = None, @@ -372,8 +386,11 @@ def _call_kernel( are declared as custom-call results so XLA owns their lifetime, then omitted from this method's return value. Descriptors with a non-``None`` ``init_value`` are materialized as JAX inputs and aliased to their - corresponding custom-call results. Inputs are a flat tuple of arrays, - and all explicit specs are native CUTLASS ``TensorSpec`` values. + corresponding custom-call results. ``output_seeds`` supplies existing + JAX values for outputs that must be updated in place by the kernel; + these operands are likewise aliased to the corresponding results. + Inputs are a flat tuple of arrays, and all explicit specs are native + CUTLASS ``TensorSpec`` values. """ import cutlass.jax as cutlass_jax @@ -383,8 +400,8 @@ def _call_kernel( raise TypeError(f"launch must be callable, got {type(launch).__name__}") inputs = tuple(inputs) - for index, value in enumerate(inputs): - _require_array_metadata(value, f"input #{index}") + for input_index, value in enumerate(inputs): + _require_array_metadata(value, f"input #{input_index}") outputs = tuple(output_descs) workspaces = tuple(workspace_descs) if not outputs: @@ -425,18 +442,57 @@ def resolve_specs( buffer_specs = output_specs + workspace_specs buffer_metadata = tuple(self._materialize_tensor_desc(desc, mode=getattr(spec, "mode", None)) for desc, spec in zip(buffers, buffer_specs)) - initialized = tuple(index for index, desc in enumerate(buffers) if desc.init_value is not None) - uninitialized = tuple(index for index, desc in enumerate(buffers) if desc.init_value is None) - seed_inputs = tuple( - jnp.full( - buffer_metadata[index].shape, - buffers[index].init_value, - dtype=buffer_metadata[index].dtype, - ) - for index in initialized + if output_seeds is None: + supplied_output_seeds = (None,) * len(outputs) + else: + supplied_output_seeds = tuple(output_seeds) + if len(supplied_output_seeds) != len(outputs): + raise ValueError( + f"Expected {len(outputs)} output seeds, got {len(supplied_output_seeds)}" + ) + + seed_by_buffer: dict[int, Any] = {} + for output_index, seed in enumerate(supplied_output_seeds): + if seed is None: + continue + output_name = outputs[output_index].name or f"output #{output_index}" + if outputs[output_index].init_value is not None: + raise ValueError( + f"{output_name} cannot have both an explicit seed and init_value" + ) + _require_array_metadata(seed, f"output seed #{output_index}") + metadata = buffer_metadata[output_index] + if tuple(seed.shape) != tuple(metadata.shape): + raise ValueError( + f"{output_name} seed shape mismatch: " + f"expected {tuple(metadata.shape)}, got {tuple(seed.shape)}" + ) + if seed.dtype != metadata.dtype: + raise ValueError( + f"{output_name} seed dtype mismatch: " + f"expected {metadata.dtype}, got {seed.dtype}" + ) + seed_by_buffer[output_index] = seed + + for buffer_index, desc in enumerate(buffers): + if desc.init_value is not None: + seed_by_buffer[buffer_index] = jnp.full( + buffer_metadata[buffer_index].shape, + desc.init_value, + dtype=buffer_metadata[buffer_index].dtype, + ) + + initialized = tuple(sorted(seed_by_buffer)) + uninitialized = tuple( + buffer_index + for buffer_index in range(len(buffers)) + if buffer_index not in seed_by_buffer ) + seed_inputs = tuple(seed_by_buffer[buffer_index] for buffer_index in initialized) aliases = {len(inputs) + seed_index: result_index for seed_index, result_index in enumerate(initialized)} - call_input_specs = input_specs + tuple(buffer_specs[index] for index in initialized) + call_input_specs = input_specs + tuple( + buffer_specs[buffer_index] for buffer_index in initialized + ) input_count = len(inputs) seed_count = len(seed_inputs) @@ -449,10 +505,10 @@ def launcher(stream, *args): raise RuntimeError(f"Kernel received {len(allocated_buffers)} allocated buffers; expected {len(uninitialized)}") ordered_buffers: list[Any] = [None] * len(buffers) - for index, value in zip(initialized, seeded_buffers): - ordered_buffers[index] = value - for index, value in zip(uninitialized, allocated_buffers): - ordered_buffers[index] = value + for buffer_index, value in zip(initialized, seeded_buffers): + ordered_buffers[buffer_index] = value + for buffer_index, value in zip(uninitialized, allocated_buffers): + ordered_buffers[buffer_index] = value output_buffers = tuple(ordered_buffers[: len(outputs)]) workspace_buffers = tuple(ordered_buffers[len(outputs) :]) launch(stream, *kernel_inputs, *output_buffers, *workspace_buffers) diff --git a/python/cudnn/_jax/layout.py b/python/cudnn/_jax/layout.py index 14dd0b7e9..780729316 100644 --- a/python/cudnn/_jax/layout.py +++ b/python/cudnn/_jax/layout.py @@ -23,13 +23,17 @@ def mode_from_layout(layout: str, *, kernel_axes: str) -> tuple[int, ...]: if not isinstance(kernel_axes, str): raise TypeError(f"kernel_axes must be a string, got {kernel_axes!r}") if len(layout) != len(kernel_axes): - raise ValueError(f"layout rank must match kernel_axes: got {layout!r} and {kernel_axes!r}") + raise ValueError( + f"layout rank must match kernel_axes: got {layout!r} and {kernel_axes!r}" + ) if len(set(layout)) != len(layout): raise ValueError(f"layout axes must be unique, got {layout!r}") if len(set(kernel_axes)) != len(kernel_axes): raise ValueError(f"kernel_axes must be unique, got {kernel_axes!r}") if set(layout) != set(kernel_axes): - raise ValueError(f"layout must contain exactly the axes in {kernel_axes!r}, got {layout!r}") + raise ValueError( + f"layout must contain exactly the axes in {kernel_axes!r}, got {layout!r}" + ) return tuple(layout.index(axis) for axis in kernel_axes) @@ -62,7 +66,9 @@ def normalize_mode(rank: int, mode: tuple[int, ...] | None = None) -> tuple[int, raise TypeError(f"mode entries must be integers, got {axis!r}") from error normalized_mode = tuple(normalized) if tuple(sorted(normalized_mode)) != tuple(range(rank)): - raise ValueError(f"mode must be a permutation of [0, {rank - 1}], got {normalized_mode}") + raise ValueError( + f"mode must be a permutation of [0, {rank - 1}], got {normalized_mode}" + ) return normalized_mode @@ -88,10 +94,15 @@ def to_public_axes( canonical_axis_by_public_axis = [0] * len(mode) for canonical_axis, public_axis in enumerate(mode): canonical_axis_by_public_axis[public_axis] = canonical_axis - return tuple(canonical_values[canonical_axis_by_public_axis[public_axis]] for public_axis in range(len(mode))) + return tuple( + canonical_values[canonical_axis_by_public_axis[public_axis]] + for public_axis in range(len(mode)) + ) -def compact_stride(shape: tuple[int, ...], stride_order: tuple[int, ...]) -> tuple[int, ...]: +def compact_stride( + shape: tuple[int, ...], stride_order: tuple[int, ...] +) -> tuple[int, ...]: """Return compact strides for dimensions ordered fastest to slowest.""" stride = [0] * len(shape) @@ -102,6 +113,24 @@ def compact_stride(shape: tuple[int, ...], stride_order: tuple[int, ...]) -> tup return tuple(stride) +def stride_order_to_public( + canonical_stride_order: tuple[int, ...], + mode: tuple[int, ...] | None = None, +) -> tuple[int, ...]: + """Map a fastest-to-slowest canonical axis order to public array axes.""" + + canonical_stride_order = tuple(canonical_stride_order) + mode = normalize_mode(len(canonical_stride_order), mode) + if tuple(sorted(canonical_stride_order)) != tuple( + range(len(canonical_stride_order)) + ): + raise ValueError( + "canonical_stride_order must be a permutation of " + f"[0, {len(canonical_stride_order) - 1}], got {canonical_stride_order}" + ) + return tuple(mode[canonical_axis] for canonical_axis in canonical_stride_order) + + def to_cutlass_layout( shape: tuple[int, ...], stride: tuple[int, ...], @@ -117,7 +146,9 @@ def to_cutlass_layout( """ if stride != compact_stride(shape, stride_order): - raise ValueError(f"JAX TensorSpec cannot represent non-compact stride {stride} for {name}") + raise ValueError( + f"JAX TensorSpec cannot represent non-compact stride {stride} for {name}" + ) layout = [0] * len(stride_order) for rank, dimension in enumerate(stride_order): @@ -129,6 +160,7 @@ def to_cutlass_layout( "compact_stride", "mode_from_layout", "normalize_mode", + "stride_order_to_public", "to_canonical_axes", "to_cutlass_layout", "to_public_axes", diff --git a/python/cudnn/block_sparse_attention/__init__.py b/python/cudnn/block_sparse_attention/__init__.py index 172f703e8..ab3f43dd8 100644 --- a/python/cudnn/block_sparse_attention/__init__.py +++ b/python/cudnn/block_sparse_attention/__init__.py @@ -6,8 +6,14 @@ _SYMBOLS = { "block_sparse_attention_forward": (".api", "block_sparse_attention_forward"), "block_sparse_attention_backward": (".api", "block_sparse_attention_backward"), + "BlockSparseAttentionForwardOp": (".op", "BlockSparseAttentionForwardOp"), + "BlockSparseAttentionBackwardOp": (".op", "BlockSparseAttentionBackwardOp"), } +_BSA_SYMBOLS = frozenset( + {"block_sparse_attention_forward", "block_sparse_attention_backward"} +) + def _load_symbol(name): module_name, symbol_name = _SYMBOLS[name] @@ -27,11 +33,11 @@ def __getattr__(name): class BSANamespace: def __getattr__(self, name): - if name in _SYMBOLS: + if name in _BSA_SYMBOLS: return _load_symbol(name) raise AttributeError(f"BSA has no attribute {name!r}") BSA = BSANamespace() -__all__ = ["BSA", *_SYMBOLS.keys()] +__all__ = ["BSA", *_SYMBOLS] diff --git a/python/cudnn/block_sparse_attention/csrc/bwd/bucketed_k2q_csr.py b/python/cudnn/block_sparse_attention/csrc/bwd/bucketed_k2q_csr.py index e11ed68df..5df63c40d 100644 --- a/python/cudnn/block_sparse_attention/csrc/bwd/bucketed_k2q_csr.py +++ b/python/cudnn/block_sparse_attention/csrc/bwd/bucketed_k2q_csr.py @@ -1,11 +1,14 @@ -from typing import Optional, Tuple +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional, Tuple + +if TYPE_CHECKING: + import torch import cutlass import cutlass.cute as cute from cutlass import Int32, const_expr -from cutlass.cute.runtime import from_dlpack import cuda.bindings.driver as cuda -import torch class BucketedK2QCsrUniversal: @@ -240,6 +243,8 @@ def _bucketed_k2q_csr_compile_key( def _to_cute_tensor(tensor: torch.Tensor) -> cute.Tensor: + from cutlass.cute.runtime import from_dlpack + return from_dlpack( tensor.detach(), assumed_align=4, @@ -256,6 +261,8 @@ def build_bucketed_k2q_csr_cutedsl( q2k_block_nums: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, int, int]: """Build bucketed K-to-Q CSR metadata with CuTe DSL kernels.""" + import torch + assert q2k_block_index.dtype == torch.int32 assert q2k_block_index.is_cuda assert q2k_block_index.ndim == 4 diff --git a/python/cudnn/block_sparse_attention/csrc/bwd/sm100_blk128/bsa_bwd_sm100.py b/python/cudnn/block_sparse_attention/csrc/bwd/sm100_blk128/bsa_bwd_sm100.py index 899ef48cb..dd3756a7a 100644 --- a/python/cudnn/block_sparse_attention/csrc/bwd/sm100_blk128/bsa_bwd_sm100.py +++ b/python/cudnn/block_sparse_attention/csrc/bwd/sm100_blk128/bsa_bwd_sm100.py @@ -1,11 +1,14 @@ # Copyright (c) 2025, Ted Zadouri, Markus Hoehnerbach, Jay Shah, Tri Dao. +from __future__ import annotations + import math -from typing import Callable, NamedTuple, Optional, Tuple +from typing import TYPE_CHECKING, Callable, NamedTuple, Optional, Tuple from functools import partial import cuda.bindings.driver as cuda -import torch +if TYPE_CHECKING: + import torch import cutlass import cutlass.cute as cute @@ -28,11 +31,6 @@ from cudnn.block_sparse_attention.csrc.utils.tcgen05_mma_helpers import gemm_w_idx, gemm_ptx_w_idx from cudnn.block_sparse_attention.csrc.utils.seqlen_info import SeqlenInfoQK from cudnn.block_sparse_attention.csrc.utils.block_info import BlockInfo -from cudnn.block_sparse_attention.csrc.bwd.bsa_bwd_prepost import ( - _bwd_postprocess_convert, - _bwd_preprocess, - _get_device_arch, -) from cudnn.block_sparse_attention.csrc.utils.cute_dsl_utils import ParamsBase, sub_packed_f32x2 from cudnn.block_sparse_attention.csrc.utils.tile_scheduler import ( TileSchedulerArguments, @@ -2193,6 +2191,14 @@ def bsa_sm100_blk128_bwd_bucketed_k2q_csr( dv: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """SM100/SM110 blk128 bwd entry receiving BSA bucketed k2q CSR.""" + import torch + + from cudnn.block_sparse_attention.csrc.bwd.bsa_bwd_prepost import ( + _bwd_postprocess_convert, + _bwd_preprocess, + _get_device_arch, + ) + assert q.dtype == torch.bfloat16, "SM100 blk128 bwd only supports bfloat16" assert q.dtype == k.dtype == v.dtype == out.dtype == dout.dtype assert lse.dtype == torch.float32 diff --git a/python/cudnn/block_sparse_attention/csrc/utils/cute_dsl_utils.py b/python/cudnn/block_sparse_attention/csrc/utils/cute_dsl_utils.py index a22e7c7a7..736115a6d 100644 --- a/python/cudnn/block_sparse_attention/csrc/utils/cute_dsl_utils.py +++ b/python/cudnn/block_sparse_attention/csrc/utils/cute_dsl_utils.py @@ -5,18 +5,25 @@ # this file are adapted from quack-kernels 0.4.1 (Apache-2.0) and modified for # cudnn-frontend's CUTLASS DSL 4.5 integration. +from __future__ import annotations + from dataclasses import dataclass, fields from functools import partial from typing import Tuple, get_origin -import torch - import cutlass import cutlass.cute as cute from cutlass._mlir.dialects import nvvm from cutlass.base_dsl.tvm_ffi_builder import spec from cutlass.cutlass_dsl import NumericMeta -from cutlass.cute.runtime import from_dlpack + +try: + import torch +except ModuleNotFoundError: + # Kernel-only users, including cudnn.jax, must not require PyTorch. The + # Torch conversion helpers below fail explicitly if they are invoked + # without the optional dependency. + torch = None # Python scalars and CuTe numeric types are compile-time values. Everything # else in a ParamsBase dataclass is flattened into MLIR values at JIT time. @@ -51,13 +58,17 @@ def convert_single_arg(arg, arg_name, arg_type, ctx): _install_constexpr_tvm_ffi_converter() -torch2cute_dtype_map = { - torch.float16: cutlass.Float16, - torch.bfloat16: cutlass.BFloat16, - torch.float32: cutlass.Float32, - torch.float8_e4m3fn: cutlass.Float8E4M3FN, - torch.float8_e5m2: cutlass.Float8E5M2, -} +torch2cute_dtype_map = ( + {} + if torch is None + else { + torch.float16: cutlass.Float16, + torch.bfloat16: cutlass.BFloat16, + torch.float32: cutlass.Float32, + torch.float8_e4m3fn: cutlass.Float8E4M3FN, + torch.float8_e5m2: cutlass.Float8E5M2, + } +) def _partition_param_fields(obj): @@ -139,6 +150,11 @@ def assume_tensor_aligned(t): def to_cute_tensor(t, assumed_align=16, leading_dim=-1, fully_dynamic=False, enable_tvm_ffi=True): """Convert torch tensor to cute tensor for TVM FFI. leading_dim=-1 defaults to t.ndim-1.""" + if torch is None: + raise ImportError("to_cute_tensor requires PyTorch") + + from cutlass.cute.runtime import from_dlpack + # NOTE: torch 2.9.1 doesn't support fp8 via DLPack but 2.11.0 nightly does # currently export raw bytes as uint8 and tell cutlass correct type # can directly export as fp8 when torch supports it diff --git a/python/cudnn/block_sparse_attention/jax.py b/python/cudnn/block_sparse_attention/jax.py new file mode 100644 index 000000000..4d4a5c53f --- /dev/null +++ b/python/cudnn/block_sparse_attention/jax.py @@ -0,0 +1,1428 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX adapters for the CuTe DSL block-sparse attention kernels.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +import jax +import jax.numpy as jnp + +from .. import data_type +from .._cute_compiler import compile_options_for_target +from .._jax import JaxApiBase, JaxTensorDesc, TupleDict +from .._jax.layout import mode_from_layout, to_public_axes +from .op import ( + BlockSparseAttentionBackwardOp, + BlockSparseAttentionForwardOp, + SUPPORTED_BACKWARD_COMPUTE_CAPABILITIES, + SUPPORTED_FORWARD_COMPUTE_CAPABILITIES, + ceil_div, + compute_capability_family, +) + + +_DATA_LAYOUTS = ("bhsd", "bshd") + + +def _data_mode(layout: str) -> tuple[int, ...]: + if not isinstance(layout, str): + raise TypeError(f"layout must be a string, got {type(layout).__name__}") + normalized = "".join( + character for character in layout.lower() if character.isalpha() + ) + if normalized not in _DATA_LAYOUTS: + raise ValueError(f"layout must be one of {_DATA_LAYOUTS}, got {layout!r}") + return mode_from_layout(normalized.upper(), kernel_axes="BHSD") + + +def _shape_sample( + shape: tuple[int, ...], dtype: Any, mode: tuple[int, ...] | None = None +) -> Any: + return jax.ShapeDtypeStruct(to_public_axes(shape, mode), dtype) + + +def _optional_signature( + value: Any | None, desc: JaxTensorDesc | None, name: str +) -> None: + if (value is None) != (desc is None): + expected = "omitted" if desc is None else "provided" + raise ValueError(f"{name} must be {expected} for this specialized callable") + if desc is not None: + JaxApiBase._check_tensor_signature(value, desc) + + +def _default_forward_sparse_block_size(target: int) -> int: + return 64 if compute_capability_family(target) in {90, 120} else 128 + + +def _default_backward_sparse_block_size(target: int) -> int: + return 64 if compute_capability_family(target) == 90 else 128 + + +def _default_bucket_size( + target: int, sparse_block_size: int, num_q_blocks: int, num_heads: int +) -> int: + if compute_capability_family(target) == 90: + return 384 + if sparse_block_size == 64: + if num_q_blocks >= 3000: + return 1024 + return 1088 if num_q_blocks < 2048 or num_q_blocks >= 8192 else 1152 + if num_q_blocks >= 4096 and num_heads <= 1: + return 256 + if num_q_blocks >= 2048: + return 512 + return 384 + + +def _round_up(value: int, alignment: int) -> int: + return ceil_div(value, alignment) * alignment + + +def _sm100_blk64_use_clc( + *, + batch: int, + heads: int, + seqlen_q: int, + block_sparse_num: int, + has_variable_block_nums: bool, +) -> bool: + if has_variable_block_nums: + return True + num_m_blocks = ceil_div(seqlen_q, 64) + if num_m_blocks >= 8192 and block_sparse_num >= 512: + return True + if heads == 1: + return False + return ( + num_m_blocks >= 128 + and batch * heads * num_m_blocks >= 512 + and block_sparse_num <= (64 if heads == 2 else 128) + ) + + +def _resolve_sm100_blk64_use_clc( + *, + kv_splits: int, + requested: bool | None, + batch: int, + heads: int, + seqlen_q: int, + block_sparse_num: int, + has_variable_block_nums: bool, +) -> bool: + if kv_splits > 1: + return False + if requested is not None: + return requested + return _sm100_blk64_use_clc( + batch=batch, + heads=heads, + seqlen_q=seqlen_q, + block_sparse_num=block_sparse_num, + has_variable_block_nums=has_variable_block_nums, + ) + + +def _sm100_blk64_auto_kv_splits(block_sparse_num: int) -> int: + if block_sparse_num >= 900: + return 8 + if block_sparse_num >= 450: + return 4 + if block_sparse_num >= 256: + return 2 + return 1 + + +def _sm100_blk64_uses_int64_kv_strides(desc: JaxTensorDesc) -> bool: + coord_stride_limit = 1 << 27 + batch, heads, seqlen_k, _ = desc.shape + stride_b, stride_h, stride_s, stride_d = desc.stride + rank6_shape = (64, 64, 2, heads, ceil_div(seqlen_k, 64), batch) + rank6_stride = ( + stride_s, + stride_d, + 64 * stride_d, + stride_h, + 64 * stride_s, + stride_b, + ) + if any(stride < 0 or stride > jnp.iinfo(jnp.int32).max for stride in rank6_stride): + return True + return (rank6_shape[4] > 1 and rank6_stride[4] >= coord_stride_limit) or ( + rank6_shape[5] > 1 and rank6_stride[5] >= coord_stride_limit + ) + + +class BlockSparseAttentionForward(JaxApiBase): + """JAX callable specialized from block-sparse attention metadata.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_q2k_block_index: Any, + *, + sample_block_sizes: Any | None = None, + sample_q2k_block_nums: Any | None = None, + sample_o: Any | None = None, + sample_lse: Any | None = None, + block_sparse_num: int | None = None, + sparse_block_size: int | None = None, + allow_empty_block_nums: bool = False, + softmax_scale: float | None = None, + pack_gqa: bool | None = None, + layout: str = "bhsd", + kv_splits: int | str = 1, + use_clc: bool | None = None, + target_compute_capability: int | None = None, + ) -> None: + self.layout = layout + self.data_mode = _data_mode(layout) + self.compute_capability = self._resolve_compute_capability( + target_compute_capability, + SUPPORTED_FORWARD_COMPUTE_CAPABILITIES, + "BlockSparseAttentionForward", + ) + self.compute_capability_family = compute_capability_family( + self.compute_capability + ) + self.q_desc = self._to_tensor_desc(sample_q, "sample_q", mode=self.data_mode) + self.k_desc = self._to_tensor_desc(sample_k, "sample_k", mode=self.data_mode) + self.v_desc = self._to_tensor_desc(sample_v, "sample_v", mode=self.data_mode) + self.block_index_desc = self._to_tensor_desc( + sample_q2k_block_index, "sample_q2k_block_index" + ) + self.block_sizes_desc = ( + None + if sample_block_sizes is None + else self._to_tensor_desc(sample_block_sizes, "sample_block_sizes") + ) + self.block_nums_desc = ( + None + if sample_q2k_block_nums is None + else self._to_tensor_desc(sample_q2k_block_nums, "sample_q2k_block_nums") + ) + + if self.q_desc.ndim != 4 or self.v_desc.ndim != 4: + raise ValueError("sample_q and sample_v must be rank 4") + batch, num_q_heads, seqlen_q, head_dim = self.q_desc.shape + value_dim = self.v_desc.shape[3] + canonical_output_shape = (batch, num_q_heads, seqlen_q, value_dim) + if sample_o is None: + self.o_desc = self._to_tensor_desc( + _shape_sample( + canonical_output_shape, self.q_desc.dtype, self.data_mode + ), + "sample_o", + mode=self.data_mode, + ) + else: + self.o_desc = self._to_tensor_desc( + sample_o, "sample_o", mode=self.data_mode + ) + if sample_lse is None: + self.lse_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((batch, num_q_heads, seqlen_q), jnp.float32), + "sample_lse", + ) + else: + self.lse_desc = self._to_tensor_desc(sample_lse, "sample_lse") + + self.sparse_block_size = ( + _default_forward_sparse_block_size(self.compute_capability) + if sparse_block_size is None + else sparse_block_size + ) + capacity = ( + self.block_index_desc.shape[3] if self.block_index_desc.ndim == 4 else 0 + ) + self.block_sparse_num = ( + capacity if block_sparse_num is None else block_sparse_num + ) + self.softmax_scale = ( + head_dim**-0.5 if softmax_scale is None else float(softmax_scale) + ) + if isinstance(kv_splits, str): + if kv_splits != "auto": + raise ValueError("kv_splits string value must be 'auto'") + if not ( + self.compute_capability_family == 100 and self.sparse_block_size == 64 + ): + raise ValueError( + "kv_splits='auto' is available only on the SM100-family blk64 path" + ) + auto_block_count = ( + self.block_sparse_num if self.block_sparse_num > 0 else capacity + ) + self.kv_splits = _sm100_blk64_auto_kv_splits(auto_block_count) + else: + self.kv_splits = kv_splits + self.use_clc = use_clc + self._op = BlockSparseAttentionForwardOp( + q=self.q_desc, + k=self.k_desc, + v=self.v_desc, + block_index=self.block_index_desc, + output=self.o_desc, + lse=self.lse_desc, + block_sizes=self.block_sizes_desc, + block_nums=self.block_nums_desc, + block_sparse_num=self.block_sparse_num, + sparse_block_size=self.sparse_block_size, + softmax_scale=self.softmax_scale, + pack_gqa=pack_gqa, + allow_empty_block_nums=allow_empty_block_nums, + kv_splits=kv_splits, + use_clc=use_clc, + target_compute_capability=self.compute_capability, + ) + self.check_support() + + def check_support(self) -> bool: + return self._op.check_support() + + def __call__( + self, + q: Any, + k: Any, + v: Any, + q2k_block_index: Any, + block_sizes: Any | None = None, + q2k_block_nums: Any | None = None, + ) -> TupleDict: + self.check_support() + self._check_tensor_signature(q, self.q_desc, mode=self.data_mode) + self._check_tensor_signature(k, self.k_desc, mode=self.data_mode) + self._check_tensor_signature(v, self.v_desc, mode=self.data_mode) + self._check_tensor_signature(q2k_block_index, self.block_index_desc) + _optional_signature(block_sizes, self.block_sizes_desc, "block_sizes") + _optional_signature(q2k_block_nums, self.block_nums_desc, "q2k_block_nums") + + inputs: list[Any] = [q, k, v, q2k_block_index] + bindings: list[tuple[JaxTensorDesc, tuple[int, ...] | None]] = [ + (self.q_desc, self.data_mode), + (self.k_desc, self.data_mode), + (self.v_desc, self.data_mode), + (self.block_index_desc, None), + ] + if self.block_sizes_desc is not None: + inputs.append(block_sizes) + bindings.append((self.block_sizes_desc, None)) + + effective_block_nums = q2k_block_nums + effective_block_nums_desc = self.block_nums_desc + needs_materialized_block_nums = ( + self.compute_capability_family == 90 and effective_block_nums is None + ) + if needs_materialized_block_nums: + block_nums_shape = ( + self._op.batch, + self.block_index_desc.shape[1], + self._op.num_q_blocks, + ) + effective_block_nums = jnp.full( + block_nums_shape, self.block_sparse_num, dtype=jnp.int32 + ) + effective_block_nums_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct(block_nums_shape, jnp.int32), + "effective_q2k_block_nums", + ) + if effective_block_nums is not None: + inputs.append(effective_block_nums) + bindings.append((effective_block_nums_desc, None)) + + split_offsets_desc = None + if self.kv_splits > 1: + if effective_block_nums is None: + valid_kv = jnp.full( + ( + self._op.batch, + self.block_index_desc.shape[1], + self._op.num_q_blocks, + ), + self.block_sparse_num, + dtype=jnp.int32, + ) + else: + valid_kv = jnp.maximum(effective_block_nums.astype(jnp.int32), 0) + split_ids = jnp.arange(self.kv_splits + 1, dtype=jnp.int32) + quotient = valid_kv // self.kv_splits + remainder_even = valid_kv - quotient * self.kv_splits + even_offsets = ( + quotient[..., None] * split_ids + + (remainder_even[..., None] * split_ids + self.kv_splits - 1) + // self.kv_splits + ) + aligned_base = (quotient // 8) * 8 + aligned_remainder = valid_kv - aligned_base * self.kv_splits + aligned_offsets = aligned_base[..., None] * split_ids + jnp.minimum( + aligned_remainder[..., None], + split_ids * 8, + ) + split_offsets = jnp.where( + (aligned_base == 0)[..., None], + even_offsets, + jnp.minimum(aligned_offsets, valid_kv[..., None]), + ).astype(jnp.int32) + split_offsets_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct(split_offsets.shape, split_offsets.dtype), + "split_offsets", + ) + inputs.append(split_offsets) + bindings.append((split_offsets_desc, None)) + + self._forward_has_effective_block_nums = effective_block_nums is not None + self._forward_has_split_offsets = split_offsets_desc is not None + self._forward_input_count = len(inputs) + + workspace_descs: tuple[JaxTensorDesc, ...] = () + if self.kv_splits > 1: + partial_shape = ( + self._op.batch, + self.kv_splits * self._op.num_q_heads, + self._op.seqlen_q, + self._op.value_dim, + ) + partial_lse_shape = partial_shape[:3] + partial_o_desc = self.q_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=partial_shape, + name="partial_o_workspace", + ) + partial_lse_desc = self.lse_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=partial_lse_shape, + name="partial_lse_workspace", + ) + workspace_descs = (partial_o_desc, partial_lse_desc) + + o, lse = self._call_kernel( + tuple(inputs), + launch=self._launch, + output_descs=(self.o_desc, self.lse_desc), + workspace_descs=workspace_descs, + input_spec=tuple( + self._to_tensor_spec(desc, mode=mode) for desc, mode in bindings + ), + output_spec=( + self._to_tensor_spec(self.o_desc, mode=self.data_mode), + self._to_tensor_spec(self.lse_desc), + ), + compile_options=compile_options_for_target(self.compute_capability), + ) + return TupleDict(o_tensor=o, lse_tensor=lse) + + def _launch(self, stream: Any, *arguments: Any) -> None: + import cutlass + import cutlass.cute as cute + + inputs = list(arguments[: self._forward_input_count]) + buffers = list(arguments[self._forward_input_count :]) + q, k, v, block_index = inputs[:4] + cursor = 4 + block_sizes = None + if self.block_sizes_desc is not None: + block_sizes = inputs[cursor] + cursor += 1 + block_nums = None + if self._forward_has_effective_block_nums: + block_nums = inputs[cursor] + cursor += 1 + split_offsets = None + if self._forward_has_split_offsets: + split_offsets = inputs[cursor] + cursor += 1 + if cursor != len(inputs): + raise RuntimeError("Unexpected BSA forward kernel inputs") + + output, lse = buffers[:2] + if self.kv_splits > 1: + partial_output, partial_lse = buffers[2:] + kernel_output, kernel_lse = partial_output, partial_lse + else: + kernel_output, kernel_lse = output, lse + + if self.compute_capability_family == 90: + from .csrc.fwd.sm90_blk64.bsa_fwd_sm90 import ( + BlockSparseAttnForwardSm90Blk64, + ) + + def select(tensor: Any, modes: tuple[int, ...]) -> Any: + return cute.make_tensor( + tensor.iterator, cute.select(tensor.layout, mode=modes) + ) + + q_kernel = select(q, (2, 3, 1, 0)) + k_kernel = select(k, (2, 3, 1, 0)) + v_kernel = select(v, (3, 2, 1, 0)) + o_kernel = select(kernel_output, (2, 3, 1, 0)) + lse_kernel = select(kernel_lse, (2, 1, 0)) + index_kernel = select(block_index, (3, 2, 1, 0)) + nums_kernel = select(block_nums, (2, 1, 0)) + if block_sizes is None: + sizes_kernel = nums_kernel + elif self.block_sizes_desc.ndim == 1: + sizes_kernel = cute.make_tensor( + block_sizes.iterator, + cute.make_layout( + ( + self._op.num_kv_blocks, + self.block_index_desc.shape[1], + self._op.batch, + ), + stride=(block_sizes.stride[0], 0, 0), + ), + ) + elif self.block_sizes_desc.ndim == 2: + sizes_kernel = cute.make_tensor( + block_sizes.iterator, + cute.make_layout( + ( + self._op.num_kv_blocks, + self.block_index_desc.shape[1], + self._op.batch, + ), + stride=(block_sizes.stride[1], 0, block_sizes.stride[0]), + ), + ) + else: + sizes_kernel = select(block_sizes, (2, 1, 0)) + split_kernel = ( + nums_kernel + if split_offsets is None + else select(split_offsets, (3, 2, 1, 0)) + ) + kernel = BlockSparseAttnForwardSm90Blk64( + gqa_ratio=self._op.gqa_ratio, + head_dim=self._op.head_dim, + value_dim=self._op.value_dim, + blocksparse_blocksize_q=64, + blocksparse_blocksize_k=64, + dtype=q.element_type, + acc_dtype=cutlass.Float32, + has_block_sizes=block_sizes is not None, + num_splits=self.kv_splits, + allow_empty_block_nums=self._op.allow_empty_block_nums + and self.block_nums_desc is not None, + ) + kernel( + q_kernel, + k_kernel, + v_kernel, + o_kernel, + lse_kernel, + index_kernel, + nums_kernel, + sizes_kernel, + split_kernel, + cutlass.Float32(self.softmax_scale), + stream, + ) + elif self.compute_capability_family == 120: + from .csrc.fwd.sm120_blk64.bsa_fwd_sm120 import ( + BlockSparseAttnForwardSm120Blk64, + ) + + def select(tensor: Any, modes: tuple[int, ...]) -> Any: + return cute.make_tensor( + tensor.iterator, cute.select(tensor.layout, mode=modes) + ) + + q_kernel = select(q, (2, 3, 1, 0)) + k_kernel = select(k, (2, 3, 1, 0)) + v_kernel = select(v, (3, 2, 1, 0)) + o_kernel = select(kernel_output, (2, 3, 1, 0)) + lse_kernel = select(kernel_lse, (2, 1, 0)) + index_kernel = select(block_index, (3, 2, 1, 0)) + nums_kernel = ( + index_kernel if block_nums is None else select(block_nums, (2, 1, 0)) + ) + block_sizes_mode = 0 + sizes_kernel = nums_kernel + if block_sizes is not None: + block_sizes_mode = self.block_sizes_desc.ndim + sizes_kernel = ( + block_sizes + if block_sizes_mode == 1 + else select(block_sizes, tuple(reversed(range(block_sizes_mode)))) + ) + kernel = BlockSparseAttnForwardSm120Blk64( + gqa_ratio=self._op.gqa_ratio, + head_dim=self._op.head_dim, + value_dim=self._op.value_dim, + blocksparse_blocksize_q=64, + blocksparse_blocksize_k=64, + dtype=q.element_type, + acc_dtype=cutlass.Float32, + has_block_sizes=block_sizes is not None, + has_block_nums=self.block_nums_desc is not None, + block_sizes_mode=block_sizes_mode, + ) + kernel( + q_kernel, + k_kernel, + v_kernel, + o_kernel, + lse_kernel, + index_kernel, + nums_kernel, + cutlass.Int32( + 0 if self.block_nums_desc is not None else self.block_sparse_num + ), + sizes_kernel, + cutlass.Float32(self.softmax_scale), + stream, + ) + elif self.sparse_block_size == 64: + from .csrc.fwd.sm100_blk64.bsa_fwd_sm100 import ( + BlockSparseAttnForwardSm100Blk64, + ) + + use_clc = _resolve_sm100_blk64_use_clc( + kv_splits=self.kv_splits, + requested=self.use_clc, + batch=self._op.batch, + heads=self._op.num_q_heads, + seqlen_q=self._op.seqlen_q, + block_sparse_num=self.block_sparse_num, + has_variable_block_nums=self.block_nums_desc is not None, + ) + kernel = BlockSparseAttnForwardSm100Blk64( + self._op.head_dim, + self._op.value_dim, + qhead_per_kvhead=1, + pack_gqa=False, + m_block_size=64, + n_block_size=256, + sparse_block_size=64, + is_persistent=use_clc, + use_clc_scheduler=use_clc, + allow_empty_block_nums=self._op.allow_empty_block_nums + and self.block_nums_desc is not None, + has_block_sizes=block_sizes is not None, + num_splits=self.kv_splits, + use_int64_kv_strides=( + _sm100_blk64_uses_int64_kv_strides(self.k_desc) + or _sm100_blk64_uses_int64_kv_strides(self.v_desc) + ), + ) + kernel( + q, + k, + v, + kernel_output, + kernel_lse, + cutlass.Float32(self.softmax_scale), + block_index, + block_sizes, + cutlass.Int32( + 0 if self.block_nums_desc is not None else self.block_sparse_num + ), + block_nums, + split_offsets, + stream, + ) + else: + from .csrc.fwd.sm100_blk128.bsa_fwd_sm100 import ( + BlockSparseAttnForwardSm100Blk128, + ) + + kernel = BlockSparseAttnForwardSm100Blk128( + self._op.head_dim, + self._op.value_dim, + qhead_per_kvhead=self._op.gqa_ratio, + pack_gqa=self._op.pack_gqa, + allow_empty_block_nums=self._op.allow_empty_block_nums + and self.block_nums_desc is not None, + has_block_sizes=block_sizes is not None, + ) + kernel( + q, + k, + v, + output, + lse, + cutlass.Float32(self.softmax_scale), + block_index, + block_sizes, + cutlass.Int32(self.block_sparse_num), + block_nums, + stream, + ) + + if self.kv_splits > 1: + from .csrc.fwd.sm100_blk64.bsa_fwd_combine import ( + BlockSparseAttnForwardCombine, + ) + + split_heads = self.kv_splits * self._op.num_q_heads + o_partial = cute.make_tensor( + kernel_output.iterator, + cute.make_layout( + ( + self.kv_splits, + self._op.batch, + self._op.seqlen_q, + self._op.num_q_heads, + self._op.value_dim, + ), + stride=( + self._op.num_q_heads * self._op.seqlen_q * self._op.value_dim, + self._op.seqlen_q * split_heads * self._op.value_dim, + self._op.value_dim, + self._op.seqlen_q * self._op.value_dim, + 1, + ), + ), + ) + lse_partial = cute.make_tensor( + kernel_lse.iterator, + cute.make_layout( + ( + self.kv_splits, + self._op.batch, + self._op.seqlen_q, + self._op.num_q_heads, + ), + stride=( + self._op.num_q_heads * self._op.seqlen_q, + self._op.seqlen_q * split_heads, + 1, + self._op.seqlen_q, + ), + ), + ) + output_bshd = cute.make_tensor( + output.iterator, cute.select(output.layout, mode=(0, 2, 1, 3)) + ) + lse_bsh = cute.make_tensor( + lse.iterator, cute.select(lse.layout, mode=(0, 2, 1)) + ) + combine = BlockSparseAttnForwardCombine( + dtype=output.element_type, + head_dim=self._op.value_dim, + tile_m=16, + k_block_size=64, + log_max_splits=(self.kv_splits - 1).bit_length(), + num_threads=128, + stages=4, + ) + combine( + o_partial, + lse_partial, + output_bshd, + lse_bsh, + None, + None, + None, + None, + None, + stream, + ) + + +class BlockSparseAttentionBackward(JaxApiBase): + """JAX callable specialized from BSA backward tensor metadata.""" + + def __init__( + self, + sample_do: Any, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_o: Any, + sample_lse: Any, + sample_q2k_block_index: Any, + *, + sample_block_sizes: Any | None = None, + sample_q2k_block_nums: Any | None = None, + sample_dq: Any | None = None, + sample_dk: Any | None = None, + sample_dv: Any | None = None, + block_sparse_num: int | None = None, + bucket_size_blocks: int | None = None, + sparse_block_size: int | None = None, + softmax_scale: float | None = None, + layout: str = "bhsd", + target_compute_capability: int | None = None, + ) -> None: + self.layout = layout + self.data_mode = _data_mode(layout) + self.compute_capability = self._resolve_compute_capability( + target_compute_capability, + SUPPORTED_BACKWARD_COMPUTE_CAPABILITIES, + "BlockSparseAttentionBackward", + ) + self.compute_capability_family = compute_capability_family( + self.compute_capability + ) + self.do_desc = self._to_tensor_desc(sample_do, "sample_do", mode=self.data_mode) + self.q_desc = self._to_tensor_desc(sample_q, "sample_q", mode=self.data_mode) + self.k_desc = self._to_tensor_desc(sample_k, "sample_k", mode=self.data_mode) + self.v_desc = self._to_tensor_desc(sample_v, "sample_v", mode=self.data_mode) + self.o_desc = self._to_tensor_desc(sample_o, "sample_o", mode=self.data_mode) + self.lse_desc = self._to_tensor_desc(sample_lse, "sample_lse") + self.block_index_desc = self._to_tensor_desc( + sample_q2k_block_index, "sample_q2k_block_index" + ) + self.block_sizes_desc = ( + None + if sample_block_sizes is None + else self._to_tensor_desc(sample_block_sizes, "sample_block_sizes") + ) + self.block_nums_desc = ( + None + if sample_q2k_block_nums is None + else self._to_tensor_desc(sample_q2k_block_nums, "sample_q2k_block_nums") + ) + + def output_desc( + sample: Any | None, source: JaxTensorDesc, name: str + ) -> JaxTensorDesc: + if sample is not None: + return self._to_tensor_desc(sample, name, mode=self.data_mode) + return self._to_tensor_desc( + _shape_sample(source.shape, source.dtype, self.data_mode), + name, + mode=self.data_mode, + ) + + self.dq_desc = output_desc(sample_dq, self.q_desc, "sample_dq") + self.dk_desc = output_desc(sample_dk, self.k_desc, "sample_dk") + self.dv_desc = output_desc(sample_dv, self.v_desc, "sample_dv") + self.sparse_block_size = ( + _default_backward_sparse_block_size(self.compute_capability) + if sparse_block_size is None + else sparse_block_size + ) + capacity = ( + self.block_index_desc.shape[3] if self.block_index_desc.ndim == 4 else 0 + ) + self.block_sparse_num = ( + capacity if block_sparse_num is None else block_sparse_num + ) + head_dim = self.q_desc.shape[3] if self.q_desc.ndim == 4 else 1 + self.softmax_scale = ( + head_dim**-0.5 if softmax_scale is None else float(softmax_scale) + ) + num_q_blocks = ( + ceil_div(self.q_desc.shape[2], self.sparse_block_size) + if self.q_desc.ndim == 4 + else 0 + ) + self.bucket_size_blocks = ( + _default_bucket_size( + self.compute_capability, + self.sparse_block_size, + num_q_blocks, + self.q_desc.shape[1], + ) + if bucket_size_blocks is None + else bucket_size_blocks + ) + self._op = BlockSparseAttentionBackwardOp( + dout=self.do_desc, + q=self.q_desc, + k=self.k_desc, + v=self.v_desc, + output=self.o_desc, + lse=self.lse_desc, + block_index=self.block_index_desc, + dq=self.dq_desc, + dk=self.dk_desc, + dv=self.dv_desc, + block_sizes=self.block_sizes_desc, + block_nums=self.block_nums_desc, + block_sparse_num=self.block_sparse_num, + sparse_block_size=self.sparse_block_size, + softmax_scale=self.softmax_scale, + bucket_size_blocks=self.bucket_size_blocks, + target_compute_capability=self.compute_capability, + ) + self.check_support() + + def check_support(self) -> bool: + return self._op.check_support() + + def _workspace_descs(self) -> tuple[tuple[str, JaxTensorDesc], ...]: + int_source = self.block_index_desc + float_source = self.lse_desc + b, h = self._op.batch, self._op.num_heads + g, nk = self._op.num_q_groups, self._op.num_kv_blocks + workspaces: list[tuple[str, JaxTensorDesc]] = [ + ( + "counts", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, h, g, nk), + name="counts_workspace", + init_value=0, + ), + ), + ( + "local_offsets", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, h, g, nk + 1), + name="local_offsets_workspace", + ), + ), + ( + "group_totals", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, h, g), + name="group_totals_workspace", + ), + ), + ( + "bucket_offsets", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, h, g, nk + 1), + name="bucket_offsets_workspace", + ), + ), + ( + "cursors", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, h, g, nk), + name="cursors_workspace", + ), + ), + ( + "bucket_indices", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, h, self._op.max_edges), + name="bucket_indices_workspace", + ), + ), + ] + if self.block_nums_desc is None: + workspaces.append( + ( + "block_nums_placeholder", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, h, self._op.num_q_blocks), + name="block_nums_placeholder", + init_value=self.block_sparse_num, + ), + ) + ) + if self.block_sizes_desc is None: + workspaces.append( + ( + "block_sizes_placeholder", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, nk), + name="block_sizes_placeholder", + init_value=0, + ), + ) + ) + + if self.compute_capability_family == 100 and self.sparse_block_size == 128: + q_rounded = self._op.num_q_blocks * 128 + k_rounded = self._op.num_kv_blocks * 128 + d_rounded = _round_up(self._op.head_dim, 32) + workspaces.extend( + ( + ( + "dpsum", + float_source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(b, h, q_rounded), + name="dpsum_workspace", + ), + ), + ( + "lse_log2", + float_source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(b, h, q_rounded), + name="lse_log2_workspace", + ), + ), + ( + "dq_accum", + float_source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(b, h, q_rounded * d_rounded), + name="dq_accum_workspace", + ), + ), + ) + ) + if self._op.num_q_groups > 1: + workspaces.extend( + ( + ( + "dk_accum", + float_source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(b, h, k_rounded * d_rounded), + name="dk_accum_workspace", + init_value=0.0, + ), + ), + ( + "dv_accum", + float_source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(b, h, k_rounded * d_rounded), + name="dv_accum_workspace", + init_value=0.0, + ), + ), + ) + ) + else: + q_round_to = 64 if self.compute_capability_family == 90 else 8 + k_round_to = 64 if self.compute_capability_family == 90 else 8 + d_round_to = 32 if self.compute_capability_family == 90 else 8 + q_rounded = _round_up(self._op.seqlen_q, q_round_to) + k_rounded = _round_up(self._op.seqlen_k, k_round_to) + d_rounded = _round_up(self._op.head_dim, d_round_to) + elems_per_bh = ( + 2 * q_rounded + q_rounded * d_rounded + 2 * k_rounded * d_rounded + ) + workspaces.append( + ( + "bwd_workspace", + float_source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(b, h, elems_per_bh), + name="bwd_workspace", + init_value=0.0, + ), + ) + ) + return tuple(workspaces) + + def __call__( + self, + do: Any, + q: Any, + k: Any, + v: Any, + o: Any, + lse: Any, + q2k_block_index: Any, + block_sizes: Any | None = None, + q2k_block_nums: Any | None = None, + ) -> TupleDict: + self.check_support() + for value, desc in ( + (do, self.do_desc), + (q, self.q_desc), + (k, self.k_desc), + (v, self.v_desc), + (o, self.o_desc), + ): + self._check_tensor_signature(value, desc, mode=self.data_mode) + self._check_tensor_signature(lse, self.lse_desc) + self._check_tensor_signature(q2k_block_index, self.block_index_desc) + _optional_signature(block_sizes, self.block_sizes_desc, "block_sizes") + _optional_signature(q2k_block_nums, self.block_nums_desc, "q2k_block_nums") + + inputs: list[Any] = [do, q, k, v, o, lse, q2k_block_index] + bindings: list[tuple[JaxTensorDesc, tuple[int, ...] | None]] = [ + (self.do_desc, self.data_mode), + (self.q_desc, self.data_mode), + (self.k_desc, self.data_mode), + (self.v_desc, self.data_mode), + (self.o_desc, self.data_mode), + (self.lse_desc, None), + (self.block_index_desc, None), + ] + if self.block_sizes_desc is not None: + inputs.append(block_sizes) + bindings.append((self.block_sizes_desc, None)) + if self.block_nums_desc is not None: + inputs.append(q2k_block_nums) + bindings.append((self.block_nums_desc, None)) + self._backward_input_count = len(inputs) + workspace_items = self._workspace_descs() + self._backward_workspace_names = tuple(name for name, _ in workspace_items) + workspace_descs = tuple(desc for _, desc in workspace_items) + + dq, dk, dv = self._call_kernel( + tuple(inputs), + launch=self._launch, + output_descs=(self.dq_desc, self.dk_desc, self.dv_desc), + workspace_descs=workspace_descs, + input_spec=tuple( + self._to_tensor_spec(desc, mode=mode) for desc, mode in bindings + ), + output_spec=( + self._to_tensor_spec(self.dq_desc, mode=self.data_mode), + self._to_tensor_spec(self.dk_desc, mode=self.data_mode), + self._to_tensor_spec(self.dv_desc, mode=self.data_mode), + ), + compile_options=compile_options_for_target(self.compute_capability), + ) + return TupleDict(dq_tensor=dq, dk_tensor=dk, dv_tensor=dv) + + def _launch(self, stream: Any, *arguments: Any) -> None: + import cutlass + import cutlass.cute as cute + + inputs = list(arguments[: self._backward_input_count]) + dq, dk, dv = arguments[ + self._backward_input_count : self._backward_input_count + 3 + ] + workspaces = dict( + zip( + self._backward_workspace_names, + arguments[self._backward_input_count + 3 :], + ) + ) + do, q, k, v, o, lse, block_index = inputs[:7] + cursor = 7 + block_sizes = None + if self.block_sizes_desc is not None: + block_sizes = inputs[cursor] + cursor += 1 + block_nums = None + if self.block_nums_desc is not None: + block_nums = inputs[cursor] + cursor += 1 + if cursor != len(inputs): + raise RuntimeError("Unexpected BSA backward kernel inputs") + if block_nums is None: + block_nums_kernel = workspaces["block_nums_placeholder"] + else: + block_nums_kernel = block_nums + + from .csrc.bwd.bucketed_k2q_csr import BucketedK2QCsrUniversal + + builder = BucketedK2QCsrUniversal( + self.block_sparse_num, + self.bucket_size_blocks, + self.block_nums_desc is not None, + self.block_index_desc.shape[3], + ) + builder( + workspaces["counts"], + workspaces["local_offsets"], + workspaces["group_totals"], + workspaces["bucket_offsets"], + workspaces["cursors"], + workspaces["bucket_indices"], + block_index, + block_nums_kernel, + stream, + ) + + if block_sizes is None: + variable_block_sizes = workspaces["block_sizes_placeholder"] + elif self.block_sizes_desc.ndim == 1: + variable_block_sizes = cute.make_tensor( + block_sizes.iterator, + cute.make_layout( + (self._op.batch, self._op.num_kv_blocks), + stride=(0, block_sizes.stride[0]), + ), + ) + else: + variable_block_sizes = block_sizes + + if self.compute_capability_family == 90: + from .csrc.bwd.sm90_blk64.bsa_bwd_sm90 import ( + BlockSparseAttnBackwardSm90Blk64, + ) + + kernel = BlockSparseAttnBackwardSm90Blk64( + q.element_type, self._op.head_dim, self._op.head_dim + ) + kernel( + ( + cutlass.Int32(self._op.seqlen_q), + cutlass.Int32(self._op.seqlen_k), + cutlass.Int32(self._op.head_dim), + (cutlass.Int32(self._op.num_heads), cutlass.Int32(self._op.batch)), + ), + do, + o, + q, + k, + v, + lse, + dq, + dk, + dv, + workspaces["bucket_offsets"], + workspaces["bucket_indices"], + variable_block_sizes if self.block_sizes_desc is not None else None, + workspaces["bwd_workspace"], + cutlass.Float32(self.softmax_scale), + stream, + ) + return + + if self.sparse_block_size == 64: + from .csrc.bwd.sm100_blk64.bsa_bwd_sm100 import ( + BlockSparseAttnBackwardSm100Blk64, + ) + + kernel = BlockSparseAttnBackwardSm100Blk64( + sparse_block_size=64, + has_block_sizes=self.block_sizes_desc is not None, + ) + kernel( + ( + cutlass.Int32(self._op.seqlen_q), + cutlass.Int32(self._op.seqlen_k), + cutlass.Int32(self._op.head_dim), + (cutlass.Int32(self._op.num_heads), cutlass.Int32(self._op.batch)), + ), + do, + o, + q, + k, + v, + lse, + dq, + dk, + dv, + workspaces["bucket_offsets"], + workspaces["bucket_indices"], + variable_block_sizes, + workspaces["bwd_workspace"], + cutlass.Float32(self.softmax_scale), + stream, + ) + return + + from .csrc.bwd.bsa_bwd_postprocess import BlockSparseAttnBackwardPostprocess + from .csrc.bwd.bsa_bwd_preprocess import BlockSparseAttnBackwardPreprocess + from .csrc.bwd.sm100_blk128.bsa_bwd_sm100 import ( + BsaK2qCsrTensors, + BlockSparseAttnBackwardSm100Blk128, + ) + + def bhsd_to_bshd(tensor: Any) -> Any: + return cute.make_tensor( + tensor.iterator, cute.select(tensor.layout, mode=(0, 2, 1, 3)) + ) + + do_bshd, q_bshd, k_bshd, v_bshd, o_bshd = [ + bhsd_to_bshd(tensor) for tensor in (do, q, k, v, o) + ] + dq_bshd, dk_bshd, dv_bshd = [bhsd_to_bshd(tensor) for tensor in (dq, dk, dv)] + preprocess = BlockSparseAttnBackwardPreprocess( + q.element_type, + self._op.head_dim, + self._op.head_dim, + 128, + ) + preprocess( + o_bshd, + do_bshd, + workspaces["dpsum"], + lse, + workspaces["lse_log2"], + workspaces["dq_accum"], + None, + None, + None, + stream, + ) + use_dkv_postprocess = self._op.num_q_groups > 1 + core = BlockSparseAttnBackwardSm100Blk128( + self._op.head_dim, + force_dkv_postprocess=use_dkv_postprocess, + ) + core( + q_bshd, + k_bshd, + v_bshd, + do_bshd, + workspaces["lse_log2"], + workspaces["dpsum"], + workspaces["dq_accum"], + workspaces["dk_accum"] if use_dkv_postprocess else dk_bshd, + workspaces["dv_accum"] if use_dkv_postprocess else dv_bshd, + cutlass.Float32(self.softmax_scale), + BsaK2qCsrTensors( + workspaces["bucket_offsets"], + workspaces["bucket_indices"], + ), + stream, + ) + postprocess = BlockSparseAttnBackwardPostprocess( + q.element_type, + self._op.head_dim, + self.compute_capability, + 128, + False, + ) + postprocess( + workspaces["dq_accum"], + dq_bshd, + cutlass.Float32(self.softmax_scale), + None, + None, + stream, + ) + if use_dkv_postprocess: + postprocess( + workspaces["dk_accum"], + dk_bshd, + cutlass.Float32(self.softmax_scale), + None, + None, + stream, + ) + postprocess( + workspaces["dv_accum"], + dv_bshd, + cutlass.Float32(1.0), + None, + None, + stream, + ) + + +@partial( + jax.jit, + static_argnames=( + "block_sparse_num", + "sparse_block_size", + "allow_empty_block_nums", + "softmax_scale", + "pack_gqa", + "layout", + "kv_splits", + "use_clc", + "target_compute_capability", + ), +) +def block_sparse_attention_forward( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + q2k_block_index: Any, + block_sparse_num: int | None = None, + block_sizes: Any | None = None, + q2k_block_nums: Any | None = None, + *, + sparse_block_size: int | None = None, + allow_empty_block_nums: bool = False, + softmax_scale: float | None = None, + pack_gqa: bool | None = None, + layout: str = "bhsd", + kv_splits: int | str = 1, + use_clc: bool | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + """Run non-causal block-sparse attention from JAX arrays.""" + + def sample(value: Any | None) -> Any | None: + return None if value is None else jax.ShapeDtypeStruct(value.shape, value.dtype) + + return BlockSparseAttentionForward( + sample(q_tensor), + sample(k_tensor), + sample(v_tensor), + sample(q2k_block_index), + sample_block_sizes=sample(block_sizes), + sample_q2k_block_nums=sample(q2k_block_nums), + block_sparse_num=block_sparse_num, + sparse_block_size=sparse_block_size, + allow_empty_block_nums=allow_empty_block_nums, + softmax_scale=softmax_scale, + pack_gqa=pack_gqa, + layout=layout, + kv_splits=kv_splits, + use_clc=use_clc, + target_compute_capability=target_compute_capability, + )( + q_tensor, + k_tensor, + v_tensor, + q2k_block_index, + block_sizes, + q2k_block_nums, + ) + + +@partial( + jax.jit, + static_argnames=( + "block_sparse_num", + "bucket_size_blocks", + "sparse_block_size", + "softmax_scale", + "layout", + "target_compute_capability", + ), +) +def block_sparse_attention_backward( + do_tensor: Any, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + o_tensor: Any, + lse_tensor: Any, + q2k_block_index: Any, + block_sparse_num: int | None = None, + block_sizes: Any | None = None, + q2k_block_nums: Any | None = None, + *, + softmax_scale: float | None = None, + bucket_size_blocks: int | None = None, + sparse_block_size: int | None = None, + layout: str = "bhsd", + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute explicit block-sparse attention gradients from JAX arrays.""" + + def sample(value: Any | None) -> Any | None: + return None if value is None else jax.ShapeDtypeStruct(value.shape, value.dtype) + + return BlockSparseAttentionBackward( + sample(do_tensor), + sample(q_tensor), + sample(k_tensor), + sample(v_tensor), + sample(o_tensor), + sample(lse_tensor), + sample(q2k_block_index), + sample_block_sizes=sample(block_sizes), + sample_q2k_block_nums=sample(q2k_block_nums), + block_sparse_num=block_sparse_num, + bucket_size_blocks=bucket_size_blocks, + sparse_block_size=sparse_block_size, + softmax_scale=softmax_scale, + layout=layout, + target_compute_capability=target_compute_capability, + )( + do_tensor, + q_tensor, + k_tensor, + v_tensor, + o_tensor, + lse_tensor, + q2k_block_index, + block_sizes, + q2k_block_nums, + ) + + +__all__ = [ + "BlockSparseAttentionBackward", + "BlockSparseAttentionForward", + "block_sparse_attention_backward", + "block_sparse_attention_forward", +] diff --git a/python/cudnn/block_sparse_attention/op.py b/python/cudnn/block_sparse_attention/op.py new file mode 100644 index 000000000..bec4ecf31 --- /dev/null +++ b/python/cudnn/block_sparse_attention/op.py @@ -0,0 +1,481 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral contracts for block-sparse attention kernels.""" + +from __future__ import annotations + +from typing import Any + +from .. import data_type +from .._op import Op +from .._tensor_desc import TensorDesc + + +SUPPORTED_FORWARD_COMPUTE_CAPABILITIES = (90, 100, 103, 107, 110, 120) +SUPPORTED_BACKWARD_COMPUTE_CAPABILITIES = (90, 100, 103, 107, 110) + + +def ceil_div(value: int, divisor: int) -> int: + return (value + divisor - 1) // divisor + + +def compute_capability_family(compute_capability: int) -> int: + if compute_capability == 90: + return 90 + if 100 <= compute_capability < 120: + return 100 + if compute_capability == 120: + return 120 + raise ValueError( + f"Unsupported block-sparse attention target SM{compute_capability}" + ) + + +def _require_shape(desc: TensorDesc[Any], expected: tuple[int, ...]) -> None: + if desc.shape != expected: + raise ValueError( + f"{desc.name or 'tensor'} shape must be {expected}, got {desc.shape}" + ) + + +def _require_dtype(desc: TensorDesc[Any], expected: data_type) -> None: + if desc.cudnn_dtype != expected: + raise ValueError( + f"{desc.name or 'tensor'} dtype must be {expected}, got {desc.cudnn_dtype}" + ) + + +class BlockSparseAttentionForwardOp(Op): + """Complete, framework-neutral signature for BSA forward.""" + + def __init__( + self, + *, + q: TensorDesc[Any], + k: TensorDesc[Any], + v: TensorDesc[Any], + block_index: TensorDesc[Any], + output: TensorDesc[Any], + lse: TensorDesc[Any], + block_sizes: TensorDesc[Any] | None, + block_nums: TensorDesc[Any] | None, + block_sparse_num: int, + sparse_block_size: int, + softmax_scale: float, + pack_gqa: bool | None, + allow_empty_block_nums: bool, + kv_splits: int, + use_clc: bool | None, + target_compute_capability: int, + ) -> None: + self.q = q + self.k = k + self.v = v + self.block_index = block_index + self.output = output + self.lse = lse + self.block_sizes = block_sizes + self.block_nums = block_nums + self.block_sparse_num = block_sparse_num + self.sparse_block_size = sparse_block_size + self.softmax_scale = softmax_scale + self.pack_gqa = pack_gqa + self.allow_empty_block_nums = allow_empty_block_nums + self.kv_splits = kv_splits + self.use_clc = use_clc + self.target_compute_capability = target_compute_capability + + self.batch = 0 + self.num_q_heads = 0 + self.num_kv_heads = 0 + self.seqlen_q = 0 + self.seqlen_k = 0 + self.head_dim = 0 + self.value_dim = 0 + self.num_q_blocks = 0 + self.num_kv_blocks = 0 + self.gqa_ratio = 0 + self.pack_gqa_effective = False + + def check_support(self) -> bool: + if self.target_compute_capability not in SUPPORTED_FORWARD_COMPUTE_CAPABILITIES: + raise ValueError( + f"BlockSparseAttentionForward has no kernel for SM{self.target_compute_capability}; " + f"supported targets are {SUPPORTED_FORWARD_COMPUTE_CAPABILITIES}" + ) + family = compute_capability_family(self.target_compute_capability) + if self.q.ndim != 4 or self.k.ndim != 4 or self.v.ndim != 4: + raise ValueError("q, k, and v must all be rank-4 tensors") + + batch, num_q_heads, seqlen_q, head_dim = self.q.shape + batch_k, num_kv_heads, seqlen_k, head_dim_k = self.k.shape + batch_v, num_v_heads, seqlen_v, value_dim = self.v.shape + if ( + min( + batch, + num_q_heads, + num_kv_heads, + seqlen_q, + seqlen_k, + head_dim, + value_dim, + ) + <= 0 + ): + raise ValueError("q, k, and v dimensions must all be positive") + if batch_k != batch or batch_v != batch: + raise ValueError("q, k, and v batch dimensions must match") + if num_v_heads != num_kv_heads or seqlen_v != seqlen_k: + raise ValueError("k and v head and sequence dimensions must match") + if head_dim_k != head_dim: + raise ValueError("q and k head dimensions must match") + if num_q_heads % num_kv_heads: + raise ValueError( + "the number of query heads must be divisible by the number of KV heads" + ) + if self.q.cudnn_dtype not in {data_type.HALF, data_type.BFLOAT16}: + raise ValueError("q, k, and v must use float16 or bfloat16") + if ( + self.k.cudnn_dtype != self.q.cudnn_dtype + or self.v.cudnn_dtype != self.q.cudnn_dtype + ): + raise ValueError("q, k, and v must have the same dtype") + if self.q.stride[3] != 1 or self.k.stride[3] != 1 or self.v.stride[3] != 1: + raise ValueError("q, k, and v must have a contiguous head dimension") + + if self.sparse_block_size not in {64, 128}: + raise ValueError("sparse_block_size must be 64 or 128") + if family in {90, 120} and self.sparse_block_size != 64: + raise ValueError( + f"SM{self.target_compute_capability} forward requires sparse_block_size=64" + ) + if family == 90: + if head_dim not in {64, 96, 128} or value_dim not in {64, 96, 128}: + raise ValueError( + "SM90 forward supports QK and V dimensions 64, 96, or 128" + ) + if seqlen_q % 64: + raise ValueError( + "SM90 forward requires seqlen_q to be a multiple of 64" + ) + elif family == 120: + if (head_dim, value_dim) != (128, 128): + raise ValueError("SM120 forward requires QK and V dimensions of 128") + elif self.sparse_block_size == 64: + if self.q.cudnn_dtype != data_type.BFLOAT16 or (head_dim, value_dim) != ( + 128, + 128, + ): + raise ValueError( + "SM100-family blk64 forward requires BF16 and QK=V=128" + ) + if num_q_heads != num_kv_heads: + raise ValueError("SM100-family blk64 forward supports MHA only") + elif (head_dim, value_dim) not in {(64, 64), (96, 96), (128, 128)}: + raise ValueError( + "SM100-family blk128 forward supports (QK, V) dimensions " + f"(64, 64), (96, 96), or (128, 128); got ({head_dim}, {value_dim})" + ) + + gqa_ratio = num_q_heads // num_kv_heads + if self.pack_gqa is True and not ( + family == 100 and self.sparse_block_size == 128 + ): + raise ValueError( + "pack_gqa is available only on the SM100-family blk128 path" + ) + if self.pack_gqa is True and 128 % gqa_ratio: + raise ValueError("pack_gqa=True requires the GQA ratio to divide 128") + pack_gqa_effective = ( + family == 100 + and self.sparse_block_size == 128 + and (gqa_ratio > 1 if self.pack_gqa is None else self.pack_gqa) + and 128 % gqa_ratio == 0 + ) + metadata_heads = num_kv_heads if pack_gqa_effective else num_q_heads + metadata_tokens = seqlen_q * gqa_ratio if pack_gqa_effective else seqlen_q + num_q_blocks = ceil_div(metadata_tokens, self.sparse_block_size) + num_kv_blocks = ceil_div(seqlen_k, self.sparse_block_size) + + if self.block_index.ndim != 4: + raise ValueError("block_index must be rank 4") + _require_dtype(self.block_index, data_type.INT32) + if self.block_index.shape[:3] != (batch, metadata_heads, num_q_blocks): + raise ValueError( + "block_index shape prefix must be " + f"{(batch, metadata_heads, num_q_blocks)}, got {self.block_index.shape[:3]}" + ) + capacity = self.block_index.shape[3] + if capacity <= 0: + raise ValueError("block_index must have non-empty KV-block capacity") + + if self.block_nums is None: + minimum = 2 if family == 100 and self.sparse_block_size == 128 else 1 + if not minimum <= self.block_sparse_num <= capacity: + raise ValueError( + f"block_sparse_num must be in [{minimum}, {capacity}], got {self.block_sparse_num}" + ) + if minimum == 2 and self.block_sparse_num % 2: + raise ValueError( + "block_sparse_num must be even for the SM100-family blk128 kernel" + ) + else: + _require_dtype(self.block_nums, data_type.INT32) + _require_shape(self.block_nums, (batch, metadata_heads, num_q_blocks)) + + if self.block_sizes is not None: + _require_dtype(self.block_sizes, data_type.INT32) + allowed_ranks = {1, 2, 3} if family in {90, 120} else {1} + if self.block_sizes.ndim not in allowed_ranks: + raise ValueError( + f"block_sizes rank must be one of {tuple(sorted(allowed_ranks))}" + ) + expected = { + 1: (num_kv_blocks,), + 2: (batch, num_kv_blocks), + 3: (batch, metadata_heads, num_kv_blocks), + }[self.block_sizes.ndim] + _require_shape(self.block_sizes, expected) + + if ( + isinstance(self.kv_splits, bool) + or not isinstance(self.kv_splits, int) + or not 1 <= self.kv_splits <= 256 + ): + raise ValueError( + f"kv_splits must be an integer in [1, 256], got {self.kv_splits!r}" + ) + if self.use_clc is not None and not isinstance(self.use_clc, bool): + raise TypeError( + f"use_clc must be a bool or None, got {type(self.use_clc).__name__}" + ) + if self.use_clc is not None and not ( + family == 100 and self.sparse_block_size == 64 + ): + raise ValueError("use_clc is available only on the SM100-family blk64 path") + if family == 120 and self.kv_splits != 1: + raise ValueError("SM120 forward does not support kv_splits") + if family == 100 and self.sparse_block_size == 128 and self.kv_splits != 1: + raise ValueError("SM100-family blk128 forward does not support kv_splits") + if self.kv_splits > 1 and self.use_clc: + raise ValueError("kv_splits > 1 does not support use_clc=True") + + _require_shape(self.output, (batch, num_q_heads, seqlen_q, value_dim)) + _require_dtype(self.output, self.q.cudnn_dtype) + _require_shape(self.lse, (batch, num_q_heads, seqlen_q)) + _require_dtype(self.lse, data_type.FLOAT) + + self.batch = batch + self.num_q_heads = num_q_heads + self.num_kv_heads = num_kv_heads + self.seqlen_q = seqlen_q + self.seqlen_k = seqlen_k + self.head_dim = head_dim + self.value_dim = value_dim + self.num_q_blocks = num_q_blocks + self.num_kv_blocks = num_kv_blocks + self.gqa_ratio = gqa_ratio + self.pack_gqa_effective = pack_gqa_effective + return True + + +class BlockSparseAttentionBackwardOp(Op): + """Complete, framework-neutral signature for BSA backward.""" + + def __init__( + self, + *, + dout: TensorDesc[Any], + q: TensorDesc[Any], + k: TensorDesc[Any], + v: TensorDesc[Any], + output: TensorDesc[Any], + lse: TensorDesc[Any], + block_index: TensorDesc[Any], + dq: TensorDesc[Any], + dk: TensorDesc[Any], + dv: TensorDesc[Any], + block_sizes: TensorDesc[Any] | None, + block_nums: TensorDesc[Any] | None, + block_sparse_num: int, + sparse_block_size: int, + softmax_scale: float, + bucket_size_blocks: int, + target_compute_capability: int, + ) -> None: + self.dout = dout + self.q = q + self.k = k + self.v = v + self.output = output + self.lse = lse + self.block_index = block_index + self.dq = dq + self.dk = dk + self.dv = dv + self.block_sizes = block_sizes + self.block_nums = block_nums + self.block_sparse_num = block_sparse_num + self.sparse_block_size = sparse_block_size + self.softmax_scale = softmax_scale + self.bucket_size_blocks = bucket_size_blocks + self.target_compute_capability = target_compute_capability + + self.batch = 0 + self.num_heads = 0 + self.seqlen_q = 0 + self.seqlen_k = 0 + self.head_dim = 0 + self.num_q_blocks = 0 + self.num_kv_blocks = 0 + self.num_q_groups = 0 + self.max_edges = 0 + + def check_support(self) -> bool: + if ( + self.target_compute_capability + not in SUPPORTED_BACKWARD_COMPUTE_CAPABILITIES + ): + raise ValueError( + f"BlockSparseAttentionBackward has no kernel for SM{self.target_compute_capability}; " + f"supported targets are {SUPPORTED_BACKWARD_COMPUTE_CAPABILITIES}" + ) + family = compute_capability_family(self.target_compute_capability) + if any( + desc.ndim != 4 + for desc in ( + self.dout, + self.q, + self.k, + self.v, + self.output, + self.dq, + self.dk, + self.dv, + ) + ): + raise ValueError("BSA backward data tensors must all be rank 4") + batch, num_heads, seqlen_q, head_dim = self.q.shape + batch_k, num_kv_heads, seqlen_k, head_dim_k = self.k.shape + if min(batch, num_heads, seqlen_q, seqlen_k, head_dim) <= 0: + raise ValueError("BSA backward dimensions must be positive") + if batch_k != batch or num_kv_heads != num_heads or head_dim_k != head_dim: + raise ValueError( + "BSA backward supports MHA with matching q and k dimensions" + ) + _require_shape(self.v, self.k.shape) + _require_shape(self.dout, self.q.shape) + _require_shape(self.output, self.q.shape) + _require_shape(self.dq, self.q.shape) + _require_shape(self.dk, self.k.shape) + _require_shape(self.dv, self.v.shape) + if self.q.cudnn_dtype != data_type.BFLOAT16: + raise ValueError("BSA backward requires bfloat16 data tensors") + for desc in (self.dout, self.k, self.v, self.output, self.dq, self.dk, self.dv): + _require_dtype(desc, data_type.BFLOAT16) + if any( + desc.stride[3] != 1 + for desc in ( + self.dout, + self.q, + self.k, + self.v, + self.output, + self.dq, + self.dk, + self.dv, + ) + ): + raise ValueError( + "BSA backward data tensors must have a contiguous head dimension" + ) + _require_shape(self.lse, (batch, num_heads, seqlen_q)) + _require_dtype(self.lse, data_type.FLOAT) + + if family == 90: + if self.sparse_block_size != 64 or head_dim != 128: + raise ValueError( + "SM90 backward requires sparse_block_size=64 and head_dim=128" + ) + elif self.sparse_block_size == 64: + if head_dim != 128: + raise ValueError("SM100-family blk64 backward requires head_dim=128") + elif self.sparse_block_size == 128: + if head_dim not in {64, 128}: + raise ValueError( + "SM100-family blk128 backward requires head_dim=64 or 128" + ) + if self.block_sizes is not None: + raise ValueError( + "SM100-family blk128 backward does not support block_sizes" + ) + else: + raise ValueError("sparse_block_size must be 64 or 128") + + num_q_blocks = ceil_div(seqlen_q, self.sparse_block_size) + num_kv_blocks = ceil_div(seqlen_k, self.sparse_block_size) + if self.block_index.ndim != 4: + raise ValueError("block_index must be rank 4") + _require_dtype(self.block_index, data_type.INT32) + if self.block_index.shape[:3] != (batch, num_heads, num_q_blocks): + raise ValueError( + f"block_index shape prefix must be {(batch, num_heads, num_q_blocks)}, got {self.block_index.shape[:3]}" + ) + capacity = self.block_index.shape[3] + if capacity <= 0: + raise ValueError("block_index must have non-empty KV-block capacity") + if self.block_nums is None: + if not 1 <= self.block_sparse_num <= capacity: + raise ValueError( + f"block_sparse_num must be in [1, {capacity}], got {self.block_sparse_num}" + ) + if self.sparse_block_size == 128 and self.block_sparse_num % 2: + raise ValueError( + "block_sparse_num must be even for the blk128 backward kernel" + ) + max_edges = num_q_blocks * self.block_sparse_num + else: + _require_dtype(self.block_nums, data_type.INT32) + _require_shape(self.block_nums, (batch, num_heads, num_q_blocks)) + max_edges = num_q_blocks * capacity + + if self.block_sizes is not None: + _require_dtype(self.block_sizes, data_type.INT32) + if self.block_sizes.ndim not in {1, 2}: + raise ValueError("block_sizes must have rank 1 or 2 for BSA backward") + expected = ( + (num_kv_blocks,) + if self.block_sizes.ndim == 1 + else (batch, num_kv_blocks) + ) + _require_shape(self.block_sizes, expected) + if ( + isinstance(self.bucket_size_blocks, bool) + or not isinstance(self.bucket_size_blocks, int) + or self.bucket_size_blocks <= 0 + ): + raise ValueError( + f"bucket_size_blocks must be positive, got {self.bucket_size_blocks!r}" + ) + + self.batch = batch + self.num_heads = num_heads + self.seqlen_q = seqlen_q + self.seqlen_k = seqlen_k + self.head_dim = head_dim + self.num_q_blocks = num_q_blocks + self.num_kv_blocks = num_kv_blocks + self.num_q_groups = ceil_div(num_q_blocks, self.bucket_size_blocks) + self.max_edges = max(1, max_edges) + return True + + +__all__ = [ + "BlockSparseAttentionBackwardOp", + "BlockSparseAttentionForwardOp", + "SUPPORTED_BACKWARD_COMPUTE_CAPABILITIES", + "SUPPORTED_FORWARD_COMPUTE_CAPABILITIES", + "ceil_div", + "compute_capability_family", +] diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py index dee253739..d10b99ef1 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py @@ -13,6 +13,7 @@ from ... import data_type from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..._jax.layout import mode_from_layout from ..utils.compiler import compile_options_for_target from .op import DEFAULT_BLOCK_I, DenseIndexerBackwardOp, IndexerBackwardOp @@ -20,6 +21,21 @@ _SUPPORTED_COMPUTE_CAPABILITY_FAMILIES = (90, 100) +def _resolve_layout( + name: str, + layout: str | None, + *, + default: str, + kernel_axes: str, + supported: tuple[str, ...], +) -> tuple[str, tuple[int, ...]]: + layout = default if layout is None else layout + if layout not in supported: + choices = ", ".join(repr(value) for value in supported) + raise ValueError(f"{name} must be one of ({choices}), got {layout!r}") + return layout, mode_from_layout(layout, kernel_axes=kernel_axes) + + def _grad_loss_operand(value: Any) -> Any: value = jnp.asarray(value, dtype=jnp.float32) if value.shape == (): @@ -53,10 +69,16 @@ def _resolve_target(self, operation_name: str) -> None: @property def _architecture_family(self) -> int: if self.compute_capability is None: - raise RuntimeError("check_support() must resolve the compute capability before lowering") - family = self._compute_capability_family(self.compute_capability, _SUPPORTED_COMPUTE_CAPABILITY_FAMILIES) + raise RuntimeError( + "check_support() must resolve the compute capability before lowering" + ) + family = self._compute_capability_family( + self.compute_capability, _SUPPORTED_COMPUTE_CAPABILITY_FAMILIES + ) if family is None: - raise RuntimeError(f"No indexer-backward kernel for SM{self.compute_capability}") + raise RuntimeError( + f"No indexer-backward kernel for SM{self.compute_capability}" + ) return family @staticmethod @@ -67,14 +89,26 @@ def _output_desc( cudnn_dtype: data_type, shape: tuple[int, ...], name: str, + mode: tuple[int, ...], ) -> JaxTensorDesc: if sample is not None: - return JaxApiBase._to_tensor_desc(sample, name) - return source.compact_like(cudnn_dtype=cudnn_dtype, shape=shape, name=name) + return JaxApiBase._to_tensor_desc(sample, name, mode=mode) + return source.compact_like( + cudnn_dtype=cudnn_dtype, + shape=shape, + stride_order=source.stride_order, + name=name, + ) - def _last_axis_spec(self, desc: JaxTensorDesc, alignment: int) -> Any: + def _last_axis_spec( + self, + desc: JaxTensorDesc, + alignment: int, + mode: tuple[int, ...], + ) -> Any: return self._to_tensor_spec( desc, + mode=mode, divisibility=(None,) * (desc.ndim - 1) + (alignment,), ) @@ -103,20 +137,77 @@ def __init__( loss_coeff: float = 1.0, block_I: int = DEFAULT_BLOCK_I, topk_indices_global: bool = False, + q_layout: str | None = None, + w_layout: str | None = None, + k_layout: str | None = None, + score_layout: str | None = None, target_compute_capability: int | None = None, ) -> None: - self.iq_desc = self._to_tensor_desc(sample_index_q, "sample_index_q") - self.w_desc = self._to_tensor_desc(sample_weights, "sample_weights") - self.ik_desc = self._to_tensor_desc(sample_index_k, "sample_index_k") - self.attn_desc = self._to_tensor_desc(sample_attn_score, "sample_attn_score") - self.idx_score_desc = self._to_tensor_desc(sample_index_score, "sample_index_score") - self.topk_desc = self._to_tensor_desc(sample_topk_indices, "sample_topk_indices") + is_thd = len(tuple(sample_index_q.shape)) == 3 + if is_thd: + q_defaults = ("THD", "THD", ("THD",)) + w_defaults = ("TH", "TH", ("TH",)) + k_defaults = ("TD", "TD", ("TD",)) + score_defaults = ("TK", "TK", ("TK",)) + else: + q_defaults = ("BSHD", "BSHD", ("BSHD", "SBHD")) + w_defaults = ("BSH", "BSH", ("BSH", "SBH")) + k_defaults = ("BSD", "BSD", ("BSD", "SBD")) + score_defaults = ("BSK", "BSK", ("BSK", "SBK")) + self.q_layout, self.q_mode = _resolve_layout( + "q_layout", + q_layout, + default=q_defaults[0], + kernel_axes=q_defaults[1], + supported=q_defaults[2], + ) + self.w_layout, self.w_mode = _resolve_layout( + "w_layout", + w_layout, + default=w_defaults[0], + kernel_axes=w_defaults[1], + supported=w_defaults[2], + ) + self.k_layout, self.k_mode = _resolve_layout( + "k_layout", + k_layout, + default=k_defaults[0], + kernel_axes=k_defaults[1], + supported=k_defaults[2], + ) + self.score_layout, self.score_mode = _resolve_layout( + "score_layout", + score_layout, + default=score_defaults[0], + kernel_axes=score_defaults[1], + supported=score_defaults[2], + ) + + self.iq_desc = self._to_tensor_desc( + sample_index_q, "sample_index_q", mode=self.q_mode + ) + self.w_desc = self._to_tensor_desc( + sample_weights, "sample_weights", mode=self.w_mode + ) + self.ik_desc = self._to_tensor_desc( + sample_index_k, "sample_index_k", mode=self.k_mode + ) + self.attn_desc = self._to_tensor_desc( + sample_attn_score, "sample_attn_score", mode=self.score_mode + ) + self.idx_score_desc = self._to_tensor_desc( + sample_index_score, "sample_index_score", mode=self.score_mode + ) + self.topk_desc = self._to_tensor_desc( + sample_topk_indices, "sample_topk_indices", mode=self.score_mode + ) self.diq_desc = self._output_desc( sample_d_index_q, source=self.iq_desc, cudnn_dtype=data_type.BFLOAT16, shape=self.iq_desc.shape, name="sample_d_index_q", + mode=self.q_mode, ) self.dw_desc = self._output_desc( sample_d_weights, @@ -124,6 +215,7 @@ def __init__( cudnn_dtype=data_type.BFLOAT16, shape=self.w_desc.shape, name="sample_d_weights", + mode=self.w_mode, ) self.dik_desc = self._output_desc( sample_d_index_k, @@ -131,6 +223,7 @@ def __init__( cudnn_dtype=data_type.BFLOAT16, shape=self.ik_desc.shape, name="sample_d_index_k", + mode=self.k_mode, ) self.grad_loss_desc = self.iq_desc.compact_like( cudnn_dtype=data_type.FLOAT, @@ -171,49 +264,64 @@ def __call__( grad_loss: Any = 1.0, ) -> TupleDict: self.check_support() - for value, expected in ( - (index_q, self.iq_desc), - (weights, self.w_desc), - (index_k, self.ik_desc), - (attn_score, self.attn_desc), - (index_score, self.idx_score_desc), - (topk_indices, self.topk_desc), + for value, expected, mode in ( + (index_q, self.iq_desc, self.q_mode), + (weights, self.w_desc, self.w_mode), + (index_k, self.ik_desc, self.k_mode), + (attn_score, self.attn_desc, self.score_mode), + (index_score, self.idx_score_desc, self.score_mode), + (topk_indices, self.topk_desc, self.score_mode), ): - self._check_tensor_signature(value, expected) + self._check_tensor_signature(value, expected, mode=mode) grad_loss = _grad_loss_operand(grad_loss) self._check_tensor_signature(grad_loss, self.grad_loss_desc) dk_accum_desc = self.ik_desc.compact_like( cudnn_dtype=data_type.FLOAT, shape=self.ik_desc.shape, + stride_order=self.ik_desc.stride_order, name="d_index_k_accum", init_value=0.0, ) grad_signal_desc = self.attn_desc.compact_like( cudnn_dtype=data_type.FLOAT, shape=self.attn_desc.shape, + stride_order=self.attn_desc.stride_order, name="grad_signal", ) d_index_q, d_weights, d_index_k_accum = self._call_kernel( - (index_q, weights, index_k, attn_score, index_score, topk_indices, grad_loss), + ( + index_q, + weights, + index_k, + attn_score, + index_score, + topk_indices, + grad_loss, + ), launch=self._launch_kernel, output_descs=(self.diq_desc, self.dw_desc, dk_accum_desc), workspace_descs=(grad_signal_desc,), input_spec=( - self._last_axis_spec(self.iq_desc, self._op.head_dim), - None, - self._last_axis_spec(self.ik_desc, self._op.head_dim), - None, - None, - None, - None, + self._last_axis_spec(self.iq_desc, self._op.head_dim, self.q_mode), + self._to_tensor_spec(self.w_desc, mode=self.w_mode), + self._last_axis_spec(self.ik_desc, self._op.head_dim, self.k_mode), + self._to_tensor_spec(self.attn_desc, mode=self.score_mode), + self._to_tensor_spec(self.idx_score_desc, mode=self.score_mode), + self._to_tensor_spec(self.topk_desc, mode=self.score_mode), + self._to_tensor_spec(self.grad_loss_desc), ), output_spec=( - self._last_axis_spec(self.diq_desc, self._op.head_dim), - None, - self._last_axis_spec(dk_accum_desc, self._op.head_dim), + self._last_axis_spec(self.diq_desc, self._op.head_dim, self.q_mode), + self._to_tensor_spec(self.dw_desc, mode=self.w_mode), + self._last_axis_spec(dk_accum_desc, self._op.head_dim, self.k_mode), + ), + workspace_spec=( + self._to_tensor_spec(grad_signal_desc, mode=self.score_mode), + ), + compile_options=compile_options_for_target( + self.compute_capability, "--opt-level 3" ), - compile_options=compile_options_for_target(self.compute_capability, "--opt-level 3"), ) return TupleDict( d_index_q=d_index_q, @@ -358,21 +466,92 @@ def __init__( loss_coeff: float = 1.0, block_I: int = DEFAULT_BLOCK_I, ratio: int = 1, + q_layout: str | None = None, + w_layout: str | None = None, + k_layout: str | None = None, + score_layout: str | None = None, + denom_layout: str | None = None, target_compute_capability: int | None = None, ) -> None: - self.iq_desc = self._to_tensor_desc(sample_index_q, "sample_index_q") - self.w_desc = self._to_tensor_desc(sample_weights, "sample_weights") - self.ik_desc = self._to_tensor_desc(sample_index_k, "sample_index_k") - self.attn_desc = self._to_tensor_desc(sample_attn_score, "sample_attn_score") - self.attn_denom_desc = self._to_tensor_desc(sample_attn_l1norm, "sample_attn_l1norm") - self.idx_score_desc = self._to_tensor_desc(sample_index_score, "sample_index_score") - self.idx_lse_desc = self._to_tensor_desc(sample_index_lse, "sample_index_lse") + is_thd = len(tuple(sample_index_q.shape)) == 3 + if is_thd: + q_defaults = ("THD", "THD", ("THD",)) + w_defaults = ("TH", "TH", ("TH",)) + k_defaults = ("TD", "TD", ("TD",)) + score_defaults = ("TK", "TK", ("TK",)) + denom_defaults = ("T", "T", ("T",)) + else: + q_defaults = ("BSHD", "BSHD", ("BSHD", "SBHD")) + w_defaults = ("BSH", "BSH", ("BSH", "SBH")) + k_defaults = ("BSD", "BSD", ("BSD", "SBD")) + score_defaults = ("BSK", "BSK", ("BSK", "SBK")) + denom_defaults = ("BS", "BS", ("BS", "SB")) + self.q_layout, self.q_mode = _resolve_layout( + "q_layout", + q_layout, + default=q_defaults[0], + kernel_axes=q_defaults[1], + supported=q_defaults[2], + ) + self.w_layout, self.w_mode = _resolve_layout( + "w_layout", + w_layout, + default=w_defaults[0], + kernel_axes=w_defaults[1], + supported=w_defaults[2], + ) + self.k_layout, self.k_mode = _resolve_layout( + "k_layout", + k_layout, + default=k_defaults[0], + kernel_axes=k_defaults[1], + supported=k_defaults[2], + ) + self.score_layout, self.score_mode = _resolve_layout( + "score_layout", + score_layout, + default=score_defaults[0], + kernel_axes=score_defaults[1], + supported=score_defaults[2], + ) + self.denom_layout, self.denom_mode = _resolve_layout( + "denom_layout", + denom_layout, + default=denom_defaults[0], + kernel_axes=denom_defaults[1], + supported=denom_defaults[2], + ) + + self.iq_desc = self._to_tensor_desc( + sample_index_q, "sample_index_q", mode=self.q_mode + ) + self.w_desc = self._to_tensor_desc( + sample_weights, "sample_weights", mode=self.w_mode + ) + self.ik_desc = self._to_tensor_desc( + sample_index_k, "sample_index_k", mode=self.k_mode + ) + self.attn_desc = self._to_tensor_desc( + sample_attn_score, "sample_attn_score", mode=self.score_mode + ) + self.attn_denom_desc = self._to_tensor_desc( + sample_attn_l1norm, + "sample_attn_l1norm", + mode=self.denom_mode, + ) + self.idx_score_desc = self._to_tensor_desc( + sample_index_score, "sample_index_score", mode=self.score_mode + ) + self.idx_lse_desc = self._to_tensor_desc( + sample_index_lse, "sample_index_lse", mode=self.denom_mode + ) self.diq_desc = self._output_desc( sample_d_index_q, source=self.iq_desc, cudnn_dtype=data_type.BFLOAT16, shape=self.iq_desc.shape, name="sample_d_index_q", + mode=self.q_mode, ) self.dw_desc = self._output_desc( sample_d_weights, @@ -380,6 +559,7 @@ def __init__( cudnn_dtype=data_type.BFLOAT16, shape=self.w_desc.shape, name="sample_d_weights", + mode=self.w_mode, ) self.dik_desc = self._output_desc( sample_d_index_k, @@ -387,15 +567,30 @@ def __init__( cudnn_dtype=data_type.BFLOAT16, shape=self.ik_desc.shape, name="sample_d_index_k", + mode=self.k_mode, ) self.grad_loss_desc = self.iq_desc.compact_like( cudnn_dtype=data_type.FLOAT, shape=(1,), name="grad_loss", ) - self.cuq_desc = None if sample_cu_seqlens_q is None else self._to_tensor_desc(sample_cu_seqlens_q, "sample_cu_seqlens_q") - self.cuk_desc = None if sample_cu_seqlens_k is None else self._to_tensor_desc(sample_cu_seqlens_k, "sample_cu_seqlens_k") - self.q_offsets_desc = None if sample_q_causal_offsets is None else self._to_tensor_desc(sample_q_causal_offsets, "sample_q_causal_offsets") + self.cuq_desc = ( + None + if sample_cu_seqlens_q is None + else self._to_tensor_desc(sample_cu_seqlens_q, "sample_cu_seqlens_q") + ) + self.cuk_desc = ( + None + if sample_cu_seqlens_k is None + else self._to_tensor_desc(sample_cu_seqlens_k, "sample_cu_seqlens_k") + ) + self.q_offsets_desc = ( + None + if sample_q_causal_offsets is None + else self._to_tensor_desc( + sample_q_causal_offsets, "sample_q_causal_offsets" + ) + ) self._op = DenseIndexerBackwardOp( index_q=self.iq_desc, weights=self.w_desc, @@ -440,16 +635,16 @@ def __call__( q_causal_offsets: Any | None = None, ) -> TupleDict: self.check_support() - for value, expected in ( - (index_q, self.iq_desc), - (weights, self.w_desc), - (index_k, self.ik_desc), - (attn_score, self.attn_desc), - (attn_l1norm, self.attn_denom_desc), - (index_score, self.idx_score_desc), - (index_lse, self.idx_lse_desc), + for value, expected, mode in ( + (index_q, self.iq_desc, self.q_mode), + (weights, self.w_desc, self.w_mode), + (index_k, self.ik_desc, self.k_mode), + (attn_score, self.attn_desc, self.score_mode), + (attn_l1norm, self.attn_denom_desc, self.denom_mode), + (index_score, self.idx_score_desc, self.score_mode), + (index_lse, self.idx_lse_desc, self.denom_mode), ): - self._check_tensor_signature(value, expected) + self._check_tensor_signature(value, expected, mode=mode) for value, expected, name in ( (cu_seqlens_q, self.cuq_desc, "cu_seqlens_q"), (cu_seqlens_k, self.cuk_desc, "cu_seqlens_k"), @@ -462,19 +657,30 @@ def __call__( grad_loss = _grad_loss_operand(grad_loss) self._check_tensor_signature(grad_loss, self.grad_loss_desc) - inputs = (index_q, weights, index_k, attn_score, attn_l1norm, index_score, index_lse, grad_loss) + inputs = ( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + grad_loss, + ) for optional in (cu_seqlens_q, cu_seqlens_k, q_causal_offsets): if optional is not None: inputs += (optional,) dk_accum_desc = self.ik_desc.compact_like( cudnn_dtype=data_type.FLOAT, shape=self.ik_desc.shape, + stride_order=self.ik_desc.stride_order, name="d_index_k_accum", init_value=0.0, ) grad_signal_desc = self.idx_score_desc.compact_like( cudnn_dtype=data_type.FLOAT, shape=self.idx_score_desc.shape, + stride_order=self.idx_score_desc.stride_order, name="grad_signal", ) d_index_q, d_weights, d_index_k_accum = self._call_kernel( @@ -483,17 +689,35 @@ def __call__( output_descs=(self.diq_desc, self.dw_desc, dk_accum_desc), workspace_descs=(grad_signal_desc,), input_spec=( - self._last_axis_spec(self.iq_desc, self._op.head_dim), - None, - self._last_axis_spec(self.ik_desc, self._op.head_dim), + self._last_axis_spec(self.iq_desc, self._op.head_dim, self.q_mode), + self._to_tensor_spec(self.w_desc, mode=self.w_mode), + self._last_axis_spec(self.ik_desc, self._op.head_dim, self.k_mode), + self._to_tensor_spec(self.attn_desc, mode=self.score_mode), + self._to_tensor_spec(self.attn_denom_desc, mode=self.denom_mode), + self._to_tensor_spec(self.idx_score_desc, mode=self.score_mode), + self._to_tensor_spec(self.idx_lse_desc, mode=self.denom_mode), + self._to_tensor_spec(self.grad_loss_desc), ) - + (None,) * (len(inputs) - 3), + + tuple( + self._to_tensor_spec(desc) + for desc in ( + self.cuq_desc, + self.cuk_desc, + self.q_offsets_desc, + ) + if desc is not None + ), output_spec=( - self._last_axis_spec(self.diq_desc, self._op.head_dim), - None, - self._last_axis_spec(dk_accum_desc, self._op.head_dim), + self._last_axis_spec(self.diq_desc, self._op.head_dim, self.q_mode), + self._to_tensor_spec(self.dw_desc, mode=self.w_mode), + self._last_axis_spec(dk_accum_desc, self._op.head_dim, self.k_mode), + ), + workspace_spec=( + self._to_tensor_spec(grad_signal_desc, mode=self.score_mode), + ), + compile_options=compile_options_for_target( + self.compute_capability, "--opt-level 3" ), - compile_options=compile_options_for_target(self.compute_capability, "--opt-level 3"), ) return TupleDict( d_index_q=d_index_q, @@ -505,7 +729,17 @@ def _launch_kernel(self, stream, *arguments) -> None: import cutlass *inputs, d_index_q, d_weights, d_index_k_accum, grad_signal = arguments - index_q, weights, index_k, attn_score, attn_l1norm, index_score, index_lse, grad_loss, *optional = inputs + ( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + grad_loss, + *optional, + ) = inputs cursor = 0 cu_seqlens_q = cu_seqlens_k = q_causal_offsets = None if self._op.is_thd: @@ -517,7 +751,10 @@ def _launch_kernel(self, stream, *arguments) -> None: grad_scale = self.loss_coeff / max(self._op.normalization_tokens, 1) if self._architecture_family == 100: - from .dense_indexer_backward_sm100 import DenseIndexerBackward2QGemmSm100, ScoreGradDense + from .dense_indexer_backward_sm100 import ( + DenseIndexerBackward2QGemmSm100, + ScoreGradDense, + ) ScoreGradDense(ratio=self._op.ratio, block_I=self._op.block_i)( index_score, @@ -601,7 +838,19 @@ def _launch_kernel(self, stream, *arguments) -> None: ) -@jax.jit(static_argnames=("sm_scale", "loss_coeff", "block_I", "topk_indices_global", "target_compute_capability")) +@jax.jit( + static_argnames=( + "sm_scale", + "loss_coeff", + "block_I", + "topk_indices_global", + "q_layout", + "w_layout", + "k_layout", + "score_layout", + "target_compute_capability", + ) +) def indexer_backward_wrapper( index_q: Any, weights: Any, @@ -614,6 +863,10 @@ def indexer_backward_wrapper( grad_loss: Any = 1.0, block_I: int = DEFAULT_BLOCK_I, topk_indices_global: bool = False, + q_layout: str | None = None, + w_layout: str | None = None, + k_layout: str | None = None, + score_layout: str | None = None, target_compute_capability: int | None = None, ) -> TupleDict: """Compute sparse indexer gradients for fixed BSHD or packed THD inputs. @@ -621,13 +874,20 @@ def indexer_backward_wrapper( Packed THD callers must pass ``topk_indices_global=True``. """ - samples = tuple(jax.ShapeDtypeStruct(value.shape, value.dtype) for value in (index_q, weights, index_k, attn_score, index_score, topk_indices)) + samples = tuple( + jax.ShapeDtypeStruct(value.shape, value.dtype) + for value in (index_q, weights, index_k, attn_score, index_score, topk_indices) + ) return IndexerBackward( *samples, sm_scale=sm_scale, loss_coeff=loss_coeff, block_I=block_I, topk_indices_global=topk_indices_global, + q_layout=q_layout, + w_layout=w_layout, + k_layout=k_layout, + score_layout=score_layout, target_compute_capability=target_compute_capability, )(index_q, weights, index_k, attn_score, index_score, topk_indices, grad_loss) @@ -640,6 +900,11 @@ def indexer_backward_wrapper( "ratio", "max_seqlen_q", "max_seqlen_k", + "q_layout", + "w_layout", + "k_layout", + "score_layout", + "denom_layout", "target_compute_capability", ) ) @@ -661,13 +926,30 @@ def dense_indexer_backward_wrapper( max_seqlen_q: int | None = None, max_seqlen_k: int | None = None, q_causal_offsets: Any | None = None, + q_layout: str | None = None, + w_layout: str | None = None, + k_layout: str | None = None, + score_layout: str | None = None, + denom_layout: str | None = None, target_compute_capability: int | None = None, ) -> TupleDict: core = (index_q, weights, index_k, attn_score, attn_l1norm, index_score, index_lse) samples = tuple(jax.ShapeDtypeStruct(value.shape, value.dtype) for value in core) - sample_cuq = None if cu_seqlens_q is None else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype) - sample_cuk = None if cu_seqlens_k is None else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype) - sample_offsets = None if q_causal_offsets is None else jax.ShapeDtypeStruct(q_causal_offsets.shape, q_causal_offsets.dtype) + sample_cuq = ( + None + if cu_seqlens_q is None + else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype) + ) + sample_cuk = ( + None + if cu_seqlens_k is None + else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype) + ) + sample_offsets = ( + None + if q_causal_offsets is None + else jax.ShapeDtypeStruct(q_causal_offsets.shape, q_causal_offsets.dtype) + ) return DenseIndexerBackward( *samples, sample_cu_seqlens_q=sample_cuq, @@ -679,6 +961,11 @@ def dense_indexer_backward_wrapper( loss_coeff=loss_coeff, block_I=block_I, ratio=ratio, + q_layout=q_layout, + w_layout=w_layout, + k_layout=k_layout, + score_layout=score_layout, + denom_layout=denom_layout, target_compute_capability=target_compute_capability, )( *core, diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/op.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/op.py index b19120c71..49b6e54cd 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/op.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/op.py @@ -17,7 +17,9 @@ MIN_HEADS = 64 -def _require_desc(name: str, desc: TensorDesc[Any] | None, *, optional: bool = False) -> None: +def _require_desc( + name: str, desc: TensorDesc[Any] | None, *, optional: bool = False +) -> None: if desc is None and optional: return if not isinstance(desc, TensorDesc): @@ -30,7 +32,9 @@ def _require_rank(desc: TensorDesc[Any], rank: int, name: str) -> None: raise ValueError(f"{name} must have rank {rank}, got shape {desc.shape}") -def _require_dtype(desc: TensorDesc[Any], expected: data_type | tuple[data_type, ...], name: str) -> None: +def _require_dtype( + desc: TensorDesc[Any], expected: data_type | tuple[data_type, ...], name: str +) -> None: expected_values = expected if isinstance(expected, tuple) else (expected,) if desc.cudnn_dtype not in expected_values: expected_text = " or ".join(value.name.lower() for value in expected_values) @@ -42,9 +46,22 @@ def _require_shape(desc: TensorDesc[Any], expected: tuple[int, ...], name: str) raise ValueError(f"{name} must have shape {expected}, got {desc.shape}") -def _require_row_major(desc: TensorDesc[Any]) -> None: - if not desc.is_compact(tuple(reversed(range(desc.ndim)))): - raise ValueError(f"{desc.name or 'tensor'} must be row-major contiguous, got stride {desc.stride}") +def _require_compact(desc: TensorDesc[Any]) -> None: + if not desc.is_compact(): + raise ValueError( + f"{desc.name or 'tensor'} must be compact, got stride {desc.stride}" + ) + + +def _require_contiguous_tail(desc: TensorDesc[Any], rank: int) -> None: + expected_stride = 1 + for axis in range(desc.ndim - 1, desc.ndim - rank - 1, -1): + if desc.shape[axis] != 1 and desc.stride[axis] != expected_stride: + raise ValueError( + f"{desc.name or 'tensor'} must have its final {rank} axes contiguous, " + f"got shape {desc.shape} and stride {desc.stride}" + ) + expected_stride *= max(desc.shape[axis], 1) def _resolve_scale(value: float, name: str) -> float: @@ -114,11 +131,15 @@ def __init__( def check_support(self) -> bool: self.is_thd = None - self.batch = self.seqlen_q = self.seqlen_k = self.heads = self.head_dim = self.topk = None + self.batch = self.seqlen_q = self.seqlen_k = self.heads = self.head_dim = ( + self.topk + ) = None self.sm_scale = None if self.index_q.ndim not in (3, 4): - raise ValueError(f"index_q must use BSHD rank 4 or packed THD rank 3, got shape {self.index_q.shape}") + raise ValueError( + f"index_q must use BSHD rank 4 or packed THD rank 3, got shape {self.index_q.shape}" + ) is_thd = self.index_q.ndim == 3 q_rank = 3 if is_thd else 4 auxiliary_rank = q_rank - 1 @@ -135,11 +156,17 @@ def check_support(self) -> bool: ): _require_rank(desc, rank, name) - for desc, name in ((self.index_q, "index_q"), (self.weights, "weights"), (self.index_k, "index_k")): + for desc, name in ( + (self.index_q, "index_q"), + (self.weights, "weights"), + (self.index_k, "index_k"), + ): _require_dtype(desc, data_type.BFLOAT16, name) _require_dtype(self.d_index_q, data_type.BFLOAT16, "d_index_q") _require_dtype(self.d_weights, data_type.BFLOAT16, "d_weights") - _require_dtype(self.d_index_k, (data_type.BFLOAT16, data_type.FLOAT), "d_index_k") + _require_dtype( + self.d_index_k, (data_type.BFLOAT16, data_type.FLOAT), "d_index_k" + ) _require_dtype(self.attn_score, data_type.FLOAT, "attn_score") _require_dtype(self.index_score, data_type.FLOAT, "index_score") _require_dtype(self.topk_indices, data_type.INT32, "topk_indices") @@ -149,24 +176,34 @@ def check_support(self) -> bool: seqlen_k, head_dim_k = self.index_k.shape batch = batch_k = 1 if not self.topk_indices_global: - raise ValueError("Packed THD IndexerBackward requires topk_indices_global=True") + raise ValueError( + "Packed THD IndexerBackward requires topk_indices_global=True" + ) else: batch, seqlen_q, heads, head_dim = self.index_q.shape batch_k, seqlen_k, head_dim_k = self.index_k.shape topk = self.topk_indices.shape[-1] dimensions = (batch, seqlen_q, seqlen_k, heads, head_dim, topk) if any(value <= 0 for value in dimensions): - raise ValueError(f"Indexer-backward dimensions must be positive, got {dimensions}") + raise ValueError( + f"Indexer-backward dimensions must be positive, got {dimensions}" + ) if batch_k != batch or head_dim_k != head_dim: raise ValueError("index_k batch and head dimensions must match index_q") if heads < MIN_HEADS: - raise ValueError(f"IndexerBackward requires heads >= {MIN_HEADS}, got {heads}") + raise ValueError( + f"IndexerBackward requires heads >= {MIN_HEADS}, got {heads}" + ) if head_dim != SUPPORTED_HEAD_DIM: - raise ValueError(f"IndexerBackward requires head_dim={SUPPORTED_HEAD_DIM}, got {head_dim}") + raise ValueError( + f"IndexerBackward requires head_dim={SUPPORTED_HEAD_DIM}, got {head_dim}" + ) if self.block_i != DEFAULT_BLOCK_I: raise ValueError(f"block_i must be {DEFAULT_BLOCK_I}, got {self.block_i}") if topk % self.block_i != 0: - raise ValueError(f"topk ({topk}) must be divisible by block_i ({self.block_i})") + raise ValueError( + f"topk ({topk}) must be divisible by block_i ({self.block_i})" + ) weights_shape = (seqlen_q, heads) if is_thd else (batch, seqlen_q, heads) score_shape = (seqlen_q, topk) if is_thd else (batch, seqlen_q, topk) @@ -188,7 +225,19 @@ def check_support(self) -> bool: self.index_score, self.topk_indices, ): - _require_row_major(desc) + _require_compact(desc) + for desc in (self.index_q, self.d_index_q): + _require_contiguous_tail(desc, 2) + for desc in ( + self.weights, + self.index_k, + self.d_weights, + self.d_index_k, + self.attn_score, + self.index_score, + self.topk_indices, + ): + _require_contiguous_tail(desc, 1) self.is_thd = is_thd self.batch = batch @@ -240,7 +289,11 @@ def __init__( ) for name, desc in descriptors: _require_desc(name, desc) - for name, desc in (("cu_seqlens_q", cu_seqlens_q), ("cu_seqlens_k", cu_seqlens_k), ("q_causal_offsets", q_causal_offsets)): + for name, desc in ( + ("cu_seqlens_q", cu_seqlens_q), + ("cu_seqlens_k", cu_seqlens_k), + ("q_causal_offsets", q_causal_offsets), + ): _require_desc(name, desc, optional=True) self.index_q = index_q @@ -274,7 +327,9 @@ def __init__( def check_support(self) -> bool: self.is_thd = None - self.batch = self.normalization_tokens = self.total_k = self.heads = self.head_dim = None + self.batch = self.normalization_tokens = self.total_k = self.heads = ( + self.head_dim + ) = None self.max_seqlen_q = self.max_seqlen_k = None self.sm_scale = None @@ -282,15 +337,25 @@ def check_support(self) -> bool: raise ValueError("cu_seqlens_q and cu_seqlens_k must be provided together") is_thd = self.cu_seqlens_q is not None if is_thd: - batch, normalization_tokens, total_k, heads, head_dim, max_q, max_k = self._check_thd() + batch, normalization_tokens, total_k, heads, head_dim, max_q, max_k = ( + self._check_thd() + ) else: - batch, normalization_tokens, total_k, heads, head_dim, max_q, max_k = self._check_bshd() + batch, normalization_tokens, total_k, heads, head_dim, max_q, max_k = ( + self._check_bshd() + ) - for desc, name in ((self.index_q, "index_q"), (self.weights, "weights"), (self.index_k, "index_k")): + for desc, name in ( + (self.index_q, "index_q"), + (self.weights, "weights"), + (self.index_k, "index_k"), + ): _require_dtype(desc, data_type.BFLOAT16, name) _require_dtype(self.d_index_q, data_type.BFLOAT16, "d_index_q") _require_dtype(self.d_weights, data_type.BFLOAT16, "d_weights") - _require_dtype(self.d_index_k, (data_type.BFLOAT16, data_type.FLOAT), "d_index_k") + _require_dtype( + self.d_index_k, (data_type.BFLOAT16, data_type.FLOAT), "d_index_k" + ) for desc, name in ( (self.attn_score, "attn_score"), (self.attn_l1norm, "attn_l1norm"), @@ -298,14 +363,22 @@ def check_support(self) -> bool: (self.index_lse, "index_lse"), ): _require_dtype(desc, data_type.FLOAT, name) - for desc, name in ((self.cu_seqlens_q, "cu_seqlens_q"), (self.cu_seqlens_k, "cu_seqlens_k"), (self.q_causal_offsets, "q_causal_offsets")): + for desc, name in ( + (self.cu_seqlens_q, "cu_seqlens_q"), + (self.cu_seqlens_k, "cu_seqlens_k"), + (self.q_causal_offsets, "q_causal_offsets"), + ): if desc is not None: _require_dtype(desc, data_type.INT32, name) if heads < MIN_HEADS: - raise ValueError(f"DenseIndexerBackward requires heads >= {MIN_HEADS}, got {heads}") + raise ValueError( + f"DenseIndexerBackward requires heads >= {MIN_HEADS}, got {heads}" + ) if head_dim != SUPPORTED_HEAD_DIM: - raise ValueError(f"DenseIndexerBackward requires head_dim={SUPPORTED_HEAD_DIM}, got {head_dim}") + raise ValueError( + f"DenseIndexerBackward requires head_dim={SUPPORTED_HEAD_DIM}, got {head_dim}" + ) if self.block_i != DEFAULT_BLOCK_I: raise ValueError(f"block_i must be {DEFAULT_BLOCK_I}, got {self.block_i}") if self.ratio < 1: @@ -327,7 +400,18 @@ def check_support(self) -> bool: self.q_causal_offsets, ): if desc is not None: - _require_row_major(desc) + _require_compact(desc) + for desc in (self.index_q, self.d_index_q): + _require_contiguous_tail(desc, 2) + for desc in ( + self.weights, + self.index_k, + self.d_weights, + self.d_index_k, + self.attn_score, + self.index_score, + ): + _require_contiguous_tail(desc, 1) self.is_thd = is_thd self.batch = batch @@ -375,10 +459,20 @@ def _check_bshd(self) -> tuple[int, int, int, int, int, int, int]: _require_rank(self.q_causal_offsets, 1, "q_causal_offsets") _require_shape(self.q_causal_offsets, (batch,), "q_causal_offsets") - max_q = seqlen_q if self.requested_max_seqlen_q is None else int(self.requested_max_seqlen_q) - max_k = seqlen_k if self.requested_max_seqlen_k is None else int(self.requested_max_seqlen_k) + max_q = ( + seqlen_q + if self.requested_max_seqlen_q is None + else int(self.requested_max_seqlen_q) + ) + max_k = ( + seqlen_k + if self.requested_max_seqlen_k is None + else int(self.requested_max_seqlen_k) + ) if (max_q, max_k) != (seqlen_q, seqlen_k): - raise ValueError(f"BSHD max_seqlen_q/k must match tensor extents {(seqlen_q, seqlen_k)}, got {(max_q, max_k)}") + raise ValueError( + f"BSHD max_seqlen_q/k must match tensor extents {(seqlen_q, seqlen_k)}, got {(max_q, max_k)}" + ) return batch, batch * seqlen_q, batch * seqlen_k, heads, head_dim, max_q, max_k def _check_thd(self) -> tuple[int, int, int, int, int, int, int]: @@ -404,11 +498,15 @@ def _check_thd(self) -> tuple[int, int, int, int, int, int, int]: if min(total_q, total_k, batch, heads, head_dim) <= 0: raise ValueError("Dense THD indexer-backward dimensions must be positive") if self.cu_seqlens_k.shape != (batch + 1,): - raise ValueError("cu_seqlens_q and cu_seqlens_k must describe the same batch") + raise ValueError( + "cu_seqlens_q and cu_seqlens_k must describe the same batch" + ) if head_dim_k != head_dim: raise ValueError("index_k head dimension must match index_q") if self.requested_max_seqlen_q is None or self.requested_max_seqlen_k is None: - raise ValueError("THD dense indexer backward requires max_seqlen_q and max_seqlen_k") + raise ValueError( + "THD dense indexer backward requires max_seqlen_q and max_seqlen_k" + ) max_q = int(self.requested_max_seqlen_q) max_k = int(self.requested_max_seqlen_k) if max_q <= 0 or max_k <= 0: diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py b/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py index 5c689a6fa..c46053e20 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py @@ -13,6 +13,7 @@ from ... import data_type from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..._jax.layout import mode_from_layout, to_public_axes from .config import DenseScoreKernelConfig, SparseScoreKernelConfig from .op import ( DenseScoreRecomputeOp, @@ -29,18 +30,46 @@ def _compile_options(target_compute_capability: int) -> str: return compile_options_for_target(target_compute_capability) -def _optional_desc(api: JaxApiBase, value: Any | None, name: str) -> JaxTensorDesc | None: - return None if value is None else api._to_tensor_desc(value, name) - - -def _check_optional_signature(api: JaxApiBase, value: Any | None, expected: JaxTensorDesc | None, name: str) -> None: +def _resolve_layout( + name: str, + layout: str | None, + *, + default: str, + kernel_axes: str, + supported: tuple[str, ...], +) -> tuple[str, tuple[int, ...]]: + layout = default if layout is None else layout + if layout not in supported: + choices = ", ".join(repr(value) for value in supported) + raise ValueError(f"{name} must be one of ({choices}), got {layout!r}") + return layout, mode_from_layout(layout, kernel_axes=kernel_axes) + + +def _optional_desc( + api: JaxApiBase, + value: Any | None, + name: str, + *, + mode: tuple[int, ...] | None = None, +) -> JaxTensorDesc | None: + return None if value is None else api._to_tensor_desc(value, name, mode=mode) + + +def _check_optional_signature( + api: JaxApiBase, + value: Any | None, + expected: JaxTensorDesc | None, + name: str, + *, + mode: tuple[int, ...] | None = None, +) -> None: if expected is None: if value is not None: raise ValueError(f"{name} was not part of the sample signature") return if value is None: raise ValueError(f"{name} is required by the sample signature") - api._check_tensor_signature(value, expected) + api._check_tensor_signature(value, expected, mode=mode) class _SparseScoreRecompute(JaxApiBase): @@ -57,14 +86,72 @@ def __init__( sample_topk_length: Any | None, qhead_per_kv_head: int | None, topk_indices_global: bool, + q_layout: str | None, + k_layout: str | None, + per_head_layout: str | None, + score_layout: str | None, + output_layout: str | None, target_compute_capability: int | None, ) -> None: - self.q_desc = self._to_tensor_desc(sample_q, "sample_q") - self.k_desc = self._to_tensor_desc(sample_k, "sample_k") - self.per_head_desc = self._to_tensor_desc(sample_per_head, "sample_per_head") - self.topk_indices_desc = self._to_tensor_desc(sample_topk_indices, "sample_topk_indices") - self.topk_length_desc = _optional_desc(self, sample_topk_length, "sample_topk_length") - self.out_desc = self._default_output_desc() if sample_out is None else self._to_tensor_desc(sample_out, "sample_out") + self.q_layout, self.q_mode = _resolve_layout( + "q_layout", + q_layout, + default="BSHD", + kernel_axes="BSHD", + supported=("BSHD", "SBHD"), + ) + self.k_layout, self.k_mode = _resolve_layout( + "k_layout", + k_layout, + default="BSD", + kernel_axes="BSD", + supported=("BSD", "SBD"), + ) + self.per_head_layout, self.per_head_mode = _resolve_layout( + "per_head_layout", + per_head_layout, + default="BSH", + kernel_axes="BSH", + supported=("BSH", "SBH"), + ) + self.score_layout, self.score_mode = _resolve_layout( + "score_layout", + score_layout, + default="BSK", + kernel_axes="BSK", + supported=("BSK", "SBK"), + ) + self.output_layout, self.output_mode = _resolve_layout( + "output_layout", + output_layout, + default="BSK", + kernel_axes="BSK", + supported=("BSK", "SBK"), + ) + self.length_layout = self.score_layout.replace("K", "") + self.length_mode = mode_from_layout(self.length_layout, kernel_axes="BS") + + self.q_desc = self._to_tensor_desc(sample_q, "sample_q", mode=self.q_mode) + self.k_desc = self._to_tensor_desc(sample_k, "sample_k", mode=self.k_mode) + self.per_head_desc = self._to_tensor_desc( + sample_per_head, "sample_per_head", mode=self.per_head_mode + ) + self.topk_indices_desc = self._to_tensor_desc( + sample_topk_indices, + "sample_topk_indices", + mode=self.score_mode, + ) + self.topk_length_desc = _optional_desc( + self, + sample_topk_length, + "sample_topk_length", + mode=self.length_mode, + ) + self.out_desc = ( + self._default_output_desc() + if sample_out is None + else self._to_tensor_desc(sample_out, "sample_out", mode=self.output_mode) + ) self.target_compute_capability = self._resolve_compute_capability( target_compute_capability, SUPPORTED_COMPUTE_CAPABILITIES, @@ -86,37 +173,75 @@ def __init__( def _default_output_desc(self) -> JaxTensorDesc: if self.topk_indices_desc.ndim != 3: - raise ValueError(f"topk_indices must have rank 3, got shape {self.topk_indices_desc.shape}") - return self.topk_indices_desc.compact_like( - cudnn_dtype=data_type.FLOAT, - shape=self.topk_indices_desc.shape, - name="sample_out", + raise ValueError( + f"topk_indices must have rank 3, got shape {self.topk_indices_desc.shape}" + ) + return self._to_tensor_desc( + jax.ShapeDtypeStruct( + to_public_axes(self.topk_indices_desc.shape, self.output_mode), + jnp.float32, + ), + "sample_out", + mode=self.output_mode, ) def check_support(self) -> bool: return self._op.check_support() - def _call(self, q: Any, k: Any, per_head: Any, topk_indices: Any, topk_length: Any | None) -> Any: + def _call( + self, q: Any, k: Any, per_head: Any, topk_indices: Any, topk_length: Any | None + ) -> Any: self.check_support() - for value, expected in ( - (q, self.q_desc), - (k, self.k_desc), - (per_head, self.per_head_desc), - (topk_indices, self.topk_indices_desc), + for value, expected, mode in ( + (q, self.q_desc, self.q_mode), + (k, self.k_desc, self.k_mode), + (per_head, self.per_head_desc, self.per_head_mode), + (topk_indices, self.topk_indices_desc, self.score_mode), ): - self._check_tensor_signature(value, expected) - _check_optional_signature(self, topk_length, self.topk_length_desc, "topk_length") + self._check_tensor_signature(value, expected, mode=mode) + _check_optional_signature( + self, + topk_length, + self.topk_length_desc, + "topk_length", + mode=self.length_mode, + ) if self.target_compute_capability == 90: - kernel_k = jnp.expand_dims(k, axis=2) - kernel_per_head = jnp.transpose(per_head, (0, 2, 1)) + kernel_k = jnp.expand_dims(k, axis=self.k_layout.index("D")) + kernel_k_layout = self.k_layout.replace("D", "HD") + kernel_k_mode = mode_from_layout(kernel_k_layout, kernel_axes="BSHD") + kernel_k_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct(kernel_k.shape, kernel_k.dtype), + "kernel_k", + mode=kernel_k_mode, + ) + per_head_axes = tuple(self.per_head_layout.index(axis) for axis in "BHS") + kernel_per_head = jnp.transpose(per_head, per_head_axes) + kernel_per_head_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct(kernel_per_head.shape, kernel_per_head.dtype), + "kernel_per_head", + ) inputs = (q, kernel_k, kernel_per_head, topk_indices) + input_bindings = ( + (self.q_desc, self.q_mode), + (kernel_k_desc, kernel_k_mode), + (kernel_per_head_desc, None), + (self.topk_indices_desc, self.score_mode), + ) else: inputs = (q, k, per_head, topk_indices) + input_bindings = ( + (self.q_desc, self.q_mode), + (self.k_desc, self.k_mode), + (self.per_head_desc, self.per_head_mode), + (self.topk_indices_desc, self.score_mode), + ) workspace_descs = () if topk_length is not None: inputs += (topk_length,) + input_bindings += ((self.topk_length_desc, self.length_mode),) elif self.target_compute_capability != 90: workspace_descs = ( self.topk_indices_desc.compact_like( @@ -131,6 +256,10 @@ def _call(self, q: Any, k: Any, per_head: Any, topk_indices: Any, topk_length: A launch=self._launch_kernel, output_descs=(self.out_desc,), workspace_descs=workspace_descs, + input_spec=tuple( + self._to_tensor_spec(desc, mode=mode) for desc, mode in input_bindings + ), + output_spec=(self._to_tensor_spec(self.out_desc, mode=self.output_mode),), compile_options=_compile_options(self.target_compute_capability), ) return output @@ -140,7 +269,9 @@ def _launch_kernel( stream: Any, *arguments: Any, ) -> None: - uses_workspace = self.topk_length_desc is None and self.target_compute_capability != 90 + uses_workspace = ( + self.topk_length_desc is None and self.target_compute_capability != 90 + ) if uses_workspace: *inputs, output, topk_length_workspace = arguments else: @@ -194,7 +325,16 @@ def _launch_kernel( topk_in_smem=config.topk_in_smem, topk_indices_global=self._op.topk_indices_global, ) - kernel(q, k, per_head, topk_indices, output, topk_length, cutlass.Float32(self._op.softmax_scale), stream) + kernel( + q, + k, + per_head, + topk_indices, + output, + topk_length, + cutlass.Float32(self._op.softmax_scale), + stream, + ) class SparseIndexerScoreRecompute(_SparseScoreRecompute): @@ -210,6 +350,11 @@ def __init__( sample_topk_length: Any | None = None, qhead_per_kv_head: int | None = None, topk_indices_global: bool = False, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + score_layout: str | None = None, + output_layout: str | None = None, target_compute_capability: int | None = None, ) -> None: super().__init__( @@ -223,6 +368,11 @@ def __init__( sample_topk_length=sample_topk_length, qhead_per_kv_head=qhead_per_kv_head, topk_indices_global=topk_indices_global, + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + score_layout=score_layout, + output_layout=output_layout, target_compute_capability=target_compute_capability, ) @@ -234,7 +384,9 @@ def __call__( topk_indices: Any, topk_length: Any | None = None, ) -> TupleDict: - return TupleDict(predict=self._call(q_indexer, k_indexer, weights, topk_indices, topk_length)) + return TupleDict( + predict=self._call(q_indexer, k_indexer, weights, topk_indices, topk_length) + ) class SparseAttnScoreRecompute(_SparseScoreRecompute): @@ -251,6 +403,11 @@ def __init__( sample_topk_length: Any | None = None, qhead_per_kv_head: int | None = None, topk_indices_global: bool = False, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + score_layout: str | None = None, + output_layout: str | None = None, target_compute_capability: int | None = None, ) -> None: super().__init__( @@ -264,6 +421,11 @@ def __init__( sample_topk_length=sample_topk_length, qhead_per_kv_head=qhead_per_kv_head, topk_indices_global=topk_indices_global, + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + score_layout=score_layout, + output_layout=output_layout, target_compute_capability=target_compute_capability, ) @@ -275,7 +437,9 @@ def __call__( topk_indices: Any, topk_length: Any | None = None, ) -> TupleDict: - return TupleDict(target=self._call(q_attn, k_attn, lse, topk_indices, topk_length)) + return TupleDict( + target=self._call(q_attn, k_attn, lse, topk_indices, topk_length) + ) class _DenseScoreRecompute(JaxApiBase): @@ -296,21 +460,99 @@ def __init__( max_seqlen_q: int | None, max_seqlen_k: int | None, sample_q_causal_offsets: Any | None, + q_layout: str | None, + k_layout: str | None, + per_head_layout: str | None, + output_layout: str | None, + denom_layout: str | None, target_compute_capability: int | None, ) -> None: - self.q_desc = self._to_tensor_desc(sample_q, "sample_q") - self.k_desc = self._to_tensor_desc(sample_k, "sample_k") - self.per_head_desc = self._to_tensor_desc(sample_per_head, "sample_per_head") - self.cu_seqlens_q_desc = _optional_desc(self, sample_cu_seqlens_q, "sample_cu_seqlens_q") - self.cu_seqlens_k_desc = _optional_desc(self, sample_cu_seqlens_k, "sample_cu_seqlens_k") - self.q_causal_offsets_desc = _optional_desc(self, sample_q_causal_offsets, "sample_q_causal_offsets") - self.is_thd = self.q_desc.ndim == 3 + self.is_thd = len(tuple(sample_q.shape)) == 3 + if self.is_thd: + q_defaults = ("THD", "THD", ("THD",)) + k_defaults = ("THD", "THD", ("THD",)) + per_head_defaults = ("TH", "TH", ("TH",)) + output_defaults = ("TK", "TK", ("TK",)) + denom_defaults = ("T", "T", ("T",)) + else: + q_defaults = ("BSHD", "BSHD", ("BSHD", "SBHD")) + k_defaults = ("BSHD", "BSHD", ("BSHD", "SBHD")) + per_head_defaults = ("BSH", "BSH", ("BSH", "SBH")) + output_defaults = ("BSK", "BSK", ("BSK", "SBK")) + denom_defaults = ("BS", "BS", ("BS", "SB")) + self.q_layout, self.q_mode = _resolve_layout( + "q_layout", + q_layout, + default=q_defaults[0], + kernel_axes=q_defaults[1], + supported=q_defaults[2], + ) + self.k_layout, self.k_mode = _resolve_layout( + "k_layout", + k_layout, + default=k_defaults[0], + kernel_axes=k_defaults[1], + supported=k_defaults[2], + ) + self.per_head_layout, self.per_head_mode = _resolve_layout( + "per_head_layout", + per_head_layout, + default=per_head_defaults[0], + kernel_axes=per_head_defaults[1], + supported=per_head_defaults[2], + ) + self.output_layout, self.output_mode = _resolve_layout( + "output_layout", + output_layout, + default=output_defaults[0], + kernel_axes=output_defaults[1], + supported=output_defaults[2], + ) + self.denom_layout, self.denom_mode = _resolve_layout( + "denom_layout", + denom_layout, + default=denom_defaults[0], + kernel_axes=denom_defaults[1], + supported=denom_defaults[2], + ) + + self.q_desc = self._to_tensor_desc(sample_q, "sample_q", mode=self.q_mode) + self.k_desc = self._to_tensor_desc(sample_k, "sample_k", mode=self.k_mode) + self.per_head_desc = self._to_tensor_desc( + sample_per_head, "sample_per_head", mode=self.per_head_mode + ) + self.cu_seqlens_q_desc = _optional_desc( + self, sample_cu_seqlens_q, "sample_cu_seqlens_q" + ) + self.cu_seqlens_k_desc = _optional_desc( + self, sample_cu_seqlens_k, "sample_cu_seqlens_k" + ) + self.q_causal_offsets_desc = _optional_desc( + self, sample_q_causal_offsets, "sample_q_causal_offsets" + ) self.requested_max_seqlen_q = max_seqlen_q self.requested_max_seqlen_k = max_seqlen_k default_out, default_denom = self._default_output_descs(max_seqlen_k) - self.out_desc = default_out if sample_out is None else self._to_tensor_desc(sample_out, "sample_out", init_value=float("-inf")) - self.denom_desc = default_denom if sample_denom_out is None else self._to_tensor_desc(sample_denom_out, "sample_denom_out") + self.out_desc = ( + default_out + if sample_out is None + else self._to_tensor_desc( + sample_out, + "sample_out", + mode=self.output_mode, + init_value=float("-inf"), + ) + ) + self.denom_desc = ( + default_denom + if sample_denom_out is None + else self._to_tensor_desc( + sample_denom_out, + "sample_denom_out", + mode=self.denom_mode, + ) + ) self.target_compute_capability = self._resolve_compute_capability( target_compute_capability, SUPPORTED_COMPUTE_CAPABILITIES, @@ -335,28 +577,42 @@ def __init__( target_compute_capability=self.target_compute_capability, ) - def _default_output_descs(self, max_seqlen_k: int | None) -> tuple[JaxTensorDesc, JaxTensorDesc]: + def _default_output_descs( + self, max_seqlen_k: int | None + ) -> tuple[JaxTensorDesc, JaxTensorDesc]: if self.is_thd: if max_seqlen_k is None: - raise ValueError("max_seqlen_k is required to infer THD dense score outputs") + raise ValueError( + "max_seqlen_k is required to infer THD dense score outputs" + ) out_shape = (self.q_desc.shape[0], int(max_seqlen_k)) denom_shape = (self.q_desc.shape[0],) elif self.q_desc.ndim == 4 and self.k_desc.ndim == 4: - out_shape = (self.q_desc.shape[0], self.q_desc.shape[1], self.k_desc.shape[1]) + out_shape = ( + self.q_desc.shape[0], + self.q_desc.shape[1], + self.k_desc.shape[1], + ) denom_shape = self.q_desc.shape[:2] else: - raise ValueError(f"Dense score inputs must both use BSHD or THD layout, got Q={self.q_desc.shape}, K={self.k_desc.shape}") + raise ValueError( + f"Dense score inputs must both use BSHD or THD layout, got Q={self.q_desc.shape}, K={self.k_desc.shape}" + ) return ( - self.q_desc.compact_like( - cudnn_dtype=data_type.FLOAT, - shape=out_shape, - name="sample_out", + self._to_tensor_desc( + jax.ShapeDtypeStruct( + to_public_axes(out_shape, self.output_mode), jnp.float32 + ), + "sample_out", + mode=self.output_mode, init_value=float("-inf"), ), - self.q_desc.compact_like( - cudnn_dtype=data_type.FLOAT, - shape=denom_shape, - name="sample_denom_out", + self._to_tensor_desc( + jax.ShapeDtypeStruct( + to_public_axes(denom_shape, self.denom_mode), jnp.float32 + ), + "sample_denom_out", + mode=self.denom_mode, ), ) @@ -373,26 +629,63 @@ def _call( q_causal_offsets: Any | None, ) -> TupleDict: self.check_support() - for value, expected in ((q, self.q_desc), (k, self.k_desc), (per_head, self.per_head_desc)): - self._check_tensor_signature(value, expected) - _check_optional_signature(self, cu_seqlens_q, self.cu_seqlens_q_desc, "cu_seqlens_q") - _check_optional_signature(self, cu_seqlens_k, self.cu_seqlens_k_desc, "cu_seqlens_k") - _check_optional_signature(self, q_causal_offsets, self.q_causal_offsets_desc, "q_causal_offsets") + for value, expected, mode in ( + (q, self.q_desc, self.q_mode), + (k, self.k_desc, self.k_mode), + (per_head, self.per_head_desc, self.per_head_mode), + ): + self._check_tensor_signature(value, expected, mode=mode) + _check_optional_signature( + self, cu_seqlens_q, self.cu_seqlens_q_desc, "cu_seqlens_q" + ) + _check_optional_signature( + self, cu_seqlens_k, self.cu_seqlens_k_desc, "cu_seqlens_k" + ) + _check_optional_signature( + self, q_causal_offsets, self.q_causal_offsets_desc, "q_causal_offsets" + ) if self.target_compute_capability == 90: - kernel_per_head = jnp.transpose(per_head, (0, 2, 1)) + per_head_axes = tuple(self.per_head_layout.index(axis) for axis in "BHS") + kernel_per_head = jnp.transpose(per_head, per_head_axes) + kernel_per_head_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct(kernel_per_head.shape, kernel_per_head.dtype), + "kernel_per_head", + ) inputs = (q, k, kernel_per_head) + input_bindings = ( + (self.q_desc, self.q_mode), + (self.k_desc, self.k_mode), + (kernel_per_head_desc, None), + ) else: inputs = (q, k, per_head) + input_bindings = ( + (self.q_desc, self.q_mode), + (self.k_desc, self.k_mode), + (self.per_head_desc, self.per_head_mode), + ) if self.is_thd: inputs += (cu_seqlens_q, cu_seqlens_k) + input_bindings += ( + (self.cu_seqlens_q_desc, None), + (self.cu_seqlens_k_desc, None), + ) if q_causal_offsets is not None: inputs += (q_causal_offsets,) + input_bindings += ((self.q_causal_offsets_desc, None),) output, denominator = self._call_kernel( inputs, launch=self._launch_kernel, output_descs=(self.out_desc, self.denom_desc), + input_spec=tuple( + self._to_tensor_spec(desc, mode=mode) for desc, mode in input_bindings + ), + output_spec=( + self._to_tensor_spec(self.out_desc, mode=self.output_mode), + self._to_tensor_spec(self.denom_desc, mode=self.denom_mode), + ), compile_options=_compile_options(self.target_compute_capability), ) return TupleDict(out=output, denom=denominator) @@ -428,7 +721,17 @@ def _launch_kernel( num_head_tiles=config.num_head_tiles, ratio=self._op.ratio, ) - kernel(q, k, None, stream, output, per_head, None, denominator, q_causal_offsets) + kernel( + q, + k, + None, + stream, + output, + per_head, + None, + denominator, + q_causal_offsets, + ) return import cutlass @@ -486,6 +789,11 @@ def __init__( max_seqlen_q: int | None = None, max_seqlen_k: int | None = None, sample_q_causal_offsets: Any | None = None, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + output_layout: str | None = None, + denom_layout: str | None = None, target_compute_capability: int | None = None, ) -> None: super().__init__( @@ -503,6 +811,11 @@ def __init__( max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, sample_q_causal_offsets=sample_q_causal_offsets, + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + output_layout=output_layout, + denom_layout=denom_layout, target_compute_capability=target_compute_capability, ) @@ -536,6 +849,11 @@ def __init__( max_seqlen_q: int | None = None, max_seqlen_k: int | None = None, sample_q_causal_offsets: Any | None = None, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + output_layout: str | None = None, + denom_layout: str | None = None, target_compute_capability: int | None = None, ) -> None: super().__init__( @@ -553,6 +871,11 @@ def __init__( max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, sample_q_causal_offsets=sample_q_causal_offsets, + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + output_layout=output_layout, + denom_layout=denom_layout, target_compute_capability=target_compute_capability, ) @@ -570,7 +893,16 @@ def __call__( @partial( jax.jit, - static_argnames=("qhead_per_kv_head", "topk_indices_global", "target_compute_capability"), + static_argnames=( + "qhead_per_kv_head", + "topk_indices_global", + "q_layout", + "k_layout", + "per_head_layout", + "score_layout", + "output_layout", + "target_compute_capability", + ), ) def sparse_indexer_score_recompute_wrapper( q_indexer: Any, @@ -580,6 +912,11 @@ def sparse_indexer_score_recompute_wrapper( qhead_per_kv_head: int | None = None, topk_length: Any | None = None, topk_indices_global: bool = False, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + score_layout: str | None = None, + output_layout: str | None = None, target_compute_capability: int | None = None, ) -> TupleDict: return SparseIndexerScoreRecompute( @@ -587,16 +924,33 @@ def sparse_indexer_score_recompute_wrapper( jax.ShapeDtypeStruct(k_indexer.shape, k_indexer.dtype), jax.ShapeDtypeStruct(weights.shape, weights.dtype), jax.ShapeDtypeStruct(topk_indices.shape, topk_indices.dtype), - sample_topk_length=None if topk_length is None else jax.ShapeDtypeStruct(topk_length.shape, topk_length.dtype), + sample_topk_length=None + if topk_length is None + else jax.ShapeDtypeStruct(topk_length.shape, topk_length.dtype), qhead_per_kv_head=qhead_per_kv_head, topk_indices_global=topk_indices_global, + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + score_layout=score_layout, + output_layout=output_layout, target_compute_capability=target_compute_capability, )(q_indexer, k_indexer, weights, topk_indices, topk_length) @partial( jax.jit, - static_argnames=("softmax_scale", "qhead_per_kv_head", "topk_indices_global", "target_compute_capability"), + static_argnames=( + "softmax_scale", + "qhead_per_kv_head", + "topk_indices_global", + "q_layout", + "k_layout", + "per_head_layout", + "score_layout", + "output_layout", + "target_compute_capability", + ), ) def sparse_attn_score_recompute_wrapper( q_attn: Any, @@ -607,6 +961,11 @@ def sparse_attn_score_recompute_wrapper( qhead_per_kv_head: int | None = None, topk_length: Any | None = None, topk_indices_global: bool = False, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + score_layout: str | None = None, + output_layout: str | None = None, target_compute_capability: int | None = None, ) -> TupleDict: return SparseAttnScoreRecompute( @@ -615,9 +974,16 @@ def sparse_attn_score_recompute_wrapper( jax.ShapeDtypeStruct(lse.shape, lse.dtype), jax.ShapeDtypeStruct(topk_indices.shape, topk_indices.dtype), softmax_scale, - sample_topk_length=None if topk_length is None else jax.ShapeDtypeStruct(topk_length.shape, topk_length.dtype), + sample_topk_length=None + if topk_length is None + else jax.ShapeDtypeStruct(topk_length.shape, topk_length.dtype), qhead_per_kv_head=qhead_per_kv_head, topk_indices_global=topk_indices_global, + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + score_layout=score_layout, + output_layout=output_layout, target_compute_capability=target_compute_capability, )(q_attn, k_attn, lse, topk_indices, topk_length) @@ -630,6 +996,11 @@ def sparse_attn_score_recompute_wrapper( "ratio", "max_seqlen_q", "max_seqlen_k", + "q_layout", + "k_layout", + "per_head_layout", + "output_layout", + "denom_layout", "target_compute_capability", ), ) @@ -645,6 +1016,11 @@ def dense_indexer_score_recompute_wrapper( max_seqlen_q: int | None = None, max_seqlen_k: int | None = None, q_causal_offsets: Any | None = None, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + output_layout: str | None = None, + denom_layout: str | None = None, target_compute_capability: int | None = None, ) -> TupleDict: return DenseIndexerScoreRecompute( @@ -654,11 +1030,22 @@ def dense_indexer_score_recompute_wrapper( qhead_per_kv_head=qhead_per_kv_head, sm_scale=sm_scale, ratio=ratio, - sample_cu_seqlens_q=None if cu_seqlens_q is None else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype), - sample_cu_seqlens_k=None if cu_seqlens_k is None else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype), + sample_cu_seqlens_q=None + if cu_seqlens_q is None + else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype), + sample_cu_seqlens_k=None + if cu_seqlens_k is None + else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype), max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, - sample_q_causal_offsets=None if q_causal_offsets is None else jax.ShapeDtypeStruct(q_causal_offsets.shape, q_causal_offsets.dtype), + sample_q_causal_offsets=None + if q_causal_offsets is None + else jax.ShapeDtypeStruct(q_causal_offsets.shape, q_causal_offsets.dtype), + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + output_layout=output_layout, + denom_layout=denom_layout, target_compute_capability=target_compute_capability, )(q, k, weights, cu_seqlens_q, cu_seqlens_k, q_causal_offsets) @@ -671,6 +1058,11 @@ def dense_indexer_score_recompute_wrapper( "ratio", "max_seqlen_q", "max_seqlen_k", + "q_layout", + "k_layout", + "per_head_layout", + "output_layout", + "denom_layout", "target_compute_capability", ), ) @@ -686,6 +1078,11 @@ def dense_attn_score_recompute_wrapper( max_seqlen_q: int | None = None, max_seqlen_k: int | None = None, q_causal_offsets: Any | None = None, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + output_layout: str | None = None, + denom_layout: str | None = None, target_compute_capability: int | None = None, ) -> TupleDict: return DenseAttnScoreRecompute( @@ -695,11 +1092,22 @@ def dense_attn_score_recompute_wrapper( softmax_scale, qhead_per_kv_head=qhead_per_kv_head, ratio=ratio, - sample_cu_seqlens_q=None if cu_seqlens_q is None else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype), - sample_cu_seqlens_k=None if cu_seqlens_k is None else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype), + sample_cu_seqlens_q=None + if cu_seqlens_q is None + else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype), + sample_cu_seqlens_k=None + if cu_seqlens_k is None + else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype), max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, - sample_q_causal_offsets=None if q_causal_offsets is None else jax.ShapeDtypeStruct(q_causal_offsets.shape, q_causal_offsets.dtype), + sample_q_causal_offsets=None + if q_causal_offsets is None + else jax.ShapeDtypeStruct(q_causal_offsets.shape, q_causal_offsets.dtype), + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + output_layout=output_layout, + denom_layout=denom_layout, target_compute_capability=target_compute_capability, )(q, k, lse, cu_seqlens_q, cu_seqlens_k, q_causal_offsets) diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/op.py b/python/cudnn/deepseek_sparse_attention/score_recompute/op.py index e700f96b5..f1357c126 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/op.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/op.py @@ -13,7 +13,12 @@ from ... import data_type from ..._op import Op from ..._tensor_desc import TensorDesc -from .config import DenseScoreKernelConfig, SparseScoreKernelConfig, resolve_dense_score_kernel_config, resolve_sparse_score_kernel_config +from .config import ( + DenseScoreKernelConfig, + SparseScoreKernelConfig, + resolve_dense_score_kernel_config, + resolve_sparse_score_kernel_config, +) SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) @@ -36,7 +41,9 @@ class DenseScoreSm90Config: num_head_tiles: int -def _require_desc(value: Any, name: str, *, optional: bool = False) -> TensorDesc[Any] | None: +def _require_desc( + value: Any, name: str, *, optional: bool = False +) -> TensorDesc[Any] | None: if optional and value is None: return None if not isinstance(value, TensorDesc): @@ -55,12 +62,25 @@ def _require_rank(desc: TensorDesc[Any], rank: int, name: str) -> None: def _require_compact(desc: TensorDesc[Any], name: str) -> None: - if not desc.is_compact(tuple(reversed(range(desc.ndim)))): - raise ValueError(f"{name} must be row-major compact, got stride {desc.stride}") + if not desc.is_compact(): + raise ValueError(f"{name} must be compact, got stride {desc.stride}") + + +def _require_contiguous_tail(desc: TensorDesc[Any], rank: int, name: str) -> None: + expected_stride = 1 + for axis in range(desc.ndim - 1, desc.ndim - rank - 1, -1): + if desc.shape[axis] != 1 and desc.stride[axis] != expected_stride: + raise ValueError( + f"{name} must have its final {rank} axes contiguous, " + f"got shape {desc.shape} and stride {desc.stride}" + ) + expected_stride *= max(desc.shape[axis], 1) def _positive_dimensions(named_dimensions: dict[str, int], operation: str) -> None: - invalid = ", ".join(f"{name}={value}" for name, value in named_dimensions.items() if value <= 0) + invalid = ", ".join( + f"{name}={value}" for name, value in named_dimensions.items() if value <= 0 + ) if invalid: raise ValueError(f"{operation} dimensions must be positive, got {invalid}") @@ -74,7 +94,9 @@ def _normalize_target(target_compute_capability: int) -> int: raise TypeError("target_compute_capability must be an integer") from error if target not in SUPPORTED_COMPUTE_CAPABILITIES: supported = ", ".join(f"SM{value}" for value in SUPPORTED_COMPUTE_CAPABILITIES) - raise ValueError(f"Unsupported score-recompute target SM{target}; supported targets are {supported}") + raise ValueError( + f"Unsupported score-recompute target SM{target}; supported targets are {supported}" + ) return target @@ -109,9 +131,13 @@ def __init__( self.output = _require_desc(output, "output") self.topk_length = _require_desc(topk_length, "topk_length", optional=True) if score_type not in ("indexer", "attention"): - raise ValueError(f"score_type must be 'indexer' or 'attention', got {score_type!r}") + raise ValueError( + f"score_type must be 'indexer' or 'attention', got {score_type!r}" + ) if not isinstance(topk_indices_global, bool): - raise TypeError(f"topk_indices_global must be a bool, got {type(topk_indices_global).__name__}") + raise TypeError( + f"topk_indices_global must be a bool, got {type(topk_indices_global).__name__}" + ) self.score_type = score_type self.softmax_scale = _normalize_scale(softmax_scale, "softmax_scale") @@ -128,7 +154,9 @@ def check_support(self) -> bool: _require_rank(self.q, 4, "Q") _require_rank(self.k, 3, "K") - _require_rank(self.per_head, 3, "weights" if self.score_type == "indexer" else "LSE") + _require_rank( + self.per_head, 3, "weights" if self.score_type == "indexer" else "LSE" + ) _require_rank(self.topk_indices, 3, "topk_indices") _require_rank(self.output, 3, "output") _require_dtype(self.q, data_type.BFLOAT16, "Q") @@ -148,6 +176,17 @@ def check_support(self) -> bool: (self.output, "output"), ): _require_compact(desc, name) + _require_contiguous_tail(self.q, 2, "Q") + for desc, name in ( + (self.k, "K"), + ( + self.per_head, + "weights" if self.score_type == "indexer" else "LSE", + ), + (self.topk_indices, "topk_indices"), + (self.output, "output"), + ): + _require_contiguous_tail(desc, 1, name) batch, seqlen_q, num_query_heads, head_dim = self.q.shape k_batch, seqlen_k, k_head_dim = self.k.shape @@ -164,23 +203,37 @@ def check_support(self) -> bool: "Sparse score-recompute", ) if k_batch != batch or k_head_dim != head_dim: - raise ValueError(f"K shape must be {(batch, seqlen_k, head_dim)}, got {self.k.shape}") + raise ValueError( + f"K shape must be {(batch, seqlen_k, head_dim)}, got {self.k.shape}" + ) expected_per_head = (batch, seqlen_q, num_query_heads) if self.per_head.shape != expected_per_head: - raise ValueError(f"per-head tensor must have shape {expected_per_head}, got {self.per_head.shape}") + raise ValueError( + f"per-head tensor must have shape {expected_per_head}, got {self.per_head.shape}" + ) if (topk_batch, topk_seqlen_q) != (batch, seqlen_q): - raise ValueError(f"topk_indices leading dimensions must be {(batch, seqlen_q)}, got {self.topk_indices.shape[:2]}") + raise ValueError( + f"topk_indices leading dimensions must be {(batch, seqlen_q)}, got {self.topk_indices.shape[:2]}" + ) if self.output.shape != self.topk_indices.shape: - raise ValueError(f"output must have shape {self.topk_indices.shape}, got {self.output.shape}") + raise ValueError( + f"output must have shape {self.topk_indices.shape}, got {self.output.shape}" + ) if self.topk_length is not None: _require_rank(self.topk_length, 2, "topk_length") _require_dtype(self.topk_length, data_type.INT32, "topk_length") _require_compact(self.topk_length, "topk_length") if self.topk_length.shape != (batch, seqlen_q): - raise ValueError(f"topk_length must have shape {(batch, seqlen_q)}, got {self.topk_length.shape}") + raise ValueError( + f"topk_length must have shape {(batch, seqlen_q)}, got {self.topk_length.shape}" + ) - qhead_per_kv_head = num_query_heads if self.requested_qhead_per_kv_head is None else self.requested_qhead_per_kv_head + qhead_per_kv_head = ( + num_query_heads + if self.requested_qhead_per_kv_head is None + else self.requested_qhead_per_kv_head + ) if isinstance(qhead_per_kv_head, bool): raise TypeError("qhead_per_kv_head must be an integer") try: @@ -188,22 +241,32 @@ def check_support(self) -> bool: except TypeError as error: raise TypeError("qhead_per_kv_head must be an integer") from error if qhead_per_kv_head != num_query_heads: - raise ValueError(f"qhead_per_kv_head must equal H_q ({num_query_heads}) for sparse MQA, got {qhead_per_kv_head}") + raise ValueError( + f"qhead_per_kv_head must equal H_q ({num_query_heads}) for sparse MQA, got {qhead_per_kv_head}" + ) if self.target_compute_capability == 90: if topk % 128: - raise ValueError(f"SM90 sparse score recompute requires topk to be a multiple of 128, got {topk}") + raise ValueError( + f"SM90 sparse score recompute requires topk to be a multiple of 128, got {topk}" + ) if qhead_per_kv_head <= 1: - raise ValueError("SM90 sparse score recompute requires qhead_per_kv_head > 1") + raise ValueError( + "SM90 sparse score recompute requires qhead_per_kv_head > 1" + ) tile_m = min(qhead_per_kv_head, 64) if qhead_per_kv_head % tile_m: - raise ValueError(f"qhead_per_kv_head ({qhead_per_kv_head}) must be divisible by tile_m ({tile_m})") - config: SparseScoreKernelConfig | SparseScoreSm90Config = SparseScoreSm90Config( - tile_m=tile_m, - tile_n=64, - kv_stage=2, - num_threads=256, - num_head_tiles=qhead_per_kv_head // tile_m, + raise ValueError( + f"qhead_per_kv_head ({qhead_per_kv_head}) must be divisible by tile_m ({tile_m})" + ) + config: SparseScoreKernelConfig | SparseScoreSm90Config = ( + SparseScoreSm90Config( + tile_m=tile_m, + tile_n=64, + kv_stage=2, + num_threads=256, + num_head_tiles=qhead_per_kv_head // tile_m, + ) ) else: config = resolve_sparse_score_kernel_config( @@ -249,9 +312,13 @@ def __init__( self.denominator = _require_desc(denominator, "denominator") self.cu_seqlens_q = _require_desc(cu_seqlens_q, "cu_seqlens_q", optional=True) self.cu_seqlens_k = _require_desc(cu_seqlens_k, "cu_seqlens_k", optional=True) - self.q_causal_offsets = _require_desc(q_causal_offsets, "q_causal_offsets", optional=True) + self.q_causal_offsets = _require_desc( + q_causal_offsets, "q_causal_offsets", optional=True + ) if score_type not in ("indexer", "attention"): - raise ValueError(f"score_type must be 'indexer' or 'attention', got {score_type!r}") + raise ValueError( + f"score_type must be 'indexer' or 'attention', got {score_type!r}" + ) if not isinstance(is_thd, bool): raise TypeError(f"is_thd must be a bool, got {type(is_thd).__name__}") if isinstance(ratio, bool): @@ -278,7 +345,9 @@ def __init__( @staticmethod def _positive_static(value: int | None, name: str) -> int: if value is None or isinstance(value, bool): - raise ValueError(f"{name} must be provided as a positive integer for THD inputs") + raise ValueError( + f"{name} must be provided as a positive integer for THD inputs" + ) try: value = index(value) except TypeError as error: @@ -298,7 +367,11 @@ def check_support(self) -> bool: expected_rank = 3 if self.is_thd else 4 _require_rank(self.q, expected_rank, "Q") _require_rank(self.k, expected_rank, "K") - _require_rank(self.per_head, expected_rank - 1, "weights" if self.score_type == "indexer" else "LSE") + _require_rank( + self.per_head, + expected_rank - 1, + "weights" if self.score_type == "indexer" else "LSE", + ) _require_rank(self.output, expected_rank - 1, "output") _require_rank(self.denominator, expected_rank - 2, "denominator") _require_dtype(self.q, data_type.BFLOAT16, "Q") @@ -318,29 +391,60 @@ def check_support(self) -> bool: (self.denominator, "denominator"), ): _require_compact(desc, name) + for desc, name in ((self.q, "Q"), (self.k, "K")): + _require_contiguous_tail(desc, 2, name) + for desc, name in ( + ( + self.per_head, + "weights" if self.score_type == "indexer" else "LSE", + ), + (self.output, "output"), + ): + _require_contiguous_tail(desc, 1, name) if self.is_thd: total_q, num_query_heads, head_dim = self.q.shape total_k, num_kv_heads, k_head_dim = self.k.shape _positive_dimensions( - {"total_q": total_q, "total_k": total_k, "H_q": num_query_heads, "H_kv": num_kv_heads, "D": head_dim}, + { + "total_q": total_q, + "total_k": total_k, + "H_q": num_query_heads, + "H_kv": num_kv_heads, + "D": head_dim, + }, "Dense THD score-recompute", ) if self.per_head.shape != (total_q, num_query_heads): - raise ValueError(f"THD per-head tensor must have shape {(total_q, num_query_heads)}, got {self.per_head.shape}") + raise ValueError( + f"THD per-head tensor must have shape {(total_q, num_query_heads)}, got {self.per_head.shape}" + ) if self.cu_seqlens_q is None or self.cu_seqlens_k is None: - raise ValueError("THD dense score recompute requires both cu_seqlens_q and cu_seqlens_k") - for desc, name in ((self.cu_seqlens_q, "cu_seqlens_q"), (self.cu_seqlens_k, "cu_seqlens_k")): + raise ValueError( + "THD dense score recompute requires both cu_seqlens_q and cu_seqlens_k" + ) + for desc, name in ( + (self.cu_seqlens_q, "cu_seqlens_q"), + (self.cu_seqlens_k, "cu_seqlens_k"), + ): _require_rank(desc, 1, name) _require_dtype(desc, data_type.INT32, name) _require_compact(desc, name) if self.cu_seqlens_q.shape != self.cu_seqlens_k.shape: - raise ValueError(f"cu_seqlens_q and cu_seqlens_k shapes must match, got {self.cu_seqlens_q.shape} and {self.cu_seqlens_k.shape}") + raise ValueError( + f"cu_seqlens_q and cu_seqlens_k shapes must match, got {self.cu_seqlens_q.shape} and {self.cu_seqlens_k.shape}" + ) batch = self.cu_seqlens_q.shape[0] - 1 if batch <= 0: - raise ValueError("cumulative sequence-length tensors must contain at least two entries") - max_seqlen_q = self._positive_static(self.requested_max_seqlen_q, "max_seqlen_q") - max_seqlen_k = self._positive_static(self.requested_max_seqlen_k, "max_seqlen_k") + raise ValueError( + "cumulative sequence-length tensors must contain at least two entries" + ) + max_seqlen_q = self._positive_static( + self.requested_max_seqlen_q, "max_seqlen_q" + ) + max_seqlen_k = self._positive_static( + self.requested_max_seqlen_k, "max_seqlen_k" + ) expected_output = (total_q, max_seqlen_k) expected_denominator = (total_q,) if self.target_compute_capability == 90: @@ -349,17 +453,28 @@ def check_support(self) -> bool: ) else: if self.cu_seqlens_q is not None or self.cu_seqlens_k is not None: - raise ValueError("BSHD dense score recompute does not accept cu_seqlens_q or cu_seqlens_k") + raise ValueError( + "BSHD dense score recompute does not accept cu_seqlens_q or cu_seqlens_k" + ) batch, seqlen_q, num_query_heads, head_dim = self.q.shape k_batch, seqlen_k, num_kv_heads, k_head_dim = self.k.shape _positive_dimensions( - {"B": batch, "S_q": seqlen_q, "S_k": seqlen_k, "H_q": num_query_heads, "H_kv": num_kv_heads, "D": head_dim}, + { + "B": batch, + "S_q": seqlen_q, + "S_k": seqlen_k, + "H_q": num_query_heads, + "H_kv": num_kv_heads, + "D": head_dim, + }, "Dense BSHD score-recompute", ) if k_batch != batch: raise ValueError(f"K batch dimension must be {batch}, got {k_batch}") if self.per_head.shape != (batch, seqlen_q, num_query_heads): - raise ValueError(f"BSHD per-head tensor must have shape {(batch, seqlen_q, num_query_heads)}, got {self.per_head.shape}") + raise ValueError( + f"BSHD per-head tensor must have shape {(batch, seqlen_q, num_query_heads)}, got {self.per_head.shape}" + ) max_seqlen_q, max_seqlen_k = seqlen_q, seqlen_k expected_output = (batch, seqlen_q, seqlen_k) expected_denominator = (batch, seqlen_q) @@ -367,21 +482,33 @@ def check_support(self) -> bool: if k_head_dim != head_dim: raise ValueError(f"K head dimension must be {head_dim}, got {k_head_dim}") if num_kv_heads != 1: - raise ValueError(f"Dense score recompute currently requires MQA with H_kv=1, got {num_kv_heads}") + raise ValueError( + f"Dense score recompute currently requires MQA with H_kv=1, got {num_kv_heads}" + ) if self.output.shape != expected_output: - raise ValueError(f"output must have shape {expected_output}, got {self.output.shape}") + raise ValueError( + f"output must have shape {expected_output}, got {self.output.shape}" + ) if self.denominator.shape != expected_denominator: - raise ValueError(f"denominator must have shape {expected_denominator}, got {self.denominator.shape}") + raise ValueError( + f"denominator must have shape {expected_denominator}, got {self.denominator.shape}" + ) if self.q_causal_offsets is not None: _require_rank(self.q_causal_offsets, 1, "q_causal_offsets") _require_dtype(self.q_causal_offsets, data_type.INT32, "q_causal_offsets") _require_compact(self.q_causal_offsets, "q_causal_offsets") if self.q_causal_offsets.shape != (batch,): - raise ValueError(f"q_causal_offsets must have shape {(batch,)}, got {self.q_causal_offsets.shape}") + raise ValueError( + f"q_causal_offsets must have shape {(batch,)}, got {self.q_causal_offsets.shape}" + ) inferred_qhead_per_kv_head = num_query_heads // num_kv_heads - qhead_per_kv_head = inferred_qhead_per_kv_head if self.requested_qhead_per_kv_head is None else self.requested_qhead_per_kv_head + qhead_per_kv_head = ( + inferred_qhead_per_kv_head + if self.requested_qhead_per_kv_head is None + else self.requested_qhead_per_kv_head + ) if isinstance(qhead_per_kv_head, bool): raise TypeError("qhead_per_kv_head must be an integer") try: @@ -389,20 +516,28 @@ def check_support(self) -> bool: except TypeError as error: raise TypeError("qhead_per_kv_head must be an integer") from error if qhead_per_kv_head != inferred_qhead_per_kv_head: - raise ValueError(f"qhead_per_kv_head must equal H_q / H_kv ({inferred_qhead_per_kv_head}), got {qhead_per_kv_head}") + raise ValueError( + f"qhead_per_kv_head must equal H_q / H_kv ({inferred_qhead_per_kv_head}), got {qhead_per_kv_head}" + ) if self.target_compute_capability == 90: if qhead_per_kv_head <= 1: - raise ValueError("SM90 dense score recompute requires qhead_per_kv_head > 1") + raise ValueError( + "SM90 dense score recompute requires qhead_per_kv_head > 1" + ) tile_m = min(qhead_per_kv_head, 64) if qhead_per_kv_head % tile_m: - raise ValueError(f"qhead_per_kv_head ({qhead_per_kv_head}) must be divisible by tile_m ({tile_m})") - config: DenseScoreKernelConfig | DenseScoreSm90Config = DenseScoreSm90Config( - tile_m=tile_m, - tile_n=64, - kv_stage=2, - num_threads=384, - num_head_tiles=qhead_per_kv_head // tile_m, + raise ValueError( + f"qhead_per_kv_head ({qhead_per_kv_head}) must be divisible by tile_m ({tile_m})" + ) + config: DenseScoreKernelConfig | DenseScoreSm90Config = ( + DenseScoreSm90Config( + tile_m=tile_m, + tile_n=64, + kv_stage=2, + num_threads=384, + num_head_tiles=qhead_per_kv_head // tile_m, + ) ) else: config = resolve_dense_score_kernel_config( diff --git a/python/cudnn/discrete_grouped_gemm/__init__.py b/python/cudnn/discrete_grouped_gemm/__init__.py index db09418bf..8ee1e9372 100644 --- a/python/cudnn/discrete_grouped_gemm/__init__.py +++ b/python/cudnn/discrete_grouped_gemm/__init__.py @@ -1,26 +1,27 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -""" -Discrete-weight Grouped GEMM GLU Kernel Module +"""Lazy Torch exports for discrete-weight grouped GEMM operations. -This module provides forward and backward discrete-weight grouped GEMM with GLU -activation (SwiGLU/GeGLU) for MoE (Mixture of Experts) workloads on SM100+ GPUs. +JAX APIs are exported through :mod:`cudnn.jax`. """ -from .discrete_grouped_gemm_swiglu import ( - DiscreteGroupedGemmSwigluSm100, - discrete_grouped_gemm_swiglu_wrapper_sm100, -) +from .._operation_api import make_operation_api -from .discrete_grouped_gemm_dswiglu import ( - DiscreteGroupedGemmDswigluSm100, - discrete_grouped_gemm_dswiglu_wrapper_sm100, +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "discrete_grouped_gemm_swiglu": ( + "DiscreteGroupedGemmSwigluSm100", + "discrete_grouped_gemm_swiglu_wrapper_sm100", + ), + "discrete_grouped_gemm_dswiglu": ( + "DiscreteGroupedGemmDswigluSm100", + "discrete_grouped_gemm_dswiglu_wrapper_sm100", + ), + }, + submodules=( + "discrete_grouped_gemm_dswiglu", + "discrete_grouped_gemm_swiglu", + ), ) - -__all__ = [ - "DiscreteGroupedGemmSwigluSm100", - "discrete_grouped_gemm_swiglu_wrapper_sm100", - "DiscreteGroupedGemmDswigluSm100", - "discrete_grouped_gemm_dswiglu_wrapper_sm100", -] diff --git a/python/cudnn/discrete_grouped_gemm/_jax_common.py b/python/cudnn/discrete_grouped_gemm/_jax_common.py new file mode 100644 index 000000000..a4a4d82a9 --- /dev/null +++ b/python/cudnn/discrete_grouped_gemm/_jax_common.py @@ -0,0 +1,321 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared JAX metadata and validation for discrete grouped GEMMs.""" + +from __future__ import annotations + +import os +from typing import Any + +import jax +import jax.numpy as jnp + +from .. import data_type +from .._dense_gemm import ( + block_scale_shape, + require_16_byte_alignment, + require_block_scale_layout, + require_cluster_shape, + require_compact_major, + require_mma_tiler, + require_tensor_shape, +) +from .._jax import JaxApiBase, JaxTensorDesc +from .._jax.datatypes import jax_to_cudnn_dtype, normalize_jax_dtype +from .._jax.gemm import ( + BLOCK_SCALE_MODE, + PROBABILITY_MODE, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, +) +from .._jax.layout import mode_from_layout, to_public_axes + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) +FIX_PAD_SIZE = 256 +MAX_EXPERTS = 1024 + +FP8_DTYPES = frozenset({data_type.FP8_E4M3, data_type.FP8_E5M2}) +AB_DTYPES = frozenset({data_type.FP4_E2M1, *FP8_DTYPES}) +SF_DTYPES = frozenset({data_type.FP8_E4M3, data_type.FP8_E8M0}) + + +def _require_dtype( + desc: JaxTensorDesc, allowed: frozenset[data_type], label: str +) -> data_type: + dtype = desc.cudnn_dtype + if dtype not in allowed: + expected = ", ".join(sorted(str(value) for value in allowed)) + raise ValueError( + f"{label} dtype must be one of {{{expected}}}, got {desc.dtype}" + ) + return dtype + + +class DiscreteGroupedGemmJaxBase(JaxApiBase): + """Common signature for the stacked-expert JAX binding. + + The Torch API accepts device pointer tables because Torch exposes stable + allocation addresses. JAX arrays do not expose such an ABI, and pointer + tables would hide the referenced buffers from XLA's liveness analysis. + This binding therefore accepts compact stacked B/SFB arrays as explicit + custom-call operands. The existing discrete kernel derives one pointer + per expert while initializing its TMA descriptors. + """ + + def _initialize_common( + self, + sample_a: Any, + sample_b: Any, + sample_sfa: Any, + sample_sfb: Any, + sample_padded_offsets: Any, + sample_alpha: Any, + *, + acc_dtype: Any | None, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int] | None, + sf_vec_size: int, + vector_f32: bool, + m_aligned: int, + use_dynamic_sched: bool, + a_layout: str, + b_layout: str, + output_layout: str, + ) -> None: + self.a_layout = a_layout + self.b_layout = b_layout + self.output_layout = output_layout + self.a_mode = gemm_a_mode(a_layout) + self.b_mode = gemm_b_mode(b_layout) + self.output_mode = gemm_output_mode(output_layout, name="output_layout") + self.scale_mode = BLOCK_SCALE_MODE + self.probability_mode = PROBABILITY_MODE + self.bias_mode = mode_from_layout("LN", kernel_axes="NL") + + self.compute_capability: int | None = None + self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=self.a_mode) + self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=self.b_mode) + self.sfa_desc = self._to_tensor_desc( + sample_sfa, "sample_sfa", mode=self.scale_mode + ) + self.sfb_desc = self._to_tensor_desc( + sample_sfb, "sample_sfb", mode=self.scale_mode + ) + self.padded_offsets_desc = self._to_tensor_desc( + sample_padded_offsets, "sample_padded_offsets" + ) + self.alpha_desc = self._to_tensor_desc(sample_alpha, "sample_alpha") + + self.acc_dtype = normalize_jax_dtype(acc_dtype, jnp.float32, "acc_dtype") + self.requested_mma_tiler_mn = mma_tiler_mn + self.requested_cluster_shape_mn = cluster_shape_mn + self.sf_vec_size = sf_vec_size + self.vector_f32 = vector_f32 + self.m_aligned = m_aligned + self.use_dynamic_sched = use_dynamic_sched + self.num_cluster_overlap_margin = int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ) + + self.m = self.n = self.k = self.expert_cnt = None + self.ab_dtype = self.sf_dtype = None + self.mma_tiler_mn = self.cluster_shape_mn = None + + def _check_common(self) -> None: + if self.a_desc.ndim != 3: + raise ValueError(f"A must have rank 3, got shape {self.a_desc.shape}") + if self.b_desc.ndim != 3: + raise ValueError(f"B must have rank 3, got shape {self.b_desc.shape}") + + m, k, a_groups = self.a_desc.shape + n, b_k, experts = self.b_desc.shape + if m < 0 or min(n, k, experts) <= 0: + raise ValueError( + "M must be non-negative and N, K, and expert count must be " + f"positive, got {(m, n, k, experts)}" + ) + if a_groups != 1: + raise ValueError( + f"A must flatten all expert rows into a singleton group dimension, got {a_groups}" + ) + if b_k != k: + raise ValueError(f"B K dimension must match A, got {b_k} and {k}") + if experts > MAX_EXPERTS: + raise ValueError( + f"expert count must not exceed {MAX_EXPERTS}, got {experts}" + ) + if m % FIX_PAD_SIZE: + raise ValueError( + f"A M dimension must be divisible by {FIX_PAD_SIZE}, got {m}" + ) + if n % 64: + raise ValueError(f"B N dimension must be divisible by 64, got {n}") + + require_tensor_shape( + self.padded_offsets_desc, (experts,), label="padded_offsets" + ) + require_tensor_shape(self.alpha_desc, (experts,), label="alpha") + if self.padded_offsets_desc.cudnn_dtype != data_type.INT32: + raise ValueError( + f"padded_offsets must have int32 dtype, got {self.padded_offsets_desc.dtype}" + ) + if self.alpha_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"alpha must have float32 dtype, got {self.alpha_desc.dtype}" + ) + + ab_dtype = _require_dtype(self.a_desc, AB_DTYPES, "A") + if self.b_desc.cudnn_dtype != ab_dtype: + raise ValueError( + f"A and B must have the same dtype, got {self.a_desc.dtype} and {self.b_desc.dtype}" + ) + sf_dtype = _require_dtype(self.sfa_desc, SF_DTYPES, "SFA") + if self.sfb_desc.cudnn_dtype != sf_dtype: + raise ValueError( + f"SFA and SFB must have the same dtype, got {self.sfa_desc.dtype} and {self.sfb_desc.dtype}" + ) + if self.sf_vec_size not in (16, 32): + raise ValueError(f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}") + if sf_dtype == data_type.FP8_E4M3 and self.sf_vec_size == 32: + raise ValueError("FP8_E4M3 scale factors require sf_vec_size=16") + if ab_dtype in FP8_DTYPES and self.sf_vec_size != 32: + raise ValueError("FP8 A and B require sf_vec_size=32") + + require_tensor_shape( + self.sfa_desc, block_scale_shape(m, k, 1, self.sf_vec_size), label="SFA" + ) + require_tensor_shape( + self.sfb_desc, + block_scale_shape(n, k, experts, self.sf_vec_size), + label="SFB", + ) + require_block_scale_layout(self.sfa_desc, "SFA") + require_block_scale_layout(self.sfb_desc, "SFB") + + a_major = require_compact_major(self.a_desc, "m", "k") + b_major = require_compact_major(self.b_desc, "n", "k") + if a_major != "k" or b_major != "k": + raise ValueError( + "The discrete grouped kernels currently require K-major A and B layouts" + ) + require_16_byte_alignment(self.a_desc) + require_16_byte_alignment(self.b_desc) + + if jax_to_cudnn_dtype(self.acc_dtype) != data_type.FLOAT: + raise ValueError(f"acc_dtype must be float32, got {self.acc_dtype}") + if not isinstance(self.vector_f32, bool): + raise TypeError( + f"vector_f32 must be a bool, got {type(self.vector_f32).__name__}" + ) + if not isinstance(self.use_dynamic_sched, bool): + raise TypeError( + f"use_dynamic_sched must be a bool, got {type(self.use_dynamic_sched).__name__}" + ) + if self.m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {self.m_aligned}") + + mma_tiler_mn = require_mma_tiler( + self.requested_mma_tiler_mn, + allowed_m=(128, 256), + allowed_n=(256,), + ) + cta_group_size = 2 if mma_tiler_mn[0] == 256 else 1 + default_cluster = (2, 1) if cta_group_size == 2 else (1, 1) + cluster_shape_mn = require_cluster_shape( + default_cluster + if self.requested_cluster_shape_mn is None + else self.requested_cluster_shape_mn, + cta_group_size=cta_group_size, + ) + if cluster_shape_mn[0] > 4 or cluster_shape_mn[1] > 4: + raise ValueError( + f"cluster_shape_mn entries must not exceed 4, got {cluster_shape_mn}" + ) + cluster_tile_m = cluster_shape_mn[0] // cta_group_size * mma_tiler_mn[0] + if cluster_tile_m not in (128, 256): + raise ValueError(f"cluster M tile must be 128 or 256, got {cluster_tile_m}") + + self.m, self.n, self.k, self.expert_cnt = m, n, k, experts + self.ab_dtype, self.sf_dtype = ab_dtype, sf_dtype + self.mma_tiler_mn, self.cluster_shape_mn = mma_tiler_mn, cluster_shape_mn + if m > 0: + self.compute_capability = self._resolve_compute_capability( + None, + SUPPORTED_COMPUTE_CAPABILITIES, + type(self).__name__, + ) + + def _canonical_desc( + self, + shape: tuple[int, ...], + dtype: Any, + name: str, + *, + mode: tuple[int, ...] | None = None, + init_value: bool | int | float | None = None, + ) -> JaxTensorDesc: + return self._to_tensor_desc( + jax.ShapeDtypeStruct(to_public_axes(shape, mode), dtype), + name, + mode=mode, + init_value=init_value, + ) + + def _workspace_desc(self, workspace_bytes: int) -> JaxTensorDesc: + if workspace_bytes <= 0: + raise ValueError( + f"kernel workspace size must be positive, got {workspace_bytes}" + ) + return self._canonical_desc((workspace_bytes,), jnp.uint8, "workspace") + + def _materialize_output_desc( + self, + desc: JaxTensorDesc | None, + *, + mode: tuple[int, ...] | None = None, + ) -> Any | None: + """Materialize an inferred result without launching an empty GEMM.""" + + if desc is None: + return None + metadata = self._materialize_tensor_desc(desc, mode=mode) + if desc.init_value is None: + return jnp.empty(metadata.shape, dtype=metadata.dtype) + return jnp.full(metadata.shape, desc.init_value, dtype=metadata.dtype) + + @staticmethod + def _workspace_tensor_spec() -> Any: + """Require the alignment used by the kernel's TMA descriptor stores.""" + + from cutlass.jax import TensorSpec + + return TensorSpec(ptr_assumed_align=128) + + def _check_runtime_common( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + ) -> None: + self._check_tensor_signature(a_tensor, self.a_desc, mode=self.a_mode) + self._check_tensor_signature(b_tensor, self.b_desc, mode=self.b_mode) + self._check_tensor_signature(sfa_tensor, self.sfa_desc, mode=self.scale_mode) + self._check_tensor_signature(sfb_tensor, self.sfb_desc, mode=self.scale_mode) + self._check_tensor_signature(padded_offsets, self.padded_offsets_desc) + self._check_tensor_signature(alpha_tensor, self.alpha_desc) + + +__all__ = [ + "AB_DTYPES", + "DiscreteGroupedGemmJaxBase", + "FIX_PAD_SIZE", + "FP8_DTYPES", + "MAX_EXPERTS", + "SF_DTYPES", + "SUPPORTED_COMPUTE_CAPABILITIES", +] diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/__init__.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/__init__.py index e0355406a..6b6ca6cf6 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/__init__.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/__init__.py @@ -1,12 +1,17 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - DiscreteGroupedGemmDswigluSm100, - discrete_grouped_gemm_dswiglu_wrapper_sm100, -) +"""Lazy Torch API exports for discrete grouped GEMM dSwiGLU.""" + +from ..._operation_api import make_operation_api -__all__ = [ - "DiscreteGroupedGemmDswigluSm100", - "discrete_grouped_gemm_dswiglu_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "api": ( + "DiscreteGroupedGemmDswigluSm100", + "discrete_grouped_gemm_dswiglu_wrapper_sm100", + ), + }, + submodules=("api",), +) diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/discrete_B_blockscaled_grouped_gemm_dglu_dbias.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/discrete_B_blockscaled_grouped_gemm_dglu_dbias.py index d52cb859c..e24c2a59b 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/discrete_B_blockscaled_grouped_gemm_dglu_dbias.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/discrete_B_blockscaled_grouped_gemm_dglu_dbias.py @@ -270,6 +270,7 @@ def __init__( expert_cnt: int, use_dynamic_sched: bool = False, act_func: str = "dswiglu", + stacked_expert_inputs: bool = False, ): """Initializes the configuration for a Blackwell blockscaled dense GEMM kernel. @@ -371,6 +372,10 @@ def __init__( self.vectorized_f32 = vectorized_f32 self.use_dynamic_sched = use_dynamic_sched + # See the forward kernel for the JAX liveness rationale. Torch keeps + # using device pointer tables; JAX passes stacked B/SFB operands whose + # expert addresses are derived here. + self.stacked_expert_inputs = stacked_expert_inputs # Amax reduction configuration self.num_epilog_warps = len(self.epilog_warp_id) @@ -638,10 +643,15 @@ def desc_init_kernel_device_ptrs( expert_idx = cute.arch.block_idx()[0] - b_ptr_tensor = cute.make_tensor(cute.make_ptr(cutlass.Int64, ptrs_b.toint(), AddressSpace.gmem, assumed_align=8), cute.make_layout((self.expert_cnt,))) - sfb_ptr_tensor = cute.make_tensor( - cute.make_ptr(cutlass.Int64, ptrs_sfb.toint(), AddressSpace.gmem, assumed_align=8), cute.make_layout((self.expert_cnt,)) - ) + if cutlass.const_expr(not self.stacked_expert_inputs): + b_ptr_tensor = cute.make_tensor( + cute.make_ptr(cutlass.Int64, ptrs_b.toint(), AddressSpace.gmem, assumed_align=8), + cute.make_layout((self.expert_cnt,)), + ) + sfb_ptr_tensor = cute.make_tensor( + cute.make_ptr(cutlass.Int64, ptrs_sfb.toint(), AddressSpace.gmem, assumed_align=8), + cute.make_layout((self.expert_cnt,)), + ) c1 = cutlass.Int32(1) c0 = cutlass.Int64(0) @@ -653,8 +663,12 @@ def desc_init_kernel_device_ptrs( stride_n = c1_64 stride_k = b_stride_size - b_ptr_val = b_ptr_tensor[expert_idx] - b_ptr = cute.make_ptr(self.b_dtype, b_ptr_val, AddressSpace.gmem) + if cutlass.const_expr(self.stacked_expert_inputs): + b_base_ptr = cute.make_ptr(self.b_dtype, ptrs_b.toint(), AddressSpace.gmem, assumed_align=16) + b_ptr = b_base_ptr + cutlass.Int64(expert_idx) * cutlass.Int64(n) * cutlass.Int64(k) + else: + b_ptr_val = b_ptr_tensor[expert_idx] + b_ptr = cute.make_ptr(self.b_dtype, b_ptr_val, AddressSpace.gmem) b_tensor_i = cute.make_tensor( b_ptr, cute.make_layout((n, k, c1), stride=(stride_n, stride_k, c0)), @@ -671,9 +685,13 @@ def desc_init_kernel_device_ptrs( workspace = TensormapWorkspace(workspace_ptr, ["b", "sfb"]) store_tma_desc(tma_atom_b, workspace.get_ptr("b", expert_idx)) - sfb_ptr_val = sfb_ptr_tensor[expert_idx] - sfb_ptr = cute.make_ptr(self.sf_dtype, sfb_ptr_val, AddressSpace.gmem) sfb_layout = blockscaled_utils.tile_atom_to_shape_SF((n, k, c1), self.sf_vec_size) + if cutlass.const_expr(self.stacked_expert_inputs): + sfb_base_ptr = cute.make_ptr(self.sf_dtype, ptrs_sfb.toint(), AddressSpace.gmem, assumed_align=16) + sfb_ptr = sfb_base_ptr + cutlass.Int64(expert_idx) * cutlass.Int64(cute.cosize(sfb_layout)) + else: + sfb_ptr_val = sfb_ptr_tensor[expert_idx] + sfb_ptr = cute.make_ptr(self.sf_dtype, sfb_ptr_val, AddressSpace.gmem) sfb_tensor_i = cute.make_tensor(sfb_ptr, sfb_layout) tma_atom_sfb, _ = cute.nvgpu.make_tiled_tma_atom_B( sfb_tma_op_arg, diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/jax.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/jax.py new file mode 100644 index 000000000..5d9814acc --- /dev/null +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/jax.py @@ -0,0 +1,601 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for the discrete-weight grouped GEMM dSwiGLU kernel.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._cute_compiler import compile_options_for_target +from ..._dense_gemm import ( + block_scale_shape, + require_16_byte_alignment, + require_compact_major, + require_tensor_shape, +) +from ..._jax import TupleDict +from ..._jax.datatypes import normalize_jax_dtype +from .._jax_common import ( + DiscreteGroupedGemmJaxBase, + FP8_DTYPES, + SUPPORTED_COMPUTE_CAPABILITIES, +) + + +class DiscreteGroupedGemmDswigluSm100(DiscreteGroupedGemmJaxBase): + """JAX callable for the discrete grouped dSwiGLU/dGeGLU kernel. + + B and SFB are compact stacked expert arrays rather than opaque device + pointer tables. Keeping them as custom-call operands gives XLA complete + buffer-liveness information while the kernel still initializes one TMA + descriptor pair per expert. + """ + + def __init__( + self, + sample_a: Any, + sample_b: Any, + sample_c: Any, + sample_sfa: Any, + sample_sfb: Any, + sample_padded_offsets: Any, + sample_alpha: Any, + sample_beta: Any, + sample_prob: Any, + *, + sample_norm_const: Any | None = None, + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + generate_dbias: bool = False, + act_func: str = "dswiglu", + epilogue_op: str | None = None, + use_dynamic_sched: bool = False, + a_layout: str = "LMK", + b_layout: str = "LNK", + output_layout: str = "LMN", + ) -> None: + self._initialize_common( + sample_a, + sample_b, + sample_sfa, + sample_sfb, + sample_padded_offsets, + sample_alpha, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + use_dynamic_sched=use_dynamic_sched, + a_layout=a_layout, + b_layout=b_layout, + output_layout=output_layout, + ) + self.c_desc = self._to_tensor_desc(sample_c, "sample_c", mode=self.output_mode) + self.beta_desc = self._to_tensor_desc(sample_beta, "sample_beta") + self.prob_desc = self._to_tensor_desc( + sample_prob, "sample_prob", mode=self.probability_mode + ) + self.norm_const_desc = ( + None + if sample_norm_const is None + else self._to_tensor_desc(sample_norm_const, "sample_norm_const") + ) + self.d_dtype = normalize_jax_dtype(d_dtype, jnp.bfloat16, "d_dtype") + self.discrete_col_sfd = discrete_col_sfd + self.generate_dbias = generate_dbias + self.act_func = act_func + self.epilogue_op = epilogue_op + + self.generate_sfd = False + self.d_row_desc = self.d_col_desc = self.dprob_desc = None + self.dbias_desc = self.sfd_row_desc = self.sfd_col_desc = self.amax_desc = None + + def check_support(self) -> bool: + self._check_common() + if self.act_func not in ("dswiglu", "dgeglu"): + raise ValueError( + f"act_func must be 'dswiglu' or 'dgeglu', got {self.act_func!r}" + ) + if self.epilogue_op not in (None, "none", "identity", "relu", "srelu"): + raise ValueError( + f"epilogue_op must be None, 'identity', 'relu', or 'srelu', got {self.epilogue_op!r}" + ) + for name, value in ( + ("discrete_col_sfd", self.discrete_col_sfd), + ("generate_dbias", self.generate_dbias), + ): + if not isinstance(value, bool): + raise TypeError(f"{name} must be a bool, got {type(value).__name__}") + + n_out = 2 * self.n + require_tensor_shape(self.c_desc, (self.m, n_out, 1), label="C") + if self.c_desc.cudnn_dtype not in ( + data_type.FLOAT, + data_type.HALF, + data_type.BFLOAT16, + ): + raise ValueError( + f"C must have float32, float16, or bfloat16 dtype, got {self.c_desc.dtype}" + ) + if require_compact_major(self.c_desc, "m", "n") != "n": + raise ValueError("C must use an N-major output layout") + require_16_byte_alignment(self.c_desc) + + require_tensor_shape(self.beta_desc, (self.expert_cnt,), label="beta") + if self.beta_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"beta must have float32 dtype, got {self.beta_desc.dtype}" + ) + require_tensor_shape(self.prob_desc, (self.m, 1, 1), label="prob") + if self.prob_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"prob must have float32 dtype, got {self.prob_desc.dtype}" + ) + if self.norm_const_desc is not None: + require_tensor_shape(self.norm_const_desc, (1,), label="norm_const") + if self.norm_const_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"norm_const must have float32 dtype, got {self.norm_const_desc.dtype}" + ) + + d_cudnn_dtype = self._output_cudnn_dtype(self.d_dtype) + if self.ab_dtype == data_type.FP4_E2M1: + if d_cudnn_dtype not in ( + data_type.HALF, + data_type.BFLOAT16, + data_type.FLOAT, + ): + raise ValueError(f"FP4 A and B do not support D dtype {self.d_dtype}") + elif d_cudnn_dtype not in ( + data_type.HALF, + data_type.BFLOAT16, + *FP8_DTYPES, + data_type.FP4_E2M1, + ): + raise ValueError(f"FP8 A and B do not support D dtype {self.d_dtype}") + + self.generate_sfd = ( + self.ab_dtype in FP8_DTYPES + and self.sf_dtype == data_type.FP8_E8M0 + and d_cudnn_dtype in FP8_DTYPES + ) + if self.generate_sfd and self.norm_const_desc is None: + raise ValueError( + "FP8 A/B and FP8 D require sample_norm_const so the kernel can generate SFD" + ) + if not self.generate_sfd and self.norm_const_desc is not None: + raise ValueError( + "sample_norm_const is only used by FP8 A/B to FP8 D configurations" + ) + if self.discrete_col_sfd and not self.generate_sfd: + raise ValueError("discrete_col_sfd requires generated SFD outputs") + + self.d_row_desc = self._canonical_desc( + (self.m, n_out, 1), self.d_dtype, "d_row_tensor", mode=self.output_mode + ) + self.d_col_desc = self._canonical_desc( + (self.m, n_out, 1), self.d_dtype, "d_col_tensor", mode=self.output_mode + ) + for desc, label in ((self.d_row_desc, "D_row"), (self.d_col_desc, "D_col")): + if require_compact_major(desc, "m", "n") != "n": + raise ValueError(f"{label} must use an N-major output layout") + require_16_byte_alignment(desc) + + self.dprob_desc = self._canonical_desc( + (self.m, 1, 1), + jnp.float32, + "dprob_tensor", + mode=self.probability_mode, + init_value=0.0, + ) + self.dbias_desc = ( + self._canonical_desc( + (self.expert_cnt, n_out, 1), + jnp.bfloat16, + "dbias_tensor", + init_value=0.0, + ) + if self.generate_dbias + else None + ) + if self.generate_sfd: + self.sfd_row_desc = self._canonical_desc( + block_scale_shape(self.m, n_out, 1, self.sf_vec_size), + self.sfa_desc.dtype, + "sfd_row_tensor", + mode=self.scale_mode, + ) + self.sfd_col_desc = self._canonical_desc( + block_scale_shape(n_out, self.m, 1, self.sf_vec_size), + self.sfa_desc.dtype, + "sfd_col_tensor", + mode=self.scale_mode, + ) + else: + self.sfd_row_desc = self.sfd_col_desc = None + self.amax_desc = ( + self._canonical_desc( + (self.expert_cnt, 2, 1), + jnp.float32, + "amax_tensor", + init_value=float("-inf"), + ) + if d_cudnn_dtype in (data_type.HALF, data_type.BFLOAT16) + else None + ) + return True + + @staticmethod + def _output_cudnn_dtype(dtype: Any) -> data_type: + from ..._jax.datatypes import jax_to_cudnn_dtype + + resolved = jax_to_cudnn_dtype(dtype) + allowed = { + data_type.FLOAT, + data_type.HALF, + data_type.BFLOAT16, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + data_type.FP4_E2M1, + } + if resolved not in allowed: + raise ValueError(f"d_dtype has unsupported dtype {dtype}") + return resolved + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Any, + prob_tensor: Any, + norm_const_tensor: Any | None = None, + *, + linear_offset: float | None = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + ) -> TupleDict: + self.check_support() + self._check_runtime_common( + a_tensor, b_tensor, sfa_tensor, sfb_tensor, padded_offsets, alpha_tensor + ) + self._check_tensor_signature(c_tensor, self.c_desc, mode=self.output_mode) + self._check_tensor_signature(beta_tensor, self.beta_desc) + self._check_tensor_signature( + prob_tensor, self.prob_desc, mode=self.probability_mode + ) + if (norm_const_tensor is None) != (self.norm_const_desc is None): + raise ValueError("norm_const_tensor presence must match sample_norm_const") + if self.norm_const_desc is not None: + self._check_tensor_signature(norm_const_tensor, self.norm_const_desc) + if linear_offset is None: + linear_offset = 1.0 if self.act_func == "dgeglu" else 0.0 + + if self.m == 0: + return TupleDict( + d_row_tensor=self._materialize_output_desc( + self.d_row_desc, mode=self.output_mode + ), + d_col_tensor=self._materialize_output_desc( + self.d_col_desc, mode=self.output_mode + ), + dprob_tensor=self._materialize_output_desc( + self.dprob_desc, mode=self.probability_mode + ), + dbias_tensor=self._materialize_output_desc(self.dbias_desc), + amax_tensor=self._materialize_output_desc(self.amax_desc), + sfd_row_tensor=self._materialize_output_desc( + self.sfd_row_desc, mode=self.scale_mode + ), + sfd_col_tensor=self._materialize_output_desc( + self.sfd_col_desc, mode=self.scale_mode + ), + ) + + import cutlass + import cutlass.cute as cute + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from .discrete_B_blockscaled_grouped_gemm_dglu_dbias import ( + BlockScaledDiscreteWeightDgluDbiasGroupedGemmKernel, + ) + + kernel = BlockScaledDiscreteWeightDgluDbiasGroupedGemmKernel( + sf_vec_size=self.sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(self.acc_dtype), + use_2cta_instrs=self.mma_tiler_mn[0] == 256, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=self.cluster_shape_mn, + vectorized_f32=self.vector_f32, + discrete_col_sfd=self.discrete_col_sfd, + expert_cnt=self.expert_cnt, + use_dynamic_sched=self.use_dynamic_sched, + act_func=self.act_func, + stacked_expert_inputs=True, + ) + max_active_clusters = self._get_max_active_clusters( + self.cluster_shape_mn[0] * self.cluster_shape_mn[1], + overlap_margin=self.num_cluster_overlap_margin, + ) + workspace_desc = self._workspace_desc(kernel.get_workspace_bytes()) + + inputs = [ + a_tensor, + b_tensor, + sfb_tensor, + c_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + ] + input_specs = [ + self._to_tensor_spec(self.a_desc, mode=self.a_mode), + self._to_tensor_spec(self.b_desc, mode=self.b_mode), + self._to_tensor_spec(self.sfb_desc, mode=self.scale_mode), + self._to_tensor_spec(self.c_desc, mode=self.output_mode), + self._to_tensor_spec(self.sfa_desc, mode=self.scale_mode), + self._to_tensor_spec(self.padded_offsets_desc), + self._to_tensor_spec(self.alpha_desc), + self._to_tensor_spec(self.beta_desc), + self._to_tensor_spec(self.prob_desc, mode=self.probability_mode), + ] + if self.norm_const_desc is not None: + inputs.append(norm_const_tensor) + input_specs.append(self._to_tensor_spec(self.norm_const_desc)) + + output_descs = [self.d_row_desc, self.d_col_desc, self.dprob_desc] + output_specs = [ + self._to_tensor_spec(self.d_row_desc, mode=self.output_mode), + self._to_tensor_spec(self.d_col_desc, mode=self.output_mode), + self._to_tensor_spec(self.dprob_desc, mode=self.probability_mode), + ] + for desc, mode in ( + (self.dbias_desc, None), + (self.sfd_row_desc, self.scale_mode), + (self.sfd_col_desc, self.scale_mode), + (self.amax_desc, None), + ): + if desc is not None: + output_descs.append(desc) + output_specs.append(self._to_tensor_spec(desc, mode=mode)) + + has_norm = self.norm_const_desc is not None + has_dbias = self.dbias_desc is not None + has_sfd = self.sfd_row_desc is not None + has_amax = self.amax_desc is not None + + def launch(stream: Any, *args: Any) -> None: + arg_index = 0 + + def take() -> Any: + nonlocal arg_index + value = args[arg_index] + arg_index += 1 + return value + + a, b, sfb, c, sfa, offsets, alpha, beta, prob = (take() for _ in range(9)) + norm_const = take() if has_norm else None + d_row, d_col, dprob = (take() for _ in range(3)) + dbias = take() if has_dbias else None + sfd_row = take() if has_sfd else None + sfd_col = take() if has_sfd else None + amax = take() if has_amax else None + workspace = take() + if arg_index != len(args): + raise RuntimeError( + f"Unexpected discrete grouped dSwiGLU buffer count: consumed {arg_index}, received {len(args)}" + ) + + if self.epilogue_op in (None, "none", "identity"): + + def epilogue(x): + return x + elif self.epilogue_op == "relu": + + def epilogue(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) + else: + + def epilogue(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) ** 2 + + kernel( + a, + b.iterator, + sfb.iterator, + cutlass.Int32(self.n), + cutlass.Int32(self.k), + cutlass.Int64(self.k), + OperandMajorMode.K, + workspace.iterator, + c, + d_row, + d_col, + sfa, + sfd_row, + sfd_col, + amax, + norm_const, + offsets, + alpha, + beta, + prob, + dprob, + cutlass.Float32(linear_offset), + dbias, + max_active_clusters, + stream, + epilogue, + cutlass.Float32(geglu_alpha), + cutlass.Float32(glu_clamp_max), + cutlass.Float32(glu_clamp_min), + ) + + results = self._call_kernel( + tuple(inputs), + launch=launch, + output_descs=tuple(output_descs), + workspace_descs=(workspace_desc,), + input_spec=tuple(input_specs), + output_spec=tuple(output_specs), + workspace_spec=(self._workspace_tensor_spec(),), + compile_options=compile_options_for_target(self.compute_capability), + ) + result_index = 0 + + def result() -> Any: + nonlocal result_index + value = results[result_index] + result_index += 1 + return value + + d_row_result, d_col_result, dprob_result = result(), result(), result() + dbias_result = result() if has_dbias else None + sfd_row_result = result() if has_sfd else None + sfd_col_result = result() if has_sfd else None + amax_result = result() if has_amax else None + return TupleDict( + d_row_tensor=d_row_result, + d_col_tensor=d_col_result, + dprob_tensor=dprob_result, + dbias_tensor=dbias_result, + amax_tensor=amax_result, + sfd_row_tensor=sfd_row_result, + sfd_col_tensor=sfd_col_result, + ) + + +@partial( + jax.jit, + static_argnames=( + "d_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "generate_dbias", + "act_func", + "epilogue_op", + "use_dynamic_sched", + "a_layout", + "b_layout", + "output_layout", + "linear_offset", + "geglu_alpha", + "glu_clamp_max", + "glu_clamp_min", + ), +) +def discrete_grouped_gemm_dswiglu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Any, + prob_tensor: Any, + norm_const_tensor: Any | None = None, + *, + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + generate_dbias: bool = False, + act_func: str = "dswiglu", + epilogue_op: str | None = None, + use_dynamic_sched: bool = False, + a_layout: str = "LMK", + b_layout: str = "LNK", + output_layout: str = "LMN", + linear_offset: float | None = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, +) -> TupleDict: + """Run discrete dSwiGLU with XLA-owned stacked expert operands.""" + + operation = DiscreteGroupedGemmDswigluSm100( + jax.ShapeDtypeStruct(a_tensor.shape, a_tensor.dtype), + jax.ShapeDtypeStruct(b_tensor.shape, b_tensor.dtype), + jax.ShapeDtypeStruct(c_tensor.shape, c_tensor.dtype), + jax.ShapeDtypeStruct(sfa_tensor.shape, sfa_tensor.dtype), + jax.ShapeDtypeStruct(sfb_tensor.shape, sfb_tensor.dtype), + jax.ShapeDtypeStruct(padded_offsets.shape, padded_offsets.dtype), + jax.ShapeDtypeStruct(alpha_tensor.shape, alpha_tensor.dtype), + jax.ShapeDtypeStruct(beta_tensor.shape, beta_tensor.dtype), + jax.ShapeDtypeStruct(prob_tensor.shape, prob_tensor.dtype), + sample_norm_const=( + None + if norm_const_tensor is None + else jax.ShapeDtypeStruct(norm_const_tensor.shape, norm_const_tensor.dtype) + ), + d_dtype=d_dtype, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + generate_dbias=generate_dbias, + act_func=act_func, + epilogue_op=epilogue_op, + use_dynamic_sched=use_dynamic_sched, + a_layout=a_layout, + b_layout=b_layout, + output_layout=output_layout, + ) + return operation( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + norm_const_tensor, + linear_offset=linear_offset, + geglu_alpha=geglu_alpha, + glu_clamp_max=glu_clamp_max, + glu_clamp_min=glu_clamp_min, + ) + + +__all__ = [ + "DiscreteGroupedGemmDswigluSm100", + "SUPPORTED_COMPUTE_CAPABILITIES", + "discrete_grouped_gemm_dswiglu_wrapper_sm100", +] diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/__init__.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/__init__.py index dafdfc1c2..d6d7a3da6 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/__init__.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/__init__.py @@ -1,12 +1,17 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - DiscreteGroupedGemmSwigluSm100, - discrete_grouped_gemm_swiglu_wrapper_sm100, -) +"""Lazy Torch API exports for discrete grouped GEMM + SwiGLU.""" + +from ..._operation_api import make_operation_api -__all__ = [ - "DiscreteGroupedGemmSwigluSm100", - "discrete_grouped_gemm_swiglu_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "api": ( + "DiscreteGroupedGemmSwigluSm100", + "discrete_grouped_gemm_swiglu_wrapper_sm100", + ), + }, + submodules=("api",), +) diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/discrete_B_blockscaled_grouped_gemm_glu_bias.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/discrete_B_blockscaled_grouped_gemm_glu_bias.py index 316cf3228..ca55a97a2 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/discrete_B_blockscaled_grouped_gemm_glu_bias.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/discrete_B_blockscaled_grouped_gemm_glu_bias.py @@ -277,6 +277,7 @@ def __init__( use_dynamic_sched: bool = False, act_func: str = "swiglu", enable_bias: bool = False, + stacked_expert_inputs: bool = False, ): """Initializes the configuration for a Blackwell blockscaled dense GEMM kernel. @@ -331,6 +332,14 @@ def __init__( self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE self.enable_bias = enable_bias + # JAX cannot safely pass an integer table containing addresses of + # otherwise-unreferenced buffers: XLA would not know that those + # buffers must remain live. In stacked mode, ``b_ptrs`` and + # ``sfb_ptrs`` are instead the base pointers of explicit B/SFB custom + # call operands. The descriptor initializer derives each expert's + # address from the compact outer expert dimension. Torch retains the + # original pointer-table ABI by using the default value. + self.stacked_expert_inputs = stacked_expert_inputs self.occupancy = 1 self.epilog_warp_id = (0, 1, 2, 3) self.mma_warp_id = 4 @@ -641,10 +650,15 @@ def desc_init_kernel_device_ptrs( expert_idx = cute.arch.block_idx()[0] - b_ptr_tensor = cute.make_tensor(cute.make_ptr(cutlass.Int64, ptrs_b.toint(), AddressSpace.gmem, assumed_align=8), cute.make_layout((self.expert_cnt,))) - sfb_ptr_tensor = cute.make_tensor( - cute.make_ptr(cutlass.Int64, ptrs_sfb.toint(), AddressSpace.gmem, assumed_align=8), cute.make_layout((self.expert_cnt,)) - ) + if cutlass.const_expr(not self.stacked_expert_inputs): + b_ptr_tensor = cute.make_tensor( + cute.make_ptr(cutlass.Int64, ptrs_b.toint(), AddressSpace.gmem, assumed_align=8), + cute.make_layout((self.expert_cnt,)), + ) + sfb_ptr_tensor = cute.make_tensor( + cute.make_ptr(cutlass.Int64, ptrs_sfb.toint(), AddressSpace.gmem, assumed_align=8), + cute.make_layout((self.expert_cnt,)), + ) c1 = cutlass.Int32(1) c0 = cutlass.Int64(0) @@ -657,8 +671,12 @@ def desc_init_kernel_device_ptrs( stride_k = b_stride_size # --- B TMA descriptor --- - b_ptr_val = b_ptr_tensor[expert_idx] - b_ptr = cute.make_ptr(self.b_dtype, b_ptr_val, AddressSpace.gmem) + if cutlass.const_expr(self.stacked_expert_inputs): + b_base_ptr = cute.make_ptr(self.b_dtype, ptrs_b.toint(), AddressSpace.gmem, assumed_align=16) + b_ptr = b_base_ptr + cutlass.Int64(expert_idx) * cutlass.Int64(n) * cutlass.Int64(k) + else: + b_ptr_val = b_ptr_tensor[expert_idx] + b_ptr = cute.make_ptr(self.b_dtype, b_ptr_val, AddressSpace.gmem) b_tensor_i = cute.make_tensor( b_ptr, cute.make_layout((n, k, c1), stride=(stride_n, stride_k, c0)), @@ -676,9 +694,13 @@ def desc_init_kernel_device_ptrs( store_tma_desc(tma_atom_b, workspace.get_ptr("b", expert_idx)) # --- SFB TMA descriptor --- - sfb_ptr_val = sfb_ptr_tensor[expert_idx] - sfb_ptr = cute.make_ptr(self.sf_dtype, sfb_ptr_val, AddressSpace.gmem) sfb_layout = blockscaled_utils.tile_atom_to_shape_SF((n, k, c1), self.sf_vec_size) + if cutlass.const_expr(self.stacked_expert_inputs): + sfb_base_ptr = cute.make_ptr(self.sf_dtype, ptrs_sfb.toint(), AddressSpace.gmem, assumed_align=16) + sfb_ptr = sfb_base_ptr + cutlass.Int64(expert_idx) * cutlass.Int64(cute.cosize(sfb_layout)) + else: + sfb_ptr_val = sfb_ptr_tensor[expert_idx] + sfb_ptr = cute.make_ptr(self.sf_dtype, sfb_ptr_val, AddressSpace.gmem) sfb_tensor_i = cute.make_tensor(sfb_ptr, sfb_layout) tma_atom_sfb, _ = cute.nvgpu.make_tiled_tma_atom_B( sfb_tma_op_arg, diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/jax.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/jax.py new file mode 100644 index 000000000..1cef279d6 --- /dev/null +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/jax.py @@ -0,0 +1,574 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for the discrete-weight grouped GEMM + SwiGLU kernel.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._cute_compiler import compile_options_for_target +from ..._dense_gemm import ( + block_scale_shape, + require_16_byte_alignment, + require_compact_major, + require_tensor_shape, +) +from ..._jax import JaxTensorDesc, TupleDict +from ..._jax.datatypes import normalize_jax_dtype +from .._jax_common import ( + DiscreteGroupedGemmJaxBase, + FP8_DTYPES, + SUPPORTED_COMPUTE_CAPABILITIES, +) + + +class DiscreteGroupedGemmSwigluSm100(DiscreteGroupedGemmJaxBase): + """JAX callable for grouped GEMM + SwiGLU using stacked expert weights. + + ``sample_b`` and ``sample_sfb`` retain an explicit expert dimension. They + are passed to XLA as ordinary operands, so their storage remains live for + the custom call. During lowering the discrete kernel derives per-expert + addresses and builds the same TMA descriptor workspace used by the Torch + pointer-table API. + """ + + def __init__( + self, + sample_a: Any, + sample_b: Any, + sample_sfa: Any, + sample_sfb: Any, + sample_padded_offsets: Any, + sample_alpha: Any, + *, + sample_norm_const: Any | None = None, + sample_prob: Any | None = None, + sample_bias: Any | None = None, + c_dtype: Any | None = None, + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "swiglu", + use_dynamic_sched: bool = False, + a_layout: str = "LMK", + b_layout: str = "LNK", + output_layout: str = "LMN", + ) -> None: + self._initialize_common( + sample_a, + sample_b, + sample_sfa, + sample_sfb, + sample_padded_offsets, + sample_alpha, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + use_dynamic_sched=use_dynamic_sched, + a_layout=a_layout, + b_layout=b_layout, + output_layout=output_layout, + ) + self.norm_const_desc = ( + None + if sample_norm_const is None + else self._to_tensor_desc(sample_norm_const, "sample_norm_const") + ) + self._uses_implicit_prob = sample_prob is None + if sample_prob is None: + sample_prob = jax.ShapeDtypeStruct( + (1, 1, self.a_desc.shape[0]), jnp.float32 + ) + self.prob_desc = self._to_tensor_desc( + sample_prob, "sample_prob", mode=self.probability_mode + ) + self.bias_desc = ( + None + if sample_bias is None + else self._to_tensor_desc(sample_bias, "sample_bias", mode=self.bias_mode) + ) + self.c_dtype = normalize_jax_dtype(c_dtype, jnp.bfloat16, "c_dtype") + self.d_dtype = normalize_jax_dtype(d_dtype, jnp.bfloat16, "d_dtype") + self.discrete_col_sfd = discrete_col_sfd + self.act_func = act_func + + self.generate_sfd = self.norm_const_desc is not None + self.c_desc = self.d_desc = self.d_col_desc = None + self.sfd_row_desc = self.sfd_col_desc = self.amax_desc = None + + def check_support(self) -> bool: + self._check_common() + if self.act_func not in ("swiglu", "geglu"): + raise ValueError( + f"act_func must be 'swiglu' or 'geglu', got {self.act_func!r}" + ) + if not isinstance(self.discrete_col_sfd, bool): + raise TypeError( + f"discrete_col_sfd must be a bool, got {type(self.discrete_col_sfd).__name__}" + ) + if self.discrete_col_sfd and not self.generate_sfd: + raise ValueError( + "discrete_col_sfd requires sample_norm_const and generated SFD outputs" + ) + + if self.norm_const_desc is not None: + require_tensor_shape(self.norm_const_desc, (1,), label="norm_const") + if self.norm_const_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"norm_const must have float32 dtype, got {self.norm_const_desc.dtype}" + ) + require_tensor_shape(self.prob_desc, (self.m, 1, 1), label="prob") + if self.prob_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"prob must have float32 dtype, got {self.prob_desc.dtype}" + ) + if self.bias_desc is not None: + require_tensor_shape( + self.bias_desc, (self.n, self.expert_cnt), label="bias" + ) + if self.bias_desc.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError( + f"bias must have float16 or bfloat16 dtype, got {self.bias_desc.dtype}" + ) + + c_cudnn_dtype = self._output_cudnn_dtype(self.c_dtype, "c_dtype") + d_cudnn_dtype = self._output_cudnn_dtype(self.d_dtype, "d_dtype") + if self.ab_dtype == data_type.FP4_E2M1: + if c_cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise NotImplementedError("FP4 A and B require float16 or bfloat16 C") + if d_cudnn_dtype not in ( + data_type.HALF, + data_type.BFLOAT16, + data_type.FLOAT, + ): + raise ValueError(f"FP4 A and B do not support D dtype {self.d_dtype}") + if self.sf_vec_size == 16 and d_cudnn_dtype == data_type.FLOAT: + raise NotImplementedError( + "FP4 A and B with sf_vec_size=16 do not support float32 D" + ) + elif d_cudnn_dtype not in ( + data_type.HALF, + data_type.BFLOAT16, + *FP8_DTYPES, + data_type.FP4_E2M1, + ): + raise ValueError(f"FP8 A and B do not support D dtype {self.d_dtype}") + + output_n = self.n // 2 + self.c_desc = self._canonical_desc( + (self.m, self.n, 1), self.c_dtype, "c_tensor", mode=self.output_mode + ) + self.d_desc = self._canonical_desc( + (self.m, output_n, 1), self.d_dtype, "d_tensor", mode=self.output_mode + ) + self.d_col_desc = self._canonical_desc( + (self.m, output_n, 1), self.d_dtype, "d_col_tensor", mode=self.output_mode + ) + for desc, label in ( + (self.c_desc, "C"), + (self.d_desc, "D"), + (self.d_col_desc, "D_col"), + ): + if require_compact_major(desc, "m", "n") != "n": + raise ValueError(f"{label} must use an N-major output layout") + require_16_byte_alignment(desc) + + if self.generate_sfd: + self.sfd_row_desc = self._canonical_desc( + block_scale_shape(self.m, output_n, 1, self.sf_vec_size), + self.sfa_desc.dtype, + "sfd_row_tensor", + mode=self.scale_mode, + ) + self.sfd_col_desc = self._canonical_desc( + block_scale_shape(output_n, self.m, 1, self.sf_vec_size), + self.sfa_desc.dtype, + "sfd_col_tensor", + mode=self.scale_mode, + ) + else: + self.sfd_row_desc = self.sfd_col_desc = None + self.amax_desc = ( + self._canonical_desc( + (self.expert_cnt, 1), + jnp.float32, + "amax_tensor", + init_value=float("-inf"), + ) + if d_cudnn_dtype in (data_type.HALF, data_type.BFLOAT16) + else None + ) + return True + + @staticmethod + def _output_cudnn_dtype(dtype: Any, name: str) -> data_type: + from ..._jax.datatypes import jax_to_cudnn_dtype + + resolved = jax_to_cudnn_dtype(dtype) + allowed = { + data_type.FLOAT, + data_type.HALF, + data_type.BFLOAT16, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + data_type.FP4_E2M1, + } + if resolved not in allowed: + raise ValueError(f"{name} has unsupported dtype {dtype}") + return resolved + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + norm_const_tensor: Any | None = None, + prob_tensor: Any | None = None, + bias_tensor: Any | None = None, + *, + linear_offset: float | None = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + ) -> TupleDict: + self.check_support() + self._check_runtime_common( + a_tensor, b_tensor, sfa_tensor, sfb_tensor, padded_offsets, alpha_tensor + ) + self._check_optional_input(norm_const_tensor, self.norm_const_desc) + if self._uses_implicit_prob: + if prob_tensor is not None: + raise ValueError( + "prob_tensor must be None when sample_prob was omitted" + ) + prob_tensor = jnp.ones((1, 1, self.m), dtype=jnp.float32) + else: + self._check_tensor_signature( + prob_tensor, self.prob_desc, mode=self.probability_mode + ) + self._check_optional_input(bias_tensor, self.bias_desc, mode=self.bias_mode) + if linear_offset is None: + linear_offset = 1.0 if self.act_func == "geglu" else 0.0 + + if self.m == 0: + return TupleDict( + c_tensor=self._materialize_output_desc( + self.c_desc, mode=self.output_mode + ), + d_tensor=self._materialize_output_desc( + self.d_desc, mode=self.output_mode + ), + d_col_tensor=self._materialize_output_desc( + self.d_col_desc, mode=self.output_mode + ), + amax_tensor=self._materialize_output_desc(self.amax_desc), + sfd_row_tensor=self._materialize_output_desc( + self.sfd_row_desc, mode=self.scale_mode + ), + sfd_col_tensor=self._materialize_output_desc( + self.sfd_col_desc, mode=self.scale_mode + ), + ) + + import cutlass + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from .discrete_B_blockscaled_grouped_gemm_glu_bias import ( + BlockScaledDiscreteWeightGroupedGemmBiasKernel, + ) + + kernel = BlockScaledDiscreteWeightGroupedGemmBiasKernel( + sf_vec_size=self.sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(self.acc_dtype), + use_2cta_instrs=self.mma_tiler_mn[0] == 256, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=self.cluster_shape_mn, + vectorized_f32=self.vector_f32, + generate_sfd=self.generate_sfd, + discrete_col_sfd=self.discrete_col_sfd, + expert_cnt=self.expert_cnt, + use_dynamic_sched=self.use_dynamic_sched, + act_func=self.act_func, + enable_bias=self.bias_desc is not None, + stacked_expert_inputs=True, + ) + max_active_clusters = self._get_max_active_clusters( + self.cluster_shape_mn[0] * self.cluster_shape_mn[1], + overlap_margin=self.num_cluster_overlap_margin, + ) + workspace_desc = self._workspace_desc(kernel.get_workspace_bytes()) + + inputs = [ + a_tensor, + b_tensor, + sfb_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + ] + input_specs = [ + self._to_tensor_spec(self.a_desc, mode=self.a_mode), + self._to_tensor_spec(self.b_desc, mode=self.b_mode), + self._to_tensor_spec(self.sfb_desc, mode=self.scale_mode), + self._to_tensor_spec(self.sfa_desc, mode=self.scale_mode), + self._to_tensor_spec(self.padded_offsets_desc), + self._to_tensor_spec(self.alpha_desc), + ] + for value, desc, mode in ( + (norm_const_tensor, self.norm_const_desc, None), + (prob_tensor, self.prob_desc, self.probability_mode), + (bias_tensor, self.bias_desc, self.bias_mode), + ): + if desc is not None: + inputs.append(value) + input_specs.append(self._to_tensor_spec(desc, mode=mode)) + + output_descs = [self.c_desc, self.d_desc, self.d_col_desc] + output_specs = [ + self._to_tensor_spec(self.c_desc, mode=self.output_mode), + self._to_tensor_spec(self.d_desc, mode=self.output_mode), + self._to_tensor_spec(self.d_col_desc, mode=self.output_mode), + ] + for desc, mode in ( + (self.sfd_row_desc, self.scale_mode), + (self.sfd_col_desc, self.scale_mode), + (self.amax_desc, None), + ): + if desc is not None: + output_descs.append(desc) + output_specs.append(self._to_tensor_spec(desc, mode=mode)) + + has_norm = self.norm_const_desc is not None + has_prob = True + has_bias = self.bias_desc is not None + has_sfd = self.sfd_row_desc is not None + has_amax = self.amax_desc is not None + + def launch(stream: Any, *args: Any) -> None: + arg_index = 0 + + def take() -> Any: + nonlocal arg_index + value = args[arg_index] + arg_index += 1 + return value + + a, b, sfb, sfa, offsets, alpha = (take() for _ in range(6)) + norm_const = take() if has_norm else None + prob = take() if has_prob else None + bias = take() if has_bias else None + c, d, d_col = (take() for _ in range(3)) + sfd_row = take() if has_sfd else None + sfd_col = take() if has_sfd else None + amax = take() if has_amax else None + workspace = take() + if arg_index != len(args): + raise RuntimeError( + f"Unexpected discrete grouped SwiGLU buffer count: consumed {arg_index}, received {len(args)}" + ) + + kernel( + a, + b.iterator, + sfb.iterator, + cutlass.Int32(self.n), + cutlass.Int32(self.k), + cutlass.Int64(self.k), + OperandMajorMode.K, + workspace.iterator, + c, + d, + d_col, + sfa, + sfd_row, + sfd_col, + amax, + norm_const, + offsets, + alpha, + prob, + bias, + max_active_clusters, + stream, + lambda x: x, + cutlass.Float32(linear_offset), + cutlass.Float32(geglu_alpha), + cutlass.Float32(glu_clamp_max), + cutlass.Float32(glu_clamp_min), + ) + + results = self._call_kernel( + tuple(inputs), + launch=launch, + output_descs=tuple(output_descs), + workspace_descs=(workspace_desc,), + input_spec=tuple(input_specs), + output_spec=tuple(output_specs), + workspace_spec=(self._workspace_tensor_spec(),), + compile_options=compile_options_for_target(self.compute_capability), + ) + result_index = 0 + + def result() -> Any: + nonlocal result_index + value = results[result_index] + result_index += 1 + return value + + c_result, d_result, d_col_result = result(), result(), result() + sfd_row_result = result() if has_sfd else None + sfd_col_result = result() if has_sfd else None + amax_result = result() if has_amax else None + return TupleDict( + c_tensor=c_result, + d_tensor=d_result, + d_col_tensor=d_col_result, + amax_tensor=amax_result, + sfd_row_tensor=sfd_row_result, + sfd_col_tensor=sfd_col_result, + ) + + def _check_optional_input( + self, + value: Any | None, + desc: JaxTensorDesc | None, + *, + mode: tuple[int, ...] | None = None, + ) -> None: + if (value is None) != (desc is None): + name = "optional input" if desc is None else desc.name + raise ValueError( + f"{name} presence must match the sample passed to the constructor" + ) + if desc is not None: + self._check_tensor_signature(value, desc, mode=mode) + + +@partial( + jax.jit, + static_argnames=( + "c_dtype", + "d_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "act_func", + "use_dynamic_sched", + "a_layout", + "b_layout", + "output_layout", + "linear_offset", + "geglu_alpha", + "glu_clamp_max", + "glu_clamp_min", + ), +) +def discrete_grouped_gemm_swiglu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + norm_const_tensor: Any | None = None, + prob_tensor: Any | None = None, + bias_tensor: Any | None = None, + *, + c_dtype: Any | None = None, + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "swiglu", + use_dynamic_sched: bool = False, + a_layout: str = "LMK", + b_layout: str = "LNK", + output_layout: str = "LMN", + linear_offset: float | None = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, +) -> TupleDict: + """Run the discrete kernel with XLA-owned stacked expert operands.""" + + operation = DiscreteGroupedGemmSwigluSm100( + jax.ShapeDtypeStruct(a_tensor.shape, a_tensor.dtype), + jax.ShapeDtypeStruct(b_tensor.shape, b_tensor.dtype), + jax.ShapeDtypeStruct(sfa_tensor.shape, sfa_tensor.dtype), + jax.ShapeDtypeStruct(sfb_tensor.shape, sfb_tensor.dtype), + jax.ShapeDtypeStruct(padded_offsets.shape, padded_offsets.dtype), + jax.ShapeDtypeStruct(alpha_tensor.shape, alpha_tensor.dtype), + sample_norm_const=( + None + if norm_const_tensor is None + else jax.ShapeDtypeStruct(norm_const_tensor.shape, norm_const_tensor.dtype) + ), + sample_prob=None + if prob_tensor is None + else jax.ShapeDtypeStruct(prob_tensor.shape, prob_tensor.dtype), + sample_bias=None + if bias_tensor is None + else jax.ShapeDtypeStruct(bias_tensor.shape, bias_tensor.dtype), + c_dtype=c_dtype, + d_dtype=d_dtype, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + act_func=act_func, + use_dynamic_sched=use_dynamic_sched, + a_layout=a_layout, + b_layout=b_layout, + output_layout=output_layout, + ) + return operation( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + norm_const_tensor, + prob_tensor, + bias_tensor, + linear_offset=linear_offset, + geglu_alpha=geglu_alpha, + glu_clamp_max=glu_clamp_max, + glu_clamp_min=glu_clamp_min, + ) + + +__all__ = [ + "DiscreteGroupedGemmSwigluSm100", + "SUPPORTED_COMPUTE_CAPABILITIES", + "discrete_grouped_gemm_swiglu_wrapper_sm100", +] diff --git a/python/cudnn/discrete_grouped_gemm/discrete_kernel_utils.py b/python/cudnn/discrete_grouped_gemm/discrete_kernel_utils.py index 992747d0c..28d14bb01 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_kernel_utils.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_kernel_utils.py @@ -38,9 +38,12 @@ - Kernel helper functions that don't depend on kernel instance state """ -from typing import Type, Tuple, Union +from __future__ import annotations -import torch +from typing import TYPE_CHECKING, Any, Type, Tuple, Union + +if TYPE_CHECKING: + import torch import cutlass import cutlass.cute as cute @@ -69,7 +72,9 @@ # --------------------------------------------------------------------------- -def _require_pointer_tensor(ptrs: torch.Tensor, name: str, expected_len: int | None = None) -> None: +def _require_pointer_tensor(ptrs: Any, name: str, expected_len: int | None = None) -> None: + import torch + if ptrs.dtype != torch.int64: raise ValueError(f"{name} must be int64, got {ptrs.dtype}") if ptrs.ndim != 1: @@ -399,6 +404,8 @@ def silu_f32_geglu_scaled(a: Union[float, Float32], fastmath: bool = False) -> U def sigmoid(x): """PyTorch reference sigmoid using exp2 for numerical consistency.""" + import torch + LOG2_E = 1.4426950408889634 exp_x = torch.exp2(x * (-LOG2_E)) ret = 1.0 / (exp_x + 1.0) @@ -415,6 +422,8 @@ def compute_reference_amax(output_tensor: torch.Tensor) -> float: Returns: float: reference amax value """ + import torch + if output_tensor.dtype != torch.float32: output_fp32 = output_tensor.float() else: @@ -444,6 +453,8 @@ def compare_and_report_mismatches( rtol: Relative tolerance max_mismatches: Maximum number of mismatches to report """ + import torch + if gpu_tensor.is_cuda: gpu_data = gpu_tensor.cpu() else: diff --git a/python/cudnn/grouped_gemm/__init__.py b/python/cudnn/grouped_gemm/__init__.py index 818a1ab99..4a0401461 100644 --- a/python/cudnn/grouped_gemm/__init__.py +++ b/python/cudnn/grouped_gemm/__init__.py @@ -1,68 +1,30 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .grouped_gemm_swiglu.api import ( - GroupedGemmSwigluSm100, - grouped_gemm_swiglu_wrapper_sm100, +"""Lazy Torch API exports for the grouped GEMM operation family.""" + +from .._operation_api import make_operation_api + +_OPERATIONS = ( + "swiglu", + "dswiglu", + "quant", + "srelu", + "dsrelu", + "glu", + "glu_hadamard", + "dglu", + "wgrad", +) + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + f"grouped_gemm_{operation}.api": ( + f"GroupedGemm{''.join(part.capitalize() for part in operation.split('_'))}Sm100", + f"grouped_gemm_{operation}_wrapper_sm100", + ) + for operation in _OPERATIONS + }, + submodules=tuple(f"grouped_gemm_{operation}" for operation in _OPERATIONS), ) - -from .grouped_gemm_dswiglu.api import ( - GroupedGemmDswigluSm100, - grouped_gemm_dswiglu_wrapper_sm100, -) - -from .grouped_gemm_quant.api import ( - GroupedGemmQuantSm100, - grouped_gemm_quant_wrapper_sm100, -) - -from .grouped_gemm_srelu.api import ( - GroupedGemmSreluSm100, - grouped_gemm_srelu_wrapper_sm100, -) - -from .grouped_gemm_dsrelu.api import ( - GroupedGemmDsreluSm100, - grouped_gemm_dsrelu_wrapper_sm100, -) - -from .grouped_gemm_glu.api import ( - GroupedGemmGluSm100, - grouped_gemm_glu_wrapper_sm100, -) - -from .grouped_gemm_glu_hadamard.api import ( - GroupedGemmGluHadamardSm100, - grouped_gemm_glu_hadamard_wrapper_sm100, -) - -from .grouped_gemm_dglu.api import ( - GroupedGemmDgluSm100, - grouped_gemm_dglu_wrapper_sm100, -) - -from .grouped_gemm_wgrad.api import ( - GroupedGemmWgradSm100, - grouped_gemm_wgrad_wrapper_sm100, -) - -__all__ = [ - "GroupedGemmSwigluSm100", - "grouped_gemm_swiglu_wrapper_sm100", - "GroupedGemmDswigluSm100", - "grouped_gemm_dswiglu_wrapper_sm100", - "GroupedGemmQuantSm100", - "grouped_gemm_quant_wrapper_sm100", - "GroupedGemmSreluSm100", - "grouped_gemm_srelu_wrapper_sm100", - "GroupedGemmDsreluSm100", - "grouped_gemm_dsrelu_wrapper_sm100", - "GroupedGemmGluSm100", - "grouped_gemm_glu_wrapper_sm100", - "GroupedGemmGluHadamardSm100", - "grouped_gemm_glu_hadamard_wrapper_sm100", - "GroupedGemmDgluSm100", - "grouped_gemm_dglu_wrapper_sm100", - "GroupedGemmWgradSm100", - "grouped_gemm_wgrad_wrapper_sm100", -] diff --git a/python/cudnn/grouped_gemm/_jax_api.py b/python/cudnn/grouped_gemm/_jax_api.py new file mode 100644 index 000000000..03d01dcfb --- /dev/null +++ b/python/cudnn/grouped_gemm/_jax_api.py @@ -0,0 +1,721 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared JAX metadata and lowering helpers for contiguous grouped GEMMs.""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any + +import jax +import jax.numpy as jnp +from cutlass.jax import TensorSpec + +from .._cute_compiler import compile_options_for_target +from .._dense_gemm import ( + block_scale_shape as _canonical_block_scale_shape, + data_type_bits, +) +from .._jax import JaxApiBase, JaxTensorDesc, TupleDict +from .._jax.gemm import BLOCK_SCALE_MODE, gemm_a_mode, gemm_b_mode, gemm_output_mode +from .._jax.layout import normalize_mode, to_public_axes + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) +MAX_EXPERTS = 1024 +FP8_SF_VEC_SIZE = 32 +TWO_CTA_MMA_TILER_M = 256 +FIX_PAD_SIZE = 256 +HADAMARD_SIZE = 16 +SF_VEC_SIZES = (16, 32) + +_NO_DEFAULT = object() + + +def as_dtype(value: Any) -> Any: + """Return a normalized JAX dtype without retaining an array value.""" + + if not isinstance(value, type) and hasattr(value, "dtype"): + value = value.dtype + return jnp.dtype(value) + + +def is_fp4_dtype(value: Any) -> bool: + """Return whether ``value`` is JAX's native logical FP4 dtype.""" + + return as_dtype(value) == jnp.dtype(jnp.float4_e2m1fn) + + +def is_fp8_dtype(value: Any) -> bool: + """Return whether ``value`` is one of the supported JAX FP8 data dtypes.""" + + return as_dtype(value) in { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + } + + +def is_low_precision_output_dtype(value: Any) -> bool: + """Return whether an output uses a native FP4 or FP8 dtype.""" + + return is_fp4_dtype(value) or is_fp8_dtype(value) + + +def require_dtype( + value: Any, + valid_dtypes: Iterable[Any], + *, + name: str = "dtype", + default: Any = _NO_DEFAULT, +) -> Any: + """Normalize and validate a JAX dtype argument.""" + + if value is None: + if default is _NO_DEFAULT: + raise ValueError(f"{name} must not be None") + value = default + dtype = as_dtype(value) + valid = tuple(as_dtype(item) for item in valid_dtypes) + if dtype not in valid: + supported = ", ".join(item.name for item in valid) + raise ValueError(f"{name} must be one of {{{supported}}}, got {dtype}") + return dtype + + +def require_array( + value: Any, + *, + name: str, + rank: int | Iterable[int] | None = None, + shape: Sequence[int] | None = None, + dtype: Any | Iterable[Any] | None = None, +) -> tuple[int, ...]: + """Validate abstract array metadata and return its public shape.""" + + if not hasattr(value, "shape") or not hasattr(value, "dtype"): + raise TypeError(f"{name} must have shape and dtype metadata") + metadata_shape = tuple(value.shape) + # Sample descriptors store canonical kernel axes. Exact shape checks in + # this adapter describe the public JAX value, so compare against the + # descriptor's inverse-mapped array shape when one is available. + actual_shape = ( + tuple(value.array_shape) + if shape is not None and hasattr(value, "array_shape") + else metadata_shape + ) + if rank is not None: + ranks = (rank,) if isinstance(rank, int) else tuple(rank) + if len(actual_shape) not in ranks: + raise ValueError( + f"{name} must have rank in {ranks}, got shape {actual_shape}" + ) + if shape is not None and actual_shape != tuple(shape): + raise ValueError(f"{name} must have shape {tuple(shape)}, got {actual_shape}") + if dtype is not None: + if ( + isinstance(dtype, Iterable) + and not isinstance(dtype, (str, bytes, type)) + and not hasattr(dtype, "dtype") + ): + valid_dtypes = tuple(dtype) + else: + valid_dtypes = (dtype,) + require_dtype(value, valid_dtypes, name=f"{name}.dtype") + return actual_shape + + +def require_layout(name: str, value: str, supported: tuple[str, ...]) -> str: + """Validate an explicit public axis-order string.""" + + if not isinstance(value, str): + raise TypeError(f"{name} must be a string, got {type(value).__name__}") + if value not in supported: + choices = ", ".join(repr(item) for item in supported) + raise ValueError(f"{name} must be one of ({choices}), got {value!r}") + return value + + +def _row_major_spec(rank: int, mode: tuple[int, ...]) -> TensorSpec: + return TensorSpec(layout=tuple(reversed(range(rank))), mode=mode) + + +def gemm_a_tensor_spec(layout: str) -> TensorSpec: + return _row_major_spec(3, gemm_a_mode(layout)) + + +def gemm_b_tensor_spec(layout: str) -> TensorSpec: + return _row_major_spec(3, gemm_b_mode(layout)) + + +def gemm_c_tensor_spec(layout: str, *, name: str = "c_layout") -> TensorSpec: + return _row_major_spec( + 3, gemm_output_mode(require_layout(name, layout, ("LMN", "LNM")), name=name) + ) + + +def block_scale_tensor_spec() -> TensorSpec: + """Describe public row-major ``L,tiles,tiles,32,4,4`` scale arrays.""" + + return _row_major_spec(6, BLOCK_SCALE_MODE) + + +def probability_tensor_spec() -> TensorSpec: + """Map public row-major ``(L, 1, M)`` to canonical ``(M, 1, L)``.""" + + return _row_major_spec(3, (2, 1, 0)) + + +def grouped_bias_tensor_spec() -> TensorSpec: + """Map public row-major ``(L, N)`` to canonical ``(N, L)``.""" + + return _row_major_spec(2, (1, 0)) + + +def grouped_workspace_tensor_spec() -> TensorSpec: + return TensorSpec(ptr_assumed_align=128) + + +@dataclass(frozen=True) +class _SampleDesc(JaxTensorDesc): + """JAX descriptor retaining the native CUTLASS lowering specification.""" + + tensor_spec: Any = None + jax_mode: tuple[int, ...] = field(default_factory=tuple) + + @property + def array_shape(self) -> tuple[int, ...]: + return to_public_axes(self.shape, self.jax_mode) + + +def _make_desc( + value: Any, + *, + tensor_spec: TensorSpec | None, + name: str, + init_value: bool | int | float | None = None, +) -> _SampleDesc: + public_shape = require_array(value, name=name) + mode = normalize_mode( + len(public_shape), None if tensor_spec is None else tensor_spec.mode + ) + desc = JaxApiBase._to_tensor_desc( + value, + name, + mode=mode, + init_value=init_value, + ) + return _SampleDesc( + dtype=desc.dtype, + shape=desc.shape, + stride=desc.stride, + stride_order=desc.stride_order, + name=desc.name, + init_value=desc.init_value, + tensor_spec=tensor_spec, + jax_mode=mode, + ) + + +def make_buffer_desc( + name: str, + shape: Sequence[int], + dtype: Any, + *, + tensor_spec: TensorSpec | None = None, + init_value: bool | int | float | None = None, +) -> _SampleDesc: + """Describe an inferred JAX output or workspace buffer.""" + + if not name: + raise ValueError("buffer descriptor name must not be empty") + return _make_desc( + jax.ShapeDtypeStruct(tuple(shape), dtype), + tensor_spec=tensor_spec, + name=name, + init_value=init_value, + ) + + +def as_gemm_tensor_desc(name: str, value: Any, tensor_spec: TensorSpec) -> _SampleDesc: + if isinstance(value, _SampleDesc): + if value.jax_mode != normalize_mode(value.ndim, tensor_spec.mode): + raise ValueError( + f"{name} descriptor mode does not match the requested layout" + ) + return value + return _make_desc(value, tensor_spec=tensor_spec, name=name) + + +class ApiBaseJax(JaxApiBase): + """Sample-signature-bound grouped GEMM callable using ``JaxApiBase``.""" + + def __init__(self) -> None: + self._is_supported = False + + def check_support(self) -> bool: + if not self._is_supported: + self._check_support() + self._is_supported = True + return True + + def _check_support(self) -> None: + raise NotImplementedError + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + self.check_support() + return self._call_impl(*args, **kwargs) + + def _call_impl(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError + + def make_tensor_desc( + self, value: Any, *, tensor_spec: TensorSpec | None = None, name: str = "" + ) -> _SampleDesc: + return _make_desc(value, tensor_spec=tensor_spec, name=name or "value") + + def make_optional_tensor_desc( + self, + value: Any | None, + *, + tensor_spec: TensorSpec | None = None, + name: str = "", + ) -> _SampleDesc | None: + return ( + None + if value is None + else self.make_tensor_desc(value, tensor_spec=tensor_spec, name=name) + ) + + @staticmethod + def as_optional_dtype(value: Any | None) -> Any | None: + return None if value is None else as_dtype(value) + + @staticmethod + def freeze_mapping(values: Mapping[str, Any]) -> Mapping[str, Any]: + # JAX captures the current configuration while tracing. Keep ordinary + # Python state so advanced users may reconfigure an un-jitted class + # instance and let the surrounding JIT establish its own cache key. + return dict(values) + + def check_tensor_signatures( + self, + expected: Mapping[str, _SampleDesc | None], + values: Mapping[str, Any], + ) -> None: + for name, expected_desc in expected.items(): + value = values[name] + if value is None or expected_desc is None: + if value is not None or expected_desc is not None: + raise ValueError( + f"{name} presence does not match the sample signature" + ) + continue + actual_shape = require_array(value, name=name) + if actual_shape != expected_desc.array_shape: + raise ValueError( + f"{name} tensor shape mismatch: expected {expected_desc.array_shape}, got {actual_shape}" + ) + if as_dtype(value) != as_dtype(expected_desc.dtype): + raise ValueError( + f"{name} tensor dtype mismatch: expected {expected_desc.dtype}, got {value.dtype}" + ) + + +class _GroupedKernelCaller(JaxApiBase): + def check_support(self) -> bool: + return True + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + raise TypeError("_GroupedKernelCaller is an internal lowering helper") + + +_CALLER = _GroupedKernelCaller() + + +def call_cutedsl( + fn: Any, + inputs: Sequence[Any], + *, + outputs: Sequence[JaxTensorDesc], + output_seeds: Sequence[Any | None] | None = None, + workspaces: Sequence[JaxTensorDesc] = (), + input_specs: Sequence[TensorSpec | None] | None = None, + static_args: Mapping[str, Any] | None = None, + allow_cuda_graph: bool = True, + use_static_tensors: bool = True, +) -> tuple[Any, ...]: + """Lower a grouped kernel with canonical stream/input/output/workspace order.""" + + static = dict(static_args or {}) + cluster_shape = static.get("cluster_shape_mn") + margin = int(static.pop("cluster_overlap_margin", 0)) + + output_specs = tuple(outputs) + workspace_specs = tuple(workspaces) + if output_seeds is None: + supplied_output_seeds = (None,) * len(output_specs) + else: + supplied_output_seeds = tuple(output_seeds) + if len(supplied_output_seeds) != len(output_specs): + raise ValueError( + f"Expected {len(output_specs)} output seeds, " + f"got {len(supplied_output_seeds)}" + ) + + if any(0 in desc.shape for desc in output_specs): + results = [] + for desc, seed in zip(output_specs, supplied_output_seeds): + spec = getattr(desc, "tensor_spec", None) + mode = None if spec is None else spec.mode + if seed is not None: + if desc.init_value is not None: + raise ValueError( + f"{desc.name} cannot have both an explicit seed and init_value" + ) + _CALLER._check_tensor_signature(seed, desc, mode=mode) + results.append(seed) + continue + metadata = _CALLER._materialize_tensor_desc( + desc, + mode=mode, + ) + if desc.init_value is None: + results.append(jnp.empty(metadata.shape, dtype=metadata.dtype)) + else: + results.append( + jnp.full( + metadata.shape, + desc.init_value, + dtype=metadata.dtype, + ) + ) + return tuple(results) + + if cluster_shape is not None: + static["max_active_clusters"] = _CALLER._get_max_active_clusters( + int(cluster_shape[0]) * int(cluster_shape[1]), + overlap_margin=margin, + ) + + def launch(stream: Any, *args: Any) -> None: + fn(stream, *args, **static) + + compute_capability = _CALLER._resolve_compute_capability( + None, + SUPPORTED_COMPUTE_CAPABILITIES, + "grouped GEMM", + ) + return _CALLER._call_kernel( + tuple(inputs), + launch=launch, + output_descs=output_specs, + output_seeds=supplied_output_seeds, + workspace_descs=workspace_specs, + input_spec=None if input_specs is None else tuple(input_specs), + output_spec=tuple(getattr(desc, "tensor_spec", None) for desc in output_specs), + workspace_spec=tuple( + getattr(desc, "tensor_spec", None) for desc in workspace_specs + ), + allow_cuda_graph=allow_cuda_graph, + compile_options=compile_options_for_target(compute_capability), + use_static_tensors=use_static_tensors, + ) + + +def block_scale_shape( + rows: int, k: int, batch: int, sf_vec_size: int +) -> tuple[int, ...]: + """Return the public row-major scale shape used by grouped JAX APIs.""" + + return to_public_axes( + _canonical_block_scale_shape(rows, k, batch, sf_vec_size), + BLOCK_SCALE_MODE, + ) + + +def require_grouped_gemm_inputs( + a_tensor: Any, + b_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + *, + max_experts: int = MAX_EXPERTS, + valid_ab_dtypes: Iterable[Any] | None = None, +) -> tuple[int, int, int, int, Any]: + """Validate canonical grouped GEMM matrix metadata.""" + + if valid_ab_dtypes is None: + valid_ab_dtypes = ( + jnp.float4_e2m1fn, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + ) + a_shape = require_array(a_tensor, name="a_tensor", rank=3, dtype=valid_ab_dtypes) + b_shape = require_array(b_tensor, name="b_tensor", rank=3, dtype=as_dtype(a_tensor)) + m, k, a_batch = a_shape + n, b_k, experts = b_shape + if a_batch != 1 or b_k != k: + raise ValueError( + f"Grouped GEMM expects A=(M,K,1) and B=(N,K,L), got {a_shape} and {b_shape}" + ) + if m < 0 or any(value <= 0 for value in (n, k, experts)): + raise ValueError( + "Grouped GEMM requires M >= 0 and positive N, K, and L; " + f"got M={m}, N={n}, K={k}, L={experts}" + ) + if experts > max_experts: + raise ValueError( + f"The number of experts must be at most {max_experts}, got {experts}" + ) + require_array( + padded_offsets, name="padded_offsets", shape=(experts,), dtype=jnp.int32 + ) + require_array( + alpha_tensor, name="alpha_tensor", shape=(experts,), dtype=jnp.float32 + ) + return m, n, k, experts, as_dtype(a_tensor) + + +def require_grouped_block_scales( + sfa_tensor: Any, + sfb_tensor: Any, + *, + m: int, + n: int, + k: int, + experts: int, + sf_vec_size: int, + valid_dtypes: Iterable[Any], +) -> Any: + require_array( + sfa_tensor, + name="sfa_tensor", + shape=block_scale_shape(m, k, 1, sf_vec_size), + dtype=valid_dtypes, + ) + sf_dtype = as_dtype(sfa_tensor) + require_array( + sfb_tensor, + name="sfb_tensor", + shape=block_scale_shape(n, k, experts, sf_vec_size), + dtype=sf_dtype, + ) + return sf_dtype + + +def require_grouped_fp8_scales( + sfa_tensor: Any, + sfb_tensor: Any, + *, + m: int, + n: int, + k: int, + experts: int, + sf_vec_size: int, +) -> Any: + return require_grouped_block_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + valid_dtypes=(jnp.float8_e8m0fnu,), + ) + + +def require_grouped_input_scales( + sfa_tensor: Any, + sfb_tensor: Any, + *, + m: int, + n: int, + k: int, + experts: int, + sf_vec_size: int, + ab_dtype: Any, +) -> Any: + """Validate MXFP8, MXFP4, or NVFP4 scale metadata. + + Native JAX FP4 arrays use logical element shapes. This helper deliberately + does not accept raw ``uint8`` payloads or emulate Torch's packed-storage + reinterpretation. + """ + + if is_fp4_dtype(ab_dtype): + if sf_vec_size not in SF_VEC_SIZES: + raise ValueError( + f"FP4 grouped GEMM requires sf_vec_size in {SF_VEC_SIZES}, got {sf_vec_size}" + ) + sf_dtype = require_grouped_block_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + valid_dtypes=(jnp.float8_e8m0fnu, jnp.float8_e4m3fn), + ) + if sf_dtype == jnp.dtype(jnp.float8_e4m3fn) and sf_vec_size != 16: + raise ValueError("float8_e4m3fn scales require sf_vec_size=16") + return sf_dtype + + if is_fp8_dtype(ab_dtype): + if sf_vec_size != FP8_SF_VEC_SIZE: + raise ValueError( + f"FP8 grouped GEMM requires sf_vec_size={FP8_SF_VEC_SIZE}, got {sf_vec_size}" + ) + return require_grouped_fp8_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ) + + raise ValueError(f"Unsupported grouped GEMM input dtype {as_dtype(ab_dtype)}") + + +def require_grouped_vector( + name: str, tensor: Any, *, length: int, dtype: Any = None +) -> Any: + require_array( + tensor, + name=name, + shape=(length,), + dtype=jnp.float32 if dtype is None else dtype, + ) + return as_dtype(tensor) + + +def require_grouped_probability(name: str, tensor: Any, *, m: int) -> None: + require_array(tensor, name=name, shape=(1, 1, m), dtype=jnp.float32) + + +def require_16_byte_extent(name: str, elements: int, dtype: Any) -> None: + from .._jax.datatypes import jax_to_cudnn_dtype + + require_contiguous_alignment( + name, elements, data_type_bits(jax_to_cudnn_dtype(dtype)) + ) + + +def require_contiguous_alignment(name: str, elements: int, element_bits: int) -> None: + if elements * element_bits % 128: + raise ValueError(f"{name} contiguous extent must span a multiple of 16 bytes") + + +def require_grouped_mma_tiler( + value: Sequence[int], + *, + allowed_m: tuple[int, ...] = (64, 128, 256), + allowed_n: tuple[int, ...] = (128, 256), +) -> tuple[int, int]: + value = tuple(value) + if len(value) != 2 or value[0] not in allowed_m or value[1] not in allowed_n: + raise ValueError(f"Unsupported grouped GEMM mma_tiler_mn {value}") + return value + + +def require_grouped_cluster_shape( + value: Sequence[int], *, mma_tiler_mn: tuple[int, int] +) -> tuple[int, int]: + value = tuple(value) + if len(value) != 2 or any(item <= 0 or item & (item - 1) for item in value): + raise ValueError( + f"cluster_shape_mn entries must be positive powers of two, got {value}" + ) + if any(item > 4 for item in value) or value[0] * value[1] > 16: + raise ValueError(f"cluster_shape_mn product must be at most 16, got {value}") + cta_group = 2 if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else 1 + if value[0] % cta_group: + raise ValueError( + f"cluster_shape_mn[0] must be divisible by {cta_group}, got {value[0]}" + ) + cluster_m = value[0] // cta_group * mma_tiler_mn[0] + if cluster_m not in (128, 256): + raise ValueError( + f"Grouped GEMM cluster M tile must be 128 or 256, got {cluster_m}" + ) + return value + + +def dense_workspace_bytes(use_dynamic_sched: bool) -> int: + return 4 if use_dynamic_sched else 0 + + +def grouped_wgrad_workspace_bytes(expert_cnt: int, input_order: str) -> int: + from .moe_utils import MoEWeightMode, WGradInputOrder, WgradSfTensormapConstructor + + return WgradSfTensormapConstructor.get_workspace_size( + WGradInputOrder(input_order), + MoEWeightMode.DENSE, + expert_cnt, + ) + + +def normalize_wgrad_input_order(input_order: Any) -> Any: + from .moe_utils import WGradInputOrder + + if isinstance(input_order, WGradInputOrder): + return input_order + try: + return WGradInputOrder(input_order) + except ValueError: + choices = ", ".join(order.value for order in WGradInputOrder) + raise ValueError( + f"input_order must be one of {{{choices}}}, got {input_order!r}" + ) from None + + +def ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +__all__ = [ + "ApiBaseJax", + "FIX_PAD_SIZE", + "FP8_SF_VEC_SIZE", + "HADAMARD_SIZE", + "MAX_EXPERTS", + "SF_VEC_SIZES", + "TWO_CTA_MMA_TILER_M", + "TupleDict", + "as_dtype", + "as_gemm_tensor_desc", + "block_scale_shape", + "block_scale_tensor_spec", + "call_cutedsl", + "ceil_div", + "dense_workspace_bytes", + "gemm_a_tensor_spec", + "gemm_b_tensor_spec", + "gemm_c_tensor_spec", + "grouped_bias_tensor_spec", + "grouped_wgrad_workspace_bytes", + "grouped_workspace_tensor_spec", + "is_fp4_dtype", + "is_fp8_dtype", + "is_low_precision_output_dtype", + "make_buffer_desc", + "normalize_wgrad_input_order", + "probability_tensor_spec", + "require_16_byte_extent", + "require_array", + "require_contiguous_alignment", + "require_dtype", + "require_grouped_block_scales", + "require_grouped_cluster_shape", + "require_grouped_fp8_scales", + "require_grouped_gemm_inputs", + "require_grouped_input_scales", + "require_grouped_mma_tiler", + "require_grouped_probability", + "require_grouped_vector", + "require_layout", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py index 748246694..4c9037d77 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py @@ -1,12 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmDgluSm100, - grouped_gemm_dglu_wrapper_sm100, -) +from ..._operation_api import make_operation_api -__all__ = [ - "GroupedGemmDgluSm100", - "grouped_gemm_dglu_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmDgluSm100", "grouped_gemm_dglu_wrapper_sm100")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dglu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_dglu/jax.py new file mode 100644 index 000000000..6d78601ff --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_dglu/jax.py @@ -0,0 +1,752 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for dense-weight grouped GEMM + dGLU on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + make_buffer_desc, + FIX_PAD_SIZE, + MAX_EXPERTS, + TWO_CTA_MMA_TILER_M, + TupleDict, + as_dtype, + as_gemm_tensor_desc, + block_scale_shape, + block_scale_tensor_spec, + call_cutedsl, + dense_workspace_bytes, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + grouped_workspace_tensor_spec, + is_fp4_dtype, + is_fp8_dtype, + probability_tensor_spec, + require_16_byte_extent, + require_dtype, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_input_scales, + require_grouped_mma_tiler, + require_grouped_probability, + require_grouped_vector, + require_layout, +) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + discrete_col_sfd: bool, + expert_cnt: int, + generate_dbias: bool, + generate_sfd: bool, + has_amax: bool, + use_dynamic_sched: bool, + act_func: str, + epilogue_op: str, + linear_offset: float, + geglu_alpha: float, + glu_clamp_max: float, + glu_clamp_min: float, + max_active_clusters: int, +): + import cutlass + import cutlass.cute as cute + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from ..moe_utils import MoEWeightMode + from .moe_blockscaled_grouped_gemm_dglu_dbias import ( + BlockScaledMoEGroupedGemmDgluDbiasKernel, + ) + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a = take() + b = take() + c = take() + sfa = take() + sfb = take() + padded_offsets = take() + alpha = take() + beta = take() + prob = take() + norm_const = take() if generate_sfd else None + + d_row = take() + d_col = take() + dprob = take() + dbias = take() if generate_dbias else None + amax = take() if has_amax else None + sfd_row = take() if generate_sfd else None + sfd_col = take() if generate_sfd else None + workspace = take() + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped GEMM argument count: consumed {arg_idx}, received {len(args)}" + ) + + if epilogue_op == "relu": + + def epilogue(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) + + elif epilogue_op == "srelu": + + def epilogue(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) ** 2 + + else: + + def epilogue(x): + return x + + kernel = BlockScaledMoEGroupedGemmDgluDbiasKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + discrete_col_sfd=discrete_col_sfd, + expert_cnt=expert_cnt, + weight_mode=MoEWeightMode.DENSE, + use_dynamic_sched=use_dynamic_sched, + act_func=act_func, + ) + kernel( + a, + b, + sfb, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int64(0), + OperandMajorMode.K, + workspace.iterator, + c, + d_row, + d_col, + sfa, + sfd_row, + sfd_col, + amax, + norm_const, + padded_offsets, + alpha, + beta, + prob, + dprob, + dbias, + max_active_clusters, + stream, + epilogue_op=epilogue, + linear_offset=linear_offset, + geglu_alpha=geglu_alpha, + glu_clamp_max=glu_clamp_max, + glu_clamp_min=glu_clamp_min, + ) + + +def _grouped_gemm_dglu_impl( + a_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Any, + prob_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + generate_dbias: bool = False, + norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "dswiglu", + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + epilogue_op: Optional[str] = None, + use_dynamic_sched: bool = False, + cluster_overlap_margin: int = 0, + *, + b_layout: str = "LNK", + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute a dense-weight block-scaled grouped GEMM with fused dGLU. + + ``dprob_tensor`` and optional ``dbias_tensor`` are fresh, zero-initialized + JAX results. Native FP4 inputs use logical array shapes; raw-byte + reinterpretation remains outside the JAX API. + """ + + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_spec = gemm_a_tensor_spec("LMK") + b_spec = gemm_b_tensor_spec(b_layout) + output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) + c_desc = as_gemm_tensor_desc("c_tensor", c_tensor, output_spec) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + ) + if n % 32: + raise ValueError(f"b_tensor N must be divisible by 32 for dGLU, got {n}") + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_input_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ab_dtype=ab_dtype, + ) + + output_n = 2 * n + expected_c_shape = (m, output_n, 1) + if c_desc.shape != expected_c_shape: + raise ValueError( + f"c_tensor must have canonical shape {expected_c_shape}, got {c_desc.shape}" + ) + require_dtype( + c_desc, + ( + jnp.float32, + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + ), + name="c_tensor.dtype", + ) + c_dtype = as_dtype(c_desc) + fp8_dtypes = { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + } + if vector_f32 and c_dtype in fp8_dtypes: + raise ValueError("vector_f32 does not support an FP8 c_tensor") + require_grouped_probability("prob_tensor", prob_tensor, m=m) + require_grouped_vector("beta_tensor", beta_tensor, length=experts) + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + if is_fp4_dtype(ab_dtype): + valid_d_dtypes = (jnp.float16, jnp.bfloat16, jnp.float32) + default_d_dtype = jnp.bfloat16 + else: + valid_d_dtypes = (jnp.float8_e4m3fn, jnp.float8_e5m2) + default_d_dtype = ab_dtype + d_dtype = require_dtype( + d_dtype, valid_d_dtypes, name="d_dtype", default=default_d_dtype + ) + generate_sfd = is_fp8_dtype(ab_dtype) + has_amax = d_dtype in {jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)} + if generate_sfd: + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required for FP8 inputs") + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + discrete_col_sfd = False + if is_fp4_dtype(ab_dtype) and b_layout != "LNK": + raise ValueError("Native FP4 B must use the K-major LNK layout") + if ( + is_fp4_dtype(ab_dtype) + and sf_vec_size == 16 + and d_dtype == jnp.dtype(jnp.float32) + and not generate_dbias + ): + raise NotImplementedError( + "FP4 with sf_vec_size=16 and float32 D requires generate_dbias=True" + ) + if act_func not in ("dswiglu", "dgeglu"): + raise ValueError(f"act_func must be 'dswiglu' or 'dgeglu', got {act_func!r}") + if linear_offset is None: + linear_offset = 1.0 if act_func == "dgeglu" else 0.0 + normalized_epilogue = ( + "identity" if epilogue_op in (None, "none", "identity") else epilogue_op + ) + if normalized_epilogue not in ("identity", "relu", "srelu"): + raise ValueError( + f"epilogue_op must be None, 'none', 'identity', 'relu', or 'srelu', got {epilogue_op!r}" + ) + + mma_tiler_mn = require_grouped_mma_tiler( + mma_tiler_mn, allowed_m=(128, 256), allowed_n=(256,) + ) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + + scale_spec = block_scale_tensor_spec() + require_16_byte_extent("a_tensor", k, ab_dtype) + require_16_byte_extent("b_tensor", n if b_layout == "LKN" else k, ab_dtype) + require_16_byte_extent("c_tensor", output_n, c_dtype) + require_16_byte_extent("d_tensor", output_n, d_dtype) + + if _validate_only: + return { + "acc_dtype": acc_dtype, + "d_dtype": d_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "linear_offset": linear_offset, + "epilogue_op": normalized_epilogue, + } + + outputs = [ + make_buffer_desc( + "d_row_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + ), + make_buffer_desc( + "d_col_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + ), + make_buffer_desc( + "dprob_tensor", + (1, 1, m), + jnp.float32, + tensor_spec=probability_tensor_spec(), + init_value=0.0, + ), + ] + if generate_dbias: + outputs.append( + make_buffer_desc( + "dbias_tensor", + (experts, output_n, 1), + jnp.bfloat16, + init_value=0.0, + ) + ) + if has_amax: + outputs.append( + make_buffer_desc( + "amax_tensor", + (experts, 2, 1), + jnp.float32, + init_value=-float("inf"), + ) + ) + if generate_sfd: + outputs.extend( + ( + make_buffer_desc( + "sfd_row_tensor", + block_scale_shape(m, output_n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + make_buffer_desc( + "sfd_col_tensor", + block_scale_shape(output_n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + ) + ) + + workspace_bytes = max(dense_workspace_bytes(bool(use_dynamic_sched)), 1) + inputs = [ + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + ] + input_specs = [ + a_spec, + b_spec, + output_spec, + scale_spec, + scale_spec, + None, + None, + None, + probability_tensor_spec(), + ] + if generate_sfd: + inputs.append(norm_const_tensor) + input_specs.append(None) + results = call_cutedsl( + _launch, + inputs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "discrete_col_sfd": bool(discrete_col_sfd), + "expert_cnt": experts, + "generate_dbias": bool(generate_dbias), + "generate_sfd": generate_sfd, + "has_amax": has_amax, + "use_dynamic_sched": bool(use_dynamic_sched), + "act_func": act_func, + "epilogue_op": normalized_epilogue, + "linear_offset": float(linear_offset), + "geglu_alpha": float(geglu_alpha), + "glu_clamp_max": float(glu_clamp_max), + "glu_clamp_min": float(glu_clamp_min), + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + workspaces=( + make_buffer_desc( + "workspace", + (workspace_bytes,), + jnp.uint8, + tensor_spec=grouped_workspace_tensor_spec(), + ), + ), + input_specs=input_specs, + ) + result_idx = 0 + d_row_tensor = results[result_idx] + result_idx += 1 + d_col_tensor = results[result_idx] + result_idx += 1 + dprob_tensor = results[result_idx] + result_idx += 1 + dbias_tensor = results[result_idx] if generate_dbias else None + result_idx += int(bool(generate_dbias)) + amax_tensor = results[result_idx] if has_amax else None + result_idx += int(has_amax) + sfd_row_tensor = results[result_idx] if generate_sfd else None + sfd_col_tensor = results[result_idx + 1] if generate_sfd else None + return TupleDict( + d_row_tensor=d_row_tensor, + d_col_tensor=d_col_tensor, + dprob_tensor=dprob_tensor, + dbias_tensor=dbias_tensor, + amax_tensor=amax_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + ) + + +class GroupedGemmDgluSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM + dGLU.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_c_tensor: Any, + sample_sfa_tensor: Any, + sample_padded_offsets: Any, + sample_alpha_tensor: Any, + sample_beta_tensor: Any, + sample_prob_tensor: Any, + sample_b_tensor: Any, + sample_sfb_tensor: Any, + generate_dbias: bool = False, + sample_norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "dswiglu", + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + epilogue_op: Optional[str] = None, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", + ) -> None: + super().__init__() + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_spec = gemm_a_tensor_spec("LMK") + b_spec = gemm_b_tensor_spec(b_layout) + output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") + scale_spec = block_scale_tensor_spec() + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + ), + "c_tensor": self.make_tensor_desc( + sample_c_tensor, tensor_spec=output_spec, name="sample_c_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + ), + "padded_offsets": self.make_tensor_desc( + sample_padded_offsets, name="sample_padded_offsets" + ), + "alpha_tensor": self.make_tensor_desc( + sample_alpha_tensor, name="sample_alpha_tensor" + ), + "beta_tensor": self.make_tensor_desc( + sample_beta_tensor, name="sample_beta_tensor" + ), + "prob_tensor": self.make_tensor_desc( + sample_prob_tensor, + tensor_spec=probability_tensor_spec(), + name="sample_prob_tensor", + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + ), + "norm_const_tensor": self.make_optional_tensor_desc( + sample_norm_const_tensor, name="sample_norm_const_tensor" + ), + } + self._config = { + "generate_dbias": generate_dbias, + "acc_dtype": self.as_optional_dtype(acc_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "discrete_col_sfd": discrete_col_sfd, + "act_func": act_func, + "linear_offset": linear_offset, + "geglu_alpha": geglu_alpha, + "glu_clamp_max": glu_clamp_max, + "glu_clamp_min": glu_clamp_min, + "epilogue_op": epilogue_op, + "use_dynamic_sched": use_dynamic_sched, + "b_layout": b_layout, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_dglu_impl( + self._sample_descs["a_tensor"], + self._sample_descs["c_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["beta_tensor"], + self._sample_descs["prob_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["sfb_tensor"], + norm_const_tensor=self._sample_descs["norm_const_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Any, + prob_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + norm_const_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + c_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + b_tensor, + sfb_tensor, + norm_const_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Any, + prob_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + norm_const_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "c_tensor": c_tensor, + "sfa_tensor": sfa_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "beta_tensor": beta_tensor, + "prob_tensor": prob_tensor, + "b_tensor": b_tensor, + "sfb_tensor": sfb_tensor, + "norm_const_tensor": norm_const_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_dglu_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "generate_dbias", + "acc_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "act_func", + "linear_offset", + "geglu_alpha", + "glu_clamp_max", + "glu_clamp_min", + "epilogue_op", + "use_dynamic_sched", + "b_layout", + ), +) +def grouped_gemm_dglu_wrapper_sm100( + a_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Any, + prob_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + generate_dbias: bool = False, + norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "dswiglu", + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + epilogue_op: Optional[str] = None, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", +) -> TupleDict: + """Compute an MXFP8 dense-weight grouped GEMM with fused dGLU.""" + + op = GroupedGemmDgluSm100( + a_tensor, + c_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + b_tensor, + sfb_tensor, + generate_dbias=generate_dbias, + sample_norm_const_tensor=norm_const_tensor, + acc_dtype=acc_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + act_func=act_func, + linear_offset=linear_offset, + geglu_alpha=geglu_alpha, + glu_clamp_max=glu_clamp_max, + glu_clamp_min=glu_clamp_min, + epilogue_op=epilogue_op, + use_dynamic_sched=use_dynamic_sched, + b_layout=b_layout, + ) + return op( + a_tensor, + c_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + b_tensor, + sfb_tensor, + norm_const_tensor, + ) + + +__all__ = [ + "GroupedGemmDgluSm100", + "grouped_gemm_dglu_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py index e15137794..ce6431284 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py @@ -1,12 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmDsreluSm100, - grouped_gemm_dsrelu_wrapper_sm100, -) +from ..._operation_api import make_operation_api -__all__ = [ - "GroupedGemmDsreluSm100", - "grouped_gemm_dsrelu_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmDsreluSm100", "grouped_gemm_dsrelu_wrapper_sm100")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/jax.py new file mode 100644 index 000000000..3a0a689c9 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/jax.py @@ -0,0 +1,670 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for dense-weight grouped GEMM + dSReLU on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + make_buffer_desc, + FIX_PAD_SIZE, + MAX_EXPERTS, + TWO_CTA_MMA_TILER_M, + TupleDict, + as_dtype, + as_gemm_tensor_desc, + block_scale_shape, + block_scale_tensor_spec, + call_cutedsl, + dense_workspace_bytes, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + grouped_workspace_tensor_spec, + is_fp4_dtype, + is_fp8_dtype, + probability_tensor_spec, + require_16_byte_extent, + require_dtype, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_input_scales, + require_grouped_mma_tiler, + require_grouped_probability, + require_grouped_vector, + require_layout, +) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + discrete_col_sfd: bool, + expert_cnt: int, + generate_dbias: bool, + generate_sfd: bool, + has_amax: bool, + use_dynamic_sched: bool, + use_dsrelu_reuse: bool, + max_active_clusters: int, +): + import cutlass + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from ..moe_utils import MoEWeightMode + from .moe_blockscaled_grouped_gemm_dsrelu_quant import ( + BlockScaledMoEGroupedGemmQuantBwdKernel, + EpilogueType, + ) + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a = take() + b = take() + c = take() + sfa = take() + sfb = take() + padded_offsets = take() + alpha = take() + prob = take() + norm_const = take() if generate_sfd else None + + d_row = take() + d_col = take() + d_srelu = take() + dprob = take() + dbias = take() if generate_dbias else None + amax = take() if has_amax else None + sfd_row = take() if generate_sfd else None + sfd_col = take() if generate_sfd else None + sfd_col_d_srelu = take() if generate_sfd else None + workspace = take() + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped GEMM argument count: consumed {arg_idx}, received {len(args)}" + ) + + kernel = BlockScaledMoEGroupedGemmQuantBwdKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + generate_sfd=generate_sfd, + discrete_col_sfd=discrete_col_sfd, + expert_cnt=expert_cnt, + weight_mode=MoEWeightMode.DENSE, + use_dynamic_sched=use_dynamic_sched, + epilogue_type=EpilogueType.DSRELU.value, + generate_dbias=generate_dbias, + generate_d_srelu=True, + use_dsrelu_reuse=use_dsrelu_reuse, + ) + kernel( + a, + b, + sfb, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int64(0), + OperandMajorMode.K, + workspace.iterator, + c, + d_row, + d_col, + sfa, + sfd_row, + sfd_col, + amax, + norm_const, + padded_offsets, + alpha, + prob, + dprob, + dbias, + d_srelu, + sfd_col_d_srelu, + max_active_clusters, + stream, + ) + + +def _grouped_gemm_dsrelu_impl( + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + generate_dbias: bool = False, + norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + use_dsrelu_reuse: bool = False, + cluster_overlap_margin: int = 0, + *, + b_layout: str = "LNK", + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute a dense-weight block-scaled grouped GEMM with fused dSReLU. + + The mutable Torch outputs are represented as functional JAX results. + ``dprob_tensor`` and optional ``dbias_tensor`` start at zero before the + kernel's atomic reductions. Native FP4 arrays use logical JAX shapes. + """ + + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_spec = gemm_a_tensor_spec("LMK") + b_spec = gemm_b_tensor_spec(b_layout) + output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) + c_desc = as_gemm_tensor_desc("c_tensor", c_tensor, output_spec) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + ) + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_input_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ab_dtype=ab_dtype, + ) + + expected_c_shape = (m, n, 1) + if c_desc.shape != expected_c_shape: + raise ValueError( + f"c_tensor must have canonical shape {expected_c_shape}, got {c_desc.shape}" + ) + require_dtype( + c_desc, + ( + jnp.float32, + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + ), + name="c_tensor.dtype", + ) + c_dtype = as_dtype(c_desc) + fp8_dtypes = { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + } + if vector_f32 and c_dtype in fp8_dtypes: + raise ValueError("vector_f32 does not support an FP8 c_tensor") + require_grouped_probability("prob_tensor", prob_tensor, m=m) + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + if is_fp4_dtype(ab_dtype): + valid_d_dtypes = (jnp.float16, jnp.bfloat16, jnp.float32) + default_d_dtype = jnp.bfloat16 + else: + valid_d_dtypes = (jnp.float8_e4m3fn, jnp.float8_e5m2) + default_d_dtype = ab_dtype + d_dtype = require_dtype( + d_dtype, valid_d_dtypes, name="d_dtype", default=default_d_dtype + ) + generate_sfd = is_fp8_dtype(ab_dtype) + has_amax = d_dtype in {jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)} + if generate_sfd: + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required for FP8 inputs") + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + discrete_col_sfd = False + if is_fp4_dtype(ab_dtype) and b_layout != "LNK": + raise ValueError("Native FP4 B must use the K-major LNK layout") + if ( + is_fp4_dtype(ab_dtype) + and sf_vec_size == 16 + and d_dtype == jnp.dtype(jnp.float32) + and not generate_dbias + ): + raise NotImplementedError( + "FP4 with sf_vec_size=16 and float32 D requires generate_dbias=True" + ) + mma_tiler_mn = require_grouped_mma_tiler( + mma_tiler_mn, allowed_m=(128, 256), allowed_n=(256,) + ) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + + scale_spec = block_scale_tensor_spec() + require_16_byte_extent("a_tensor", k, ab_dtype) + require_16_byte_extent("b_tensor", n if b_layout == "LKN" else k, ab_dtype) + require_16_byte_extent("c_tensor", n, c_dtype) + require_16_byte_extent("d_tensor", n, d_dtype) + + if _validate_only: + return { + "acc_dtype": acc_dtype, + "d_dtype": d_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + } + + outputs = [ + make_buffer_desc("d_row_tensor", (1, m, n), d_dtype, tensor_spec=output_spec), + make_buffer_desc("d_col_tensor", (1, m, n), d_dtype, tensor_spec=output_spec), + make_buffer_desc("d_srelu_tensor", (1, m, n), d_dtype, tensor_spec=output_spec), + make_buffer_desc( + "dprob_tensor", + (1, 1, m), + jnp.float32, + tensor_spec=probability_tensor_spec(), + init_value=0.0, + ), + ] + if generate_dbias: + outputs.append( + make_buffer_desc( + "dbias_tensor", + (experts, n, 1), + jnp.bfloat16, + init_value=0.0, + ) + ) + if has_amax: + outputs.append( + make_buffer_desc( + "amax_tensor", (experts, 1), jnp.float32, init_value=-float("inf") + ) + ) + if generate_sfd: + outputs.extend( + ( + make_buffer_desc( + "sfd_row_tensor", + block_scale_shape(m, n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + make_buffer_desc( + "sfd_col_tensor", + block_scale_shape(n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + make_buffer_desc( + "sfd_col_d_srelu_tensor", + block_scale_shape(n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + ) + ) + + workspace_bytes = max(dense_workspace_bytes(bool(use_dynamic_sched)), 1) + inputs = [ + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + ] + input_specs = [ + a_spec, + b_spec, + output_spec, + scale_spec, + scale_spec, + None, + None, + probability_tensor_spec(), + ] + if generate_sfd: + inputs.append(norm_const_tensor) + input_specs.append(None) + results = call_cutedsl( + _launch, + inputs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "discrete_col_sfd": bool(discrete_col_sfd), + "expert_cnt": experts, + "generate_dbias": bool(generate_dbias), + "generate_sfd": generate_sfd, + "has_amax": has_amax, + "use_dynamic_sched": bool(use_dynamic_sched), + "use_dsrelu_reuse": bool(use_dsrelu_reuse), + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + workspaces=( + make_buffer_desc( + "workspace", + (workspace_bytes,), + jnp.uint8, + tensor_spec=grouped_workspace_tensor_spec(), + ), + ), + input_specs=input_specs, + ) + result_idx = 0 + d_row_tensor = results[result_idx] + result_idx += 1 + d_col_tensor = results[result_idx] + result_idx += 1 + d_srelu_tensor = results[result_idx] + result_idx += 1 + dprob_tensor = results[result_idx] + result_idx += 1 + dbias_tensor = results[result_idx] if generate_dbias else None + result_idx += int(bool(generate_dbias)) + amax_tensor = results[result_idx] if has_amax else None + result_idx += int(has_amax) + sfd_row_tensor = results[result_idx] if generate_sfd else None + sfd_col_tensor = results[result_idx + 1] if generate_sfd else None + sfd_col_d_srelu_tensor = results[result_idx + 2] if generate_sfd else None + return TupleDict( + d_row_tensor=d_row_tensor, + d_col_tensor=d_col_tensor, + d_srelu_tensor=d_srelu_tensor, + dprob_tensor=dprob_tensor, + dbias_tensor=dbias_tensor, + amax_tensor=amax_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + sfd_col_d_srelu_tensor=sfd_col_d_srelu_tensor, + ) + + +class GroupedGemmDsreluSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM + dSReLU.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_b_tensor: Any, + sample_c_tensor: Any, + sample_sfa_tensor: Any, + sample_sfb_tensor: Any, + sample_padded_offsets: Any, + sample_alpha_tensor: Any, + sample_prob_tensor: Any, + generate_dbias: bool = False, + sample_norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + use_dsrelu_reuse: bool = False, + *, + b_layout: str = "LNK", + ) -> None: + super().__init__() + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_spec = gemm_a_tensor_spec("LMK") + b_spec = gemm_b_tensor_spec(b_layout) + output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") + scale_spec = block_scale_tensor_spec() + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + ), + "c_tensor": self.make_tensor_desc( + sample_c_tensor, tensor_spec=output_spec, name="sample_c_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + ), + "padded_offsets": self.make_tensor_desc( + sample_padded_offsets, name="sample_padded_offsets" + ), + "alpha_tensor": self.make_tensor_desc( + sample_alpha_tensor, name="sample_alpha_tensor" + ), + "prob_tensor": self.make_tensor_desc( + sample_prob_tensor, + tensor_spec=probability_tensor_spec(), + name="sample_prob_tensor", + ), + "norm_const_tensor": self.make_optional_tensor_desc( + sample_norm_const_tensor, name="sample_norm_const_tensor" + ), + } + self._config = { + "generate_dbias": generate_dbias, + "acc_dtype": self.as_optional_dtype(acc_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "discrete_col_sfd": discrete_col_sfd, + "use_dynamic_sched": use_dynamic_sched, + "use_dsrelu_reuse": use_dsrelu_reuse, + "b_layout": b_layout, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_dsrelu_impl( + self._sample_descs["a_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["c_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["sfb_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["prob_tensor"], + norm_const_tensor=self._sample_descs["norm_const_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + norm_const_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + norm_const_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + norm_const_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "b_tensor": b_tensor, + "c_tensor": c_tensor, + "sfa_tensor": sfa_tensor, + "sfb_tensor": sfb_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "prob_tensor": prob_tensor, + "norm_const_tensor": norm_const_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_dsrelu_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "generate_dbias", + "acc_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "use_dynamic_sched", + "use_dsrelu_reuse", + "b_layout", + ), +) +def grouped_gemm_dsrelu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + generate_dbias: bool = False, + norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + use_dsrelu_reuse: bool = False, + *, + b_layout: str = "LNK", +) -> TupleDict: + """Compute an MXFP8 dense-weight grouped GEMM with fused dSReLU.""" + + op = GroupedGemmDsreluSm100( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + generate_dbias=generate_dbias, + sample_norm_const_tensor=norm_const_tensor, + acc_dtype=acc_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + use_dynamic_sched=use_dynamic_sched, + use_dsrelu_reuse=use_dsrelu_reuse, + b_layout=b_layout, + ) + return op( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + norm_const_tensor, + ) + + +__all__ = [ + "GroupedGemmDsreluSm100", + "grouped_gemm_dsrelu_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py index f0261951a..91f534d34 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py @@ -1,12 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmDswigluSm100, - grouped_gemm_dswiglu_wrapper_sm100, -) +from ..._operation_api import make_operation_api -__all__ = [ - "GroupedGemmDswigluSm100", - "grouped_gemm_dswiglu_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmDswigluSm100", "grouped_gemm_dswiglu_wrapper_sm100")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/jax.py new file mode 100644 index 000000000..4e07bdf13 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/jax.py @@ -0,0 +1,625 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for contiguous grouped GEMM + dSwiGLU on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + make_buffer_desc, + FIX_PAD_SIZE, + MAX_EXPERTS, + TWO_CTA_MMA_TILER_M, + TupleDict, + as_dtype, + as_gemm_tensor_desc, + block_scale_shape, + block_scale_tensor_spec, + call_cutedsl, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + is_fp4_dtype, + is_fp8_dtype, + probability_tensor_spec, + require_16_byte_extent, + require_dtype, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_input_scales, + require_grouped_mma_tiler, + require_grouped_probability, + require_grouped_vector, + require_layout, +) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + discrete_col_sfd: bool, + expert_cnt: int, + epilogue_op: str, + generate_sfd: bool, + has_amax: bool, + max_active_clusters: int, +): + import cutlass.cute as cute + from cutlass.jax import jax_to_cutlass_dtype + + from .grouped_gemm_dswiglu_quant import BlockScaledContiguousGroupedGemmKernel + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a, b, c, sfa, sfb, padded_offsets, alpha, beta, prob = (take() for _ in range(9)) + norm_const = take() if generate_sfd else None + d_row, d_col, dprob = (take() for _ in range(3)) + amax = take() if has_amax else None + sfd_row = take() if generate_sfd else None + sfd_col = take() if generate_sfd else None + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped dSwiGLU argument count: consumed {arg_idx}, received {len(args)}" + ) + + if epilogue_op == "relu": + + def epilogue(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) + + elif epilogue_op == "srelu": + + def epilogue(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) ** 2 + + else: + + def epilogue(x): + return x + + kernel = BlockScaledContiguousGroupedGemmKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + discrete_col_sfd=discrete_col_sfd, + expert_cnt=expert_cnt, + use_mono_increase_expert_idx=True, + ) + kernel( + a, + b, + c, + d_row, + d_col, + sfa, + sfb, + sfd_row, + sfd_col, + amax, + norm_const, + padded_offsets, + alpha, + beta, + prob, + dprob, + max_active_clusters, + stream, + epilogue_op=epilogue, + ) + + +def _grouped_gemm_dswiglu_impl( + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Optional[Any], + prob_tensor: Any, + norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + epilogue_op: Optional[str] = None, + b_layout: str = "LNK", + cluster_overlap_margin: int = 0, + *, + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute the contiguous block-scaled grouped dSwiGLU fusion. + + ``dprob_tensor`` is modeled as a fresh zero-initialized JAX result instead + of a caller-owned mutable buffer. If ``beta_tensor`` is ``None``, a vector + of ones is supplied, matching the convenience behavior of the Torch API. + Native FP4 inputs use logical JAX shapes; raw-byte reinterpretation is not + part of this adapter. Configuration values are static under JAX tracing. + """ + + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + a_spec = gemm_a_tensor_spec("LMK") + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + b_spec = gemm_b_tensor_spec(b_layout) + output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) + c_desc = as_gemm_tensor_desc("c_tensor", c_tensor, output_spec) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + ) + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_input_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ab_dtype=ab_dtype, + ) + + output_n = 2 * n + expected_c_shape = (m, output_n, 1) + if c_desc.shape != expected_c_shape: + raise ValueError( + f"c_tensor must have canonical shape {expected_c_shape}, got {c_desc.shape}" + ) + require_dtype( + c_desc, + ( + jnp.float32, + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + ), + name="c_tensor.dtype", + ) + c_dtype = as_dtype(c_desc) + if vector_f32 and c_dtype in { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + }: + raise ValueError("vector_f32 does not support an FP8 c_tensor") + require_grouped_probability("prob_tensor", prob_tensor, m=m) + if beta_tensor is None: + if not _validate_only: + beta_tensor = jnp.ones((experts,), dtype=jnp.float32) + else: + require_grouped_vector("beta_tensor", beta_tensor, length=experts) + + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + if is_fp4_dtype(ab_dtype): + valid_d_dtypes = (jnp.float16, jnp.bfloat16, jnp.float32) + default_d_dtype = jnp.bfloat16 + else: + valid_d_dtypes = (jnp.float8_e4m3fn, jnp.float8_e5m2) + default_d_dtype = ab_dtype + d_dtype = require_dtype( + d_dtype, valid_d_dtypes, name="d_dtype", default=default_d_dtype + ) + generate_sfd = is_fp8_dtype(ab_dtype) + has_amax = d_dtype in {jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)} + if generate_sfd: + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required for FP8 inputs") + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + discrete_col_sfd = False + if is_fp4_dtype(ab_dtype) and b_layout != "LNK": + raise ValueError("Native FP4 B must use the K-major LNK layout") + normalized_epilogue = ( + "identity" if epilogue_op in (None, "none", "identity") else epilogue_op + ) + if normalized_epilogue not in ("identity", "relu", "srelu"): + raise ValueError( + f"epilogue_op must be None, 'none', 'identity', 'relu', or 'srelu', got {epilogue_op!r}" + ) + + mma_tiler_mn = require_grouped_mma_tiler(mma_tiler_mn) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + + scale_spec = block_scale_tensor_spec() + require_16_byte_extent("a_tensor", k, ab_dtype) + require_16_byte_extent("b_tensor", k, ab_dtype) + require_16_byte_extent("c_tensor", output_n, c_dtype) + require_16_byte_extent("d_tensor", output_n, d_dtype) + + if _validate_only: + return { + "acc_dtype": acc_dtype, + "d_dtype": d_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "epilogue_op": normalized_epilogue, + } + + inputs = [ + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + ] + input_specs = [ + a_spec, + b_spec, + output_spec, + scale_spec, + scale_spec, + None, + None, + None, + probability_tensor_spec(), + ] + if generate_sfd: + inputs.append(norm_const_tensor) + input_specs.append(None) + + outputs = [ + make_buffer_desc( + "d_row_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + ), + make_buffer_desc( + "d_col_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + ), + make_buffer_desc( + "dprob_tensor", + (1, 1, m), + jnp.float32, + tensor_spec=probability_tensor_spec(), + init_value=0.0, + ), + ] + if has_amax: + outputs.append( + make_buffer_desc( + "amax_tensor", + (experts, 2, 1), + jnp.float32, + init_value=-float("inf"), + ) + ) + if generate_sfd: + outputs.extend( + ( + make_buffer_desc( + "sfd_row_tensor", + block_scale_shape(m, output_n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + make_buffer_desc( + "sfd_col_tensor", + block_scale_shape(output_n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + ) + ) + + results = call_cutedsl( + _launch, + inputs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "discrete_col_sfd": bool(discrete_col_sfd), + "expert_cnt": experts, + "epilogue_op": normalized_epilogue, + "generate_sfd": generate_sfd, + "has_amax": has_amax, + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + input_specs=input_specs, + ) + result_idx = 3 + d_row_tensor, d_col_tensor, dprob_tensor = results[:3] + amax_tensor = results[result_idx] if has_amax else None + result_idx += int(has_amax) + sfd_row_tensor = results[result_idx] if generate_sfd else None + sfd_col_tensor = results[result_idx + 1] if generate_sfd else None + return TupleDict( + d_row_tensor=d_row_tensor, + d_col_tensor=d_col_tensor, + dprob_tensor=dprob_tensor, + amax_tensor=amax_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + ) + + +class GroupedGemmDswigluSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM + dSwiGLU.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_b_tensor: Any, + sample_c_tensor: Any, + sample_sfa_tensor: Any, + sample_sfb_tensor: Any, + sample_padded_offsets: Any, + sample_alpha_tensor: Any, + sample_beta_tensor: Optional[Any], + sample_prob_tensor: Any, + sample_norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + epilogue_op: Optional[str] = None, + b_layout: str = "LNK", + ) -> None: + super().__init__() + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + a_spec = gemm_a_tensor_spec("LMK") + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + b_spec = gemm_b_tensor_spec(b_layout) + output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") + scale_spec = block_scale_tensor_spec() + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + ), + "c_tensor": self.make_tensor_desc( + sample_c_tensor, tensor_spec=output_spec, name="sample_c_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + ), + "padded_offsets": self.make_tensor_desc( + sample_padded_offsets, name="sample_padded_offsets" + ), + "alpha_tensor": self.make_tensor_desc( + sample_alpha_tensor, name="sample_alpha_tensor" + ), + "beta_tensor": self.make_optional_tensor_desc( + sample_beta_tensor, name="sample_beta_tensor" + ), + "prob_tensor": self.make_tensor_desc( + sample_prob_tensor, + tensor_spec=probability_tensor_spec(), + name="sample_prob_tensor", + ), + "norm_const_tensor": self.make_optional_tensor_desc( + sample_norm_const_tensor, name="sample_norm_const_tensor" + ), + } + self._config = { + "acc_dtype": self.as_optional_dtype(acc_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "discrete_col_sfd": discrete_col_sfd, + "epilogue_op": epilogue_op, + "b_layout": b_layout, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_dswiglu_impl( + self._sample_descs["a_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["c_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["sfb_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["beta_tensor"], + self._sample_descs["prob_tensor"], + self._sample_descs["norm_const_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Optional[Any], + prob_tensor: Any, + norm_const_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + norm_const_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Optional[Any], + prob_tensor: Any, + norm_const_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "b_tensor": b_tensor, + "c_tensor": c_tensor, + "sfa_tensor": sfa_tensor, + "sfb_tensor": sfb_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "beta_tensor": beta_tensor, + "prob_tensor": prob_tensor, + "norm_const_tensor": norm_const_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_dswiglu_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "acc_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "epilogue_op", + "b_layout", + ), +) +def grouped_gemm_dswiglu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Optional[Any], + prob_tensor: Any, + norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + epilogue_op: Optional[str] = None, + b_layout: str = "LNK", +) -> TupleDict: + """Compute the contiguous block-scaled grouped dSwiGLU fusion.""" + + op = GroupedGemmDswigluSm100( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + norm_const_tensor, + acc_dtype=acc_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + epilogue_op=epilogue_op, + b_layout=b_layout, + ) + return op( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + norm_const_tensor, + ) + + +__all__ = [ + "GroupedGemmDswigluSm100", + "grouped_gemm_dswiglu_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_glu/__init__.py index 177db7f07..8e38a750c 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu/__init__.py @@ -1,12 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmGluSm100, - grouped_gemm_glu_wrapper_sm100, -) +from ..._operation_api import make_operation_api -__all__ = [ - "GroupedGemmGluSm100", - "grouped_gemm_glu_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmGluSm100", "grouped_gemm_glu_wrapper_sm100")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_glu/jax.py new file mode 100644 index 000000000..56800496c --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu/jax.py @@ -0,0 +1,687 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for dense-weight grouped GEMM + GLU on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + make_buffer_desc, + FIX_PAD_SIZE, + MAX_EXPERTS, + TWO_CTA_MMA_TILER_M, + TupleDict, + as_gemm_tensor_desc, + block_scale_shape, + block_scale_tensor_spec, + call_cutedsl, + dense_workspace_bytes, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + grouped_bias_tensor_spec, + grouped_workspace_tensor_spec, + is_fp4_dtype, + is_fp8_dtype, + probability_tensor_spec, + require_16_byte_extent, + require_array, + require_dtype, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_input_scales, + require_grouped_mma_tiler, + require_grouped_probability, + require_grouped_vector, + require_layout, +) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + discrete_col_sfd: bool, + expert_cnt: int, + has_prob: bool, + has_bias: bool, + generate_sfd: bool, + has_amax: bool, + use_dynamic_sched: bool, + act_func: str, + linear_offset: float, + geglu_alpha: float, + glu_clamp_max: float, + glu_clamp_min: float, + max_active_clusters: int, +): + import cutlass + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from ..moe_utils import MoEWeightMode + from .moe_blockscaled_grouped_gemm_glu_bias import ( + BlockScaledMoEGroupedGemmGluBiasKernel, + ) + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a = take() + b = take() + sfa = take() + sfb = take() + padded_offsets = take() + alpha = take() + prob = take() if has_prob else None + bias = take() if has_bias else None + norm_const = take() if generate_sfd else None + + c = take() + d = take() + d_col = take() + amax = take() if has_amax else None + sfd_row = take() if generate_sfd else None + sfd_col = take() if generate_sfd else None + workspace = take() + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped GEMM argument count: consumed {arg_idx}, received {len(args)}" + ) + + kernel = BlockScaledMoEGroupedGemmGluBiasKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + generate_sfd=generate_sfd, + discrete_col_sfd=discrete_col_sfd, + expert_cnt=expert_cnt, + weight_mode=MoEWeightMode.DENSE, + use_dynamic_sched=use_dynamic_sched, + act_func=act_func, + enable_bias=has_bias, + ) + kernel( + a, + b, + sfb, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int64(0), + OperandMajorMode.K, + workspace.iterator, + c, + d, + d_col, + sfa, + sfd_row, + sfd_col, + amax, + norm_const, + padded_offsets, + alpha, + prob, + bias, + max_active_clusters, + stream, + linear_offset=cutlass.Float32(linear_offset), + geglu_alpha=cutlass.Float32(geglu_alpha), + glu_clamp_max=cutlass.Float32(glu_clamp_max), + glu_clamp_min=cutlass.Float32(glu_clamp_min), + ) + + +def _grouped_gemm_glu_impl( + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "swiglu", + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + use_dynamic_sched: bool = False, + cluster_overlap_margin: int = 0, + *, + b_layout: str = "LNK", + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute a dense-weight block-scaled grouped GEMM with fused GLU. + + Dense expert weights use public shape ``(L, N, K)``. ``swiglu`` and ``geglu`` + produce ``N / 2`` output columns. Native FP4 arrays cover MXFP4/NVFP4; + raw-byte reinterpretation is intentionally not accepted. Low-precision + outputs from FP8 inputs return scale factors. + """ + + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK",)) + a_spec = gemm_a_tensor_spec("LMK") + b_spec = gemm_b_tensor_spec(b_layout) + output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + ) + if n % 64: + raise ValueError(f"b_tensor N must be divisible by 64 for GLU, got {n}") + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_input_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ab_dtype=ab_dtype, + ) + if prob_tensor is not None: + require_grouped_probability("prob_tensor", prob_tensor, m=m) + if bias_tensor is not None: + require_array( + bias_tensor, + name="bias_tensor", + shape=(experts, n), + dtype=(jnp.float16, jnp.bfloat16, jnp.float32), + ) + + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + c_dtype = require_dtype( + c_dtype, + ( + jnp.float32, + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + jnp.float4_e2m1fn, + ), + name="c_dtype", + default=jnp.bfloat16, + ) + if is_fp4_dtype(ab_dtype): + valid_d_dtypes = (jnp.float16, jnp.bfloat16, jnp.float32) + else: + valid_d_dtypes = ( + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + jnp.float4_e2m1fn, + ) + d_dtype = require_dtype( + d_dtype, valid_d_dtypes, name="d_dtype", default=jnp.bfloat16 + ) + if ( + is_fp4_dtype(ab_dtype) + and sf_vec_size == 16 + and d_dtype == jnp.dtype(jnp.float32) + and bias_tensor is None + ): + raise NotImplementedError( + "FP4 with sf_vec_size=16 and float32 D requires a bias tensor" + ) + fp8_dtypes = { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + } + # The Torch GLU wrapper requests row/column scales for every FP8-input + # configuration, including a high-precision D output. + generate_sfd = is_fp8_dtype(ab_dtype) + has_amax = d_dtype in {jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)} + if generate_sfd: + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required for FP8 inputs") + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + discrete_col_sfd = False + if vector_f32 and c_dtype in fp8_dtypes: + raise ValueError("vector_f32 does not support an FP8 c_dtype") + if act_func not in ("swiglu", "geglu"): + raise ValueError(f"act_func must be 'swiglu' or 'geglu', got {act_func!r}") + if linear_offset is None: + linear_offset = 1.0 if act_func == "geglu" else 0.0 + + mma_tiler_mn = require_grouped_mma_tiler( + mma_tiler_mn, allowed_m=(128, 256), allowed_n=(256,) + ) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + + output_n = n // 2 + scale_spec = block_scale_tensor_spec() + require_16_byte_extent("a_tensor", k, ab_dtype) + require_16_byte_extent("b_tensor", k, ab_dtype) + require_16_byte_extent("c_tensor", n, c_dtype) + require_16_byte_extent("d_tensor", output_n, d_dtype) + + if _validate_only: + return { + "acc_dtype": acc_dtype, + "c_dtype": c_dtype, + "d_dtype": d_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "linear_offset": linear_offset, + } + + inputs = [ + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + ] + input_specs = [a_spec, b_spec, scale_spec, scale_spec, None, None] + if prob_tensor is not None: + inputs.append(prob_tensor) + input_specs.append(probability_tensor_spec()) + if bias_tensor is not None: + inputs.append(bias_tensor) + input_specs.append(grouped_bias_tensor_spec()) + if generate_sfd: + inputs.append(norm_const_tensor) + input_specs.append(None) + + outputs = [ + make_buffer_desc("c_tensor", (1, m, n), c_dtype, tensor_spec=output_spec), + make_buffer_desc( + "d_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + ), + ] + outputs.append( + make_buffer_desc( + "d_col_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + ) + ) + if has_amax: + outputs.append( + make_buffer_desc( + "amax_tensor", (experts, 1), jnp.float32, init_value=-float("inf") + ) + ) + if generate_sfd: + outputs.extend( + ( + make_buffer_desc( + "sfd_row_tensor", + block_scale_shape(m, output_n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + make_buffer_desc( + "sfd_col_tensor", + block_scale_shape(output_n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + ) + ) + + workspace_bytes = max(dense_workspace_bytes(bool(use_dynamic_sched)), 1) + workspaces = [ + make_buffer_desc( + "workspace", + (workspace_bytes,), + jnp.uint8, + tensor_spec=grouped_workspace_tensor_spec(), + ) + ] + results = call_cutedsl( + _launch, + inputs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "discrete_col_sfd": bool(discrete_col_sfd), + "expert_cnt": experts, + "has_prob": prob_tensor is not None, + "has_bias": bias_tensor is not None, + "generate_sfd": generate_sfd, + "has_amax": has_amax, + "use_dynamic_sched": bool(use_dynamic_sched), + "act_func": act_func, + "linear_offset": float(linear_offset), + "geglu_alpha": float(geglu_alpha), + "glu_clamp_max": float(glu_clamp_max), + "glu_clamp_min": float(glu_clamp_min), + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + workspaces=workspaces, + input_specs=input_specs, + ) + c_tensor, d_tensor, d_col_tensor = results[:3] + result_idx = 3 + amax_tensor = results[result_idx] if has_amax else None + result_idx += int(has_amax) + sfd_row_tensor = results[result_idx] if generate_sfd else None + sfd_col_tensor = results[result_idx + 1] if generate_sfd else None + return TupleDict( + c_tensor=c_tensor, + d_tensor=d_tensor, + d_col_tensor=d_col_tensor, + amax_tensor=amax_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + ) + + +class GroupedGemmGluSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM + GLU.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_sfa_tensor: Any, + sample_padded_offsets: Any, + sample_alpha_tensor: Any, + sample_b_tensor: Any, + sample_sfb_tensor: Any, + sample_bias_tensor: Optional[Any] = None, + sample_norm_const_tensor: Optional[Any] = None, + sample_prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "swiglu", + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", + ) -> None: + super().__init__() + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK",)) + a_spec = gemm_a_tensor_spec("LMK") + b_spec = gemm_b_tensor_spec(b_layout) + scale_spec = block_scale_tensor_spec() + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + ), + "padded_offsets": self.make_tensor_desc( + sample_padded_offsets, name="sample_padded_offsets" + ), + "alpha_tensor": self.make_tensor_desc( + sample_alpha_tensor, name="sample_alpha_tensor" + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + ), + "bias_tensor": self.make_optional_tensor_desc( + sample_bias_tensor, + tensor_spec=grouped_bias_tensor_spec(), + name="sample_bias_tensor", + ), + "norm_const_tensor": self.make_optional_tensor_desc( + sample_norm_const_tensor, name="sample_norm_const_tensor" + ), + "prob_tensor": self.make_optional_tensor_desc( + sample_prob_tensor, + tensor_spec=probability_tensor_spec(), + name="sample_prob_tensor", + ), + } + self._config = { + "acc_dtype": self.as_optional_dtype(acc_dtype), + "c_dtype": self.as_optional_dtype(c_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "discrete_col_sfd": discrete_col_sfd, + "act_func": act_func, + "linear_offset": linear_offset, + "geglu_alpha": geglu_alpha, + "glu_clamp_max": glu_clamp_max, + "glu_clamp_min": glu_clamp_min, + "use_dynamic_sched": use_dynamic_sched, + "b_layout": b_layout, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_glu_impl( + self._sample_descs["a_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["sfb_tensor"], + self._sample_descs["bias_tensor"], + self._sample_descs["norm_const_tensor"], + self._sample_descs["prob_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + b_tensor, + sfb_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "sfa_tensor": sfa_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "b_tensor": b_tensor, + "sfb_tensor": sfb_tensor, + "bias_tensor": bias_tensor, + "norm_const_tensor": norm_const_tensor, + "prob_tensor": prob_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_glu_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "acc_dtype", + "c_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "act_func", + "linear_offset", + "geglu_alpha", + "glu_clamp_max", + "glu_clamp_min", + "use_dynamic_sched", + "b_layout", + ), +) +def grouped_gemm_glu_wrapper_sm100( + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "swiglu", + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", +) -> TupleDict: + """Compute an MXFP8 dense-weight grouped GEMM with fused GLU.""" + + op = GroupedGemmGluSm100( + a_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + b_tensor, + sfb_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + acc_dtype=acc_dtype, + c_dtype=c_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + act_func=act_func, + linear_offset=linear_offset, + geglu_alpha=geglu_alpha, + glu_clamp_max=glu_clamp_max, + glu_clamp_min=glu_clamp_min, + use_dynamic_sched=use_dynamic_sched, + b_layout=b_layout, + ) + return op( + a_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + b_tensor, + sfb_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + ) + + +__all__ = [ + "GroupedGemmGluSm100", + "grouped_gemm_glu_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/__init__.py index 2ef06aa02..1d81b6710 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/__init__.py @@ -1,12 +1,15 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmGluHadamardSm100, - grouped_gemm_glu_hadamard_wrapper_sm100, -) +from ..._operation_api import make_operation_api -__all__ = [ - "GroupedGemmGluHadamardSm100", - "grouped_gemm_glu_hadamard_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "api": ( + "GroupedGemmGluHadamardSm100", + "grouped_gemm_glu_hadamard_wrapper_sm100", + ) + }, + submodules=("api", "jax"), +) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/hadamard_utils.py b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/hadamard_utils.py index ec55574bc..08b7b056d 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/hadamard_utils.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/hadamard_utils.py @@ -10,7 +10,6 @@ from cutlass.cute.nvgpu import tcgen05 from cutlass.cute.nvgpu import OperandMajorMode from cutlass.cute.nvgpu.tcgen05 import OperandSource -import torch HADAMARD_SIZE = 16 TMEM_ROW_STRIDE = 1 << 16 @@ -235,6 +234,8 @@ def hadamard_smem_transpose_fwht_amax(sD, d_buffer, feature_offset: cutlass.Cons def hadamard_matrix(n, dtype=None, device=None): + import torch + if dtype is None: dtype = torch.float32 if n < 1: diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/jax.py new file mode 100644 index 000000000..386131692 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/jax.py @@ -0,0 +1,648 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for dense grouped GEMM + GLU + Hadamard on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + make_buffer_desc, + FIX_PAD_SIZE, + HADAMARD_SIZE, + MAX_EXPERTS, + SF_VEC_SIZES, + TupleDict, + as_gemm_tensor_desc, + block_scale_tensor_spec, + call_cutedsl, + dense_workspace_bytes, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + grouped_bias_tensor_spec, + grouped_workspace_tensor_spec, + probability_tensor_spec, + require_array, + require_contiguous_alignment, + require_dtype, + require_grouped_block_scales, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_mma_tiler, + require_grouped_probability, + require_layout, +) + + +def hadamard_values(size: int) -> tuple[tuple[int, ...], ...]: + """Return an unnormalized Sylvester Hadamard matrix as Python values.""" + + if size < 1 or size & (size - 1): + raise ValueError(f"Hadamard size must be a positive power of two, got {size}") + matrix = ((1,),) + while len(matrix) < size: + matrix = tuple(row + row for row in matrix) + tuple( + row + tuple(-value for value in row) for row in matrix + ) + return matrix + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + expert_cnt: int, + act_func: str, + has_bias: bool, + has_hadamard: bool, + use_dynamic_sched: bool, + use_tmem_post_rht_amax: bool, + max_active_clusters: int, +): + import cutlass + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from ..moe_utils import MoEWeightMode + from .moe_blockscaled_grouped_gemm_glu_hadamard import ( + BlockScaledMoEGroupedGemmGluHadamardKernel, + ) + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a = take() + b = take() + sfa = take() + sfb = take() + padded_offsets = take() + alpha = take() + prob = take() + hadamard = take() if has_hadamard else None + bias = take() if has_bias else None + c = take() + d = take() + amax = take() + post_rht_amax = take() + workspace = take() + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped GEMM argument count: consumed {arg_idx}, received {len(args)}" + ) + + kernel = BlockScaledMoEGroupedGemmGluHadamardKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=True, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + expert_cnt=expert_cnt, + weight_mode=MoEWeightMode.DENSE, + use_dynamic_sched=use_dynamic_sched, + act_func=act_func, + enable_bias=has_bias, + use_tmem_post_rht_amax=use_tmem_post_rht_amax, + ) + kernel( + a, + b, + sfa, + sfb, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int64(0), + OperandMajorMode.K, + workspace.iterator, + c, + d, + amax, + post_rht_amax, + padded_offsets, + alpha, + prob, + hadamard, + bias, + max_active_clusters, + stream, + linear_offset=cutlass.Float32(1.0 if act_func == "geglu" else 0.0), + ) + + +def _grouped_gemm_glu_hadamard_impl( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + bias_tensor: Optional[Any] = None, + hadamard_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + m_aligned: int = 256, + act_func: str = "swiglu", + use_dynamic_sched: bool = False, + use_tmem_post_rht_amax: bool = False, + cluster_overlap_margin: int = 0, + *, + b_layout: str = "LNK", + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute dense native-FP4 grouped GEMM with GLU and Hadamard amax. + + A and B use JAX's native ``float4_e2m1fn`` dtype with public shapes + ``(1, M, K)`` and ``(L, N, K)``. Raw ``uint8`` FP4 payloads and discrete + per-expert pointer arrays are intentionally not part of this API. + + Runtime ``padded_offsets`` values must be nondecreasing multiples of 256, + must not exceed M, and must end at the number of rows whose outputs are + meaningful. Rows beyond the final offset are returned as zero. Empty + experts produce zero amax values. These value constraints are trusted + while tracing with :func:`jax.jit`. + """ + + from cutlass import Float32 + from cutlass.jax import jax_to_cutlass_dtype + + from .moe_blockscaled_grouped_gemm_glu_hadamard import ( + BlockScaledMoEGroupedGemmGluHadamardKernel, + ) + + kernel = BlockScaledMoEGroupedGemmGluHadamardKernel + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK",)) + a_spec = gemm_a_tensor_spec("LMK") + b_spec = gemm_b_tensor_spec(b_layout) + output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + valid_ab_dtypes=(jnp.float4_e2m1fn,), + ) + if m % FIX_PAD_SIZE: + raise ValueError(f"M must be divisible by {FIX_PAD_SIZE}, got {m}") + if n % 64: + raise ValueError(f"N must be divisible by 64, got {n}") + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + if sf_vec_size not in SF_VEC_SIZES: + raise ValueError( + f"sf_vec_size must be one of {SF_VEC_SIZES}, got {sf_vec_size}" + ) + if act_func not in ("swiglu", "geglu", "srelu"): + raise ValueError( + f"act_func must be 'swiglu', 'geglu', or 'srelu', got {act_func!r}" + ) + sf_dtype = require_grouped_block_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + valid_dtypes=(jnp.float8_e8m0fnu, jnp.float8_e4m3fn), + ) + if sf_dtype == jnp.dtype(jnp.float8_e4m3fn) and sf_vec_size == 32: + raise ValueError("float8_e4m3fn scales require sf_vec_size=16") + require_grouped_probability("prob_tensor", prob_tensor, m=m) + + if bias_tensor is not None: + require_array( + bias_tensor, + name="bias_tensor", + shape=(experts, n), + dtype=(jnp.float16, jnp.bfloat16, jnp.float32), + ) + + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + c_dtype = require_dtype( + c_dtype, (jnp.float16, jnp.bfloat16), name="c_dtype", default=jnp.bfloat16 + ) + d_dtype = require_dtype( + d_dtype, (jnp.float16, jnp.bfloat16), name="d_dtype", default=jnp.bfloat16 + ) + + mma_tiler_mn = require_grouped_mma_tiler( + mma_tiler_mn, allowed_m=(256,), allowed_n=(256,) + ) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + + output_n = n if act_func == "srelu" else n // 2 + if output_n % HADAMARD_SIZE: + raise ValueError( + f"D's N dimension must be divisible by {HADAMARD_SIZE}, got {output_n}" + ) + require_contiguous_alignment("a_tensor", k, 4) + require_contiguous_alignment("b_tensor", k, 4) + require_contiguous_alignment("c_tensor", n, c_dtype.itemsize * 8) + require_contiguous_alignment("d_tensor", output_n, d_dtype.itemsize * 8) + + if not kernel.can_implement( + jax_to_cutlass_dtype(ab_dtype), + jax_to_cutlass_dtype(sf_dtype), + sf_vec_size, + Float32, + jax_to_cutlass_dtype(d_dtype), + True, + mma_tiler_mn, + cluster_shape_mn, + m, + n, + k, + experts, + "k", + "k", + "n", + m_aligned, + ): + raise ValueError("Unsupported grouped GEMM GLU Hadamard configuration") + + has_hadamard = bool(use_tmem_post_rht_amax) + if has_hadamard and hadamard_tensor is not None: + require_array( + hadamard_tensor, + name="hadamard_tensor", + shape=(HADAMARD_SIZE, HADAMARD_SIZE), + dtype=jnp.bfloat16, + ) + elif not has_hadamard and hadamard_tensor is not None: + raise ValueError( + "hadamard_tensor is used only when use_tmem_post_rht_amax=True" + ) + + if _validate_only: + return { + "acc_dtype": acc_dtype, + "c_dtype": c_dtype, + "d_dtype": d_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + } + + scale_spec = block_scale_tensor_spec() + inputs = [ + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + ] + input_specs = [ + a_spec, + b_spec, + scale_spec, + scale_spec, + None, + None, + probability_tensor_spec(), + ] + + if has_hadamard: + if hadamard_tensor is None: + hadamard_tensor = jnp.asarray( + hadamard_values(HADAMARD_SIZE), dtype=jnp.bfloat16 + ) + inputs.append(hadamard_tensor) + input_specs.append(None) + + if bias_tensor is not None: + inputs.append(bias_tensor) + input_specs.append(grouped_bias_tensor_spec()) + + workspace_bytes = max(dense_workspace_bytes(bool(use_dynamic_sched)), 1) + c_tensor, d_tensor, amax_tensor, post_rht_amax_tensor = call_cutedsl( + _launch, + inputs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "expert_cnt": experts, + "act_func": act_func, + "has_bias": bias_tensor is not None, + "has_hadamard": has_hadamard, + "use_dynamic_sched": bool(use_dynamic_sched), + "use_tmem_post_rht_amax": bool(use_tmem_post_rht_amax), + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=( + make_buffer_desc( + "c_tensor", (1, m, n), c_dtype, tensor_spec=output_spec, init_value=0 + ), + make_buffer_desc( + "d_tensor", + (1, m, output_n), + d_dtype, + tensor_spec=output_spec, + init_value=0, + ), + make_buffer_desc("amax_tensor", (experts, 1), jnp.float32, init_value=0.0), + make_buffer_desc( + "post_rht_amax_tensor", (experts, 1), jnp.float32, init_value=0.0 + ), + ), + workspaces=( + make_buffer_desc( + "workspace", + (workspace_bytes,), + jnp.uint8, + tensor_spec=grouped_workspace_tensor_spec(), + init_value=0, + ), + ), + input_specs=input_specs, + ) + return TupleDict( + c_tensor=c_tensor, + d_tensor=d_tensor, + amax_tensor=amax_tensor, + post_rht_amax_tensor=post_rht_amax_tensor, + ) + + +class GroupedGemmGluHadamardSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM + GLU + Hadamard.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_b_tensor: Any, + sample_sfa_tensor: Any, + sample_sfb_tensor: Any, + sample_padded_offsets: Any, + sample_alpha_tensor: Any, + sample_prob_tensor: Any, + sample_bias_tensor: Optional[Any] = None, + sample_hadamard_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + m_aligned: int = 256, + act_func: str = "swiglu", + use_dynamic_sched: bool = False, + use_tmem_post_rht_amax: bool = False, + *, + b_layout: str = "LNK", + ) -> None: + super().__init__() + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK",)) + a_spec = gemm_a_tensor_spec("LMK") + b_spec = gemm_b_tensor_spec(b_layout) + scale_spec = block_scale_tensor_spec() + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + ), + "padded_offsets": self.make_tensor_desc( + sample_padded_offsets, name="sample_padded_offsets" + ), + "alpha_tensor": self.make_tensor_desc( + sample_alpha_tensor, name="sample_alpha_tensor" + ), + "prob_tensor": self.make_tensor_desc( + sample_prob_tensor, + tensor_spec=probability_tensor_spec(), + name="sample_prob_tensor", + ), + "bias_tensor": self.make_optional_tensor_desc( + sample_bias_tensor, + tensor_spec=grouped_bias_tensor_spec(), + name="sample_bias_tensor", + ), + "hadamard_tensor": self.make_optional_tensor_desc( + sample_hadamard_tensor, name="sample_hadamard_tensor" + ), + } + self._config = { + "acc_dtype": self.as_optional_dtype(acc_dtype), + "c_dtype": self.as_optional_dtype(c_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "act_func": act_func, + "use_dynamic_sched": use_dynamic_sched, + "use_tmem_post_rht_amax": use_tmem_post_rht_amax, + "b_layout": b_layout, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_glu_hadamard_impl( + self._sample_descs["a_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["sfb_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["prob_tensor"], + self._sample_descs["bias_tensor"], + self._sample_descs["hadamard_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + bias_tensor: Optional[Any] = None, + hadamard_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + bias_tensor, + hadamard_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + bias_tensor: Optional[Any] = None, + hadamard_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "b_tensor": b_tensor, + "sfa_tensor": sfa_tensor, + "sfb_tensor": sfb_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "prob_tensor": prob_tensor, + "bias_tensor": bias_tensor, + "hadamard_tensor": hadamard_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_glu_hadamard_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "acc_dtype", + "c_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "act_func", + "use_dynamic_sched", + "use_tmem_post_rht_amax", + "b_layout", + ), +) +def grouped_gemm_glu_hadamard_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + bias_tensor: Optional[Any] = None, + hadamard_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + m_aligned: int = 256, + act_func: str = "swiglu", + use_dynamic_sched: bool = False, + use_tmem_post_rht_amax: bool = False, + *, + b_layout: str = "LNK", +) -> TupleDict: + """Compute dense native-FP4 grouped GEMM with GLU and Hadamard amax.""" + + op = GroupedGemmGluHadamardSm100( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + bias_tensor, + hadamard_tensor, + acc_dtype=acc_dtype, + c_dtype=c_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + act_func=act_func, + use_dynamic_sched=use_dynamic_sched, + use_tmem_post_rht_amax=use_tmem_post_rht_amax, + b_layout=b_layout, + ) + return op( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + bias_tensor, + hadamard_tensor, + ) + + +__all__ = [ + "GroupedGemmGluHadamardSm100", + "grouped_gemm_glu_hadamard_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py index 9f88a8d00..a7030ac7e 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py @@ -1,20 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -""" -Grouped GEMM Quant Kernel Module +from ..._operation_api import make_operation_api -This module provides the contiguous grouped GEMM with output quantization -for MoE (Mixture of Experts) workloads on SM100+ GPUs. -Used for FC2 (forward down-projection) and dFC1 (backward FC1 GEMMs). -""" - -from .api import ( - GroupedGemmQuantSm100, - grouped_gemm_quant_wrapper_sm100, +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmQuantSm100", "grouped_gemm_quant_wrapper_sm100")}, + submodules=("api", "jax"), ) - -__all__ = [ - "GroupedGemmQuantSm100", - "grouped_gemm_quant_wrapper_sm100", -] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_quant/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_quant/jax.py new file mode 100644 index 000000000..423163851 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_quant/jax.py @@ -0,0 +1,626 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for dense-weight grouped GEMM quantization on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + make_buffer_desc, + FIX_PAD_SIZE, + MAX_EXPERTS, + TWO_CTA_MMA_TILER_M, + TupleDict, + as_gemm_tensor_desc, + block_scale_shape, + block_scale_tensor_spec, + call_cutedsl, + dense_workspace_bytes, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + grouped_bias_tensor_spec, + grouped_workspace_tensor_spec, + is_fp4_dtype, + is_fp8_dtype, + is_low_precision_output_dtype, + probability_tensor_spec, + require_16_byte_extent, + require_array, + require_dtype, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_input_scales, + require_grouped_mma_tiler, + require_grouped_probability, + require_grouped_vector, + require_layout, +) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + discrete_col_sfd: bool, + expert_cnt: int, + has_bias: bool, + has_row_scale: bool, + generate_sfd: bool, + has_amax: bool, + use_dynamic_sched: bool, + max_active_clusters: int, +): + import cutlass + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from .grouped_gemm_quant import BlockScaledMoEGroupedGemmQuantKernel + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a = take() + b = take() + sfa = take() + sfb = take() + padded_offsets = take() + alpha = take() + prob = take() + row_scale = take() if has_row_scale else None + bias = take() if has_bias else None + norm_const = take() if generate_sfd else None + d = take() + d_col = take() if generate_sfd else None + amax = take() if has_amax else None + sfd_row = take() if generate_sfd else None + sfd_col = take() if generate_sfd else None + workspace = take() + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped GEMM argument count: consumed {arg_idx}, received {len(args)}" + ) + + kernel = BlockScaledMoEGroupedGemmQuantKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + generate_sfd=generate_sfd, + discrete_col_sfd=discrete_col_sfd, + enable_bias=has_bias, + expert_cnt=expert_cnt, + use_dynamic_sched=use_dynamic_sched, + ) + kernel( + a, + b, + sfb, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int64(0), + OperandMajorMode.K, + workspace.iterator, + d, + d_col, + sfa, + sfd_row, + sfd_col, + amax, + norm_const, + padded_offsets, + alpha, + row_scale, + bias, + prob, + max_active_clusters, + stream, + ) + + +def _grouped_gemm_quant_impl( + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + row_scale_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + cluster_overlap_margin: int = 0, + *, + b_layout: str = "LNK", + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute a dense-weight block-scaled grouped GEMM with quantization. + + The grouped dimension is represented by B's leading ``L`` dimension and by the + runtime ``padded_offsets`` tensor. FP8 outputs return row/column E8M0 scale + factors; native FP4 inputs use logical JAX shapes and never raw-byte + reinterpretation. FP16/BF16 outputs return an initialized amax reduction. + """ + + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_spec = gemm_a_tensor_spec("LMK") + b_spec = gemm_b_tensor_spec(b_layout) + output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + ) + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_input_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ab_dtype=ab_dtype, + ) + if prob_tensor is None: + raise ValueError("prob_tensor is required; pass ones when no gating is needed") + require_grouped_probability("prob_tensor", prob_tensor, m=m) + if row_scale_tensor is not None: + require_grouped_vector("row_scale_tensor", row_scale_tensor, length=m) + if bias_tensor is not None: + require_array( + bias_tensor, + name="bias_tensor", + shape=(experts, n), + dtype=(jnp.float16, jnp.bfloat16, jnp.float32), + ) + + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + if is_fp4_dtype(ab_dtype): + valid_d_dtypes = (jnp.float16, jnp.bfloat16, jnp.float32) + else: + valid_d_dtypes = ( + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + jnp.float4_e2m1fn, + ) + d_dtype = require_dtype( + d_dtype, valid_d_dtypes, name="d_dtype", default=jnp.bfloat16 + ) + if ( + is_fp4_dtype(ab_dtype) + and sf_vec_size == 16 + and d_dtype == jnp.dtype(jnp.float32) + ): + raise NotImplementedError( + "FP4 with sf_vec_size=16 does not support a float32 D output" + ) + generate_sfd = is_fp8_dtype(ab_dtype) and is_low_precision_output_dtype(d_dtype) + has_amax = d_dtype in {jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)} + if generate_sfd: + if norm_const_tensor is None: + raise ValueError( + "norm_const_tensor is required for an FP8/FP4 output from FP8 inputs" + ) + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + discrete_col_sfd = False + if is_fp4_dtype(ab_dtype) and b_layout != "LNK": + raise ValueError("Native FP4 B must use the K-major LNK layout") + mma_tiler_mn = require_grouped_mma_tiler( + mma_tiler_mn, allowed_m=(128, 256), allowed_n=(256,) + ) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + + scale_spec = block_scale_tensor_spec() + require_16_byte_extent("a_tensor", k, ab_dtype) + require_16_byte_extent("b_tensor", n if b_layout == "LKN" else k, ab_dtype) + require_16_byte_extent("d_tensor", n, d_dtype) + + if _validate_only: + return { + "acc_dtype": acc_dtype, + "d_dtype": d_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + } + + inputs = [ + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + ] + input_specs = [ + a_spec, + b_spec, + scale_spec, + scale_spec, + None, + None, + probability_tensor_spec(), + ] + if row_scale_tensor is not None: + inputs.append(row_scale_tensor) + input_specs.append(None) + if bias_tensor is not None: + inputs.append(bias_tensor) + input_specs.append(grouped_bias_tensor_spec()) + if generate_sfd: + inputs.append(norm_const_tensor) + input_specs.append(None) + + outputs = [ + make_buffer_desc("d_tensor", (1, m, n), d_dtype, tensor_spec=output_spec) + ] + if generate_sfd: + outputs.append( + make_buffer_desc( + "d_col_tensor", (1, m, n), d_dtype, tensor_spec=output_spec + ) + ) + if has_amax: + outputs.append( + make_buffer_desc( + "amax_tensor", (experts, 1), jnp.float32, init_value=-float("inf") + ) + ) + if generate_sfd: + outputs.extend( + ( + make_buffer_desc( + "sfd_row_tensor", + block_scale_shape(m, n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + make_buffer_desc( + "sfd_col_tensor", + block_scale_shape(n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + ) + ) + + workspace_bytes = max(dense_workspace_bytes(bool(use_dynamic_sched)), 1) + results = call_cutedsl( + _launch, + inputs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "discrete_col_sfd": bool(discrete_col_sfd), + "expert_cnt": experts, + "has_bias": bias_tensor is not None, + "has_row_scale": row_scale_tensor is not None, + "generate_sfd": generate_sfd, + "has_amax": has_amax, + "use_dynamic_sched": bool(use_dynamic_sched), + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + workspaces=( + make_buffer_desc( + "workspace", + (workspace_bytes,), + jnp.uint8, + tensor_spec=grouped_workspace_tensor_spec(), + ), + ), + input_specs=input_specs, + ) + result_idx = 1 + d_tensor = results[0] + d_col_tensor = results[result_idx] if generate_sfd else None + result_idx += int(generate_sfd) + amax_tensor = results[result_idx] if has_amax else None + result_idx += int(has_amax) + sfd_row_tensor = results[result_idx] if generate_sfd else None + sfd_col_tensor = results[result_idx + 1] if generate_sfd else None + return TupleDict( + d_tensor=d_tensor, + d_col_tensor=d_col_tensor, + amax_tensor=amax_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + ) + + +class GroupedGemmQuantSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM quantization.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_sfa_tensor: Any, + sample_padded_offsets: Any, + sample_alpha_tensor: Any, + sample_b_tensor: Any, + sample_sfb_tensor: Any, + sample_bias_tensor: Optional[Any] = None, + sample_norm_const_tensor: Optional[Any] = None, + sample_prob_tensor: Optional[Any] = None, + sample_row_scale_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", + ) -> None: + super().__init__() + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_spec = gemm_a_tensor_spec("LMK") + b_spec = gemm_b_tensor_spec(b_layout) + scale_spec = block_scale_tensor_spec() + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + ), + "padded_offsets": self.make_tensor_desc( + sample_padded_offsets, name="sample_padded_offsets" + ), + "alpha_tensor": self.make_tensor_desc( + sample_alpha_tensor, name="sample_alpha_tensor" + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + ), + "bias_tensor": self.make_optional_tensor_desc( + sample_bias_tensor, + tensor_spec=grouped_bias_tensor_spec(), + name="sample_bias_tensor", + ), + "norm_const_tensor": self.make_optional_tensor_desc( + sample_norm_const_tensor, name="sample_norm_const_tensor" + ), + "prob_tensor": self.make_optional_tensor_desc( + sample_prob_tensor, + tensor_spec=probability_tensor_spec(), + name="sample_prob_tensor", + ), + "row_scale_tensor": self.make_optional_tensor_desc( + sample_row_scale_tensor, name="sample_row_scale_tensor" + ), + } + self._config = { + "acc_dtype": self.as_optional_dtype(acc_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "discrete_col_sfd": discrete_col_sfd, + "use_dynamic_sched": use_dynamic_sched, + "b_layout": b_layout, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_quant_impl( + self._sample_descs["a_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["sfb_tensor"], + self._sample_descs["bias_tensor"], + self._sample_descs["norm_const_tensor"], + self._sample_descs["prob_tensor"], + self._sample_descs["row_scale_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + row_scale_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + b_tensor, + sfb_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + row_scale_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + row_scale_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "sfa_tensor": sfa_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "b_tensor": b_tensor, + "sfb_tensor": sfb_tensor, + "bias_tensor": bias_tensor, + "norm_const_tensor": norm_const_tensor, + "prob_tensor": prob_tensor, + "row_scale_tensor": row_scale_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_quant_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "acc_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "use_dynamic_sched", + "b_layout", + ), +) +def grouped_gemm_quant_wrapper_sm100( + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + row_scale_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", +) -> TupleDict: + """Compute an MXFP8 dense-weight grouped GEMM with optional quantization.""" + + op = GroupedGemmQuantSm100( + a_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + b_tensor, + sfb_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + row_scale_tensor, + acc_dtype=acc_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + use_dynamic_sched=use_dynamic_sched, + b_layout=b_layout, + ) + return op( + a_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + b_tensor, + sfb_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + row_scale_tensor, + ) + + +__all__ = [ + "GroupedGemmQuantSm100", + "grouped_gemm_quant_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_srelu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_srelu/__init__.py index 70b2bfe7d..65ad86054 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_srelu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_srelu/__init__.py @@ -1,19 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -""" -Grouped GEMM SReLU Kernel Module +from ..._operation_api import make_operation_api -This module provides the forward grouped GEMM with SReLU activation -for MoE (Mixture of Experts) workloads on SM100+ GPUs. -""" - -from .api import ( - GroupedGemmSreluSm100, - grouped_gemm_srelu_wrapper_sm100, +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmSreluSm100", "grouped_gemm_srelu_wrapper_sm100")}, + submodules=("api", "jax"), ) - -__all__ = [ - "GroupedGemmSreluSm100", - "grouped_gemm_srelu_wrapper_sm100", -] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_srelu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_srelu/jax.py new file mode 100644 index 000000000..aac5dd0c5 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_srelu/jax.py @@ -0,0 +1,627 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for dense-weight grouped GEMM + squared ReLU on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + make_buffer_desc, + FIX_PAD_SIZE, + MAX_EXPERTS, + TWO_CTA_MMA_TILER_M, + TupleDict, + as_gemm_tensor_desc, + block_scale_shape, + block_scale_tensor_spec, + call_cutedsl, + dense_workspace_bytes, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + grouped_bias_tensor_spec, + grouped_workspace_tensor_spec, + is_fp4_dtype, + is_fp8_dtype, + is_low_precision_output_dtype, + probability_tensor_spec, + require_16_byte_extent, + require_array, + require_dtype, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_input_scales, + require_grouped_mma_tiler, + require_grouped_probability, + require_grouped_vector, + require_layout, +) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + discrete_col_sfd: bool, + expert_cnt: int, + has_bias: bool, + generate_sfd: bool, + has_amax: bool, + use_dynamic_sched: bool, + max_active_clusters: int, +): + import cutlass + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from .moe_blockscaled_grouped_gemm_srelu_quant import ( + BlockScaledMoEGroupedGemmQuantKernel, + EpilogueType, + ) + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a = take() + b = take() + sfa = take() + sfb = take() + padded_offsets = take() + alpha = take() + prob = take() + bias = take() if has_bias else None + norm_const = take() if generate_sfd else None + c = take() + d = take() + d_col = take() if generate_sfd else None + amax = take() if has_amax else None + sfd_row = take() if generate_sfd else None + sfd_col = take() if generate_sfd else None + workspace = take() + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped GEMM argument count: consumed {arg_idx}, received {len(args)}" + ) + + kernel = BlockScaledMoEGroupedGemmQuantKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + generate_sfd=generate_sfd, + discrete_col_sfd=discrete_col_sfd, + generate_c=True, + enable_bias=has_bias, + expert_cnt=expert_cnt, + use_dynamic_sched=use_dynamic_sched, + epilogue_type=EpilogueType.SRELU.value, + ) + kernel( + a, + b, + sfb, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int64(0), + OperandMajorMode.K, + workspace.iterator, + c, + d, + d_col, + sfa, + sfd_row, + sfd_col, + amax, + norm_const, + padded_offsets, + alpha, + bias, + prob, + max_active_clusters, + stream, + ) + + +def _grouped_gemm_srelu_impl( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + cluster_overlap_margin: int = 0, + *, + b_layout: str = "LNK", + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute a dense-weight block-scaled grouped GEMM and squared-ReLU fusion. + + Native FP4 inputs cover MXFP4/NVFP4 without raw-byte reinterpretation. + Low-precision outputs from FP8 inputs return row/column scale factors; + FP16/BF16 outputs return an initialized per-expert amax reduction. + """ + + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_spec = gemm_a_tensor_spec("LMK") + b_spec = gemm_b_tensor_spec(b_layout) + output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + ) + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_input_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ab_dtype=ab_dtype, + ) + if prob_tensor is None: + raise ValueError("prob_tensor is required; pass ones when no gating is needed") + require_grouped_probability("prob_tensor", prob_tensor, m=m) + if bias_tensor is not None: + require_array( + bias_tensor, + name="bias_tensor", + shape=(experts, n), + dtype=(jnp.float16, jnp.bfloat16, jnp.float32), + ) + + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + c_dtype = require_dtype( + c_dtype, + (jnp.float32, jnp.float16, jnp.bfloat16, jnp.float8_e4m3fn, jnp.float8_e5m2), + name="c_dtype", + default=jnp.bfloat16, + ) + if is_fp4_dtype(ab_dtype): + valid_d_dtypes = (jnp.float16, jnp.bfloat16, jnp.float32) + else: + valid_d_dtypes = ( + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + jnp.float4_e2m1fn, + ) + d_dtype = require_dtype( + d_dtype, valid_d_dtypes, name="d_dtype", default=jnp.bfloat16 + ) + if ( + is_fp4_dtype(ab_dtype) + and sf_vec_size == 16 + and d_dtype == jnp.dtype(jnp.float32) + ): + raise NotImplementedError( + "FP4 with sf_vec_size=16 does not support a float32 D output" + ) + generate_sfd = is_fp8_dtype(ab_dtype) and is_low_precision_output_dtype(d_dtype) + has_amax = d_dtype in {jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)} + if generate_sfd: + if norm_const_tensor is None: + raise ValueError( + "norm_const_tensor is required for an FP8/FP4 output from FP8 inputs" + ) + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + discrete_col_sfd = False + if is_fp4_dtype(ab_dtype) and b_layout != "LNK": + raise ValueError("Native FP4 B must use the K-major LNK layout") + mma_tiler_mn = require_grouped_mma_tiler( + mma_tiler_mn, allowed_m=(128, 256), allowed_n=(256,) + ) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + + scale_spec = block_scale_tensor_spec() + require_16_byte_extent("a_tensor", k, ab_dtype) + require_16_byte_extent("b_tensor", n if b_layout == "LKN" else k, ab_dtype) + require_16_byte_extent("c_tensor", n, c_dtype) + require_16_byte_extent("d_tensor", n, d_dtype) + + if _validate_only: + return { + "acc_dtype": acc_dtype, + "c_dtype": c_dtype, + "d_dtype": d_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + } + + inputs = [ + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + ] + input_specs = [ + a_spec, + b_spec, + scale_spec, + scale_spec, + None, + None, + probability_tensor_spec(), + ] + if bias_tensor is not None: + inputs.append(bias_tensor) + input_specs.append(grouped_bias_tensor_spec()) + if generate_sfd: + inputs.append(norm_const_tensor) + input_specs.append(None) + + outputs = [ + make_buffer_desc("c_tensor", (1, m, n), c_dtype, tensor_spec=output_spec), + make_buffer_desc("d_tensor", (1, m, n), d_dtype, tensor_spec=output_spec), + ] + if generate_sfd: + outputs.append( + make_buffer_desc( + "d_col_tensor", (1, m, n), d_dtype, tensor_spec=output_spec + ) + ) + if has_amax: + outputs.append( + make_buffer_desc( + "amax_tensor", (experts, 1), jnp.float32, init_value=-float("inf") + ) + ) + if generate_sfd: + outputs.extend( + ( + make_buffer_desc( + "sfd_row_tensor", + block_scale_shape(m, n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + make_buffer_desc( + "sfd_col_tensor", + block_scale_shape(n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + ) + ) + + workspace_bytes = max(dense_workspace_bytes(bool(use_dynamic_sched)), 1) + results = call_cutedsl( + _launch, + inputs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "discrete_col_sfd": bool(discrete_col_sfd), + "expert_cnt": experts, + "has_bias": bias_tensor is not None, + "generate_sfd": generate_sfd, + "has_amax": has_amax, + "use_dynamic_sched": bool(use_dynamic_sched), + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + workspaces=( + make_buffer_desc( + "workspace", + (workspace_bytes,), + jnp.uint8, + tensor_spec=grouped_workspace_tensor_spec(), + ), + ), + input_specs=input_specs, + ) + result_idx = 0 + c_tensor, d_tensor = results[:2] + result_idx = 2 + d_col_tensor = results[result_idx] if generate_sfd else None + result_idx += int(generate_sfd) + amax_tensor = results[result_idx] if has_amax else None + result_idx += int(has_amax) + sfd_row_tensor = results[result_idx] if generate_sfd else None + sfd_col_tensor = results[result_idx + 1] if generate_sfd else None + return TupleDict( + c_tensor=c_tensor, + d_tensor=d_tensor, + d_col_tensor=d_col_tensor, + amax_tensor=amax_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + ) + + +class GroupedGemmSreluSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM + SReLU.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_b_tensor: Any, + sample_sfa_tensor: Any, + sample_sfb_tensor: Any, + sample_padded_offsets: Any, + sample_alpha_tensor: Any, + sample_bias_tensor: Optional[Any] = None, + sample_norm_const_tensor: Optional[Any] = None, + sample_prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", + ) -> None: + super().__init__() + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_spec = gemm_a_tensor_spec("LMK") + b_spec = gemm_b_tensor_spec(b_layout) + scale_spec = block_scale_tensor_spec() + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + ), + "padded_offsets": self.make_tensor_desc( + sample_padded_offsets, name="sample_padded_offsets" + ), + "alpha_tensor": self.make_tensor_desc( + sample_alpha_tensor, name="sample_alpha_tensor" + ), + "bias_tensor": self.make_optional_tensor_desc( + sample_bias_tensor, + tensor_spec=grouped_bias_tensor_spec(), + name="sample_bias_tensor", + ), + "norm_const_tensor": self.make_optional_tensor_desc( + sample_norm_const_tensor, name="sample_norm_const_tensor" + ), + "prob_tensor": self.make_optional_tensor_desc( + sample_prob_tensor, + tensor_spec=probability_tensor_spec(), + name="sample_prob_tensor", + ), + } + self._config = { + "acc_dtype": self.as_optional_dtype(acc_dtype), + "c_dtype": self.as_optional_dtype(c_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "discrete_col_sfd": discrete_col_sfd, + "use_dynamic_sched": use_dynamic_sched, + "b_layout": b_layout, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_srelu_impl( + self._sample_descs["a_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["sfb_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["bias_tensor"], + self._sample_descs["norm_const_tensor"], + self._sample_descs["prob_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "b_tensor": b_tensor, + "sfa_tensor": sfa_tensor, + "sfb_tensor": sfb_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "bias_tensor": bias_tensor, + "norm_const_tensor": norm_const_tensor, + "prob_tensor": prob_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_srelu_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "acc_dtype", + "c_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "use_dynamic_sched", + "b_layout", + ), +) +def grouped_gemm_srelu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", +) -> TupleDict: + """Compute an MXFP8 dense-weight grouped GEMM and squared-ReLU fusion.""" + + op = GroupedGemmSreluSm100( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + acc_dtype=acc_dtype, + c_dtype=c_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + use_dynamic_sched=use_dynamic_sched, + b_layout=b_layout, + ) + return op( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + ) + + +__all__ = [ + "GroupedGemmSreluSm100", + "grouped_gemm_srelu_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py index 6bcd935f2..d2829021b 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py @@ -1,19 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -""" -Grouped GEMM SwiGLU Kernel Module +from ..._operation_api import make_operation_api -This module provides the forward grouped GEMM with SwiGLU activation -for MoE (Mixture of Experts) workloads on SM100+ GPUs. -""" - -from .api import ( - GroupedGemmSwigluSm100, - grouped_gemm_swiglu_wrapper_sm100, +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmSwigluSm100", "grouped_gemm_swiglu_wrapper_sm100")}, + submodules=("api", "jax"), ) - -__all__ = [ - "GroupedGemmSwigluSm100", - "grouped_gemm_swiglu_wrapper_sm100", -] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/jax.py new file mode 100644 index 000000000..ad147731c --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/jax.py @@ -0,0 +1,564 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for contiguous grouped GEMM + SwiGLU on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + make_buffer_desc, + FIX_PAD_SIZE, + MAX_EXPERTS, + TWO_CTA_MMA_TILER_M, + TupleDict, + as_gemm_tensor_desc, + block_scale_shape, + block_scale_tensor_spec, + call_cutedsl, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + is_fp4_dtype, + is_fp8_dtype, + probability_tensor_spec, + require_16_byte_extent, + require_dtype, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_input_scales, + require_grouped_mma_tiler, + require_grouped_probability, + require_grouped_vector, + require_layout, +) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + discrete_col_sfd: bool, + expert_cnt: int, + has_prob: bool, + generate_sfd: bool, + has_amax: bool, + max_active_clusters: int, +): + from cutlass.jax import jax_to_cutlass_dtype + + from .grouped_gemm_swiglu_quant import BlockScaledContiguousGroupedGemmKernel + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a = take() + b = take() + sfa = take() + sfb = take() + padded_offsets = take() + alpha = take() + norm_const = take() if generate_sfd else None + prob = take() if has_prob else None + c = take() + d = take() + d_col = take() + amax = take() if has_amax else None + sfd_row = take() if generate_sfd else None + sfd_col = take() if generate_sfd else None + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped GEMM argument count: consumed {arg_idx}, received {len(args)}" + ) + + kernel = BlockScaledContiguousGroupedGemmKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vector_f32=vector_f32, + generate_sfd=generate_sfd, + discrete_col_sfd=discrete_col_sfd, + expert_cnt=expert_cnt, + use_mono_increase_expert_idx=True, + ) + kernel( + a, + b, + c, + d, + d_col, + sfa, + sfb, + sfd_row, + sfd_col, + amax, + norm_const, + padded_offsets, + alpha, + prob, + max_active_clusters, + stream, + ) + + +def _grouped_gemm_swiglu_impl( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + cluster_overlap_margin: int = 0, + *, + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute a contiguous block-scaled grouped GEMM and fused SwiGLU. + + ``a_tensor`` has public shape ``(1, M, K)`` and contains all padded expert rows. + ``b_tensor`` has public shape ``(L, N, K)``, where ``L`` is the expert count. + ``padded_offsets`` supplies each expert's cumulative padded row boundary at + runtime. Matrix outputs use public ``LMN`` order. + + FP4 inputs use JAX's native logical dtype; raw packed ``uint8`` storage is + deliberately excluded. Configuration values are static under JAX tracing. + """ + + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + a_spec = gemm_a_tensor_spec("LMK") + b_spec = gemm_b_tensor_spec("LNK") + output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + ) + if n % 64: + raise ValueError(f"b_tensor N must be divisible by 64 for SwiGLU, got {n}") + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_input_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ab_dtype=ab_dtype, + ) + if prob_tensor is not None: + require_grouped_probability("prob_tensor", prob_tensor, m=m) + + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + c_dtype = require_dtype( + c_dtype, + ( + jnp.float32, + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + jnp.float4_e2m1fn, + ), + name="c_dtype", + default=jnp.bfloat16, + ) + if is_fp4_dtype(ab_dtype): + valid_d_dtypes = (jnp.float16, jnp.bfloat16, jnp.float32) + else: + valid_d_dtypes = ( + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + jnp.float4_e2m1fn, + ) + d_dtype = require_dtype( + d_dtype, valid_d_dtypes, name="d_dtype", default=jnp.bfloat16 + ) + if ( + is_fp4_dtype(ab_dtype) + and sf_vec_size == 16 + and d_dtype == jnp.dtype(jnp.float32) + ): + raise NotImplementedError( + "FP4 with sf_vec_size=16 does not support a float32 D output" + ) + generate_sfd = is_fp8_dtype(ab_dtype) + if generate_sfd: + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required for FP8 inputs") + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + discrete_col_sfd = False + if vector_f32 and c_dtype in { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + }: + raise ValueError("vector_f32 does not support an FP8 c_dtype") + mma_tiler_mn = require_grouped_mma_tiler(mma_tiler_mn) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + if mma_tiler_mn[1] == 128 and d_dtype in { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + }: + raise NotImplementedError("FP8 output requires mma_tiler_mn[1] == 256") + + output_n = n // 2 + scale_spec = block_scale_tensor_spec() + require_16_byte_extent("a_tensor", k, ab_dtype) + require_16_byte_extent("b_tensor", k, ab_dtype) + require_16_byte_extent("c_tensor", n, c_dtype) + require_16_byte_extent("d_tensor", output_n, d_dtype) + + if _validate_only: + return { + "acc_dtype": acc_dtype, + "c_dtype": c_dtype, + "d_dtype": d_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + } + + inputs = [ + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + ] + input_specs = [a_spec, b_spec, scale_spec, scale_spec, None, None] + if generate_sfd: + inputs.append(norm_const_tensor) + input_specs.append(None) + if prob_tensor is not None: + inputs.append(prob_tensor) + input_specs.append(probability_tensor_spec()) + + has_amax = d_dtype in {jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)} + output_specs = [ + make_buffer_desc("c_tensor", (1, m, n), c_dtype, tensor_spec=output_spec), + make_buffer_desc( + "d_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + ), + make_buffer_desc( + "d_col_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + ), + ] + if has_amax: + output_specs.append( + make_buffer_desc( + "amax_tensor", (experts, 1), jnp.float32, init_value=-float("inf") + ) + ) + if generate_sfd: + output_specs.extend( + ( + make_buffer_desc( + "sfd_row_tensor", + block_scale_shape(m, output_n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + make_buffer_desc( + "sfd_col_tensor", + block_scale_shape(output_n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + ) + ) + results = call_cutedsl( + _launch, + inputs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "discrete_col_sfd": bool(discrete_col_sfd), + "expert_cnt": experts, + "has_prob": prob_tensor is not None, + "generate_sfd": generate_sfd, + "has_amax": has_amax, + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=output_specs, + input_specs=input_specs, + ) + result_idx = 0 + c_tensor = results[result_idx] + result_idx += 1 + d_tensor = results[result_idx] + result_idx += 1 + d_col_tensor = results[result_idx] + result_idx += 1 + amax_tensor = results[result_idx] if has_amax else None + result_idx += int(has_amax) + sfd_row_tensor = results[result_idx] if generate_sfd else None + sfd_col_tensor = results[result_idx + 1] if generate_sfd else None + return TupleDict( + c_tensor=c_tensor, + d_tensor=d_tensor, + d_col_tensor=d_col_tensor, + amax_tensor=amax_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + ) + + +class GroupedGemmSwigluSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM + SwiGLU.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_b_tensor: Any, + sample_sfa_tensor: Any, + sample_sfb_tensor: Any, + sample_padded_offsets: Any, + sample_alpha_tensor: Any, + sample_norm_const_tensor: Optional[Any] = None, + sample_prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + ) -> None: + super().__init__() + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + a_spec = gemm_a_tensor_spec("LMK") + b_spec = gemm_b_tensor_spec("LNK") + scale_spec = block_scale_tensor_spec() + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + ), + "padded_offsets": self.make_tensor_desc( + sample_padded_offsets, name="sample_padded_offsets" + ), + "alpha_tensor": self.make_tensor_desc( + sample_alpha_tensor, name="sample_alpha_tensor" + ), + "norm_const_tensor": self.make_optional_tensor_desc( + sample_norm_const_tensor, name="sample_norm_const_tensor" + ), + "prob_tensor": self.make_optional_tensor_desc( + sample_prob_tensor, + tensor_spec=probability_tensor_spec(), + name="sample_prob_tensor", + ), + } + self._config = { + "acc_dtype": self.as_optional_dtype(acc_dtype), + "c_dtype": self.as_optional_dtype(c_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "discrete_col_sfd": discrete_col_sfd, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_swiglu_impl( + self._sample_descs["a_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["sfb_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["norm_const_tensor"], + self._sample_descs["prob_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + norm_const_tensor, + prob_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "b_tensor": b_tensor, + "sfa_tensor": sfa_tensor, + "sfb_tensor": sfb_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "norm_const_tensor": norm_const_tensor, + "prob_tensor": prob_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_swiglu_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "acc_dtype", + "c_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + ), +) +def grouped_gemm_swiglu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, +) -> TupleDict: + """Compute a contiguous block-scaled grouped GEMM and fused SwiGLU.""" + + return GroupedGemmSwigluSm100( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + norm_const_tensor, + prob_tensor, + acc_dtype=acc_dtype, + c_dtype=c_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + )( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + norm_const_tensor, + prob_tensor, + ) + + +__all__ = [ + "GroupedGemmSwigluSm100", + "grouped_gemm_swiglu_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/__init__.py index 17a7ed70e..b823f92bc 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/__init__.py @@ -1,12 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmWgradSm100, - grouped_gemm_wgrad_wrapper_sm100, -) +from ..._operation_api import make_operation_api -__all__ = [ - "GroupedGemmWgradSm100", - "grouped_gemm_wgrad_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmWgradSm100", "grouped_gemm_wgrad_wrapper_sm100")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/jax.py new file mode 100644 index 000000000..dc2d6d5a1 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/jax.py @@ -0,0 +1,577 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for dense-output grouped GEMM weight gradients on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp +from cutlass.jax import TensorSpec + +from .._jax_api import ( + ApiBaseJax, + make_buffer_desc, + TWO_CTA_MMA_TILER_M, + TupleDict, + as_dtype, + call_cutedsl, + ceil_div, + grouped_wgrad_workspace_bytes, + grouped_workspace_tensor_spec, + normalize_wgrad_input_order, + require_16_byte_extent, + require_array, + require_dtype, + require_grouped_cluster_shape, + require_grouped_mma_tiler, +) + + +def wgrad_scale_shape( + rows: int, + tokens: int, + sf_vec_size: int, +) -> tuple[int, int]: + """Return the packed 2Dx2D scale-buffer shape.""" + + return ( + ceil_div(rows, 128) * 128, + ceil_div(ceil_div(tokens, sf_vec_size), 4) * 4, + ) + + +def wgrad_a_tensor_spec() -> TensorSpec: + """Describe ``A=(hidden,tokens)`` with tokens contiguous.""" + + return TensorSpec(layout=(1, 0), mode=(0, 1), ptr_assumed_align=16) + + +def wgrad_b_tensor_spec() -> TensorSpec: + """Describe ``B=(tokens,intermediate)`` with tokens contiguous.""" + + return TensorSpec(layout=(0, 1), mode=(0, 1), ptr_assumed_align=16) + + +def wgrad_scale_tensor_spec() -> TensorSpec: + """Describe the packed two-dimensional scale buffers.""" + + return TensorSpec(layout=(1, 0), mode=(0, 1), ptr_assumed_align=16) + + +def wgrad_output_tensor_spec() -> TensorSpec: + """Describe contiguous ``(experts,hidden,intermediate)`` output.""" + + return TensorSpec(layout=(2, 1, 0), mode=(0, 1, 2), ptr_assumed_align=16) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + accumulate_on_output: bool, + expert_cnt: int, + input_order: str, + has_global_scale: bool, + max_active_clusters: int, +): + from cutlass.jax import jax_to_cutlass_dtype + + from ..moe_utils import MoEWeightMode, WGradInputOrder + from .moe_blockscaled_grouped_gemm_wgrad import ( + BlockScaledMoEGroupedGemmWgradKernel, + ) + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a_tensor = take() + b_tensor = take() + sfa_tensor = take() + sfb_tensor = take() + offsets_tensor = take() + global_scale_a = take() if has_global_scale else None + global_scale_b = take() if has_global_scale else None + wgrad_tensor = take() + workspace = take() + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped wgrad argument count: consumed {arg_idx}, " + f"received {len(args)}" + ) + + kernel = BlockScaledMoEGroupedGemmWgradKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=(mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M), + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + accumulate_on_output=accumulate_on_output, + expert_cnt=expert_cnt, + weight_mode=MoEWeightMode.DENSE, + input_order=WGradInputOrder(input_order), + ) + kernel( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + wgrad_tensor, + offsets_tensor, + workspace, + max_active_clusters, + stream, + global_scale_a, + global_scale_b, + None, + ) + + +def _grouped_gemm_wgrad_impl( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + offsets_tensor: Any, + global_scale_a: Optional[Any] = None, + global_scale_b: Optional[Any] = None, + wgrad_tensor: Optional[Any] = None, + acc_dtype: Any = None, + wgrad_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + accumulate_on_output: bool = False, + input_order: str = "tensor2d", + cluster_overlap_margin: int = 0, + *, + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute dense expert weight gradients using the SM100 grouped kernel. + + ``a_tensor`` has shape ``(hidden,tokens)`` and ``b_tensor`` has shape + ``(tokens,intermediate)``. ``offsets_tensor`` contains each expert's + cumulative token end offset and determines grouping at runtime; its static + length determines the output shape + ``(experts,hidden,intermediate)``. Offset values must be non-decreasing, + aligned to ``sf_vec_size``, and no greater than ``tokens``. + + The JAX surface exposes only the dense output mode; raw per-expert pointer + outputs are not accepted. When ``accumulate_on_output=True``, an optional + ``wgrad_tensor`` seeds the reducing epilogue and the returned value contains + the seed plus the newly computed gradients. Omitting the seed preserves the + simple fresh-output behavior by starting from zero. + + ``input_order='tensor_ragged'`` is supported when A and B have already been + packed expert-by-expert according to the runtime offsets. JAX cannot infer + or perform that value-dependent packing from abstract shapes. + """ + + a_shape = require_array( + a_tensor, + name="a_tensor", + rank=2, + dtype=(jnp.float8_e4m3fn, jnp.float8_e5m2), + ) + ab_dtype = as_dtype(a_tensor) + b_shape = require_array(b_tensor, name="b_tensor", rank=2, dtype=ab_dtype) + hidden, tokens = a_shape + b_tokens, intermediate = b_shape + if b_tokens != tokens: + raise ValueError( + "a_tensor and b_tensor token dimensions must match, got " + f"{a_shape} and {b_shape}" + ) + dimensions = { + "hidden": hidden, + "tokens": tokens, + "intermediate": intermediate, + } + nonpositive = [ + f"{name}={value}" for name, value in dimensions.items() if value <= 0 + ] + if nonpositive: + raise ValueError( + "Grouped-wgrad dimensions must be positive, got " + ", ".join(nonpositive) + ) + + if sf_vec_size != 32: + raise ValueError( + f"The JAX FP8 grouped-wgrad path requires sf_vec_size=32, got {sf_vec_size}" + ) + + require_array( + sfa_tensor, + name="sfa_tensor", + shape=wgrad_scale_shape(hidden, tokens, sf_vec_size), + dtype=jnp.float8_e8m0fnu, + ) + sf_dtype = as_dtype(sfa_tensor) + require_array( + sfb_tensor, + name="sfb_tensor", + shape=wgrad_scale_shape(intermediate, tokens, sf_vec_size), + dtype=sf_dtype, + ) + + offsets_shape = require_array( + offsets_tensor, + name="offsets_tensor", + rank=1, + dtype=jnp.int32, + ) + (expert_cnt,) = offsets_shape + if expert_cnt <= 0: + raise ValueError( + f"offsets_tensor must contain at least one expert, got {expert_cnt}" + ) + output_shape = (expert_cnt, hidden, intermediate) + + if wgrad_tensor is not None and not accumulate_on_output: + raise ValueError("wgrad_tensor is only valid when accumulate_on_output=True") + + has_global_scale = global_scale_a is not None or global_scale_b is not None + if has_global_scale: + if global_scale_a is None or global_scale_b is None: + raise ValueError( + "global_scale_a and global_scale_b must be provided together" + ) + require_array( + global_scale_a, + name="global_scale_a", + shape=(expert_cnt,), + dtype=jnp.float32, + ) + require_array( + global_scale_b, + name="global_scale_b", + shape=(expert_cnt,), + dtype=jnp.float32, + ) + + acc_dtype = require_dtype( + acc_dtype, + (jnp.float32,), + name="acc_dtype", + default=jnp.float32, + ) + valid_wgrad_dtypes = (jnp.bfloat16, jnp.float16, jnp.float32) + if wgrad_tensor is not None: + require_array( + wgrad_tensor, + name="wgrad_tensor", + shape=output_shape, + dtype=valid_wgrad_dtypes, + ) + seed_dtype = as_dtype(wgrad_tensor) + wgrad_dtype = require_dtype( + wgrad_dtype, + valid_wgrad_dtypes, + name="wgrad_dtype", + default=seed_dtype, + ) + if wgrad_dtype != seed_dtype: + raise ValueError( + "wgrad_dtype must match wgrad_tensor.dtype, got " + f"{wgrad_dtype} and {seed_dtype}" + ) + else: + wgrad_dtype = require_dtype( + wgrad_dtype, + valid_wgrad_dtypes, + name="wgrad_dtype", + default=jnp.bfloat16, + ) + mma_tiler_mn = require_grouped_mma_tiler(mma_tiler_mn, allowed_m=(128, 256)) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, + mma_tiler_mn=mma_tiler_mn, + ) + input_order_value = normalize_wgrad_input_order(input_order).value + + require_16_byte_extent("a_tensor", tokens, ab_dtype) + require_16_byte_extent("b_tensor", tokens, ab_dtype) + require_16_byte_extent("wgrad_tensor", intermediate, wgrad_dtype) + + if _validate_only: + return { + "acc_dtype": acc_dtype, + "wgrad_dtype": wgrad_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + } + + inputs = [a_tensor, b_tensor, sfa_tensor, sfb_tensor, offsets_tensor] + input_specs = [ + wgrad_a_tensor_spec(), + wgrad_b_tensor_spec(), + wgrad_scale_tensor_spec(), + wgrad_scale_tensor_spec(), + None, + ] + if has_global_scale: + inputs.extend((global_scale_a, global_scale_b)) + input_specs.extend((None, None)) + + (result_wgrad_tensor,) = call_cutedsl( + _launch, + inputs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "accumulate_on_output": bool(accumulate_on_output), + "expert_cnt": expert_cnt, + "input_order": input_order_value, + "has_global_scale": has_global_scale, + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=( + make_buffer_desc( + "wgrad_tensor", + output_shape, + wgrad_dtype, + tensor_spec=wgrad_output_tensor_spec(), + init_value=( + 0.0 if accumulate_on_output and wgrad_tensor is None else None + ), + ), + ), + output_seeds=(wgrad_tensor,), + workspaces=( + make_buffer_desc( + "workspace", + ( + grouped_wgrad_workspace_bytes( + expert_cnt, + input_order_value, + ), + ), + jnp.uint8, + tensor_spec=grouped_workspace_tensor_spec(), + ), + ), + input_specs=input_specs, + ) + return TupleDict(wgrad_tensor=result_wgrad_tensor) + + +class GroupedGemmWgradSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM weight gradients.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_b_tensor: Any, + sample_sfa_tensor: Any, + sample_sfb_tensor: Any, + sample_offsets_tensor: Any, + sample_global_scale_a: Optional[Any] = None, + sample_global_scale_b: Optional[Any] = None, + acc_dtype: Any = None, + wgrad_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + accumulate_on_output: bool = False, + input_order: str = "tensor2d", + *, + sample_wgrad_tensor: Optional[Any] = None, + ) -> None: + super().__init__() + a_spec = wgrad_a_tensor_spec() + b_spec = wgrad_b_tensor_spec() + scale_spec = wgrad_scale_tensor_spec() + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + ), + "offsets_tensor": self.make_tensor_desc( + sample_offsets_tensor, name="sample_offsets_tensor" + ), + "global_scale_a": self.make_optional_tensor_desc( + sample_global_scale_a, name="sample_global_scale_a" + ), + "global_scale_b": self.make_optional_tensor_desc( + sample_global_scale_b, name="sample_global_scale_b" + ), + "wgrad_tensor": self.make_optional_tensor_desc( + sample_wgrad_tensor, + tensor_spec=wgrad_output_tensor_spec(), + name="sample_wgrad_tensor", + ), + } + self._config = { + "acc_dtype": self.as_optional_dtype(acc_dtype), + "wgrad_dtype": self.as_optional_dtype(wgrad_dtype), + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "accumulate_on_output": accumulate_on_output, + "input_order": input_order, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_wgrad_impl( + self._sample_descs["a_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["sfb_tensor"], + self._sample_descs["offsets_tensor"], + self._sample_descs["global_scale_a"], + self._sample_descs["global_scale_b"], + self._sample_descs["wgrad_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + offsets_tensor: Any, + global_scale_a: Optional[Any] = None, + global_scale_b: Optional[Any] = None, + wgrad_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + offsets_tensor, + global_scale_a, + global_scale_b, + wgrad_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + offsets_tensor: Any, + global_scale_a: Optional[Any] = None, + global_scale_b: Optional[Any] = None, + wgrad_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "b_tensor": b_tensor, + "sfa_tensor": sfa_tensor, + "sfb_tensor": sfb_tensor, + "offsets_tensor": offsets_tensor, + "global_scale_a": global_scale_a, + "global_scale_b": global_scale_b, + "wgrad_tensor": wgrad_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_wgrad_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "acc_dtype", + "wgrad_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "accumulate_on_output", + "input_order", + ), +) +def grouped_gemm_wgrad_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + offsets_tensor: Any, + global_scale_a: Optional[Any] = None, + global_scale_b: Optional[Any] = None, + acc_dtype: Any = None, + wgrad_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + accumulate_on_output: bool = False, + input_order: str = "tensor2d", + *, + wgrad_tensor: Optional[Any] = None, +) -> TupleDict: + """Compute dense expert weight gradients using the SM100 grouped kernel. + + Set ``accumulate_on_output=True`` and pass ``wgrad_tensor`` to accumulate + into an existing JAX value. If accumulation is enabled without a seed, the + returned gradient starts from zero. Pointer-table outputs are not supported. + """ + + op = GroupedGemmWgradSm100( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + offsets_tensor, + global_scale_a, + global_scale_b, + acc_dtype=acc_dtype, + wgrad_dtype=wgrad_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + accumulate_on_output=accumulate_on_output, + input_order=input_order, + sample_wgrad_tensor=wgrad_tensor, + ) + return op( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + offsets_tensor, + global_scale_a, + global_scale_b, + wgrad_tensor, + ) + + +__all__ = [ + "GroupedGemmWgradSm100", + "grouped_gemm_wgrad_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/moe_kernel_helpers.py b/python/cudnn/grouped_gemm/moe_kernel_helpers.py index ab41414da..bcf4adbf4 100644 --- a/python/cudnn/grouped_gemm/moe_kernel_helpers.py +++ b/python/cudnn/grouped_gemm/moe_kernel_helpers.py @@ -38,9 +38,12 @@ - Kernel helper functions that don't depend on kernel instance state """ -from typing import Type, Tuple, Union +from __future__ import annotations -import torch +from typing import TYPE_CHECKING, Type, Tuple, Union + +if TYPE_CHECKING: + import torch import cutlass import cutlass.cute as cute @@ -455,6 +458,8 @@ def silu_f32_geglu_scaled(a: Union[float, Float32], fastmath: bool = False) -> U def sigmoid(x): """PyTorch reference sigmoid using exp2 for numerical consistency.""" + import torch + LOG2_E = 1.4426950408889634 exp_x = torch.exp2(x * (-LOG2_E)) ret = 1.0 / (exp_x + 1.0) @@ -471,6 +476,8 @@ def compute_reference_amax(output_tensor: torch.Tensor) -> float: Returns: float: reference amax value """ + import torch + if output_tensor.dtype != torch.float32: output_fp32 = output_tensor.float() else: @@ -500,6 +507,8 @@ def compare_and_report_mismatches( rtol: Relative tolerance max_mismatches: Maximum number of mismatches to report """ + import torch + if gpu_tensor.is_cuda: gpu_data = gpu_tensor.cpu() else: diff --git a/python/cudnn/grouped_gemm/utils.py b/python/cudnn/grouped_gemm/utils.py index 1e58680ef..9302da871 100644 --- a/python/cudnn/grouped_gemm/utils.py +++ b/python/cudnn/grouped_gemm/utils.py @@ -1,5 +1,6 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT + # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions are met: @@ -33,7 +34,12 @@ the forward (grouped_gemm_swiglu) and backward (grouped_gemm_dswiglu) kernels. """ -from typing import Tuple, Union +from __future__ import annotations + +from typing import TYPE_CHECKING, Tuple, Union + +if TYPE_CHECKING: + import torch from cutlass.cutlass_dsl import ( Boolean, @@ -52,7 +58,6 @@ from cutlass.cute.typing import Float32, Int32 import cutlass.cute as cute import cutlass -import torch import cutlass.pipeline as pipeline from cutlass.pipeline import ( Agent, @@ -228,6 +233,8 @@ def warp_redux_sync_fmax( def logical_shape_fp4x2_aware(tensor: torch.Tensor) -> Tuple[int, ...]: """Return correct shapes for NVFP4 tensor.""" + import torch + if tensor.dtype == torch.float4_e2m1fn_x2: innermost_dim_index = next((i for i, s in enumerate(tensor.stride()) if s == 1), None) if innermost_dim_index is None: diff --git a/python/cudnn/jax/__init__.py b/python/cudnn/jax/__init__.py index ae7f5aef6..c4dc9e379 100644 --- a/python/cudnn/jax/__init__.py +++ b/python/cudnn/jax/__init__.py @@ -19,12 +19,19 @@ import jax as _jax except ModuleNotFoundError as error: missing_name = error.name - if missing_name is None or not any(missing_name == prefix or missing_name.startswith(f"{prefix}.") for prefix in _OPTIONAL_DEPENDENCY_PREFIXES): + if missing_name is None or not any( + missing_name == prefix or missing_name.startswith(f"{prefix}.") + for prefix in _OPTIONAL_DEPENDENCY_PREFIXES + ): raise - raise ImportError(f"cuDNN JAX APIs require the optional JAX dependencies. Install them with `{_INSTALL_HINT}`.") from error + raise ImportError( + f"cuDNN JAX APIs require the optional JAX dependencies. Install them with `{_INSTALL_HINT}`." + ) from error if not _cutlass_jax.is_available(): - minimum_version = ".".join(str(part) for part in _cutlass_jax.CUTE_DSL_MIN_SUPPORTED_JAX_VERSION) + minimum_version = ".".join( + str(part) for part in _cutlass_jax.CUTE_DSL_MIN_SUPPORTED_JAX_VERSION + ) installed_version = getattr(_jax, "__version__", "unknown") raise ImportError( f"CUTLASS JAX support is unavailable with JAX {installed_version}; " @@ -45,44 +52,237 @@ "gemm_dsrelu_wrapper_sm100": ("..gemm_dsrelu.jax", "gemm_dsrelu_wrapper_sm100"), "GemmAmaxSm100": ("..gemm_amax.jax", "GemmAmaxSm100"), "gemm_amax_wrapper_sm100": ("..gemm_amax.jax", "gemm_amax_wrapper_sm100"), - "SparseAttentionBackward": ("..deepseek_sparse_attention.sparse_attention_backward.jax", "SparseAttentionBackward"), + "SparseAttentionBackward": ( + "..deepseek_sparse_attention.sparse_attention_backward.jax", + "SparseAttentionBackward", + ), "sparse_attention_backward_wrapper": ( "..deepseek_sparse_attention.sparse_attention_backward.jax", "sparse_attention_backward_wrapper", ), - "IndexerForward": ("..deepseek_sparse_attention.indexer_forward.jax", "IndexerForward"), - "indexer_forward_wrapper": ("..deepseek_sparse_attention.indexer_forward.jax", "indexer_forward_wrapper"), + "IndexerForward": ( + "..deepseek_sparse_attention.indexer_forward.jax", + "IndexerForward", + ), + "indexer_forward_wrapper": ( + "..deepseek_sparse_attention.indexer_forward.jax", + "indexer_forward_wrapper", + ), "IndexerTopK": ("..deepseek_sparse_attention.indexer_top_k.jax", "IndexerTopK"), - "indexer_top_k_wrapper": ("..deepseek_sparse_attention.indexer_top_k.jax", "indexer_top_k_wrapper"), - "local_to_global_wrapper": ("..deepseek_sparse_attention.indexer_top_k.jax", "local_to_global_wrapper"), - "compactify_wrapper": ("..deepseek_sparse_attention.indexer_top_k.jax", "compactify_wrapper"), - "SparseIndexerScoreRecompute": ("..deepseek_sparse_attention.score_recompute.jax", "SparseIndexerScoreRecompute"), + "indexer_top_k_wrapper": ( + "..deepseek_sparse_attention.indexer_top_k.jax", + "indexer_top_k_wrapper", + ), + "local_to_global_wrapper": ( + "..deepseek_sparse_attention.indexer_top_k.jax", + "local_to_global_wrapper", + ), + "compactify_wrapper": ( + "..deepseek_sparse_attention.indexer_top_k.jax", + "compactify_wrapper", + ), + "SparseIndexerScoreRecompute": ( + "..deepseek_sparse_attention.score_recompute.jax", + "SparseIndexerScoreRecompute", + ), "sparse_indexer_score_recompute_wrapper": ( "..deepseek_sparse_attention.score_recompute.jax", "sparse_indexer_score_recompute_wrapper", ), - "SparseAttnScoreRecompute": ("..deepseek_sparse_attention.score_recompute.jax", "SparseAttnScoreRecompute"), + "SparseAttnScoreRecompute": ( + "..deepseek_sparse_attention.score_recompute.jax", + "SparseAttnScoreRecompute", + ), "sparse_attn_score_recompute_wrapper": ( "..deepseek_sparse_attention.score_recompute.jax", "sparse_attn_score_recompute_wrapper", ), - "DenseIndexerScoreRecompute": ("..deepseek_sparse_attention.score_recompute.jax", "DenseIndexerScoreRecompute"), + "DenseIndexerScoreRecompute": ( + "..deepseek_sparse_attention.score_recompute.jax", + "DenseIndexerScoreRecompute", + ), "dense_indexer_score_recompute_wrapper": ( "..deepseek_sparse_attention.score_recompute.jax", "dense_indexer_score_recompute_wrapper", ), - "DenseAttnScoreRecompute": ("..deepseek_sparse_attention.score_recompute.jax", "DenseAttnScoreRecompute"), + "DenseAttnScoreRecompute": ( + "..deepseek_sparse_attention.score_recompute.jax", + "DenseAttnScoreRecompute", + ), "dense_attn_score_recompute_wrapper": ( "..deepseek_sparse_attention.score_recompute.jax", "dense_attn_score_recompute_wrapper", ), - "IndexerBackward": ("..deepseek_sparse_attention.indexer_backward.jax", "IndexerBackward"), - "indexer_backward_wrapper": ("..deepseek_sparse_attention.indexer_backward.jax", "indexer_backward_wrapper"), - "DenseIndexerBackward": ("..deepseek_sparse_attention.indexer_backward.jax", "DenseIndexerBackward"), - "dense_indexer_backward_wrapper": ("..deepseek_sparse_attention.indexer_backward.jax", "dense_indexer_backward_wrapper"), + "IndexerBackward": ( + "..deepseek_sparse_attention.indexer_backward.jax", + "IndexerBackward", + ), + "indexer_backward_wrapper": ( + "..deepseek_sparse_attention.indexer_backward.jax", + "indexer_backward_wrapper", + ), + "DenseIndexerBackward": ( + "..deepseek_sparse_attention.indexer_backward.jax", + "DenseIndexerBackward", + ), + "dense_indexer_backward_wrapper": ( + "..deepseek_sparse_attention.indexer_backward.jax", + "dense_indexer_backward_wrapper", + ), + "GroupedGemmSwigluSm100": ( + "..grouped_gemm.grouped_gemm_swiglu.jax", + "GroupedGemmSwigluSm100", + ), + "grouped_gemm_swiglu_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_swiglu.jax", + "grouped_gemm_swiglu_wrapper_sm100", + ), + "GroupedGemmDswigluSm100": ( + "..grouped_gemm.grouped_gemm_dswiglu.jax", + "GroupedGemmDswigluSm100", + ), + "grouped_gemm_dswiglu_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_dswiglu.jax", + "grouped_gemm_dswiglu_wrapper_sm100", + ), + "GroupedGemmQuantSm100": ( + "..grouped_gemm.grouped_gemm_quant.jax", + "GroupedGemmQuantSm100", + ), + "grouped_gemm_quant_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_quant.jax", + "grouped_gemm_quant_wrapper_sm100", + ), + "GroupedGemmSreluSm100": ( + "..grouped_gemm.grouped_gemm_srelu.jax", + "GroupedGemmSreluSm100", + ), + "grouped_gemm_srelu_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_srelu.jax", + "grouped_gemm_srelu_wrapper_sm100", + ), + "GroupedGemmDsreluSm100": ( + "..grouped_gemm.grouped_gemm_dsrelu.jax", + "GroupedGemmDsreluSm100", + ), + "grouped_gemm_dsrelu_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_dsrelu.jax", + "grouped_gemm_dsrelu_wrapper_sm100", + ), + "GroupedGemmGluSm100": ( + "..grouped_gemm.grouped_gemm_glu.jax", + "GroupedGemmGluSm100", + ), + "grouped_gemm_glu_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_glu.jax", + "grouped_gemm_glu_wrapper_sm100", + ), + "GroupedGemmGluHadamardSm100": ( + "..grouped_gemm.grouped_gemm_glu_hadamard.jax", + "GroupedGemmGluHadamardSm100", + ), + "grouped_gemm_glu_hadamard_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_glu_hadamard.jax", + "grouped_gemm_glu_hadamard_wrapper_sm100", + ), + "GroupedGemmDgluSm100": ( + "..grouped_gemm.grouped_gemm_dglu.jax", + "GroupedGemmDgluSm100", + ), + "grouped_gemm_dglu_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_dglu.jax", + "grouped_gemm_dglu_wrapper_sm100", + ), + "GroupedGemmWgradSm100": ( + "..grouped_gemm.grouped_gemm_wgrad.jax", + "GroupedGemmWgradSm100", + ), + "grouped_gemm_wgrad_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_wgrad.jax", + "grouped_gemm_wgrad_wrapper_sm100", + ), + "DiscreteGroupedGemmSwigluSm100": ( + "..discrete_grouped_gemm.discrete_grouped_gemm_swiglu.jax", + "DiscreteGroupedGemmSwigluSm100", + ), + "discrete_grouped_gemm_swiglu_wrapper_sm100": ( + "..discrete_grouped_gemm.discrete_grouped_gemm_swiglu.jax", + "discrete_grouped_gemm_swiglu_wrapper_sm100", + ), + "DiscreteGroupedGemmDswigluSm100": ( + "..discrete_grouped_gemm.discrete_grouped_gemm_dswiglu.jax", + "DiscreteGroupedGemmDswigluSm100", + ), + "discrete_grouped_gemm_dswiglu_wrapper_sm100": ( + "..discrete_grouped_gemm.discrete_grouped_gemm_dswiglu.jax", + "discrete_grouped_gemm_dswiglu_wrapper_sm100", + ), + "SelectionAttention": ( + "..native_sparse_attention.selection.jax", + "SelectionAttention", + ), + "selection_attention_wrapper": ( + "..native_sparse_attention.selection.jax", + "selection_attention_wrapper", + ), + "CompressionAttention": ( + "..native_sparse_attention.compression.jax", + "CompressionAttention", + ), + "compression_attention_wrapper": ( + "..native_sparse_attention.compression.jax", + "compression_attention_wrapper", + ), + "SlidingWindowAttention": ( + "..native_sparse_attention.sliding_window_attention.jax", + "SlidingWindowAttention", + ), + "sliding_window_attention_wrapper": ( + "..native_sparse_attention.sliding_window_attention.jax", + "sliding_window_attention_wrapper", + ), + "TopKReduction": ("..native_sparse_attention.top_k.jax", "TopKReduction"), + "topk_reduction_wrapper": ( + "..native_sparse_attention.top_k.jax", + "topk_reduction_wrapper", + ), + "SdpafwdSm100D256": ("..sdpa.fwd.jax", "SdpafwdSm100D256"), + "sdpa_fwd_wrapper_sm100_d256": ("..sdpa.fwd.jax", "sdpa_fwd_wrapper_sm100_d256"), + "SdpabwdSm100D256": ("..sdpa.bwd.jax", "SdpabwdSm100D256"), + "sdpa_bwd_wrapper_sm100_d256": ("..sdpa.bwd.jax", "sdpa_bwd_wrapper_sm100_d256"), + "BlockSparseAttentionForward": ( + "..block_sparse_attention.jax", + "BlockSparseAttentionForward", + ), + "block_sparse_attention_forward": ( + "..block_sparse_attention.jax", + "block_sparse_attention_forward", + ), + "BlockSparseAttentionBackward": ( + "..block_sparse_attention.jax", + "BlockSparseAttentionBackward", + ), + "block_sparse_attention_backward": ( + "..block_sparse_attention.jax", + "block_sparse_attention_backward", + ), } -_DSA_EXPORTS = frozenset(name for name, (module_name, _) in _OPERATION_EXPORTS.items() if module_name.startswith("..deepseek_sparse_attention.")) +_DSA_EXPORTS = frozenset( + name + for name, (module_name, _) in _OPERATION_EXPORTS.items() + if module_name.startswith("..deepseek_sparse_attention.") +) +_NSA_EXPORTS = frozenset( + name + for name, (module_name, _) in _OPERATION_EXPORTS.items() + if module_name.startswith("..native_sparse_attention.") +) +_BSA_EXPORTS = frozenset( + name + for name, (module_name, _) in _OPERATION_EXPORTS.items() + if module_name.startswith("..block_sparse_attention.") +) def _load_operation(name: str) -> Any: @@ -92,21 +292,27 @@ def _load_operation(name: str) -> Any: return value -class _DSANamespace: - """Lazy JAX view of the DeepSeek sparse-attention operation family.""" +class _OperationNamespace: + """Lazy view of one JAX operation family.""" + + def __init__(self, qualified_name: str, exports: frozenset[str]) -> None: + self._qualified_name = qualified_name + self._exports = exports def __getattr__(self, name: str) -> Any: - if name not in _DSA_EXPORTS: - raise AttributeError(f"cudnn.jax.DSA has no attribute {name!r}") + if name not in self._exports: + raise AttributeError(f"{self._qualified_name} has no attribute {name!r}") value = _load_operation(name) setattr(self, name, value) return value def __dir__(self) -> list[str]: - return sorted((*vars(self), *_DSA_EXPORTS)) + return sorted((*vars(self), *self._exports)) -DSA = _DSANamespace() +DSA = _OperationNamespace("cudnn.jax.DSA", _DSA_EXPORTS) +NSA = _OperationNamespace("cudnn.jax.NSA", _NSA_EXPORTS) +BSA = _OperationNamespace("cudnn.jax.BSA", _BSA_EXPORTS) def __getattr__(name: str) -> Any: @@ -120,9 +326,11 @@ def __dir__() -> list[str]: __all__ = [ + "BSA", "DSA", "JaxApiBase", "JaxTensorDesc", + "NSA", "TupleDict", *_OPERATION_EXPORTS, ] diff --git a/python/cudnn/native_sparse_attention/__init__.py b/python/cudnn/native_sparse_attention/__init__.py index 6d4f8220b..fb3e4867d 100644 --- a/python/cudnn/native_sparse_attention/__init__.py +++ b/python/cudnn/native_sparse_attention/__init__.py @@ -1,28 +1,37 @@ -from .selection import SelectionAttention, selection_attention_wrapper -from .compression import CompressionAttention, compression_attention_wrapper -from .sliding_window_attention import ( - SlidingWindowAttention, - sliding_window_attention_wrapper, +"""Lazy Torch namespace for native sparse-attention operations.""" + +from typing import Any + +from .._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "selection": ("SelectionAttention", "selection_attention_wrapper"), + "compression": ("CompressionAttention", "compression_attention_wrapper"), + "sliding_window_attention": ( + "SlidingWindowAttention", + "sliding_window_attention_wrapper", + ), + "top_k": ("TopKReduction", "topk_reduction_wrapper"), + }, ) -from .top_k import TopKReduction, topk_reduction_wrapper +_OPERATION_EXPORTS = frozenset(__all__) class NSANamespace: - SelectionAttention = staticmethod(SelectionAttention) - selection_attention_wrapper = staticmethod(selection_attention_wrapper) + def __getattr__(self, name: str) -> Any: + if name not in _OPERATION_EXPORTS: + raise AttributeError(f"NSA has no attribute {name!r}") + value = __getattr__(name) + setattr(self, name, value) + return value - SlidingWindowAttention = staticmethod(SlidingWindowAttention) - sliding_window_attention_wrapper = staticmethod(sliding_window_attention_wrapper) - - CompressionAttention = staticmethod(CompressionAttention) - compression_attention_wrapper = staticmethod(compression_attention_wrapper) - - TopKReduction = staticmethod(TopKReduction) - topk_reduction_wrapper = staticmethod(topk_reduction_wrapper) + def __dir__(self) -> list[str]: + return sorted((*vars(self), *_OPERATION_EXPORTS)) NSA = NSANamespace() -__all__ = [ - "NSA", -] +# Preserve the historical root contract: operations are accessed through NSA. +__all__ = ["NSA"] diff --git a/python/cudnn/native_sparse_attention/compression/__init__.py b/python/cudnn/native_sparse_attention/compression/__init__.py index b44780956..5e76bf348 100644 --- a/python/cudnn/native_sparse_attention/compression/__init__.py +++ b/python/cudnn/native_sparse_attention/compression/__init__.py @@ -1,6 +1,9 @@ -from .api import CompressionAttention, compression_attention_wrapper +"""Lazy Torch compression-attention API exports.""" -__all__ = [ - "CompressionAttention", - "compression_attention_wrapper", -] +from ..._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("CompressionAttention", "compression_attention_wrapper")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/native_sparse_attention/compression/fmha.py b/python/cudnn/native_sparse_attention/compression/fmha.py index a28eede15..20878b5da 100644 --- a/python/cudnn/native_sparse_attention/compression/fmha.py +++ b/python/cudnn/native_sparse_attention/compression/fmha.py @@ -104,6 +104,7 @@ def __init__( mma_tiler: Tuple[int, int, int], is_persistent: bool, mask_type: fmha_utils.MaskType, + persistent_sm_count: Optional[int] = None, ): """Initializes the configuration for a Blackwell Fused Multi-Head Attention (FMHA) kernel. @@ -155,6 +156,7 @@ def __init__( self.cluster_shape_mn = (1, 1) self.is_persistent = is_persistent self.mask_type = mask_type + self.persistent_sm_count = persistent_sm_count self.softmax0_warp_ids = (0, 1, 2, 3) self.softmax1_warp_ids = (4, 5, 6, 7) @@ -355,6 +357,7 @@ def __call__( cute.shape((s_q, d, ((h_r, h_k), b))), self.cta_tiler, self.is_persistent, + self.persistent_sm_count, ) self.q_major_mode = utils.LayoutEnum.from_tensor(q).mma_major_mode() diff --git a/python/cudnn/native_sparse_attention/compression/fmha_helpers.py b/python/cudnn/native_sparse_attention/compression/fmha_helpers.py index 51627cb80..cb6753cce 100644 --- a/python/cudnn/native_sparse_attention/compression/fmha_helpers.py +++ b/python/cudnn/native_sparse_attention/compression/fmha_helpers.py @@ -137,6 +137,7 @@ def __init__( def get_grid_shape( params: FmhaStaticTileSchedulerParams, *, + persistent_sm_count: Optional[int] = None, loc=None, ip=None, ) -> cute.Shape: @@ -154,10 +155,13 @@ def get_grid_shape( :rtype: cute.Shape """ if params.is_persistent: - hardware_info = HardwareInfo() - sm_count = hardware_info.get_device_multiprocessor_count() + if persistent_sm_count is None: + hardware_info = HardwareInfo() + persistent_sm_count = hardware_info.get_device_multiprocessor_count() + if persistent_sm_count <= 0: + raise ValueError(f"persistent_sm_count must be positive, got {persistent_sm_count}") return ( - min(sm_count, cute.size(params.problem_shape_mbh, loc=loc, ip=ip)), + min(persistent_sm_count, cute.size(params.problem_shape_mbh, loc=loc, ip=ip)), 1, 1, ) @@ -297,6 +301,7 @@ def compute_grid( o_shape: cute.Shape, cta_tiler: Tuple[int, int, int], is_persistent: bool, + persistent_sm_count: Optional[int] = None, ) -> Tuple[FmhaStaticTileSchedulerParams, Tuple[int, int, int]]: """ Compute grid parameters for FMHA operation. @@ -330,7 +335,10 @@ def compute_grid( cute.size(o_shape[2][1]), ), ) - grid = FmhaStaticTileScheduler.get_grid_shape(tile_sched_params) + grid = FmhaStaticTileScheduler.get_grid_shape( + tile_sched_params, + persistent_sm_count=persistent_sm_count, + ) return tile_sched_params, grid diff --git a/python/cudnn/native_sparse_attention/compression/jax.py b/python/cudnn/native_sparse_attention/compression/jax.py new file mode 100644 index 000000000..20cdb529d --- /dev/null +++ b/python/cudnn/native_sparse_attention/compression/jax.py @@ -0,0 +1,559 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for fixed BHSD/BSHD and packed THD NSA compression attention.""" + +from __future__ import annotations + +from functools import partial +import math +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._cute_compiler import compile_options_for_target +from ..._jax import JaxApiBase, TupleDict +from ..jax_utils import ( + BHS_TO_BSH_MODE, + FIXED_LAYOUTS, + describe_bhs_as_bsh, + describe_fixed_data, + fixed_data_mode, + make_fixed_output, + normalize_attention_layout, + normalize_supported_dtype, + require_fixed_qkv, +) + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) + + +class CompressionAttention(JaxApiBase): + """JAX callable specialized from fixed BHSD/BSHD or packed THD metadata.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_cum_seqlen_q: Any | None = None, + sample_cum_seqlen_k: Any | None = None, + enable_lse: bool = False, + o_dtype: Any = None, + qk_acc_dtype: Any = None, + pv_acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + is_persistent: bool = False, + scale_q: float = 1.0, + scale_k: float = 1.0, + scale_v: float = 1.0, + inv_scale_o: float = 1.0, + scale_softmax: float | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + input_dtypes = ( + data_type.HALF, + data_type.BFLOAT16, + data_type.FP8_E4M3, + ) + ranks = tuple(len(sample.shape) for sample in (sample_q, sample_k, sample_v)) + if len(set(ranks)) != 1: + raise ValueError(f"Q, K, and V must all use the same rank, got {ranks}") + self.input_layout = normalize_attention_layout(layout, ranks[0]) + if self.input_layout in FIXED_LAYOUTS: + self.data_mode = fixed_data_mode(self.input_layout, kernel_axes="BSHD") + if sample_cum_seqlen_q is not None or sample_cum_seqlen_k is not None: + raise ValueError( + "cumulative sequence lengths are only valid for packed THD inputs" + ) + if max_s_q is not None or max_s_k is not None: + raise ValueError("max_s_q and max_s_k are only valid for packed THD") + self.q_desc = describe_fixed_data( + sample_q, + "sample_q", + layout=self.input_layout, + kernel_axes="BSHD", + ) + self.k_desc = describe_fixed_data( + sample_k, + "sample_k", + layout=self.input_layout, + kernel_axes="BSHD", + ) + self.v_desc = describe_fixed_data( + sample_v, + "sample_v", + layout=self.input_layout, + kernel_axes="BSHD", + ) + ( + self.batch, + self.num_query_heads, + self.num_kv_heads, + self.seqlen_q, + self.seqlen_k, + self.head_dim, + ) = require_fixed_qkv( + self.q_desc, + self.k_desc, + self.v_desc, + operation_name="CompressionAttention", + kernel_axes="BSHD", + input_dtypes=input_dtypes, + ) + self.cum_q_desc = self.cum_k_desc = None + self.lse_extent = self.seqlen_q + self.q_kernel_desc = self.q_desc + self.k_kernel_desc = self.k_desc + self.v_kernel_desc = self.v_desc + elif self.input_layout == "THD": + self.data_mode = None + if len(sample_k.shape) != 3 or len(sample_v.shape) != 3: + raise ValueError("Q, K, and V must all use the same rank") + if sample_cum_seqlen_q is None or sample_cum_seqlen_k is None: + raise ValueError( + "packed THD inputs require cumulative Q and K sequence lengths" + ) + if max_s_q is None or max_s_k is None: + raise ValueError("packed THD inputs require max_s_q and max_s_k") + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.k_desc = self._to_tensor_desc(sample_k, "sample_k") + self.v_desc = self._to_tensor_desc(sample_v, "sample_v") + total_q, self.num_query_heads, self.head_dim = self.q_desc.shape + total_k, self.num_kv_heads, k_head_dim = self.k_desc.shape + v_total, value_heads, value_dim = self.v_desc.shape + if min(total_q, total_k, self.num_query_heads, self.num_kv_heads) <= 0: + raise ValueError("packed THD dimensions must be positive") + if (v_total, value_heads, value_dim) != ( + total_k, + self.num_kv_heads, + self.head_dim, + ) or k_head_dim != self.head_dim: + raise ValueError("packed K and V shapes must match Q/K head metadata") + if self.q_desc.cudnn_dtype not in input_dtypes: + raise ValueError(f"unsupported Q dtype {self.q_desc.dtype}") + if ( + self.k_desc.cudnn_dtype != self.q_desc.cudnn_dtype + or self.v_desc.cudnn_dtype != self.q_desc.cudnn_dtype + ): + raise ValueError("Q, K, and V must have the same dtype") + if self.head_dim not in (32, 64, 128): + raise ValueError("head dimension must be 32, 64, or 128") + if self.num_query_heads % self.num_kv_heads: + raise ValueError("H_q must be divisible by H_kv") + self.cum_q_desc = self._to_tensor_desc( + sample_cum_seqlen_q, "sample_cum_seqlen_q" + ) + self.cum_k_desc = self._to_tensor_desc( + sample_cum_seqlen_k, "sample_cum_seqlen_k" + ) + if ( + self.cum_q_desc.ndim != 1 + or self.cum_q_desc.shape != self.cum_k_desc.shape + or self.cum_q_desc.shape[0] < 2 + ): + raise ValueError( + "cumulative Q and K sequence lengths must have matching " + "(B + 1,) shapes" + ) + if ( + self.cum_q_desc.cudnn_dtype not in (data_type.INT32, data_type.INT64) + or self.cum_k_desc.cudnn_dtype != self.cum_q_desc.cudnn_dtype + ): + raise ValueError( + "cumulative sequence lengths must share int32 or int64 dtype" + ) + self.batch = self.cum_q_desc.shape[0] - 1 + self.seqlen_q = int(max_s_q) + self.seqlen_k = int(max_s_k) + self.lse_extent = total_q + if ( + self.seqlen_q <= 0 + or self.seqlen_k <= 0 + or self.seqlen_q > total_q + or self.seqlen_k > total_k + ): + raise ValueError( + "max_s_q and max_s_k must be positive and no larger than " + "their packed token counts" + ) + self.q_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *sample_q.shape), sample_q.dtype), + "q_tensor", + public_stride_order=(3, 2, 0, 1), + ) + self.k_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *sample_k.shape), sample_k.dtype), + "k_tensor", + public_stride_order=(3, 2, 0, 1), + ) + self.v_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *sample_v.shape), sample_v.dtype), + "v_tensor", + public_stride_order=(3, 2, 0, 1), + ) + if self.seqlen_q < self.seqlen_k or self.seqlen_q % self.seqlen_k: + raise ValueError( + "Compression attention requires S_q to be an integer multiple of " + f"S_k, got S_q={self.seqlen_q} and S_k={self.seqlen_k}" + ) + if tuple(mma_tiler_mn) != (128, 128): + raise ValueError(f"mma_tiler_mn must be (128, 128), got {mma_tiler_mn}") + + self.enable_lse = bool(enable_lse) + self.is_persistent = bool(is_persistent) + self.output_dtype = normalize_supported_dtype( + o_dtype, + sample_q.dtype, + "o_dtype", + (jnp.float16, jnp.bfloat16, jnp.float8_e4m3fn), + ) + normalize_supported_dtype( + qk_acc_dtype, jnp.float32, "qk_acc_dtype", (jnp.float32,) + ) + normalize_supported_dtype( + pv_acc_dtype, jnp.float32, "pv_acc_dtype", (jnp.float32,) + ) + base_scale = ( + 1.0 / math.sqrt(self.head_dim) + if scale_softmax is None + else float(scale_softmax) + ) + self.scale_softmax = float(scale_q) * float(scale_k) * base_scale + self.scale_output = float(scale_v) * float(inv_scale_o) + self.target_compute_capability = target_compute_capability + self.compute_capability: int | None = None + self.persistent_sm_count: int | None = None + + if self.input_layout in FIXED_LAYOUTS: + self.o_desc = make_fixed_output( + tuple(sample_q.shape), + self.output_dtype, + "o_tensor", + layout=self.input_layout, + kernel_axes="BSHD", + ) + self.o_kernel_desc = self.o_desc + self.lse_desc = ( + describe_bhs_as_bsh( + jax.ShapeDtypeStruct( + (self.batch, self.num_query_heads, self.seqlen_q), + jnp.float32, + ), + "lse_tensor", + ) + if self.enable_lse + else None + ) + self.lse_kernel_desc = self.lse_desc + else: + output_shape = ( + self.q_desc.shape[0], + self.num_query_heads, + self.head_dim, + ) + self.o_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct(output_shape, self.output_dtype), "o_tensor" + ) + self.o_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *output_shape), self.output_dtype), + "o_kernel_tensor", + public_stride_order=(3, 2, 0, 1), + ) + self.lse_desc = ( + self._to_tensor_desc( + jax.ShapeDtypeStruct(output_shape[:2], jnp.float32), + "lse_tensor", + ) + if self.enable_lse + else None + ) + self.lse_kernel_desc = ( + self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *output_shape[:2]), jnp.float32), + "lse_kernel_tensor", + ) + if self.enable_lse + else None + ) + + def check_support(self) -> bool: + self.compute_capability = self._resolve_compute_capability( + self.target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "CompressionAttention", + ) + if self.is_persistent: + self.persistent_sm_count = self._get_device_multiprocessor_count() + return True + + def __call__( + self, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + ) -> TupleDict: + self.check_support() + signature_mode = self.data_mode + for value, desc in ( + (q_tensor, self.q_desc), + (k_tensor, self.k_desc), + (v_tensor, self.v_desc), + ): + self._check_tensor_signature(value, desc, mode=signature_mode) + + if self.input_layout in FIXED_LAYOUTS: + if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: + raise ValueError( + "cumulative sequence lengths must be omitted for fixed inputs" + ) + inputs = (q_tensor, k_tensor, v_tensor) + input_specs = ( + self._to_tensor_spec(self.q_kernel_desc, mode=self.data_mode), + self._to_tensor_spec(self.k_kernel_desc, mode=self.data_mode), + self._to_tensor_spec(self.v_kernel_desc, mode=self.data_mode), + ) + output_specs = ( + self._to_tensor_spec(self.o_kernel_desc, mode=self.data_mode), + ) + if self.lse_kernel_desc is not None: + output_specs += ( + self._to_tensor_spec(self.lse_kernel_desc, mode=BHS_TO_BSH_MODE), + ) + launch = self._launch_kernel + else: + if cum_seqlen_q_tensor is None or cum_seqlen_k_tensor is None: + raise ValueError( + "packed THD inputs require cumulative Q and K sequence lengths" + ) + self._check_tensor_signature(cum_seqlen_q_tensor, self.cum_q_desc) + self._check_tensor_signature(cum_seqlen_k_tensor, self.cum_k_desc) + q_tensor = jnp.reshape(q_tensor, self.q_kernel_desc.shape) + k_tensor = jnp.reshape(k_tensor, self.k_kernel_desc.shape) + v_tensor = jnp.reshape(v_tensor, self.v_kernel_desc.shape) + inputs = ( + q_tensor, + k_tensor, + v_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + input_specs = ( + self._to_tensor_spec(self.q_kernel_desc), + self._to_tensor_spec(self.k_kernel_desc), + self._to_tensor_spec(self.v_kernel_desc), + self._to_tensor_spec(self.cum_q_desc), + self._to_tensor_spec(self.cum_k_desc), + ) + output_specs = (self._to_tensor_spec(self.o_kernel_desc),) + if self.lse_kernel_desc is not None: + output_specs += (self._to_tensor_spec(self.lse_kernel_desc),) + launch = self._launch_varlen_kernel + + results = self._call_kernel( + inputs, + launch=launch, + output_descs=(self.o_kernel_desc,) + if self.lse_kernel_desc is None + else (self.o_kernel_desc, self.lse_kernel_desc), + input_spec=input_specs, + output_spec=output_specs, + compile_options=compile_options_for_target(self.compute_capability), + ) + output = results[0] + lse = results[1] if self.enable_lse else None + if self.input_layout == "THD": + output = jnp.reshape(output, self.o_desc.shape) + if lse is not None: + lse = jnp.reshape(lse, self.lse_desc.shape) + return TupleDict( + o_tensor=output, + lse_tensor=lse, + ) + + def _launch_kernel( + self, stream: Any, q: Any, k: Any, v: Any, output: Any, *optional: Any + ) -> None: + lse = optional[0] if optional else None + self._run_kernel(stream, q, k, v, output, lse, None, None) + + def _launch_varlen_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + cum_seqlen_q: Any, + cum_seqlen_k: Any, + output: Any, + *optional: Any, + ) -> None: + lse = optional[0] if optional else None + self._run_kernel( + stream, + q, + k, + v, + output, + lse, + cum_seqlen_q, + cum_seqlen_k, + ) + + def _run_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + output: Any, + lse: Any, + cum_seqlen_q: Any, + cum_seqlen_k: Any, + ) -> None: + from cutlass import Float32, Int32 + + from .fmha import BlackwellFusedMultiHeadAttentionForward + from .fmha_helpers import MaskType + + kernel = BlackwellFusedMultiHeadAttentionForward( + qk_acc_dtype=Float32, + pv_acc_dtype=Float32, + mma_tiler=(128, 128, self.head_dim), + is_persistent=self.is_persistent, + mask_type=MaskType.COMPRESSED_CAUSAL_MASK, + persistent_sm_count=self.persistent_sm_count, + ) + problem_size = tuple( + Int32(value) + for value in ( + self.batch, + self.seqlen_q, + self.lse_extent, + self.seqlen_k, + self.num_query_heads, + self.num_kv_heads, + self.head_dim, + ) + ) + kernel( + q, + k, + v, + output, + problem_size, + cum_seqlen_q, + cum_seqlen_k, + lse, + Float32(self.scale_softmax * math.log2(math.e)), + Float32(self.scale_softmax), + Float32(self.scale_output), + None, + Int32(0), + stream, + ) + + +@partial( + jax.jit, + static_argnames=( + "enable_lse", + "o_dtype", + "qk_acc_dtype", + "pv_acc_dtype", + "mma_tiler_mn", + "is_persistent", + "scale_q", + "scale_k", + "scale_v", + "inv_scale_o", + "scale_softmax", + "max_s_q", + "max_s_k", + "layout", + "target_compute_capability", + ), +) +def compression_attention_wrapper( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + enable_lse: bool = False, + o_dtype: Any = None, + qk_acc_dtype: Any = None, + pv_acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + is_persistent: bool = False, + scale_q: float = 1.0, + scale_k: float = 1.0, + scale_v: float = 1.0, + inv_scale_o: float = 1.0, + scale_softmax: float | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + layout: str | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute fixed BHSD/BSHD or packed THD compression attention. + + Packed THD calls must provide both cumulative sequence-length arrays and + static ``max_s_q``/``max_s_k`` bounds because JAX traces array values. + Fixed outputs follow ``layout``; LSE remains ``(B, H, S)``. + """ + + qkv_samples = tuple( + jax.ShapeDtypeStruct(value.shape, value.dtype) + for value in (q_tensor, k_tensor, v_tensor) + ) + cum_q_sample = ( + None + if cum_seqlen_q_tensor is None + else jax.ShapeDtypeStruct(cum_seqlen_q_tensor.shape, cum_seqlen_q_tensor.dtype) + ) + cum_k_sample = ( + None + if cum_seqlen_k_tensor is None + else jax.ShapeDtypeStruct(cum_seqlen_k_tensor.shape, cum_seqlen_k_tensor.dtype) + ) + return CompressionAttention( + *qkv_samples, + cum_q_sample, + cum_k_sample, + enable_lse=enable_lse, + o_dtype=o_dtype, + qk_acc_dtype=qk_acc_dtype, + pv_acc_dtype=pv_acc_dtype, + mma_tiler_mn=mma_tiler_mn, + is_persistent=is_persistent, + scale_q=scale_q, + scale_k=scale_k, + scale_v=scale_v, + inv_scale_o=inv_scale_o, + scale_softmax=scale_softmax, + max_s_q=max_s_q, + max_s_k=max_s_k, + layout=layout, + target_compute_capability=target_compute_capability, + )( + q_tensor, + k_tensor, + v_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + + +__all__ = [ + "CompressionAttention", + "SUPPORTED_COMPUTE_CAPABILITIES", + "compression_attention_wrapper", +] diff --git a/python/cudnn/native_sparse_attention/jax_utils.py b/python/cudnn/native_sparse_attention/jax_utils.py new file mode 100644 index 000000000..9d9a25df0 --- /dev/null +++ b/python/cudnn/native_sparse_attention/jax_utils.py @@ -0,0 +1,207 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared tensor metadata and validation for JAX NSA adapters.""" + +from __future__ import annotations + +from typing import Any + +import jax +import jax.numpy as jnp + +from .. import data_type +from .._jax import JaxApiBase, JaxTensorDesc +from .._jax.datatypes import normalize_jax_dtype +from .._jax.layout import mode_from_layout, stride_order_to_public, to_public_axes + +FIXED_LAYOUTS = ("BHSD", "BSHD") +BHS_TO_BSH_MODE = mode_from_layout("BHS", kernel_axes="BSH") + + +def normalize_attention_layout( + layout: str | None, + rank: int, + *, + allow_packed: bool = True, +) -> str: + """Normalize a fixed BHSD/BSHD or packed THD public layout.""" + + if layout is None: + if rank == 4: + return "BHSD" + if allow_packed and rank == 3: + return "THD" + raise ValueError(f"attention layout cannot be inferred from rank {rank}") + if not isinstance(layout, str): + raise TypeError(f"layout must be a string or None, got {layout!r}") + normalized = "".join( + character for character in layout.upper() if character.isalpha() + ) + allowed = (*FIXED_LAYOUTS, "THD") if allow_packed else FIXED_LAYOUTS + if normalized not in allowed: + choices = ", ".join(allowed) + raise ValueError(f"layout must be one of {choices}, got {layout!r}") + expected_rank = 3 if normalized == "THD" else 4 + if rank != expected_rank: + raise ValueError( + f"{normalized} layout requires rank-{expected_rank} tensors, got rank {rank}" + ) + return normalized + + +def fixed_data_mode(layout: str, *, kernel_axes: str) -> tuple[int, ...]: + if layout not in FIXED_LAYOUTS: + raise ValueError(f"fixed layout must be one of {FIXED_LAYOUTS}, got {layout!r}") + return mode_from_layout(layout, kernel_axes=kernel_axes) + + +def describe_fixed_data( + sample: Any, + name: str, + *, + layout: str, + kernel_axes: str, + kernel_stride_order: tuple[int, ...] | None = None, + init_value: bool | int | float | None = None, +) -> JaxTensorDesc: + """Describe public fixed-shape metadata in canonical kernel axes.""" + + mode = fixed_data_mode(layout, kernel_axes=kernel_axes) + public_stride_order = ( + None + if kernel_stride_order is None + else stride_order_to_public(kernel_stride_order, mode) + ) + return JaxApiBase._to_tensor_desc( + sample, + name, + mode=mode, + public_stride_order=public_stride_order, + init_value=init_value, + ) + + +def describe_bhs_as_bsh( + sample: Any, + name: str, + *, + init_value: bool | int | float | None = None, +) -> JaxTensorDesc: + """Describe public BHS metadata as the BSH modes consumed by a kernel.""" + + return JaxApiBase._to_tensor_desc( + sample, + name, + mode=BHS_TO_BSH_MODE, + init_value=init_value, + ) + + +def make_fixed_output( + public_shape: tuple[int, ...], + dtype: Any, + name: str, + *, + layout: str, + kernel_axes: str, + kernel_stride_order: tuple[int, ...] | None = None, + init_value: bool | int | float | None = None, +) -> JaxTensorDesc: + return describe_fixed_data( + jax.ShapeDtypeStruct(public_shape, dtype), + name, + layout=layout, + kernel_axes=kernel_axes, + kernel_stride_order=kernel_stride_order, + init_value=init_value, + ) + + +def require_fixed_qkv( + q_desc: JaxTensorDesc, + k_desc: JaxTensorDesc, + v_desc: JaxTensorDesc | None = None, + *, + operation_name: str, + head_dims: tuple[int, ...] = (32, 64, 128), + kernel_axes: str = "BHSD", + input_dtypes: tuple[data_type, ...] = (data_type.HALF, data_type.BFLOAT16), +) -> tuple[int, int, int, int, int, int]: + """Validate fixed-shape public BHSD Q/K/V signatures.""" + + for desc in (q_desc, k_desc, v_desc): + if desc is not None and desc.ndim != 4: + raise ValueError( + f"{desc.name} must have rank 4 (B, H, S, D), got {desc.shape}" + ) + if q_desc.cudnn_dtype not in input_dtypes: + expected = ", ".join(str(dtype) for dtype in input_dtypes) + raise ValueError( + f"{operation_name} requires one of {{{expected}}}, got {q_desc.dtype}" + ) + for desc in (k_desc, v_desc): + if desc is not None and desc.cudnn_dtype != q_desc.cudnn_dtype: + raise ValueError(f"{desc.name} must have the same dtype as {q_desc.name}") + + logical_bhsd_mode = mode_from_layout("BHSD", kernel_axes=kernel_axes) + q_shape = to_public_axes(q_desc.shape, logical_bhsd_mode) + k_shape = to_public_axes(k_desc.shape, logical_bhsd_mode) + batch, num_query_heads, seqlen_q, head_dim = q_shape + k_batch, num_kv_heads, seqlen_k, k_head_dim = k_shape + dimensions = (batch, num_query_heads, num_kv_heads, seqlen_q, seqlen_k, head_dim) + if any(value <= 0 for value in dimensions): + raise ValueError( + f"{operation_name} dimensions must be positive, got {dimensions}" + ) + if k_batch != batch: + raise ValueError( + f"Q and K batch dimensions must match, got {batch} and {k_batch}" + ) + if k_head_dim != head_dim: + raise ValueError( + f"Q and K head dimensions must match, got {head_dim} and {k_head_dim}" + ) + if head_dim not in head_dims: + expected = ", ".join(str(value) for value in head_dims) + raise ValueError(f"head dimension must be one of {expected}, got {head_dim}") + if num_query_heads % num_kv_heads: + raise ValueError( + f"H_q ({num_query_heads}) must be divisible by H_kv ({num_kv_heads})" + ) + + if v_desc is not None: + v_batch, num_value_heads, v_seqlen, value_dim = to_public_axes( + v_desc.shape, logical_bhsd_mode + ) + if (v_batch, num_value_heads, v_seqlen) != (batch, num_kv_heads, seqlen_k): + raise ValueError("K and V batch, head, and sequence dimensions must match") + if value_dim != head_dim: + raise ValueError( + f"V head dimension must match Q/K ({head_dim}), got {value_dim}" + ) + return dimensions + + +def normalize_supported_dtype( + value: Any | None, default: Any, name: str, allowed: tuple[Any, ...] +) -> Any: + dtype = normalize_jax_dtype(value, default, name) + allowed_dtypes = tuple(jnp.dtype(item) for item in allowed) + if dtype not in allowed_dtypes: + expected = ", ".join(str(item) for item in allowed_dtypes) + raise ValueError(f"{name} must be one of {expected}, got {dtype}") + return dtype + + +__all__ = [ + "BHS_TO_BSH_MODE", + "FIXED_LAYOUTS", + "describe_bhs_as_bsh", + "describe_fixed_data", + "fixed_data_mode", + "make_fixed_output", + "normalize_attention_layout", + "normalize_supported_dtype", + "require_fixed_qkv", +] diff --git a/python/cudnn/native_sparse_attention/selection/__init__.py b/python/cudnn/native_sparse_attention/selection/__init__.py index 576723053..231338d4e 100644 --- a/python/cudnn/native_sparse_attention/selection/__init__.py +++ b/python/cudnn/native_sparse_attention/selection/__init__.py @@ -1,6 +1,9 @@ -from .api import SelectionAttention, selection_attention_wrapper +"""Lazy Torch selection-attention API exports.""" -__all__ = [ - "SelectionAttention", - "selection_attention_wrapper", -] +from ..._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("SelectionAttention", "selection_attention_wrapper")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/native_sparse_attention/selection/jax.py b/python/cudnn/native_sparse_attention/selection/jax.py new file mode 100644 index 000000000..d7082a0a9 --- /dev/null +++ b/python/cudnn/native_sparse_attention/selection/jax.py @@ -0,0 +1,343 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for packed selection attention.""" + +from __future__ import annotations + +from functools import partial +import math +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._cute_compiler import compile_options_for_target +from ..._jax import JaxApiBase, TupleDict +from ..jax_utils import normalize_supported_dtype + +SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) + + +class SelectionAttention(JaxApiBase): + """JAX callable specialized from packed THD selection metadata.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_block_indices: Any, + sample_block_counts: Any, + sample_cum_seqlen: Any, + *, + max_s_q: int, + max_s_k: int, + block_size: int = 64, + scale_softmax: float | None = None, + o_dtype: Any = None, + acc_dtype: Any = None, + target_compute_capability: int | None = None, + ) -> None: + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.k_desc = self._to_tensor_desc(sample_k, "sample_k") + self.v_desc = self._to_tensor_desc(sample_v, "sample_v") + self.block_indices_desc = self._to_tensor_desc( + sample_block_indices, "sample_block_indices" + ) + self.block_counts_desc = self._to_tensor_desc( + sample_block_counts, "sample_block_counts" + ) + self.cum_seqlen_desc = self._to_tensor_desc( + sample_cum_seqlen, "sample_cum_seqlen" + ) + + if self.q_desc.ndim != 3 or self.k_desc.ndim != 3 or self.v_desc.ndim != 3: + raise ValueError( + "SelectionAttention requires packed rank-3 (T, H, D) Q, K, and V" + ) + if self.q_desc.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError( + f"Q must have dtype float16 or bfloat16, got {self.q_desc.dtype}" + ) + if ( + self.k_desc.cudnn_dtype != self.q_desc.cudnn_dtype + or self.v_desc.cudnn_dtype != self.q_desc.cudnn_dtype + ): + raise ValueError("Q, K, and V must have the same dtype") + + self.total_tokens, self.num_query_heads, self.head_dim = self.q_desc.shape + k_tokens, self.num_kv_heads, k_head_dim = self.k_desc.shape + v_tokens, num_value_heads, self.value_dim = self.v_desc.shape + dimensions = ( + self.total_tokens, + self.num_query_heads, + self.num_kv_heads, + self.head_dim, + self.value_dim, + ) + if any(value <= 0 for value in dimensions): + raise ValueError( + f"SelectionAttention dimensions must be positive, got {dimensions}" + ) + if (k_tokens, v_tokens) != (self.total_tokens, self.total_tokens): + raise ValueError("Q, K, and V must have the same packed token count") + if k_head_dim != self.head_dim: + raise ValueError("Q and K head dimensions must match") + if num_value_heads != self.num_kv_heads: + raise ValueError("K and V head counts must match") + if self.head_dim % 16 or self.value_dim % 16: + raise ValueError("Q/K and V head dimensions must be multiples of 16") + if self.num_query_heads % self.num_kv_heads: + raise ValueError("H_q must be divisible by H_kv") + self.gqa_group_size = self.num_query_heads // self.num_kv_heads + if self.gqa_group_size not in (1, 2, 4, 8, 16): + raise ValueError( + f"H_q / H_kv must be one of {{1, 2, 4, 8, 16}}, got {self.gqa_group_size}" + ) + + if self.block_indices_desc.ndim != 3 or self.block_indices_desc.shape[:2] != ( + self.total_tokens, + self.num_kv_heads, + ): + raise ValueError("block_indices must have shape (T, H_kv, K)") + if self.block_indices_desc.shape[2] <= 0: + raise ValueError("block_indices K dimension must be positive") + if self.block_counts_desc.shape != (self.total_tokens, self.num_kv_heads): + raise ValueError("block_counts must have shape (T, H_kv)") + if ( + self.block_indices_desc.cudnn_dtype != data_type.INT32 + or self.block_counts_desc.cudnn_dtype != data_type.INT32 + ): + raise ValueError("block_indices and block_counts must have dtype int32") + if self.cum_seqlen_desc.ndim != 1 or self.cum_seqlen_desc.shape[0] < 2: + raise ValueError("cum_seqlen must have shape (B + 1,) with B > 0") + if self.cum_seqlen_desc.cudnn_dtype not in ( + data_type.INT32, + data_type.INT64, + ): + raise ValueError( + "cumulative sequence lengths must have dtype int32 or int64" + ) + + self.max_s_q = int(max_s_q) + self.max_s_k = int(max_s_k) + if self.max_s_q <= 0 or self.max_s_k <= 0: + raise ValueError("max_s_q and max_s_k must be positive") + if self.max_s_q != self.max_s_k: + raise ValueError( + "SelectionAttention requires max_s_q and max_s_k to be identical" + ) + self.block_size = int(block_size) + if self.block_size not in (16, 32, 64): + raise ValueError("block_size must be 16, 32, or 64") + + self.output_dtype = normalize_supported_dtype( + o_dtype, + sample_q.dtype, + "o_dtype", + (sample_q.dtype,), + ) + normalize_supported_dtype(acc_dtype, jnp.float32, "acc_dtype", (jnp.float32,)) + self.scale_softmax = ( + 1.0 / math.sqrt(self.head_dim) + if scale_softmax is None + else float(scale_softmax) + ) + self.target_compute_capability = target_compute_capability + self.compute_capability: int | None = None + + self.q_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *self.q_desc.shape), self.q_desc.dtype), + "q_tensor", + ) + self.k_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *self.k_desc.shape), self.k_desc.dtype), + "k_tensor", + ) + self.v_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *self.v_desc.shape), self.v_desc.dtype), + "v_tensor", + ) + self.o_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct( + (1, self.total_tokens, self.num_query_heads, self.value_dim), + self.output_dtype, + ), + "o_tensor", + init_value=0, + ) + self.l_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct( + (1, self.total_tokens, self.num_query_heads), jnp.float32 + ), + "l_tensor", + init_value=0.0, + ) + self.m_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct( + (1, self.total_tokens, self.num_query_heads), jnp.float32 + ), + "m_tensor", + init_value=float("-inf"), + ) + + def check_support(self) -> bool: + self.compute_capability = self._resolve_compute_capability( + self.target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "SelectionAttention", + ) + return True + + def __call__( + self, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + block_indices_tensor: Any, + block_counts_tensor: Any, + cum_seqlen_tensor: Any, + ) -> TupleDict: + self.check_support() + for value, desc in ( + (q_tensor, self.q_desc), + (k_tensor, self.k_desc), + (v_tensor, self.v_desc), + (block_indices_tensor, self.block_indices_desc), + (block_counts_tensor, self.block_counts_desc), + (cum_seqlen_tensor, self.cum_seqlen_desc), + ): + self._check_tensor_signature(value, desc) + + q_storage = jnp.reshape(q_tensor, self.q_kernel_desc.shape) + k_storage = jnp.reshape(k_tensor, self.k_kernel_desc.shape) + v_storage = jnp.reshape(v_tensor, self.v_kernel_desc.shape) + output, lse_sum, row_max = self._call_kernel( + ( + q_storage, + k_storage, + v_storage, + block_indices_tensor, + block_counts_tensor, + cum_seqlen_tensor, + ), + launch=self._launch_kernel, + output_descs=(self.o_desc, self.l_desc, self.m_desc), + input_spec=( + self._to_tensor_spec(self.q_kernel_desc), + self._to_tensor_spec(self.k_kernel_desc), + self._to_tensor_spec(self.v_kernel_desc), + self._to_tensor_spec(self.block_indices_desc), + self._to_tensor_spec(self.block_counts_desc), + self._to_tensor_spec(self.cum_seqlen_desc), + ), + compile_options=compile_options_for_target(self.compute_capability), + ) + return TupleDict( + o_tensor=jnp.reshape( + output, (self.total_tokens, self.num_query_heads, self.value_dim) + ), + l_tensor=jnp.reshape(lse_sum, (self.total_tokens, self.num_query_heads, 1)), + m_tensor=jnp.reshape(row_max, (self.total_tokens, self.num_query_heads, 1)), + ) + + def _launch_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + block_indices: Any, + block_counts: Any, + cum_seqlen: Any, + output: Any, + lse_sum: Any, + row_max: Any, + ) -> None: + from cutlass import Float32 + from cutlass.jax import jax_to_cutlass_dtype + + from .NSA_select_attn_fwd_hmma import HopperSelectAttentionFwd + + kernel = HopperSelectAttentionFwd( + head_dim=self.head_dim, + value_dim=self.value_dim, + GQA_group_size=self.gqa_group_size, + block_size=self.block_size, + dtype=jax_to_cutlass_dtype(self.q_desc.dtype), + acc_dtype=Float32, + ) + kernel( + q, + k, + v, + output, + lse_sum, + row_max, + block_indices, + block_counts, + self.max_s_q, + cum_seqlen, + Float32(self.scale_softmax), + stream, + ) + + +@partial( + jax.jit, + static_argnames=( + "max_s_q", + "max_s_k", + "block_size", + "scale_softmax", + "o_dtype", + "acc_dtype", + "target_compute_capability", + ), +) +def selection_attention_wrapper( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + block_indices_tensor: Any, + block_counts_tensor: Any, + cum_seqlen_tensor: Any, + *, + max_s_q: int, + max_s_k: int, + block_size: int = 64, + scale_softmax: float | None = None, + o_dtype: Any = None, + acc_dtype: Any = None, + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute packed THD selection attention.""" + + values = ( + q_tensor, + k_tensor, + v_tensor, + block_indices_tensor, + block_counts_tensor, + cum_seqlen_tensor, + ) + samples = tuple(jax.ShapeDtypeStruct(value.shape, value.dtype) for value in values) + return SelectionAttention( + *samples, + max_s_q=max_s_q, + max_s_k=max_s_k, + block_size=block_size, + scale_softmax=scale_softmax, + o_dtype=o_dtype, + acc_dtype=acc_dtype, + target_compute_capability=target_compute_capability, + )(*values) + + +__all__ = [ + "SUPPORTED_COMPUTE_CAPABILITIES", + "SelectionAttention", + "selection_attention_wrapper", +] diff --git a/python/cudnn/native_sparse_attention/sliding_window_attention/__init__.py b/python/cudnn/native_sparse_attention/sliding_window_attention/__init__.py index 7c0dc8fc7..503009633 100644 --- a/python/cudnn/native_sparse_attention/sliding_window_attention/__init__.py +++ b/python/cudnn/native_sparse_attention/sliding_window_attention/__init__.py @@ -1,6 +1,9 @@ -from .api import SlidingWindowAttention, sliding_window_attention_wrapper +"""Lazy Torch sliding-window-attention API exports.""" -__all__ = [ - "SlidingWindowAttention", - "sliding_window_attention_wrapper", -] +from ..._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("SlidingWindowAttention", "sliding_window_attention_wrapper")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/native_sparse_attention/sliding_window_attention/jax.py b/python/cudnn/native_sparse_attention/sliding_window_attention/jax.py new file mode 100644 index 000000000..5125b7f8d --- /dev/null +++ b/python/cudnn/native_sparse_attention/sliding_window_attention/jax.py @@ -0,0 +1,369 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for fixed and fully packed NSA sliding-window attention.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._jax import JaxApiBase, TupleDict +from ..jax_utils import ( + FIXED_LAYOUTS, + describe_fixed_data, + fixed_data_mode, + normalize_attention_layout, + normalize_supported_dtype, + require_fixed_qkv, +) + + +class SlidingWindowAttention(JaxApiBase): + """JAX callable for fixed BHSD/BSHD or fully packed THD attention. + + Packed inputs use cumulative sequence lengths and static maxima. Their + dynamic values must start at zero, be monotonic, end at the packed token + count, and not exceed the supplied maxima. + """ + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_cum_seqlen_q: Any | None = None, + sample_cum_seqlen_k: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + left_bound: int = 1, + right_bound: int = 0, + is_infer: bool = False, + attn_scale: float | None = None, + o_dtype: Any = None, + layout: str | None = None, + ) -> None: + ranks = tuple(len(sample.shape) for sample in (sample_q, sample_k, sample_v)) + if len(set(ranks)) != 1: + raise ValueError(f"Q, K, and V must use the same rank, got {ranks}") + self.input_layout = normalize_attention_layout(layout, ranks[0]) + if self.input_layout in FIXED_LAYOUTS: + if sample_cum_seqlen_q is not None or sample_cum_seqlen_k is not None: + raise ValueError( + "cumulative sequence lengths must be omitted for fixed inputs" + ) + if max_s_q is not None or max_s_k is not None: + raise ValueError("max_s_q and max_s_k are only valid for THD layout") + self._init_fixed(sample_q, sample_k, sample_v) + else: + self._init_packed( + sample_q, + sample_k, + sample_v, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + max_s_q, + max_s_k, + ) + self.left_bound = int(left_bound) + self.right_bound = int(right_bound) + self.is_infer = bool(is_infer) + self.attn_scale = None if attn_scale is None else float(attn_scale) + self.output_dtype = normalize_supported_dtype( + o_dtype, + sample_q.dtype, + "o_dtype", + (jnp.float16, jnp.bfloat16), + ) + if self.left_bound < 1: + raise ValueError(f"left_bound must be at least 1, got {self.left_bound}") + if self.right_bound < 0: + raise ValueError( + f"right_bound must be non-negative, got {self.right_bound}" + ) + + def _init_fixed(self, sample_q: Any, sample_k: Any, sample_v: Any) -> None: + self.data_mode = fixed_data_mode(self.input_layout, kernel_axes="BHSD") + self.q_desc = describe_fixed_data( + sample_q, + "sample_q", + layout=self.input_layout, + kernel_axes="BHSD", + ) + self.k_desc = describe_fixed_data( + sample_k, + "sample_k", + layout=self.input_layout, + kernel_axes="BHSD", + ) + self.v_desc = describe_fixed_data( + sample_v, + "sample_v", + layout=self.input_layout, + kernel_axes="BHSD", + ) + require_fixed_qkv( + self.q_desc, + self.k_desc, + self.v_desc, + operation_name="SlidingWindowAttention", + kernel_axes="BHSD", + ) + self.cum_q_desc = self.cum_k_desc = None + self.max_s_q = self.max_s_k = None + + def _init_packed( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_cum_seqlen_q: Any | None, + sample_cum_seqlen_k: Any | None, + max_s_q: int | None, + max_s_k: int | None, + ) -> None: + if sample_cum_seqlen_q is None or sample_cum_seqlen_k is None: + raise ValueError( + "packed THD inputs require cumulative Q and K sequence lengths" + ) + if max_s_q is None or max_s_k is None: + raise ValueError("packed THD inputs require max_s_q and max_s_k") + + self.data_mode = None + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.k_desc = self._to_tensor_desc(sample_k, "sample_k") + self.v_desc = self._to_tensor_desc(sample_v, "sample_v") + self.cum_q_desc = self._to_tensor_desc( + sample_cum_seqlen_q, "sample_cum_seqlen_q" + ) + self.cum_k_desc = self._to_tensor_desc( + sample_cum_seqlen_k, "sample_cum_seqlen_k" + ) + self.total_q, self.num_query_heads, self.head_dim = self.q_desc.shape + self.total_k, self.num_kv_heads, k_head_dim = self.k_desc.shape + v_total, v_heads, value_dim = self.v_desc.shape + dimensions = ( + self.total_q, + self.total_k, + self.num_query_heads, + self.num_kv_heads, + self.head_dim, + ) + if any(value <= 0 for value in dimensions): + raise ValueError( + f"SlidingWindowAttention dimensions must be positive, got {dimensions}" + ) + if (v_total, v_heads, value_dim) != ( + self.total_k, + self.num_kv_heads, + self.head_dim, + ) or k_head_dim != self.head_dim: + raise ValueError("packed K and V metadata must match Q/K head dimensions") + if self.q_desc.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError("packed Q, K, and V must use float16 or bfloat16") + if ( + self.k_desc.cudnn_dtype != self.q_desc.cudnn_dtype + or self.v_desc.cudnn_dtype != self.q_desc.cudnn_dtype + ): + raise ValueError("packed Q, K, and V must have the same dtype") + if self.num_query_heads % self.num_kv_heads: + raise ValueError("H_q must be divisible by H_kv") + if ( + self.cum_q_desc.ndim != 1 + or self.cum_q_desc.shape != self.cum_k_desc.shape + or self.cum_q_desc.shape[0] < 2 + ): + raise ValueError( + "cumulative Q and K sequence lengths must have matching (B + 1,) shapes" + ) + if ( + self.cum_q_desc.cudnn_dtype != data_type.INT32 + or self.cum_k_desc.cudnn_dtype != data_type.INT32 + ): + raise ValueError("cumulative sequence lengths must use int32") + self.batch = self.cum_q_desc.shape[0] - 1 + self.max_s_q = int(max_s_q) + self.max_s_k = int(max_s_k) + if ( + self.max_s_q <= 0 + or self.max_s_k <= 0 + or self.max_s_q > self.total_q + or self.max_s_k > self.total_k + ): + raise ValueError( + "max_s_q and max_s_k must be positive and no larger than packed token counts" + ) + + def check_support(self) -> bool: + return True + + def __call__( + self, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + ) -> TupleDict: + self.check_support() + for value, desc in ( + (q_tensor, self.q_desc), + (k_tensor, self.k_desc), + (v_tensor, self.v_desc), + ): + self._check_tensor_signature(value, desc, mode=self.data_mode) + + if self.input_layout == "BHSD": + if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: + raise ValueError( + "cumulative sequence lengths must be omitted for fixed inputs" + ) + q_btnh = jnp.transpose(q_tensor, (0, 2, 1, 3)) + k_btnh = jnp.transpose(k_tensor, (0, 2, 1, 3)) + v_btnh = jnp.transpose(v_tensor, (0, 2, 1, 3)) + query_lengths = key_value_lengths = None + elif self.input_layout == "BSHD": + if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: + raise ValueError( + "cumulative sequence lengths must be omitted for fixed inputs" + ) + q_btnh, k_btnh, v_btnh = q_tensor, k_tensor, v_tensor + query_lengths = key_value_lengths = None + else: + if cum_seqlen_q_tensor is None or cum_seqlen_k_tensor is None: + raise ValueError( + "packed THD inputs require cumulative Q and K sequence lengths" + ) + self._check_tensor_signature(cum_seqlen_q_tensor, self.cum_q_desc) + self._check_tensor_signature(cum_seqlen_k_tensor, self.cum_k_desc) + query_lengths = jnp.diff(cum_seqlen_q_tensor) + key_value_lengths = jnp.diff(cum_seqlen_k_tensor) + q_btnh = self._pad_packed(q_tensor, cum_seqlen_q_tensor[:-1], self.max_s_q) + k_btnh = self._pad_packed(k_tensor, cum_seqlen_k_tensor[:-1], self.max_s_k) + v_btnh = self._pad_packed(v_tensor, cum_seqlen_k_tensor[:-1], self.max_s_k) + attention_result = jax.nn.dot_product_attention( + q_btnh, + k_btnh, + v_btnh, + scale=self.attn_scale, + query_seq_lengths=query_lengths, + key_value_seq_lengths=key_value_lengths, + local_window_size=(self.left_bound - 1, self.right_bound), + implementation="cudnn", + return_residual=not self.is_infer, + ) + if self.is_infer: + output_btnh = attention_result + stats = None + else: + output_btnh, residual_btn = attention_result + stats = residual_btn + if self.input_layout == "THD": + output = self._unpad_packed( + output_btnh, cum_seqlen_q_tensor, self.total_q + ).astype(self.output_dtype) + if stats is not None: + stats = self._unpad_packed( + stats[..., None], cum_seqlen_q_tensor, self.total_q + ) + else: + output = ( + jnp.transpose(output_btnh, (0, 2, 1, 3)) + if self.input_layout == "BHSD" + else output_btnh + ).astype(self.output_dtype) + if stats is not None: + stats = jnp.transpose(stats, (0, 2, 1))[..., None] + return TupleDict(o_tensor=output, stats_tensor=stats) + + @staticmethod + def _pad_packed(values: Any, starts: Any, max_seqlen: int) -> Any: + padded = jnp.pad(values, ((0, max_seqlen), (0, 0), (0, 0))) + slice_shape = (max_seqlen, values.shape[1], values.shape[2]) + return jax.vmap( + lambda start: jax.lax.dynamic_slice(padded, (start, 0, 0), slice_shape) + )(starts) + + @staticmethod + def _unpad_packed(values: Any, cumulative: Any, total_tokens: int) -> Any: + token = jnp.arange(total_tokens, dtype=cumulative.dtype) + batch = jnp.searchsorted(cumulative[1:], token, side="right") + position = token - cumulative[batch] + return values[batch, position] + + +@partial( + jax.jit, + static_argnames=( + "left_bound", + "right_bound", + "is_infer", + "attn_scale", + "o_dtype", + "max_s_q", + "max_s_k", + "layout", + ), +) +def sliding_window_attention_wrapper( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + left_bound: int = 1, + right_bound: int = 0, + is_infer: bool = False, + attn_scale: float | None = None, + o_dtype: Any = None, + layout: str | None = None, +) -> TupleDict: + """Compute fixed BHSD/BSHD or fully packed THD attention. + + Packed THD calls require cumulative sequence lengths and static + ``max_s_q``/``max_s_k`` bounds. Arbitrary ragged byte-offset tensors are + intentionally outside this functional JAX API. + """ + + samples = tuple( + jax.ShapeDtypeStruct(value.shape, value.dtype) + for value in (q_tensor, k_tensor, v_tensor) + ) + sample_cum_seqlen_q = ( + None + if cum_seqlen_q_tensor is None + else jax.ShapeDtypeStruct(cum_seqlen_q_tensor.shape, cum_seqlen_q_tensor.dtype) + ) + sample_cum_seqlen_k = ( + None + if cum_seqlen_k_tensor is None + else jax.ShapeDtypeStruct(cum_seqlen_k_tensor.shape, cum_seqlen_k_tensor.dtype) + ) + return SlidingWindowAttention( + *samples, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + max_s_q=max_s_q, + max_s_k=max_s_k, + left_bound=left_bound, + right_bound=right_bound, + is_infer=is_infer, + attn_scale=attn_scale, + o_dtype=o_dtype, + layout=layout, + )( + q_tensor, + k_tensor, + v_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + + +__all__ = ["SlidingWindowAttention", "sliding_window_attention_wrapper"] diff --git a/python/cudnn/native_sparse_attention/top_k/__init__.py b/python/cudnn/native_sparse_attention/top_k/__init__.py index 06dc32184..d60e6bd78 100644 --- a/python/cudnn/native_sparse_attention/top_k/__init__.py +++ b/python/cudnn/native_sparse_attention/top_k/__init__.py @@ -1,6 +1,9 @@ -from .api import TopKReduction, topk_reduction_wrapper +"""Lazy Torch top-K API exports.""" -__all__ = [ - "TopKReduction", - "topk_reduction_wrapper", -] +from ..._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("TopKReduction", "topk_reduction_wrapper")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/native_sparse_attention/top_k/jax.py b/python/cudnn/native_sparse_attention/top_k/jax.py new file mode 100644 index 000000000..b0e1062c6 --- /dev/null +++ b/python/cudnn/native_sparse_attention/top_k/jax.py @@ -0,0 +1,630 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for SM100 NSA top-K reduction.""" + +from __future__ import annotations + +from functools import partial +import math +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._cute_compiler import compile_options_for_target +from ..._jax import JaxApiBase, TupleDict +from ..._jax.layout import to_public_axes +from ..jax_utils import ( + FIXED_LAYOUTS, + describe_fixed_data, + fixed_data_mode, + make_fixed_output, + normalize_attention_layout, + normalize_supported_dtype, + require_fixed_qkv, +) + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) +_PACKED_LAYOUT = "THD" +_KERNEL_DATA_AXES = "BHSD" +_KERNEL_DATA_STRIDE_ORDER = (3, 1, 2, 0) + + +class TopKReduction(JaxApiBase): + """JAX callable specialized from fixed BHSD/BSHD or packed THD metadata. + + Packed inputs use logical ``(T, H, D)`` Q/K arrays and require both + cumulative sequence-length arrays plus static ``max_s_q`` and ``max_s_k``. + The cumulative arrays remain explicit custom-call operands. Their contents, + including the final packed-token offsets, are therefore validated by the + caller rather than read while JAX traces the operation. + """ + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_lse: Any, + sample_cum_seqlen_q: Any | None = None, + sample_cum_seqlen_k: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + acc_dtype: Any = None, + k_value: int = 16, + selection_block_size: int = 64, + compress_stride: int = 32, + is_causal: bool = True, + mma_tiler_mn: tuple[int, int] = (128, 128), + scale_softmax: float | None = None, + layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + q_rank = len(tuple(sample_q.shape)) + k_rank = len(tuple(sample_k.shape)) + if q_rank != k_rank: + raise ValueError( + f"Q and K must use the same rank, got {q_rank} and {k_rank}" + ) + self.input_layout = normalize_attention_layout(layout, q_rank) + + if self.input_layout in FIXED_LAYOUTS: + self.data_mode = fixed_data_mode( + self.input_layout, kernel_axes=_KERNEL_DATA_AXES + ) + self._init_fixed_shape( + sample_q, + sample_k, + sample_lse, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + ) + else: + self.data_mode = None + self._init_packed_shape( + sample_q, + sample_k, + sample_lse, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + max_s_q, + max_s_k, + ) + + normalize_supported_dtype(acc_dtype, jnp.float32, "acc_dtype", (jnp.float32,)) + self.k_value = int(k_value) + self.selection_block_size = int(selection_block_size) + self.compress_stride = int(compress_stride) + self.is_causal = bool(is_causal) + self.mma_tiler_mn = tuple(mma_tiler_mn) + self._check_configuration() + self.scale_softmax = ( + 1.0 / math.sqrt(self.head_dim) + if scale_softmax is None + else float(scale_softmax) + ) + self.target_compute_capability = target_compute_capability + self.compute_capability: int | None = None + + if self.input_layout in FIXED_LAYOUTS: + canonical_output_shape = ( + self.batch, + self.num_kv_heads, + self.seqlen_q, + self.k_value, + ) + output_shape = to_public_axes(canonical_output_shape, self.data_mode) + self.scores_desc = make_fixed_output( + output_shape, + jnp.float32, + "topk_scores_tensor", + layout=self.input_layout, + kernel_axes=_KERNEL_DATA_AXES, + kernel_stride_order=_KERNEL_DATA_STRIDE_ORDER, + init_value=float("-inf"), + ) + self.indices_desc = make_fixed_output( + output_shape, + jnp.int32, + "topk_indices_tensor", + layout=self.input_layout, + kernel_axes=_KERNEL_DATA_AXES, + kernel_stride_order=_KERNEL_DATA_STRIDE_ORDER, + init_value=-1, + ) + else: + output_shape = (self.total_q_tokens, self.num_kv_heads, self.k_value) + self.scores_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct(output_shape, jnp.float32), + "topk_scores_tensor", + init_value=float("-inf"), + ) + self.indices_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct(output_shape, jnp.int32), + "topk_indices_tensor", + init_value=-1, + ) + + def _init_fixed_shape( + self, + sample_q: Any, + sample_k: Any, + sample_lse: Any, + sample_cum_seqlen_q: Any | None, + sample_cum_seqlen_k: Any | None, + ) -> None: + if sample_cum_seqlen_q is not None or sample_cum_seqlen_k is not None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must both be omitted for BHSD layout" + ) + + # The kernel consumes BHSD modes with the H-inside-S storage used by + # the Torch wrapper's transposed BSHD tensors. + self.q_desc = describe_fixed_data( + sample_q, + "sample_q", + layout=self.input_layout, + kernel_axes=_KERNEL_DATA_AXES, + kernel_stride_order=_KERNEL_DATA_STRIDE_ORDER, + ) + self.k_desc = describe_fixed_data( + sample_k, + "sample_k", + layout=self.input_layout, + kernel_axes=_KERNEL_DATA_AXES, + kernel_stride_order=_KERNEL_DATA_STRIDE_ORDER, + ) + ( + self.batch, + self.num_query_heads, + self.num_kv_heads, + self.seqlen_q, + self.seqlen_k, + self.head_dim, + ) = require_fixed_qkv( + self.q_desc, + self.k_desc, + operation_name="TopKReduction", + kernel_axes=_KERNEL_DATA_AXES, + ) + self.total_q_tokens = self.batch * self.seqlen_q + self.lse_desc = self._to_tensor_desc(sample_lse, "sample_lse") + expected_lse_shape = (self.batch, self.num_query_heads, self.seqlen_q) + if self.lse_desc.shape != expected_lse_shape: + raise ValueError( + f"sample_lse must have shape {expected_lse_shape}, got {self.lse_desc.shape}" + ) + if self.lse_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"sample_lse must have dtype float32, got {self.lse_desc.dtype}" + ) + + self.max_s_q = self.seqlen_q + self.max_s_k = self.seqlen_k + self.cum_seqlen_q_desc = None + self.cum_seqlen_k_desc = None + self.q_kernel_desc = self.q_desc + self.k_kernel_desc = self.k_desc + self.lse_kernel_desc = self.lse_desc + + def _init_packed_shape( + self, + sample_q: Any, + sample_k: Any, + sample_lse: Any, + sample_cum_seqlen_q: Any | None, + sample_cum_seqlen_k: Any | None, + max_s_q: int | None, + max_s_k: int | None, + ) -> None: + if sample_cum_seqlen_q is None or sample_cum_seqlen_k is None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k are both required for THD layout" + ) + if max_s_q is None or max_s_k is None: + raise ValueError("max_s_q and max_s_k are both required for THD layout") + + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.k_desc = self._to_tensor_desc(sample_k, "sample_k") + self.lse_desc = self._to_tensor_desc(sample_lse, "sample_lse") + self.cum_seqlen_q_desc = self._to_tensor_desc( + sample_cum_seqlen_q, "sample_cum_seqlen_q" + ) + self.cum_seqlen_k_desc = self._to_tensor_desc( + sample_cum_seqlen_k, "sample_cum_seqlen_k" + ) + + self.total_q_tokens, self.num_query_heads, self.head_dim = self.q_desc.shape + self.total_k_tokens, self.num_kv_heads, k_head_dim = self.k_desc.shape + dimensions = ( + self.total_q_tokens, + self.total_k_tokens, + self.num_query_heads, + self.num_kv_heads, + self.head_dim, + ) + if any(value <= 0 for value in dimensions): + raise ValueError( + f"TopKReduction dimensions must be positive, got {dimensions}" + ) + if self.q_desc.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError( + f"TopKReduction requires float16 or bfloat16 inputs, got {self.q_desc.dtype}" + ) + if self.k_desc.cudnn_dtype != self.q_desc.cudnn_dtype: + raise ValueError("sample_k must have the same dtype as sample_q") + if k_head_dim != self.head_dim: + raise ValueError( + f"Q and K head dimensions must match, got {self.head_dim} and {k_head_dim}" + ) + if self.head_dim not in (32, 64, 128): + raise ValueError( + f"head dimension must be one of 32, 64, 128, got {self.head_dim}" + ) + if self.num_query_heads % self.num_kv_heads: + raise ValueError( + f"H_q ({self.num_query_heads}) must be divisible by H_kv ({self.num_kv_heads})" + ) + + expected_lse_shapes = ( + (self.total_q_tokens, self.num_query_heads), + (self.total_q_tokens, self.num_query_heads, 1), + ) + if self.lse_desc.shape not in expected_lse_shapes: + raise ValueError( + "sample_lse must have shape " + f"{expected_lse_shapes[0]} or {expected_lse_shapes[1]}, got {self.lse_desc.shape}" + ) + if self.lse_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"sample_lse must have dtype float32, got {self.lse_desc.dtype}" + ) + + if ( + self.cum_seqlen_q_desc.ndim != 1 + or self.cum_seqlen_k_desc.ndim != 1 + or self.cum_seqlen_q_desc.shape != self.cum_seqlen_k_desc.shape + or self.cum_seqlen_q_desc.shape[0] < 2 + ): + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must have the same shape (B + 1,) with B > 0" + ) + if ( + self.cum_seqlen_q_desc.cudnn_dtype != data_type.INT32 + or self.cum_seqlen_k_desc.cudnn_dtype != data_type.INT32 + ): + raise ValueError("cum_seqlen_q and cum_seqlen_k must have dtype int32") + + self.batch = self.cum_seqlen_q_desc.shape[0] - 1 + self.seqlen_q = self.total_q_tokens + self.seqlen_k = self.total_k_tokens + self.max_s_q = int(max_s_q) + self.max_s_k = int(max_s_k) + if ( + self.max_s_q <= 0 + or self.max_s_k <= 0 + or self.max_s_q > self.total_q_tokens + or self.max_s_k > self.total_k_tokens + ): + raise ValueError( + "max_s_q and max_s_k must be positive and no larger than " + "their packed token counts" + ) + + # These virtual rank-4 views preserve the packed THD byte order while + # exposing the B,H,T,D shapes read by the CuTe kernel. LSE is likewise + # presented as B,H,T, matching the Torch wrapper's transpose. + self.q_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct( + (1, self.num_query_heads, self.total_q_tokens, self.head_dim), + self.q_desc.dtype, + ), + "q_tensor", + public_stride_order=_KERNEL_DATA_STRIDE_ORDER, + ) + self.k_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct( + (1, self.num_kv_heads, self.total_k_tokens, self.head_dim), + self.k_desc.dtype, + ), + "k_tensor", + public_stride_order=_KERNEL_DATA_STRIDE_ORDER, + ) + self.lse_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct( + (1, self.num_query_heads, self.total_q_tokens), jnp.float32 + ), + "lse_tensor", + ) + + def _check_configuration(self) -> None: + if self.mma_tiler_mn != (128, 128): + raise ValueError( + f"mma_tiler_mn must be (128, 128), got {self.mma_tiler_mn}" + ) + if self.selection_block_size <= 0 or self.compress_stride <= 0: + raise ValueError( + "selection_block_size and compress_stride must be positive" + ) + if self.selection_block_size % self.compress_stride: + raise ValueError( + "selection_block_size must be divisible by compress_stride" + ) + reduction_width = self.selection_block_size // self.compress_stride + if ( + reduction_width > self.mma_tiler_mn[1] + or self.mma_tiler_mn[1] % reduction_width + ): + raise ValueError( + "selection_block_size / compress_stride must divide the MMA N tile" + ) + candidate_count = self.mma_tiler_mn[1] // reduction_width + if self.k_value <= 0 or self.k_value % 4 or self.k_value > candidate_count: + raise ValueError( + "k_value must be a positive multiple of 4 no greater than the " + f"per-tile candidate count ({candidate_count}), got {self.k_value}" + ) + + def check_support(self) -> bool: + self.compute_capability = self._resolve_compute_capability( + self.target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "TopKReduction", + ) + return True + + def __call__( + self, + q_tensor: Any, + k_tensor: Any, + lse_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + ) -> TupleDict: + self.check_support() + self._check_tensor_signature(q_tensor, self.q_desc, mode=self.data_mode) + self._check_tensor_signature(k_tensor, self.k_desc, mode=self.data_mode) + self._check_tensor_signature(lse_tensor, self.lse_desc) + + if self.input_layout in FIXED_LAYOUTS: + if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must be omitted for fixed layout" + ) + inputs = (q_tensor, k_tensor, lse_tensor) + input_spec = ( + self._to_tensor_spec(self.q_kernel_desc, mode=self.data_mode), + self._to_tensor_spec(self.k_kernel_desc, mode=self.data_mode), + self._to_tensor_spec(self.lse_kernel_desc), + ) + launch = self._launch_kernel + else: + if cum_seqlen_q_tensor is None or cum_seqlen_k_tensor is None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k are both required for THD layout" + ) + self._check_tensor_signature(cum_seqlen_q_tensor, self.cum_seqlen_q_desc) + self._check_tensor_signature(cum_seqlen_k_tensor, self.cum_seqlen_k_desc) + q_storage = jnp.transpose( + jnp.reshape( + q_tensor, + (1, self.total_q_tokens, self.num_query_heads, self.head_dim), + ), + (0, 2, 1, 3), + ) + k_storage = jnp.transpose( + jnp.reshape( + k_tensor, + (1, self.total_k_tokens, self.num_kv_heads, self.head_dim), + ), + (0, 2, 1, 3), + ) + lse_storage = jnp.transpose( + jnp.reshape(lse_tensor, (self.total_q_tokens, self.num_query_heads)), + (1, 0), + )[None, ...] + inputs = ( + q_storage, + k_storage, + lse_storage, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + input_spec = ( + self._to_tensor_spec(self.q_kernel_desc), + self._to_tensor_spec(self.k_kernel_desc), + self._to_tensor_spec(self.lse_kernel_desc), + self._to_tensor_spec(self.cum_seqlen_q_desc), + self._to_tensor_spec(self.cum_seqlen_k_desc), + ) + launch = self._launch_packed_kernel + + output_mode = self.data_mode if self.input_layout in FIXED_LAYOUTS else None + scores, indices = self._call_kernel( + inputs, + launch=launch, + output_descs=(self.scores_desc, self.indices_desc), + input_spec=input_spec, + output_spec=( + self._to_tensor_spec(self.scores_desc, mode=output_mode), + self._to_tensor_spec(self.indices_desc, mode=output_mode), + ), + compile_options=compile_options_for_target(self.compute_capability), + ) + return TupleDict(topk_scores_tensor=scores, topk_indices_tensor=indices) + + def _launch_kernel( + self, + stream: Any, + q: Any, + k: Any, + lse: Any, + topk_scores: Any, + topk_indices: Any, + ) -> None: + self._invoke_kernel( + stream, + q, + k, + lse, + topk_scores, + topk_indices, + None, + None, + ) + + def _launch_packed_kernel( + self, + stream: Any, + q: Any, + k: Any, + lse: Any, + cum_seqlen_q: Any, + cum_seqlen_k: Any, + topk_scores: Any, + topk_indices: Any, + ) -> None: + self._invoke_kernel( + stream, + q, + k, + lse, + topk_scores, + topk_indices, + cum_seqlen_q, + cum_seqlen_k, + ) + + def _invoke_kernel( + self, + stream: Any, + q: Any, + k: Any, + lse: Any, + topk_scores: Any, + topk_indices: Any, + cum_seqlen_q: Any | None, + cum_seqlen_k: Any | None, + ) -> None: + from cutlass import Float32, Int32 + from cutlass.jax import jax_to_cutlass_dtype + + from .nsa_top_k_reduction_fwd import FineGrainedReductionQK + + kernel = FineGrainedReductionQK( + element_dtype=jax_to_cutlass_dtype(self.q_desc.dtype), + acc_dtype=Float32, + k_value=self.k_value, + selection_block_size=self.selection_block_size, + compress_block_sliding_stride=self.compress_stride, + mma_tiler=(128, 128, self.head_dim), + is_causal=self.is_causal, + ) + problem_size = tuple( + Int32(value) + for value in ( + self.batch, + self.max_s_q, + self.max_s_k, + self.num_query_heads, + self.num_kv_heads, + self.head_dim, + ) + ) + kernel( + problem_size, + q, + k, + lse, + topk_scores, + topk_indices, + Float32(self.scale_softmax * math.log2(math.e)), + cum_seqlen_q, + cum_seqlen_k, + stream, + ) + + +@partial( + jax.jit, + static_argnames=( + "max_s_q", + "max_s_k", + "acc_dtype", + "k_value", + "selection_block_size", + "compress_stride", + "is_causal", + "mma_tiler_mn", + "scale_softmax", + "layout", + "target_compute_capability", + ), +) +def topk_reduction_wrapper( + q_tensor: Any, + k_tensor: Any, + lse_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + acc_dtype: Any = None, + k_value: int = 16, + selection_block_size: int = 64, + compress_stride: int = 32, + is_causal: bool = True, + mma_tiler_mn: tuple[int, int] = (128, 128), + scale_softmax: float | None = None, + layout: str | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + """Select blocks for fixed BHSD/BSHD or packed THD inputs. + + Fixed outputs follow ``layout``; fixed LSE remains ``(B, H, S)``. + """ + + samples = tuple( + jax.ShapeDtypeStruct(value.shape, value.dtype) + for value in (q_tensor, k_tensor, lse_tensor) + ) + sample_cum_seqlen_q = ( + None + if cum_seqlen_q_tensor is None + else jax.ShapeDtypeStruct(cum_seqlen_q_tensor.shape, cum_seqlen_q_tensor.dtype) + ) + sample_cum_seqlen_k = ( + None + if cum_seqlen_k_tensor is None + else jax.ShapeDtypeStruct(cum_seqlen_k_tensor.shape, cum_seqlen_k_tensor.dtype) + ) + return TopKReduction( + *samples, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + max_s_q=max_s_q, + max_s_k=max_s_k, + acc_dtype=acc_dtype, + k_value=k_value, + selection_block_size=selection_block_size, + compress_stride=compress_stride, + is_causal=is_causal, + mma_tiler_mn=mma_tiler_mn, + scale_softmax=scale_softmax, + layout=layout, + target_compute_capability=target_compute_capability, + )( + q_tensor, + k_tensor, + lse_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + + +__all__ = [ + "SUPPORTED_COMPUTE_CAPABILITIES", + "TopKReduction", + "topk_reduction_wrapper", +] diff --git a/python/cudnn/sdpa/__init__.py b/python/cudnn/sdpa/__init__.py index b200dcbda..f743f26dc 100644 --- a/python/cudnn/sdpa/__init__.py +++ b/python/cudnn/sdpa/__init__.py @@ -1,12 +1,14 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .bwd import SdpabwdSm100D256, sdpa_bwd_wrapper_sm100_d256 -from .fwd import SdpafwdSm100D256, sdpa_fwd_wrapper_sm100_d256 +"""Lazy Torch API exports for the CuTe SDPA operations.""" -__all__ = [ - "SdpafwdSm100D256", - "sdpa_fwd_wrapper_sm100_d256", - "SdpabwdSm100D256", - "sdpa_bwd_wrapper_sm100_d256", -] +from .._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "fwd": ("SdpafwdSm100D256", "sdpa_fwd_wrapper_sm100_d256"), + "bwd": ("SdpabwdSm100D256", "sdpa_bwd_wrapper_sm100_d256"), + }, +) diff --git a/python/cudnn/sdpa/bwd/__init__.py b/python/cudnn/sdpa/bwd/__init__.py index 0d0c368ed..9568c8ca7 100644 --- a/python/cudnn/sdpa/bwd/__init__.py +++ b/python/cudnn/sdpa/bwd/__init__.py @@ -1,9 +1,12 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import SdpabwdSm100D256, sdpa_bwd_wrapper_sm100_d256 +"""Lazy Torch SDPA-backward API exports.""" -__all__ = [ - "SdpabwdSm100D256", - "sdpa_bwd_wrapper_sm100_d256", -] +from ..._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("SdpabwdSm100D256", "sdpa_bwd_wrapper_sm100_d256")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/sdpa/bwd/jax.py b/python/cudnn/sdpa/bwd/jax.py new file mode 100644 index 000000000..bad2dc364 --- /dev/null +++ b/python/cudnn/sdpa/bwd/jax.py @@ -0,0 +1,723 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for fixed BHSD/BSHD and packed THD SM100 d=256 SDPA backward.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._cute_compiler import compile_options_for_target +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..jax_utils import ( + FIXED_LAYOUTS, + describe_fixed_data, + fixed_data_mode, + make_fixed_output, + normalize_sdpa_layout, + require_fixed_qkv, + require_float32_dtype, + resolve_sdpa_config, +) + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) +_PACKED_LAYOUT = "THD" + + +class SdpabwdSm100D256(JaxApiBase): + """JAX callable specialized from fixed BHSD/BSHD or packed THD metadata. + + Packed THD inputs require cumulative sequence lengths and explicit static + ``max_s_q``/``max_s_k`` bounds. JAX samples expose array metadata rather + than cumulative-length values, so these bounds cannot be inferred while + tracing. + """ + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_o: Any, + sample_do: Any, + sample_lse: Any, + sample_cum_seqlen_q: Any | None = None, + sample_cum_seqlen_k: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + dkdv_mma_tiler_mn: tuple[int, int] = (128, 64), + is_causal: bool = False, + window_size: tuple[int, int] = (-1, -1), + scale_softmax: float | None = None, + layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + ranks = tuple( + len(tuple(sample.shape)) + for sample in (sample_q, sample_k, sample_v, sample_o, sample_do) + ) + if len(set(ranks)) != 1: + raise ValueError(f"Q, K, V, O, and dO must use the same rank, got {ranks}") + self.input_layout = normalize_sdpa_layout(layout, ranks[0]) + if self.input_layout in FIXED_LAYOUTS: + if max_s_q is not None or max_s_k is not None: + raise ValueError("max_s_q and max_s_k are only valid for THD layout") + self.data_mode = fixed_data_mode(self.input_layout) + self._init_fixed_shape( + sample_q, + sample_k, + sample_v, + sample_o, + sample_do, + sample_lse, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + ) + elif self.input_layout == _PACKED_LAYOUT: + self.data_mode = None + self._init_packed_shape( + sample_q, + sample_k, + sample_v, + sample_o, + sample_do, + sample_lse, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + max_s_q, + max_s_k, + ) + + require_float32_dtype(acc_dtype, "acc_dtype") + if tuple(mma_tiler_mn) != (128, 128): + raise ValueError(f"mma_tiler_mn must be (128, 128), got {mma_tiler_mn}") + if tuple(dkdv_mma_tiler_mn) != (128, 64): + raise ValueError( + f"dkdv_mma_tiler_mn must be (128, 64), got {dkdv_mma_tiler_mn}" + ) + + self.is_causal = bool(is_causal) + ( + self.scale_softmax, + self.window_size_left, + self.window_size_right, + self.mask_kind, + ) = resolve_sdpa_config( + seqlen_q=self.max_s_q, + seqlen_k=self.max_s_k, + # The Torch API always selects the residual mask for non-causal + # varlen inputs, independent of the packed token count. + tile_extent=self.seqlen_q if self.input_layout in FIXED_LAYOUTS else 1, + is_causal=self.is_causal, + window_size=tuple(window_size), + scale_softmax=scale_softmax, + ) + self.target_compute_capability = target_compute_capability + self.compute_capability: int | None = None + + if self.input_layout in FIXED_LAYOUTS: + self.dq_desc = make_fixed_output( + tuple(sample_q.shape), + sample_q.dtype, + "dq_tensor", + layout=self.input_layout, + ) + self.dk_desc = make_fixed_output( + tuple(sample_k.shape), + sample_k.dtype, + "dk_tensor", + layout=self.input_layout, + ) + self.dv_desc = make_fixed_output( + tuple(sample_v.shape), + sample_v.dtype, + "dv_tensor", + layout=self.input_layout, + ) + self.dq_kernel_desc = self.dq_desc + self.dk_kernel_desc = self.dk_desc + self.dv_kernel_desc = self.dv_desc + self.workspace_seqlen_q = self.seqlen_q + self.workspace_batch = self.batch + else: + self.dq_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct(tuple(sample_q.shape), sample_q.dtype), + "dq_tensor", + ) + self.dk_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct(tuple(sample_k.shape), sample_k.dtype), + "dk_tensor", + ) + self.dv_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct(tuple(sample_v.shape), sample_v.dtype), + "dv_tensor", + ) + self.dq_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *sample_q.shape), sample_q.dtype), + "dq_kernel_tensor", + ) + self.dk_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *sample_k.shape), sample_k.dtype), + "dk_kernel_tensor", + ) + self.dv_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *sample_v.shape), sample_v.dtype), + "dv_kernel_tensor", + ) + # Workspace size follows the promoted (1, T, H, D) ABI, matching + # the Torch adapter's use of s_qo=T and b=1. + self.workspace_seqlen_q = self.total_q_tokens + self.workspace_batch = 1 + + def _init_fixed_shape( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_o: Any, + sample_do: Any, + sample_lse: Any, + sample_cum_seqlen_q: Any | None, + sample_cum_seqlen_k: Any | None, + ) -> None: + if sample_cum_seqlen_q is not None or sample_cum_seqlen_k is not None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must be omitted for fixed layout" + ) + self.q_desc = describe_fixed_data( + sample_q, "sample_q", layout=self.input_layout + ) + self.k_desc = describe_fixed_data( + sample_k, "sample_k", layout=self.input_layout + ) + self.v_desc = describe_fixed_data( + sample_v, "sample_v", layout=self.input_layout + ) + self.o_desc = describe_fixed_data( + sample_o, "sample_o", layout=self.input_layout + ) + self.do_desc = describe_fixed_data( + sample_do, "sample_do", layout=self.input_layout + ) + self.lse_desc = self._to_tensor_desc(sample_lse, "sample_lse") + ( + self.batch, + self.num_query_heads, + self.num_kv_heads, + self.seqlen_q, + self.seqlen_k, + self.head_dim, + ) = require_fixed_qkv(self.q_desc, self.k_desc, self.v_desc) + for desc in (self.o_desc, self.do_desc): + if ( + desc.shape != self.q_desc.shape + or desc.cudnn_dtype != self.q_desc.cudnn_dtype + ): + raise ValueError(f"{desc.name} must match sample_q's shape and dtype") + expected_lse = (self.batch, self.num_query_heads, self.seqlen_q) + if ( + self.lse_desc.shape != expected_lse + or self.lse_desc.cudnn_dtype != data_type.FLOAT + ): + raise ValueError( + f"sample_lse must have shape {expected_lse} and dtype float32" + ) + self.total_q_tokens = self.batch * self.seqlen_q + self.total_k_tokens = self.batch * self.seqlen_k + self.max_s_q = self.seqlen_q + self.max_s_k = self.seqlen_k + self.cum_q_desc = None + self.cum_k_desc = None + self.q_kernel_desc = self.q_desc + self.k_kernel_desc = self.k_desc + self.v_kernel_desc = self.v_desc + self.o_kernel_desc = self.o_desc + self.do_kernel_desc = self.do_desc + self.lse_kernel_desc = self.lse_desc + + def _init_packed_shape( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_o: Any, + sample_do: Any, + sample_lse: Any, + sample_cum_seqlen_q: Any | None, + sample_cum_seqlen_k: Any | None, + max_s_q: int | None, + max_s_k: int | None, + ) -> None: + if sample_cum_seqlen_q is None or sample_cum_seqlen_k is None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k are both required for THD layout" + ) + if max_s_q is None or max_s_k is None: + raise ValueError("max_s_q and max_s_k are both required for THD layout") + + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.k_desc = self._to_tensor_desc(sample_k, "sample_k") + self.v_desc = self._to_tensor_desc(sample_v, "sample_v") + self.o_desc = self._to_tensor_desc(sample_o, "sample_o") + self.do_desc = self._to_tensor_desc(sample_do, "sample_do") + self.lse_desc = self._to_tensor_desc(sample_lse, "sample_lse") + self.cum_q_desc = self._to_tensor_desc( + sample_cum_seqlen_q, "sample_cum_seqlen_q" + ) + self.cum_k_desc = self._to_tensor_desc( + sample_cum_seqlen_k, "sample_cum_seqlen_k" + ) + + self.total_q_tokens, self.num_query_heads, self.head_dim = self.q_desc.shape + self.total_k_tokens, self.num_kv_heads, k_head_dim = self.k_desc.shape + v_total, v_heads, value_dim = self.v_desc.shape + dimensions = ( + self.total_q_tokens, + self.total_k_tokens, + self.num_query_heads, + self.num_kv_heads, + self.head_dim, + ) + if any(value <= 0 for value in dimensions): + raise ValueError(f"SDPA dimensions must be positive, got {dimensions}") + if (v_total, v_heads, value_dim) != ( + self.total_k_tokens, + self.num_kv_heads, + self.head_dim, + ) or k_head_dim != self.head_dim: + raise ValueError( + "K and V must share packed token, head, and head-dim metadata" + ) + if self.q_desc.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError( + f"SDPA requires float16 or bfloat16 inputs, got {self.q_desc.dtype}" + ) + if ( + self.k_desc.cudnn_dtype != self.q_desc.cudnn_dtype + or self.v_desc.cudnn_dtype != self.q_desc.cudnn_dtype + ): + raise ValueError("Q, K, and V must have the same dtype") + for desc in (self.o_desc, self.do_desc): + if ( + desc.shape != self.q_desc.shape + or desc.cudnn_dtype != self.q_desc.cudnn_dtype + ): + raise ValueError(f"{desc.name} must match sample_q's shape and dtype") + if self.head_dim != 256: + raise ValueError(f"head dimension must be 256, got {self.head_dim}") + if self.num_query_heads % self.num_kv_heads: + raise ValueError( + f"H_q ({self.num_query_heads}) must be divisible by H_kv ({self.num_kv_heads})" + ) + expected_lse = (self.total_q_tokens, self.num_query_heads) + if ( + self.lse_desc.shape != expected_lse + or self.lse_desc.cudnn_dtype != data_type.FLOAT + ): + raise ValueError( + f"sample_lse must have shape {expected_lse} and dtype float32" + ) + if ( + self.cum_q_desc.ndim != 1 + or self.cum_q_desc.shape != self.cum_k_desc.shape + or self.cum_q_desc.shape[0] < 2 + ): + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must have the same shape (B + 1,) with B > 0" + ) + if ( + self.cum_q_desc.cudnn_dtype != data_type.INT32 + or self.cum_k_desc.cudnn_dtype != data_type.INT32 + ): + raise ValueError("cum_seqlen_q and cum_seqlen_k must have dtype int32") + + self.batch = self.cum_q_desc.shape[0] - 1 + self.seqlen_q = self.total_q_tokens + self.seqlen_k = self.total_k_tokens + self.max_s_q = int(max_s_q) + self.max_s_k = int(max_s_k) + if ( + self.max_s_q <= 0 + or self.max_s_k <= 0 + or self.max_s_q > self.total_q_tokens + or self.max_s_k > self.total_k_tokens + ): + raise ValueError( + "max_s_q and max_s_k must be positive and no larger than their packed token counts" + ) + + self.q_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *sample_q.shape), sample_q.dtype), "q_tensor" + ) + self.k_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *sample_k.shape), sample_k.dtype), "k_tensor" + ) + self.v_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *sample_v.shape), sample_v.dtype), "v_tensor" + ) + self.o_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *sample_o.shape), sample_o.dtype), "o_tensor" + ) + self.do_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *sample_do.shape), sample_do.dtype), "do_tensor" + ) + # The kernel addresses packed LSE as (1, H, T) with T contiguous. + self.lse_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct( + (1, self.num_query_heads, self.total_q_tokens), jnp.float32 + ), + "lse_tensor", + ) + + def check_support(self) -> bool: + self.compute_capability = self._resolve_compute_capability( + self.target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "SdpabwdSm100D256", + ) + return True + + def __call__( + self, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + o_tensor: Any, + do_tensor: Any, + lse_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + ) -> TupleDict: + self.check_support() + signature_mode = self.data_mode + for value, desc in ( + (q_tensor, self.q_desc), + (k_tensor, self.k_desc), + (v_tensor, self.v_desc), + (o_tensor, self.o_desc), + (do_tensor, self.do_desc), + (lse_tensor, self.lse_desc), + ): + mode = signature_mode if desc is not self.lse_desc else None + self._check_tensor_signature(value, desc, mode=mode) + + if self.input_layout in FIXED_LAYOUTS: + if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must be omitted for fixed layout" + ) + inputs = (q_tensor, k_tensor, v_tensor, o_tensor, do_tensor, lse_tensor) + input_spec = ( + self._to_tensor_spec(self.q_kernel_desc, mode=self.data_mode), + self._to_tensor_spec(self.k_kernel_desc, mode=self.data_mode), + self._to_tensor_spec(self.v_kernel_desc, mode=self.data_mode), + self._to_tensor_spec(self.o_kernel_desc, mode=self.data_mode), + self._to_tensor_spec(self.do_kernel_desc, mode=self.data_mode), + self._to_tensor_spec(self.lse_kernel_desc), + ) + output_spec = ( + self._to_tensor_spec(self.dq_kernel_desc, mode=self.data_mode), + self._to_tensor_spec(self.dk_kernel_desc, mode=self.data_mode), + self._to_tensor_spec(self.dv_kernel_desc, mode=self.data_mode), + ) + launch = self._launch_kernel + else: + if cum_seqlen_q_tensor is None or cum_seqlen_k_tensor is None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k are both required for THD layout" + ) + self._check_tensor_signature(cum_seqlen_q_tensor, self.cum_q_desc) + self._check_tensor_signature(cum_seqlen_k_tensor, self.cum_k_desc) + lse_storage = jnp.transpose( + jnp.reshape( + lse_tensor, + (1, self.total_q_tokens, self.num_query_heads), + ), + (0, 2, 1), + ) + inputs = ( + jnp.reshape(q_tensor, self.q_kernel_desc.shape), + jnp.reshape(k_tensor, self.k_kernel_desc.shape), + jnp.reshape(v_tensor, self.v_kernel_desc.shape), + jnp.reshape(o_tensor, self.o_kernel_desc.shape), + jnp.reshape(do_tensor, self.do_kernel_desc.shape), + lse_storage, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + input_spec = ( + self._to_tensor_spec(self.q_kernel_desc), + self._to_tensor_spec(self.k_kernel_desc), + self._to_tensor_spec(self.v_kernel_desc), + self._to_tensor_spec(self.o_kernel_desc), + self._to_tensor_spec(self.do_kernel_desc), + self._to_tensor_spec(self.lse_kernel_desc), + self._to_tensor_spec(self.cum_q_desc), + self._to_tensor_spec(self.cum_k_desc), + ) + output_spec = ( + self._to_tensor_spec(self.dq_kernel_desc), + self._to_tensor_spec(self.dk_kernel_desc), + self._to_tensor_spec(self.dv_kernel_desc), + ) + launch = self._launch_varlen_kernel + + dq, dk, dv = self._call_kernel( + inputs, + launch=launch, + output_descs=( + self.dq_kernel_desc, + self.dk_kernel_desc, + self.dv_kernel_desc, + ), + workspace_descs=(self._workspace_desc(),), + input_spec=input_spec, + output_spec=output_spec, + compile_options=compile_options_for_target(self.compute_capability), + ) + if self.input_layout == _PACKED_LAYOUT: + dq = jnp.reshape(dq, self.dq_desc.shape) + dk = jnp.reshape(dk, self.dk_desc.shape) + dv = jnp.reshape(dv, self.dv_desc.shape) + return TupleDict(dq_tensor=dq, dk_tensor=dk, dv_tensor=dv) + + def _workspace_desc(self) -> JaxTensorDesc: + from cutlass import Float32 + + from .fmha_backward_sm100_2kernel import ( + BlackwellFusedMultiHeadAttentionBackward, + ) + + shape = BlackwellFusedMultiHeadAttentionBackward.get_workspace_size( + self.workspace_seqlen_q, + self.head_dim, + self.num_query_heads, + self.workspace_batch, + Float32, + ) + return self.q_kernel_desc.compact_like( + cudnn_dtype=data_type.UINT8, + shape=shape, + name="workspace", + init_value=0, + ) + + def _launch_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + output: Any, + doutput: Any, + lse: Any, + dq: Any, + dk: Any, + dv: Any, + workspace: Any, + ) -> None: + self._invoke_kernel( + stream, + q, + k, + v, + output, + doutput, + lse, + dq, + dk, + dv, + workspace, + None, + None, + ) + + def _launch_varlen_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + output: Any, + doutput: Any, + lse: Any, + cum_seqlen_q: Any, + cum_seqlen_k: Any, + dq: Any, + dk: Any, + dv: Any, + workspace: Any, + ) -> None: + self._invoke_kernel( + stream, + q, + k, + v, + output, + doutput, + lse, + dq, + dk, + dv, + workspace, + cum_seqlen_q, + cum_seqlen_k, + ) + + def _invoke_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + output: Any, + doutput: Any, + lse: Any, + dq: Any, + dk: Any, + dv: Any, + workspace: Any, + cum_seqlen_q: Any | None, + cum_seqlen_k: Any | None, + ) -> None: + from cutlass import Float32, Int32 + from cutlass.jax import jax_to_cutlass_dtype + + from ..fmha_utils import MaskEnum + from .fmha_backward_sm100_2kernel import ( + BlackwellFusedMultiHeadAttentionBackward, + ) + + mask_type = { + "residual": MaskEnum.RESIDUAL_MASK, + "window": MaskEnum.WINDOW_MASK_INFERENCE, + }[self.mask_kind] + kernel = BlackwellFusedMultiHeadAttentionBackward( + element_dtype=jax_to_cutlass_dtype(self.q_desc.dtype), + acc_dtype=Float32, + mma_tiler=(128, 128, 256), + dkdv_mma_tiler=(128, 64, 256), + varlen=self.input_layout == _PACKED_LAYOUT, + is_causal=self.is_causal, + mask_type=mask_type, + window_size_left=self.window_size_left, + window_size_right=self.window_size_right, + ) + problem_shape = ( + Int32(self.max_s_q), + Int32(self.max_s_k), + Int32(256), + ( + ( + Int32(self.num_query_heads // self.num_kv_heads), + Int32(self.num_kv_heads), + ), + Int32(self.batch), + ), + ) + kernel( + problem_shape, + q, + k, + v, + output, + dq, + dk, + dv, + doutput, + lse, + cum_seqlen_q, + cum_seqlen_k, + Float32(self.scale_softmax), + workspace, + stream, + ) + + +@partial( + jax.jit, + static_argnames=( + "max_s_q", + "max_s_k", + "acc_dtype", + "mma_tiler_mn", + "dkdv_mma_tiler_mn", + "is_causal", + "window_size", + "scale_softmax", + "layout", + "target_compute_capability", + ), +) +def sdpa_bwd_wrapper_sm100_d256( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + o_tensor: Any, + do_tensor: Any, + lse_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + dkdv_mma_tiler_mn: tuple[int, int] = (128, 64), + is_causal: bool = False, + window_size: tuple[int, int] = (-1, -1), + scale_softmax: float | None = None, + layout: str | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute fixed BHSD/BSHD or packed THD SDPA gradients. + + Packed calls must provide both cumulative sequence-length arrays and static + ``max_s_q``/``max_s_k`` bounds. Fixed gradients follow ``layout``; LSE + remains ``(B, H, S)``. + """ + + values = (q_tensor, k_tensor, v_tensor, o_tensor, do_tensor, lse_tensor) + samples = tuple(jax.ShapeDtypeStruct(value.shape, value.dtype) for value in values) + sample_cum_seqlen_q = ( + None + if cum_seqlen_q_tensor is None + else jax.ShapeDtypeStruct(cum_seqlen_q_tensor.shape, cum_seqlen_q_tensor.dtype) + ) + sample_cum_seqlen_k = ( + None + if cum_seqlen_k_tensor is None + else jax.ShapeDtypeStruct(cum_seqlen_k_tensor.shape, cum_seqlen_k_tensor.dtype) + ) + return SdpabwdSm100D256( + *samples, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + max_s_q=max_s_q, + max_s_k=max_s_k, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + dkdv_mma_tiler_mn=dkdv_mma_tiler_mn, + is_causal=is_causal, + window_size=window_size, + scale_softmax=scale_softmax, + layout=layout, + target_compute_capability=target_compute_capability, + )( + *values, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + + +__all__ = [ + "SUPPORTED_COMPUTE_CAPABILITIES", + "SdpabwdSm100D256", + "sdpa_bwd_wrapper_sm100_d256", +] diff --git a/python/cudnn/sdpa/fwd/__init__.py b/python/cudnn/sdpa/fwd/__init__.py index bf95abac2..44f91d48a 100644 --- a/python/cudnn/sdpa/fwd/__init__.py +++ b/python/cudnn/sdpa/fwd/__init__.py @@ -1,9 +1,12 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import SdpafwdSm100D256, sdpa_fwd_wrapper_sm100_d256 +"""Lazy Torch SDPA-forward API exports.""" -__all__ = [ - "SdpafwdSm100D256", - "sdpa_fwd_wrapper_sm100_d256", -] +from ..._operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("SdpafwdSm100D256", "sdpa_fwd_wrapper_sm100_d256")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/sdpa/fwd/jax.py b/python/cudnn/sdpa/fwd/jax.py new file mode 100644 index 000000000..254497d1f --- /dev/null +++ b/python/cudnn/sdpa/fwd/jax.py @@ -0,0 +1,548 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for fixed BHSD/BSHD and packed THD SM100 d=256 SDPA forward.""" + +from __future__ import annotations + +from functools import partial +import math +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._cute_compiler import compile_options_for_target +from ..._jax import JaxApiBase, TupleDict +from ..jax_utils import ( + FIXED_LAYOUTS, + describe_fixed_data, + fixed_data_mode, + make_fixed_output, + normalize_sdpa_layout, + require_fixed_qkv, + require_float32_dtype, + resolve_sdpa_config, +) + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) +_PACKED_LAYOUT = "THD" + + +class SdpafwdSm100D256(JaxApiBase): + """JAX callable specialized from fixed BHSD/BSHD or packed THD metadata. + + Packed THD inputs require cumulative sequence lengths and explicit static + ``max_s_q``/``max_s_k`` bounds. JAX samples expose array metadata rather + than cumulative-length values, so these bounds cannot be inferred while + tracing. + """ + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_cum_seqlen_q: Any | None = None, + sample_cum_seqlen_k: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + qk_acc_dtype: Any = None, + pv_acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + is_causal: bool = False, + window_size: tuple[int, int] = (-1, -1), + scale_softmax: float | None = None, + scale_output: float = 1.0, + layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + ranks = tuple( + len(tuple(sample.shape)) for sample in (sample_q, sample_k, sample_v) + ) + if len(set(ranks)) != 1: + raise ValueError(f"Q, K, and V must use the same rank, got {ranks}") + self.input_layout = normalize_sdpa_layout(layout, ranks[0]) + if self.input_layout in FIXED_LAYOUTS: + if max_s_q is not None or max_s_k is not None: + raise ValueError("max_s_q and max_s_k are only valid for THD layout") + self.data_mode = fixed_data_mode(self.input_layout) + self._init_fixed_shape( + sample_q, + sample_k, + sample_v, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + ) + elif self.input_layout == _PACKED_LAYOUT: + self.data_mode = None + self._init_packed_shape( + sample_q, + sample_k, + sample_v, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + max_s_q, + max_s_k, + ) + + require_float32_dtype(qk_acc_dtype, "qk_acc_dtype") + require_float32_dtype(pv_acc_dtype, "pv_acc_dtype") + if tuple(mma_tiler_mn) != (128, 128): + raise ValueError(f"mma_tiler_mn must be (128, 128), got {mma_tiler_mn}") + self.is_causal = bool(is_causal) + ( + self.scale_softmax, + self.window_size_left, + self.window_size_right, + self.mask_kind, + ) = resolve_sdpa_config( + seqlen_q=self.max_s_q, + seqlen_k=self.max_s_k, + # The Torch API always selects the residual mask for non-causal + # varlen inputs, independent of the packed token count. + tile_extent=self.seqlen_k if self.input_layout in FIXED_LAYOUTS else 1, + is_causal=self.is_causal, + window_size=tuple(window_size), + scale_softmax=scale_softmax, + ) + self.scale_output = float(scale_output) + self.target_compute_capability = target_compute_capability + self.compute_capability: int | None = None + + if self.input_layout in FIXED_LAYOUTS: + self.o_desc = make_fixed_output( + tuple(sample_q.shape), + sample_q.dtype, + "o_tensor", + layout=self.input_layout, + ) + self.o_kernel_desc = self.o_desc + self.lse_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct( + (self.batch, self.num_query_heads, self.seqlen_q), + jnp.float32, + ), + "lse_tensor", + ) + self.lse_kernel_desc = self.lse_desc + else: + self.o_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct(tuple(sample_q.shape), sample_q.dtype), + "o_tensor", + ) + self.o_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *sample_q.shape), sample_q.dtype), + "o_kernel_tensor", + ) + self.lse_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct( + (self.total_q_tokens, self.num_query_heads), jnp.float32 + ), + "lse_tensor", + # The kernel writes sequence-contiguous LSE. TensorSpec asks + # XLA for that physical order while preserving public (T, H). + public_stride_order=(0, 1), + ) + # Forward consumes packed LSE through its iterator, just as the + # Torch wrapper passes the rank-2 (T, H) buffer directly. + self.lse_kernel_desc = self.lse_desc + + def _init_fixed_shape( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_cum_seqlen_q: Any | None, + sample_cum_seqlen_k: Any | None, + ) -> None: + if sample_cum_seqlen_q is not None or sample_cum_seqlen_k is not None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must be omitted for fixed layout" + ) + self.q_desc = describe_fixed_data( + sample_q, "sample_q", layout=self.input_layout + ) + self.k_desc = describe_fixed_data( + sample_k, "sample_k", layout=self.input_layout + ) + self.v_desc = describe_fixed_data( + sample_v, "sample_v", layout=self.input_layout + ) + ( + self.batch, + self.num_query_heads, + self.num_kv_heads, + self.seqlen_q, + self.seqlen_k, + self.head_dim, + ) = require_fixed_qkv(self.q_desc, self.k_desc, self.v_desc) + self.total_q_tokens = self.batch * self.seqlen_q + self.total_k_tokens = self.batch * self.seqlen_k + self.max_s_q = self.seqlen_q + self.max_s_k = self.seqlen_k + self.cum_q_desc = None + self.cum_k_desc = None + self.q_kernel_desc = self.q_desc + self.k_kernel_desc = self.k_desc + self.v_kernel_desc = self.v_desc + + def _init_packed_shape( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_cum_seqlen_q: Any | None, + sample_cum_seqlen_k: Any | None, + max_s_q: int | None, + max_s_k: int | None, + ) -> None: + if sample_cum_seqlen_q is None or sample_cum_seqlen_k is None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k are both required for THD layout" + ) + if max_s_q is None or max_s_k is None: + raise ValueError("max_s_q and max_s_k are both required for THD layout") + + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.k_desc = self._to_tensor_desc(sample_k, "sample_k") + self.v_desc = self._to_tensor_desc(sample_v, "sample_v") + self.cum_q_desc = self._to_tensor_desc( + sample_cum_seqlen_q, "sample_cum_seqlen_q" + ) + self.cum_k_desc = self._to_tensor_desc( + sample_cum_seqlen_k, "sample_cum_seqlen_k" + ) + + self.total_q_tokens, self.num_query_heads, self.head_dim = self.q_desc.shape + self.total_k_tokens, self.num_kv_heads, k_head_dim = self.k_desc.shape + v_total, v_heads, value_dim = self.v_desc.shape + dimensions = ( + self.total_q_tokens, + self.total_k_tokens, + self.num_query_heads, + self.num_kv_heads, + self.head_dim, + ) + if any(value <= 0 for value in dimensions): + raise ValueError(f"SDPA dimensions must be positive, got {dimensions}") + if (v_total, v_heads, value_dim) != ( + self.total_k_tokens, + self.num_kv_heads, + self.head_dim, + ) or k_head_dim != self.head_dim: + raise ValueError( + "K and V must share packed token, head, and head-dim metadata" + ) + if self.q_desc.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError( + f"SDPA requires float16 or bfloat16 inputs, got {self.q_desc.dtype}" + ) + if ( + self.k_desc.cudnn_dtype != self.q_desc.cudnn_dtype + or self.v_desc.cudnn_dtype != self.q_desc.cudnn_dtype + ): + raise ValueError("Q, K, and V must have the same dtype") + if self.head_dim != 256: + raise ValueError(f"head dimension must be 256, got {self.head_dim}") + if self.num_query_heads % self.num_kv_heads: + raise ValueError( + f"H_q ({self.num_query_heads}) must be divisible by H_kv ({self.num_kv_heads})" + ) + if ( + self.cum_q_desc.ndim != 1 + or self.cum_q_desc.shape != self.cum_k_desc.shape + or self.cum_q_desc.shape[0] < 2 + ): + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must have the same shape (B + 1,) with B > 0" + ) + if ( + self.cum_q_desc.cudnn_dtype != data_type.INT32 + or self.cum_k_desc.cudnn_dtype != data_type.INT32 + ): + raise ValueError("cum_seqlen_q and cum_seqlen_k must have dtype int32") + + self.batch = self.cum_q_desc.shape[0] - 1 + self.seqlen_q = self.total_q_tokens + self.seqlen_k = self.total_k_tokens + self.max_s_q = int(max_s_q) + self.max_s_k = int(max_s_k) + if ( + self.max_s_q <= 0 + or self.max_s_k <= 0 + or self.max_s_q > self.total_q_tokens + or self.max_s_k > self.total_k_tokens + ): + raise ValueError( + "max_s_q and max_s_k must be positive and no larger than their packed token counts" + ) + + # Public THD remains row-major; adding the singleton batch dimension + # yields the exact (1, T, H, D) ABI used by the Torch wrapper. + self.q_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *sample_q.shape), sample_q.dtype), "q_tensor" + ) + self.k_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *sample_k.shape), sample_k.dtype), "k_tensor" + ) + self.v_kernel_desc = self._to_tensor_desc( + jax.ShapeDtypeStruct((1, *sample_v.shape), sample_v.dtype), "v_tensor" + ) + + def check_support(self) -> bool: + self.compute_capability = self._resolve_compute_capability( + self.target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "SdpafwdSm100D256", + ) + return True + + def __call__( + self, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + ) -> TupleDict: + self.check_support() + signature_mode = self.data_mode + for value, desc in ( + (q_tensor, self.q_desc), + (k_tensor, self.k_desc), + (v_tensor, self.v_desc), + ): + self._check_tensor_signature(value, desc, mode=signature_mode) + + if self.input_layout in FIXED_LAYOUTS: + if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must be omitted for fixed layout" + ) + inputs = (q_tensor, k_tensor, v_tensor) + input_spec = ( + self._to_tensor_spec(self.q_kernel_desc, mode=self.data_mode), + self._to_tensor_spec(self.k_kernel_desc, mode=self.data_mode), + self._to_tensor_spec(self.v_kernel_desc, mode=self.data_mode), + ) + output_spec = ( + self._to_tensor_spec(self.o_kernel_desc, mode=self.data_mode), + self._to_tensor_spec(self.lse_kernel_desc), + ) + launch = self._launch_kernel + else: + if cum_seqlen_q_tensor is None or cum_seqlen_k_tensor is None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k are both required for THD layout" + ) + self._check_tensor_signature(cum_seqlen_q_tensor, self.cum_q_desc) + self._check_tensor_signature(cum_seqlen_k_tensor, self.cum_k_desc) + inputs = ( + jnp.reshape(q_tensor, self.q_kernel_desc.shape), + jnp.reshape(k_tensor, self.k_kernel_desc.shape), + jnp.reshape(v_tensor, self.v_kernel_desc.shape), + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + input_spec = ( + self._to_tensor_spec(self.q_kernel_desc), + self._to_tensor_spec(self.k_kernel_desc), + self._to_tensor_spec(self.v_kernel_desc), + self._to_tensor_spec(self.cum_q_desc), + self._to_tensor_spec(self.cum_k_desc), + ) + output_spec = ( + self._to_tensor_spec(self.o_kernel_desc), + self._to_tensor_spec(self.lse_kernel_desc), + ) + launch = self._launch_varlen_kernel + + output, lse = self._call_kernel( + inputs, + launch=launch, + output_descs=(self.o_kernel_desc, self.lse_kernel_desc), + input_spec=input_spec, + output_spec=output_spec, + compile_options=compile_options_for_target(self.compute_capability), + ) + if self.input_layout == _PACKED_LAYOUT: + output = jnp.reshape(output, self.o_desc.shape) + return TupleDict(o_tensor=output, lse_tensor=lse) + + def _launch_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + output: Any, + lse: Any, + ) -> None: + self._invoke_kernel(stream, q, k, v, output, lse, None, None) + + def _launch_varlen_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + cum_seqlen_q: Any, + cum_seqlen_k: Any, + output: Any, + lse: Any, + ) -> None: + self._invoke_kernel( + stream, + q, + k, + v, + output, + lse, + cum_seqlen_q, + cum_seqlen_k, + ) + + def _invoke_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + output: Any, + lse: Any, + cum_seqlen_q: Any | None, + cum_seqlen_k: Any | None, + ) -> None: + from cutlass import Float32, Int32 + + from ..fmha_utils import MaskEnum + from .fmha_forward_sm100_d256 import BlackwellFusedMultiHeadAttentionForward + + mask_type = { + "residual": MaskEnum.RESIDUAL_MASK, + "window": MaskEnum.WINDOW_MASK_INFERENCE, + }[self.mask_kind] + kernel = BlackwellFusedMultiHeadAttentionForward( + qk_acc_dtype=Float32, + pv_acc_dtype=Float32, + mma_tiler=(128, 128, 256), + is_persistent=False, + mask_type=mask_type, + ) + problem_size = tuple( + Int32(value) + for value in ( + self.batch, + self.max_s_q, + self.max_s_k, + self.num_query_heads, + self.num_kv_heads, + 256, + ) + ) + left = None if self.window_size_left < 0 else Int32(self.window_size_left) + right = None if self.window_size_right < 0 else Int32(self.window_size_right) + kernel( + q, + k, + v, + output, + problem_size, + cum_seqlen_q, + cum_seqlen_k, + lse, + Float32(self.scale_softmax * math.log2(math.e)), + Float32(self.scale_softmax), + Float32(self.scale_output), + left, + right, + stream, + ) + + +@partial( + jax.jit, + static_argnames=( + "max_s_q", + "max_s_k", + "qk_acc_dtype", + "pv_acc_dtype", + "mma_tiler_mn", + "is_causal", + "window_size", + "scale_softmax", + "scale_output", + "layout", + "target_compute_capability", + ), +) +def sdpa_fwd_wrapper_sm100_d256( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + qk_acc_dtype: Any = None, + pv_acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + is_causal: bool = False, + window_size: tuple[int, int] = (-1, -1), + scale_softmax: float | None = None, + scale_output: float = 1.0, + layout: str | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute fixed BHSD/BSHD or packed THD SDPA forward. + + Packed calls must provide both cumulative sequence-length arrays and static + ``max_s_q``/``max_s_k`` bounds. Fixed outputs follow ``layout``; LSE remains + ``(B, H, S)``. + """ + + samples = tuple( + jax.ShapeDtypeStruct(value.shape, value.dtype) + for value in (q_tensor, k_tensor, v_tensor) + ) + sample_cum_seqlen_q = ( + None + if cum_seqlen_q_tensor is None + else jax.ShapeDtypeStruct(cum_seqlen_q_tensor.shape, cum_seqlen_q_tensor.dtype) + ) + sample_cum_seqlen_k = ( + None + if cum_seqlen_k_tensor is None + else jax.ShapeDtypeStruct(cum_seqlen_k_tensor.shape, cum_seqlen_k_tensor.dtype) + ) + return SdpafwdSm100D256( + *samples, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + max_s_q=max_s_q, + max_s_k=max_s_k, + qk_acc_dtype=qk_acc_dtype, + pv_acc_dtype=pv_acc_dtype, + mma_tiler_mn=mma_tiler_mn, + is_causal=is_causal, + window_size=window_size, + scale_softmax=scale_softmax, + scale_output=scale_output, + layout=layout, + target_compute_capability=target_compute_capability, + )( + q_tensor, + k_tensor, + v_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + + +__all__ = [ + "SUPPORTED_COMPUTE_CAPABILITIES", + "SdpafwdSm100D256", + "sdpa_fwd_wrapper_sm100_d256", +] diff --git a/python/cudnn/sdpa/jax_utils.py b/python/cudnn/sdpa/jax_utils.py new file mode 100644 index 000000000..c54fa7fd5 --- /dev/null +++ b/python/cudnn/sdpa/jax_utils.py @@ -0,0 +1,187 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared fixed-shape metadata and configuration for JAX SDPA.""" + +from __future__ import annotations + +import math +from typing import Any + +import jax +import jax.numpy as jnp + +from .. import data_type +from .._jax import JaxApiBase, JaxTensorDesc +from .._jax.datatypes import normalize_jax_dtype +from .._jax.layout import mode_from_layout, stride_order_to_public, to_public_axes + +FIXED_LAYOUTS = ("BHSD", "BSHD") +KERNEL_AXES = "BSHD" +KERNEL_STRIDE_ORDER = (3, 2, 1, 0) + + +def normalize_sdpa_layout(layout: str | None, rank: int) -> str: + """Normalize the public fixed/packed axis order for an SDPA call.""" + + if layout is None: + if rank == 4: + return "BHSD" + if rank == 3: + return "THD" + raise ValueError(f"SDPA layout cannot be inferred from rank {rank}") + if not isinstance(layout, str): + raise TypeError(f"layout must be a string or None, got {layout!r}") + normalized = "".join( + character for character in layout.upper() if character.isalpha() + ) + if normalized not in (*FIXED_LAYOUTS, "THD"): + raise ValueError(f"layout must be BHSD, BSHD, or THD, got {layout!r}") + expected_rank = 3 if normalized == "THD" else 4 + if rank != expected_rank: + raise ValueError( + f"{normalized} layout requires rank-{expected_rank} tensors, got rank {rank}" + ) + return normalized + + +def fixed_data_mode(layout: str) -> tuple[int, ...]: + if layout not in FIXED_LAYOUTS: + raise ValueError(f"fixed layout must be one of {FIXED_LAYOUTS}, got {layout!r}") + return mode_from_layout(layout, kernel_axes=KERNEL_AXES) + + +def describe_fixed_data( + sample: Any, + name: str, + *, + layout: str, + init_value: bool | int | float | None = None, +) -> JaxTensorDesc: + """Describe fixed data with the BSHD layout hard-coded by the kernel.""" + + mode = fixed_data_mode(layout) + + return JaxApiBase._to_tensor_desc( + sample, + name, + mode=mode, + public_stride_order=stride_order_to_public(KERNEL_STRIDE_ORDER, mode), + init_value=init_value, + ) + + +def make_fixed_output( + public_shape: tuple[int, ...], + dtype: Any, + name: str, + *, + layout: str, + init_value: bool | int | float | None = None, +) -> JaxTensorDesc: + return describe_fixed_data( + jax.ShapeDtypeStruct(public_shape, dtype), + name, + layout=layout, + init_value=init_value, + ) + + +def require_fixed_qkv( + q_desc: JaxTensorDesc, + k_desc: JaxTensorDesc, + v_desc: JaxTensorDesc, +) -> tuple[int, int, int, int, int, int]: + for desc in (q_desc, k_desc, v_desc): + if desc.ndim != 4: + raise ValueError( + f"{desc.name} must have rank 4 (B, H, S, D), got {desc.shape}" + ) + if q_desc.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError( + f"SDPA requires float16 or bfloat16 inputs, got {q_desc.dtype}" + ) + if ( + k_desc.cudnn_dtype != q_desc.cudnn_dtype + or v_desc.cudnn_dtype != q_desc.cudnn_dtype + ): + raise ValueError("Q, K, and V must have the same dtype") + + logical_bhsd_mode = mode_from_layout("BHSD", kernel_axes=KERNEL_AXES) + q_shape = to_public_axes(q_desc.shape, logical_bhsd_mode) + k_shape = to_public_axes(k_desc.shape, logical_bhsd_mode) + v_shape = to_public_axes(v_desc.shape, logical_bhsd_mode) + batch, num_query_heads, seqlen_q, head_dim = q_shape + k_batch, num_kv_heads, seqlen_k, k_head_dim = k_shape + v_batch, num_value_heads, v_seqlen, value_dim = v_shape + dimensions = (batch, num_query_heads, num_kv_heads, seqlen_q, seqlen_k, head_dim) + if any(value <= 0 for value in dimensions): + raise ValueError(f"SDPA dimensions must be positive, got {dimensions}") + if (k_batch, v_batch) != (batch, batch): + raise ValueError("Q, K, and V batch dimensions must match") + if (num_value_heads, v_seqlen) != (num_kv_heads, seqlen_k): + raise ValueError("K and V head and sequence dimensions must match") + if (k_head_dim, value_dim) != (head_dim, head_dim): + raise ValueError("Q, K, and V head dimensions must match") + if head_dim != 256: + raise ValueError(f"head dimension must be 256, got {head_dim}") + if num_query_heads % num_kv_heads: + raise ValueError( + f"H_q ({num_query_heads}) must be divisible by H_kv ({num_kv_heads})" + ) + return dimensions + + +def require_float32_dtype(value: Any | None, name: str) -> Any: + dtype = normalize_jax_dtype(value, jnp.float32, name) + if dtype != jnp.dtype(jnp.float32): + raise ValueError(f"{name} must be float32, got {dtype}") + return dtype + + +def resolve_sdpa_config( + *, + seqlen_q: int, + seqlen_k: int, + tile_extent: int, + is_causal: bool, + window_size: tuple[int, int], + scale_softmax: float | None, +) -> tuple[float, int, int, str]: + if len(window_size) != 2: + raise ValueError(f"window_size must contain two values, got {window_size}") + window_size_left, window_size_right = window_size + if is_causal: + window_size_right = 0 + elif (window_size_left, window_size_right) != (-1, -1): + raise NotImplementedError( + f"window_size must be (-1, -1) for non-causal mode, got {window_size}" + ) + if window_size_left >= seqlen_k - 1: + raise ValueError( + f"window_size_left must be less than S_kv - 1 ({seqlen_k - 1}), got {window_size_left}" + ) + if window_size_right >= seqlen_q - 1: + raise ValueError( + f"window_size_right must be less than S_q - 1 ({seqlen_q - 1}), got {window_size_right}" + ) + + scale = ( + 1.0 / math.sqrt(256) + if scale_softmax is None or scale_softmax == 0.0 + else float(scale_softmax) + ) + mask_kind = "residual" if not is_causal and tile_extent % 128 else "window" + return scale, window_size_left, window_size_right, mask_kind + + +__all__ = [ + "FIXED_LAYOUTS", + "describe_fixed_data", + "fixed_data_mode", + "make_fixed_output", + "normalize_sdpa_layout", + "require_fixed_qkv", + "require_float32_dtype", + "resolve_sdpa_config", +] diff --git a/test/python/fe_api/test_dsa_score_recompute_op.py b/test/python/fe_api/test_dsa_score_recompute_op.py index 0a3661275..c6d02f919 100644 --- a/test/python/fe_api/test_dsa_score_recompute_op.py +++ b/test/python/fe_api/test_dsa_score_recompute_op.py @@ -78,15 +78,21 @@ def desc(self, shape, dtype=_DataType.BFLOAT16, *, stride_order=None, name=""): name=name, ) - def sparse(self, *, target=100, score_type="indexer", with_length=False, **overrides): + def sparse( + self, *, target=100, score_type="indexer", with_length=False, **overrides + ): aux_dtype = _DataType.BFLOAT16 if score_type == "indexer" else _DataType.FLOAT arguments = { "q": self.desc((2, 4, 32, 128), name="q"), "k": self.desc((2, 256, 128), name="k"), "per_head": self.desc((2, 4, 32), aux_dtype, name="per_head"), - "topk_indices": self.desc((2, 4, 128), _DataType.INT32, name="topk_indices"), + "topk_indices": self.desc( + (2, 4, 128), _DataType.INT32, name="topk_indices" + ), "output": self.desc((2, 4, 128), _DataType.FLOAT, name="output"), - "topk_length": self.desc((2, 4), _DataType.INT32, name="topk_length") if with_length else None, + "topk_length": self.desc((2, 4), _DataType.INT32, name="topk_length") + if with_length + else None, "score_type": score_type, "softmax_scale": 0.125 if score_type == "attention" else 1.0, "target_compute_capability": target, @@ -112,7 +118,9 @@ def dense(self, *, target=100, score_type="indexer", thd=False, **overrides): "k": self.desc(k_shape, name="k"), "per_head": self.desc(aux_shape, aux_dtype, name="per_head"), "output": self.desc(output_shape, _DataType.FLOAT, name="output"), - "denominator": self.desc(denominator_shape, _DataType.FLOAT, name="denominator"), + "denominator": self.desc( + denominator_shape, _DataType.FLOAT, name="denominator" + ), "score_type": score_type, "scale": 0.125 if score_type == "attention" else 1.0, "ratio": 2, @@ -131,14 +139,23 @@ def test_sparse_resolves_sm100_and_preserves_explicit_lengths(self): self.assertTrue(operation.check_support()) self.assertEqual(operation.qhead_per_kv_head, 32) self.assertIsInstance(operation.config, self.config.SparseScoreKernelConfig) - self.assertEqual((operation.config.m_block_size, operation.config.n_block_size), (32, 128)) + self.assertEqual( + (operation.config.m_block_size, operation.config.n_block_size), (32, 128) + ) self.assertTrue(operation.config.have_topk_length) def test_sparse_resolves_sm90(self): operation = self.sparse(target=90) self.assertTrue(operation.check_support()) self.assertIsInstance(operation.config, self.op.SparseScoreSm90Config) - self.assertEqual((operation.config.tile_m, operation.config.tile_n, operation.config.num_threads), (32, 64, 256)) + self.assertEqual( + ( + operation.config.tile_m, + operation.config.tile_n, + operation.config.num_threads, + ), + (32, 64, 256), + ) def test_dense_bshd_and_thd_resolve_sm100(self): bshd = self.dense() @@ -151,7 +168,9 @@ def test_dense_bshd_and_thd_resolve_sm100(self): self.assertEqual((thd.max_seqlen_q, thd.max_seqlen_k), (4, 8)) def test_dense_sm90_thd_is_explicitly_unsupported(self): - with self.assertRaisesRegex(NotImplementedError, "host-side sequence-length reads"): + with self.assertRaisesRegex( + NotImplementedError, "host-side sequence-length reads" + ): self.dense(target=90, thd=True).check_support() def test_complete_signatures_are_cross_checked(self): @@ -164,18 +183,22 @@ def test_complete_signatures_are_cross_checked(self): with self.assertRaisesRegex(ValueError, "ratio must be >= 1"): self.dense(ratio=0).check_support() - def test_native_abi_descriptors_must_be_compact(self): - noncompact = self.tensor.TensorDesc( + def test_native_abi_feature_axes_must_be_contiguous(self): + unsupported_layout = self.tensor.TensorDesc( dtype=_DataType.BFLOAT16, shape=(2, 4, 32, 128), stride=(16384, 128, 512, 1), stride_order=(3, 1, 2, 0), ) - with self.assertRaisesRegex(ValueError, "Q must be row-major compact"): - self.sparse(q=noncompact).check_support() + with self.assertRaisesRegex( + ValueError, "Q must have its final 2 axes contiguous" + ): + self.sparse(q=unsupported_layout).check_support() def test_sm100_configuration_rejects_incompatible_topk(self): - with self.assertRaisesRegex(ValueError, "multiple of the selected n_block_size"): + with self.assertRaisesRegex( + ValueError, "multiple of the selected n_block_size" + ): self.sparse( topk_indices=self.desc((2, 4, 96), _DataType.INT32), output=self.desc((2, 4, 96), _DataType.FLOAT), diff --git a/test/python/fe_api/test_jax_api_base.py b/test/python/fe_api/test_jax_api_base.py index e3aaef2eb..dbb7bf26e 100644 --- a/test/python/fe_api/test_jax_api_base.py +++ b/test/python/fe_api/test_jax_api_base.py @@ -541,6 +541,7 @@ def invoke(*args): cutlass.jax = cutlass_jax input_value = _ArrayMetadata((2,), "float32", label="input") + output_seed = _ArrayMetadata((2,), "float32", label="seed") input_tensor_spec = _TensorSpec(layout=(0,), mode=(0,)) initialized_output_spec = _TensorSpec(layout=(2, 1, 0), mode=(2, 0, 1)) api = Adapter() @@ -557,10 +558,38 @@ def invoke(*args): (input_value,), launch=launch, output_descs=output_descs, + output_seeds=(output_seed, None), workspace_descs=workspace_descs, input_spec=(input_tensor_spec,), output_spec=(None, initialized_output_spec), ) + with self.assertRaisesRegex(ValueError, "Expected 2 output seeds, got 1"): + api._call_kernel( + (input_value,), + launch=launch, + output_descs=output_descs, + output_seeds=(output_seed,), + ) + with self.assertRaisesRegex(ValueError, "output seed shape mismatch"): + api._call_kernel( + (input_value,), + launch=launch, + output_descs=output_descs, + output_seeds=( + _ArrayMetadata((3,), "float32", label="wrong shape"), + None, + ), + ) + with self.assertRaisesRegex(ValueError, "output seed dtype mismatch"): + api._call_kernel( + (input_value,), + launch=launch, + output_descs=output_descs, + output_seeds=( + _ArrayMetadata((2,), "bfloat16", label="wrong dtype"), + None, + ), + ) with self.assertRaisesRegex(TypeError, "input #0 must have shape and dtype metadata"): api._call_kernel( ([input_value],), @@ -590,19 +619,26 @@ def invoke(*args): ((3,), 0, "float32"), ], ) - self.assertEqual(seen["call_options"]["input_output_aliases"], {1: 1, 2: 2}) - self.assertEqual(len(seen["call_options"]["input_spec"]), 3) - self.assertIs(seen["call_options"]["input_spec"][1], initialized_output_spec) - self.assertIs(seen["call_options"]["input_spec"][2], seen["call_options"]["output_spec"][2]) + self.assertEqual( + seen["call_options"]["input_output_aliases"], + {1: 0, 2: 1, 3: 2}, + ) + self.assertEqual(len(seen["call_options"]["input_spec"]), 4) + self.assertIs( + seen["call_options"]["input_spec"][1], + seen["call_options"]["output_spec"][0], + ) + self.assertIs(seen["call_options"]["input_spec"][2], initialized_output_spec) + self.assertIs(seen["call_options"]["input_spec"][3], seen["call_options"]["output_spec"][2]) stream, launch_input, output, amax, counter, scratch = seen["launch_args"] self.assertEqual(stream, "stream") self.assertIs(launch_input, input_value) - self.assertEqual(output.label, "allocated(0)") + self.assertIs(output, output_seed) self.assertIs(amax, full_calls[0][3]) self.assertIs(counter, full_calls[1][3]) self.assertEqual(scratch.label, "allocated(3)") - self.assertEqual([result.label for result in results], ["allocated(0)", "full(-inf)"]) + self.assertEqual([result.label for result in results], ["seed", "full(-inf)"]) self.assertNotIn(full_calls[0][3], vars(api).values()) self.assertNotIn(full_calls[1][3], vars(api).values()) diff --git a/test/python/fe_api/test_jax_bsa_contract.py b/test/python/fe_api/test_jax_bsa_contract.py new file mode 100644 index 000000000..1a2e77345 --- /dev/null +++ b/test/python/fe_api/test_jax_bsa_contract.py @@ -0,0 +1,389 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for the JAX block-sparse attention adapters.""" + +from __future__ import annotations + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_BSA_ROOT = _CUDNN_ROOT / "block_sparse_attention" +_PACKAGE = "cudnn_jax_bsa_contract_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + + +_DTYPE_TO_CUDNN = { + "float16": _DataType.HALF, + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, + "int32": _DataType.INT32, +} +_CUDNN_TO_DTYPE = {value: key for key, value in _DTYPE_TO_CUDNN.items()} + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + +class _TensorSpec: + def __init__(self, *, layout, mode, divisibility=None): + self.layout = layout + self.mode = mode + self.divisibility = divisibility + + +class JaxBsaContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + bsa_name = f"{_PACKAGE}.block_sparse_attention" + bsa = types.ModuleType(bsa_name) + bsa.__path__ = [str(_BSA_ROOT)] + bsa.__package__ = bsa_name + bsa.__spec__ = ModuleSpec(bsa_name, loader=None, is_package=True) + sys.modules[bsa_name] = bsa + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + layout_module = importlib.import_module(f"{internal_name}.layout") + result_module = importlib.import_module(f"{_PACKAGE}._result") + + class JaxTensorDesc(tensor_module.TensorDesc): + @property + def cudnn_dtype(self): + return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) + + def compact_like( + self, *, cudnn_dtype, shape, stride_order=None, name="", init_value=None + ): + if stride_order is None: + stride_order = tuple(reversed(range(len(shape)))) + stride = layout_module.compact_stride(tuple(shape), tuple(stride_order)) + return JaxTensorDesc( + dtype=_CUDNN_TO_DTYPE[cudnn_dtype], + shape=tuple(shape), + stride=stride, + stride_order=tuple(stride_order), + name=name, + init_value=init_value, + ) + + class JaxApiBase: + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del operation_name + resolved = 100 if target is None else target + if resolved not in supported: + raise ValueError(f"unsupported target {resolved}") + return resolved + + @staticmethod + def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): + public_shape = tuple(value.shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + public_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride(public_shape, public_order) + canonical_axis_by_public_axis = layout_module.to_public_axes( + tuple(range(len(public_shape))), mode + ) + return JaxTensorDesc( + dtype=value.dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple( + canonical_axis_by_public_axis[axis] for axis in public_order + ), + name=name, + init_value=init_value, + ) + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + actual_shape = layout_module.to_canonical_axes(tuple(value.shape), mode) + if actual_shape != expected.shape: + raise ValueError(f"{expected.name} shape mismatch") + if ( + _DTYPE_TO_CUDNN.get(value.dtype, _DataType.NOT_SET) + != expected.cudnn_dtype + ): + raise ValueError(f"{expected.name} dtype mismatch") + + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + return _TensorSpec( + layout=layout_module.to_cutlass_layout( + desc.shape, + desc.stride, + desc.stride_order, + mode=mode, + name=desc.name, + ), + mode=mode, + divisibility=divisibility, + ) + + def _call_kernel(self, inputs, *, output_descs, output_spec, **options): + self.captured_call = { + "inputs": tuple(inputs), + "output_descs": tuple(output_descs), + "output_spec": tuple(output_spec), + **options, + } + return tuple( + _Array( + layout_module.to_public_axes(desc.shape, spec.mode), + _CUDNN_TO_DTYPE[desc.cudnn_dtype], + ) + for desc, spec in zip(output_descs, output_spec) + ) + + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = result_module.TupleDict + + fake_jax = types.ModuleType("jax") + fake_jax.__path__ = [] + fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + fake_jax.ShapeDtypeStruct = _Array + cls.static_argnames = {} + + def jit(function=None, *, static_argnames=()): + def decorate(target): + cls.static_argnames[target.__name__] = tuple(static_argnames) + return target + + return decorate if function is None else decorate(function) + + fake_jax.jit = jit + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.float32 = "float32" + fake_jnp.int32 = "int32" + fake_jnp.iinfo = lambda _dtype: types.SimpleNamespace(max=2**31 - 1) + fake_jax.numpy = fake_jnp + + try: + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "torch": None, + "cutlass": None, + }, + ): + cls.module = importlib.import_module(f"{bsa_name}.jax") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + @staticmethod + def _forward_samples(layout="bhsd"): + q_shape = (2, 4, 256, 128) if layout == "bhsd" else (2, 256, 4, 128) + kv_shape = (2, 2, 512, 128) if layout == "bhsd" else (2, 512, 2, 128) + return ( + _Array(q_shape, "bfloat16"), + _Array(kv_shape, "bfloat16"), + _Array(kv_shape, "bfloat16"), + _Array((2, 2, 4, 8), "int32"), + ) + + def test_forward_infers_public_output_layout_and_cutlass_buffers(self): + q, k, v, block_index = self._forward_samples("bshd") + api = self.module.BlockSparseAttentionForward( + q, + k, + v, + block_index, + block_sparse_num=8, + layout="bshd", + target_compute_capability=100, + ) + result = api(q, k, v, block_index) + + self.assertEqual(api.data_mode, (0, 2, 1, 3)) + self.assertEqual(api.q_desc.shape, (2, 4, 256, 128)) + self.assertTrue(api._op.pack_gqa_effective) + self.assertEqual(tuple(result.keys()), ("o_tensor", "lse_tensor")) + self.assertEqual(result["o_tensor"].shape, (2, 256, 4, 128)) + self.assertEqual(result["lse_tensor"].shape, (2, 4, 256)) + self.assertEqual(api.captured_call["workspace_descs"], ()) + self.assertIn("--gpu-arch sm_100a", api.captured_call["compile_options"]) + + def test_backward_declares_csr_and_accumulator_workspaces(self): + q = _Array((2, 4, 256, 128), "bfloat16") + k = _Array((2, 4, 512, 128), "bfloat16") + lse = _Array((2, 4, 256), "float32") + block_index = _Array((2, 4, 2, 8), "int32") + api = self.module.BlockSparseAttentionBackward( + q, + q, + k, + k, + q, + lse, + block_index, + block_sparse_num=8, + target_compute_capability=100, + ) + result = api(q, q, k, k, q, lse, block_index) + + self.assertEqual(tuple(result.keys()), ("dq_tensor", "dk_tensor", "dv_tensor")) + self.assertEqual(result["dq_tensor"].shape, q.shape) + self.assertEqual(result["dk_tensor"].shape, k.shape) + workspace_names = tuple( + desc.name for desc in api.captured_call["workspace_descs"] + ) + self.assertIn("counts_workspace", workspace_names) + self.assertIn("bucket_offsets_workspace", workspace_names) + self.assertIn("dq_accum_workspace", workspace_names) + counts = next( + desc + for desc in api.captured_call["workspace_descs"] + if desc.name == "counts_workspace" + ) + self.assertEqual(counts.init_value, 0) + + def test_invalid_metadata_and_unsupported_paths_fail_before_lowering(self): + q, k, v, _ = self._forward_samples() + wrong_index = _Array((2, 4, 3, 8), "int32") + with self.assertRaisesRegex(ValueError, "block_index shape prefix"): + self.module.BlockSparseAttentionForward( + q, + k, + v, + wrong_index, + block_sparse_num=8, + pack_gqa=False, + target_compute_capability=100, + ) + + with self.assertRaisesRegex(ValueError, "sparse_block_size=64"): + self.module.BlockSparseAttentionForward( + q, + k, + v, + _Array((2, 4, 4, 8), "int32"), + block_sparse_num=8, + sparse_block_size=128, + target_compute_capability=90, + ) + + def test_split_kv_disables_the_clc_scheduler(self): + resolve = self.module._resolve_sm100_blk64_use_clc + + self.assertFalse( + resolve( + kv_splits=2, + requested=None, + batch=1, + heads=4, + seqlen_q=256, + block_sparse_num=8, + has_variable_block_nums=True, + ) + ) + + q = _Array((2, 4, 256, 128), "bfloat16") + k = _Array((2, 4, 512, 128), "bfloat16") + v = _Array((2, 4, 512, 128), "bfloat16") + block_index = _Array((2, 4, 4, 8), "int32") + with self.assertRaisesRegex(ValueError, r"\[1, 256\]"): + self.module.BlockSparseAttentionForward( + q, + k, + v, + block_index, + block_sparse_num=8, + kv_splits=257, + sparse_block_size=64, + target_compute_capability=100, + ) + + q, k, v, block_index = self._forward_samples() + with self.assertRaisesRegex(ValueError, "SM100-family blk64"): + self.module.BlockSparseAttentionForward( + q, + k, + v, + block_index, + block_sparse_num=8, + use_clc=False, + sparse_block_size=128, + target_compute_capability=100, + ) + + def test_backward_auto_bucket_matches_kernel_tuning(self): + choose = self.module._default_bucket_size + + self.assertEqual(choose(100, 64, 2048, 4), 1152) + self.assertEqual(choose(100, 64, 3000, 4), 1024) + self.assertEqual(choose(100, 128, 4096, 1), 256) + + def test_forward_auto_kv_splits_uses_static_metadata(self): + choose = self.module._sm100_blk64_auto_kv_splits + + self.assertEqual(choose(255), 1) + self.assertEqual(choose(256), 2) + self.assertEqual(choose(450), 4) + self.assertEqual(choose(900), 8) + + def test_function_wrappers_make_configuration_static(self): + self.assertIn( + "block_sparse_num", self.static_argnames["block_sparse_attention_forward"] + ) + self.assertIn("layout", self.static_argnames["block_sparse_attention_forward"]) + self.assertIn( + "bucket_size_blocks", + self.static_argnames["block_sparse_attention_backward"], + ) + self.assertIn( + "target_compute_capability", + self.static_argnames["block_sparse_attention_backward"], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_discrete_grouped_gemm_contract.py b/test/python/fe_api/test_jax_discrete_grouped_gemm_contract.py new file mode 100644 index 000000000..a4101532f --- /dev/null +++ b/test/python/fe_api/test_jax_discrete_grouped_gemm_contract.py @@ -0,0 +1,458 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for the discrete grouped GEMM JAX adapters.""" + +from __future__ import annotations + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_DISCRETE_ROOT = _CUDNN_ROOT / "discrete_grouped_gemm" +_PACKAGE = "cudnn_jax_discrete_grouped_gemm_contract_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + INT64 = auto() + UINT8 = auto() + FP8_E4M3 = auto() + FP8_E5M2 = auto() + FP8_E8M0 = auto() + FP4_E2M1 = auto() + + +_DTYPE_TO_CUDNN = { + "float16": _DataType.HALF, + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, + "int32": _DataType.INT32, + "int64": _DataType.INT64, + "uint8": _DataType.UINT8, + "float4_e2m1fn": _DataType.FP4_E2M1, + "float8_e4m3fn": _DataType.FP8_E4M3, + "float8_e5m2": _DataType.FP8_E5M2, + "float8_e8m0fnu": _DataType.FP8_E8M0, +} + + +class _Array: + def __init__(self, shape, dtype, fill_value=None): + self.shape = tuple(shape) + self.dtype = dtype + self.fill_value = fill_value + self.iterator = object() + + +class _TensorSpec: + def __init__(self, mode=None, ptr_assumed_align=None): + self.mode = mode + self.ptr_assumed_align = ptr_assumed_align + + +class JaxDiscreteGroupedGemmContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + discrete_name = f"{_PACKAGE}.discrete_grouped_gemm" + discrete = types.ModuleType(discrete_name) + discrete.__path__ = [str(_DISCRETE_ROOT)] + discrete.__package__ = discrete_name + discrete.__spec__ = ModuleSpec(discrete_name, loader=None, is_package=True) + sys.modules[discrete_name] = discrete + + for child in ("discrete_grouped_gemm_swiglu", "discrete_grouped_gemm_dswiglu"): + child_name = f"{discrete_name}.{child}" + child_module = types.ModuleType(child_name) + child_module.__path__ = [str(_DISCRETE_ROOT / child)] + child_module.__package__ = child_name + child_module.__spec__ = ModuleSpec(child_name, loader=None, is_package=True) + sys.modules[child_name] = child_module + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + layout_module = importlib.import_module(f"{internal_name}.layout") + result_module = importlib.import_module(f"{_PACKAGE}._result") + + class JaxTensorDesc(tensor_module.TensorDesc): + @property + def cudnn_dtype(self): + return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) + + class JaxApiBase: + @staticmethod + def _resolve_compute_capability(_target, _supported, _name): + return 100 + + @staticmethod + def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): + public_shape = tuple(value.shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + public_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride(public_shape, public_order) + canonical_axis_by_public_axis = layout_module.to_public_axes( + tuple(range(len(public_shape))), mode + ) + return JaxTensorDesc( + dtype=value.dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple( + canonical_axis_by_public_axis[axis] for axis in public_order + ), + name=name, + init_value=init_value, + ) + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + actual = layout_module.to_canonical_axes(tuple(value.shape), mode) + if actual != expected.shape: + raise ValueError( + f"{expected.name} shape mismatch: expected {expected.shape}, got {actual}" + ) + if ( + _DTYPE_TO_CUDNN.get(value.dtype, _DataType.NOT_SET) + != expected.cudnn_dtype + ): + raise ValueError(f"{expected.name} dtype mismatch") + + @staticmethod + def _to_tensor_spec(_desc, *, mode=None, **_unused): + return _TensorSpec(mode) + + @staticmethod + def _materialize_tensor_desc(desc, *, mode=None): + return _Array( + layout_module.to_public_axes(desc.shape, mode), desc.dtype + ) + + @staticmethod + def _get_max_active_clusters(_cluster_size, *, overlap_margin=0): + return 8 - overlap_margin + + def _call_kernel( + self, + inputs, + *, + launch, + output_descs, + workspace_descs=(), + output_spec=(), + workspace_spec=None, + **options, + ): + outputs = tuple( + _Array( + layout_module.to_public_axes(desc.shape, spec.mode), desc.dtype + ) + for desc, spec in zip(output_descs, output_spec) + ) + workspaces = tuple( + _Array(desc.shape, desc.dtype) for desc in workspace_descs + ) + launch("stream", *inputs, *outputs, *workspaces) + self.captured_kernel_call = ( + tuple(inputs), + tuple(output_descs), + tuple(workspace_descs), + options, + ) + self.captured_workspaces = workspaces + self.captured_workspace_specs = workspace_spec + return outputs + + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = result_module.TupleDict + + datatypes = types.ModuleType(f"{internal_name}.datatypes") + datatypes.jax_to_cudnn_dtype = lambda dtype: _DTYPE_TO_CUDNN.get( + dtype, _DataType.NOT_SET + ) + datatypes.normalize_jax_dtype = lambda value, default, _name: ( + default if value is None else value + ) + sys.modules[datatypes.__name__] = datatypes + + fake_jax = types.ModuleType("jax") + fake_jax.ShapeDtypeStruct = _Array + fake_jax.jit = lambda function=None, **_kwargs: ( + (lambda target: target) if function is None else function + ) + fake_jnp = types.ModuleType("jax.numpy") + for dtype in _DTYPE_TO_CUDNN: + setattr(fake_jnp, dtype, dtype) + fake_jnp.dtype = lambda dtype: dtype + fake_jnp.ones = lambda shape, dtype: _Array(shape, dtype) + fake_jnp.empty = lambda shape, dtype: _Array(shape, dtype) + fake_jnp.full = lambda shape, value, dtype: _Array(shape, dtype, value) + fake_jax.numpy = fake_jnp + + try: + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "torch": None, + "cutlass": None, + }, + ): + cls.common = importlib.import_module(f"{discrete_name}._jax_common") + cls.forward = importlib.import_module( + f"{discrete_name}.discrete_grouped_gemm_swiglu.jax" + ) + cls.backward = importlib.import_module( + f"{discrete_name}.discrete_grouped_gemm_dswiglu.jax" + ) + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + @staticmethod + def _scale_shape( + rows: int, k: int, experts: int, sf_vec_size: int = 32 + ) -> tuple[int, ...]: + canonical = ( + 32, + 4, + (rows + 127) // 128, + 4, + ((k + sf_vec_size - 1) // sf_vec_size + 3) // 4, + experts, + ) + return (experts, canonical[2], canonical[4], 32, 4, 4) + + def test_forward_layouts_map_stacked_experts_to_canonical_kernel_axes(self): + m, n, k, experts = 256, 256, 128, 4 + operation = self.forward.DiscreteGroupedGemmSwigluSm100( + _Array((1, m, k), "float8_e4m3fn"), + _Array((experts, n, k), "float8_e4m3fn"), + _Array(self._scale_shape(m, k, 1), "float8_e8m0fnu"), + _Array(self._scale_shape(n, k, experts), "float8_e8m0fnu"), + _Array((experts,), "int32"), + _Array((experts,), "float32"), + sample_norm_const=_Array((1,), "float32"), + sample_prob=_Array((1, 1, m), "float32"), + sample_bias=_Array((experts, n), "bfloat16"), + ) + + self.assertTrue(operation.check_support()) + self.assertEqual(operation.a_desc.shape, (m, k, 1)) + self.assertEqual(operation.b_desc.shape, (n, k, experts)) + self.assertEqual(operation.b_desc.stride_order, (1, 0, 2)) + self.assertEqual(operation.c_desc.shape, (m, n, 1)) + self.assertEqual(operation.d_desc.shape, (m, n // 2, 1)) + self.assertEqual(operation.d_desc.stride_order, (1, 0, 2)) + self.assertEqual(operation.sfd_row_desc.shape, (32, 4, 2, 4, 1, 1)) + self.assertEqual(operation.amax_desc.shape, (experts, 1)) + self.assertEqual(operation.amax_desc.init_value, float("-inf")) + + def test_backward_infers_initialized_reduction_outputs(self): + m, n, k, experts = 256, 256, 128, 4 + operation = self.backward.DiscreteGroupedGemmDswigluSm100( + _Array((1, m, k), "float8_e4m3fn"), + _Array((experts, n, k), "float8_e4m3fn"), + _Array((1, m, 2 * n), "bfloat16"), + _Array(self._scale_shape(m, k, 1), "float8_e8m0fnu"), + _Array(self._scale_shape(n, k, experts), "float8_e8m0fnu"), + _Array((experts,), "int32"), + _Array((experts,), "float32"), + _Array((experts,), "float32"), + _Array((1, 1, m), "float32"), + generate_dbias=True, + ) + + self.assertTrue(operation.check_support()) + self.assertEqual(operation.d_row_desc.shape, (m, 2 * n, 1)) + self.assertEqual(operation.dprob_desc.shape, (m, 1, 1)) + self.assertEqual(operation.dprob_desc.init_value, 0.0) + self.assertEqual(operation.dbias_desc.shape, (experts, 2 * n, 1)) + self.assertEqual(operation.dbias_desc.init_value, 0.0) + self.assertEqual(operation.amax_desc.shape, (experts, 2, 1)) + self.assertEqual(operation.amax_desc.init_value, float("-inf")) + + def test_non_k_major_stacked_weights_are_rejected(self): + m, n, k, experts = 256, 256, 128, 4 + operation = self.forward.DiscreteGroupedGemmSwigluSm100( + _Array((1, m, k), "float8_e4m3fn"), + _Array((experts, k, n), "float8_e4m3fn"), + _Array(self._scale_shape(m, k, 1), "float8_e8m0fnu"), + _Array(self._scale_shape(n, k, experts), "float8_e8m0fnu"), + _Array((experts,), "int32"), + _Array((experts,), "float32"), + b_layout="LKN", + ) + + with self.assertRaisesRegex(ValueError, "K-major A and B"): + operation.check_support() + + def test_launcher_passes_live_stacked_operands_before_outputs_and_workspace(self): + m, n, k, experts = 256, 256, 128, 4 + a = _Array((1, m, k), "float8_e4m3fn") + b = _Array((experts, n, k), "float8_e4m3fn") + sfa = _Array(self._scale_shape(m, k, 1), "float8_e8m0fnu") + sfb = _Array(self._scale_shape(n, k, experts), "float8_e8m0fnu") + offsets = _Array((experts,), "int32") + alpha = _Array((experts,), "float32") + operation = self.forward.DiscreteGroupedGemmSwigluSm100( + a, b, sfa, sfb, offsets, alpha + ) + + captured = {} + + class FakeKernel: + def __init__(self, **config): + captured["config"] = config + + @staticmethod + def get_workspace_bytes(): + return 512 + + def __call__(self, *args): + captured["args"] = args + + kernel_module_name = f"{_PACKAGE}.discrete_grouped_gemm.discrete_grouped_gemm_swiglu.discrete_B_blockscaled_grouped_gemm_glu_bias" + kernel_module = types.ModuleType(kernel_module_name) + kernel_module.BlockScaledDiscreteWeightGroupedGemmBiasKernel = FakeKernel + + cutlass = types.ModuleType("cutlass") + cutlass.Int32 = lambda value: ("i32", value) + cutlass.Int64 = lambda value: ("i64", value) + cutlass.Float32 = lambda value: ("f32", value) + cute = types.ModuleType("cutlass.cute") + nvgpu = types.ModuleType("cutlass.cute.nvgpu") + nvgpu.OperandMajorMode = types.SimpleNamespace(K="K") + cutlass_jax = types.ModuleType("cutlass.jax") + cutlass_jax.TensorSpec = _TensorSpec + cutlass_jax.jax_to_cutlass_dtype = lambda dtype: dtype + cutlass.cute = cute + + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + "cutlass.cute": cute, + "cutlass.cute.nvgpu": nvgpu, + "cutlass.jax": cutlass_jax, + kernel_module_name: kernel_module, + }, + ): + result = operation(a, b, sfa, sfb, offsets, alpha) + + kernel_args = captured["args"] + self.assertIs(kernel_args[0], a) + self.assertIs(kernel_args[1], b.iterator) + self.assertIs(kernel_args[2], sfb.iterator) + self.assertEqual(kernel_args[6], "K") + self.assertIs(kernel_args[7], operation.captured_workspaces[0].iterator) + self.assertEqual(operation.captured_workspaces[0].shape, (512,)) + self.assertEqual(operation.captured_workspace_specs[0].ptr_assumed_align, 128) + self.assertEqual(kernel_args[18].shape, (1, 1, m)) + self.assertEqual(kernel_args[18].dtype, "float32") + self.assertEqual( + tuple(result.keys()), + ( + "c_tensor", + "d_tensor", + "d_col_tensor", + "amax_tensor", + "sfd_row_tensor", + "sfd_col_tensor", + ), + ) + self.assertEqual(kernel_args[20], 8) + self.assertEqual(kernel_args[21], "stream") + self.assertTrue(captured["config"]["stacked_expert_inputs"]) + + def test_zero_m_materializes_results_without_launching(self): + m, n, k, experts = 0, 256, 128, 4 + a = _Array((1, m, k), "float8_e4m3fn") + b = _Array((experts, n, k), "float8_e4m3fn") + sfa = _Array(self._scale_shape(m, k, 1), "float8_e8m0fnu") + sfb = _Array(self._scale_shape(n, k, experts), "float8_e8m0fnu") + offsets = _Array((experts,), "int32") + alpha = _Array((experts,), "float32") + + forward = self.forward.DiscreteGroupedGemmSwigluSm100( + a, b, sfa, sfb, offsets, alpha + ) + forward_result = forward(a, b, sfa, sfb, offsets, alpha) + self.assertEqual(forward_result["c_tensor"].shape, (1, 0, n)) + self.assertEqual(forward_result["d_tensor"].shape, (1, 0, n // 2)) + self.assertEqual(forward_result["amax_tensor"].fill_value, float("-inf")) + + c = _Array((1, m, 2 * n), "bfloat16") + beta = _Array((experts,), "float32") + prob = _Array((1, 1, m), "float32") + backward = self.backward.DiscreteGroupedGemmDswigluSm100( + a, + b, + c, + sfa, + sfb, + offsets, + alpha, + beta, + prob, + generate_dbias=True, + ) + backward_result = backward(a, b, c, sfa, sfb, offsets, alpha, beta, prob) + self.assertEqual(backward_result["d_row_tensor"].shape, (1, 0, 2 * n)) + self.assertEqual(backward_result["dprob_tensor"].shape, (1, 1, 0)) + self.assertEqual(backward_result["dbias_tensor"].fill_value, 0.0) + self.assertEqual(backward_result["amax_tensor"].fill_value, float("-inf")) + self.assertEqual( + tuple(backward_result.keys()), + ( + "d_row_tensor", + "d_col_tensor", + "dprob_tensor", + "dbias_tensor", + "amax_tensor", + "sfd_row_tensor", + "sfd_col_tensor", + ), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py b/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py index 2dec243da..8454d0ee7 100644 --- a/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py +++ b/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py @@ -128,7 +128,9 @@ class JaxTensorDesc(tensor_module.TensorDesc): def cudnn_dtype(self): return cls.dtype_to_cudnn.get(self.dtype, _DataType.NOT_SET) - def compact_like(self, *, cudnn_dtype, shape, stride_order=None, name="", init_value=None): + def compact_like( + self, *, cudnn_dtype, shape, stride_order=None, name="", init_value=None + ): canonical = tensor_module.make_compact_tensor_desc( dtype=cudnn_dtype, shape=shape, @@ -150,18 +152,24 @@ class JaxApiBase: @staticmethod def _to_tensor_desc(value, name, *, mode=None, init_value=None): - del mode - shape = tuple(value.shape) - stride = [0] * len(shape) + public_shape = tuple(value.shape) + mode = tuple(range(len(public_shape))) if mode is None else tuple(mode) + public_stride = [0] * len(public_shape) running = 1 - for dimension in reversed(range(len(shape))): - stride[dimension] = running - running *= max(shape[dimension], 1) + for dimension in reversed(range(len(public_shape))): + public_stride[dimension] = running + running *= max(public_shape[dimension], 1) + canonical_axis_by_public = [0] * len(mode) + for canonical_axis, public_axis in enumerate(mode): + canonical_axis_by_public[public_axis] = canonical_axis return JaxTensorDesc( dtype=value.dtype, - shape=shape, - stride=tuple(stride), - stride_order=tuple(reversed(range(len(shape)))), + shape=tuple(public_shape[axis] for axis in mode), + stride=tuple(public_stride[axis] for axis in mode), + stride_order=tuple( + canonical_axis_by_public[axis] + for axis in reversed(range(len(public_shape))) + ), name=name, init_value=init_value, ) @@ -175,13 +183,19 @@ def _resolve_compute_capability(target, supported, operation_name): @staticmethod def _compute_capability_family(target, supported): - return max((value for value in supported if value <= target), default=None) + return max( + (value for value in supported if value <= target), default=None + ) @staticmethod def _check_tensor_signature(value, expected, *, mode=None): - del mode + mode = tuple(range(len(value.shape))) if mode is None else tuple(mode) actual_dtype = cls.dtype_to_cudnn.get(value.dtype, _DataType.NOT_SET) - if tuple(value.shape) != expected.shape or actual_dtype != expected.cudnn_dtype: + actual_shape = tuple(value.shape[axis] for axis in mode) + if ( + actual_shape != expected.shape + or actual_dtype != expected.cudnn_dtype + ): raise ValueError(f"{expected.name} tensor signature mismatch") @staticmethod @@ -192,9 +206,35 @@ def _to_tensor_spec(desc, *, mode=None, divisibility=None): def _call_kernel(self, inputs, *, launch, **options): options["launch"] = launch JaxApiBase.captured_call = (inputs, options) - return tuple(_Array(desc.shape, cls.cudnn_to_dtype[desc.cudnn_dtype]) for desc in options["output_descs"]) + specs = options.get("output_spec") or (None,) * len( + options["output_descs"] + ) + results = [] + for desc, spec in zip(options["output_descs"], specs): + mode = ( + tuple(range(desc.ndim)) + if spec is None or spec.mode is None + else tuple(spec.mode) + ) + canonical_axis_by_public = [0] * len(mode) + for canonical_axis, public_axis in enumerate(mode): + canonical_axis_by_public[public_axis] = canonical_axis + public_shape = tuple( + desc.shape[canonical_axis_by_public[public_axis]] + for public_axis in range(desc.ndim) + ) + results.append( + _Array( + public_shape, + cls.cudnn_to_dtype[desc.cudnn_dtype], + ) + ) + return tuple(results) internal = types.ModuleType(f"{_PACKAGE}._jax") + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal.__name__ + internal.__spec__ = ModuleSpec(internal.__name__, loader=None, is_package=True) internal.JaxApiBase = JaxApiBase internal.JaxTensorDesc = JaxTensorDesc internal.TupleDict = _TupleDict @@ -205,7 +245,9 @@ def _call_kernel(self, inputs, *, launch, **options): fake_jnp.float32 = cls.float32 fake_jnp.int32 = cls.int32 fake_jnp.dtype = lambda value: value.dtype if hasattr(value, "dtype") else value - fake_jnp.asarray = lambda value, dtype: _Array(getattr(value, "shape", ()), dtype) + fake_jnp.asarray = lambda value, dtype: _Array( + getattr(value, "shape", ()), dtype + ) fake_jnp.reshape = lambda value, shape: _Array(shape, value.dtype) fake_jax = types.ModuleType("jax") @@ -213,7 +255,9 @@ def _call_kernel(self, inputs, *, launch, **options): fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) fake_jax.numpy = fake_jnp fake_jax.ShapeDtypeStruct = _Array - fake_jax.jit = lambda function=None, **_kwargs: (lambda fn: fn) if function is None else function + fake_jax.jit = lambda function=None, **_kwargs: ( + (lambda fn: fn) if function is None else function + ) try: with mock.patch.dict(sys.modules, {"jax": fake_jax, "jax.numpy": fake_jnp}): @@ -258,6 +302,18 @@ def _sparse_thd_inputs(cls, *, topk=128): _Array(score_shape, cls.int32), ) + @classmethod + def _sparse_sequence_major_inputs(cls, *, topk=128): + score_shape = (64, 2, topk) + return ( + _Array((64, 2, 64, 128), cls.bfloat16), + _Array((64, 2, 64), cls.bfloat16), + _Array((256, 2, 128), cls.bfloat16), + _Array(score_shape, cls.float32), + _Array(score_shape, cls.float32), + _Array(score_shape, cls.int32), + ) + @classmethod def _dense_bshd_inputs(cls): return ( @@ -282,18 +338,37 @@ def _dense_thd_inputs(cls): _Array((96,), cls.float32), ) + @classmethod + def _dense_sequence_major_inputs(cls): + return ( + _Array((64, 2, 64, 128), cls.bfloat16), + _Array((64, 2, 64), cls.bfloat16), + _Array((256, 2, 128), cls.bfloat16), + _Array((64, 2, 256), cls.float32), + _Array((64, 2), cls.float32), + _Array((64, 2, 256), cls.float32), + _Array((64, 2), cls.float32), + ) + def test_sparse_wrapper_declares_functional_outputs_and_zeroed_dk(self): inputs = self._sparse_inputs() with mock.patch.dict(sys.modules, self.fake_jax_modules): - result = self.module.indexer_backward_wrapper(*inputs, target_compute_capability=103) + result = self.module.indexer_backward_wrapper( + *inputs, target_compute_capability=103 + ) - self.assertEqual(tuple(result), (result["d_index_q"], result["d_weights"], result["d_index_k"])) + self.assertEqual( + tuple(result), + (result["d_index_q"], result["d_weights"], result["d_index_k"]), + ) self.assertEqual(result["d_index_q"].shape, inputs[0].shape) self.assertEqual(result["d_weights"].shape, inputs[1].shape) self.assertEqual(result["d_index_k"].shape, inputs[2].shape) kernel_inputs, options = self.JaxApiBase.captured_call self.assertEqual(kernel_inputs[:-1], inputs) - self.assertEqual((kernel_inputs[-1].shape, kernel_inputs[-1].dtype), ((1,), self.float32)) + self.assertEqual( + (kernel_inputs[-1].shape, kernel_inputs[-1].dtype), ((1,), self.float32) + ) self.assertEqual(options["output_descs"][2].init_value, 0.0) self.assertEqual(options["workspace_descs"][0].shape, inputs[3].shape) self.assertIn("--gpu-arch sm_103a", options["compile_options"]) @@ -303,7 +378,9 @@ def test_sparse_wrapper_declares_functional_outputs_and_zeroed_dk(self): def test_sparse_packed_thd_signature_requires_global_indices(self): inputs = self._sparse_thd_inputs() with mock.patch.dict(sys.modules, self.fake_jax_modules): - with self.assertRaisesRegex(ValueError, "requires topk_indices_global=True"): + with self.assertRaisesRegex( + ValueError, "requires topk_indices_global=True" + ): self.module.indexer_backward_wrapper( *inputs, target_compute_capability=100, @@ -321,6 +398,39 @@ def test_sparse_packed_thd_signature_requires_global_indices(self): self.assertEqual(options["input_spec"][0].divisibility, (None, None, 128)) self.assertEqual(options["input_spec"][2].divisibility, (None, 128)) + def test_sparse_sequence_major_layouts_preserve_public_shapes(self): + inputs = self._sparse_sequence_major_inputs() + with mock.patch.dict(sys.modules, self.fake_jax_modules): + result = self.module.indexer_backward_wrapper( + *inputs, + q_layout="SBHD", + w_layout="SBH", + k_layout="SBD", + score_layout="SBK", + target_compute_capability=100, + ) + + self.assertEqual(result["d_index_q"].shape, inputs[0].shape) + self.assertEqual(result["d_weights"].shape, inputs[1].shape) + self.assertEqual(result["d_index_k"].shape, inputs[2].shape) + _, options = self.JaxApiBase.captured_call + self.assertEqual( + tuple(spec.mode for spec in options["input_spec"][:6]), + ( + (1, 0, 2, 3), + (1, 0, 2), + (1, 0, 2), + (1, 0, 2), + (1, 0, 2), + (1, 0, 2), + ), + ) + self.assertEqual( + tuple(spec.mode for spec in options["output_spec"]), + ((1, 0, 2, 3), (1, 0, 2), (1, 0, 2)), + ) + self.assertEqual(options["workspace_spec"][0].mode, (1, 0, 2)) + def test_grad_loss_uses_a_fixed_runtime_abi_descriptor(self): inputs = self._sparse_inputs() with mock.patch.dict(sys.modules, self.fake_jax_modules): @@ -345,7 +455,9 @@ def test_dense_bshd_declares_functional_outputs_and_runtime_grad_loss(self): self.assertEqual(result["d_weights"].shape, inputs[1].shape) self.assertEqual(result["d_index_k"].shape, inputs[2].shape) kernel_inputs, options = self.JaxApiBase.captured_call - self.assertEqual((kernel_inputs[7].shape, kernel_inputs[7].dtype), ((1,), self.float32)) + self.assertEqual( + (kernel_inputs[7].shape, kernel_inputs[7].dtype), ((1,), self.float32) + ) self.assertEqual(options["output_descs"][2].init_value, 0.0) self.assertEqual(options["workspace_descs"][0].shape, inputs[3].shape) self.assertIn("--gpu-arch sm_100a", options["compile_options"]) @@ -380,6 +492,41 @@ def test_dense_thd_signature_and_optional_runtime_operands(self): self.assertEqual(options["input_spec"][2].divisibility, (None, 128)) self.assertIn("--gpu-arch sm_90a", options["compile_options"]) + def test_dense_sequence_major_layouts_preserve_public_shapes(self): + inputs = self._dense_sequence_major_inputs() + with mock.patch.dict(sys.modules, self.fake_jax_modules): + result = self.module.dense_indexer_backward_wrapper( + *inputs, + q_layout="SBHD", + w_layout="SBH", + k_layout="SBD", + score_layout="SBK", + denom_layout="SB", + target_compute_capability=100, + ) + + self.assertEqual(result["d_index_q"].shape, inputs[0].shape) + self.assertEqual(result["d_weights"].shape, inputs[1].shape) + self.assertEqual(result["d_index_k"].shape, inputs[2].shape) + _, options = self.JaxApiBase.captured_call + self.assertEqual( + tuple(spec.mode for spec in options["input_spec"][:7]), + ( + (1, 0, 2, 3), + (1, 0, 2), + (1, 0, 2), + (1, 0, 2), + (1, 0), + (1, 0, 2), + (1, 0), + ), + ) + self.assertEqual( + tuple(spec.mode for spec in options["output_spec"]), + ((1, 0, 2, 3), (1, 0, 2), (1, 0, 2)), + ) + self.assertEqual(options["workspace_spec"][0].mode, (1, 0, 2)) + def _kernel_modules(self): package = self.module.__package__ @@ -397,7 +544,10 @@ def kernel(label): dense_100.DenseIndexerBackward2QGemmSm100 = kernel("dense_gemm_100") dense_90 = types.ModuleType(f"{package}.dense_indexer_backward_sm90") dense_90.ScoreGradDenseSm90 = kernel("dense_score_90") - return {module.__name__: module for module in (sparse_100, sparse_90, dense_100, dense_90)} + return { + module.__name__: module + for module in (sparse_100, sparse_90, dense_100, dense_90) + } @staticmethod def _fake_cutlass(): @@ -414,7 +564,9 @@ def _fake_cute(): shape=tuple(shape), stride=tuple(stride), ) - cute.make_tensor = lambda _iterator, layout: _CuteTensor(layout.shape, layout.stride) + cute.make_tensor = lambda _iterator, layout: _CuteTensor( + layout.shape, layout.stride + ) return cute @staticmethod @@ -431,15 +583,26 @@ def test_sparse_sm100_and_sm90_launch_functionally(self): kernel_inputs = tuple(object() for _ in range(7)) outputs = tuple(object() for _ in range(3)) workspace = (object(),) - modules = {**self.fake_jax_modules, "cutlass": self._fake_cutlass(), **self._kernel_modules()} + modules = { + **self.fake_jax_modules, + "cutlass": self._fake_cutlass(), + **self._kernel_modules(), + } - for target, labels in ((100, ("sparse_score_100", "sparse_gemm_100")), (90, ("sparse_score_90", "sparse_gemm_90"))): + for target, labels in ( + (100, ("sparse_score_100", "sparse_gemm_100")), + (90, ("sparse_score_90", "sparse_gemm_90")), + ): with self.subTest(target=target), mock.patch.dict(sys.modules, modules): _RecordedKernel.events.clear() - api = self.module.IndexerBackward(*inputs, target_compute_capability=target) + api = self.module.IndexerBackward( + *inputs, target_compute_capability=target + ) api.check_support() api._launch_kernel(object(), *kernel_inputs, *outputs, *workspace) - self.assertEqual(tuple(event[0] for event in _RecordedKernel.events), labels) + self.assertEqual( + tuple(event[0] for event in _RecordedKernel.events), labels + ) score_args = _RecordedKernel.events[0][2] self.assertIs(score_args[-2 if target == 100 else -1], workspace[0]) gemm_args = _RecordedKernel.events[1][2] @@ -456,7 +619,9 @@ def test_sparse_packed_thd_launch_adds_a_synthetic_batch_view(self): "cutlass.cute": cute, **self._kernel_modules(), } - kernel_inputs = tuple(self._compact_cute(value.shape) for value in inputs) + (self._compact_cute((1,)),) + kernel_inputs = tuple(self._compact_cute(value.shape) for value in inputs) + ( + self._compact_cute((1,)), + ) outputs = ( self._compact_cute(inputs[0].shape), self._compact_cute(inputs[1].shape), @@ -464,7 +629,10 @@ def test_sparse_packed_thd_launch_adds_a_synthetic_batch_view(self): ) workspace = (self._compact_cute(inputs[3].shape),) - for target, labels in ((100, ("sparse_score_100", "sparse_gemm_100")), (90, ("sparse_score_90", "sparse_gemm_90"))): + for target, labels in ( + (100, ("sparse_score_100", "sparse_gemm_100")), + (90, ("sparse_score_90", "sparse_gemm_90")), + ): with self.subTest(target=target), mock.patch.dict(sys.modules, modules): _RecordedKernel.events.clear() api = self.module.IndexerBackward( @@ -475,7 +643,9 @@ def test_sparse_packed_thd_launch_adds_a_synthetic_batch_view(self): api.check_support() api._launch_kernel(object(), *kernel_inputs, *outputs, *workspace) - self.assertEqual(tuple(event[0] for event in _RecordedKernel.events), labels) + self.assertEqual( + tuple(event[0] for event in _RecordedKernel.events), labels + ) score_args = _RecordedKernel.events[0][2] self.assertEqual(score_args[0].shape, (1, 96, 128)) self.assertEqual(score_args[1].shape, (1, 96, 128)) @@ -487,7 +657,11 @@ def test_sparse_packed_thd_launch_adds_a_synthetic_batch_view(self): self.assertEqual(gemm_args[7].shape, (1, 96, 128)) def test_dense_sm100_bshd_and_sm90_thd_launch_functionally(self): - modules = {**self.fake_jax_modules, "cutlass": self._fake_cutlass(), **self._kernel_modules()} + modules = { + **self.fake_jax_modules, + "cutlass": self._fake_cutlass(), + **self._kernel_modules(), + } outputs = tuple(object() for _ in range(3)) workspace = (object(),) @@ -497,7 +671,10 @@ def test_dense_sm100_bshd_and_sm90_thd_launch_functionally(self): api.check_support() kernel_inputs = tuple(object() for _ in range(8)) api._launch_kernel(object(), *kernel_inputs, *outputs, *workspace) - self.assertEqual(tuple(event[0] for event in _RecordedKernel.events), ("dense_score_100", "dense_gemm_100")) + self.assertEqual( + tuple(event[0] for event in _RecordedKernel.events), + ("dense_score_100", "dense_gemm_100"), + ) score_args = _RecordedKernel.events[0][2] self.assertIs(score_args[11], workspace[0]) self.assertIs(score_args[12], kernel_inputs[7]) @@ -520,7 +697,10 @@ def test_dense_sm100_bshd_and_sm90_thd_launch_functionally(self): api.check_support() kernel_inputs = tuple(object() for _ in range(11)) api._launch_kernel(object(), *kernel_inputs, *outputs, *workspace) - self.assertEqual(tuple(event[0] for event in _RecordedKernel.events), ("dense_score_90", "sparse_gemm_90")) + self.assertEqual( + tuple(event[0] for event in _RecordedKernel.events), + ("dense_score_90", "sparse_gemm_90"), + ) score_args = _RecordedKernel.events[0][2] self.assertEqual(score_args[4:7], kernel_inputs[8:11]) self.assertIs(score_args[11], workspace[0]) diff --git a/test/python/fe_api/test_jax_dsa_score_recompute_contract.py b/test/python/fe_api/test_jax_dsa_score_recompute_contract.py index cee84f8e1..c99f20c57 100644 --- a/test/python/fe_api/test_jax_dsa_score_recompute_contract.py +++ b/test/python/fe_api/test_jax_dsa_score_recompute_contract.py @@ -48,6 +48,12 @@ class _TupleDict(dict): pass +class _TensorSpec: + def __init__(self, *, mode=None, divisibility=None): + self.mode = mode + self.divisibility = divisibility + + def _identity_jit(fn=None, **_kwargs): return (lambda decorated: decorated) if fn is None else fn @@ -68,8 +74,12 @@ def setUpClass(cls) -> None: fake_jnp.bfloat16 = _DataType.BFLOAT16 fake_jnp.float32 = _DataType.FLOAT fake_jnp.int32 = _DataType.INT32 - fake_jnp.expand_dims = lambda value, axis: _Array(value.shape[:axis] + (1,) + value.shape[axis:], value.dtype) - fake_jnp.transpose = lambda value, axes: _Array(tuple(value.shape[axis] for axis in axes), value.dtype) + fake_jnp.expand_dims = lambda value, axis: _Array( + value.shape[:axis] + (1,) + value.shape[axis:], value.dtype + ) + fake_jnp.transpose = lambda value, axes: _Array( + tuple(value.shape[axis] for axis in axes), value.dtype + ) fake_jax = types.ModuleType("jax") fake_jax.__path__ = [] @@ -83,10 +93,24 @@ def setUpClass(cls) -> None: class FakeJaxApiBase: @staticmethod def _to_tensor_desc(value, name, *, mode=None, init_value=None): - del mode - return tensor_module.make_compact_tensor_desc( + public_shape = tuple(value.shape) + mode = tuple(range(len(public_shape))) if mode is None else tuple(mode) + public_stride = [0] * len(public_shape) + running = 1 + for dimension in reversed(range(len(public_shape))): + public_stride[dimension] = running + running *= max(public_shape[dimension], 1) + canonical_axis_by_public = [0] * len(mode) + for canonical_axis, public_axis in enumerate(mode): + canonical_axis_by_public[public_axis] = canonical_axis + return tensor_module.TensorDesc( dtype=value.dtype, - shape=value.shape, + shape=tuple(public_shape[axis] for axis in mode), + stride=tuple(public_stride[axis] for axis in mode), + stride_order=tuple( + canonical_axis_by_public[axis] + for axis in reversed(range(len(public_shape))) + ), name=name, init_value=init_value, ) @@ -101,13 +125,20 @@ def _resolve_compute_capability(target, supported, operation_name): @staticmethod def _check_tensor_signature(value, expected, *, mode=None): - del mode - if tuple(value.shape) != expected.shape: + mode = tuple(range(len(value.shape))) if mode is None else tuple(mode) + if tuple(value.shape[axis] for axis in mode) != expected.shape: raise ValueError(f"{expected.name} tensor shape mismatch") if value.dtype != expected.cudnn_dtype: raise ValueError(f"{expected.name} tensor dtype mismatch") - def _call_kernel(self, inputs, *, launch, output_descs, workspace_descs=(), **options): + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + del desc + return _TensorSpec(mode=mode, divisibility=divisibility) + + def _call_kernel( + self, inputs, *, launch, output_descs, workspace_descs=(), **options + ): self.captured_call = { "inputs": tuple(inputs), "outputs": tuple(output_descs), @@ -115,9 +146,34 @@ def _call_kernel(self, inputs, *, launch, output_descs, workspace_descs=(), **op "launch": launch, **options, } - return tuple(_Array(desc.shape, desc.cudnn_dtype) for desc in output_descs) + specs = options.get("output_spec") or (None,) * len(output_descs) + outputs = [] + for desc, spec in zip(output_descs, specs): + mode = ( + tuple(range(desc.ndim)) + if spec is None or spec.mode is None + else tuple(spec.mode) + ) + canonical_axis_by_public = [0] * len(mode) + for canonical_axis, public_axis in enumerate(mode): + canonical_axis_by_public[public_axis] = canonical_axis + outputs.append( + _Array( + tuple( + desc.shape[canonical_axis_by_public[public_axis]] + for public_axis in range(desc.ndim) + ), + desc.cudnn_dtype, + ) + ) + return tuple(outputs) fake_internal_jax = types.ModuleType(f"{_PACKAGE}._jax") + fake_internal_jax.__path__ = [str(_CUDNN_ROOT / "_jax")] + fake_internal_jax.__package__ = fake_internal_jax.__name__ + fake_internal_jax.__spec__ = ModuleSpec( + fake_internal_jax.__name__, loader=None, is_package=True + ) fake_internal_jax.JaxApiBase = FakeJaxApiBase fake_internal_jax.JaxTensorDesc = tensor_module.TensorDesc fake_internal_jax.TupleDict = _TupleDict @@ -177,23 +233,47 @@ def sparse_samples(self): self.array((2, 4, 128), _DataType.INT32), ) + def sequence_major_sparse_samples(self): + return ( + self.array((4, 2, 32, 128)), + self.array((256, 2, 128)), + self.array((4, 2, 32)), + self.array((4, 2, 128), _DataType.INT32), + ) + def dense_samples(self, *, attention=False): return ( self.array((2, 4, 32, 128)), self.array((2, 8, 1, 128)), - self.array((2, 4, 32), _DataType.FLOAT if attention else _DataType.BFLOAT16), + self.array( + (2, 4, 32), _DataType.FLOAT if attention else _DataType.BFLOAT16 + ), + ) + + def sequence_major_dense_samples(self): + return ( + self.array((4, 2, 32, 128)), + self.array((8, 2, 1, 128)), + self.array((4, 2, 32)), ) def test_sm100_sparse_infers_output_and_hidden_length_workspace(self): q, k, weights, indices = self.sparse_samples() - api = self.adapter.SparseIndexerScoreRecompute(q, k, weights, indices, target_compute_capability=100) + api = self.adapter.SparseIndexerScoreRecompute( + q, k, weights, indices, target_compute_capability=100 + ) result = api(q, k, weights, indices) self.assertIsInstance(result, _TupleDict) - self.assertEqual((result["predict"].shape, result["predict"].dtype), ((2, 4, 128), _DataType.FLOAT)) + self.assertEqual( + (result["predict"].shape, result["predict"].dtype), + ((2, 4, 128), _DataType.FLOAT), + ) self.assertEqual(len(api.captured_call["workspaces"]), 1) workspace = api.captured_call["workspaces"][0] - self.assertEqual((workspace.shape, workspace.cudnn_dtype), ((1, 1), _DataType.INT32)) + self.assertEqual( + (workspace.shape, workspace.cudnn_dtype), ((1, 1), _DataType.INT32) + ) self.assertEqual(api.captured_call["compile_options"], "--gpu-arch sm_100") def test_explicit_topk_length_remains_a_real_input(self): @@ -218,7 +298,9 @@ def test_explicit_topk_length_remains_a_real_input(self): def test_sm90_sparse_adapts_k_and_per_head_layouts(self): q, k, weights, indices = self.sparse_samples() - api = self.adapter.SparseIndexerScoreRecompute(q, k, weights, indices, target_compute_capability=90) + api = self.adapter.SparseIndexerScoreRecompute( + q, k, weights, indices, target_compute_capability=90 + ) api(q, k, weights, indices) _, kernel_k, kernel_weights, _ = api.captured_call["inputs"] @@ -226,6 +308,35 @@ def test_sm90_sparse_adapts_k_and_per_head_layouts(self): self.assertEqual(kernel_weights.shape, (2, 32, 4)) self.assertEqual(api.captured_call["workspaces"], ()) + def test_sparse_sequence_major_layouts_preserve_public_shapes(self): + q, k, weights, indices = self.sequence_major_sparse_samples() + api = self.adapter.SparseIndexerScoreRecompute( + q, + k, + weights, + indices, + q_layout="SBHD", + k_layout="SBD", + per_head_layout="SBH", + score_layout="SBK", + output_layout="SBK", + target_compute_capability=90, + ) + result = api(q, k, weights, indices) + + self.assertEqual(result["predict"].shape, indices.shape) + _, kernel_k, kernel_weights, _ = api.captured_call["inputs"] + self.assertEqual(kernel_k.shape, (256, 2, 1, 128)) + self.assertEqual(kernel_weights.shape, (2, 32, 4)) + self.assertEqual( + tuple(spec.mode for spec in api.captured_call["input_spec"]), + ((1, 0, 2, 3), (1, 0, 2, 3), None, (1, 0, 2)), + ) + self.assertEqual( + api.captured_call["output_spec"][0].mode, + (1, 0, 2), + ) + def test_dense_bshd_outputs_are_functional_and_initialized(self): q, k, weights = self.dense_samples() offsets = self.array((2,), _DataType.INT32) @@ -256,6 +367,32 @@ def test_dense_bshd_outputs_are_functional_and_initialized(self): self.assertEqual(explicit.out_desc.init_value, float("-inf")) self.assertIsNone(explicit.denom_desc.init_value) + def test_dense_sequence_major_layouts_preserve_public_shapes(self): + q, k, weights = self.sequence_major_dense_samples() + api = self.adapter.DenseIndexerScoreRecompute( + q, + k, + weights, + q_layout="SBHD", + k_layout="SBHD", + per_head_layout="SBH", + output_layout="SBK", + denom_layout="SB", + target_compute_capability=100, + ) + result = api(q, k, weights) + + self.assertEqual(result["out"].shape, (4, 2, 8)) + self.assertEqual(result["denom"].shape, (4, 2)) + self.assertEqual( + tuple(spec.mode for spec in api.captured_call["input_spec"]), + ((1, 0, 2, 3), (1, 0, 2, 3), (1, 0, 2)), + ) + self.assertEqual( + tuple(spec.mode for spec in api.captured_call["output_spec"]), + ((1, 0, 2), (1, 0)), + ) + def test_dense_thd_sm100_and_sm90_rejection(self): q = self.array((7, 32, 128)) k = self.array((11, 1, 128)) @@ -345,9 +482,13 @@ def __call__(self, *args): sm100_module.SparseScoreRecomputeSm100 = FakeSm100Kernel q, k, weights, indices = self.sparse_samples() - sm90 = self.adapter.SparseIndexerScoreRecompute(q, k, weights, indices, target_compute_capability=90) + sm90 = self.adapter.SparseIndexerScoreRecompute( + q, k, weights, indices, target_compute_capability=90 + ) sm90.check_support() - sm100 = self.adapter.SparseIndexerScoreRecompute(q, k, weights, indices, target_compute_capability=100) + sm100 = self.adapter.SparseIndexerScoreRecompute( + q, k, weights, indices, target_compute_capability=100 + ) sm100.check_support() with mock.patch.dict( sys.modules, @@ -358,15 +499,38 @@ def __call__(self, *args): }, ): sm90._launch_kernel("stream90", "q90", "k90", "aux90", "indices90", "out90") - sm100._launch_kernel("stream100", "q100", "k100", "aux100", "indices100", "out100", "dummy") + sm100._launch_kernel( + "stream100", "q100", "k100", "aux100", "indices100", "out100", "dummy" + ) self.assertEqual( [entry for entry in calls if entry[0].endswith("call")], [ - ("sm90_call", ("q90", "k90", "indices90", "stream90", "out90", "aux90", None, None)), + ( + "sm90_call", + ( + "q90", + "k90", + "indices90", + "stream90", + "out90", + "aux90", + None, + None, + ), + ), ( "sm100_call", - ("q100", "k100", "aux100", "indices100", "out100", "dummy", ("Float32", 1.0), "stream100"), + ( + "q100", + "k100", + "aux100", + "indices100", + "out100", + "dummy", + ("Float32", 1.0), + "stream100", + ), ), ], ) @@ -424,13 +588,30 @@ def __call__(self, *args): sm100_module.__name__: sm100_module, }, ): - sm90._launch_kernel("stream90", "q90", "k90", "aux90", "offset90", "out90", "denom90") - sm100._launch_kernel("stream100", "q100", "k100", "aux100", "offset100", "out100", "denom100") + sm90._launch_kernel( + "stream90", "q90", "k90", "aux90", "offset90", "out90", "denom90" + ) + sm100._launch_kernel( + "stream100", "q100", "k100", "aux100", "offset100", "out100", "denom100" + ) self.assertEqual( [entry for entry in calls if entry[0].endswith("call")], [ - ("sm90_call", ("q90", "k90", None, "stream90", "out90", "aux90", None, "denom90", "offset90")), + ( + "sm90_call", + ( + "q90", + "k90", + None, + "stream90", + "out90", + "aux90", + None, + "denom90", + "offset90", + ), + ), ( "sm100_call", ( @@ -455,20 +636,37 @@ def test_common_and_jax_modules_keep_framework_import_boundaries(self): def imports_framework(path, framework): tree = ast.parse(path.read_text(), filename=str(path)) for node in ast.walk(tree): - if isinstance(node, ast.Import) and any(alias.name == framework or alias.name.startswith(f"{framework}.") for alias in node.names): + if isinstance(node, ast.Import) and any( + alias.name == framework or alias.name.startswith(f"{framework}.") + for alias in node.names + ): return True - if isinstance(node, ast.ImportFrom) and node.level == 0 and (node.module == framework or (node.module or "").startswith(f"{framework}.")): + if ( + isinstance(node, ast.ImportFrom) + and node.level == 0 + and ( + node.module == framework + or (node.module or "").startswith(f"{framework}.") + ) + ): return True return False for filename in ("config.py", "op.py"): path = _SCORE_ROOT / filename for framework in ("torch", "jax", "cutlass", "cuda"): - self.assertFalse(imports_framework(path, framework), f"{filename} imports {framework}") + self.assertFalse( + imports_framework(path, framework), + f"{filename} imports {framework}", + ) self.assertFalse(imports_framework(_SCORE_ROOT / "jax.py", "torch")) init_tree = ast.parse((_SCORE_ROOT / "__init__.py").read_text()) - imported_modules = {node.module for node in ast.walk(init_tree) if isinstance(node, ast.ImportFrom)} + imported_modules = { + node.module + for node in ast.walk(init_tree) + if isinstance(node, ast.ImportFrom) + } self.assertNotIn("api", imported_modules) diff --git a/test/python/fe_api/test_jax_grouped_gemm_contract.py b/test/python/fe_api/test_jax_grouped_gemm_contract.py new file mode 100644 index 000000000..854772c06 --- /dev/null +++ b/test/python/fe_api/test_jax_grouped_gemm_contract.py @@ -0,0 +1,567 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free source contracts for JAX grouped GEMM adapters.""" + +from __future__ import annotations + +import ast +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +from unittest import mock + +import pytest + +pytestmark = pytest.mark.L0 + +_ROOT = Path(__file__).resolve().parents[3] +_GROUPED_ROOT = _ROOT / "python" / "cudnn" / "grouped_gemm" +_FAMILIES = ( + "swiglu", + "dswiglu", + "quant", + "srelu", + "dsrelu", + "glu", + "glu_hadamard", + "dglu", + "wgrad", +) + +_NATIVE_FP4_FAMILIES = ( + "swiglu", + "dswiglu", + "quant", + "srelu", + "dsrelu", + "glu", + "dglu", +) + + +def _class_stem(family: str) -> str: + return "".join(part.capitalize() for part in family.split("_")) + + +def _jax_path(family: str) -> Path: + return _GROUPED_ROOT / f"grouped_gemm_{family}" / "jax.py" + + +@pytest.mark.parametrize("family", _FAMILIES) +def test_each_torch_family_has_a_jax_class_and_wrapper(family: str): + tree = ast.parse(_jax_path(family).read_text()) + definitions = { + node.name: node + for node in tree.body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) + } + + assert f"GroupedGemm{_class_stem(family)}Sm100" in definitions + wrapper = definitions[f"grouped_gemm_{family}_wrapper_sm100"] + assert isinstance(wrapper, ast.FunctionDef) + assert wrapper.decorator_list + assert "jax.jit" in ast.unparse(wrapper.decorator_list[0]) + + +@pytest.mark.parametrize("family", _FAMILIES) +def test_launchers_receive_precomputed_occupancy_and_preserve_stream_first(family: str): + source = _jax_path(family).read_text() + tree = ast.parse(source) + launcher = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "_launch" + ) + + assert launcher.args.args[0].arg == "stream" + assert "max_active_clusters" in { + argument.arg for argument in launcher.args.kwonlyargs + } + assert "HardwareInfo" not in source + assert "resolve_max_active_clusters" not in source + + +@pytest.mark.parametrize("family", _FAMILIES) +def test_jax_module_has_no_top_level_torch_or_kernel_import(family: str): + tree = ast.parse(_jax_path(family).read_text()) + imported_modules = [] + for node in tree.body: + if isinstance(node, ast.Import): + imported_modules.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + imported_modules.append(node.module or "") + + assert not any( + module == "torch" or module.startswith("torch.") for module in imported_modules + ) + assert not any( + "grouped_gemm_" in module and module != "_jax_api" + for module in imported_modules + ) + + +def test_shared_adapter_uses_current_jax_api_base_and_call_kernel(): + source = (_GROUPED_ROOT / "_jax_api.py").read_text() + tree = ast.parse(source) + classes = {node.name: node for node in tree.body if isinstance(node, ast.ClassDef)} + + assert ast.unparse(classes["ApiBaseJax"].bases[0]) == "JaxApiBase" + assert "_CALLER._call_kernel(" in source + assert "_CALLER._get_max_active_clusters(" in source + assert source.index("_CALLER._get_max_active_clusters(") < source.index( + "def launch(stream:" + ) + + +def test_grouped_layouts_are_explicit_row_major_public_mappings(): + source = (_GROUPED_ROOT / "_jax_api.py").read_text() + + assert "return _row_major_spec(6, BLOCK_SCALE_MODE)" in source + assert "return _row_major_spec(3, (2, 1, 0))" in source + assert "return _row_major_spec(2, (1, 0))" in source + assert "_canonical_block_scale_shape(rows, k, batch, sf_vec_size)" in source + assert "BLOCK_SCALE_MODE" in source + + +@pytest.mark.parametrize("family", ("dswiglu", "dglu", "dsrelu")) +def test_dprob_descriptors_use_public_probability_shape(family: str): + tree = ast.parse(_jax_path(family).read_text()) + dprob_calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "make_buffer_desc" + and node.args + and isinstance(node.args[0], ast.Constant) + and node.args[0].value == "dprob_tensor" + ] + + assert dprob_calls + assert all(ast.unparse(call.args[1]) == "(1, 1, m)" for call in dprob_calls) + + +def test_grouped_buffers_use_tensor_descriptors_directly(): + source = (_GROUPED_ROOT / "_jax_api.py").read_text() + + assert "class BufferSpec" not in source + assert "def make_buffer_desc(" in source + assert "init_value=init_value" in source + + +@pytest.mark.parametrize("family", _NATIVE_FP4_FAMILIES) +def test_dense_grouped_adapters_use_native_fp4_recipe_validation(family: str): + source = _jax_path(family).read_text() + + assert "require_grouped_input_scales(" in source + assert "jnp.float4_e2m1fn" in source or "is_fp4_dtype(ab_dtype)" in source + assert "torch.uint8" not in source + + +@pytest.mark.parametrize("family", ("swiglu", "srelu", "quant")) +def test_native_fp4_float32_output_preserves_torch_recipe_guard(family: str): + source = _jax_path(family).read_text() + + assert "sf_vec_size == 16" in source + assert "d_dtype == jnp.dtype(jnp.float32)" in source + + +def test_native_fp4_glu_float32_output_requires_bias_for_sf_vec_16(): + source = _jax_path("glu").read_text() + + assert "sf_vec_size == 16" in source + assert "d_dtype == jnp.dtype(jnp.float32)" in source + assert "bias_tensor is None" in source + + +def test_glu_preserves_torch_sfd_generation_for_all_fp8_inputs(): + source = _jax_path("glu").read_text() + + assert "generate_sfd = is_fp8_dtype(ab_dtype)" in source + + +def test_glu_accepts_native_fp4_intermediate_output(): + tree = ast.parse(_jax_path("glu").read_text()) + c_dtype_call = next( + node.value + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == "c_dtype" + for target in node.targets + ) + and isinstance(node.value, ast.Call) + and isinstance(node.value.func, ast.Name) + and node.value.func.id == "require_dtype" + ) + + assert "jnp.float4_e2m1fn" in ast.unparse(c_dtype_call.args[1]) + + +@pytest.mark.parametrize("family", ("dswiglu", "dglu", "dsrelu")) +def test_native_fp4_backward_paths_return_amax_without_sfd(family: str): + source = _jax_path(family).read_text() + + assert "generate_sfd = is_fp8_dtype(ab_dtype)" in source + assert '"has_amax": has_amax' in source + assert '"amax_tensor"' in source + + +@pytest.mark.parametrize("family", ("dswiglu", "dglu", "dsrelu")) +def test_backward_inputs_validate_canonical_descriptor_shapes(family: str): + tree = ast.parse(_jax_path(family).read_text()) + invalid_checks = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "require_array" + and node.args + and isinstance(node.args[0], ast.Name) + and node.args[0].id == "c_desc" + ] + + assert not invalid_checks + assert "if c_desc.shape != expected_c_shape:" in _jax_path(family).read_text() + + +def test_grouped_packages_keep_torch_exports_lazy(): + root_source = (_GROUPED_ROOT / "__init__.py").read_text() + assert "make_operation_api" in root_source + assert ".api import" not in root_source + + for family in _FAMILIES: + source = (_GROUPED_ROOT / f"grouped_gemm_{family}" / "__init__.py").read_text() + assert "make_operation_api" in source + assert 'submodules=("api", "jax")' in source + assert ".api import" not in source + + +def test_grouped_package_import_does_not_load_torch_apis(): + package_name = "cudnn_frontend_grouped_lazy_contract" + package = types.ModuleType(package_name) + package.__path__ = [str(_ROOT / "python" / "cudnn")] + package.__package__ = package_name + + with mock.patch.dict(sys.modules, {package_name: package}): + grouped = importlib.import_module(f"{package_name}.grouped_gemm") + expected = { + f"GroupedGemm{_class_stem(family)}Sm100" for family in _FAMILIES + } | {f"grouped_gemm_{family}_wrapper_sm100" for family in _FAMILIES} + assert set(grouped.__all__) == expected + assert not any( + module_name.startswith(f"{package_name}.grouped_gemm.") + for module_name in sys.modules + ) + + for family in _FAMILIES: + subpackage_name = f"{package_name}.grouped_gemm.grouped_gemm_{family}" + subpackage = importlib.import_module(subpackage_name) + assert "jax" in dir(subpackage) + assert f"{subpackage_name}.api" not in sys.modules + + for module_name in tuple(sys.modules): + if module_name == package_name or module_name.startswith(f"{package_name}."): + sys.modules.pop(module_name, None) + + +def test_shared_layout_helpers_map_row_major_arrays_to_kernel_axes(): + class DataType(Enum): + NOT_SET = auto() + + class TensorSpec: + def __init__(self, *, layout=None, mode=None, **kwargs): + self.layout = layout + self.mode = mode + self.__dict__.update(kwargs) + + class ShapeDtypeStruct: + def __init__(self, shape, dtype, fill_value=None): + self.shape = tuple(shape) + self.dtype = dtype + self.fill_value = fill_value + + fake_jax = types.ModuleType("jax") + fake_jax.__path__ = [] + fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + fake_jax.ShapeDtypeStruct = ShapeDtypeStruct + fake_jax.local_devices = lambda **_kwargs: () + fake_jax.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, + ) + + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.dtype = lambda value: value + for dtype_name in ( + "float4_e2m1fn", + "float8_e4m3fn", + "float8_e5m2", + "float8_e8m0fnu", + ): + setattr(fake_jnp, dtype_name, dtype_name) + fake_jnp.empty = lambda shape, dtype: ShapeDtypeStruct(shape, dtype) + fake_jnp.full = lambda shape, fill_value, dtype: ShapeDtypeStruct( + shape, dtype, fill_value + ) + fake_jax.numpy = fake_jnp + + fake_cutlass = types.ModuleType("cutlass") + fake_cutlass.__path__ = [] + fake_cutlass.__spec__ = ModuleSpec("cutlass", loader=None, is_package=True) + fake_cutlass_jax = types.ModuleType("cutlass.jax") + fake_cutlass_jax.TensorSpec = TensorSpec + fake_cutlass.jax = fake_cutlass_jax + + package_name = "cudnn_frontend_grouped_layout_contract" + package = types.ModuleType(package_name) + package.__path__ = [str(_ROOT / "python" / "cudnn")] + package.__package__ = package_name + package.data_type = DataType + + grouped_package = types.ModuleType(f"{package_name}.grouped_gemm") + grouped_package.__path__ = [str(_GROUPED_ROOT)] + grouped_package.__package__ = grouped_package.__name__ + + modules = { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "cutlass": fake_cutlass, + "cutlass.jax": fake_cutlass_jax, + package_name: package, + grouped_package.__name__: grouped_package, + } + with mock.patch.dict(sys.modules, modules): + adapter = importlib.import_module(f"{package_name}.grouped_gemm._jax_api") + + assert adapter.gemm_a_tensor_spec("LMK").mode == (1, 2, 0) + assert adapter.gemm_b_tensor_spec("LKN").mode == (2, 1, 0) + assert adapter.probability_tensor_spec().mode == (2, 1, 0) + assert adapter.grouped_bias_tensor_spec().mode == (1, 0) + assert adapter.block_scale_shape(256, 64, 4, 32) == (4, 2, 1, 32, 4, 4) + + sample = ShapeDtypeStruct((1, 256, 64), "float8") + desc = adapter.ApiBaseJax().make_tensor_desc( + sample, + tensor_spec=adapter.gemm_a_tensor_spec("LMK"), + name="sample_a", + ) + assert desc.shape == (256, 64, 1) + assert desc.stride == (64, 1, 16384) + assert desc.array_shape == sample.shape + + fp4_sfa = ShapeDtypeStruct( + adapter.block_scale_shape(256, 256, 1, 16), + "float8_e4m3fn", + ) + fp4_sfb = ShapeDtypeStruct( + adapter.block_scale_shape(128, 256, 2, 16), + "float8_e4m3fn", + ) + assert ( + adapter.require_grouped_input_scales( + fp4_sfa, + fp4_sfb, + m=256, + n=128, + k=256, + experts=2, + sf_vec_size=16, + ab_dtype="float4_e2m1fn", + ) + == "float8_e4m3fn" + ) + + invalid_fp4_sfa = ShapeDtypeStruct( + adapter.block_scale_shape(256, 256, 1, 32), + "float8_e4m3fn", + ) + invalid_fp4_sfb = ShapeDtypeStruct( + adapter.block_scale_shape(128, 256, 2, 32), + "float8_e4m3fn", + ) + with pytest.raises(ValueError, match="sf_vec_size=16"): + adapter.require_grouped_input_scales( + invalid_fp4_sfa, + invalid_fp4_sfb, + m=256, + n=128, + k=256, + experts=2, + sf_vec_size=32, + ab_dtype="float4_e2m1fn", + ) + + mxfp4_sfa = ShapeDtypeStruct( + adapter.block_scale_shape(256, 256, 1, 32), + "float8_e8m0fnu", + ) + mxfp4_sfb = ShapeDtypeStruct( + adapter.block_scale_shape(128, 256, 2, 32), + "float8_e8m0fnu", + ) + assert ( + adapter.require_grouped_input_scales( + mxfp4_sfa, + mxfp4_sfb, + m=256, + n=128, + k=256, + experts=2, + sf_vec_size=32, + ab_dtype="float4_e2m1fn", + ) + == "float8_e8m0fnu" + ) + with pytest.raises(ValueError, match="Unsupported grouped GEMM input dtype"): + adapter.require_grouped_input_scales( + mxfp4_sfa, + mxfp4_sfb, + m=256, + n=128, + k=256, + experts=2, + sf_vec_size=32, + ab_dtype="uint8", + ) + + fp8_sfa = ShapeDtypeStruct( + adapter.block_scale_shape(256, 256, 1, 32), + "float8_e8m0fnu", + ) + fp8_sfb = ShapeDtypeStruct( + adapter.block_scale_shape(128, 256, 2, 32), + "float8_e8m0fnu", + ) + assert ( + adapter.require_grouped_input_scales( + fp8_sfa, + fp8_sfb, + m=256, + n=128, + k=256, + experts=2, + sf_vec_size=32, + ab_dtype="float8_e4m3fn", + ) + == "float8_e8m0fnu" + ) + + launched = [] + lowered = {} + + def launch(stream, input_value, output_value, workspace_value, **static): + launched.append( + (stream, input_value, output_value, workspace_value, static) + ) + + def fake_call_kernel(inputs, **options): + lowered.update(inputs=inputs, **options) + options["launch"]("stream", *inputs, "output", "workspace") + return ("result",) + + with ( + mock.patch.object( + adapter._CALLER, "_get_max_active_clusters", return_value=7 + ) as occupancy, + mock.patch.object( + adapter._CALLER, "_resolve_compute_capability", return_value=100 + ), + mock.patch.object( + adapter._CALLER, "_call_kernel", side_effect=fake_call_kernel + ), + ): + result = adapter.call_cutedsl( + launch, + (sample,), + outputs=( + adapter.make_buffer_desc( + "output", + (1, 256, 128), + "bfloat16", + tensor_spec=adapter.gemm_c_tensor_spec("LMN"), + ), + ), + workspaces=(adapter.make_buffer_desc("workspace", (4,), "uint8"),), + input_specs=(adapter.gemm_a_tensor_spec("LMK"),), + static_args={ + "cluster_shape_mn": (2, 1), + "cluster_overlap_margin": 1, + "expert_cnt": 4, + }, + ) + + assert result == ("result",) + occupancy.assert_called_once_with(2, overlap_margin=1) + assert launched == [ + ( + "stream", + sample, + "output", + "workspace", + { + "cluster_shape_mn": (2, 1), + "expert_cnt": 4, + "max_active_clusters": 7, + }, + ) + ] + assert lowered["output_descs"][0].shape == (256, 128, 1) + assert lowered["workspace_descs"][0].shape == (4,) + + def materialize(desc, *, mode=None): + return ShapeDtypeStruct( + adapter.to_public_axes(desc.shape, mode), + desc.dtype, + ) + + with ( + mock.patch.object( + adapter._CALLER, + "_get_max_active_clusters", + side_effect=AssertionError( + "empty grouped GEMM must not query occupancy" + ), + ), + mock.patch.object( + adapter._CALLER, + "_call_kernel", + side_effect=AssertionError("empty grouped GEMM must not lower"), + ), + mock.patch.object( + adapter._CALLER, + "_materialize_tensor_desc", + side_effect=materialize, + ), + ): + empty, initialized = adapter.call_cutedsl( + launch, + (sample,), + outputs=( + adapter.make_buffer_desc( + "empty", + (1, 0, 128), + "bfloat16", + tensor_spec=adapter.gemm_c_tensor_spec("LMN"), + ), + adapter.make_buffer_desc( + "amax", + (4, 1), + "float32", + init_value=float("-inf"), + ), + ), + static_args={"cluster_shape_mn": (2, 1)}, + ) + + assert empty.shape == (1, 0, 128) + assert initialized.shape == (4, 1) + assert initialized.fill_value == float("-inf") + + for module_name in tuple(sys.modules): + if module_name == package_name or module_name.startswith(f"{package_name}."): + sys.modules.pop(module_name, None) diff --git a/test/python/fe_api/test_jax_grouped_wgrad_accumulation_contract.py b/test/python/fe_api/test_jax_grouped_wgrad_accumulation_contract.py new file mode 100644 index 000000000..f8d7a6674 --- /dev/null +++ b/test/python/fe_api/test_jax_grouped_wgrad_accumulation_contract.py @@ -0,0 +1,84 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for functional JAX grouped-wgrad accumulation.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.L0 + +_ROOT = Path(__file__).resolve().parents[3] +_WGRAD_PATH = ( + _ROOT / "python" / "cudnn" / "grouped_gemm" / "grouped_gemm_wgrad" / "jax.py" +) + + +def _definitions() -> dict[str, ast.AST]: + tree = ast.parse(_WGRAD_PATH.read_text()) + return { + node.name: node + for node in tree.body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) + } + + +def _method(node: ast.ClassDef, name: str) -> ast.FunctionDef: + return next( + child + for child in node.body + if isinstance(child, ast.FunctionDef) and child.name == name + ) + + +def _argument_names(function: ast.FunctionDef) -> tuple[str, ...]: + return tuple( + argument.arg for argument in (*function.args.args, *function.args.kwonlyargs) + ) + + +def test_wgrad_seed_is_part_of_the_class_and_wrapper_signatures(): + definitions = _definitions() + operation = definitions["GroupedGemmWgradSm100"] + assert isinstance(operation, ast.ClassDef) + wrapper = definitions["grouped_gemm_wgrad_wrapper_sm100"] + assert isinstance(wrapper, ast.FunctionDef) + + assert "sample_wgrad_tensor" in _argument_names(_method(operation, "__init__")) + assert "wgrad_tensor" in _argument_names(_method(operation, "__call__")) + assert "wgrad_tensor" in _argument_names(_method(operation, "_call_impl")) + assert "wgrad_tensor" in _argument_names(wrapper) + + +def test_wgrad_seed_is_aliased_as_the_custom_call_result(): + implementation = _definitions()["_grouped_gemm_wgrad_impl"] + assert isinstance(implementation, ast.FunctionDef) + calls = [ + node + for node in ast.walk(implementation) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "call_cutedsl" + ] + assert len(calls) == 1 + + keywords = {keyword.arg: keyword.value for keyword in calls[0].keywords} + assert ast.unparse(keywords["output_seeds"]) == "(wgrad_tensor,)" + + source = ast.get_source_segment(_WGRAD_PATH.read_text(), implementation) + assert source is not None + assert "wgrad_tensor is only valid when accumulate_on_output=True" in source + assert "wgrad_dtype must match wgrad_tensor.dtype" in source + assert "accumulate_on_output and wgrad_tensor is None" in source + + +def test_fresh_output_and_seeded_accumulation_are_documented_separately(): + source = _WGRAD_PATH.read_text() + + assert "Omitting the seed preserves the" in source + assert "simple fresh-output behavior by starting from zero" in source + assert "Pointer-table outputs are not supported" in source diff --git a/test/python/fe_api/test_jax_import_contract.py b/test/python/fe_api/test_jax_import_contract.py index 7bfb03769..f001f206d 100644 --- a/test/python/fe_api/test_jax_import_contract.py +++ b/test/python/fe_api/test_jax_import_contract.py @@ -25,7 +25,11 @@ def _module_name(path): relative = path.relative_to(_CUDNN_ROOT) - parts = relative.parts[:-1] if path.name == "__init__.py" else relative.with_suffix("").parts + parts = ( + relative.parts[:-1] + if path.name == "__init__.py" + else relative.with_suffix("").parts + ) return ".".join(("cudnn", *parts)) @@ -85,7 +89,10 @@ def visit(node, torch_optional=False): visit(child, torch_optional) return if isinstance(node, ast.Try): - catches_import = any(handler.type is not None and _catches_import_error(handler.type) for handler in node.handlers) + catches_import = any( + handler.type is not None and _catches_import_error(handler.type) + for handler in node.handlers + ) for child in node.body: visit(child, torch_optional or catches_import) for handler in node.handlers: @@ -94,7 +101,9 @@ def visit(node, torch_optional=False): for child in (*node.orelse, *node.finalbody): visit(child, torch_optional) return - if not all_scopes and isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)): + if not all_scopes and isinstance( + node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda) + ): return for child in ast.iter_child_nodes(node): visit(child, torch_optional) @@ -111,8 +120,14 @@ def _import_targets(path, node, modules): if node.level: current_module = _module_name(path) - package = current_module if path.name == "__init__.py" else current_module.rpartition(".")[0] - base = importlib.util.resolve_name("." * node.level + (node.module or ""), package) + package = ( + current_module + if path.name == "__init__.py" + else current_module.rpartition(".")[0] + ) + base = importlib.util.resolve_name( + "." * node.level + (node.module or ""), package + ) else: base = node.module or "" @@ -148,51 +163,135 @@ def _jax_import_violations(root, modules): while queue: path, all_scopes, chain = queue.popleft() - if path in scanned_all_scopes or (not all_scopes and path in scanned_module_scope): + if path in scanned_all_scopes or ( + not all_scopes and path in scanned_module_scope + ): continue (scanned_all_scopes if all_scopes else scanned_module_scope).add(path) visited.add(path) for node, torch_optional in _imports(path, all_scopes=all_scopes): for target in _import_targets(path, node, modules): - chain_text = " -> ".join(str(item.relative_to(_REPO_ROOT)) for item in chain) + chain_text = " -> ".join( + str(item.relative_to(_REPO_ROOT)) for item in chain + ) if target == "torch" or target.startswith("torch."): if not torch_optional: - violations.append(f"{chain_text}:{node.lineno} imports {target}") + violations.append( + f"{chain_text}:{node.lineno} imports {target}" + ) continue dependency = modules.get(target) if dependency is None: continue if dependency.name == "api.py": - violations.append(f"{chain_text}:{node.lineno} imports Torch API module {dependency.relative_to(_REPO_ROOT)}") + violations.append( + f"{chain_text}:{node.lineno} imports Torch API module {dependency.relative_to(_REPO_ROOT)}" + ) continue - queue.extend((initializer, False, (*chain, initializer)) for initializer in _package_initializers(dependency)) + queue.extend( + (initializer, False, (*chain, initializer)) + for initializer in _package_initializers(dependency) + ) queue.append((dependency, False, (*chain, dependency))) return visited, violations class JaxImportContractTest(unittest.TestCase): - def test_jax_dsa_facade_matches_the_torch_operation_surface(self): - def literal_assignment(path, name): - tree = ast.parse(path.read_text(), filename=str(path)) - assignment = next( - node - for node in tree.body - if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == name for target in node.targets) + @staticmethod + def _literal_assignment(path, name): + tree = ast.parse(path.read_text(), filename=str(path)) + assignment = next( + node + for node in tree.body + if isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == name + for target in node.targets ) - return ast.literal_eval(assignment.value) + ) + return ast.literal_eval(assignment.value) + + @staticmethod + def _operation_exports(path): + tree = ast.parse(path.read_text(), filename=str(path)) + call = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "make_operation_api" + ) + exports = ast.literal_eval( + next(keyword.value for keyword in call.keywords if keyword.arg == "exports") + ) + return {name for module_exports in exports.values() for name in module_exports} - torch_exports = literal_assignment( + def test_jax_dsa_facade_matches_the_torch_operation_surface(self): + torch_exports = self._literal_assignment( _CUDNN_ROOT / "deepseek_sparse_attention" / "__init__.py", "_SYMBOLS", ) - jax_exports = literal_assignment(_CUDNN_ROOT / "jax" / "__init__.py", "_OPERATION_EXPORTS") - jax_dsa_exports = {name for name, (module_name, _) in jax_exports.items() if module_name.startswith("..deepseek_sparse_attention.")} + jax_exports = self._literal_assignment( + _CUDNN_ROOT / "jax" / "__init__.py", "_OPERATION_EXPORTS" + ) + jax_dsa_exports = { + name + for name, (module_name, _) in jax_exports.items() + if module_name.startswith("..deepseek_sparse_attention.") + } self.assertEqual(jax_dsa_exports, set(torch_exports)) + def test_jax_facade_covers_the_remaining_cute_operation_packages(self): + jax_exports = self._literal_assignment( + _CUDNN_ROOT / "jax" / "__init__.py", "_OPERATION_EXPORTS" + ) + grouped_families = ( + "swiglu", + "dswiglu", + "quant", + "srelu", + "dsrelu", + "glu", + "glu_hadamard", + "dglu", + "wgrad", + ) + grouped_names = { + name + for family in grouped_families + for name in ( + f"GroupedGemm{''.join(part.capitalize() for part in family.split('_'))}Sm100", + f"grouped_gemm_{family}_wrapper_sm100", + ) + } + discrete_names = { + "DiscreteGroupedGemmSwigluSm100", + "discrete_grouped_gemm_swiglu_wrapper_sm100", + "DiscreteGroupedGemmDswigluSm100", + "discrete_grouped_gemm_dswiglu_wrapper_sm100", + } + nsa_names = self._operation_exports( + _CUDNN_ROOT / "native_sparse_attention" / "__init__.py" + ) + sdpa_names = self._operation_exports(_CUDNN_ROOT / "sdpa" / "__init__.py") + bsa_names = { + name + for name in self._literal_assignment( + _CUDNN_ROOT / "block_sparse_attention" / "__init__.py", + "_SYMBOLS", + ) + if not name.endswith("Op") + } + + self.assertLessEqual( + grouped_names | discrete_names | nsa_names | sdpa_names | bsa_names, + set(jax_exports), + ) + def test_jax_import_graph_has_no_required_torch_dependency(self): modules = {_module_name(path): path for path in _CUDNN_ROOT.rglob("*.py")} roots = set(_CUDNN_ROOT.rglob("jax.py")) diff --git a/test/python/fe_api/test_jax_layout.py b/test/python/fe_api/test_jax_layout.py index 9873ea034..1e605a8d8 100644 --- a/test/python/fe_api/test_jax_layout.py +++ b/test/python/fe_api/test_jax_layout.py @@ -15,7 +15,9 @@ pytestmark = pytest.mark.L0 -_LAYOUT_FILE = Path(__file__).resolve().parents[3] / "python" / "cudnn" / "_jax" / "layout.py" +_LAYOUT_FILE = ( + Path(__file__).resolve().parents[3] / "python" / "cudnn" / "_jax" / "layout.py" +) def _load_layout_module(): @@ -119,6 +121,33 @@ def test_cutlass_layout_is_indexed_by_public_axes(self): (2, 1, 0), ) + def test_canonical_stride_order_maps_to_named_public_layout(self): + bhsd_to_bshd = self.layout.mode_from_layout("BHSD", kernel_axes="BSHD") + bshd_to_bshd = self.layout.mode_from_layout("BSHD", kernel_axes="BSHD") + + self.assertEqual( + self.layout.stride_order_to_public((3, 2, 1, 0), bhsd_to_bshd), + (3, 1, 2, 0), + ) + self.assertEqual( + self.layout.stride_order_to_public((3, 2, 1, 0), bshd_to_bshd), + (3, 2, 1, 0), + ) + + bshd_to_bhsd = self.layout.mode_from_layout("BSHD", kernel_axes="BHSD") + self.assertEqual( + self.layout.to_public_axes((2, 8, 32, 16), bshd_to_bhsd), + (2, 32, 8, 16), + ) + self.assertEqual( + self.layout.stride_order_to_public((3, 1, 2, 0), bshd_to_bhsd), + (3, 2, 1, 0), + ) + + def test_canonical_stride_order_requires_a_permutation(self): + with self.assertRaisesRegex(ValueError, "must be a permutation"): + self.layout.stride_order_to_public((2, 2, 0)) + def test_invalid_mode_is_rejected(self): with self.assertRaisesRegex(ValueError, "mode must be a permutation"): self.layout.normalize_mode(3, (0, 0, 2)) diff --git a/test/python/fe_api/test_jax_nsa_sdpa_source_contract.py b/test/python/fe_api/test_jax_nsa_sdpa_source_contract.py new file mode 100644 index 000000000..eb2eb495f --- /dev/null +++ b/test/python/fe_api/test_jax_nsa_sdpa_source_contract.py @@ -0,0 +1,355 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Source contracts shared by the JAX NSA and d=256 SDPA adapters.""" + +from __future__ import annotations + +import ast +from pathlib import Path +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_ADAPTERS = { + "selection": _CUDNN_ROOT / "native_sparse_attention" / "selection" / "jax.py", + "compression": _CUDNN_ROOT / "native_sparse_attention" / "compression" / "jax.py", + "sliding": _CUDNN_ROOT + / "native_sparse_attention" + / "sliding_window_attention" + / "jax.py", + "top_k": _CUDNN_ROOT / "native_sparse_attention" / "top_k" / "jax.py", + "sdpa_fwd": _CUDNN_ROOT / "sdpa" / "fwd" / "jax.py", + "sdpa_bwd": _CUDNN_ROOT / "sdpa" / "bwd" / "jax.py", +} + + +def _tree(path: Path) -> ast.Module: + return ast.parse(path.read_text(), filename=str(path)) + + +def _class(tree: ast.Module, name: str) -> ast.ClassDef: + return next( + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == name + ) + + +def _method(class_node: ast.ClassDef, name: str) -> ast.FunctionDef: + return next( + node + for node in class_node.body + if isinstance(node, ast.FunctionDef) and node.name == name + ) + + +class JaxNsaSdpaSourceContractTest(unittest.TestCase): + def test_layout_helpers_separate_logical_modes_from_physical_storage(self): + nsa_helpers = ( + _CUDNN_ROOT / "native_sparse_attention" / "jax_utils.py" + ).read_text() + sdpa_helpers = (_CUDNN_ROOT / "sdpa" / "jax_utils.py").read_text() + + self.assertNotIn("BHSD_STORAGE_ORDER", nsa_helpers) + self.assertNotIn("BHSD_STORAGE_ORDER", sdpa_helpers) + self.assertIn("mode_from_layout(layout, kernel_axes=kernel_axes)", nsa_helpers) + self.assertIn("kernel_stride_order is None", nsa_helpers) + self.assertIn('KERNEL_AXES = "BSHD"', sdpa_helpers) + self.assertIn("KERNEL_STRIDE_ORDER = (3, 2, 1, 0)", sdpa_helpers) + + def test_all_adapters_are_torch_independent_jax_api_classes(self): + expected_classes = { + "selection": "SelectionAttention", + "compression": "CompressionAttention", + "sliding": "SlidingWindowAttention", + "top_k": "TopKReduction", + "sdpa_fwd": "SdpafwdSm100D256", + "sdpa_bwd": "SdpabwdSm100D256", + } + for operation, path in _ADAPTERS.items(): + with self.subTest(operation=operation): + tree = _tree(path) + imports = [ + node + for node in ast.walk(tree) + if isinstance(node, (ast.Import, ast.ImportFrom)) + ] + self.assertFalse( + any( + ( + isinstance(node, ast.Import) + and any(alias.name == "torch" for alias in node.names) + ) + or (isinstance(node, ast.ImportFrom) and node.module == "torch") + for node in imports + ) + ) + adapter = _class(tree, expected_classes[operation]) + self.assertEqual( + [base.id for base in adapter.bases if isinstance(base, ast.Name)], + ["JaxApiBase"], + ) + self.assertTrue( + any( + isinstance(node, ast.FunctionDef) + and node.name == "check_support" + for node in adapter.body + ) + ) + self.assertTrue( + any( + isinstance(node, ast.FunctionDef) and node.name == "__call__" + for node in adapter.body + ) + ) + + def test_cute_launchers_follow_stream_inputs_outputs_workspaces_order(self): + expected = { + "selection": ( + "stream", + "q", + "k", + "v", + "block_indices", + "block_counts", + "cum_seqlen", + "output", + "lse_sum", + "row_max", + ), + "compression": ("stream", "q", "k", "v", "output"), + "top_k": ("stream", "q", "k", "lse", "topk_scores", "topk_indices"), + "sdpa_fwd": ("stream", "q", "k", "v", "output", "lse"), + "sdpa_bwd": ( + "stream", + "q", + "k", + "v", + "output", + "doutput", + "lse", + "dq", + "dk", + "dv", + "workspace", + ), + } + class_names = { + "selection": "SelectionAttention", + "compression": "CompressionAttention", + "top_k": "TopKReduction", + "sdpa_fwd": "SdpafwdSm100D256", + "sdpa_bwd": "SdpabwdSm100D256", + } + for operation, argument_names in expected.items(): + with self.subTest(operation=operation): + launcher = _method( + _class(_tree(_ADAPTERS[operation]), class_names[operation]), + "_launch_kernel", + ) + actual = tuple(argument.arg for argument in launcher.args.args[1:]) + self.assertEqual(actual, argument_names) + + def test_functional_buffers_preserve_required_initial_values(self): + expected_literals = { + "selection": {"0", "0.0", "float('-inf')"}, + "top_k": {"float('-inf')", "-1"}, + "sdpa_bwd": {"0"}, + } + for operation, expected in expected_literals.items(): + with self.subTest(operation=operation): + values = set() + for call in ( + node + for node in ast.walk(_tree(_ADAPTERS[operation])) + if isinstance(node, ast.Call) + ): + for keyword in call.keywords: + if keyword.arg == "init_value": + values.add(ast.unparse(keyword.value)) + self.assertTrue( + expected.issubset(values), + f"{operation} initialized buffers changed: {values}", + ) + + def test_selection_uses_one_self_attention_offset_operand(self): + tree = _tree(_ADAPTERS["selection"]) + wrapper = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) + and node.name == "selection_attention_wrapper" + ) + argument_names = tuple(argument.arg for argument in wrapper.args.args) + + self.assertIn("cum_seqlen_tensor", argument_names) + self.assertNotIn("cum_seqlen_q_tensor", argument_names) + self.assertNotIn("cum_seqlen_k_tensor", argument_names) + self.assertIn("(1, 2, 4, 8, 16)", _ADAPTERS["selection"].read_text()) + + def test_persistent_compression_resolves_hardware_before_lowering(self): + tree = _tree(_ADAPTERS["compression"]) + adapter = _class(tree, "CompressionAttention") + check_support = _method(adapter, "check_support") + runner = _method(adapter, "_run_kernel") + + self.assertIn( + "self._get_device_multiprocessor_count()", + ast.unparse(check_support), + ) + self.assertIn("persistent_sm_count", ast.unparse(runner)) + self.assertNotIn("HardwareInfo", _ADAPTERS["compression"].read_text()) + + def test_compression_supports_fixed_bhsd_and_packed_thd(self): + tree = _tree(_ADAPTERS["compression"]) + wrapper = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) + and node.name == "compression_attention_wrapper" + ) + argument_names = { + argument.arg for argument in (*wrapper.args.args, *wrapper.args.kwonlyargs) + } + varlen_launcher = _method( + _class(tree, "CompressionAttention"), "_launch_varlen_kernel" + ) + + self.assertTrue( + { + "cum_seqlen_q_tensor", + "cum_seqlen_k_tensor", + "max_s_q", + "max_s_k", + }.issubset(argument_names) + ) + self.assertEqual( + tuple(argument.arg for argument in varlen_launcher.args.args[1:]), + ( + "stream", + "q", + "k", + "v", + "cum_seqlen_q", + "cum_seqlen_k", + "output", + ), + ) + source = _ADAPTERS["compression"].read_text() + self.assertIn("normalize_attention_layout(layout, ranks[0])", source) + self.assertIn("self.input_layout in FIXED_LAYOUTS", source) + self.assertIn("jnp.float8_e4m3fn", source) + self.assertIn("self.lse_extent = total_q", source) + + def test_fixed_adapters_expose_named_bhsd_and_bshd_layouts(self): + for operation in ("compression", "sliding", "top_k", "sdpa_fwd", "sdpa_bwd"): + source = _ADAPTERS[operation].read_text() + with self.subTest(operation=operation): + self.assertIn("layout: str | None", source) + self.assertIn('"layout",', source) + + def test_compression_packed_views_duplicate_the_token_stride(self): + tree = _tree(_ADAPTERS["compression"]) + packed_orders = [ + ast.literal_eval(keyword.value) + for call in ast.walk(tree) + if isinstance(call, ast.Call) + for keyword in call.keywords + if keyword.arg == "public_stride_order" + and ast.unparse(keyword.value) == "(3, 2, 0, 1)" + ] + + self.assertEqual(len(packed_orders), 4) + token_count, heads, head_dim = 20, 8, 128 + shape = (1, token_count, heads, head_dim) + running = 1 + strides = [0] * len(shape) + for axis in packed_orders[0]: + strides[axis] = running + running *= max(shape[axis], 1) + self.assertEqual( + tuple(strides), + (heads * head_dim, heads * head_dim, head_dim, 1), + ) + + def test_sliding_window_supports_stats_two_sided_windows_and_packed_thd(self): + source = _ADAPTERS["sliding"].read_text() + + self.assertIn("return_residual=not self.is_infer", source) + self.assertIn("(self.left_bound - 1, self.right_bound)", source) + self.assertIn("query_seq_lengths=query_lengths", source) + self.assertIn("key_value_seq_lengths=key_value_lengths", source) + self.assertIn("jnp.diff(cum_seqlen_q_tensor)", source) + self.assertIn("self._pad_packed(", source) + self.assertIn("self._unpad_packed(", source) + self.assertNotIn("currently supports inference only", source) + self.assertNotIn("currently requires S_q == S_k", source) + + def test_cutlass_kernels_are_deferred_until_adapter_calls(self): + kernel_modules = { + "selection": "NSA_select_attn_fwd_hmma", + "compression": "fmha", + "top_k": "nsa_top_k_reduction_fwd", + "sdpa_fwd": "fmha_forward_sm100_d256", + "sdpa_bwd": "fmha_backward_sm100_2kernel", + } + for operation, module_name in kernel_modules.items(): + with self.subTest(operation=operation): + tree = _tree(_ADAPTERS[operation]) + top_level_imports = [ + node + for node in tree.body + if isinstance(node, (ast.Import, ast.ImportFrom)) + ] + self.assertFalse( + any( + isinstance(node, ast.ImportFrom) + and (node.module or "").endswith(module_name) + for node in top_level_imports + ) + ) + self.assertTrue( + any( + isinstance(node, ast.ImportFrom) + and (node.module or "").endswith(module_name) + for node in ast.walk(tree) + ) + ) + + def test_package_initializers_keep_torch_apis_lazy(self): + initializers = [ + _CUDNN_ROOT / "native_sparse_attention" / "__init__.py", + *( + _CUDNN_ROOT / "native_sparse_attention" / name / "__init__.py" + for name in ( + "selection", + "compression", + "sliding_window_attention", + "top_k", + ) + ), + _CUDNN_ROOT / "sdpa" / "__init__.py", + _CUDNN_ROOT / "sdpa" / "fwd" / "__init__.py", + _CUDNN_ROOT / "sdpa" / "bwd" / "__init__.py", + ] + for path in initializers: + with self.subTest(path=path.relative_to(_CUDNN_ROOT)): + direct_api_imports = [ + node + for node in _tree(path).body + if isinstance(node, ast.ImportFrom) + and node.module in {"api", ".api"} + ] + self.assertEqual(direct_api_imports, []) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_nsa_topk_thd_contract.py b/test/python/fe_api/test_jax_nsa_topk_thd_contract.py new file mode 100644 index 000000000..8d7ed478e --- /dev/null +++ b/test/python/fe_api/test_jax_nsa_topk_thd_contract.py @@ -0,0 +1,400 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for packed JAX NSA top-K reduction.""" + +from __future__ import annotations + +import ast +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_NSA_ROOT = _CUDNN_ROOT / "native_sparse_attention" +_TOPK_ROOT = _NSA_ROOT / "top_k" +_PACKAGE = "cudnn_jax_nsa_topk_thd_contract_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + INT64 = auto() + + +_DTYPE_TO_CUDNN = { + "float16": _DataType.HALF, + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, + "int32": _DataType.INT32, + "int64": _DataType.INT64, +} +_CUDNN_TO_DTYPE = {value: key for key, value in _DTYPE_TO_CUDNN.items()} + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + def __getitem__(self, key): + if key == (None, Ellipsis): + return _Array((1, *self.shape), self.dtype) + raise TypeError(f"unsupported test index {key!r}") + + +class _TensorSpec: + def __init__(self, *, layout, mode, divisibility=None): + self.layout = layout + self.mode = mode + self.divisibility = divisibility + + +class JaxNsaTopkThdContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + nsa_name = f"{_PACKAGE}.native_sparse_attention" + nsa = types.ModuleType(nsa_name) + nsa.__path__ = [str(_NSA_ROOT)] + nsa.__package__ = nsa_name + nsa.__spec__ = ModuleSpec(nsa_name, loader=None, is_package=True) + sys.modules[nsa_name] = nsa + + topk_name = f"{nsa_name}.top_k" + topk = types.ModuleType(topk_name) + topk.__path__ = [str(_TOPK_ROOT)] + topk.__package__ = topk_name + topk.__spec__ = ModuleSpec(topk_name, loader=None, is_package=True) + sys.modules[topk_name] = topk + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + layout_module = importlib.import_module(f"{internal_name}.layout") + result_module = importlib.import_module(f"{_PACKAGE}._result") + + class JaxTensorDesc(tensor_module.TensorDesc): + @property + def cudnn_dtype(self): + return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) + + class JaxApiBase: + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del operation_name + resolved = 100 if target is None else target + if resolved not in supported: + raise ValueError(f"unsupported target {resolved}") + return resolved + + @staticmethod + def _to_tensor_desc( + value, + name, + *, + mode=None, + public_stride_order=None, + init_value=None, + ): + public_shape = tuple(value.shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + if public_stride_order is None: + public_stride_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride( + public_shape, tuple(public_stride_order) + ) + canonical_axis_by_public_axis = layout_module.to_public_axes( + tuple(range(len(public_shape))), mode + ) + return JaxTensorDesc( + dtype=value.dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple( + canonical_axis_by_public_axis[axis] + for axis in public_stride_order + ), + name=name, + init_value=init_value, + ) + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + actual_shape = layout_module.to_canonical_axes(tuple(value.shape), mode) + if actual_shape != expected.shape: + raise ValueError(f"{expected.name} shape mismatch") + if ( + _DTYPE_TO_CUDNN.get(value.dtype, _DataType.NOT_SET) + != expected.cudnn_dtype + ): + raise ValueError(f"{expected.name} dtype mismatch") + + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + mode = layout_module.normalize_mode(desc.ndim, mode) + return _TensorSpec( + layout=layout_module.to_cutlass_layout( + desc.shape, + desc.stride, + desc.stride_order, + mode=mode, + name=desc.name, + ), + mode=mode, + divisibility=divisibility, + ) + + def _call_kernel( + self, + inputs, + *, + launch, + output_descs, + input_spec, + output_spec, + **options, + ): + self.captured_call = { + "inputs": tuple(inputs), + "launch": launch, + "output_descs": tuple(output_descs), + "input_spec": tuple(input_spec), + "output_spec": tuple(output_spec), + **options, + } + return tuple( + _Array(desc.shape, _CUDNN_TO_DTYPE[desc.cudnn_dtype]) + for desc in output_descs + ) + + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = result_module.TupleDict + + datatypes = types.ModuleType(f"{internal_name}.datatypes") + datatypes.normalize_jax_dtype = lambda value, default, _name: ( + default if value is None else value + ) + sys.modules[datatypes.__name__] = datatypes + + fake_jax = types.ModuleType("jax") + fake_jax.__path__ = [] + fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + fake_jax.ShapeDtypeStruct = _Array + cls.static_argnames = {} + + def jit(function=None, *, static_argnames=()): + def decorate(target): + cls.static_argnames[target.__name__] = tuple(static_argnames) + return target + + return decorate if function is None else decorate(function) + + fake_jax.jit = jit + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.float16 = "float16" + fake_jnp.bfloat16 = "bfloat16" + fake_jnp.float32 = "float32" + fake_jnp.int32 = "int32" + fake_jnp.dtype = lambda value: value + fake_jnp.reshape = lambda value, shape: _Array(shape, value.dtype) + fake_jnp.transpose = lambda value, axes: _Array( + tuple(value.shape[axis] for axis in axes), value.dtype + ) + fake_jax.numpy = fake_jnp + + try: + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "torch": None, + "cutlass": None, + }, + ): + cls.module = importlib.import_module(f"{topk_name}.jax") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + @staticmethod + def _packed_samples(cum_dtype="int32"): + return ( + _Array((20, 8, 128), "bfloat16"), + _Array((18, 2, 128), "bfloat16"), + _Array((20, 8, 1), "float32"), + _Array((3,), cum_dtype), + _Array((3,), cum_dtype), + ) + + def test_packed_thd_uses_explicit_offsets_and_virtual_kernel_views(self): + q, k, lse, cum_q, cum_k = self._packed_samples() + api = self.module.TopKReduction( + q, + k, + lse, + cum_q, + cum_k, + max_s_q=12, + max_s_k=10, + layout="T,H,D", + target_compute_capability=100, + ) + result = api(q, k, lse, cum_q, cum_k) + + self.assertEqual(api.input_layout, "THD") + self.assertEqual(api.batch, 2) + self.assertEqual(api.q_kernel_desc.shape, (1, 8, 20, 128)) + self.assertEqual(api.k_kernel_desc.shape, (1, 2, 18, 128)) + self.assertEqual(api.q_kernel_desc.stride_order, (3, 1, 2, 0)) + self.assertEqual(api.lse_kernel_desc.shape, (1, 8, 20)) + self.assertEqual( + tuple(value.shape for value in api.captured_call["inputs"]), + ((1, 8, 20, 128), (1, 2, 18, 128), (1, 8, 20), (3,), (3,)), + ) + self.assertEqual( + tuple(spec.layout for spec in api.captured_call["input_spec"]), + ((3, 1, 2, 0), (3, 1, 2, 0), (2, 1, 0), (0,), (0,)), + ) + self.assertEqual(api.captured_call["launch"].__name__, "_launch_packed_kernel") + self.assertEqual(result["topk_scores_tensor"].shape, (20, 2, 16)) + self.assertEqual(result["topk_indices_tensor"].shape, (20, 2, 16)) + self.assertEqual(api.scores_desc.init_value, float("-inf")) + self.assertEqual(api.indices_desc.init_value, -1) + + def test_fixed_bhsd_contract_is_preserved(self): + q = _Array((2, 8, 20, 128), "float16") + k = _Array((2, 2, 18, 128), "float16") + lse = _Array((2, 8, 20), "float32") + api = self.module.TopKReduction( + q, + k, + lse, + target_compute_capability=100, + ) + result = api(q, k, lse) + + self.assertEqual(api.input_layout, "BHSD") + self.assertEqual(len(api.captured_call["inputs"]), 3) + self.assertEqual(api.captured_call["launch"].__name__, "_launch_kernel") + self.assertEqual(api.captured_call["input_spec"][0].layout, (3, 1, 2, 0)) + self.assertEqual(result["topk_scores_tensor"].shape, (2, 2, 20, 16)) + self.assertEqual(result["topk_indices_tensor"].shape, (2, 2, 20, 16)) + + def test_packed_metadata_is_rejected_before_lowering(self): + q, k, lse, cum_q, cum_k = self._packed_samples() + with self.assertRaisesRegex(ValueError, "max_s_q and max_s_k"): + self.module.TopKReduction(q, k, lse, cum_q, cum_k) + with self.assertRaisesRegex(ValueError, "both required"): + self.module.TopKReduction( + q, + k, + lse, + cum_q, + None, + max_s_q=12, + max_s_k=10, + ) + + _, _, _, bad_cum_q, bad_cum_k = self._packed_samples("int64") + with self.assertRaisesRegex(ValueError, "dtype int32"): + self.module.TopKReduction( + q, + k, + lse, + bad_cum_q, + bad_cum_k, + max_s_q=12, + max_s_k=10, + ) + with self.assertRaisesRegex(ValueError, "BHSD layout requires rank-4"): + self.module.TopKReduction(q, k, lse, layout="BHSD") + + def test_wrapper_marks_packed_problem_extents_and_layout_static(self): + static = self.static_argnames["topk_reduction_wrapper"] + self.assertIn("max_s_q", static) + self.assertIn("max_s_k", static) + self.assertIn("layout", static) + + def test_packed_launcher_follows_stream_inputs_outputs_order(self): + path = _TOPK_ROOT / "jax.py" + tree = ast.parse(path.read_text(), filename=str(path)) + adapter = next( + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == "TopKReduction" + ) + launcher = next( + node + for node in adapter.body + if isinstance(node, ast.FunctionDef) + and node.name == "_launch_packed_kernel" + ) + arguments = tuple(argument.arg for argument in launcher.args.args[1:]) + self.assertEqual( + arguments, + ( + "stream", + "q", + "k", + "lse", + "cum_seqlen_q", + "cum_seqlen_k", + "topk_scores", + "topk_indices", + ), + ) + imports = [ + node + for node in ast.walk(tree) + if isinstance(node, (ast.Import, ast.ImportFrom)) + ] + self.assertFalse( + any( + ( + isinstance(node, ast.Import) + and any(alias.name == "torch" for alias in node.names) + ) + or (isinstance(node, ast.ImportFrom) and node.module == "torch") + for node in imports + ) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_sdpa_thd_contract.py b/test/python/fe_api/test_jax_sdpa_thd_contract.py new file mode 100644 index 000000000..417299148 --- /dev/null +++ b/test/python/fe_api/test_jax_sdpa_thd_contract.py @@ -0,0 +1,184 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free source contracts for packed-THD JAX SDPA.""" + +from __future__ import annotations + +import ast +from pathlib import Path +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_FORWARD = _CUDNN_ROOT / "sdpa" / "fwd" / "jax.py" +_BACKWARD = _CUDNN_ROOT / "sdpa" / "bwd" / "jax.py" + + +def _tree(path: Path) -> ast.Module: + return ast.parse(path.read_text(), filename=str(path)) + + +def _class(tree: ast.Module, name: str) -> ast.ClassDef: + return next( + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == name + ) + + +def _function(nodes: list[ast.stmt], name: str) -> ast.FunctionDef: + return next( + node + for node in nodes + if isinstance(node, ast.FunctionDef) and node.name == name + ) + + +def _argument_names(function: ast.FunctionDef) -> tuple[str, ...]: + return tuple(argument.arg for argument in function.args.args) + + +def _jit_static_argnames(function: ast.FunctionDef) -> tuple[str, ...]: + decorator = next( + node + for node in function.decorator_list + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "partial" + ) + keyword = next( + keyword for keyword in decorator.keywords if keyword.arg == "static_argnames" + ) + return ast.literal_eval(keyword.value) + + +class JaxSdpaThdContractTest(unittest.TestCase): + def test_constructors_bind_cumulative_metadata_and_static_bounds(self): + for path, class_name in ( + (_FORWARD, "SdpafwdSm100D256"), + (_BACKWARD, "SdpabwdSm100D256"), + ): + constructor = _function(_class(_tree(path), class_name).body, "__init__") + arguments = _argument_names(constructor) + with self.subTest(path=path.relative_to(_CUDNN_ROOT)): + self.assertIn("sample_cum_seqlen_q", arguments) + self.assertIn("sample_cum_seqlen_k", arguments) + self.assertIn("max_s_q", arguments) + self.assertIn("max_s_k", arguments) + + def test_functional_wrappers_keep_offsets_dynamic_and_bounds_static(self): + for path, wrapper_name in ( + (_FORWARD, "sdpa_fwd_wrapper_sm100_d256"), + (_BACKWARD, "sdpa_bwd_wrapper_sm100_d256"), + ): + wrapper = _function(_tree(path).body, wrapper_name) + arguments = _argument_names(wrapper) + static_arguments = _jit_static_argnames(wrapper) + with self.subTest(path=path.relative_to(_CUDNN_ROOT)): + self.assertIn("cum_seqlen_q_tensor", arguments) + self.assertIn("cum_seqlen_k_tensor", arguments) + self.assertNotIn("cum_seqlen_q_tensor", static_arguments) + self.assertNotIn("cum_seqlen_k_tensor", static_arguments) + self.assertIn("max_s_q", static_arguments) + self.assertIn("max_s_k", static_arguments) + + def test_varlen_launchers_follow_stream_inputs_outputs_workspaces_order(self): + expected = { + _FORWARD: ( + "stream", + "q", + "k", + "v", + "cum_seqlen_q", + "cum_seqlen_k", + "output", + "lse", + ), + _BACKWARD: ( + "stream", + "q", + "k", + "v", + "output", + "doutput", + "lse", + "cum_seqlen_q", + "cum_seqlen_k", + "dq", + "dk", + "dv", + "workspace", + ), + } + class_names = { + _FORWARD: "SdpafwdSm100D256", + _BACKWARD: "SdpabwdSm100D256", + } + for path, arguments in expected.items(): + launcher = _function( + _class(_tree(path), class_names[path]).body, + "_launch_varlen_kernel", + ) + with self.subTest(path=path.relative_to(_CUDNN_ROOT)): + self.assertEqual(_argument_names(launcher)[1:], arguments) + + def test_thd_is_promoted_without_changing_public_result_ranks(self): + forward = _FORWARD.read_text() + backward = _BACKWARD.read_text() + + for source in (forward, backward): + self.assertIn("jax.ShapeDtypeStruct((1, *sample_q.shape)", source) + self.assertIn("jnp.reshape(q_tensor, self.q_kernel_desc.shape)", source) + self.assertIn("max_s_q and max_s_k are both required for THD", source) + + self.assertIn("jnp.reshape(output, self.o_desc.shape)", forward) + self.assertIn("(self.total_q_tokens, self.num_query_heads)", forward) + self.assertIn("public_stride_order=(0, 1)", forward) + for gradient in ("dq", "dk", "dv"): + self.assertIn( + f"{gradient} = jnp.reshape({gradient}, self.{gradient}_desc.shape)", + backward, + ) + self.assertIn("(1, self.num_query_heads, self.total_q_tokens)", backward) + self.assertNotIn("public_stride_order=(1, 2, 0)", backward) + + def test_backward_workspace_is_initialized_and_adapters_do_not_import_torch(self): + backward_tree = _tree(_BACKWARD) + init_values = { + ast.unparse(keyword.value) + for call in ast.walk(backward_tree) + if isinstance(call, ast.Call) + for keyword in call.keywords + if keyword.arg == "init_value" + } + self.assertIn("0", init_values) + + for path in (_FORWARD, _BACKWARD): + imports = [ + node + for node in ast.walk(_tree(path)) + if isinstance(node, (ast.Import, ast.ImportFrom)) + ] + with self.subTest(path=path.relative_to(_CUDNN_ROOT)): + self.assertFalse( + any( + ( + isinstance(node, ast.Import) + and any(alias.name == "torch" for alias in node.names) + ) + or (isinstance(node, ast.ImportFrom) and node.module == "torch") + for node in imports + ) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_sliding_packed_contract.py b/test/python/fe_api/test_jax_sliding_packed_contract.py new file mode 100644 index 000000000..781053a0d --- /dev/null +++ b/test/python/fe_api/test_jax_sliding_packed_contract.py @@ -0,0 +1,64 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Executable contracts for packed-THD sliding-window repacking.""" + +from __future__ import annotations + +import ast +from pathlib import Path +from typing import Any + +import pytest + +pytestmark = pytest.mark.L0 + +jax = pytest.importorskip("jax") +jnp = pytest.importorskip("jax.numpy") + +_ADAPTER = ( + Path(__file__).resolve().parents[3] + / "python" + / "cudnn" + / "native_sparse_attention" + / "sliding_window_attention" + / "jax.py" +) + + +def _load_static_method(name: str): + tree = ast.parse(_ADAPTER.read_text(), filename=str(_ADAPTER)) + class_node = next( + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == "SlidingWindowAttention" + ) + method = next( + node + for node in class_node.body + if isinstance(node, ast.FunctionDef) and node.name == name + ) + method.decorator_list = [] + module = ast.fix_missing_locations(ast.Module(body=[method], type_ignores=[])) + namespace = {"Any": Any, "jax": jax, "jnp": jnp} + exec(compile(module, str(_ADAPTER), "exec"), namespace) + return namespace[name] + + +def test_packed_padding_and_gather_round_trip_under_jit(): + pad_packed = _load_static_method("_pad_packed") + unpad_packed = _load_static_method("_unpad_packed") + values = jnp.arange(5, dtype=jnp.float32).reshape(5, 1, 1) + cumulative = jnp.array([0, 2, 5], dtype=jnp.int32) + + padded = jax.jit(lambda x, offsets: pad_packed(x, offsets[:-1], 3))( + values, cumulative + ) + assert padded.shape == (2, 3, 1, 1) + assert padded[:, :, 0, 0].tolist() == [[0.0, 1.0, 2.0], [2.0, 3.0, 4.0]] + + restored = jax.jit(lambda x, offsets: unpad_packed(x, offsets, 5))( + padded, cumulative + ) + assert restored.shape == values.shape + assert restored[:, 0, 0].tolist() == values[:, 0, 0].tolist() From 8386cb6c8a53448876f53dd41bf5e097772c46fe Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Wed, 8 Jul 2026 12:30:37 -0500 Subject: [PATCH 13/13] Refactor JAX kernel calls around tensor descriptors Move framework-neutral helpers into common packages, derive CUTLASS lowering metadata and output allocation from JAX tensor descriptors, and harmonize JAX operator wrappers across GEMM and attention kernels. --- python/cudnn/__init__.py | 4 +- python/cudnn/_jax/api_base.py | 394 ++++++++++++------ python/cudnn/_jax/compiler.py | 20 + python/cudnn/_jax/gemm.py | 2 +- python/cudnn/_jax/layout.py | 2 +- python/cudnn/_jax/result.py | 2 +- python/cudnn/api_base.py | 4 +- python/cudnn/block_sparse_attention/jax.py | 145 +++---- python/cudnn/block_sparse_attention/op.py | 4 +- python/cudnn/common/__init__.py | 4 + .../cute_arch.py} | 16 +- python/cudnn/{_op.py => common/op.py} | 0 .../operation_api.py} | 0 python/cudnn/{_result.py => common/result.py} | 0 .../tensor_desc.py} | 2 +- .../indexer_backward/__init__.py | 2 +- .../indexer_backward/api.py | 2 +- .../indexer_backward/jax.py | 184 ++++---- .../indexer_backward/op.py | 4 +- .../indexer_forward/__init__.py | 2 +- .../indexer_forward/jax.py | 65 +-- .../indexer_forward/op.py | 14 +- .../indexer_top_k/__init__.py | 2 +- .../indexer_top_k/jax.py | 28 +- .../indexer_top_k/op.py | 4 +- .../score_recompute/__init__.py | 2 +- .../score_recompute/jax.py | 169 +++----- .../score_recompute/op.py | 4 +- .../sparse_score_recompute_sm100.py | 1 + .../sparse_attention_backward/__init__.py | 2 +- .../sparse_attention_backward/jax.py | 77 ++-- .../sparse_attention_backward/op.py | 4 +- .../utils/compiler.py | 2 +- .../deepseek_sparse_attention/utils/copy.py | 29 ++ .../cudnn/discrete_grouped_gemm/__init__.py | 2 +- .../discrete_grouped_gemm/_jax_common.py | 39 +- .../discrete_grouped_gemm_dswiglu/__init__.py | 2 +- .../discrete_grouped_gemm_dswiglu/jax.py | 96 ++--- .../discrete_grouped_gemm_swiglu/__init__.py | 2 +- .../discrete_grouped_gemm_swiglu/jax.py | 116 ++---- python/cudnn/gemm/__init__.py | 4 + .../cudnn/{_dense_gemm.py => gemm/helpers.py} | 4 +- python/cudnn/{_gemm_relu.py => gemm/srelu.py} | 13 +- python/cudnn/gemm_amax/__init__.py | 2 +- python/cudnn/gemm_amax/jax.py | 89 ++-- python/cudnn/gemm_amax/op.py | 6 +- python/cudnn/gemm_dsrelu/__init__.py | 2 +- python/cudnn/gemm_dsrelu/jax.py | 128 ++---- python/cudnn/gemm_dsrelu/op.py | 4 +- python/cudnn/gemm_srelu/__init__.py | 2 +- python/cudnn/gemm_srelu/jax.py | 124 ++---- python/cudnn/gemm_srelu/op.py | 4 +- python/cudnn/gemm_swiglu/__init__.py | 2 +- python/cudnn/gemm_swiglu/jax.py | 134 ++---- python/cudnn/gemm_swiglu/op.py | 6 +- python/cudnn/grouped_gemm/__init__.py | 2 +- python/cudnn/grouped_gemm/_jax_api.py | 262 ++++++------ .../grouped_gemm_dglu/__init__.py | 2 +- .../grouped_gemm/grouped_gemm_dglu/jax.py | 80 ++-- .../grouped_gemm_dsrelu/__init__.py | 2 +- .../grouped_gemm/grouped_gemm_dsrelu/jax.py | 82 ++-- .../grouped_gemm_dswiglu/__init__.py | 2 +- .../grouped_gemm/grouped_gemm_dswiglu/jax.py | 76 ++-- .../grouped_gemm/grouped_gemm_glu/__init__.py | 2 +- .../grouped_gemm/grouped_gemm_glu/jax.py | 75 ++-- .../grouped_gemm_glu_hadamard/__init__.py | 2 +- .../grouped_gemm_glu_hadamard/jax.py | 72 ++-- .../grouped_gemm_quant/__init__.py | 2 +- .../grouped_gemm/grouped_gemm_quant/jax.py | 78 ++-- .../grouped_gemm_srelu/__init__.py | 2 +- .../grouped_gemm/grouped_gemm_srelu/jax.py | 78 ++-- .../grouped_gemm_swiglu/__init__.py | 2 +- .../grouped_gemm/grouped_gemm_swiglu/jax.py | 71 ++-- .../grouped_gemm_wgrad/__init__.py | 2 +- .../grouped_gemm/grouped_gemm_wgrad/jax.py | 169 ++++---- .../cudnn/native_sparse_attention/__init__.py | 2 +- .../compression/__init__.py | 2 +- .../compression/jax.py | 133 +++--- .../native_sparse_attention/jax_utils.py | 19 +- .../selection/__init__.py | 2 +- .../native_sparse_attention/selection/jax.py | 73 ++-- .../sliding_window_attention/__init__.py | 2 +- .../sliding_window_attention/jax.py | 24 +- .../native_sparse_attention/top_k/__init__.py | 2 +- .../native_sparse_attention/top_k/jax.py | 108 ++--- python/cudnn/rmsnorm_rht_amax/__init__.py | 2 +- python/cudnn/rmsnorm_rht_amax/jax.py | 25 +- python/cudnn/rmsnorm_rht_amax/kernel.py | 30 ++ python/cudnn/rmsnorm_rht_amax/op.py | 4 +- python/cudnn/sdpa/__init__.py | 2 +- python/cudnn/sdpa/bwd/__init__.py | 2 +- python/cudnn/sdpa/bwd/jax.py | 152 +++---- python/cudnn/sdpa/fwd/__init__.py | 2 +- python/cudnn/sdpa/fwd/jax.py | 122 +++--- python/cudnn/sdpa/jax_utils.py | 12 +- test/python/fe_api/test_dense_gemm_helpers.py | 8 +- test/python/fe_api/test_dsa_compiler.py | 52 --- .../fe_api/test_dsa_indexer_backward_op.py | 6 +- test/python/fe_api/test_dsa_indexer_ops.py | 64 ++- .../fe_api/test_dsa_score_recompute_op.py | 4 +- .../test_dsa_sparse_attention_backward_op.py | 6 +- test/python/fe_api/test_gemm_amax_op.py | 6 +- test/python/fe_api/test_gemm_relu_imports.py | 4 +- test/python/fe_api/test_gemm_swiglu_op.py | 6 +- test/python/fe_api/test_jax_api_base.py | 194 +++++---- test/python/fe_api/test_jax_bsa_contract.py | 67 ++- test/python/fe_api/test_jax_compiler.py | 74 ++++ test/python/fe_api/test_jax_datatypes.py | 2 +- ...test_jax_discrete_grouped_gemm_contract.py | 75 +++- .../test_jax_dsa_indexer_backward_contract.py | 86 +++- test/python/fe_api/test_jax_dsa_indexers.py | 22 +- .../test_jax_dsa_score_recompute_contract.py | 58 ++- ..._dsa_sparse_attention_backward_contract.py | 21 +- test/python/fe_api/test_jax_gemm_amax.py | 6 +- .../fe_api/test_jax_gemm_amax_contract.py | 55 ++- test/python/fe_api/test_jax_gemm_layout.py | 2 +- test/python/fe_api/test_jax_gemm_relu.py | 15 +- .../fe_api/test_jax_gemm_relu_contract.py | 107 ++++- test/python/fe_api/test_jax_gemm_swiglu.py | 6 +- .../fe_api/test_jax_gemm_swiglu_contract.py | 79 ++-- .../fe_api/test_jax_grouped_gemm_contract.py | 89 +++- ...jax_grouped_wgrad_accumulation_contract.py | 82 +++- test/python/fe_api/test_jax_layout.py | 2 +- .../fe_api/test_jax_nsa_topk_thd_contract.py | 56 ++- .../fe_api/test_jax_sdpa_thd_contract.py | 5 +- test/python/fe_api/test_result.py | 10 +- .../fe_api/test_rmsnorm_rht_amax_imports.py | 6 +- .../fe_api/test_rmsnorm_rht_amax_kernel.py | 2 +- .../python/fe_api/test_rmsnorm_rht_amax_op.py | 6 +- test/python/fe_api/test_shared_tensor_desc.py | 4 +- 130 files changed, 2697 insertions(+), 2350 deletions(-) create mode 100644 python/cudnn/_jax/compiler.py create mode 100644 python/cudnn/common/__init__.py rename python/cudnn/{_cute_compiler.py => common/cute_arch.py} (64%) rename python/cudnn/{_op.py => common/op.py} (100%) rename python/cudnn/{_operation_api.py => common/operation_api.py} (100%) rename python/cudnn/{_result.py => common/result.py} (100%) rename python/cudnn/{_tensor_desc.py => common/tensor_desc.py} (99%) create mode 100644 python/cudnn/gemm/__init__.py rename python/cudnn/{_dense_gemm.py => gemm/helpers.py} (99%) rename python/cudnn/{_gemm_relu.py => gemm/srelu.py} (97%) delete mode 100644 test/python/fe_api/test_dsa_compiler.py create mode 100644 test/python/fe_api/test_jax_compiler.py diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 04f771209..7465c8d76 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -51,8 +51,8 @@ def is_windows(): if hasattr(_pybind_module, _optional_symbol): globals()[_optional_symbol] = getattr(_pybind_module, _optional_symbol) -from ._op import Op -from ._tensor_desc import TensorDesc +from .common.op import Op +from .common.tensor_desc import TensorDesc from .datatypes import _library_type, _is_torch_tensor __version__ = "1.26.0" diff --git a/python/cudnn/_jax/api_base.py b/python/cudnn/_jax/api_base.py index 6c4fa13ba..e708f9a6d 100644 --- a/python/cudnn/_jax/api_base.py +++ b/python/cudnn/_jax/api_base.py @@ -7,16 +7,109 @@ from abc import ABC, abstractmethod from collections.abc import Callable +from dataclasses import dataclass, field, replace from operator import index from typing import Any from .. import data_type -from .._tensor_desc import TensorDesc +from ..common.tensor_desc import TensorDesc from .layout import compact_stride, normalize_mode, to_canonical_axes, to_cutlass_layout, to_public_axes +@dataclass(frozen=True) class JaxTensorDesc(TensorDesc[Any]): - """Framework-neutral tensor metadata backed by a JAX dtype.""" + """JAX tensor metadata and CUTLASS lowering constraints.""" + + mode: tuple[int, ...] = field(default_factory=tuple) + divisibility: tuple[int | None, ...] | None = None + ptr_assumed_align: int | None = None + + def __post_init__(self) -> None: + super().__post_init__() + object.__setattr__( + self, + "mode", + normalize_mode(self.ndim, None if not self.mode else self.mode), + ) + object.__setattr__( + self, + "divisibility", + _normalize_divisibility(self.ndim, self.divisibility), + ) + object.__setattr__( + self, + "ptr_assumed_align", + _normalize_ptr_assumed_align(self.ptr_assumed_align), + ) + + @property + def array_shape(self) -> tuple[int, ...]: + """Return the shape in the axis order used by the JAX array.""" + + return to_public_axes(self.shape, self.mode) + + @classmethod + def from_array( + cls, + value: Any, + *, + name: str = "", + mode: tuple[int, ...] | None = None, + public_stride_order: tuple[int, ...] | None = None, + init_value: bool | int | float | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, + ) -> "JaxTensorDesc": + """Describe a JAX array-like value using its shape and dtype metadata.""" + + shape = _require_array_metadata(value, name or "value") + return cls.from_shape( + shape, + value.dtype, + name=name, + mode=mode, + public_stride_order=public_stride_order, + init_value=init_value, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + ) + + @classmethod + def from_shape( + cls, + shape: tuple[int, ...], + dtype: Any, + *, + name: str = "", + mode: tuple[int, ...] | None = None, + public_stride_order: tuple[int, ...] | None = None, + init_value: bool | int | float | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, + ) -> "JaxTensorDesc": + """Describe a compact JAX tensor from its user-facing shape and dtype.""" + + shape = tuple(shape) + rank = len(shape) + mode = normalize_mode(rank, mode) + public_stride_order = _normalize_stride_order(rank, public_stride_order) + public_stride = compact_stride(shape, public_stride_order) + canonical_axis_by_public_axis = to_public_axes(tuple(range(rank)), mode) + + return cls( + dtype=dtype, + shape=to_canonical_axes(shape, mode), + stride=to_canonical_axes(public_stride, mode), + stride_order=tuple( + canonical_axis_by_public_axis[axis] + for axis in public_stride_order + ), + name=name, + init_value=init_value, + mode=mode, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + ) @property def cudnn_dtype(self) -> data_type: @@ -32,6 +125,9 @@ def compact_like( stride_order: tuple[int, ...] | None = None, name: str = "", init_value: bool | int | float | None = None, + mode: tuple[int, ...] | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, ) -> "JaxTensorDesc": """Create a compact JAX descriptor from canonical output metadata.""" @@ -52,8 +148,19 @@ def compact_like( stride_order=canonical.stride_order, name=canonical.name, init_value=canonical.init_value, + mode=normalize_mode(canonical.ndim, mode), + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, ) + def with_divisibility( + self, + divisibility: tuple[int | None, ...] | None, + ) -> "JaxTensorDesc": + """Return this descriptor with canonical-axis divisibility constraints.""" + + return replace(self, divisibility=divisibility) + def _require_array_metadata(value: Any, name: str) -> tuple[Any, ...]: if not hasattr(value, "shape") or not hasattr(value, "dtype"): @@ -82,6 +189,61 @@ def _normalize_stride_order(rank: int, stride_order: tuple[int, ...] | None) -> return normalized_order +def _normalize_divisibility( + rank: int, + divisibility: tuple[int | None, ...] | None, +) -> tuple[int | None, ...] | None: + if divisibility is None: + return None + divisibility = tuple(divisibility) + if len(divisibility) != rank: + raise ValueError( + f"divisibility rank mismatch: expected {rank}, got {len(divisibility)}" + ) + + normalized = [] + for value in divisibility: + if value is None: + normalized.append(None) + continue + if isinstance(value, bool): + raise TypeError( + f"divisibility entries must be positive integers or None, got {value!r}" + ) + try: + value = index(value) + except TypeError as error: + raise TypeError( + f"divisibility entries must be positive integers or None, got {value!r}" + ) from error + if value <= 0: + raise ValueError( + f"divisibility entries must be positive integers or None, got {value}" + ) + normalized.append(value) + return tuple(normalized) + + +def _normalize_ptr_assumed_align(value: int | None) -> int | None: + if value is None: + return None + if isinstance(value, bool): + raise TypeError( + f"ptr_assumed_align must be a positive integer or None, got {value!r}" + ) + try: + value = index(value) + except TypeError as error: + raise TypeError( + f"ptr_assumed_align must be a positive integer or None, got {value!r}" + ) from error + if value <= 0: + raise ValueError( + f"ptr_assumed_align must be a positive integer or None, got {value}" + ) + return value + + class JaxApiBase(ABC): """Common tensor metadata, validation, and CuTe binding for JAX adapters.""" @@ -257,6 +419,8 @@ def _to_tensor_desc( mode: tuple[int, ...] | None = None, public_stride_order: tuple[int, ...] | None = None, init_value: bool | int | float | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, ) -> JaxTensorDesc: """Describe a public JAX array in canonical kernel-axis order. @@ -265,37 +429,32 @@ def _to_tensor_desc( to slowest. It defaults to compact row-major storage. """ - public_shape = _require_array_metadata(value, name) - rank = len(public_shape) - mode = normalize_mode(rank, mode) - public_stride_order = _normalize_stride_order(rank, public_stride_order) - public_stride = compact_stride(public_shape, public_stride_order) - canonical_axis_by_public_axis = to_public_axes(tuple(range(rank)), mode) - - return JaxTensorDesc( - dtype=value.dtype, - shape=to_canonical_axes(public_shape, mode), - stride=to_canonical_axes(public_stride, mode), - stride_order=tuple(canonical_axis_by_public_axis[axis] for axis in public_stride_order), + return JaxTensorDesc.from_array( + value, name=name, + mode=mode, + public_stride_order=public_stride_order, init_value=init_value, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, ) @staticmethod def _check_tensor_signature( value: Any, - expected: TensorDesc[Any], - *, - mode: tuple[int, ...] | None = None, + expected: JaxTensorDesc, ) -> None: """Validate a public JAX value against a canonical descriptor.""" + if not isinstance(expected, JaxTensorDesc): + raise TypeError( + f"expected must be a JaxTensorDesc, got {type(expected).__name__}" + ) name = expected.name or "value" public_shape = _require_array_metadata(value, name) - mode = normalize_mode(expected.ndim, mode) if len(public_shape) != expected.ndim: raise ValueError(f"{name} tensor shape mismatch: expected {expected.shape}, got public shape {public_shape}") - actual_shape = to_canonical_axes(public_shape, mode) + actual_shape = to_canonical_axes(public_shape, expected.mode) if actual_shape != expected.shape: raise ValueError(f"{name} tensor shape mismatch: expected {expected.shape}, got {actual_shape}") @@ -307,60 +466,71 @@ def _check_tensor_signature( @staticmethod def _to_tensor_spec( - desc: TensorDesc[Any], - *, - mode: tuple[int, ...] | None = None, - divisibility: tuple[int | None, ...] | None = None, + desc: JaxTensorDesc, ) -> Any: """Build a CUTLASS TensorSpec indexed by public JAX array axes. - ``desc`` and ``divisibility`` use canonical kernel axes. TensorSpec + ``desc.divisibility`` uses canonical kernel axes. TensorSpec layout and divisibility use public axes, while TensorSpec ``mode`` records the canonical-to-public binding. """ - mode = normalize_mode(desc.ndim, mode) + if not isinstance(desc, JaxTensorDesc): + raise TypeError( + f"desc must be a JaxTensorDesc, got {type(desc).__name__}" + ) public_layout = to_cutlass_layout( desc.shape, desc.stride, desc.stride_order, - mode=mode, + mode=desc.mode, name=desc.name or "tensor", ) - if divisibility is None: + default_mode = tuple(range(desc.ndim)) + default_layout = tuple(reversed(range(desc.ndim))) + if ( + desc.mode == default_mode + and public_layout == default_layout + and desc.divisibility is None + and desc.ptr_assumed_align is None + ): + return None + + if desc.divisibility is None: public_divisibility = None else: - divisibility = tuple(divisibility) - if len(divisibility) != desc.ndim: - raise ValueError(f"divisibility rank mismatch: expected {desc.ndim}, got {len(divisibility)}") - public_divisibility = to_public_axes(divisibility, mode) + public_divisibility = to_public_axes( + desc.divisibility, + desc.mode, + ) from cutlass.jax import TensorSpec - return TensorSpec( + options: dict[str, Any] = dict( layout=public_layout, - mode=mode, - divisibility=public_divisibility, + mode=desc.mode, ) + if public_divisibility is not None: + options["divisibility"] = public_divisibility + if desc.ptr_assumed_align is not None: + options["ptr_assumed_align"] = desc.ptr_assumed_align + return TensorSpec(**options) @staticmethod - def _materialize_tensor_desc( - desc: TensorDesc[Any], - *, - mode: tuple[int, ...] | None = None, + def _to_shape_dtype_struct( + desc: JaxTensorDesc, ) -> Any: - """Declare a canonical descriptor as a public JAX output buffer.""" + """Convert a canonical descriptor to abstract JAX output metadata.""" - if not isinstance(desc, TensorDesc): - raise TypeError(f"desc must be a TensorDesc, got {type(desc).__name__}") - mode = normalize_mode(desc.ndim, mode) + if not isinstance(desc, JaxTensorDesc): + raise TypeError(f"desc must be a JaxTensorDesc, got {type(desc).__name__}") from jax import ShapeDtypeStruct from .datatypes import cudnn_to_jax_dtype return ShapeDtypeStruct( - to_public_axes(desc.shape, mode), + desc.array_shape, cudnn_to_jax_dtype(desc.cudnn_dtype), ) @@ -369,12 +539,9 @@ def _call_kernel( inputs: tuple[Any, ...], *, launch: Callable[..., None], - output_descs: tuple[TensorDesc[Any], ...], - output_seeds: tuple[Any | None, ...] | None = None, - workspace_descs: tuple[TensorDesc[Any], ...] = (), - input_spec: tuple[Any | None, ...] | None = None, - output_spec: tuple[Any | None, ...] | None = None, - workspace_spec: tuple[Any | None, ...] | None = None, + input_descs: tuple[JaxTensorDesc, ...], + output_descs: tuple[JaxTensorDesc, ...], + workspace_descs: tuple[JaxTensorDesc, ...] = (), allow_cuda_graph: bool = True, compile_options: Any = None, use_static_tensors: bool = True, @@ -386,13 +553,13 @@ def _call_kernel( are declared as custom-call results so XLA owns their lifetime, then omitted from this method's return value. Descriptors with a non-``None`` ``init_value`` are materialized as JAX inputs and aliased to their - corresponding custom-call results. ``output_seeds`` supplies existing - JAX values for outputs that must be updated in place by the kernel; - these operands are likewise aliased to the corresponding results. - Inputs are a flat tuple of arrays, and all explicit specs are native - CUTLASS ``TensorSpec`` values. + corresponding custom-call results. Outputs without an initial value + are allocated as custom-call results by XLA. + Descriptors validate input signatures and provide all CUTLASS + ``TensorSpec`` metadata. """ + import cutlass.cute as cute import cutlass.jax as cutlass_jax import jax.numpy as jnp @@ -407,105 +574,78 @@ def _call_kernel( if not outputs: raise ValueError("A JAX operation must provide at least one output descriptor") - for label, descs in (("output", outputs), ("workspace", workspaces)): + input_descs = tuple(input_descs) + if len(input_descs) != len(inputs): + raise ValueError( + f"Expected {len(inputs)} input descriptors, got {len(input_descs)}" + ) + + for label, descs in ( + ("input", input_descs), + ("output", outputs), + ("workspace", workspaces), + ): for desc in descs: - if not isinstance(desc, TensorDesc): - raise TypeError(f"{label}_descs must contain TensorDesc values, got {type(desc).__name__}") - - if input_spec is None: - input_specs = (None,) * len(inputs) - else: - input_specs = tuple(input_spec) - if len(input_specs) != len(inputs): - raise ValueError(f"Expected {len(inputs)} input specs, got {len(input_specs)}") - if any(spec is not None and not isinstance(spec, cutlass_jax.TensorSpec) for spec in input_specs): - raise TypeError("input_spec must contain only TensorSpec or None values") - - def resolve_specs( - descs: tuple[TensorDesc[Any], ...], - supplied: tuple[Any | None, ...] | None, - label: str, - ) -> tuple[Any, ...]: - if supplied is None: - supplied = (None,) * len(descs) - else: - supplied = tuple(supplied) - if len(supplied) != len(descs): - raise ValueError(f"Expected {len(descs)} {label} specs, got {len(supplied)}") - if any(spec is not None and not isinstance(spec, cutlass_jax.TensorSpec) for spec in supplied): - raise TypeError(f"{label}_spec must contain only TensorSpec or None values") - return tuple(self._to_tensor_spec(desc) if spec is None else spec for desc, spec in zip(descs, supplied)) - - output_specs = resolve_specs(outputs, output_spec, "output") - workspace_specs = resolve_specs(workspaces, workspace_spec, "workspace") + if not isinstance(desc, JaxTensorDesc): + raise TypeError( + f"{label}_descs must contain JaxTensorDesc values, " + f"got {type(desc).__name__}" + ) + + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc) + + input_specs = tuple(self._to_tensor_spec(desc) for desc in input_descs) + output_specs = tuple(self._to_tensor_spec(desc) for desc in outputs) + workspace_specs = tuple(self._to_tensor_spec(desc) for desc in workspaces) buffers = outputs + workspaces buffer_specs = output_specs + workspace_specs - buffer_metadata = tuple(self._materialize_tensor_desc(desc, mode=getattr(spec, "mode", None)) for desc, spec in zip(buffers, buffer_specs)) - - if output_seeds is None: - supplied_output_seeds = (None,) * len(outputs) - else: - supplied_output_seeds = tuple(output_seeds) - if len(supplied_output_seeds) != len(outputs): - raise ValueError( - f"Expected {len(outputs)} output seeds, got {len(supplied_output_seeds)}" - ) - - seed_by_buffer: dict[int, Any] = {} - for output_index, seed in enumerate(supplied_output_seeds): - if seed is None: - continue - output_name = outputs[output_index].name or f"output #{output_index}" - if outputs[output_index].init_value is not None: - raise ValueError( - f"{output_name} cannot have both an explicit seed and init_value" - ) - _require_array_metadata(seed, f"output seed #{output_index}") - metadata = buffer_metadata[output_index] - if tuple(seed.shape) != tuple(metadata.shape): - raise ValueError( - f"{output_name} seed shape mismatch: " - f"expected {tuple(metadata.shape)}, got {tuple(seed.shape)}" - ) - if seed.dtype != metadata.dtype: - raise ValueError( - f"{output_name} seed dtype mismatch: " - f"expected {metadata.dtype}, got {seed.dtype}" - ) - seed_by_buffer[output_index] = seed + buffer_metadata = tuple( + self._to_shape_dtype_struct(desc) for desc in buffers + ) + initialized_buffer_by_index: dict[int, Any] = {} for buffer_index, desc in enumerate(buffers): if desc.init_value is not None: - seed_by_buffer[buffer_index] = jnp.full( + initialized_buffer_by_index[buffer_index] = jnp.full( buffer_metadata[buffer_index].shape, desc.init_value, dtype=buffer_metadata[buffer_index].dtype, ) - initialized = tuple(sorted(seed_by_buffer)) + initialized = tuple(sorted(initialized_buffer_by_index)) uninitialized = tuple( buffer_index for buffer_index in range(len(buffers)) - if buffer_index not in seed_by_buffer + if buffer_index not in initialized_buffer_by_index + ) + initialized_inputs = tuple( + initialized_buffer_by_index[buffer_index] + for buffer_index in initialized ) - seed_inputs = tuple(seed_by_buffer[buffer_index] for buffer_index in initialized) - aliases = {len(inputs) + seed_index: result_index for seed_index, result_index in enumerate(initialized)} + aliases = { + len(inputs) + initialized_input_index: result_index + for initialized_input_index, result_index in enumerate(initialized) + } call_input_specs = input_specs + tuple( buffer_specs[buffer_index] for buffer_index in initialized ) input_count = len(inputs) - seed_count = len(seed_inputs) + initialized_count = len(initialized_inputs) + @cute.jit(preprocess=False) def launcher(stream, *args): kernel_inputs = args[:input_count] - seeded_buffers = args[input_count : input_count + seed_count] - allocated_buffers = args[input_count + seed_count :] + initialized_buffers = args[ + input_count : input_count + initialized_count + ] + allocated_buffers = args[input_count + initialized_count :] if len(allocated_buffers) != len(uninitialized): raise RuntimeError(f"Kernel received {len(allocated_buffers)} allocated buffers; expected {len(uninitialized)}") ordered_buffers: list[Any] = [None] * len(buffers) - for buffer_index, value in zip(initialized, seeded_buffers): + for buffer_index, value in zip(initialized, initialized_buffers): ordered_buffers[buffer_index] = value for buffer_index, value in zip(uninitialized, allocated_buffers): ordered_buffers[buffer_index] = value @@ -523,7 +663,7 @@ def launcher(stream, *args): compile_options=compile_options, use_static_tensors=use_static_tensors, ) - results = call(*inputs, *seed_inputs) + results = call(*inputs, *initialized_inputs) if not isinstance(results, (tuple, list)): results = (results,) return tuple(results[: len(outputs)]) diff --git a/python/cudnn/_jax/compiler.py b/python/cudnn/_jax/compiler.py new file mode 100644 index 000000000..eda074de9 --- /dev/null +++ b/python/cudnn/_jax/compiler.py @@ -0,0 +1,20 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""CUTLASS JAX compilation options.""" + +from __future__ import annotations + +from ..common.cute_arch import gpu_arch_flag_for_compute_capability + + +def compile_options_for_target(compute_capability: int, extra: str = "") -> str: + """Build CUTLASS JAX compile options for an explicit target.""" + + parts = [f"--gpu-arch {gpu_arch_flag_for_compute_capability(compute_capability)}"] + if extra: + parts.append(extra) + return " ".join(parts) + + +__all__ = ["compile_options_for_target"] diff --git a/python/cudnn/_jax/gemm.py b/python/cudnn/_jax/gemm.py index 4734e89f8..f0f359420 100644 --- a/python/cudnn/_jax/gemm.py +++ b/python/cudnn/_jax/gemm.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free public-axis bindings for dense JAX GEMMs.""" +"""Public-axis bindings for dense JAX GEMMs.""" from __future__ import annotations diff --git a/python/cudnn/_jax/layout.py b/python/cudnn/_jax/layout.py index 780729316..9cdc93b13 100644 --- a/python/cudnn/_jax/layout.py +++ b/python/cudnn/_jax/layout.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free axis and layout conversion for CUTLASS JAX calls.""" +"""Axis and layout conversion for CUTLASS JAX calls.""" from __future__ import annotations diff --git a/python/cudnn/_jax/result.py b/python/cudnn/_jax/result.py index c19ddcc80..7f302ede2 100644 --- a/python/cudnn/_jax/result.py +++ b/python/cudnn/_jax/result.py @@ -10,7 +10,7 @@ import jax -from .._result import TupleDict +from ..common.result import TupleDict def _flatten(value: TupleDict) -> tuple[tuple[Any, ...], tuple[Any, ...]]: diff --git a/python/cudnn/api_base.py b/python/cudnn/api_base.py index 9cddcb85c..ed382369a 100644 --- a/python/cudnn/api_base.py +++ b/python/cudnn/api_base.py @@ -22,8 +22,8 @@ import cutlass.cute as cute from cudnn import data_type from cudnn._experimental_warnings import warn_experimental_api_once -from cudnn._tensor_desc import TensorDesc as _TensorDesc -from cudnn._result import TupleDict as TupleDict +from cudnn.common.result import TupleDict as TupleDict +from cudnn.common.tensor_desc import TensorDesc as _TensorDesc from cudnn.datatypes import ( _convert_to_cutlass_data_type, _cudnn_to_torch_data_type, diff --git a/python/cudnn/block_sparse_attention/jax.py b/python/cudnn/block_sparse_attention/jax.py index 4d4a5c53f..53147ad82 100644 --- a/python/cudnn/block_sparse_attention/jax.py +++ b/python/cudnn/block_sparse_attention/jax.py @@ -12,7 +12,7 @@ import jax.numpy as jnp from .. import data_type -from .._cute_compiler import compile_options_for_target +from .._jax.compiler import compile_options_for_target from .._jax import JaxApiBase, JaxTensorDesc, TupleDict from .._jax.layout import mode_from_layout, to_public_axes from .op import ( @@ -39,20 +39,12 @@ def _data_mode(layout: str) -> tuple[int, ...]: return mode_from_layout(normalized.upper(), kernel_axes="BHSD") -def _shape_sample( - shape: tuple[int, ...], dtype: Any, mode: tuple[int, ...] | None = None -) -> Any: - return jax.ShapeDtypeStruct(to_public_axes(shape, mode), dtype) - - def _optional_signature( value: Any | None, desc: JaxTensorDesc | None, name: str ) -> None: if (value is None) != (desc is None): expected = "omitted" if desc is None else "provided" raise ValueError(f"{name} must be {expected} for this specialized callable") - if desc is not None: - JaxApiBase._check_tensor_signature(value, desc) def _default_forward_sparse_block_size(target: int) -> int: @@ -215,11 +207,10 @@ def __init__( value_dim = self.v_desc.shape[3] canonical_output_shape = (batch, num_q_heads, seqlen_q, value_dim) if sample_o is None: - self.o_desc = self._to_tensor_desc( - _shape_sample( - canonical_output_shape, self.q_desc.dtype, self.data_mode - ), - "sample_o", + self.o_desc = JaxTensorDesc.from_shape( + to_public_axes(canonical_output_shape, self.data_mode), + self.q_desc.dtype, + name="sample_o", mode=self.data_mode, ) else: @@ -227,9 +218,10 @@ def __init__( sample_o, "sample_o", mode=self.data_mode ) if sample_lse is None: - self.lse_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((batch, num_q_heads, seqlen_q), jnp.float32), - "sample_lse", + self.lse_desc = JaxTensorDesc.from_shape( + (batch, num_q_heads, seqlen_q), + jnp.float32, + name="sample_lse", ) else: self.lse_desc = self._to_tensor_desc(sample_lse, "sample_lse") @@ -297,23 +289,19 @@ def __call__( q2k_block_nums: Any | None = None, ) -> TupleDict: self.check_support() - self._check_tensor_signature(q, self.q_desc, mode=self.data_mode) - self._check_tensor_signature(k, self.k_desc, mode=self.data_mode) - self._check_tensor_signature(v, self.v_desc, mode=self.data_mode) - self._check_tensor_signature(q2k_block_index, self.block_index_desc) _optional_signature(block_sizes, self.block_sizes_desc, "block_sizes") _optional_signature(q2k_block_nums, self.block_nums_desc, "q2k_block_nums") inputs: list[Any] = [q, k, v, q2k_block_index] - bindings: list[tuple[JaxTensorDesc, tuple[int, ...] | None]] = [ - (self.q_desc, self.data_mode), - (self.k_desc, self.data_mode), - (self.v_desc, self.data_mode), - (self.block_index_desc, None), + input_descs: list[JaxTensorDesc] = [ + self.q_desc, + self.k_desc, + self.v_desc, + self.block_index_desc, ] if self.block_sizes_desc is not None: inputs.append(block_sizes) - bindings.append((self.block_sizes_desc, None)) + input_descs.append(self.block_sizes_desc) effective_block_nums = q2k_block_nums effective_block_nums_desc = self.block_nums_desc @@ -329,13 +317,14 @@ def __call__( effective_block_nums = jnp.full( block_nums_shape, self.block_sparse_num, dtype=jnp.int32 ) - effective_block_nums_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct(block_nums_shape, jnp.int32), - "effective_q2k_block_nums", + effective_block_nums_desc = JaxTensorDesc.from_shape( + block_nums_shape, + jnp.int32, + name="effective_q2k_block_nums", ) if effective_block_nums is not None: inputs.append(effective_block_nums) - bindings.append((effective_block_nums_desc, None)) + input_descs.append(effective_block_nums_desc) split_offsets_desc = None if self.kv_splits > 1: @@ -371,11 +360,11 @@ def __call__( jnp.minimum(aligned_offsets, valid_kv[..., None]), ).astype(jnp.int32) split_offsets_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct(split_offsets.shape, split_offsets.dtype), + split_offsets, "split_offsets", ) inputs.append(split_offsets) - bindings.append((split_offsets_desc, None)) + input_descs.append(split_offsets_desc) self._forward_has_effective_block_nums = effective_block_nums is not None self._forward_has_split_offsets = split_offsets_desc is not None @@ -406,14 +395,8 @@ def __call__( tuple(inputs), launch=self._launch, output_descs=(self.o_desc, self.lse_desc), + input_descs=tuple(input_descs), workspace_descs=workspace_descs, - input_spec=tuple( - self._to_tensor_spec(desc, mode=mode) for desc, mode in bindings - ), - output_spec=( - self._to_tensor_spec(self.o_desc, mode=self.data_mode), - self._to_tensor_spec(self.lse_desc), - ), compile_options=compile_options_for_target(self.compute_capability), ) return TupleDict(o_tensor=o, lse_tensor=lse) @@ -787,9 +770,10 @@ def output_desc( ) -> JaxTensorDesc: if sample is not None: return self._to_tensor_desc(sample, name, mode=self.data_mode) - return self._to_tensor_desc( - _shape_sample(source.shape, source.dtype, self.data_mode), - name, + return JaxTensorDesc.from_shape( + to_public_axes(source.shape, self.data_mode), + source.dtype, + name=name, mode=self.data_mode, ) @@ -1022,35 +1006,25 @@ def __call__( q2k_block_nums: Any | None = None, ) -> TupleDict: self.check_support() - for value, desc in ( - (do, self.do_desc), - (q, self.q_desc), - (k, self.k_desc), - (v, self.v_desc), - (o, self.o_desc), - ): - self._check_tensor_signature(value, desc, mode=self.data_mode) - self._check_tensor_signature(lse, self.lse_desc) - self._check_tensor_signature(q2k_block_index, self.block_index_desc) _optional_signature(block_sizes, self.block_sizes_desc, "block_sizes") _optional_signature(q2k_block_nums, self.block_nums_desc, "q2k_block_nums") inputs: list[Any] = [do, q, k, v, o, lse, q2k_block_index] - bindings: list[tuple[JaxTensorDesc, tuple[int, ...] | None]] = [ - (self.do_desc, self.data_mode), - (self.q_desc, self.data_mode), - (self.k_desc, self.data_mode), - (self.v_desc, self.data_mode), - (self.o_desc, self.data_mode), - (self.lse_desc, None), - (self.block_index_desc, None), + input_descs: list[JaxTensorDesc] = [ + self.do_desc, + self.q_desc, + self.k_desc, + self.v_desc, + self.o_desc, + self.lse_desc, + self.block_index_desc, ] if self.block_sizes_desc is not None: inputs.append(block_sizes) - bindings.append((self.block_sizes_desc, None)) + input_descs.append(self.block_sizes_desc) if self.block_nums_desc is not None: inputs.append(q2k_block_nums) - bindings.append((self.block_nums_desc, None)) + input_descs.append(self.block_nums_desc) self._backward_input_count = len(inputs) workspace_items = self._workspace_descs() self._backward_workspace_names = tuple(name for name, _ in workspace_items) @@ -1060,15 +1034,8 @@ def __call__( tuple(inputs), launch=self._launch, output_descs=(self.dq_desc, self.dk_desc, self.dv_desc), + input_descs=tuple(input_descs), workspace_descs=workspace_descs, - input_spec=tuple( - self._to_tensor_spec(desc, mode=mode) for desc, mode in bindings - ), - output_spec=( - self._to_tensor_spec(self.dq_desc, mode=self.data_mode), - self._to_tensor_spec(self.dk_desc, mode=self.data_mode), - self._to_tensor_spec(self.dv_desc, mode=self.data_mode), - ), compile_options=compile_options_for_target(self.compute_capability), ) return TupleDict(dq_tensor=dq, dk_tensor=dk, dv_tensor=dv) @@ -1328,16 +1295,13 @@ def block_sparse_attention_forward( ) -> TupleDict: """Run non-causal block-sparse attention from JAX arrays.""" - def sample(value: Any | None) -> Any | None: - return None if value is None else jax.ShapeDtypeStruct(value.shape, value.dtype) - return BlockSparseAttentionForward( - sample(q_tensor), - sample(k_tensor), - sample(v_tensor), - sample(q2k_block_index), - sample_block_sizes=sample(block_sizes), - sample_q2k_block_nums=sample(q2k_block_nums), + q_tensor, + k_tensor, + v_tensor, + q2k_block_index, + sample_block_sizes=block_sizes, + sample_q2k_block_nums=q2k_block_nums, block_sparse_num=block_sparse_num, sparse_block_size=sparse_block_size, allow_empty_block_nums=allow_empty_block_nums, @@ -1388,19 +1352,16 @@ def block_sparse_attention_backward( ) -> TupleDict: """Compute explicit block-sparse attention gradients from JAX arrays.""" - def sample(value: Any | None) -> Any | None: - return None if value is None else jax.ShapeDtypeStruct(value.shape, value.dtype) - return BlockSparseAttentionBackward( - sample(do_tensor), - sample(q_tensor), - sample(k_tensor), - sample(v_tensor), - sample(o_tensor), - sample(lse_tensor), - sample(q2k_block_index), - sample_block_sizes=sample(block_sizes), - sample_q2k_block_nums=sample(q2k_block_nums), + do_tensor, + q_tensor, + k_tensor, + v_tensor, + o_tensor, + lse_tensor, + q2k_block_index, + sample_block_sizes=block_sizes, + sample_q2k_block_nums=q2k_block_nums, block_sparse_num=block_sparse_num, bucket_size_blocks=bucket_size_blocks, sparse_block_size=sparse_block_size, diff --git a/python/cudnn/block_sparse_attention/op.py b/python/cudnn/block_sparse_attention/op.py index bec4ecf31..673f2cc78 100644 --- a/python/cudnn/block_sparse_attention/op.py +++ b/python/cudnn/block_sparse_attention/op.py @@ -8,8 +8,8 @@ from typing import Any from .. import data_type -from .._op import Op -from .._tensor_desc import TensorDesc +from ..common.op import Op +from ..common.tensor_desc import TensorDesc SUPPORTED_FORWARD_COMPUTE_CAPABILITIES = (90, 100, 103, 107, 110, 120) diff --git a/python/cudnn/common/__init__.py b/python/cudnn/common/__init__.py new file mode 100644 index 000000000..c833f57f0 --- /dev/null +++ b/python/cudnn/common/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral metadata and API infrastructure.""" diff --git a/python/cudnn/_cute_compiler.py b/python/cudnn/common/cute_arch.py similarity index 64% rename from python/cudnn/_cute_compiler.py rename to python/cudnn/common/cute_arch.py index 23549f2d6..cbc678310 100644 --- a/python/cudnn/_cute_compiler.py +++ b/python/cudnn/common/cute_arch.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free CuTe compilation options for framework adapters.""" +"""CuTe architecture naming shared by framework adapters.""" from __future__ import annotations @@ -27,16 +27,4 @@ def gpu_arch_flag_for_compute_capability(compute_capability: int) -> str: raise RuntimeError(f"Unsupported GPU compute capability SM{compute_capability} for CuTe kernels; supported targets are {supported}") from error -def compile_options_for_target(compute_capability: int, extra: str = "") -> str: - """Build CuTe compile options for an explicit framework-supplied target.""" - - parts = [ - "--enable-tvm-ffi", - f"--gpu-arch {gpu_arch_flag_for_compute_capability(compute_capability)}", - ] - if extra: - parts.append(extra) - return " ".join(parts) - - -__all__ = ["compile_options_for_target", "gpu_arch_flag_for_compute_capability"] +__all__ = ["gpu_arch_flag_for_compute_capability"] diff --git a/python/cudnn/_op.py b/python/cudnn/common/op.py similarity index 100% rename from python/cudnn/_op.py rename to python/cudnn/common/op.py diff --git a/python/cudnn/_operation_api.py b/python/cudnn/common/operation_api.py similarity index 100% rename from python/cudnn/_operation_api.py rename to python/cudnn/common/operation_api.py diff --git a/python/cudnn/_result.py b/python/cudnn/common/result.py similarity index 100% rename from python/cudnn/_result.py rename to python/cudnn/common/result.py diff --git a/python/cudnn/_tensor_desc.py b/python/cudnn/common/tensor_desc.py similarity index 99% rename from python/cudnn/_tensor_desc.py rename to python/cudnn/common/tensor_desc.py index 037b4a257..6a998d059 100644 --- a/python/cudnn/_tensor_desc.py +++ b/python/cudnn/common/tensor_desc.py @@ -10,7 +10,7 @@ from operator import index from typing import Any, Generic, TypeVar -from . import data_type +from .. import data_type DTypeT = TypeVar("DTypeT") diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py index 5d5c717b2..787e3868e 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py @@ -3,7 +3,7 @@ """Lazy Torch API and framework-neutral indexer-backward exports.""" -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py index 0bda60ad8..b9440c255 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py @@ -18,7 +18,7 @@ import cuda.bindings.driver as cuda from cudnn import data_type -from cudnn._tensor_desc import make_compact_tensor_desc +from cudnn.common.tensor_desc import make_compact_tensor_desc from cudnn.api_base import APIBase, TupleDict from cudnn.deepseek_sparse_attention.utils.runtime import ( torch_stream_context as _torch_stream_context, diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py index d10b99ef1..ec1e58cf6 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py @@ -13,8 +13,8 @@ from ... import data_type from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..._jax.compiler import compile_options_for_target from ..._jax.layout import mode_from_layout -from ..utils.compiler import compile_options_for_target from .op import DEFAULT_BLOCK_I, DenseIndexerBackwardOp, IndexerBackwardOp _SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) @@ -98,18 +98,7 @@ def _output_desc( shape=shape, stride_order=source.stride_order, name=name, - ) - - def _last_axis_spec( - self, - desc: JaxTensorDesc, - alignment: int, - mode: tuple[int, ...], - ) -> Any: - return self._to_tensor_spec( - desc, mode=mode, - divisibility=(None,) * (desc.ndim - 1) + (alignment,), ) @@ -264,17 +253,7 @@ def __call__( grad_loss: Any = 1.0, ) -> TupleDict: self.check_support() - for value, expected, mode in ( - (index_q, self.iq_desc, self.q_mode), - (weights, self.w_desc, self.w_mode), - (index_k, self.ik_desc, self.k_mode), - (attn_score, self.attn_desc, self.score_mode), - (index_score, self.idx_score_desc, self.score_mode), - (topk_indices, self.topk_desc, self.score_mode), - ): - self._check_tensor_signature(value, expected, mode=mode) grad_loss = _grad_loss_operand(grad_loss) - self._check_tensor_signature(grad_loss, self.grad_loss_desc) dk_accum_desc = self.ik_desc.compact_like( cudnn_dtype=data_type.FLOAT, @@ -282,12 +261,27 @@ def __call__( stride_order=self.ik_desc.stride_order, name="d_index_k_accum", init_value=0.0, + mode=self.k_mode, ) grad_signal_desc = self.attn_desc.compact_like( cudnn_dtype=data_type.FLOAT, shape=self.attn_desc.shape, stride_order=self.attn_desc.stride_order, name="grad_signal", + mode=self.score_mode, + ) + head_dim = self._op.head_dim + iq_desc = self.iq_desc.with_divisibility( + (None,) * (self.iq_desc.ndim - 1) + (head_dim,) + ) + ik_desc = self.ik_desc.with_divisibility( + (None,) * (self.ik_desc.ndim - 1) + (head_dim,) + ) + diq_desc = self.diq_desc.with_divisibility( + (None,) * (self.diq_desc.ndim - 1) + (head_dim,) + ) + dk_accum_desc = dk_accum_desc.with_divisibility( + (None,) * (dk_accum_desc.ndim - 1) + (head_dim,) ) d_index_q, d_weights, d_index_k_accum = self._call_kernel( ( @@ -300,25 +294,17 @@ def __call__( grad_loss, ), launch=self._launch_kernel, - output_descs=(self.diq_desc, self.dw_desc, dk_accum_desc), - workspace_descs=(grad_signal_desc,), - input_spec=( - self._last_axis_spec(self.iq_desc, self._op.head_dim, self.q_mode), - self._to_tensor_spec(self.w_desc, mode=self.w_mode), - self._last_axis_spec(self.ik_desc, self._op.head_dim, self.k_mode), - self._to_tensor_spec(self.attn_desc, mode=self.score_mode), - self._to_tensor_spec(self.idx_score_desc, mode=self.score_mode), - self._to_tensor_spec(self.topk_desc, mode=self.score_mode), - self._to_tensor_spec(self.grad_loss_desc), - ), - output_spec=( - self._last_axis_spec(self.diq_desc, self._op.head_dim, self.q_mode), - self._to_tensor_spec(self.dw_desc, mode=self.w_mode), - self._last_axis_spec(dk_accum_desc, self._op.head_dim, self.k_mode), - ), - workspace_spec=( - self._to_tensor_spec(grad_signal_desc, mode=self.score_mode), + output_descs=(diq_desc, self.dw_desc, dk_accum_desc), + input_descs=( + iq_desc, + self.w_desc, + ik_desc, + self.attn_desc, + self.idx_score_desc, + self.topk_desc, + self.grad_loss_desc, ), + workspace_descs=(grad_signal_desc,), compile_options=compile_options_for_target( self.compute_capability, "--opt-level 3" ), @@ -635,16 +621,6 @@ def __call__( q_causal_offsets: Any | None = None, ) -> TupleDict: self.check_support() - for value, expected, mode in ( - (index_q, self.iq_desc, self.q_mode), - (weights, self.w_desc, self.w_mode), - (index_k, self.ik_desc, self.k_mode), - (attn_score, self.attn_desc, self.score_mode), - (attn_l1norm, self.attn_denom_desc, self.denom_mode), - (index_score, self.idx_score_desc, self.score_mode), - (index_lse, self.idx_lse_desc, self.denom_mode), - ): - self._check_tensor_signature(value, expected, mode=mode) for value, expected, name in ( (cu_seqlens_q, self.cuq_desc, "cu_seqlens_q"), (cu_seqlens_k, self.cuk_desc, "cu_seqlens_k"), @@ -652,10 +628,7 @@ def __call__( ): if (value is None) != (expected is None): raise ValueError(f"{name} presence must match its sample") - if value is not None: - self._check_tensor_signature(value, expected) grad_loss = _grad_loss_operand(grad_loss) - self._check_tensor_signature(grad_loss, self.grad_loss_desc) inputs = ( index_q, @@ -667,54 +640,64 @@ def __call__( index_lse, grad_loss, ) - for optional in (cu_seqlens_q, cu_seqlens_k, q_causal_offsets): - if optional is not None: + input_descs = ( + self.iq_desc, + self.w_desc, + self.ik_desc, + self.attn_desc, + self.attn_denom_desc, + self.idx_score_desc, + self.idx_lse_desc, + self.grad_loss_desc, + ) + for optional, desc in ( + (cu_seqlens_q, self.cuq_desc), + (cu_seqlens_k, self.cuk_desc), + (q_causal_offsets, self.q_offsets_desc), + ): + if desc is not None: inputs += (optional,) + input_descs += (desc,) dk_accum_desc = self.ik_desc.compact_like( cudnn_dtype=data_type.FLOAT, shape=self.ik_desc.shape, stride_order=self.ik_desc.stride_order, name="d_index_k_accum", init_value=0.0, + mode=self.k_mode, ) grad_signal_desc = self.idx_score_desc.compact_like( cudnn_dtype=data_type.FLOAT, shape=self.idx_score_desc.shape, stride_order=self.idx_score_desc.stride_order, name="grad_signal", + mode=self.score_mode, + ) + head_dim = self._op.head_dim + iq_desc = self.iq_desc.with_divisibility( + (None,) * (self.iq_desc.ndim - 1) + (head_dim,) + ) + ik_desc = self.ik_desc.with_divisibility( + (None,) * (self.ik_desc.ndim - 1) + (head_dim,) + ) + diq_desc = self.diq_desc.with_divisibility( + (None,) * (self.diq_desc.ndim - 1) + (head_dim,) + ) + dk_accum_desc = dk_accum_desc.with_divisibility( + (None,) * (dk_accum_desc.ndim - 1) + (head_dim,) + ) + input_descs = ( + iq_desc, + self.w_desc, + ik_desc, + *input_descs[3:], ) d_index_q, d_weights, d_index_k_accum = self._call_kernel( inputs, launch=self._launch_kernel, - output_descs=(self.diq_desc, self.dw_desc, dk_accum_desc), + output_descs=(diq_desc, self.dw_desc, dk_accum_desc), + input_descs=input_descs, workspace_descs=(grad_signal_desc,), - input_spec=( - self._last_axis_spec(self.iq_desc, self._op.head_dim, self.q_mode), - self._to_tensor_spec(self.w_desc, mode=self.w_mode), - self._last_axis_spec(self.ik_desc, self._op.head_dim, self.k_mode), - self._to_tensor_spec(self.attn_desc, mode=self.score_mode), - self._to_tensor_spec(self.attn_denom_desc, mode=self.denom_mode), - self._to_tensor_spec(self.idx_score_desc, mode=self.score_mode), - self._to_tensor_spec(self.idx_lse_desc, mode=self.denom_mode), - self._to_tensor_spec(self.grad_loss_desc), - ) - + tuple( - self._to_tensor_spec(desc) - for desc in ( - self.cuq_desc, - self.cuk_desc, - self.q_offsets_desc, - ) - if desc is not None - ), - output_spec=( - self._last_axis_spec(self.diq_desc, self._op.head_dim, self.q_mode), - self._to_tensor_spec(self.dw_desc, mode=self.w_mode), - self._last_axis_spec(dk_accum_desc, self._op.head_dim, self.k_mode), - ), - workspace_spec=( - self._to_tensor_spec(grad_signal_desc, mode=self.score_mode), - ), compile_options=compile_options_for_target( self.compute_capability, "--opt-level 3" ), @@ -874,12 +857,13 @@ def indexer_backward_wrapper( Packed THD callers must pass ``topk_indices_global=True``. """ - samples = tuple( - jax.ShapeDtypeStruct(value.shape, value.dtype) - for value in (index_q, weights, index_k, attn_score, index_score, topk_indices) - ) return IndexerBackward( - *samples, + index_q, + weights, + index_k, + attn_score, + index_score, + topk_indices, sm_scale=sm_scale, loss_coeff=loss_coeff, block_I=block_I, @@ -934,27 +918,11 @@ def dense_indexer_backward_wrapper( target_compute_capability: int | None = None, ) -> TupleDict: core = (index_q, weights, index_k, attn_score, attn_l1norm, index_score, index_lse) - samples = tuple(jax.ShapeDtypeStruct(value.shape, value.dtype) for value in core) - sample_cuq = ( - None - if cu_seqlens_q is None - else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype) - ) - sample_cuk = ( - None - if cu_seqlens_k is None - else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype) - ) - sample_offsets = ( - None - if q_causal_offsets is None - else jax.ShapeDtypeStruct(q_causal_offsets.shape, q_causal_offsets.dtype) - ) return DenseIndexerBackward( - *samples, - sample_cu_seqlens_q=sample_cuq, - sample_cu_seqlens_k=sample_cuk, - sample_q_causal_offsets=sample_offsets, + *core, + sample_cu_seqlens_q=cu_seqlens_q, + sample_cu_seqlens_k=cu_seqlens_k, + sample_q_causal_offsets=q_causal_offsets, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, sm_scale=sm_scale, diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/op.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/op.py index 49b6e54cd..a9ba6f9cf 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/op.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/op.py @@ -9,8 +9,8 @@ from typing import Any, Optional from ... import data_type -from ..._op import Op -from ..._tensor_desc import TensorDesc +from ...common.op import Op +from ...common.tensor_desc import TensorDesc DEFAULT_BLOCK_I = 128 SUPPORTED_HEAD_DIM = 128 diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py index 012c83fd9..49fbc38ac 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py @@ -3,7 +3,7 @@ JAX APIs are exported through :mod:`cudnn.jax`. """ -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py index db19225f7..2746b8539 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py @@ -12,8 +12,8 @@ import jax.numpy as jnp from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..._jax.compiler import compile_options_for_target from ..._jax.layout import mode_from_layout, to_public_axes -from ..utils.compiler import compile_options_for_target from .op import IndexerForwardOp, SUPPORTED_COMPUTE_CAPABILITIES, TMA_ALIGN_ELEMENTS @@ -200,12 +200,10 @@ def _default_output_desc(self, max_seqlen_k: int | None) -> JaxTensorDesc: padded_seqlen_k = ((logical_seqlen_k + TMA_ALIGN_ELEMENTS - 1) // TMA_ALIGN_ELEMENTS) * TMA_ALIGN_ELEMENTS canonical_shape = (*leading_shape, padded_seqlen_k) - return self._to_tensor_desc( - jax.ShapeDtypeStruct( - to_public_axes(canonical_shape, self.output_mode), - jnp.float32, - ), - "sample_out", + return JaxTensorDesc.from_shape( + to_public_axes(canonical_shape, self.output_mode), + jnp.float32, + name="sample_out", mode=self.output_mode, init_value=float("-inf"), ) @@ -224,19 +222,12 @@ def __call__( q_causal_offsets: Any | None = None, ) -> TupleDict: self.check_support() - self._check_tensor_signature(q, self.q_desc, mode=self.q_mode) - self._check_tensor_signature(k, self.k_desc, mode=self.k_mode) - self._check_tensor_signature(w, self.w_desc, mode=self.w_mode) self._check_optional_signature(cu_seqlens_q, self.cu_seqlens_q_desc, "cu_seqlens_q") self._check_optional_signature(cu_seqlens_k, self.cu_seqlens_k_desc, "cu_seqlens_k") self._check_optional_signature(q_causal_offsets, self.q_causal_offsets_desc, "q_causal_offsets") inputs = [q, k, w] - input_bindings = [ - (self.q_desc, self.q_mode), - (self.k_desc, self.k_mode), - (self.w_desc, self.w_mode), - ] + input_descs = [self.q_desc, self.k_desc, self.w_desc] for value, desc in ( (cu_seqlens_q, self.cu_seqlens_q_desc), (cu_seqlens_k, self.cu_seqlens_k_desc), @@ -244,21 +235,16 @@ def __call__( ): if desc is not None: inputs.append(value) - input_bindings.append((desc, None)) + input_descs.append(desc) - output_divisibility = (None,) * (self.o_desc.ndim - 1) + (TMA_ALIGN_ELEMENTS,) + output_desc = self.o_desc.with_divisibility( + (None,) * (self.o_desc.ndim - 1) + (TMA_ALIGN_ELEMENTS,) + ) (scores_padded,) = self._call_kernel( tuple(inputs), launch=self._launch_kernel, - output_descs=(self.o_desc,), - input_spec=tuple(self._to_tensor_spec(desc, mode=mode) for desc, mode in input_bindings), - output_spec=( - self._to_tensor_spec( - self.o_desc, - mode=self.output_mode, - divisibility=output_divisibility, - ), - ), + output_descs=(output_desc,), + input_descs=tuple(input_descs), compile_options=compile_options_for_target(self.target_compute_capability), ) if self._op.s_k is None: @@ -272,8 +258,6 @@ def _check_optional_signature(value: Any | None, desc: JaxTensorDesc | None, nam if (value is None) != (desc is None): expected = "omitted" if desc is None else "provided" raise ValueError(f"{name} must be {expected} for this specialized callable") - if desc is not None: - JaxApiBase._check_tensor_signature(value, desc) def _launch_kernel( self, @@ -383,26 +367,13 @@ def indexer_forward_wrapper( ) -> TupleDict: """Compute fixed batch/sequence-major or packed THD indexer scores.""" - sample_q = jax.ShapeDtypeStruct(q.shape, q.dtype) - sample_k = jax.ShapeDtypeStruct(k.shape, k.dtype) - sample_w = jax.ShapeDtypeStruct(w.shape, w.dtype) - sample_cu_q = None if cu_seqlens_q is None else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype) - sample_cu_k = None if cu_seqlens_k is None else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype) - sample_offsets = ( - None - if q_causal_offsets is None - else jax.ShapeDtypeStruct( - q_causal_offsets.shape, - q_causal_offsets.dtype, - ) - ) return IndexerForward( - sample_q, - sample_k, - sample_w, - sample_cu_seqlens_q=sample_cu_q, - sample_cu_seqlens_k=sample_cu_k, - sample_q_causal_offsets=sample_offsets, + q, + k, + w, + sample_cu_seqlens_q=cu_seqlens_q, + sample_cu_seqlens_k=cu_seqlens_k, + sample_q_causal_offsets=q_causal_offsets, ratio=ratio, qhead_per_kv_head=qhead_per_kv_head, max_seqlen_q=max_seqlen_q, diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/op.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/op.py index 5ca7af613..347b6a7a8 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/op.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/op.py @@ -8,8 +8,8 @@ from typing import Any from ... import data_type -from ..._op import Op -from ..._tensor_desc import TensorDesc +from ...common.op import Op +from ...common.tensor_desc import TensorDesc TMA_ALIGN_ELEMENTS = 4 SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) @@ -141,6 +141,16 @@ def check_support(self) -> bool: raise ValueError(f"IndexerForward tuning values must be positive, got {tuning}") if self.target_compute_capability < 100 and tuning != (128, 128, 2, 4): raise ValueError("SM90 IndexerForward supports only m_block_size=128, n_block_size=128, " f"q_stage=2, kv_stage=4; got {tuning}") + if self.target_compute_capability >= 100: + packed_q = s_q * qhead_per_kv_head + minimum_packed_q = self.q_stage * self.m_block_size + if packed_q < minimum_packed_q: + raise ValueError( + "SM100-family IndexerForward requires at least one complete packed query tile: max_seqlen_q * " + f"qhead_per_kv_head >= q_stage * m_block_size, got " + f"{s_q} * {qhead_per_kv_head} < {self.q_stage} * " + f"{self.m_block_size}" + ) if self.q_causal_offsets is not None: offsets = self.q_causal_offsets diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/__init__.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/__init__.py index 909029826..7bf0c95d1 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/__init__.py @@ -3,7 +3,7 @@ JAX APIs are exported through :mod:`cudnn.jax`. """ -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py index 6a424e165..83d62194c 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py @@ -13,7 +13,7 @@ from ... import data_type from ..._jax import JaxApiBase, TupleDict -from ..utils.compiler import compile_options_for_target +from ..._jax.compiler import compile_options_for_target from .op import INT32_MAX, IndexerTopKOp, SUPPORTED_COMPUTE_CAPABILITIES, SUPPORTED_INPUT_DTYPES @@ -99,16 +99,14 @@ def check_support(self) -> bool: def __call__(self, input_values: Any, seq_lens: Any) -> TupleDict: self.check_support() - self._check_tensor_signature(input_values, self.input_desc) - self._check_tensor_signature(seq_lens, self.seq_lens_desc) output_descs = (self.indices_desc,) if self.values_desc is None else (self.indices_desc, self.values_desc) results = self._call_kernel( (input_values, seq_lens), launch=self._launch_kernel, output_descs=output_descs, + input_descs=(self.input_desc, self.seq_lens_desc), workspace_descs=(self.workspace_desc,), - input_spec=(self._to_tensor_spec(self.input_desc), self._to_tensor_spec(self.seq_lens_desc)), compile_options=compile_options_for_target(self.target_compute_capability), ) indices = results[0] @@ -225,16 +223,15 @@ def __call__( cu_seqlens_k: Any | None = None, ) -> TupleDict: self.check_support() - self._check_tensor_signature(local_indices, self.local_desc) inputs = [local_indices] + input_descs = [self.local_desc] if self.is_varlen: if cu_seqlens_q is None or cu_seqlens_k is None: raise ValueError("Packed local-to-global requires both cumulative sequence tensors") if self.cu_q_desc is None or self.cu_k_desc is None: raise RuntimeError("Packed cumulative sequence descriptors were not configured") - self._check_tensor_signature(cu_seqlens_q, self.cu_q_desc) - self._check_tensor_signature(cu_seqlens_k, self.cu_k_desc) inputs.extend((cu_seqlens_q, cu_seqlens_k)) + input_descs.extend((self.cu_q_desc, self.cu_k_desc)) elif cu_seqlens_q is not None or cu_seqlens_k is not None: raise ValueError("Fixed local-to-global does not accept cumulative sequence tensors") @@ -242,6 +239,7 @@ def __call__( tuple(inputs), launch=self._launch_kernel, output_descs=(self.output_desc,), + input_descs=tuple(input_descs), compile_options=compile_options_for_target(self.target_compute_capability, "--opt-level 3"), ) return TupleDict(indices=result) @@ -302,11 +300,11 @@ def check_support(self) -> bool: def __call__(self, indices: Any) -> TupleDict: self.check_support() - self._check_tensor_signature(indices, self.input_desc) compact_indices, topk_length = self._call_kernel( (indices,), launch=self._launch_kernel, output_descs=(self.output_desc, self.length_desc), + input_descs=(self.input_desc,), compile_options=compile_options_for_target(self.target_compute_capability, "--opt-level 3"), ) return TupleDict(indices=compact_indices, topk_length=topk_length) @@ -341,8 +339,8 @@ def indexer_top_k_wrapper( """ return IndexerTopK( - jax.ShapeDtypeStruct(input_values.shape, input_values.dtype), - jax.ShapeDtypeStruct(seq_lens.shape, seq_lens.dtype), + input_values, + seq_lens, top_k, next_n=next_n, return_val=return_val, @@ -362,13 +360,11 @@ def local_to_global_wrapper( ) -> TupleDict: """Convert local top-K indices to the global flattened KV index space.""" - sample_cu_q = None if cu_seqlens_q is None else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype) - sample_cu_k = None if cu_seqlens_k is None else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype) return _LocalToGlobal( - jax.ShapeDtypeStruct(local_indices.shape, local_indices.dtype), + local_indices, seqlen_k, - sample_cu_seqlens_q=sample_cu_q, - sample_cu_seqlens_k=sample_cu_k, + sample_cu_seqlens_q=cu_seqlens_q, + sample_cu_seqlens_k=cu_seqlens_k, target_compute_capability=target_compute_capability, )( local_indices, @@ -388,7 +384,7 @@ def compactify_wrapper(indices: Any, *, target_compute_capability: int | None = rows *= extent flat_indices = jnp.reshape(indices, (rows, indices.shape[-1])) return _Compactify( - jax.ShapeDtypeStruct(flat_indices.shape, flat_indices.dtype), + flat_indices, target_compute_capability=target_compute_capability, )(flat_indices) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/op.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/op.py index 31c7664e4..543180ffe 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/op.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/op.py @@ -8,8 +8,8 @@ from typing import Any from ... import data_type -from ..._op import Op -from ..._tensor_desc import TensorDesc +from ...common.op import Op +from ...common.tensor_desc import TensorDesc INT32_MAX = (1 << 31) - 1 SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/__init__.py b/python/cudnn/deepseek_sparse_attention/score_recompute/__init__.py index ddaace4f3..7838de720 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/__init__.py @@ -3,7 +3,7 @@ JAX APIs are exported through :mod:`cudnn.jax`. """ -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py b/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py index c46053e20..f34e42a7e 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py @@ -25,7 +25,7 @@ def _compile_options(target_compute_capability: int) -> str: - from ..utils.compiler import compile_options_for_target + from ..._jax.compiler import compile_options_for_target return compile_options_for_target(target_compute_capability) @@ -69,7 +69,6 @@ def _check_optional_signature( return if value is None: raise ValueError(f"{name} is required by the sample signature") - api._check_tensor_signature(value, expected, mode=mode) class _SparseScoreRecompute(JaxApiBase): @@ -176,12 +175,10 @@ def _default_output_desc(self) -> JaxTensorDesc: raise ValueError( f"topk_indices must have rank 3, got shape {self.topk_indices_desc.shape}" ) - return self._to_tensor_desc( - jax.ShapeDtypeStruct( - to_public_axes(self.topk_indices_desc.shape, self.output_mode), - jnp.float32, - ), - "sample_out", + return JaxTensorDesc.from_shape( + to_public_axes(self.topk_indices_desc.shape, self.output_mode), + jnp.float32, + name="sample_out", mode=self.output_mode, ) @@ -192,13 +189,6 @@ def _call( self, q: Any, k: Any, per_head: Any, topk_indices: Any, topk_length: Any | None ) -> Any: self.check_support() - for value, expected, mode in ( - (q, self.q_desc, self.q_mode), - (k, self.k_desc, self.k_mode), - (per_head, self.per_head_desc, self.per_head_mode), - (topk_indices, self.topk_indices_desc, self.score_mode), - ): - self._check_tensor_signature(value, expected, mode=mode) _check_optional_signature( self, topk_length, @@ -212,36 +202,36 @@ def _call( kernel_k_layout = self.k_layout.replace("D", "HD") kernel_k_mode = mode_from_layout(kernel_k_layout, kernel_axes="BSHD") kernel_k_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct(kernel_k.shape, kernel_k.dtype), + kernel_k, "kernel_k", mode=kernel_k_mode, ) per_head_axes = tuple(self.per_head_layout.index(axis) for axis in "BHS") kernel_per_head = jnp.transpose(per_head, per_head_axes) kernel_per_head_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct(kernel_per_head.shape, kernel_per_head.dtype), + kernel_per_head, "kernel_per_head", ) inputs = (q, kernel_k, kernel_per_head, topk_indices) - input_bindings = ( - (self.q_desc, self.q_mode), - (kernel_k_desc, kernel_k_mode), - (kernel_per_head_desc, None), - (self.topk_indices_desc, self.score_mode), + input_descs = ( + self.q_desc, + kernel_k_desc, + kernel_per_head_desc, + self.topk_indices_desc, ) else: inputs = (q, k, per_head, topk_indices) - input_bindings = ( - (self.q_desc, self.q_mode), - (self.k_desc, self.k_mode), - (self.per_head_desc, self.per_head_mode), - (self.topk_indices_desc, self.score_mode), + input_descs = ( + self.q_desc, + self.k_desc, + self.per_head_desc, + self.topk_indices_desc, ) workspace_descs = () if topk_length is not None: inputs += (topk_length,) - input_bindings += ((self.topk_length_desc, self.length_mode),) + input_descs += (self.topk_length_desc,) elif self.target_compute_capability != 90: workspace_descs = ( self.topk_indices_desc.compact_like( @@ -255,11 +245,8 @@ def _call( inputs, launch=self._launch_kernel, output_descs=(self.out_desc,), + input_descs=input_descs, workspace_descs=workspace_descs, - input_spec=tuple( - self._to_tensor_spec(desc, mode=mode) for desc, mode in input_bindings - ), - output_spec=(self._to_tensor_spec(self.out_desc, mode=self.output_mode),), compile_options=_compile_options(self.target_compute_capability), ) return output @@ -599,19 +586,17 @@ def _default_output_descs( f"Dense score inputs must both use BSHD or THD layout, got Q={self.q_desc.shape}, K={self.k_desc.shape}" ) return ( - self._to_tensor_desc( - jax.ShapeDtypeStruct( - to_public_axes(out_shape, self.output_mode), jnp.float32 - ), - "sample_out", + JaxTensorDesc.from_shape( + to_public_axes(out_shape, self.output_mode), + jnp.float32, + name="sample_out", mode=self.output_mode, init_value=float("-inf"), ), - self._to_tensor_desc( - jax.ShapeDtypeStruct( - to_public_axes(denom_shape, self.denom_mode), jnp.float32 - ), - "sample_denom_out", + JaxTensorDesc.from_shape( + to_public_axes(denom_shape, self.denom_mode), + jnp.float32, + name="sample_denom_out", mode=self.denom_mode, ), ) @@ -629,12 +614,6 @@ def _call( q_causal_offsets: Any | None, ) -> TupleDict: self.check_support() - for value, expected, mode in ( - (q, self.q_desc, self.q_mode), - (k, self.k_desc, self.k_mode), - (per_head, self.per_head_desc, self.per_head_mode), - ): - self._check_tensor_signature(value, expected, mode=mode) _check_optional_signature( self, cu_seqlens_q, self.cu_seqlens_q_desc, "cu_seqlens_q" ) @@ -649,43 +628,37 @@ def _call( per_head_axes = tuple(self.per_head_layout.index(axis) for axis in "BHS") kernel_per_head = jnp.transpose(per_head, per_head_axes) kernel_per_head_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct(kernel_per_head.shape, kernel_per_head.dtype), + kernel_per_head, "kernel_per_head", ) inputs = (q, k, kernel_per_head) - input_bindings = ( - (self.q_desc, self.q_mode), - (self.k_desc, self.k_mode), - (kernel_per_head_desc, None), + input_descs = ( + self.q_desc, + self.k_desc, + kernel_per_head_desc, ) else: inputs = (q, k, per_head) - input_bindings = ( - (self.q_desc, self.q_mode), - (self.k_desc, self.k_mode), - (self.per_head_desc, self.per_head_mode), + input_descs = ( + self.q_desc, + self.k_desc, + self.per_head_desc, ) if self.is_thd: inputs += (cu_seqlens_q, cu_seqlens_k) - input_bindings += ( - (self.cu_seqlens_q_desc, None), - (self.cu_seqlens_k_desc, None), + input_descs += ( + self.cu_seqlens_q_desc, + self.cu_seqlens_k_desc, ) if q_causal_offsets is not None: inputs += (q_causal_offsets,) - input_bindings += ((self.q_causal_offsets_desc, None),) + input_descs += (self.q_causal_offsets_desc,) output, denominator = self._call_kernel( inputs, launch=self._launch_kernel, output_descs=(self.out_desc, self.denom_desc), - input_spec=tuple( - self._to_tensor_spec(desc, mode=mode) for desc, mode in input_bindings - ), - output_spec=( - self._to_tensor_spec(self.out_desc, mode=self.output_mode), - self._to_tensor_spec(self.denom_desc, mode=self.denom_mode), - ), + input_descs=input_descs, compile_options=_compile_options(self.target_compute_capability), ) return TupleDict(out=output, denom=denominator) @@ -920,13 +893,11 @@ def sparse_indexer_score_recompute_wrapper( target_compute_capability: int | None = None, ) -> TupleDict: return SparseIndexerScoreRecompute( - jax.ShapeDtypeStruct(q_indexer.shape, q_indexer.dtype), - jax.ShapeDtypeStruct(k_indexer.shape, k_indexer.dtype), - jax.ShapeDtypeStruct(weights.shape, weights.dtype), - jax.ShapeDtypeStruct(topk_indices.shape, topk_indices.dtype), - sample_topk_length=None - if topk_length is None - else jax.ShapeDtypeStruct(topk_length.shape, topk_length.dtype), + q_indexer, + k_indexer, + weights, + topk_indices, + sample_topk_length=topk_length, qhead_per_kv_head=qhead_per_kv_head, topk_indices_global=topk_indices_global, q_layout=q_layout, @@ -969,14 +940,12 @@ def sparse_attn_score_recompute_wrapper( target_compute_capability: int | None = None, ) -> TupleDict: return SparseAttnScoreRecompute( - jax.ShapeDtypeStruct(q_attn.shape, q_attn.dtype), - jax.ShapeDtypeStruct(k_attn.shape, k_attn.dtype), - jax.ShapeDtypeStruct(lse.shape, lse.dtype), - jax.ShapeDtypeStruct(topk_indices.shape, topk_indices.dtype), + q_attn, + k_attn, + lse, + topk_indices, softmax_scale, - sample_topk_length=None - if topk_length is None - else jax.ShapeDtypeStruct(topk_length.shape, topk_length.dtype), + sample_topk_length=topk_length, qhead_per_kv_head=qhead_per_kv_head, topk_indices_global=topk_indices_global, q_layout=q_layout, @@ -1024,23 +993,17 @@ def dense_indexer_score_recompute_wrapper( target_compute_capability: int | None = None, ) -> TupleDict: return DenseIndexerScoreRecompute( - jax.ShapeDtypeStruct(q.shape, q.dtype), - jax.ShapeDtypeStruct(k.shape, k.dtype), - jax.ShapeDtypeStruct(weights.shape, weights.dtype), + q, + k, + weights, qhead_per_kv_head=qhead_per_kv_head, sm_scale=sm_scale, ratio=ratio, - sample_cu_seqlens_q=None - if cu_seqlens_q is None - else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype), - sample_cu_seqlens_k=None - if cu_seqlens_k is None - else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype), + sample_cu_seqlens_q=cu_seqlens_q, + sample_cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, - sample_q_causal_offsets=None - if q_causal_offsets is None - else jax.ShapeDtypeStruct(q_causal_offsets.shape, q_causal_offsets.dtype), + sample_q_causal_offsets=q_causal_offsets, q_layout=q_layout, k_layout=k_layout, per_head_layout=per_head_layout, @@ -1086,23 +1049,17 @@ def dense_attn_score_recompute_wrapper( target_compute_capability: int | None = None, ) -> TupleDict: return DenseAttnScoreRecompute( - jax.ShapeDtypeStruct(q.shape, q.dtype), - jax.ShapeDtypeStruct(k.shape, k.dtype), - jax.ShapeDtypeStruct(lse.shape, lse.dtype), + q, + k, + lse, softmax_scale, qhead_per_kv_head=qhead_per_kv_head, ratio=ratio, - sample_cu_seqlens_q=None - if cu_seqlens_q is None - else jax.ShapeDtypeStruct(cu_seqlens_q.shape, cu_seqlens_q.dtype), - sample_cu_seqlens_k=None - if cu_seqlens_k is None - else jax.ShapeDtypeStruct(cu_seqlens_k.shape, cu_seqlens_k.dtype), + sample_cu_seqlens_q=cu_seqlens_q, + sample_cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, - sample_q_causal_offsets=None - if q_causal_offsets is None - else jax.ShapeDtypeStruct(q_causal_offsets.shape, q_causal_offsets.dtype), + sample_q_causal_offsets=q_causal_offsets, q_layout=q_layout, k_layout=k_layout, per_head_layout=per_head_layout, diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/op.py b/python/cudnn/deepseek_sparse_attention/score_recompute/op.py index f1357c126..ede5f357d 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/op.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/op.py @@ -11,8 +11,8 @@ from typing import Any from ... import data_type -from ..._op import Op -from ..._tensor_desc import TensorDesc +from ...common.op import Op +from ...common.tensor_desc import TensorDesc from .config import ( DenseScoreKernelConfig, SparseScoreKernelConfig, diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/sparse_score_recompute_sm100.py b/python/cudnn/deepseek_sparse_attention/score_recompute/sparse_score_recompute_sm100.py index bcb9a4890..eabd8853f 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/sparse_score_recompute_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/sparse_score_recompute_sm100.py @@ -348,6 +348,7 @@ def kernel( tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams, ): """Device-side kernel entry with CLC persistent scheduling.""" + mK = copy_utils.as_global_tensor(mK) seqlen_q_packed = cute.size(mQ.shape[0]) seqlen_k = cute.size(mK.shape[0]) warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/__init__.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/__init__.py index ebd33b73b..19498935a 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/__init__.py @@ -3,7 +3,7 @@ """Lazy Torch API and framework-neutral sparse-attention backward exports.""" -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py index f769b9e04..da79f0d53 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py @@ -11,7 +11,7 @@ from ... import data_type from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict -from ..utils.compiler import compile_options_for_target +from ..._jax.compiler import compile_options_for_target from .op import BLOCK_TILE, SparseAttentionBackwardOp _SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) @@ -133,53 +133,52 @@ def __call__( topk_length: Any | None = None, ) -> TupleDict: self.check_support() - for value, expected in ( - (q, self.q_desc), - (kv, self.kv_desc), - (out, self.out_desc), - (dout, self.dout_desc), - (lse, self.lse_desc), - (attn_sink, self.attn_sink_desc), - (topk_idxs, self.topk_idxs_desc), - ): - self._check_tensor_signature(value, expected) if (topk_length is None) != (self.topk_length_desc is None): raise ValueError("topk_length presence must match sample_topk_length") - if topk_length is not None: - self._check_tensor_signature(topk_length, self.topk_length_desc) inputs = (q, kv, out, dout, lse, attn_sink, topk_idxs) - input_specs = self._input_specs() + q_desc = self.q_desc.with_divisibility( + (None, None, self._op.head_dim) + ) + kv_desc = self.kv_desc.with_divisibility( + (None, self._op.head_dim) + ) + out_desc = self.out_desc.with_divisibility( + (None, None, self._op.head_dim_v) + ) + dout_desc = self.dout_desc.with_divisibility( + (None, None, self._op.head_dim_v) + ) + input_descs = ( + q_desc, + kv_desc, + out_desc, + dout_desc, + self.lse_desc, + self.attn_sink_desc, + self.topk_idxs_desc, + ) if topk_length is not None: inputs += (topk_length,) - input_specs += (None,) + input_descs += (self.topk_length_desc,) + + dq_desc = self.dq_desc.with_divisibility( + (None, None, self._op.head_dim) + ) + dkv_desc = self.dkv_desc.with_divisibility( + (None, self._op.head_dim) + ) dq, dkv, d_sink = self._call_kernel( inputs, launch=self._launch_kernel, - output_descs=(self.dq_desc, self.dkv_desc, self.d_sink_desc), + output_descs=(dq_desc, dkv_desc, self.d_sink_desc), + input_descs=input_descs, workspace_descs=self._workspace_descs(), - input_spec=input_specs, - output_spec=( - self._to_tensor_spec(self.dq_desc, divisibility=(None, None, self._op.head_dim)), - self._to_tensor_spec(self.dkv_desc, divisibility=(None, self._op.head_dim)), - None, - ), compile_options=compile_options_for_target(self.compute_capability), ) return TupleDict(dq=dq, dkv=dkv, d_sink=d_sink) - def _input_specs(self) -> tuple[Any | None, ...]: - return ( - self._to_tensor_spec(self.q_desc, divisibility=(None, None, self._op.head_dim)), - self._to_tensor_spec(self.kv_desc, divisibility=(None, self._op.head_dim)), - self._to_tensor_spec(self.out_desc, divisibility=(None, None, self._op.head_dim_v)), - self._to_tensor_spec(self.dout_desc, divisibility=(None, None, self._op.head_dim_v)), - None, - None, - None, - ) - def _workspace_descs(self) -> tuple[JaxTensorDesc, ...]: if self._architecture_family == 100: q_rounded = _round_up(self._op.total_seqlen_q, 8) @@ -447,11 +446,15 @@ def sparse_attention_backward_wrapper( ) -> TupleDict: """Compute DeepSeek sparse-attention gradients from JAX arrays.""" - samples = tuple(jax.ShapeDtypeStruct(value.shape, value.dtype) for value in (q, kv, out, dout, lse, attn_sink, topk_idxs)) - sample_topk_length = None if topk_length is None else jax.ShapeDtypeStruct(topk_length.shape, topk_length.dtype) return SparseAttentionBackward( - *samples, - sample_topk_length=sample_topk_length, + q, + kv, + out, + dout, + lse, + attn_sink, + topk_idxs, + sample_topk_length=topk_length, softmax_scale=softmax_scale, block_tile=block_tile, target_compute_capability=target_compute_capability, diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/op.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/op.py index a28762255..cfc9bbd55 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/op.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/op.py @@ -9,8 +9,8 @@ from typing import Any, Optional from ... import data_type -from ..._op import Op -from ..._tensor_desc import TensorDesc +from ...common.op import Op +from ...common.tensor_desc import TensorDesc SUPPORTED_HEAD_DIMS = (512, 576) BLOCK_TILE = 64 diff --git a/python/cudnn/deepseek_sparse_attention/utils/compiler.py b/python/cudnn/deepseek_sparse_attention/utils/compiler.py index 9ce687138..31c9bfd5c 100644 --- a/python/cudnn/deepseek_sparse_attention/utils/compiler.py +++ b/python/cudnn/deepseek_sparse_attention/utils/compiler.py @@ -18,7 +18,7 @@ from functools import lru_cache -from ..._cute_compiler import compile_options_for_target, gpu_arch_flag_for_compute_capability +from ...common.cute_arch import gpu_arch_flag_for_compute_capability @lru_cache(maxsize=None) diff --git a/python/cudnn/deepseek_sparse_attention/utils/copy.py b/python/cudnn/deepseek_sparse_attention/utils/copy.py index 82c50a477..b2526824a 100644 --- a/python/cudnn/deepseek_sparse_attention/utils/copy.py +++ b/python/cudnn/deepseek_sparse_attention/utils/copy.py @@ -19,6 +19,35 @@ def load_s2r(src: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor: return dst +@dsl_user_op +def as_global_tensor(tensor: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor: + """Cast a generic kernel tensor back to global memory for async copies.""" + + if tensor.memspace == cute.AddressSpace.gmem: + return tensor + if tensor.memspace != cute.AddressSpace.generic: + raise ValueError(f"Expected a global or generic tensor, got {tensor.memspace}") + + pointer = tensor.iterator + global_pointer = llvm.addrspacecast( + llvm.PointerType.get(cute.AddressSpace.gmem), + pointer.llvm_ptr, + loc=loc, + ip=ip, + ) + return cute.make_tensor( + cute.make_ptr( + tensor.element_type, + global_pointer, + cute.AddressSpace.gmem, + assumed_align=pointer.alignment, + loc=loc, + ip=ip, + ), + tensor.layout, + ) + + def tiled_copy_1d(dtype: Type[cutlass.Numeric], num_threads: int, num_copy_elems: int = 1, is_async: bool = False) -> cute.TiledCopy: num_copy_bits = min(128, num_copy_elems * dtype.width) copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp() diff --git a/python/cudnn/discrete_grouped_gemm/__init__.py b/python/cudnn/discrete_grouped_gemm/__init__.py index 8ee1e9372..57c00cb64 100644 --- a/python/cudnn/discrete_grouped_gemm/__init__.py +++ b/python/cudnn/discrete_grouped_gemm/__init__.py @@ -6,7 +6,7 @@ JAX APIs are exported through :mod:`cudnn.jax`. """ -from .._operation_api import make_operation_api +from ..common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/discrete_grouped_gemm/_jax_common.py b/python/cudnn/discrete_grouped_gemm/_jax_common.py index a4a4d82a9..e1448123d 100644 --- a/python/cudnn/discrete_grouped_gemm/_jax_common.py +++ b/python/cudnn/discrete_grouped_gemm/_jax_common.py @@ -8,11 +8,10 @@ import os from typing import Any -import jax import jax.numpy as jnp from .. import data_type -from .._dense_gemm import ( +from ..gemm.helpers import ( block_scale_shape, require_16_byte_alignment, require_block_scale_layout, @@ -255,12 +254,15 @@ def _canonical_desc( *, mode: tuple[int, ...] | None = None, init_value: bool | int | float | None = None, + ptr_assumed_align: int | None = None, ) -> JaxTensorDesc: - return self._to_tensor_desc( - jax.ShapeDtypeStruct(to_public_axes(shape, mode), dtype), - name, + return JaxTensorDesc.from_shape( + to_public_axes(shape, mode), + dtype, + name=name, mode=mode, init_value=init_value, + ptr_assumed_align=ptr_assumed_align, ) def _workspace_desc(self, workspace_bytes: int) -> JaxTensorDesc: @@ -268,31 +270,26 @@ def _workspace_desc(self, workspace_bytes: int) -> JaxTensorDesc: raise ValueError( f"kernel workspace size must be positive, got {workspace_bytes}" ) - return self._canonical_desc((workspace_bytes,), jnp.uint8, "workspace") + return self._canonical_desc( + (workspace_bytes,), + jnp.uint8, + "workspace", + ptr_assumed_align=128, + ) def _materialize_output_desc( self, desc: JaxTensorDesc | None, - *, - mode: tuple[int, ...] | None = None, ) -> Any | None: """Materialize an inferred result without launching an empty GEMM.""" if desc is None: return None - metadata = self._materialize_tensor_desc(desc, mode=mode) + metadata = self._to_shape_dtype_struct(desc) if desc.init_value is None: return jnp.empty(metadata.shape, dtype=metadata.dtype) return jnp.full(metadata.shape, desc.init_value, dtype=metadata.dtype) - @staticmethod - def _workspace_tensor_spec() -> Any: - """Require the alignment used by the kernel's TMA descriptor stores.""" - - from cutlass.jax import TensorSpec - - return TensorSpec(ptr_assumed_align=128) - def _check_runtime_common( self, a_tensor: Any, @@ -302,10 +299,10 @@ def _check_runtime_common( padded_offsets: Any, alpha_tensor: Any, ) -> None: - self._check_tensor_signature(a_tensor, self.a_desc, mode=self.a_mode) - self._check_tensor_signature(b_tensor, self.b_desc, mode=self.b_mode) - self._check_tensor_signature(sfa_tensor, self.sfa_desc, mode=self.scale_mode) - self._check_tensor_signature(sfb_tensor, self.sfb_desc, mode=self.scale_mode) + self._check_tensor_signature(a_tensor, self.a_desc) + self._check_tensor_signature(b_tensor, self.b_desc) + self._check_tensor_signature(sfa_tensor, self.sfa_desc) + self._check_tensor_signature(sfb_tensor, self.sfb_desc) self._check_tensor_signature(padded_offsets, self.padded_offsets_desc) self._check_tensor_signature(alpha_tensor, self.alpha_desc) diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/__init__.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/__init__.py index 6b6ca6cf6..2fcfd2688 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/__init__.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/__init__.py @@ -3,7 +3,7 @@ """Lazy Torch API exports for discrete grouped GEMM dSwiGLU.""" -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/jax.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/jax.py index 5d9814acc..f8ecc0a6e 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/jax.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/jax.py @@ -12,8 +12,8 @@ import jax.numpy as jnp from ... import data_type -from ..._cute_compiler import compile_options_for_target -from ..._dense_gemm import ( +from ..._jax.compiler import compile_options_for_target +from ...gemm.helpers import ( block_scale_shape, require_16_byte_alignment, require_compact_major, @@ -278,11 +278,9 @@ def __call__( self._check_runtime_common( a_tensor, b_tensor, sfa_tensor, sfb_tensor, padded_offsets, alpha_tensor ) - self._check_tensor_signature(c_tensor, self.c_desc, mode=self.output_mode) + self._check_tensor_signature(c_tensor, self.c_desc) self._check_tensor_signature(beta_tensor, self.beta_desc) - self._check_tensor_signature( - prob_tensor, self.prob_desc, mode=self.probability_mode - ) + self._check_tensor_signature(prob_tensor, self.prob_desc) if (norm_const_tensor is None) != (self.norm_const_desc is None): raise ValueError("norm_const_tensor presence must match sample_norm_const") if self.norm_const_desc is not None: @@ -292,23 +290,13 @@ def __call__( if self.m == 0: return TupleDict( - d_row_tensor=self._materialize_output_desc( - self.d_row_desc, mode=self.output_mode - ), - d_col_tensor=self._materialize_output_desc( - self.d_col_desc, mode=self.output_mode - ), - dprob_tensor=self._materialize_output_desc( - self.dprob_desc, mode=self.probability_mode - ), + d_row_tensor=self._materialize_output_desc(self.d_row_desc), + d_col_tensor=self._materialize_output_desc(self.d_col_desc), + dprob_tensor=self._materialize_output_desc(self.dprob_desc), dbias_tensor=self._materialize_output_desc(self.dbias_desc), amax_tensor=self._materialize_output_desc(self.amax_desc), - sfd_row_tensor=self._materialize_output_desc( - self.sfd_row_desc, mode=self.scale_mode - ), - sfd_col_tensor=self._materialize_output_desc( - self.sfd_col_desc, mode=self.scale_mode - ), + sfd_row_tensor=self._materialize_output_desc(self.sfd_row_desc), + sfd_col_tensor=self._materialize_output_desc(self.sfd_col_desc), ) import cutlass @@ -350,36 +338,30 @@ def __call__( beta_tensor, prob_tensor, ] - input_specs = [ - self._to_tensor_spec(self.a_desc, mode=self.a_mode), - self._to_tensor_spec(self.b_desc, mode=self.b_mode), - self._to_tensor_spec(self.sfb_desc, mode=self.scale_mode), - self._to_tensor_spec(self.c_desc, mode=self.output_mode), - self._to_tensor_spec(self.sfa_desc, mode=self.scale_mode), - self._to_tensor_spec(self.padded_offsets_desc), - self._to_tensor_spec(self.alpha_desc), - self._to_tensor_spec(self.beta_desc), - self._to_tensor_spec(self.prob_desc, mode=self.probability_mode), + input_descs = [ + self.a_desc, + self.b_desc, + self.sfb_desc, + self.c_desc, + self.sfa_desc, + self.padded_offsets_desc, + self.alpha_desc, + self.beta_desc, + self.prob_desc, ] if self.norm_const_desc is not None: inputs.append(norm_const_tensor) - input_specs.append(self._to_tensor_spec(self.norm_const_desc)) + input_descs.append(self.norm_const_desc) output_descs = [self.d_row_desc, self.d_col_desc, self.dprob_desc] - output_specs = [ - self._to_tensor_spec(self.d_row_desc, mode=self.output_mode), - self._to_tensor_spec(self.d_col_desc, mode=self.output_mode), - self._to_tensor_spec(self.dprob_desc, mode=self.probability_mode), - ] - for desc, mode in ( - (self.dbias_desc, None), - (self.sfd_row_desc, self.scale_mode), - (self.sfd_col_desc, self.scale_mode), - (self.amax_desc, None), + for desc in ( + self.dbias_desc, + self.sfd_row_desc, + self.sfd_col_desc, + self.amax_desc, ): if desc is not None: output_descs.append(desc) - output_specs.append(self._to_tensor_spec(desc, mode=mode)) has_norm = self.norm_const_desc is not None has_dbias = self.dbias_desc is not None @@ -457,10 +439,8 @@ def epilogue(x): tuple(inputs), launch=launch, output_descs=tuple(output_descs), + input_descs=tuple(input_descs), workspace_descs=(workspace_desc,), - input_spec=tuple(input_specs), - output_spec=tuple(output_specs), - workspace_spec=(self._workspace_tensor_spec(),), compile_options=compile_options_for_target(self.compute_capability), ) result_index = 0 @@ -546,20 +526,16 @@ def discrete_grouped_gemm_dswiglu_wrapper_sm100( """Run discrete dSwiGLU with XLA-owned stacked expert operands.""" operation = DiscreteGroupedGemmDswigluSm100( - jax.ShapeDtypeStruct(a_tensor.shape, a_tensor.dtype), - jax.ShapeDtypeStruct(b_tensor.shape, b_tensor.dtype), - jax.ShapeDtypeStruct(c_tensor.shape, c_tensor.dtype), - jax.ShapeDtypeStruct(sfa_tensor.shape, sfa_tensor.dtype), - jax.ShapeDtypeStruct(sfb_tensor.shape, sfb_tensor.dtype), - jax.ShapeDtypeStruct(padded_offsets.shape, padded_offsets.dtype), - jax.ShapeDtypeStruct(alpha_tensor.shape, alpha_tensor.dtype), - jax.ShapeDtypeStruct(beta_tensor.shape, beta_tensor.dtype), - jax.ShapeDtypeStruct(prob_tensor.shape, prob_tensor.dtype), - sample_norm_const=( - None - if norm_const_tensor is None - else jax.ShapeDtypeStruct(norm_const_tensor.shape, norm_const_tensor.dtype) - ), + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + sample_norm_const=norm_const_tensor, d_dtype=d_dtype, acc_dtype=acc_dtype, mma_tiler_mn=mma_tiler_mn, diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/__init__.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/__init__.py index d6d7a3da6..1d2297f71 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/__init__.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/__init__.py @@ -3,7 +3,7 @@ """Lazy Torch API exports for discrete grouped GEMM + SwiGLU.""" -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/jax.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/jax.py index 1cef279d6..669b6b6f0 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/jax.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/jax.py @@ -12,8 +12,8 @@ import jax.numpy as jnp from ... import data_type -from ..._cute_compiler import compile_options_for_target -from ..._dense_gemm import ( +from ..._jax.compiler import compile_options_for_target +from ...gemm.helpers import ( block_scale_shape, require_16_byte_alignment, require_compact_major, @@ -90,12 +90,16 @@ def __init__( ) self._uses_implicit_prob = sample_prob is None if sample_prob is None: - sample_prob = jax.ShapeDtypeStruct( - (1, 1, self.a_desc.shape[0]), jnp.float32 + self.prob_desc = JaxTensorDesc.from_shape( + (1, 1, self.a_desc.shape[0]), + jnp.float32, + name="sample_prob", + mode=self.probability_mode, + ) + else: + self.prob_desc = self._to_tensor_desc( + sample_prob, "sample_prob", mode=self.probability_mode ) - self.prob_desc = self._to_tensor_desc( - sample_prob, "sample_prob", mode=self.probability_mode - ) self.bias_desc = ( None if sample_bias is None @@ -260,31 +264,19 @@ def __call__( ) prob_tensor = jnp.ones((1, 1, self.m), dtype=jnp.float32) else: - self._check_tensor_signature( - prob_tensor, self.prob_desc, mode=self.probability_mode - ) - self._check_optional_input(bias_tensor, self.bias_desc, mode=self.bias_mode) + self._check_tensor_signature(prob_tensor, self.prob_desc) + self._check_optional_input(bias_tensor, self.bias_desc) if linear_offset is None: linear_offset = 1.0 if self.act_func == "geglu" else 0.0 if self.m == 0: return TupleDict( - c_tensor=self._materialize_output_desc( - self.c_desc, mode=self.output_mode - ), - d_tensor=self._materialize_output_desc( - self.d_desc, mode=self.output_mode - ), - d_col_tensor=self._materialize_output_desc( - self.d_col_desc, mode=self.output_mode - ), + c_tensor=self._materialize_output_desc(self.c_desc), + d_tensor=self._materialize_output_desc(self.d_desc), + d_col_tensor=self._materialize_output_desc(self.d_col_desc), amax_tensor=self._materialize_output_desc(self.amax_desc), - sfd_row_tensor=self._materialize_output_desc( - self.sfd_row_desc, mode=self.scale_mode - ), - sfd_col_tensor=self._materialize_output_desc( - self.sfd_col_desc, mode=self.scale_mode - ), + sfd_row_tensor=self._materialize_output_desc(self.sfd_row_desc), + sfd_col_tensor=self._materialize_output_desc(self.sfd_col_desc), ) import cutlass @@ -324,37 +316,31 @@ def __call__( padded_offsets, alpha_tensor, ] - input_specs = [ - self._to_tensor_spec(self.a_desc, mode=self.a_mode), - self._to_tensor_spec(self.b_desc, mode=self.b_mode), - self._to_tensor_spec(self.sfb_desc, mode=self.scale_mode), - self._to_tensor_spec(self.sfa_desc, mode=self.scale_mode), - self._to_tensor_spec(self.padded_offsets_desc), - self._to_tensor_spec(self.alpha_desc), + input_descs = [ + self.a_desc, + self.b_desc, + self.sfb_desc, + self.sfa_desc, + self.padded_offsets_desc, + self.alpha_desc, ] - for value, desc, mode in ( - (norm_const_tensor, self.norm_const_desc, None), - (prob_tensor, self.prob_desc, self.probability_mode), - (bias_tensor, self.bias_desc, self.bias_mode), + for value, desc in ( + (norm_const_tensor, self.norm_const_desc), + (prob_tensor, self.prob_desc), + (bias_tensor, self.bias_desc), ): if desc is not None: inputs.append(value) - input_specs.append(self._to_tensor_spec(desc, mode=mode)) + input_descs.append(desc) output_descs = [self.c_desc, self.d_desc, self.d_col_desc] - output_specs = [ - self._to_tensor_spec(self.c_desc, mode=self.output_mode), - self._to_tensor_spec(self.d_desc, mode=self.output_mode), - self._to_tensor_spec(self.d_col_desc, mode=self.output_mode), - ] - for desc, mode in ( - (self.sfd_row_desc, self.scale_mode), - (self.sfd_col_desc, self.scale_mode), - (self.amax_desc, None), + for desc in ( + self.sfd_row_desc, + self.sfd_col_desc, + self.amax_desc, ): if desc is not None: output_descs.append(desc) - output_specs.append(self._to_tensor_spec(desc, mode=mode)) has_norm = self.norm_const_desc is not None has_prob = True @@ -419,10 +405,8 @@ def take() -> Any: tuple(inputs), launch=launch, output_descs=tuple(output_descs), + input_descs=tuple(input_descs), workspace_descs=(workspace_desc,), - input_spec=tuple(input_specs), - output_spec=tuple(output_specs), - workspace_spec=(self._workspace_tensor_spec(),), compile_options=compile_options_for_target(self.compute_capability), ) result_index = 0 @@ -450,8 +434,6 @@ def _check_optional_input( self, value: Any | None, desc: JaxTensorDesc | None, - *, - mode: tuple[int, ...] | None = None, ) -> None: if (value is None) != (desc is None): name = "optional input" if desc is None else desc.name @@ -459,7 +441,7 @@ def _check_optional_input( f"{name} presence must match the sample passed to the constructor" ) if desc is not None: - self._check_tensor_signature(value, desc, mode=mode) + self._check_tensor_signature(value, desc) @partial( @@ -518,23 +500,15 @@ def discrete_grouped_gemm_swiglu_wrapper_sm100( """Run the discrete kernel with XLA-owned stacked expert operands.""" operation = DiscreteGroupedGemmSwigluSm100( - jax.ShapeDtypeStruct(a_tensor.shape, a_tensor.dtype), - jax.ShapeDtypeStruct(b_tensor.shape, b_tensor.dtype), - jax.ShapeDtypeStruct(sfa_tensor.shape, sfa_tensor.dtype), - jax.ShapeDtypeStruct(sfb_tensor.shape, sfb_tensor.dtype), - jax.ShapeDtypeStruct(padded_offsets.shape, padded_offsets.dtype), - jax.ShapeDtypeStruct(alpha_tensor.shape, alpha_tensor.dtype), - sample_norm_const=( - None - if norm_const_tensor is None - else jax.ShapeDtypeStruct(norm_const_tensor.shape, norm_const_tensor.dtype) - ), - sample_prob=None - if prob_tensor is None - else jax.ShapeDtypeStruct(prob_tensor.shape, prob_tensor.dtype), - sample_bias=None - if bias_tensor is None - else jax.ShapeDtypeStruct(bias_tensor.shape, bias_tensor.dtype), + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + sample_norm_const=norm_const_tensor, + sample_prob=prob_tensor, + sample_bias=bias_tensor, c_dtype=c_dtype, d_dtype=d_dtype, acc_dtype=acc_dtype, diff --git a/python/cudnn/gemm/__init__.py b/python/cudnn/gemm/__init__.py new file mode 100644 index 000000000..acccf6a26 --- /dev/null +++ b/python/cudnn/gemm/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral implementation shared by dense GEMM operations.""" diff --git a/python/cudnn/_dense_gemm.py b/python/cudnn/gemm/helpers.py similarity index 99% rename from python/cudnn/_dense_gemm.py rename to python/cudnn/gemm/helpers.py index 3326d1b01..c3bd54f1c 100644 --- a/python/cudnn/_dense_gemm.py +++ b/python/cudnn/gemm/helpers.py @@ -8,8 +8,8 @@ from operator import index from typing import Iterable -from . import data_type -from ._tensor_desc import TensorDesc +from .. import data_type +from ..common.tensor_desc import TensorDesc _STANDARD_MMA_M = (128, 256) _STANDARD_MMA_N = tuple(range(32, 257, 32)) diff --git a/python/cudnn/_gemm_relu.py b/python/cudnn/gemm/srelu.py similarity index 97% rename from python/cudnn/_gemm_relu.py rename to python/cudnn/gemm/srelu.py index 2720521a2..5bdd04858 100644 --- a/python/cudnn/_gemm_relu.py +++ b/python/cudnn/gemm/srelu.py @@ -7,8 +7,10 @@ from typing import Any -from . import data_type -from ._dense_gemm import ( +from .. import data_type +from ..common.op import Op +from ..common.tensor_desc import TensorDesc +from .helpers import ( block_scale_shape, require_16_byte_alignment, require_block_scale_layout, @@ -18,8 +20,6 @@ require_mma_tiler, require_tensor_shape, ) -from ._op import Op -from ._tensor_desc import TensorDesc _AB_DTYPES = frozenset( { @@ -35,7 +35,7 @@ _C_DTYPES = _WIDE_OUTPUT_DTYPES | _FP8_DTYPES -class BlockScaledGemmReluSm100Op(Op): +class BlockScaledGemmSreluSm100OpBase(Op): """Logical tensor signature shared by sReLU forward and backward.""" def __init__( @@ -215,6 +215,3 @@ def check_support(self) -> bool: self.ab_dtype = ab_dtype self.a_major, self.b_major, self.output_major = a_major, b_major, c_major return True - - -__all__ = ["BlockScaledGemmReluSm100Op"] diff --git a/python/cudnn/gemm_amax/__init__.py b/python/cudnn/gemm_amax/__init__.py index 33844177a..70648275c 100644 --- a/python/cudnn/gemm_amax/__init__.py +++ b/python/cudnn/gemm_amax/__init__.py @@ -3,7 +3,7 @@ """Lazy Torch API and framework-neutral GEMM + amax operation exports.""" -from .._operation_api import make_operation_api +from ..common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/gemm_amax/jax.py b/python/cudnn/gemm_amax/jax.py index 71436d36d..89e52c5f0 100644 --- a/python/cudnn/gemm_amax/jax.py +++ b/python/cudnn/gemm_amax/jax.py @@ -13,8 +13,8 @@ import jax.numpy as jnp from .. import data_type -from .._cute_compiler import compile_options_for_target -from .._dense_gemm import require_gemm_inputs +from .._jax.compiler import compile_options_for_target +from ..gemm.helpers import require_gemm_inputs from .._jax import JaxApiBase, JaxTensorDesc, TupleDict from .._jax.datatypes import jax_to_cudnn_dtype, normalize_jax_dtype from .._jax.gemm import BLOCK_SCALE_MODE, gemm_a_mode, gemm_b_mode, gemm_output_mode @@ -45,30 +45,26 @@ def __init__( a_layout: str = "LMK", b_layout: str = "LNK", ) -> None: - self.a_layout = a_layout - self.b_layout = b_layout - self.c_layout = c_layout - self.a_mode = gemm_a_mode(a_layout) - self.b_mode = gemm_b_mode(b_layout) - self.output_mode = gemm_output_mode(c_layout) - self.scale_mode = BLOCK_SCALE_MODE + a_mode = gemm_a_mode(a_layout) + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(c_layout) self.compute_capability = self._resolve_compute_capability( None, SUPPORTED_COMPUTE_CAPABILITIES, "GemmAmaxSm100", ) - self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=self.a_mode) - self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=self.b_mode) + self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=a_mode) + self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=b_mode) self.sfa_desc = self._to_tensor_desc( sample_sfa, "sample_sfa", - mode=self.scale_mode, + mode=BLOCK_SCALE_MODE, ) self.sfb_desc = self._to_tensor_desc( sample_sfb, "sample_sfb", - mode=self.scale_mode, + mode=BLOCK_SCALE_MODE, ) self.acc_dtype = normalize_jax_dtype(acc_dtype, jnp.float32, "acc_dtype") @@ -76,19 +72,20 @@ def __init__( raise ValueError("sample_c and sample_amax must be provided together") if sample_c is None: resolved_c_dtype = normalize_jax_dtype(c_dtype, jnp.float32, "c_dtype") - self.c_desc, self.amax_desc = self._default_output_descs(resolved_c_dtype) + self.c_desc, self.amax_desc = self._default_output_descs(resolved_c_dtype, output_mode) else: + if c_dtype is not None: + raise ValueError("c_dtype cannot be specified with sample_c and sample_amax") self.c_desc = self._to_tensor_desc( sample_c, "sample_c", - mode=self.output_mode, + mode=output_mode, ) self.amax_desc = self._to_tensor_desc( sample_amax, "sample_amax", init_value=float("-inf"), ) - self._check_requested_c_dtype(c_dtype) acc_cudnn_dtype = jax_to_cudnn_dtype(self.acc_dtype) if acc_cudnn_dtype == data_type.NOT_SET: @@ -110,32 +107,24 @@ def __init__( def _default_output_descs( self, c_dtype: Any, + output_mode: tuple[int, ...], ) -> tuple[JaxTensorDesc, JaxTensorDesc]: m, n, _, batch = require_gemm_inputs(self.a_desc, self.b_desc) return ( - self._to_tensor_desc( - jax.ShapeDtypeStruct( - to_public_axes((m, n, batch), self.output_mode), - c_dtype, - ), - "sample_c", - mode=self.output_mode, + JaxTensorDesc.from_shape( + to_public_axes((m, n, batch), output_mode), + c_dtype, + name="sample_c", + mode=output_mode, ), - self._to_tensor_desc( - jax.ShapeDtypeStruct((1, 1, 1), jnp.float32), - "sample_amax", + JaxTensorDesc.from_shape( + (1, 1, 1), + jnp.float32, + name="sample_amax", init_value=float("-inf"), ), ) - def _check_requested_c_dtype(self, requested: Any | None) -> None: - if requested is None: - return - requested_dtype = normalize_jax_dtype(requested, requested, "c_dtype") - actual_dtype = jnp.dtype(self.c_desc.dtype) - if requested_dtype != actual_dtype: - raise ValueError(f"c_dtype={requested_dtype} does not match the explicit sample dtype " f"{actual_dtype}") - def check_support(self) -> bool: """Validate the JAX dtype surface and common operation contract.""" @@ -161,15 +150,6 @@ def check_support(self) -> bool: data_type.FP8_E4M3, }: raise ValueError(f"The JAX GEMM + amax API requires E8M0 or E4M3 scale factors, " f"got {self.sfb_desc.dtype}") - if self.c_desc.cudnn_dtype not in { - data_type.FLOAT, - data_type.HALF, - data_type.BFLOAT16, - data_type.FP8_E4M3, - data_type.FP8_E5M2, - data_type.FP4_E2M1, - }: - raise ValueError(f"The JAX GEMM + amax API received unsupported C dtype " f"{self.c_desc.dtype}") return self._op.check_support() def __call__( @@ -180,10 +160,6 @@ def __call__( sfb_tensor: Any, ) -> TupleDict: self.check_support() - self._check_tensor_signature(a_tensor, self.a_desc, mode=self.a_mode) - self._check_tensor_signature(b_tensor, self.b_desc, mode=self.b_mode) - self._check_tensor_signature(sfa_tensor, self.sfa_desc, mode=self.scale_mode) - self._check_tensor_signature(sfb_tensor, self.sfb_desc, mode=self.scale_mode) max_active_clusters = self._get_max_active_clusters( self._op.cluster_shape_mn[0] * self._op.cluster_shape_mn[1], overlap_margin=self.num_cluster_overlap_margin, @@ -213,17 +189,8 @@ def launch(stream, a, b, sfa, sfb, c, amax): c_tensor, amax_tensor = self._call_kernel( (a_tensor, b_tensor, sfa_tensor, sfb_tensor), launch=launch, + input_descs=(self.a_desc, self.b_desc, self.sfa_desc, self.sfb_desc), output_descs=(self.c_desc, self.amax_desc), - input_spec=( - self._to_tensor_spec(self.a_desc, mode=self.a_mode), - self._to_tensor_spec(self.b_desc, mode=self.b_mode), - self._to_tensor_spec(self.sfa_desc, mode=self.scale_mode), - self._to_tensor_spec(self.sfb_desc, mode=self.scale_mode), - ), - output_spec=( - self._to_tensor_spec(self.c_desc, mode=self.output_mode), - self._to_tensor_spec(self.amax_desc), - ), compile_options=compile_options_for_target(self.compute_capability), ) return TupleDict(c_tensor=c_tensor, amax_tensor=amax_tensor) @@ -260,10 +227,10 @@ def gemm_amax_wrapper_sm100( """Compute a block-scaled GEMM and global max-absolute reduction.""" return GemmAmaxSm100( - jax.ShapeDtypeStruct(a_tensor.shape, a_tensor.dtype), - jax.ShapeDtypeStruct(b_tensor.shape, b_tensor.dtype), - jax.ShapeDtypeStruct(sfa_tensor.shape, sfa_tensor.dtype), - jax.ShapeDtypeStruct(sfb_tensor.shape, sfb_tensor.dtype), + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, c_layout=c_layout, c_dtype=c_dtype, acc_dtype=acc_dtype, diff --git a/python/cudnn/gemm_amax/op.py b/python/cudnn/gemm_amax/op.py index f05b09f89..4ef637851 100644 --- a/python/cudnn/gemm_amax/op.py +++ b/python/cudnn/gemm_amax/op.py @@ -8,7 +8,7 @@ from typing import Any, Iterable, Optional from .. import data_type -from .._dense_gemm import ( +from ..gemm.helpers import ( block_scale_shape, data_type_bits, require_block_scale_layout, @@ -18,8 +18,8 @@ require_mma_tiler, require_tensor_shape, ) -from .._op import Op -from .._tensor_desc import TensorDesc +from ..common.op import Op +from ..common.tensor_desc import TensorDesc _AB_DTYPES = frozenset( { diff --git a/python/cudnn/gemm_dsrelu/__init__.py b/python/cudnn/gemm_dsrelu/__init__.py index 5fb97a9eb..b0be8c758 100644 --- a/python/cudnn/gemm_dsrelu/__init__.py +++ b/python/cudnn/gemm_dsrelu/__init__.py @@ -3,7 +3,7 @@ """Lazy Torch API and framework-neutral GEMM + dsReLU exports.""" -from .._operation_api import make_operation_api +from ..common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/gemm_dsrelu/jax.py b/python/cudnn/gemm_dsrelu/jax.py index b809958d9..f761fe8c7 100644 --- a/python/cudnn/gemm_dsrelu/jax.py +++ b/python/cudnn/gemm_dsrelu/jax.py @@ -13,8 +13,8 @@ import jax.numpy as jnp from .. import data_type -from .._cute_compiler import compile_options_for_target -from .._dense_gemm import require_gemm_inputs +from .._jax.compiler import compile_options_for_target +from ..gemm.helpers import require_gemm_inputs from .._jax import JaxApiBase, JaxTensorDesc, TupleDict from .._jax.datatypes import jax_to_cudnn_dtype, normalize_jax_dtype from .._jax.gemm import BLOCK_SCALE_MODE, PROBABILITY_MODE, gemm_a_mode, gemm_b_mode, gemm_output_mode @@ -40,7 +40,6 @@ def __init__( *, sample_d: Any | None = None, sample_dprob: Any | None = None, - sample_norm_const: Any | None = None, alpha: float = 1.0, d_layout: str = "LMN", d_dtype: Any | None = None, @@ -52,39 +51,38 @@ def __init__( a_layout: str = "LMK", b_layout: str = "LNK", ) -> None: - self.a_mode = gemm_a_mode(a_layout) - self.b_mode = gemm_b_mode(b_layout) - self.output_mode = gemm_output_mode(d_layout, name="d_layout") - self.scale_mode = BLOCK_SCALE_MODE - self.probability_mode = PROBABILITY_MODE + a_mode = gemm_a_mode(a_layout) + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(d_layout, name="d_layout") self.compute_capability = self._resolve_compute_capability( None, SUPPORTED_COMPUTE_CAPABILITIES, "GemmDsreluSm100", ) - self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=self.a_mode) - self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=self.b_mode) - self.c_desc = self._to_tensor_desc(sample_c, "sample_c", mode=self.output_mode) - self.sfa_desc = self._to_tensor_desc(sample_sfa, "sample_sfa", mode=self.scale_mode) - self.sfb_desc = self._to_tensor_desc(sample_sfb, "sample_sfb", mode=self.scale_mode) - self.prob_desc = self._to_tensor_desc(sample_prob, "sample_prob", mode=self.probability_mode) + self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=a_mode) + self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=b_mode) + self.c_desc = self._to_tensor_desc(sample_c, "sample_c", mode=output_mode) + self.sfa_desc = self._to_tensor_desc(sample_sfa, "sample_sfa", mode=BLOCK_SCALE_MODE) + self.sfb_desc = self._to_tensor_desc(sample_sfb, "sample_sfb", mode=BLOCK_SCALE_MODE) + self.prob_desc = self._to_tensor_desc(sample_prob, "sample_prob", mode=PROBABILITY_MODE) self.acc_dtype = normalize_jax_dtype(acc_dtype, jnp.float32, "acc_dtype") - resolved_d_dtype = normalize_jax_dtype(d_dtype, jnp.bfloat16, "d_dtype") if (sample_d is None) != (sample_dprob is None): raise ValueError("sample_d and sample_dprob must be provided together") if sample_d is None: - self.d_desc, self.dprob_desc = self._default_output_descs(resolved_d_dtype) + resolved_d_dtype = normalize_jax_dtype(d_dtype, jnp.bfloat16, "d_dtype") + self.d_desc, self.dprob_desc = self._default_output_descs(resolved_d_dtype, output_mode) else: - self.d_desc = self._to_tensor_desc(sample_d, "sample_d", mode=self.output_mode) + if d_dtype is not None: + raise ValueError("d_dtype cannot be specified with sample_d and sample_dprob") + self.d_desc = self._to_tensor_desc(sample_d, "sample_d", mode=output_mode) self.dprob_desc = self._to_tensor_desc( sample_dprob, "sample_dprob", - mode=self.probability_mode, + mode=PROBABILITY_MODE, init_value=0.0, ) - self._check_requested_output_dtype(d_dtype, self.d_desc, "d_dtype") self.amax_desc = None if self.a_desc.cudnn_dtype == data_type.FP4_E2M1 and self.d_desc.cudnn_dtype in _WIDE_OUTPUT_DTYPES: @@ -94,13 +92,6 @@ def __init__( name="amax_tensor", init_value=float("-inf"), ) - self.norm_const_desc = None - self.sfd_desc = None - if self.d_desc.cudnn_dtype in {data_type.FP8_E4M3, data_type.FP8_E5M2}: - raise NotImplementedError("FP8 D output is unavailable because the current dsReLU kernel does not implement SFD generation") - if sample_norm_const is not None: - raise ValueError("sample_norm_const is only used with FP8 D output, which is not implemented") - acc_cudnn_dtype = jax_to_cudnn_dtype(self.acc_dtype) if acc_cudnn_dtype == data_type.NOT_SET: raise ValueError(f"Unsupported JAX accumulator dtype {self.acc_dtype}") @@ -116,9 +107,9 @@ def __init__( sfb=self.sfb_desc, prob=self.prob_desc, dprob=self.dprob_desc, - sfd=self.sfd_desc, + sfd=None, amax=self.amax_desc, - norm_const=self.norm_const_desc, + norm_const=None, alpha=alpha, acc_dtype=acc_cudnn_dtype, mma_tiler_mn=mma_tiler_mn, @@ -128,39 +119,34 @@ def __init__( self.vector_f32 = vector_f32 self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) - def _default_output_descs(self, d_dtype: Any) -> tuple[JaxTensorDesc, JaxTensorDesc]: + def _default_output_descs( + self, + d_dtype: Any, + output_mode: tuple[int, ...], + ) -> tuple[JaxTensorDesc, JaxTensorDesc]: m, n, _, batch = require_gemm_inputs(self.a_desc, self.b_desc) - d_shape = to_public_axes((m, n, batch), self.output_mode) - dprob_shape = to_public_axes((m, 1, batch), self.probability_mode) + d_shape = to_public_axes((m, n, batch), output_mode) + dprob_shape = to_public_axes((m, 1, batch), PROBABILITY_MODE) return ( - self._to_tensor_desc( - jax.ShapeDtypeStruct(d_shape, d_dtype), - "sample_d", - mode=self.output_mode, + JaxTensorDesc.from_shape( + d_shape, + d_dtype, + name="sample_d", + mode=output_mode, ), - self._to_tensor_desc( - jax.ShapeDtypeStruct(dprob_shape, jnp.float32), - "sample_dprob", - mode=self.probability_mode, + JaxTensorDesc.from_shape( + dprob_shape, + jnp.float32, + name="sample_dprob", + mode=PROBABILITY_MODE, init_value=0.0, ), ) - @staticmethod - def _check_requested_output_dtype(requested: Any | None, desc: JaxTensorDesc, name: str) -> None: - if requested is None: - return - requested_dtype = normalize_jax_dtype(requested, requested, name) - actual_dtype = jnp.dtype(desc.dtype) - if requested_dtype != actual_dtype: - raise ValueError(f"{name}={requested_dtype} does not match the explicit sample dtype {actual_dtype}") - def check_support(self) -> bool: self._op.check_support() if self.a_desc.cudnn_dtype not in _JAX_INPUT_DTYPES: raise NotImplementedError("The JAX GEMM + dsReLU API supports native FP4 and FP8 inputs") - if self.c_desc.cudnn_dtype not in _WIDE_OUTPUT_DTYPES or self.d_desc.cudnn_dtype not in _WIDE_OUTPUT_DTYPES: - raise NotImplementedError("The JAX GEMM + dsReLU API received an unsupported output dtype") return True def __call__( @@ -171,20 +157,8 @@ def __call__( sfa_tensor: Any, sfb_tensor: Any, prob_tensor: Any, - norm_const_tensor: Any | None = None, ) -> TupleDict: self.check_support() - for value, desc, mode in ( - (a_tensor, self.a_desc, self.a_mode), - (b_tensor, self.b_desc, self.b_mode), - (c_tensor, self.c_desc, self.output_mode), - (sfa_tensor, self.sfa_desc, self.scale_mode), - (sfb_tensor, self.sfb_desc, self.scale_mode), - (prob_tensor, self.prob_desc, self.probability_mode), - ): - self._check_tensor_signature(value, desc, mode=mode) - if norm_const_tensor is not None: - raise ValueError("norm_const_tensor is only used with FP8 D output, which is not implemented") max_active_clusters = self._get_max_active_clusters( self._op.cluster_shape_mn[0] * self._op.cluster_shape_mn[1], @@ -232,23 +206,11 @@ def squared_relu_backward(x, upstream): optional_output_descs = () if self.amax_desc is None else (self.amax_desc,) output_descs = (self.d_desc, self.dprob_desc) + optional_output_descs - output_specs = ( - self._to_tensor_spec(self.d_desc, mode=self.output_mode), - self._to_tensor_spec(self.dprob_desc, mode=self.probability_mode), - ) + tuple(self._to_tensor_spec(desc) for desc in optional_output_descs) results = self._call_kernel( (a_tensor, b_tensor, c_tensor, sfa_tensor, sfb_tensor, prob_tensor), launch=launch, + input_descs=(self.a_desc, self.b_desc, self.c_desc, self.sfa_desc, self.sfb_desc, self.prob_desc), output_descs=output_descs, - input_spec=( - self._to_tensor_spec(self.a_desc, mode=self.a_mode), - self._to_tensor_spec(self.b_desc, mode=self.b_mode), - self._to_tensor_spec(self.c_desc, mode=self.output_mode), - self._to_tensor_spec(self.sfa_desc, mode=self.scale_mode), - self._to_tensor_spec(self.sfb_desc, mode=self.scale_mode), - self._to_tensor_spec(self.prob_desc, mode=self.probability_mode), - ), - output_spec=output_specs, compile_options=compile_options_for_target(self.compute_capability), ) d_tensor, dprob_tensor, *optional_results = results @@ -288,7 +250,6 @@ def gemm_dsrelu_wrapper_sm100( acc_dtype: Any | None = None, mma_tiler_mn: tuple[int, int] = (256, 256), cluster_shape_mn: tuple[int, int] | None = None, - norm_const_tensor: Any | None = None, sf_vec_size: int = 16, vector_f32: bool = False, *, @@ -298,13 +259,12 @@ def gemm_dsrelu_wrapper_sm100( """Compute the block-scaled squared-ReLU backward fusion.""" return GemmDsreluSm100( - jax.ShapeDtypeStruct(a_tensor.shape, a_tensor.dtype), - jax.ShapeDtypeStruct(b_tensor.shape, b_tensor.dtype), - jax.ShapeDtypeStruct(c_tensor.shape, c_tensor.dtype), - jax.ShapeDtypeStruct(sfa_tensor.shape, sfa_tensor.dtype), - jax.ShapeDtypeStruct(sfb_tensor.shape, sfb_tensor.dtype), - jax.ShapeDtypeStruct(prob_tensor.shape, prob_tensor.dtype), - sample_norm_const=(None if norm_const_tensor is None else jax.ShapeDtypeStruct(norm_const_tensor.shape, norm_const_tensor.dtype)), + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + prob_tensor, alpha=alpha, d_layout=d_layout, d_dtype=d_dtype, @@ -315,7 +275,7 @@ def gemm_dsrelu_wrapper_sm100( vector_f32=vector_f32, a_layout=a_layout, b_layout=b_layout, - )(a_tensor, b_tensor, c_tensor, sfa_tensor, sfb_tensor, prob_tensor, norm_const_tensor) + )(a_tensor, b_tensor, c_tensor, sfa_tensor, sfb_tensor, prob_tensor) __all__ = [ diff --git a/python/cudnn/gemm_dsrelu/op.py b/python/cudnn/gemm_dsrelu/op.py index 6046628ab..608c5f134 100644 --- a/python/cudnn/gemm_dsrelu/op.py +++ b/python/cudnn/gemm_dsrelu/op.py @@ -4,10 +4,10 @@ """Framework-neutral dense GEMM + squared-ReLU backward operation.""" from .. import data_type -from .._gemm_relu import BlockScaledGemmReluSm100Op +from ..gemm.srelu import BlockScaledGemmSreluSm100OpBase -class GemmDsreluSm100Op(BlockScaledGemmReluSm100Op): +class GemmDsreluSm100Op(BlockScaledGemmSreluSm100OpBase): """Logical signature for block-scaled GEMM + squared-ReLU backward.""" def check_support(self) -> bool: diff --git a/python/cudnn/gemm_srelu/__init__.py b/python/cudnn/gemm_srelu/__init__.py index 11913e2a5..59aa5689f 100644 --- a/python/cudnn/gemm_srelu/__init__.py +++ b/python/cudnn/gemm_srelu/__init__.py @@ -3,7 +3,7 @@ """Lazy Torch API and framework-neutral GEMM + sReLU exports.""" -from .._operation_api import make_operation_api +from ..common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/gemm_srelu/jax.py b/python/cudnn/gemm_srelu/jax.py index 75bfd8b06..efab492d8 100644 --- a/python/cudnn/gemm_srelu/jax.py +++ b/python/cudnn/gemm_srelu/jax.py @@ -13,8 +13,8 @@ import jax.numpy as jnp from .. import data_type -from .._cute_compiler import compile_options_for_target -from .._dense_gemm import require_gemm_inputs +from .._jax.compiler import compile_options_for_target +from ..gemm.helpers import require_gemm_inputs from .._jax import JaxApiBase, JaxTensorDesc, TupleDict from .._jax.datatypes import jax_to_cudnn_dtype, normalize_jax_dtype from .._jax.gemm import BLOCK_SCALE_MODE, PROBABILITY_MODE, gemm_a_mode, gemm_b_mode, gemm_output_mode @@ -24,7 +24,6 @@ SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) _JAX_INPUT_DTYPES = frozenset({data_type.FP4_E2M1, data_type.FP8_E4M3, data_type.FP8_E5M2}) _WIDE_OUTPUT_DTYPES = frozenset({data_type.HALF, data_type.BFLOAT16, data_type.FLOAT}) -_JAX_C_DTYPES = _WIDE_OUTPUT_DTYPES | {data_type.FP8_E4M3, data_type.FP8_E5M2} class GemmSreluSm100(JaxApiBase): @@ -40,7 +39,6 @@ def __init__( *, sample_c: Any | None = None, sample_d: Any | None = None, - sample_norm_const: Any | None = None, alpha: float = 1.0, c_layout: str = "LMN", c_dtype: Any | None = None, @@ -53,38 +51,37 @@ def __init__( a_layout: str = "LMK", b_layout: str = "LNK", ) -> None: - self.a_mode = gemm_a_mode(a_layout) - self.b_mode = gemm_b_mode(b_layout) - self.output_mode = gemm_output_mode(c_layout) - self.scale_mode = BLOCK_SCALE_MODE - self.probability_mode = PROBABILITY_MODE + a_mode = gemm_a_mode(a_layout) + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(c_layout) self.compute_capability = self._resolve_compute_capability( None, SUPPORTED_COMPUTE_CAPABILITIES, "GemmSreluSm100", ) - self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=self.a_mode) - self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=self.b_mode) - self.sfa_desc = self._to_tensor_desc(sample_sfa, "sample_sfa", mode=self.scale_mode) - self.sfb_desc = self._to_tensor_desc(sample_sfb, "sample_sfb", mode=self.scale_mode) - self.prob_desc = self._to_tensor_desc(sample_prob, "sample_prob", mode=self.probability_mode) + self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=a_mode) + self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=b_mode) + self.sfa_desc = self._to_tensor_desc(sample_sfa, "sample_sfa", mode=BLOCK_SCALE_MODE) + self.sfb_desc = self._to_tensor_desc(sample_sfb, "sample_sfb", mode=BLOCK_SCALE_MODE) + self.prob_desc = self._to_tensor_desc(sample_prob, "sample_prob", mode=PROBABILITY_MODE) self.acc_dtype = normalize_jax_dtype(acc_dtype, jnp.float32, "acc_dtype") - resolved_c_dtype = normalize_jax_dtype(c_dtype, jnp.bfloat16, "c_dtype") - resolved_d_dtype = normalize_jax_dtype(d_dtype, jnp.bfloat16, "d_dtype") if (sample_c is None) != (sample_d is None): raise ValueError("sample_c and sample_d must be provided together") if sample_c is None: + resolved_c_dtype = normalize_jax_dtype(c_dtype, jnp.bfloat16, "c_dtype") + resolved_d_dtype = normalize_jax_dtype(d_dtype, jnp.bfloat16, "d_dtype") self.c_desc, self.d_desc = self._default_output_descs( resolved_c_dtype, resolved_d_dtype, + output_mode, ) else: - self.c_desc = self._to_tensor_desc(sample_c, "sample_c", mode=self.output_mode) - self.d_desc = self._to_tensor_desc(sample_d, "sample_d", mode=self.output_mode) - self._check_requested_output_dtype(c_dtype, self.c_desc, "c_dtype") - self._check_requested_output_dtype(d_dtype, self.d_desc, "d_dtype") + if c_dtype is not None or d_dtype is not None: + raise ValueError("c_dtype and d_dtype cannot be specified with sample_c and sample_d") + self.c_desc = self._to_tensor_desc(sample_c, "sample_c", mode=output_mode) + self.d_desc = self._to_tensor_desc(sample_d, "sample_d", mode=output_mode) self.amax_desc = None if self.a_desc.cudnn_dtype == data_type.FP4_E2M1 and self.d_desc.cudnn_dtype in _WIDE_OUTPUT_DTYPES: @@ -94,13 +91,6 @@ def __init__( name="amax_tensor", init_value=float("-inf"), ) - self.norm_const_desc = None - self.sfd_desc = None - if self.d_desc.cudnn_dtype in {data_type.FP8_E4M3, data_type.FP8_E5M2}: - raise NotImplementedError("FP8 D output is unavailable because the current sReLU kernel does not implement SFD generation") - if sample_norm_const is not None: - raise ValueError("sample_norm_const is only used with FP8 D output, which is not implemented") - acc_cudnn_dtype = jax_to_cudnn_dtype(self.acc_dtype) if acc_cudnn_dtype == data_type.NOT_SET: raise ValueError(f"Unsupported JAX accumulator dtype {self.acc_dtype}") @@ -115,9 +105,9 @@ def __init__( sfa=self.sfa_desc, sfb=self.sfb_desc, prob=self.prob_desc, - sfd=self.sfd_desc, + sfd=None, amax=self.amax_desc, - norm_const=self.norm_const_desc, + norm_const=None, alpha=alpha, acc_dtype=acc_cudnn_dtype, mma_tiler_mn=mma_tiler_mn, @@ -127,37 +117,33 @@ def __init__( self.vector_f32 = vector_f32 self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) - def _default_output_descs(self, c_dtype: Any, d_dtype: Any) -> tuple[JaxTensorDesc, JaxTensorDesc]: + def _default_output_descs( + self, + c_dtype: Any, + d_dtype: Any, + output_mode: tuple[int, ...], + ) -> tuple[JaxTensorDesc, JaxTensorDesc]: m, n, _, batch = require_gemm_inputs(self.a_desc, self.b_desc) - public_shape = to_public_axes((m, n, batch), self.output_mode) + public_shape = to_public_axes((m, n, batch), output_mode) return ( - self._to_tensor_desc( - jax.ShapeDtypeStruct(public_shape, c_dtype), - "sample_c", - mode=self.output_mode, + JaxTensorDesc.from_shape( + public_shape, + c_dtype, + name="sample_c", + mode=output_mode, ), - self._to_tensor_desc( - jax.ShapeDtypeStruct(public_shape, d_dtype), - "sample_d", - mode=self.output_mode, + JaxTensorDesc.from_shape( + public_shape, + d_dtype, + name="sample_d", + mode=output_mode, ), ) - @staticmethod - def _check_requested_output_dtype(requested: Any | None, desc: JaxTensorDesc, name: str) -> None: - if requested is None: - return - requested_dtype = normalize_jax_dtype(requested, requested, name) - actual_dtype = jnp.dtype(desc.dtype) - if requested_dtype != actual_dtype: - raise ValueError(f"{name}={requested_dtype} does not match the explicit sample dtype {actual_dtype}") - def check_support(self) -> bool: self._op.check_support() if self.a_desc.cudnn_dtype not in _JAX_INPUT_DTYPES: raise NotImplementedError("The JAX GEMM + sReLU API supports native FP4 and FP8 inputs") - if self.c_desc.cudnn_dtype not in _JAX_C_DTYPES or self.d_desc.cudnn_dtype not in _WIDE_OUTPUT_DTYPES: - raise NotImplementedError("The JAX GEMM + sReLU API received an unsupported output dtype") return True def __call__( @@ -167,19 +153,8 @@ def __call__( sfa_tensor: Any, sfb_tensor: Any, prob_tensor: Any, - norm_const_tensor: Any | None = None, ) -> TupleDict: self.check_support() - for value, desc, mode in ( - (a_tensor, self.a_desc, self.a_mode), - (b_tensor, self.b_desc, self.b_mode), - (sfa_tensor, self.sfa_desc, self.scale_mode), - (sfb_tensor, self.sfb_desc, self.scale_mode), - (prob_tensor, self.prob_desc, self.probability_mode), - ): - self._check_tensor_signature(value, desc, mode=mode) - if norm_const_tensor is not None: - raise ValueError("norm_const_tensor is only used with FP8 D output, which is not implemented") max_active_clusters = self._get_max_active_clusters( self._op.cluster_shape_mn[0] * self._op.cluster_shape_mn[1], @@ -226,22 +201,11 @@ def squared_relu(x): optional_output_descs = () if self.amax_desc is None else (self.amax_desc,) output_descs = (self.c_desc, self.d_desc) + optional_output_descs - output_specs = ( - self._to_tensor_spec(self.c_desc, mode=self.output_mode), - self._to_tensor_spec(self.d_desc, mode=self.output_mode), - ) + tuple(self._to_tensor_spec(desc) for desc in optional_output_descs) results = self._call_kernel( (a_tensor, b_tensor, sfa_tensor, sfb_tensor, prob_tensor), launch=launch, + input_descs=(self.a_desc, self.b_desc, self.sfa_desc, self.sfb_desc, self.prob_desc), output_descs=output_descs, - input_spec=( - self._to_tensor_spec(self.a_desc, mode=self.a_mode), - self._to_tensor_spec(self.b_desc, mode=self.b_mode), - self._to_tensor_spec(self.sfa_desc, mode=self.scale_mode), - self._to_tensor_spec(self.sfb_desc, mode=self.scale_mode), - self._to_tensor_spec(self.prob_desc, mode=self.probability_mode), - ), - output_spec=output_specs, compile_options=compile_options_for_target(self.compute_capability), ) c_tensor, d_tensor, *optional_results = results @@ -282,7 +246,6 @@ def gemm_srelu_wrapper_sm100( acc_dtype: Any | None = None, mma_tiler_mn: tuple[int, int] = (256, 256), cluster_shape_mn: tuple[int, int] | None = None, - norm_const_tensor: Any | None = None, sf_vec_size: int = 16, vector_f32: bool = False, *, @@ -292,12 +255,11 @@ def gemm_srelu_wrapper_sm100( """Compute block-scaled GEMM and ``D = prob * relu(C) ** 2``.""" return GemmSreluSm100( - jax.ShapeDtypeStruct(a_tensor.shape, a_tensor.dtype), - jax.ShapeDtypeStruct(b_tensor.shape, b_tensor.dtype), - jax.ShapeDtypeStruct(sfa_tensor.shape, sfa_tensor.dtype), - jax.ShapeDtypeStruct(sfb_tensor.shape, sfb_tensor.dtype), - jax.ShapeDtypeStruct(prob_tensor.shape, prob_tensor.dtype), - sample_norm_const=(None if norm_const_tensor is None else jax.ShapeDtypeStruct(norm_const_tensor.shape, norm_const_tensor.dtype)), + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + prob_tensor, alpha=alpha, c_layout=c_layout, c_dtype=c_dtype, @@ -309,7 +271,7 @@ def gemm_srelu_wrapper_sm100( vector_f32=vector_f32, a_layout=a_layout, b_layout=b_layout, - )(a_tensor, b_tensor, sfa_tensor, sfb_tensor, prob_tensor, norm_const_tensor) + )(a_tensor, b_tensor, sfa_tensor, sfb_tensor, prob_tensor) __all__ = [ diff --git a/python/cudnn/gemm_srelu/op.py b/python/cudnn/gemm_srelu/op.py index 538aee587..b8f3194a4 100644 --- a/python/cudnn/gemm_srelu/op.py +++ b/python/cudnn/gemm_srelu/op.py @@ -3,10 +3,10 @@ """Framework-neutral dense GEMM + squared-ReLU operation.""" -from .._gemm_relu import BlockScaledGemmReluSm100Op +from ..gemm.srelu import BlockScaledGemmSreluSm100OpBase -class GemmSreluSm100Op(BlockScaledGemmReluSm100Op): +class GemmSreluSm100Op(BlockScaledGemmSreluSm100OpBase): """Logical signature for block-scaled GEMM + squared ReLU.""" def check_support(self) -> bool: diff --git a/python/cudnn/gemm_swiglu/__init__.py b/python/cudnn/gemm_swiglu/__init__.py index 917c2fa19..239e6c8b2 100644 --- a/python/cudnn/gemm_swiglu/__init__.py +++ b/python/cudnn/gemm_swiglu/__init__.py @@ -3,7 +3,7 @@ """Lazy Torch API and framework-neutral GEMM + SwiGLU operation exports.""" -from .._operation_api import make_operation_api +from ..common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/gemm_swiglu/jax.py b/python/cudnn/gemm_swiglu/jax.py index 1b138c7f1..47d998071 100644 --- a/python/cudnn/gemm_swiglu/jax.py +++ b/python/cudnn/gemm_swiglu/jax.py @@ -13,8 +13,8 @@ import jax.numpy as jnp from .. import data_type -from .._cute_compiler import compile_options_for_target -from .._dense_gemm import require_gemm_inputs +from .._jax.compiler import compile_options_for_target +from ..gemm.helpers import require_gemm_inputs from .._jax import JaxApiBase, JaxTensorDesc, TupleDict from .._jax.datatypes import jax_to_cudnn_dtype, normalize_jax_dtype from .._jax.gemm import BLOCK_SCALE_MODE, gemm_a_mode, gemm_b_mode, gemm_output_mode @@ -45,9 +45,7 @@ def __init__( sample_c: Any | None = None, sample_sfa: Any | None = None, sample_sfb: Any | None = None, - sample_sfc: Any | None = None, sample_amax: Any | None = None, - sample_norm_const: Any | None = None, alpha: float = 1.0, c_layout: str = "LMN", ab12_dtype: Any | None = None, @@ -61,25 +59,18 @@ def __init__( a_layout: str = "LMK", b_layout: str = "LNK", ) -> None: - self.a_layout = a_layout - self.b_layout = b_layout - self.c_layout = c_layout - self.a_mode = gemm_a_mode(a_layout) - self.b_mode = gemm_b_mode(b_layout) - self.output_mode = gemm_output_mode(c_layout) - self.scale_mode = BLOCK_SCALE_MODE - self.sf_vec_size = sf_vec_size + a_mode = gemm_a_mode(a_layout) + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(c_layout) self.compute_capability = self._resolve_compute_capability( None, SUPPORTED_COMPUTE_CAPABILITIES, "GemmSwigluSm100", ) - self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=self.a_mode) - self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=self.b_mode) + self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=a_mode) + self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=b_mode) self.acc_dtype = normalize_jax_dtype(acc_dtype, jnp.float32, "acc_dtype") - if sample_sfc is not None or sample_norm_const is not None: - raise NotImplementedError("The JAX GEMM + SwiGLU API does not expose FP8 C or SFC because the current supported input combinations reject FP8 C") self.is_block_scaled = any( sample is not None @@ -98,9 +89,8 @@ def __init__( raise ValueError("vector_f32 only applies to block-scaled GEMM + SwiGLU") if ab12_stages != 4: raise ValueError("ab12_stages only applies to block-scaled GEMM + SwiGLU") - self.sfa_desc = None if sample_sfa is None else self._to_tensor_desc(sample_sfa, "sample_sfa", mode=self.scale_mode) - self.sfb_desc = None if sample_sfb is None else self._to_tensor_desc(sample_sfb, "sample_sfb", mode=self.scale_mode) - self.norm_const_desc = None + self.sfa_desc = None if sample_sfa is None else self._to_tensor_desc(sample_sfa, "sample_sfa", mode=BLOCK_SCALE_MODE) + self.sfb_desc = None if sample_sfb is None else self._to_tensor_desc(sample_sfb, "sample_sfb", mode=BLOCK_SCALE_MODE) if (sample_ab12 is None) != (sample_c is None): raise ValueError("sample_ab12 and sample_c must be provided together") @@ -110,25 +100,23 @@ def __init__( self.ab12_desc, self.c_desc = self._default_output_descs( resolved_ab12_dtype, resolved_c_dtype, + output_mode, ) else: + if ab12_dtype is not None or c_dtype is not None: + raise ValueError("ab12_dtype and c_dtype cannot be specified with sample_ab12 and sample_c") self.ab12_desc = self._to_tensor_desc( sample_ab12, "sample_ab12", - mode=self.output_mode, + mode=output_mode, ) self.c_desc = self._to_tensor_desc( sample_c, "sample_c", - mode=self.output_mode, + mode=output_mode, ) - self._check_requested_output_dtype(ab12_dtype, self.ab12_desc, "ab12_dtype") - self._check_requested_output_dtype(c_dtype, self.c_desc, "c_dtype") - self.sfc_desc = None self.amax_desc = None - if self.c_desc.cudnn_dtype in {data_type.FP8_E4M3, data_type.FP8_E5M2}: - raise NotImplementedError("The JAX GEMM + SwiGLU API does not support FP8 C output") if self.is_block_scaled: needs_amax = self.a_desc.cudnn_dtype == data_type.FP4_E2M1 and self.c_desc.cudnn_dtype == data_type.BFLOAT16 if sample_amax is not None: @@ -148,9 +136,9 @@ def __init__( sfb=self.sfb_desc, ab12=self.ab12_desc, c=self.c_desc, - sfc=self.sfc_desc, + sfc=None, amax=self.amax_desc, - norm_const=self.norm_const_desc, + norm_const=None, alpha=alpha, acc_dtype=acc_cudnn_dtype, mma_tiler_mn=mma_tiler_mn, @@ -176,50 +164,35 @@ def _default_output_descs( self, ab12_dtype: Any, c_dtype: Any, + output_mode: tuple[int, ...], ) -> tuple[JaxTensorDesc, JaxTensorDesc]: m, n, _, batch = require_gemm_inputs(self.a_desc, self.b_desc) if n % 2: raise ValueError(f"SwiGLU requires an even N dimension, got {n}") return ( - self._to_tensor_desc( - jax.ShapeDtypeStruct( - to_public_axes((m, n, batch), self.output_mode), - ab12_dtype, - ), - "sample_ab12", - mode=self.output_mode, + JaxTensorDesc.from_shape( + to_public_axes((m, n, batch), output_mode), + ab12_dtype, + name="sample_ab12", + mode=output_mode, ), - self._to_tensor_desc( - jax.ShapeDtypeStruct( - to_public_axes((m, n // 2, batch), self.output_mode), - c_dtype, - ), - "sample_c", - mode=self.output_mode, + JaxTensorDesc.from_shape( + to_public_axes((m, n // 2, batch), output_mode), + c_dtype, + name="sample_c", + mode=output_mode, ), ) def _default_amax_desc(self) -> JaxTensorDesc: - return self._to_tensor_desc( - jax.ShapeDtypeStruct((1,), jnp.float32), - "sample_amax", + return JaxTensorDesc.from_shape( + (1,), + jnp.float32, + name="sample_amax", init_value=float("-inf"), ) - @staticmethod - def _check_requested_output_dtype( - requested: Any | None, - desc: JaxTensorDesc, - name: str, - ) -> None: - if requested is None: - return - requested_dtype = normalize_jax_dtype(requested, requested, name) - actual_dtype = jnp.dtype(desc.dtype) - if requested_dtype != actual_dtype: - raise ValueError(f"{name}={requested_dtype} does not match the explicit sample dtype {actual_dtype}") - def check_support(self) -> bool: return self._op.check_support() @@ -229,19 +202,12 @@ def __call__( b_tensor: Any, sfa_tensor: Any | None = None, sfb_tensor: Any | None = None, - norm_const_tensor: Any | None = None, ) -> TupleDict: self.check_support() - self._check_tensor_signature(a_tensor, self.a_desc, mode=self.a_mode) - self._check_tensor_signature(b_tensor, self.b_desc, mode=self.b_mode) if self.is_block_scaled: if sfa_tensor is None or sfb_tensor is None: raise ValueError("sfa_tensor and sfb_tensor are required for this block-scaled callable") - self._check_tensor_signature(sfa_tensor, self.sfa_desc, mode=self.scale_mode) - self._check_tensor_signature(sfb_tensor, self.sfb_desc, mode=self.scale_mode) - if norm_const_tensor is not None: - raise NotImplementedError("norm_const_tensor is unavailable because the JAX API does not support FP8 C output") - elif sfa_tensor is not None or sfb_tensor is not None or norm_const_tensor is not None: + elif sfa_tensor is not None or sfb_tensor is not None: raise ValueError("Scale-factor tensors are not part of this standard GEMM + SwiGLU callable") max_active_clusters = self._get_max_active_clusters( @@ -282,15 +248,8 @@ def launch( ab12_tensor, c_tensor = self._call_kernel( (a_tensor, b_tensor), launch=launch, + input_descs=(self.a_desc, self.b_desc), output_descs=(self.ab12_desc, self.c_desc), - input_spec=( - self._to_tensor_spec(self.a_desc, mode=self.a_mode), - self._to_tensor_spec(self.b_desc, mode=self.b_mode), - ), - output_spec=( - self._to_tensor_spec(self.ab12_desc, mode=self.output_mode), - self._to_tensor_spec(self.c_desc, mode=self.output_mode), - ), compile_options=compile_options_for_target(self.compute_capability), ) return TupleDict( @@ -301,21 +260,11 @@ def launch( ) kernel_inputs = (a_tensor, b_tensor, sfa_tensor, sfb_tensor) - input_specs = ( - self._to_tensor_spec(self.a_desc, mode=self.a_mode), - self._to_tensor_spec(self.b_desc, mode=self.b_mode), - self._to_tensor_spec(self.sfa_desc, mode=self.scale_mode), - self._to_tensor_spec(self.sfb_desc, mode=self.scale_mode), - ) + input_descs = (self.a_desc, self.b_desc, self.sfa_desc, self.sfb_desc) output_descs = (self.ab12_desc, self.c_desc) - output_specs = ( - self._to_tensor_spec(self.ab12_desc, mode=self.output_mode), - self._to_tensor_spec(self.c_desc, mode=self.output_mode), - ) if self.amax_desc is not None: output_descs += (self.amax_desc,) - output_specs += (self._to_tensor_spec(self.amax_desc),) def launch_block_scaled( stream: Any, @@ -361,9 +310,8 @@ def launch_block_scaled( results = self._call_kernel( kernel_inputs, launch=launch_block_scaled, + input_descs=input_descs, output_descs=output_descs, - input_spec=input_specs, - output_spec=output_specs, compile_options=compile_options_for_target(self.compute_capability), ) ab12_tensor, c_tensor, *optional_results = results @@ -405,7 +353,6 @@ def gemm_swiglu_wrapper_sm100( cluster_shape_mn: tuple[int, int] | None = None, sfa_tensor: Any | None = None, sfb_tensor: Any | None = None, - norm_const_tensor: Any | None = None, sf_vec_size: int = 16, vector_f32: bool = False, ab12_stages: int = 4, @@ -416,11 +363,10 @@ def gemm_swiglu_wrapper_sm100( """Compute standard or block-scaled GEMM + SwiGLU on a local SM100-family GPU.""" return GemmSwigluSm100( - jax.ShapeDtypeStruct(a_tensor.shape, a_tensor.dtype), - jax.ShapeDtypeStruct(b_tensor.shape, b_tensor.dtype), - sample_sfa=None if sfa_tensor is None else jax.ShapeDtypeStruct(sfa_tensor.shape, sfa_tensor.dtype), - sample_sfb=None if sfb_tensor is None else jax.ShapeDtypeStruct(sfb_tensor.shape, sfb_tensor.dtype), - sample_norm_const=None if norm_const_tensor is None else jax.ShapeDtypeStruct(norm_const_tensor.shape, norm_const_tensor.dtype), + a_tensor, + b_tensor, + sample_sfa=sfa_tensor, + sample_sfb=sfb_tensor, alpha=alpha, c_layout=c_layout, ab12_dtype=ab12_dtype, @@ -433,7 +379,7 @@ def gemm_swiglu_wrapper_sm100( ab12_stages=ab12_stages, a_layout=a_layout, b_layout=b_layout, - )(a_tensor, b_tensor, sfa_tensor, sfb_tensor, norm_const_tensor) + )(a_tensor, b_tensor, sfa_tensor, sfb_tensor) __all__ = [ diff --git a/python/cudnn/gemm_swiglu/op.py b/python/cudnn/gemm_swiglu/op.py index 81f5cc423..3dfde75a1 100644 --- a/python/cudnn/gemm_swiglu/op.py +++ b/python/cudnn/gemm_swiglu/op.py @@ -8,7 +8,7 @@ from typing import Any, Optional from .. import data_type -from .._dense_gemm import ( +from ..gemm.helpers import ( block_scale_shape, require_16_byte_alignment, require_block_scale_layout, @@ -18,8 +18,8 @@ require_mma_tiler, require_tensor_shape, ) -from .._op import Op -from .._tensor_desc import TensorDesc +from ..common.op import Op +from ..common.tensor_desc import TensorDesc _STANDARD_AB_DTYPES = frozenset( { diff --git a/python/cudnn/grouped_gemm/__init__.py b/python/cudnn/grouped_gemm/__init__.py index 4a0401461..9830fbd84 100644 --- a/python/cudnn/grouped_gemm/__init__.py +++ b/python/cudnn/grouped_gemm/__init__.py @@ -3,7 +3,7 @@ """Lazy Torch API exports for the grouped GEMM operation family.""" -from .._operation_api import make_operation_api +from ..common.operation_api import make_operation_api _OPERATIONS = ( "swiglu", diff --git a/python/cudnn/grouped_gemm/_jax_api.py b/python/cudnn/grouped_gemm/_jax_api.py index 03d01dcfb..82db2a9ee 100644 --- a/python/cudnn/grouped_gemm/_jax_api.py +++ b/python/cudnn/grouped_gemm/_jax_api.py @@ -6,15 +6,12 @@ from __future__ import annotations from collections.abc import Iterable, Mapping, Sequence -from dataclasses import dataclass, field from typing import Any -import jax import jax.numpy as jnp -from cutlass.jax import TensorSpec -from .._cute_compiler import compile_options_for_target -from .._dense_gemm import ( +from .._jax.compiler import compile_options_for_target +from ..gemm.helpers import ( block_scale_shape as _canonical_block_scale_shape, data_type_bits, ) @@ -136,85 +133,9 @@ def require_layout(name: str, value: str, supported: tuple[str, ...]) -> str: return value -def _row_major_spec(rank: int, mode: tuple[int, ...]) -> TensorSpec: - return TensorSpec(layout=tuple(reversed(range(rank))), mode=mode) - - -def gemm_a_tensor_spec(layout: str) -> TensorSpec: - return _row_major_spec(3, gemm_a_mode(layout)) - - -def gemm_b_tensor_spec(layout: str) -> TensorSpec: - return _row_major_spec(3, gemm_b_mode(layout)) - - -def gemm_c_tensor_spec(layout: str, *, name: str = "c_layout") -> TensorSpec: - return _row_major_spec( - 3, gemm_output_mode(require_layout(name, layout, ("LMN", "LNM")), name=name) - ) - - -def block_scale_tensor_spec() -> TensorSpec: - """Describe public row-major ``L,tiles,tiles,32,4,4`` scale arrays.""" - - return _row_major_spec(6, BLOCK_SCALE_MODE) - - -def probability_tensor_spec() -> TensorSpec: - """Map public row-major ``(L, 1, M)`` to canonical ``(M, 1, L)``.""" - - return _row_major_spec(3, (2, 1, 0)) - - -def grouped_bias_tensor_spec() -> TensorSpec: - """Map public row-major ``(L, N)`` to canonical ``(N, L)``.""" - - return _row_major_spec(2, (1, 0)) - - -def grouped_workspace_tensor_spec() -> TensorSpec: - return TensorSpec(ptr_assumed_align=128) - - -@dataclass(frozen=True) -class _SampleDesc(JaxTensorDesc): - """JAX descriptor retaining the native CUTLASS lowering specification.""" - - tensor_spec: Any = None - jax_mode: tuple[int, ...] = field(default_factory=tuple) - - @property - def array_shape(self) -> tuple[int, ...]: - return to_public_axes(self.shape, self.jax_mode) - - -def _make_desc( - value: Any, - *, - tensor_spec: TensorSpec | None, - name: str, - init_value: bool | int | float | None = None, -) -> _SampleDesc: - public_shape = require_array(value, name=name) - mode = normalize_mode( - len(public_shape), None if tensor_spec is None else tensor_spec.mode - ) - desc = JaxApiBase._to_tensor_desc( - value, - name, - mode=mode, - init_value=init_value, - ) - return _SampleDesc( - dtype=desc.dtype, - shape=desc.shape, - stride=desc.stride, - stride_order=desc.stride_order, - name=desc.name, - init_value=desc.init_value, - tensor_spec=tensor_spec, - jax_mode=mode, - ) +PROBABILITY_MODE = (2, 1, 0) +GROUPED_BIAS_MODE = (1, 0) +GROUPED_WORKSPACE_ALIGNMENT = 128 def make_buffer_desc( @@ -222,29 +143,75 @@ def make_buffer_desc( shape: Sequence[int], dtype: Any, *, - tensor_spec: TensorSpec | None = None, + mode: tuple[int, ...] | None = None, + public_stride_order: tuple[int, ...] | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, init_value: bool | int | float | None = None, -) -> _SampleDesc: +) -> JaxTensorDesc: """Describe an inferred JAX output or workspace buffer.""" if not name: raise ValueError("buffer descriptor name must not be empty") - return _make_desc( - jax.ShapeDtypeStruct(tuple(shape), dtype), - tensor_spec=tensor_spec, + return JaxTensorDesc.from_shape( + tuple(shape), + dtype, name=name, + mode=mode, + public_stride_order=public_stride_order, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, init_value=init_value, ) -def as_gemm_tensor_desc(name: str, value: Any, tensor_spec: TensorSpec) -> _SampleDesc: - if isinstance(value, _SampleDesc): - if value.jax_mode != normalize_mode(value.ndim, tensor_spec.mode): +def as_gemm_tensor_desc( + name: str, + value: Any, + *, + mode: tuple[int, ...] | None = None, + public_stride_order: tuple[int, ...] | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, +) -> JaxTensorDesc: + if isinstance(value, JaxTensorDesc): + if value.mode != normalize_mode(value.ndim, mode): raise ValueError( f"{name} descriptor mode does not match the requested layout" ) + expected = JaxTensorDesc.from_shape( + value.array_shape, + value.dtype, + name=name, + mode=mode, + public_stride_order=public_stride_order, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + ) + if (value.stride, value.stride_order) != ( + expected.stride, + expected.stride_order, + ): + raise ValueError( + f"{name} descriptor stride order does not match the requested layout" + ) + if value.divisibility != expected.divisibility: + raise ValueError( + f"{name} descriptor divisibility does not match the requested layout" + ) + if value.ptr_assumed_align != expected.ptr_assumed_align: + raise ValueError( + f"{name} descriptor pointer alignment does not match the requested layout" + ) return value - return _make_desc(value, tensor_spec=tensor_spec, name=name) + return JaxTensorDesc.from_array( + value, + name=name, + mode=mode, + public_stride_order=public_stride_order, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + ) class ApiBaseJax(JaxApiBase): @@ -270,21 +237,45 @@ def _call_impl(self, *args: Any, **kwargs: Any) -> Any: raise NotImplementedError def make_tensor_desc( - self, value: Any, *, tensor_spec: TensorSpec | None = None, name: str = "" - ) -> _SampleDesc: - return _make_desc(value, tensor_spec=tensor_spec, name=name or "value") + self, + value: Any, + *, + mode: tuple[int, ...] | None = None, + public_stride_order: tuple[int, ...] | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, + name: str = "", + ) -> JaxTensorDesc: + return JaxTensorDesc.from_array( + value, + name=name or "value", + mode=mode, + public_stride_order=public_stride_order, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + ) def make_optional_tensor_desc( self, value: Any | None, *, - tensor_spec: TensorSpec | None = None, + mode: tuple[int, ...] | None = None, + public_stride_order: tuple[int, ...] | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, name: str = "", - ) -> _SampleDesc | None: + ) -> JaxTensorDesc | None: return ( None if value is None - else self.make_tensor_desc(value, tensor_spec=tensor_spec, name=name) + else self.make_tensor_desc( + value, + mode=mode, + public_stride_order=public_stride_order, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + name=name, + ) ) @staticmethod @@ -300,7 +291,7 @@ def freeze_mapping(values: Mapping[str, Any]) -> Mapping[str, Any]: def check_tensor_signatures( self, - expected: Mapping[str, _SampleDesc | None], + expected: Mapping[str, JaxTensorDesc | None], values: Mapping[str, Any], ) -> None: for name, expected_desc in expected.items(): @@ -337,49 +328,35 @@ def call_cutedsl( fn: Any, inputs: Sequence[Any], *, + input_descs: Sequence[JaxTensorDesc], outputs: Sequence[JaxTensorDesc], - output_seeds: Sequence[Any | None] | None = None, workspaces: Sequence[JaxTensorDesc] = (), - input_specs: Sequence[TensorSpec | None] | None = None, static_args: Mapping[str, Any] | None = None, allow_cuda_graph: bool = True, use_static_tensors: bool = True, ) -> tuple[Any, ...]: - """Lower a grouped kernel with canonical stream/input/output/workspace order.""" + """Lower a grouped kernel with canonical stream/input/output/workspace order. + + Tensor descriptors retain framework metadata and all CUTLASS lowering + constraints. + """ static = dict(static_args or {}) cluster_shape = static.get("cluster_shape_mn") margin = int(static.pop("cluster_overlap_margin", 0)) - output_specs = tuple(outputs) - workspace_specs = tuple(workspaces) - if output_seeds is None: - supplied_output_seeds = (None,) * len(output_specs) - else: - supplied_output_seeds = tuple(output_seeds) - if len(supplied_output_seeds) != len(output_specs): - raise ValueError( - f"Expected {len(output_specs)} output seeds, " - f"got {len(supplied_output_seeds)}" - ) + input_descs = tuple(input_descs) + output_descs = tuple(outputs) + workspace_descs = tuple(workspaces) + if len(input_descs) != len(inputs): + raise ValueError( + f"Expected {len(inputs)} input descriptors, got {len(input_descs)}" + ) - if any(0 in desc.shape for desc in output_specs): + if any(0 in desc.shape for desc in output_descs): results = [] - for desc, seed in zip(output_specs, supplied_output_seeds): - spec = getattr(desc, "tensor_spec", None) - mode = None if spec is None else spec.mode - if seed is not None: - if desc.init_value is not None: - raise ValueError( - f"{desc.name} cannot have both an explicit seed and init_value" - ) - _CALLER._check_tensor_signature(seed, desc, mode=mode) - results.append(seed) - continue - metadata = _CALLER._materialize_tensor_desc( - desc, - mode=mode, - ) + for desc in output_descs: + metadata = _CALLER._to_shape_dtype_struct(desc) if desc.init_value is None: results.append(jnp.empty(metadata.shape, dtype=metadata.dtype)) else: @@ -409,14 +386,9 @@ def launch(stream: Any, *args: Any) -> None: return _CALLER._call_kernel( tuple(inputs), launch=launch, - output_descs=output_specs, - output_seeds=supplied_output_seeds, - workspace_descs=workspace_specs, - input_spec=None if input_specs is None else tuple(input_specs), - output_spec=tuple(getattr(desc, "tensor_spec", None) for desc in output_specs), - workspace_spec=tuple( - getattr(desc, "tensor_spec", None) for desc in workspace_specs - ), + input_descs=input_descs, + output_descs=output_descs, + workspace_descs=workspace_descs, allow_cuda_graph=allow_cuda_graph, compile_options=compile_options_for_target(compute_capability), use_static_tensors=use_static_tensors, @@ -679,9 +651,12 @@ def ceil_div(a: int, b: int) -> int: __all__ = [ "ApiBaseJax", + "BLOCK_SCALE_MODE", "FIX_PAD_SIZE", "FP8_SF_VEC_SIZE", "HADAMARD_SIZE", + "GROUPED_BIAS_MODE", + "GROUPED_WORKSPACE_ALIGNMENT", "MAX_EXPERTS", "SF_VEC_SIZES", "TWO_CTA_MMA_TILER_M", @@ -689,22 +664,19 @@ def ceil_div(a: int, b: int) -> int: "as_dtype", "as_gemm_tensor_desc", "block_scale_shape", - "block_scale_tensor_spec", "call_cutedsl", "ceil_div", "dense_workspace_bytes", - "gemm_a_tensor_spec", - "gemm_b_tensor_spec", - "gemm_c_tensor_spec", - "grouped_bias_tensor_spec", + "gemm_a_mode", + "gemm_b_mode", + "gemm_output_mode", "grouped_wgrad_workspace_bytes", - "grouped_workspace_tensor_spec", "is_fp4_dtype", "is_fp8_dtype", "is_low_precision_output_dtype", "make_buffer_desc", "normalize_wgrad_input_order", - "probability_tensor_spec", + "PROBABILITY_MODE", "require_16_byte_extent", "require_array", "require_contiguous_alignment", diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py index 4c9037d77..b6c91eddd 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py @@ -1,7 +1,7 @@ # Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dglu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_dglu/jax.py index 6d78601ff..811888516 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dglu/jax.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dglu/jax.py @@ -14,24 +14,24 @@ from .._jax_api import ( ApiBaseJax, + BLOCK_SCALE_MODE, make_buffer_desc, FIX_PAD_SIZE, MAX_EXPERTS, TWO_CTA_MMA_TILER_M, TupleDict, + GROUPED_WORKSPACE_ALIGNMENT, + PROBABILITY_MODE, as_dtype, as_gemm_tensor_desc, block_scale_shape, - block_scale_tensor_spec, call_cutedsl, dense_workspace_bytes, - gemm_a_tensor_spec, - gemm_b_tensor_spec, - gemm_c_tensor_spec, - grouped_workspace_tensor_spec, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, is_fp4_dtype, is_fp8_dtype, - probability_tensor_spec, require_16_byte_extent, require_dtype, require_grouped_cluster_shape, @@ -211,12 +211,12 @@ def _grouped_gemm_dglu_impl( output_layout = require_layout("output_layout", output_layout, ("LMN",)) b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) - a_spec = gemm_a_tensor_spec("LMK") - b_spec = gemm_b_tensor_spec(b_layout) - output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") - a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) - b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) - c_desc = as_gemm_tensor_desc("c_tensor", c_tensor, output_spec) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) + c_desc = as_gemm_tensor_desc("c_tensor", c_tensor, mode=output_mode) m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( a_desc, b_desc, @@ -318,7 +318,6 @@ def _grouped_gemm_dglu_impl( cluster_shape_mn, mma_tiler_mn=mma_tiler_mn ) - scale_spec = block_scale_tensor_spec() require_16_byte_extent("a_tensor", k, ab_dtype) require_16_byte_extent("b_tensor", n if b_layout == "LKN" else k, ab_dtype) require_16_byte_extent("c_tensor", output_n, c_dtype) @@ -336,16 +335,16 @@ def _grouped_gemm_dglu_impl( outputs = [ make_buffer_desc( - "d_row_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + "d_row_tensor", (1, m, output_n), d_dtype, mode=output_mode ), make_buffer_desc( - "d_col_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + "d_col_tensor", (1, m, output_n), d_dtype, mode=output_mode ), make_buffer_desc( "dprob_tensor", (1, 1, m), jnp.float32, - tensor_spec=probability_tensor_spec(), + mode=PROBABILITY_MODE, init_value=0.0, ), ] @@ -374,13 +373,13 @@ def _grouped_gemm_dglu_impl( "sfd_row_tensor", block_scale_shape(m, output_n, 1, sf_vec_size), jnp.float8_e8m0fnu, - tensor_spec=scale_spec, + mode=BLOCK_SCALE_MODE, ), make_buffer_desc( "sfd_col_tensor", block_scale_shape(output_n, m, 1, sf_vec_size), jnp.float8_e8m0fnu, - tensor_spec=scale_spec, + mode=BLOCK_SCALE_MODE, ), ) ) @@ -397,23 +396,24 @@ def _grouped_gemm_dglu_impl( beta_tensor, prob_tensor, ] - input_specs = [ - a_spec, - b_spec, - output_spec, - scale_spec, - scale_spec, - None, - None, - None, - probability_tensor_spec(), + input_descs = [ + a_desc, + b_desc, + c_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + as_gemm_tensor_desc("beta_tensor", beta_tensor), + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE), ] if generate_sfd: inputs.append(norm_const_tensor) - input_specs.append(None) + input_descs.append(as_gemm_tensor_desc("norm_const_tensor", norm_const_tensor)) results = call_cutedsl( _launch, inputs, + input_descs=input_descs, static_args={ "acc_dtype": acc_dtype, "mma_tiler_mn": mma_tiler_mn, @@ -440,10 +440,9 @@ def _grouped_gemm_dglu_impl( "workspace", (workspace_bytes,), jnp.uint8, - tensor_spec=grouped_workspace_tensor_spec(), + ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT, ), ), - input_specs=input_specs, ) result_idx = 0 d_row_tensor = results[result_idx] @@ -507,19 +506,18 @@ def __init__( super().__init__() output_layout = require_layout("output_layout", output_layout, ("LMN",)) b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) - a_spec = gemm_a_tensor_spec("LMK") - b_spec = gemm_b_tensor_spec(b_layout) - output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") - scale_spec = block_scale_tensor_spec() + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") self._sample_descs = { "a_tensor": self.make_tensor_desc( - sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + sample_a_tensor, mode=a_mode, name="sample_a_tensor" ), "c_tensor": self.make_tensor_desc( - sample_c_tensor, tensor_spec=output_spec, name="sample_c_tensor" + sample_c_tensor, mode=output_mode, name="sample_c_tensor" ), "sfa_tensor": self.make_tensor_desc( - sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfa_tensor" ), "padded_offsets": self.make_tensor_desc( sample_padded_offsets, name="sample_padded_offsets" @@ -532,14 +530,14 @@ def __init__( ), "prob_tensor": self.make_tensor_desc( sample_prob_tensor, - tensor_spec=probability_tensor_spec(), + mode=PROBABILITY_MODE, name="sample_prob_tensor", ), "b_tensor": self.make_tensor_desc( - sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + sample_b_tensor, mode=b_mode, name="sample_b_tensor" ), "sfb_tensor": self.make_tensor_desc( - sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfb_tensor" ), "norm_const_tensor": self.make_optional_tensor_desc( sample_norm_const_tensor, name="sample_norm_const_tensor" diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py index ce6431284..af74a57ab 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py @@ -1,7 +1,7 @@ # Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/jax.py index 3a0a689c9..9594646d9 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/jax.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/jax.py @@ -14,24 +14,24 @@ from .._jax_api import ( ApiBaseJax, + BLOCK_SCALE_MODE, make_buffer_desc, FIX_PAD_SIZE, MAX_EXPERTS, TWO_CTA_MMA_TILER_M, TupleDict, + GROUPED_WORKSPACE_ALIGNMENT, + PROBABILITY_MODE, as_dtype, as_gemm_tensor_desc, block_scale_shape, - block_scale_tensor_spec, call_cutedsl, dense_workspace_bytes, - gemm_a_tensor_spec, - gemm_b_tensor_spec, - gemm_c_tensor_spec, - grouped_workspace_tensor_spec, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, is_fp4_dtype, is_fp8_dtype, - probability_tensor_spec, require_16_byte_extent, require_dtype, require_grouped_cluster_shape, @@ -186,12 +186,12 @@ def _grouped_gemm_dsrelu_impl( output_layout = require_layout("output_layout", output_layout, ("LMN",)) b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) - a_spec = gemm_a_tensor_spec("LMK") - b_spec = gemm_b_tensor_spec(b_layout) - output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") - a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) - b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) - c_desc = as_gemm_tensor_desc("c_tensor", c_tensor, output_spec) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) + c_desc = as_gemm_tensor_desc("c_tensor", c_tensor, mode=output_mode) m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( a_desc, b_desc, @@ -277,7 +277,6 @@ def _grouped_gemm_dsrelu_impl( cluster_shape_mn, mma_tiler_mn=mma_tiler_mn ) - scale_spec = block_scale_tensor_spec() require_16_byte_extent("a_tensor", k, ab_dtype) require_16_byte_extent("b_tensor", n if b_layout == "LKN" else k, ab_dtype) require_16_byte_extent("c_tensor", n, c_dtype) @@ -292,14 +291,14 @@ def _grouped_gemm_dsrelu_impl( } outputs = [ - make_buffer_desc("d_row_tensor", (1, m, n), d_dtype, tensor_spec=output_spec), - make_buffer_desc("d_col_tensor", (1, m, n), d_dtype, tensor_spec=output_spec), - make_buffer_desc("d_srelu_tensor", (1, m, n), d_dtype, tensor_spec=output_spec), + make_buffer_desc("d_row_tensor", (1, m, n), d_dtype, mode=output_mode), + make_buffer_desc("d_col_tensor", (1, m, n), d_dtype, mode=output_mode), + make_buffer_desc("d_srelu_tensor", (1, m, n), d_dtype, mode=output_mode), make_buffer_desc( "dprob_tensor", (1, 1, m), jnp.float32, - tensor_spec=probability_tensor_spec(), + mode=PROBABILITY_MODE, init_value=0.0, ), ] @@ -325,19 +324,19 @@ def _grouped_gemm_dsrelu_impl( "sfd_row_tensor", block_scale_shape(m, n, 1, sf_vec_size), jnp.float8_e8m0fnu, - tensor_spec=scale_spec, + mode=BLOCK_SCALE_MODE, ), make_buffer_desc( "sfd_col_tensor", block_scale_shape(n, m, 1, sf_vec_size), jnp.float8_e8m0fnu, - tensor_spec=scale_spec, + mode=BLOCK_SCALE_MODE, ), make_buffer_desc( "sfd_col_d_srelu_tensor", block_scale_shape(n, m, 1, sf_vec_size), jnp.float8_e8m0fnu, - tensor_spec=scale_spec, + mode=BLOCK_SCALE_MODE, ), ) ) @@ -353,22 +352,23 @@ def _grouped_gemm_dsrelu_impl( alpha_tensor, prob_tensor, ] - input_specs = [ - a_spec, - b_spec, - output_spec, - scale_spec, - scale_spec, - None, - None, - probability_tensor_spec(), + input_descs = [ + a_desc, + b_desc, + c_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE), ] if generate_sfd: inputs.append(norm_const_tensor) - input_specs.append(None) + input_descs.append(as_gemm_tensor_desc("norm_const_tensor", norm_const_tensor)) results = call_cutedsl( _launch, inputs, + input_descs=input_descs, static_args={ "acc_dtype": acc_dtype, "mma_tiler_mn": mma_tiler_mn, @@ -390,10 +390,9 @@ def _grouped_gemm_dsrelu_impl( "workspace", (workspace_bytes,), jnp.uint8, - tensor_spec=grouped_workspace_tensor_spec(), + ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT, ), ), - input_specs=input_specs, ) result_idx = 0 d_row_tensor = results[result_idx] @@ -456,25 +455,24 @@ def __init__( super().__init__() output_layout = require_layout("output_layout", output_layout, ("LMN",)) b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) - a_spec = gemm_a_tensor_spec("LMK") - b_spec = gemm_b_tensor_spec(b_layout) - output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") - scale_spec = block_scale_tensor_spec() + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") self._sample_descs = { "a_tensor": self.make_tensor_desc( - sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + sample_a_tensor, mode=a_mode, name="sample_a_tensor" ), "b_tensor": self.make_tensor_desc( - sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + sample_b_tensor, mode=b_mode, name="sample_b_tensor" ), "c_tensor": self.make_tensor_desc( - sample_c_tensor, tensor_spec=output_spec, name="sample_c_tensor" + sample_c_tensor, mode=output_mode, name="sample_c_tensor" ), "sfa_tensor": self.make_tensor_desc( - sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfa_tensor" ), "sfb_tensor": self.make_tensor_desc( - sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfb_tensor" ), "padded_offsets": self.make_tensor_desc( sample_padded_offsets, name="sample_padded_offsets" @@ -484,7 +482,7 @@ def __init__( ), "prob_tensor": self.make_tensor_desc( sample_prob_tensor, - tensor_spec=probability_tensor_spec(), + mode=PROBABILITY_MODE, name="sample_prob_tensor", ), "norm_const_tensor": self.make_optional_tensor_desc( diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py index 91f534d34..541bb1986 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py @@ -1,7 +1,7 @@ # Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/jax.py index 4e07bdf13..92d2cacf3 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/jax.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/jax.py @@ -14,22 +14,22 @@ from .._jax_api import ( ApiBaseJax, + BLOCK_SCALE_MODE, make_buffer_desc, FIX_PAD_SIZE, MAX_EXPERTS, TWO_CTA_MMA_TILER_M, TupleDict, + PROBABILITY_MODE, as_dtype, as_gemm_tensor_desc, block_scale_shape, - block_scale_tensor_spec, call_cutedsl, - gemm_a_tensor_spec, - gemm_b_tensor_spec, - gemm_c_tensor_spec, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, is_fp4_dtype, is_fp8_dtype, - probability_tensor_spec, require_16_byte_extent, require_dtype, require_grouped_cluster_shape, @@ -166,13 +166,13 @@ def _grouped_gemm_dswiglu_impl( """ output_layout = require_layout("output_layout", output_layout, ("LMN",)) - a_spec = gemm_a_tensor_spec("LMK") + a_mode = gemm_a_mode("LMK") b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) - b_spec = gemm_b_tensor_spec(b_layout) - output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") - a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) - b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) - c_desc = as_gemm_tensor_desc("c_tensor", c_tensor, output_spec) + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) + c_desc = as_gemm_tensor_desc("c_tensor", c_tensor, mode=output_mode) m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( a_desc, b_desc, @@ -261,7 +261,6 @@ def _grouped_gemm_dswiglu_impl( cluster_shape_mn, mma_tiler_mn=mma_tiler_mn ) - scale_spec = block_scale_tensor_spec() require_16_byte_extent("a_tensor", k, ab_dtype) require_16_byte_extent("b_tensor", k, ab_dtype) require_16_byte_extent("c_tensor", output_n, c_dtype) @@ -287,33 +286,33 @@ def _grouped_gemm_dswiglu_impl( beta_tensor, prob_tensor, ] - input_specs = [ - a_spec, - b_spec, - output_spec, - scale_spec, - scale_spec, - None, - None, - None, - probability_tensor_spec(), + input_descs = [ + a_desc, + b_desc, + c_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + as_gemm_tensor_desc("beta_tensor", beta_tensor), + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE), ] if generate_sfd: inputs.append(norm_const_tensor) - input_specs.append(None) + input_descs.append(as_gemm_tensor_desc("norm_const_tensor", norm_const_tensor)) outputs = [ make_buffer_desc( - "d_row_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + "d_row_tensor", (1, m, output_n), d_dtype, mode=output_mode ), make_buffer_desc( - "d_col_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + "d_col_tensor", (1, m, output_n), d_dtype, mode=output_mode ), make_buffer_desc( "dprob_tensor", (1, 1, m), jnp.float32, - tensor_spec=probability_tensor_spec(), + mode=PROBABILITY_MODE, init_value=0.0, ), ] @@ -333,13 +332,13 @@ def _grouped_gemm_dswiglu_impl( "sfd_row_tensor", block_scale_shape(m, output_n, 1, sf_vec_size), jnp.float8_e8m0fnu, - tensor_spec=scale_spec, + mode=BLOCK_SCALE_MODE, ), make_buffer_desc( "sfd_col_tensor", block_scale_shape(output_n, m, 1, sf_vec_size), jnp.float8_e8m0fnu, - tensor_spec=scale_spec, + mode=BLOCK_SCALE_MODE, ), ) ) @@ -347,6 +346,7 @@ def _grouped_gemm_dswiglu_impl( results = call_cutedsl( _launch, inputs, + input_descs=input_descs, static_args={ "acc_dtype": acc_dtype, "mma_tiler_mn": mma_tiler_mn, @@ -361,7 +361,6 @@ def _grouped_gemm_dswiglu_impl( "cluster_overlap_margin": int(cluster_overlap_margin), }, outputs=outputs, - input_specs=input_specs, ) result_idx = 3 d_row_tensor, d_col_tensor, dprob_tensor = results[:3] @@ -408,26 +407,25 @@ def __init__( ) -> None: super().__init__() output_layout = require_layout("output_layout", output_layout, ("LMN",)) - a_spec = gemm_a_tensor_spec("LMK") + a_mode = gemm_a_mode("LMK") b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) - b_spec = gemm_b_tensor_spec(b_layout) - output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") - scale_spec = block_scale_tensor_spec() + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") self._sample_descs = { "a_tensor": self.make_tensor_desc( - sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + sample_a_tensor, mode=a_mode, name="sample_a_tensor" ), "b_tensor": self.make_tensor_desc( - sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + sample_b_tensor, mode=b_mode, name="sample_b_tensor" ), "c_tensor": self.make_tensor_desc( - sample_c_tensor, tensor_spec=output_spec, name="sample_c_tensor" + sample_c_tensor, mode=output_mode, name="sample_c_tensor" ), "sfa_tensor": self.make_tensor_desc( - sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfa_tensor" ), "sfb_tensor": self.make_tensor_desc( - sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfb_tensor" ), "padded_offsets": self.make_tensor_desc( sample_padded_offsets, name="sample_padded_offsets" @@ -440,7 +438,7 @@ def __init__( ), "prob_tensor": self.make_tensor_desc( sample_prob_tensor, - tensor_spec=probability_tensor_spec(), + mode=PROBABILITY_MODE, name="sample_prob_tensor", ), "norm_const_tensor": self.make_optional_tensor_desc( diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_glu/__init__.py index 8e38a750c..46a787dc2 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu/__init__.py @@ -1,7 +1,7 @@ # Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_glu/jax.py index 56800496c..45e9c3e46 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu/jax.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu/jax.py @@ -14,24 +14,24 @@ from .._jax_api import ( ApiBaseJax, + BLOCK_SCALE_MODE, make_buffer_desc, FIX_PAD_SIZE, MAX_EXPERTS, TWO_CTA_MMA_TILER_M, TupleDict, + GROUPED_BIAS_MODE, + GROUPED_WORKSPACE_ALIGNMENT, + PROBABILITY_MODE, as_gemm_tensor_desc, block_scale_shape, - block_scale_tensor_spec, call_cutedsl, dense_workspace_bytes, - gemm_a_tensor_spec, - gemm_b_tensor_spec, - gemm_c_tensor_spec, - grouped_bias_tensor_spec, - grouped_workspace_tensor_spec, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, is_fp4_dtype, is_fp8_dtype, - probability_tensor_spec, require_16_byte_extent, require_array, require_dtype, @@ -192,11 +192,11 @@ def _grouped_gemm_glu_impl( output_layout = require_layout("output_layout", output_layout, ("LMN",)) b_layout = require_layout("b_layout", b_layout, ("LNK",)) - a_spec = gemm_a_tensor_spec("LMK") - b_spec = gemm_b_tensor_spec(b_layout) - output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") - a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) - b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( a_desc, b_desc, @@ -298,7 +298,6 @@ def _grouped_gemm_glu_impl( ) output_n = n // 2 - scale_spec = block_scale_tensor_spec() require_16_byte_extent("a_tensor", k, ab_dtype) require_16_byte_extent("b_tensor", k, ab_dtype) require_16_byte_extent("c_tensor", n, c_dtype) @@ -322,26 +321,37 @@ def _grouped_gemm_glu_impl( padded_offsets, alpha_tensor, ] - input_specs = [a_spec, b_spec, scale_spec, scale_spec, None, None] + input_descs = [ + a_desc, + b_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + ] if prob_tensor is not None: inputs.append(prob_tensor) - input_specs.append(probability_tensor_spec()) + input_descs.append( + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE) + ) if bias_tensor is not None: inputs.append(bias_tensor) - input_specs.append(grouped_bias_tensor_spec()) + input_descs.append( + as_gemm_tensor_desc("bias_tensor", bias_tensor, mode=GROUPED_BIAS_MODE) + ) if generate_sfd: inputs.append(norm_const_tensor) - input_specs.append(None) + input_descs.append(as_gemm_tensor_desc("norm_const_tensor", norm_const_tensor)) outputs = [ - make_buffer_desc("c_tensor", (1, m, n), c_dtype, tensor_spec=output_spec), + make_buffer_desc("c_tensor", (1, m, n), c_dtype, mode=output_mode), make_buffer_desc( - "d_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + "d_tensor", (1, m, output_n), d_dtype, mode=output_mode ), ] outputs.append( make_buffer_desc( - "d_col_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + "d_col_tensor", (1, m, output_n), d_dtype, mode=output_mode ) ) if has_amax: @@ -357,13 +367,13 @@ def _grouped_gemm_glu_impl( "sfd_row_tensor", block_scale_shape(m, output_n, 1, sf_vec_size), jnp.float8_e8m0fnu, - tensor_spec=scale_spec, + mode=BLOCK_SCALE_MODE, ), make_buffer_desc( "sfd_col_tensor", block_scale_shape(output_n, m, 1, sf_vec_size), jnp.float8_e8m0fnu, - tensor_spec=scale_spec, + mode=BLOCK_SCALE_MODE, ), ) ) @@ -374,12 +384,13 @@ def _grouped_gemm_glu_impl( "workspace", (workspace_bytes,), jnp.uint8, - tensor_spec=grouped_workspace_tensor_spec(), + ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT, ) ] results = call_cutedsl( _launch, inputs, + input_descs=input_descs, static_args={ "acc_dtype": acc_dtype, "mma_tiler_mn": mma_tiler_mn, @@ -402,7 +413,6 @@ def _grouped_gemm_glu_impl( }, outputs=outputs, workspaces=workspaces, - input_specs=input_specs, ) c_tensor, d_tensor, d_col_tensor = results[:3] result_idx = 3 @@ -456,15 +466,14 @@ def __init__( super().__init__() output_layout = require_layout("output_layout", output_layout, ("LMN",)) b_layout = require_layout("b_layout", b_layout, ("LNK",)) - a_spec = gemm_a_tensor_spec("LMK") - b_spec = gemm_b_tensor_spec(b_layout) - scale_spec = block_scale_tensor_spec() + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) self._sample_descs = { "a_tensor": self.make_tensor_desc( - sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + sample_a_tensor, mode=a_mode, name="sample_a_tensor" ), "sfa_tensor": self.make_tensor_desc( - sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfa_tensor" ), "padded_offsets": self.make_tensor_desc( sample_padded_offsets, name="sample_padded_offsets" @@ -473,14 +482,14 @@ def __init__( sample_alpha_tensor, name="sample_alpha_tensor" ), "b_tensor": self.make_tensor_desc( - sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + sample_b_tensor, mode=b_mode, name="sample_b_tensor" ), "sfb_tensor": self.make_tensor_desc( - sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfb_tensor" ), "bias_tensor": self.make_optional_tensor_desc( sample_bias_tensor, - tensor_spec=grouped_bias_tensor_spec(), + mode=GROUPED_BIAS_MODE, name="sample_bias_tensor", ), "norm_const_tensor": self.make_optional_tensor_desc( @@ -488,7 +497,7 @@ def __init__( ), "prob_tensor": self.make_optional_tensor_desc( sample_prob_tensor, - tensor_spec=probability_tensor_spec(), + mode=PROBABILITY_MODE, name="sample_prob_tensor", ), } diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/__init__.py index 1d81b6710..89974fb93 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/__init__.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/jax.py index 386131692..552e78f6b 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/jax.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/jax.py @@ -14,22 +14,22 @@ from .._jax_api import ( ApiBaseJax, + BLOCK_SCALE_MODE, make_buffer_desc, FIX_PAD_SIZE, HADAMARD_SIZE, MAX_EXPERTS, SF_VEC_SIZES, TupleDict, + GROUPED_BIAS_MODE, + GROUPED_WORKSPACE_ALIGNMENT, + PROBABILITY_MODE, as_gemm_tensor_desc, - block_scale_tensor_spec, call_cutedsl, dense_workspace_bytes, - gemm_a_tensor_spec, - gemm_b_tensor_spec, - gemm_c_tensor_spec, - grouped_bias_tensor_spec, - grouped_workspace_tensor_spec, - probability_tensor_spec, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, require_array, require_contiguous_alignment, require_dtype, @@ -196,11 +196,11 @@ def _grouped_gemm_glu_hadamard_impl( kernel = BlockScaledMoEGroupedGemmGluHadamardKernel output_layout = require_layout("output_layout", output_layout, ("LMN",)) b_layout = require_layout("b_layout", b_layout, ("LNK",)) - a_spec = gemm_a_tensor_spec("LMK") - b_spec = gemm_b_tensor_spec(b_layout) - output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") - a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) - b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( a_desc, b_desc, @@ -316,7 +316,6 @@ def _grouped_gemm_glu_hadamard_impl( "cluster_shape_mn": cluster_shape_mn, } - scale_spec = block_scale_tensor_spec() inputs = [ a_tensor, b_tensor, @@ -326,14 +325,14 @@ def _grouped_gemm_glu_hadamard_impl( alpha_tensor, prob_tensor, ] - input_specs = [ - a_spec, - b_spec, - scale_spec, - scale_spec, - None, - None, - probability_tensor_spec(), + input_descs = [ + a_desc, + b_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE), ] if has_hadamard: @@ -342,16 +341,19 @@ def _grouped_gemm_glu_hadamard_impl( hadamard_values(HADAMARD_SIZE), dtype=jnp.bfloat16 ) inputs.append(hadamard_tensor) - input_specs.append(None) + input_descs.append(as_gemm_tensor_desc("hadamard_tensor", hadamard_tensor)) if bias_tensor is not None: inputs.append(bias_tensor) - input_specs.append(grouped_bias_tensor_spec()) + input_descs.append( + as_gemm_tensor_desc("bias_tensor", bias_tensor, mode=GROUPED_BIAS_MODE) + ) workspace_bytes = max(dense_workspace_bytes(bool(use_dynamic_sched)), 1) c_tensor, d_tensor, amax_tensor, post_rht_amax_tensor = call_cutedsl( _launch, inputs, + input_descs=input_descs, static_args={ "acc_dtype": acc_dtype, "mma_tiler_mn": mma_tiler_mn, @@ -368,13 +370,13 @@ def _grouped_gemm_glu_hadamard_impl( }, outputs=( make_buffer_desc( - "c_tensor", (1, m, n), c_dtype, tensor_spec=output_spec, init_value=0 + "c_tensor", (1, m, n), c_dtype, mode=output_mode, init_value=0 ), make_buffer_desc( "d_tensor", (1, m, output_n), d_dtype, - tensor_spec=output_spec, + mode=output_mode, init_value=0, ), make_buffer_desc("amax_tensor", (experts, 1), jnp.float32, init_value=0.0), @@ -387,11 +389,10 @@ def _grouped_gemm_glu_hadamard_impl( "workspace", (workspace_bytes,), jnp.uint8, - tensor_spec=grouped_workspace_tensor_spec(), + ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT, init_value=0, ), ), - input_specs=input_specs, ) return TupleDict( c_tensor=c_tensor, @@ -433,21 +434,20 @@ def __init__( super().__init__() output_layout = require_layout("output_layout", output_layout, ("LMN",)) b_layout = require_layout("b_layout", b_layout, ("LNK",)) - a_spec = gemm_a_tensor_spec("LMK") - b_spec = gemm_b_tensor_spec(b_layout) - scale_spec = block_scale_tensor_spec() + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) self._sample_descs = { "a_tensor": self.make_tensor_desc( - sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + sample_a_tensor, mode=a_mode, name="sample_a_tensor" ), "b_tensor": self.make_tensor_desc( - sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + sample_b_tensor, mode=b_mode, name="sample_b_tensor" ), "sfa_tensor": self.make_tensor_desc( - sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfa_tensor" ), "sfb_tensor": self.make_tensor_desc( - sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfb_tensor" ), "padded_offsets": self.make_tensor_desc( sample_padded_offsets, name="sample_padded_offsets" @@ -457,12 +457,12 @@ def __init__( ), "prob_tensor": self.make_tensor_desc( sample_prob_tensor, - tensor_spec=probability_tensor_spec(), + mode=PROBABILITY_MODE, name="sample_prob_tensor", ), "bias_tensor": self.make_optional_tensor_desc( sample_bias_tensor, - tensor_spec=grouped_bias_tensor_spec(), + mode=GROUPED_BIAS_MODE, name="sample_bias_tensor", ), "hadamard_tensor": self.make_optional_tensor_desc( diff --git a/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py index a7030ac7e..2806ae24b 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py @@ -1,7 +1,7 @@ # Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/grouped_gemm/grouped_gemm_quant/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_quant/jax.py index 423163851..7c0d62798 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_quant/jax.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_quant/jax.py @@ -14,25 +14,25 @@ from .._jax_api import ( ApiBaseJax, + BLOCK_SCALE_MODE, make_buffer_desc, FIX_PAD_SIZE, MAX_EXPERTS, TWO_CTA_MMA_TILER_M, TupleDict, + GROUPED_BIAS_MODE, + GROUPED_WORKSPACE_ALIGNMENT, + PROBABILITY_MODE, as_gemm_tensor_desc, block_scale_shape, - block_scale_tensor_spec, call_cutedsl, dense_workspace_bytes, - gemm_a_tensor_spec, - gemm_b_tensor_spec, - gemm_c_tensor_spec, - grouped_bias_tensor_spec, - grouped_workspace_tensor_spec, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, is_fp4_dtype, is_fp8_dtype, is_low_precision_output_dtype, - probability_tensor_spec, require_16_byte_extent, require_array, require_dtype, @@ -173,11 +173,11 @@ def _grouped_gemm_quant_impl( output_layout = require_layout("output_layout", output_layout, ("LMN",)) b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) - a_spec = gemm_a_tensor_spec("LMK") - b_spec = gemm_b_tensor_spec(b_layout) - output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") - a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) - b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( a_desc, b_desc, @@ -256,7 +256,6 @@ def _grouped_gemm_quant_impl( cluster_shape_mn, mma_tiler_mn=mma_tiler_mn ) - scale_spec = block_scale_tensor_spec() require_16_byte_extent("a_tensor", k, ab_dtype) require_16_byte_extent("b_tensor", n if b_layout == "LKN" else k, ab_dtype) require_16_byte_extent("d_tensor", n, d_dtype) @@ -278,32 +277,34 @@ def _grouped_gemm_quant_impl( alpha_tensor, prob_tensor, ] - input_specs = [ - a_spec, - b_spec, - scale_spec, - scale_spec, - None, - None, - probability_tensor_spec(), + input_descs = [ + a_desc, + b_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE), ] if row_scale_tensor is not None: inputs.append(row_scale_tensor) - input_specs.append(None) + input_descs.append(as_gemm_tensor_desc("row_scale_tensor", row_scale_tensor)) if bias_tensor is not None: inputs.append(bias_tensor) - input_specs.append(grouped_bias_tensor_spec()) + input_descs.append( + as_gemm_tensor_desc("bias_tensor", bias_tensor, mode=GROUPED_BIAS_MODE) + ) if generate_sfd: inputs.append(norm_const_tensor) - input_specs.append(None) + input_descs.append(as_gemm_tensor_desc("norm_const_tensor", norm_const_tensor)) outputs = [ - make_buffer_desc("d_tensor", (1, m, n), d_dtype, tensor_spec=output_spec) + make_buffer_desc("d_tensor", (1, m, n), d_dtype, mode=output_mode) ] if generate_sfd: outputs.append( make_buffer_desc( - "d_col_tensor", (1, m, n), d_dtype, tensor_spec=output_spec + "d_col_tensor", (1, m, n), d_dtype, mode=output_mode ) ) if has_amax: @@ -319,13 +320,13 @@ def _grouped_gemm_quant_impl( "sfd_row_tensor", block_scale_shape(m, n, 1, sf_vec_size), jnp.float8_e8m0fnu, - tensor_spec=scale_spec, + mode=BLOCK_SCALE_MODE, ), make_buffer_desc( "sfd_col_tensor", block_scale_shape(n, m, 1, sf_vec_size), jnp.float8_e8m0fnu, - tensor_spec=scale_spec, + mode=BLOCK_SCALE_MODE, ), ) ) @@ -334,6 +335,7 @@ def _grouped_gemm_quant_impl( results = call_cutedsl( _launch, inputs, + input_descs=input_descs, static_args={ "acc_dtype": acc_dtype, "mma_tiler_mn": mma_tiler_mn, @@ -355,10 +357,9 @@ def _grouped_gemm_quant_impl( "workspace", (workspace_bytes,), jnp.uint8, - tensor_spec=grouped_workspace_tensor_spec(), + ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT, ), ), - input_specs=input_specs, ) result_idx = 1 d_tensor = results[0] @@ -408,15 +409,14 @@ def __init__( super().__init__() output_layout = require_layout("output_layout", output_layout, ("LMN",)) b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) - a_spec = gemm_a_tensor_spec("LMK") - b_spec = gemm_b_tensor_spec(b_layout) - scale_spec = block_scale_tensor_spec() + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) self._sample_descs = { "a_tensor": self.make_tensor_desc( - sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + sample_a_tensor, mode=a_mode, name="sample_a_tensor" ), "sfa_tensor": self.make_tensor_desc( - sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfa_tensor" ), "padded_offsets": self.make_tensor_desc( sample_padded_offsets, name="sample_padded_offsets" @@ -425,14 +425,14 @@ def __init__( sample_alpha_tensor, name="sample_alpha_tensor" ), "b_tensor": self.make_tensor_desc( - sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + sample_b_tensor, mode=b_mode, name="sample_b_tensor" ), "sfb_tensor": self.make_tensor_desc( - sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfb_tensor" ), "bias_tensor": self.make_optional_tensor_desc( sample_bias_tensor, - tensor_spec=grouped_bias_tensor_spec(), + mode=GROUPED_BIAS_MODE, name="sample_bias_tensor", ), "norm_const_tensor": self.make_optional_tensor_desc( @@ -440,7 +440,7 @@ def __init__( ), "prob_tensor": self.make_optional_tensor_desc( sample_prob_tensor, - tensor_spec=probability_tensor_spec(), + mode=PROBABILITY_MODE, name="sample_prob_tensor", ), "row_scale_tensor": self.make_optional_tensor_desc( diff --git a/python/cudnn/grouped_gemm/grouped_gemm_srelu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_srelu/__init__.py index 65ad86054..b13e70906 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_srelu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_srelu/__init__.py @@ -1,7 +1,7 @@ # Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/grouped_gemm/grouped_gemm_srelu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_srelu/jax.py index aac5dd0c5..854d2e904 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_srelu/jax.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_srelu/jax.py @@ -14,25 +14,25 @@ from .._jax_api import ( ApiBaseJax, + BLOCK_SCALE_MODE, make_buffer_desc, FIX_PAD_SIZE, MAX_EXPERTS, TWO_CTA_MMA_TILER_M, TupleDict, + GROUPED_BIAS_MODE, + GROUPED_WORKSPACE_ALIGNMENT, + PROBABILITY_MODE, as_gemm_tensor_desc, block_scale_shape, - block_scale_tensor_spec, call_cutedsl, dense_workspace_bytes, - gemm_a_tensor_spec, - gemm_b_tensor_spec, - gemm_c_tensor_spec, - grouped_bias_tensor_spec, - grouped_workspace_tensor_spec, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, is_fp4_dtype, is_fp8_dtype, is_low_precision_output_dtype, - probability_tensor_spec, require_16_byte_extent, require_array, require_dtype, @@ -176,11 +176,11 @@ def _grouped_gemm_srelu_impl( output_layout = require_layout("output_layout", output_layout, ("LMN",)) b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) - a_spec = gemm_a_tensor_spec("LMK") - b_spec = gemm_b_tensor_spec(b_layout) - output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") - a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) - b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( a_desc, b_desc, @@ -263,7 +263,6 @@ def _grouped_gemm_srelu_impl( cluster_shape_mn, mma_tiler_mn=mma_tiler_mn ) - scale_spec = block_scale_tensor_spec() require_16_byte_extent("a_tensor", k, ab_dtype) require_16_byte_extent("b_tensor", n if b_layout == "LKN" else k, ab_dtype) require_16_byte_extent("c_tensor", n, c_dtype) @@ -287,30 +286,32 @@ def _grouped_gemm_srelu_impl( alpha_tensor, prob_tensor, ] - input_specs = [ - a_spec, - b_spec, - scale_spec, - scale_spec, - None, - None, - probability_tensor_spec(), + input_descs = [ + a_desc, + b_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE), ] if bias_tensor is not None: inputs.append(bias_tensor) - input_specs.append(grouped_bias_tensor_spec()) + input_descs.append( + as_gemm_tensor_desc("bias_tensor", bias_tensor, mode=GROUPED_BIAS_MODE) + ) if generate_sfd: inputs.append(norm_const_tensor) - input_specs.append(None) + input_descs.append(as_gemm_tensor_desc("norm_const_tensor", norm_const_tensor)) outputs = [ - make_buffer_desc("c_tensor", (1, m, n), c_dtype, tensor_spec=output_spec), - make_buffer_desc("d_tensor", (1, m, n), d_dtype, tensor_spec=output_spec), + make_buffer_desc("c_tensor", (1, m, n), c_dtype, mode=output_mode), + make_buffer_desc("d_tensor", (1, m, n), d_dtype, mode=output_mode), ] if generate_sfd: outputs.append( make_buffer_desc( - "d_col_tensor", (1, m, n), d_dtype, tensor_spec=output_spec + "d_col_tensor", (1, m, n), d_dtype, mode=output_mode ) ) if has_amax: @@ -326,13 +327,13 @@ def _grouped_gemm_srelu_impl( "sfd_row_tensor", block_scale_shape(m, n, 1, sf_vec_size), jnp.float8_e8m0fnu, - tensor_spec=scale_spec, + mode=BLOCK_SCALE_MODE, ), make_buffer_desc( "sfd_col_tensor", block_scale_shape(n, m, 1, sf_vec_size), jnp.float8_e8m0fnu, - tensor_spec=scale_spec, + mode=BLOCK_SCALE_MODE, ), ) ) @@ -341,6 +342,7 @@ def _grouped_gemm_srelu_impl( results = call_cutedsl( _launch, inputs, + input_descs=input_descs, static_args={ "acc_dtype": acc_dtype, "mma_tiler_mn": mma_tiler_mn, @@ -361,10 +363,9 @@ def _grouped_gemm_srelu_impl( "workspace", (workspace_bytes,), jnp.uint8, - tensor_spec=grouped_workspace_tensor_spec(), + ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT, ), ), - input_specs=input_specs, ) result_idx = 0 c_tensor, d_tensor = results[:2] @@ -416,21 +417,20 @@ def __init__( super().__init__() output_layout = require_layout("output_layout", output_layout, ("LMN",)) b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) - a_spec = gemm_a_tensor_spec("LMK") - b_spec = gemm_b_tensor_spec(b_layout) - scale_spec = block_scale_tensor_spec() + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) self._sample_descs = { "a_tensor": self.make_tensor_desc( - sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + sample_a_tensor, mode=a_mode, name="sample_a_tensor" ), "b_tensor": self.make_tensor_desc( - sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + sample_b_tensor, mode=b_mode, name="sample_b_tensor" ), "sfa_tensor": self.make_tensor_desc( - sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfa_tensor" ), "sfb_tensor": self.make_tensor_desc( - sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfb_tensor" ), "padded_offsets": self.make_tensor_desc( sample_padded_offsets, name="sample_padded_offsets" @@ -440,7 +440,7 @@ def __init__( ), "bias_tensor": self.make_optional_tensor_desc( sample_bias_tensor, - tensor_spec=grouped_bias_tensor_spec(), + mode=GROUPED_BIAS_MODE, name="sample_bias_tensor", ), "norm_const_tensor": self.make_optional_tensor_desc( @@ -448,7 +448,7 @@ def __init__( ), "prob_tensor": self.make_optional_tensor_desc( sample_prob_tensor, - tensor_spec=probability_tensor_spec(), + mode=PROBABILITY_MODE, name="sample_prob_tensor", ), } diff --git a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py index d2829021b..c62d1240f 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py @@ -1,7 +1,7 @@ # Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/jax.py index ad147731c..901489cbc 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/jax.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/jax.py @@ -14,21 +14,21 @@ from .._jax_api import ( ApiBaseJax, + BLOCK_SCALE_MODE, make_buffer_desc, FIX_PAD_SIZE, MAX_EXPERTS, TWO_CTA_MMA_TILER_M, TupleDict, + PROBABILITY_MODE, as_gemm_tensor_desc, block_scale_shape, - block_scale_tensor_spec, call_cutedsl, - gemm_a_tensor_spec, - gemm_b_tensor_spec, - gemm_c_tensor_spec, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, is_fp4_dtype, is_fp8_dtype, - probability_tensor_spec, require_16_byte_extent, require_dtype, require_grouped_cluster_shape, @@ -154,11 +154,11 @@ def _grouped_gemm_swiglu_impl( """ output_layout = require_layout("output_layout", output_layout, ("LMN",)) - a_spec = gemm_a_tensor_spec("LMK") - b_spec = gemm_b_tensor_spec("LNK") - output_spec = gemm_c_tensor_spec(output_layout, name="output_layout") - a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, a_spec) - b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, b_spec) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode("LNK") + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( a_desc, b_desc, @@ -246,7 +246,6 @@ def _grouped_gemm_swiglu_impl( raise NotImplementedError("FP8 output requires mma_tiler_mn[1] == 256") output_n = n // 2 - scale_spec = block_scale_tensor_spec() require_16_byte_extent("a_tensor", k, ab_dtype) require_16_byte_extent("b_tensor", k, ab_dtype) require_16_byte_extent("c_tensor", n, c_dtype) @@ -269,50 +268,60 @@ def _grouped_gemm_swiglu_impl( padded_offsets, alpha_tensor, ] - input_specs = [a_spec, b_spec, scale_spec, scale_spec, None, None] + input_descs = [ + a_desc, + b_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + ] if generate_sfd: inputs.append(norm_const_tensor) - input_specs.append(None) + input_descs.append(as_gemm_tensor_desc("norm_const_tensor", norm_const_tensor)) if prob_tensor is not None: inputs.append(prob_tensor) - input_specs.append(probability_tensor_spec()) + input_descs.append( + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE) + ) has_amax = d_dtype in {jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)} - output_specs = [ - make_buffer_desc("c_tensor", (1, m, n), c_dtype, tensor_spec=output_spec), + outputs = [ + make_buffer_desc("c_tensor", (1, m, n), c_dtype, mode=output_mode), make_buffer_desc( - "d_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + "d_tensor", (1, m, output_n), d_dtype, mode=output_mode ), make_buffer_desc( - "d_col_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec + "d_col_tensor", (1, m, output_n), d_dtype, mode=output_mode ), ] if has_amax: - output_specs.append( + outputs.append( make_buffer_desc( "amax_tensor", (experts, 1), jnp.float32, init_value=-float("inf") ) ) if generate_sfd: - output_specs.extend( + outputs.extend( ( make_buffer_desc( "sfd_row_tensor", block_scale_shape(m, output_n, 1, sf_vec_size), jnp.float8_e8m0fnu, - tensor_spec=scale_spec, + mode=BLOCK_SCALE_MODE, ), make_buffer_desc( "sfd_col_tensor", block_scale_shape(output_n, m, 1, sf_vec_size), jnp.float8_e8m0fnu, - tensor_spec=scale_spec, + mode=BLOCK_SCALE_MODE, ), ) ) results = call_cutedsl( _launch, inputs, + input_descs=input_descs, static_args={ "acc_dtype": acc_dtype, "mma_tiler_mn": mma_tiler_mn, @@ -326,8 +335,7 @@ def _grouped_gemm_swiglu_impl( "has_amax": has_amax, "cluster_overlap_margin": int(cluster_overlap_margin), }, - outputs=output_specs, - input_specs=input_specs, + outputs=outputs, ) result_idx = 0 c_tensor = results[result_idx] @@ -376,21 +384,20 @@ def __init__( ) -> None: super().__init__() output_layout = require_layout("output_layout", output_layout, ("LMN",)) - a_spec = gemm_a_tensor_spec("LMK") - b_spec = gemm_b_tensor_spec("LNK") - scale_spec = block_scale_tensor_spec() + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode("LNK") self._sample_descs = { "a_tensor": self.make_tensor_desc( - sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + sample_a_tensor, mode=a_mode, name="sample_a_tensor" ), "b_tensor": self.make_tensor_desc( - sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + sample_b_tensor, mode=b_mode, name="sample_b_tensor" ), "sfa_tensor": self.make_tensor_desc( - sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfa_tensor" ), "sfb_tensor": self.make_tensor_desc( - sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfb_tensor" ), "padded_offsets": self.make_tensor_desc( sample_padded_offsets, name="sample_padded_offsets" @@ -403,7 +410,7 @@ def __init__( ), "prob_tensor": self.make_optional_tensor_desc( sample_prob_tensor, - tensor_spec=probability_tensor_spec(), + mode=PROBABILITY_MODE, name="sample_prob_tensor", ), } diff --git a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/__init__.py index b823f92bc..01f13bc67 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/__init__.py @@ -1,7 +1,7 @@ # Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/jax.py index dc2d6d5a1..e6b8b9a6c 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/jax.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/jax.py @@ -11,18 +11,18 @@ import jax import jax.numpy as jnp -from cutlass.jax import TensorSpec from .._jax_api import ( ApiBaseJax, + GROUPED_WORKSPACE_ALIGNMENT, make_buffer_desc, TWO_CTA_MMA_TILER_M, TupleDict, as_dtype, + as_gemm_tensor_desc, call_cutedsl, ceil_div, grouped_wgrad_workspace_bytes, - grouped_workspace_tensor_spec, normalize_wgrad_input_order, require_16_byte_extent, require_array, @@ -45,28 +45,10 @@ def wgrad_scale_shape( ) -def wgrad_a_tensor_spec() -> TensorSpec: - """Describe ``A=(hidden,tokens)`` with tokens contiguous.""" - - return TensorSpec(layout=(1, 0), mode=(0, 1), ptr_assumed_align=16) - - -def wgrad_b_tensor_spec() -> TensorSpec: - """Describe ``B=(tokens,intermediate)`` with tokens contiguous.""" - - return TensorSpec(layout=(0, 1), mode=(0, 1), ptr_assumed_align=16) - - -def wgrad_scale_tensor_spec() -> TensorSpec: - """Describe the packed two-dimensional scale buffers.""" - - return TensorSpec(layout=(1, 0), mode=(0, 1), ptr_assumed_align=16) - - -def wgrad_output_tensor_spec() -> TensorSpec: - """Describe contiguous ``(experts,hidden,intermediate)`` output.""" - - return TensorSpec(layout=(2, 1, 0), mode=(0, 1, 2), ptr_assumed_align=16) +WGRAD_ALIGNMENT = 16 +WGRAD_A_STRIDE_ORDER = (1, 0) +WGRAD_B_STRIDE_ORDER = (0, 1) +WGRAD_OUTPUT_STRIDE_ORDER = (2, 1, 0) def _launch( @@ -147,7 +129,6 @@ def _grouped_gemm_wgrad_impl( offsets_tensor: Any, global_scale_a: Optional[Any] = None, global_scale_b: Optional[Any] = None, - wgrad_tensor: Optional[Any] = None, acc_dtype: Any = None, wgrad_dtype: Any = None, mma_tiler_mn: tuple[int, int] = (256, 256), @@ -169,10 +150,9 @@ def _grouped_gemm_wgrad_impl( aligned to ``sf_vec_size``, and no greater than ``tokens``. The JAX surface exposes only the dense output mode; raw per-expert pointer - outputs are not accepted. When ``accumulate_on_output=True``, an optional - ``wgrad_tensor`` seeds the reducing epilogue and the returned value contains - the seed plus the newly computed gradients. Omitting the seed preserves the - simple fresh-output behavior by starting from zero. + outputs are not accepted. The output is inferred from the inputs and + allocated by JAX. When ``accumulate_on_output=True``, the reducing epilogue + starts from a zero-initialized output. ``input_order='tensor_ragged'`` is supported when A and B have already been packed expert-by-expert according to the runtime offsets. JAX cannot infer @@ -239,9 +219,6 @@ def _grouped_gemm_wgrad_impl( ) output_shape = (expert_cnt, hidden, intermediate) - if wgrad_tensor is not None and not accumulate_on_output: - raise ValueError("wgrad_tensor is only valid when accumulate_on_output=True") - has_global_scale = global_scale_a is not None or global_scale_b is not None if has_global_scale: if global_scale_a is None or global_scale_b is None: @@ -268,32 +245,12 @@ def _grouped_gemm_wgrad_impl( default=jnp.float32, ) valid_wgrad_dtypes = (jnp.bfloat16, jnp.float16, jnp.float32) - if wgrad_tensor is not None: - require_array( - wgrad_tensor, - name="wgrad_tensor", - shape=output_shape, - dtype=valid_wgrad_dtypes, - ) - seed_dtype = as_dtype(wgrad_tensor) - wgrad_dtype = require_dtype( - wgrad_dtype, - valid_wgrad_dtypes, - name="wgrad_dtype", - default=seed_dtype, - ) - if wgrad_dtype != seed_dtype: - raise ValueError( - "wgrad_dtype must match wgrad_tensor.dtype, got " - f"{wgrad_dtype} and {seed_dtype}" - ) - else: - wgrad_dtype = require_dtype( - wgrad_dtype, - valid_wgrad_dtypes, - name="wgrad_dtype", - default=jnp.bfloat16, - ) + wgrad_dtype = require_dtype( + wgrad_dtype, + valid_wgrad_dtypes, + name="wgrad_dtype", + default=jnp.bfloat16, + ) mma_tiler_mn = require_grouped_mma_tiler(mma_tiler_mn, allowed_m=(128, 256)) if cluster_shape_mn is None: cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) @@ -316,20 +273,46 @@ def _grouped_gemm_wgrad_impl( } inputs = [a_tensor, b_tensor, sfa_tensor, sfb_tensor, offsets_tensor] - input_specs = [ - wgrad_a_tensor_spec(), - wgrad_b_tensor_spec(), - wgrad_scale_tensor_spec(), - wgrad_scale_tensor_spec(), - None, + input_descs = [ + as_gemm_tensor_desc( + "a_tensor", + a_tensor, + public_stride_order=WGRAD_A_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + ), + as_gemm_tensor_desc( + "b_tensor", + b_tensor, + public_stride_order=WGRAD_B_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + ), + as_gemm_tensor_desc( + "sfa_tensor", + sfa_tensor, + public_stride_order=WGRAD_A_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + ), + as_gemm_tensor_desc( + "sfb_tensor", + sfb_tensor, + public_stride_order=WGRAD_A_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + ), + as_gemm_tensor_desc("offsets_tensor", offsets_tensor), ] if has_global_scale: inputs.extend((global_scale_a, global_scale_b)) - input_specs.extend((None, None)) + input_descs.extend( + ( + as_gemm_tensor_desc("global_scale_a", global_scale_a), + as_gemm_tensor_desc("global_scale_b", global_scale_b), + ) + ) (result_wgrad_tensor,) = call_cutedsl( _launch, inputs, + input_descs=input_descs, static_args={ "acc_dtype": acc_dtype, "mma_tiler_mn": mma_tiler_mn, @@ -346,13 +329,11 @@ def _grouped_gemm_wgrad_impl( "wgrad_tensor", output_shape, wgrad_dtype, - tensor_spec=wgrad_output_tensor_spec(), - init_value=( - 0.0 if accumulate_on_output and wgrad_tensor is None else None - ), + public_stride_order=WGRAD_OUTPUT_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + init_value=0.0 if accumulate_on_output else None, ), ), - output_seeds=(wgrad_tensor,), workspaces=( make_buffer_desc( "workspace", @@ -363,10 +344,9 @@ def _grouped_gemm_wgrad_impl( ), ), jnp.uint8, - tensor_spec=grouped_workspace_tensor_spec(), + ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT, ), ), - input_specs=input_specs, ) return TupleDict(wgrad_tensor=result_wgrad_tensor) @@ -390,25 +370,32 @@ def __init__( sf_vec_size: int = 32, accumulate_on_output: bool = False, input_order: str = "tensor2d", - *, - sample_wgrad_tensor: Optional[Any] = None, ) -> None: super().__init__() - a_spec = wgrad_a_tensor_spec() - b_spec = wgrad_b_tensor_spec() - scale_spec = wgrad_scale_tensor_spec() self._sample_descs = { "a_tensor": self.make_tensor_desc( - sample_a_tensor, tensor_spec=a_spec, name="sample_a_tensor" + sample_a_tensor, + public_stride_order=WGRAD_A_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + name="sample_a_tensor", ), "b_tensor": self.make_tensor_desc( - sample_b_tensor, tensor_spec=b_spec, name="sample_b_tensor" + sample_b_tensor, + public_stride_order=WGRAD_B_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + name="sample_b_tensor", ), "sfa_tensor": self.make_tensor_desc( - sample_sfa_tensor, tensor_spec=scale_spec, name="sample_sfa_tensor" + sample_sfa_tensor, + public_stride_order=WGRAD_A_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + name="sample_sfa_tensor", ), "sfb_tensor": self.make_tensor_desc( - sample_sfb_tensor, tensor_spec=scale_spec, name="sample_sfb_tensor" + sample_sfb_tensor, + public_stride_order=WGRAD_A_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + name="sample_sfb_tensor", ), "offsets_tensor": self.make_tensor_desc( sample_offsets_tensor, name="sample_offsets_tensor" @@ -419,11 +406,6 @@ def __init__( "global_scale_b": self.make_optional_tensor_desc( sample_global_scale_b, name="sample_global_scale_b" ), - "wgrad_tensor": self.make_optional_tensor_desc( - sample_wgrad_tensor, - tensor_spec=wgrad_output_tensor_spec(), - name="sample_wgrad_tensor", - ), } self._config = { "acc_dtype": self.as_optional_dtype(acc_dtype), @@ -452,7 +434,6 @@ def _check_support(self) -> None: self._sample_descs["offsets_tensor"], self._sample_descs["global_scale_a"], self._sample_descs["global_scale_b"], - self._sample_descs["wgrad_tensor"], **self._config, _validate_only=True, ) @@ -467,7 +448,6 @@ def __call__( offsets_tensor: Any, global_scale_a: Optional[Any] = None, global_scale_b: Optional[Any] = None, - wgrad_tensor: Optional[Any] = None, ) -> TupleDict: return super().__call__( a_tensor, @@ -477,7 +457,6 @@ def __call__( offsets_tensor, global_scale_a, global_scale_b, - wgrad_tensor, ) def _call_impl( @@ -489,7 +468,6 @@ def _call_impl( offsets_tensor: Any, global_scale_a: Optional[Any] = None, global_scale_b: Optional[Any] = None, - wgrad_tensor: Optional[Any] = None, ) -> TupleDict: values = { "a_tensor": a_tensor, @@ -499,7 +477,6 @@ def _call_impl( "offsets_tensor": offsets_tensor, "global_scale_a": global_scale_a, "global_scale_b": global_scale_b, - "wgrad_tensor": wgrad_tensor, } self.check_tensor_signatures(self._sample_descs, values) return _grouped_gemm_wgrad_impl(**values, **self._config) @@ -532,14 +509,12 @@ def grouped_gemm_wgrad_wrapper_sm100( sf_vec_size: int = 32, accumulate_on_output: bool = False, input_order: str = "tensor2d", - *, - wgrad_tensor: Optional[Any] = None, ) -> TupleDict: """Compute dense expert weight gradients using the SM100 grouped kernel. - Set ``accumulate_on_output=True`` and pass ``wgrad_tensor`` to accumulate - into an existing JAX value. If accumulation is enabled without a seed, the - returned gradient starts from zero. Pointer-table outputs are not supported. + The output is inferred and allocated by JAX. Set + ``accumulate_on_output=True`` to use the reducing epilogue with a + zero-initialized output. Pointer-table outputs are not supported. """ op = GroupedGemmWgradSm100( @@ -557,7 +532,6 @@ def grouped_gemm_wgrad_wrapper_sm100( sf_vec_size=sf_vec_size, accumulate_on_output=accumulate_on_output, input_order=input_order, - sample_wgrad_tensor=wgrad_tensor, ) return op( a_tensor, @@ -567,7 +541,6 @@ def grouped_gemm_wgrad_wrapper_sm100( offsets_tensor, global_scale_a, global_scale_b, - wgrad_tensor, ) diff --git a/python/cudnn/native_sparse_attention/__init__.py b/python/cudnn/native_sparse_attention/__init__.py index fb3e4867d..fcb30fdf3 100644 --- a/python/cudnn/native_sparse_attention/__init__.py +++ b/python/cudnn/native_sparse_attention/__init__.py @@ -2,7 +2,7 @@ from typing import Any -from .._operation_api import make_operation_api +from ..common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/native_sparse_attention/compression/__init__.py b/python/cudnn/native_sparse_attention/compression/__init__.py index 5e76bf348..75c7e5c34 100644 --- a/python/cudnn/native_sparse_attention/compression/__init__.py +++ b/python/cudnn/native_sparse_attention/compression/__init__.py @@ -1,6 +1,6 @@ """Lazy Torch compression-attention API exports.""" -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/native_sparse_attention/compression/jax.py b/python/cudnn/native_sparse_attention/compression/jax.py index 20cdb529d..5215e515d 100644 --- a/python/cudnn/native_sparse_attention/compression/jax.py +++ b/python/cudnn/native_sparse_attention/compression/jax.py @@ -13,12 +13,11 @@ import jax.numpy as jnp from ... import data_type -from ..._cute_compiler import compile_options_for_target -from ..._jax import JaxApiBase, TupleDict +from ..._jax.compiler import compile_options_for_target +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict from ..jax_utils import ( BHS_TO_BSH_MODE, FIXED_LAYOUTS, - describe_bhs_as_bsh, describe_fixed_data, fixed_data_mode, make_fixed_output, @@ -182,19 +181,22 @@ def __init__( "max_s_q and max_s_k must be positive and no larger than " "their packed token counts" ) - self.q_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *sample_q.shape), sample_q.dtype), - "q_tensor", + self.q_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_q.shape), + sample_q.dtype, + name="q_tensor", public_stride_order=(3, 2, 0, 1), ) - self.k_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *sample_k.shape), sample_k.dtype), - "k_tensor", + self.k_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_k.shape), + sample_k.dtype, + name="k_tensor", public_stride_order=(3, 2, 0, 1), ) - self.v_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *sample_v.shape), sample_v.dtype), - "v_tensor", + self.v_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_v.shape), + sample_v.dtype, + name="v_tensor", public_stride_order=(3, 2, 0, 1), ) if self.seqlen_q < self.seqlen_k or self.seqlen_q % self.seqlen_k: @@ -240,12 +242,11 @@ def __init__( ) self.o_kernel_desc = self.o_desc self.lse_desc = ( - describe_bhs_as_bsh( - jax.ShapeDtypeStruct( - (self.batch, self.num_query_heads, self.seqlen_q), - jnp.float32, - ), - "lse_tensor", + JaxTensorDesc.from_shape( + (self.batch, self.num_query_heads, self.seqlen_q), + jnp.float32, + name="lse_tensor", + mode=BHS_TO_BSH_MODE, ) if self.enable_lse else None @@ -257,26 +258,29 @@ def __init__( self.num_query_heads, self.head_dim, ) - self.o_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct(output_shape, self.output_dtype), "o_tensor" + self.o_desc = JaxTensorDesc.from_shape( + output_shape, self.output_dtype, name="o_tensor" ) - self.o_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *output_shape), self.output_dtype), - "o_kernel_tensor", + self.o_kernel_desc = JaxTensorDesc.from_shape( + (1, *output_shape), + self.output_dtype, + name="o_kernel_tensor", public_stride_order=(3, 2, 0, 1), ) self.lse_desc = ( - self._to_tensor_desc( - jax.ShapeDtypeStruct(output_shape[:2], jnp.float32), - "lse_tensor", + JaxTensorDesc.from_shape( + output_shape[:2], + jnp.float32, + name="lse_tensor", ) if self.enable_lse else None ) self.lse_kernel_desc = ( - self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *output_shape[:2]), jnp.float32), - "lse_kernel_tensor", + JaxTensorDesc.from_shape( + (1, *output_shape[:2]), + jnp.float32, + name="lse_kernel_tensor", ) if self.enable_lse else None @@ -301,40 +305,29 @@ def __call__( cum_seqlen_k_tensor: Any | None = None, ) -> TupleDict: self.check_support() - signature_mode = self.data_mode - for value, desc in ( - (q_tensor, self.q_desc), - (k_tensor, self.k_desc), - (v_tensor, self.v_desc), - ): - self._check_tensor_signature(value, desc, mode=signature_mode) - if self.input_layout in FIXED_LAYOUTS: if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: raise ValueError( "cumulative sequence lengths must be omitted for fixed inputs" ) inputs = (q_tensor, k_tensor, v_tensor) - input_specs = ( - self._to_tensor_spec(self.q_kernel_desc, mode=self.data_mode), - self._to_tensor_spec(self.k_kernel_desc, mode=self.data_mode), - self._to_tensor_spec(self.v_kernel_desc, mode=self.data_mode), + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.v_kernel_desc, ) - output_specs = ( - self._to_tensor_spec(self.o_kernel_desc, mode=self.data_mode), - ) - if self.lse_kernel_desc is not None: - output_specs += ( - self._to_tensor_spec(self.lse_kernel_desc, mode=BHS_TO_BSH_MODE), - ) launch = self._launch_kernel else: if cum_seqlen_q_tensor is None or cum_seqlen_k_tensor is None: raise ValueError( "packed THD inputs require cumulative Q and K sequence lengths" ) - self._check_tensor_signature(cum_seqlen_q_tensor, self.cum_q_desc) - self._check_tensor_signature(cum_seqlen_k_tensor, self.cum_k_desc) + for value, desc in ( + (q_tensor, self.q_desc), + (k_tensor, self.k_desc), + (v_tensor, self.v_desc), + ): + self._check_tensor_signature(value, desc) q_tensor = jnp.reshape(q_tensor, self.q_kernel_desc.shape) k_tensor = jnp.reshape(k_tensor, self.k_kernel_desc.shape) v_tensor = jnp.reshape(v_tensor, self.v_kernel_desc.shape) @@ -345,16 +338,13 @@ def __call__( cum_seqlen_q_tensor, cum_seqlen_k_tensor, ) - input_specs = ( - self._to_tensor_spec(self.q_kernel_desc), - self._to_tensor_spec(self.k_kernel_desc), - self._to_tensor_spec(self.v_kernel_desc), - self._to_tensor_spec(self.cum_q_desc), - self._to_tensor_spec(self.cum_k_desc), + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.v_kernel_desc, + self.cum_q_desc, + self.cum_k_desc, ) - output_specs = (self._to_tensor_spec(self.o_kernel_desc),) - if self.lse_kernel_desc is not None: - output_specs += (self._to_tensor_spec(self.lse_kernel_desc),) launch = self._launch_varlen_kernel results = self._call_kernel( @@ -363,8 +353,7 @@ def __call__( output_descs=(self.o_kernel_desc,) if self.lse_kernel_desc is None else (self.o_kernel_desc, self.lse_kernel_desc), - input_spec=input_specs, - output_spec=output_specs, + input_descs=input_descs, compile_options=compile_options_for_target(self.compute_capability), ) output = results[0] @@ -510,24 +499,12 @@ def compression_attention_wrapper( Fixed outputs follow ``layout``; LSE remains ``(B, H, S)``. """ - qkv_samples = tuple( - jax.ShapeDtypeStruct(value.shape, value.dtype) - for value in (q_tensor, k_tensor, v_tensor) - ) - cum_q_sample = ( - None - if cum_seqlen_q_tensor is None - else jax.ShapeDtypeStruct(cum_seqlen_q_tensor.shape, cum_seqlen_q_tensor.dtype) - ) - cum_k_sample = ( - None - if cum_seqlen_k_tensor is None - else jax.ShapeDtypeStruct(cum_seqlen_k_tensor.shape, cum_seqlen_k_tensor.dtype) - ) return CompressionAttention( - *qkv_samples, - cum_q_sample, - cum_k_sample, + q_tensor, + k_tensor, + v_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, enable_lse=enable_lse, o_dtype=o_dtype, qk_acc_dtype=qk_acc_dtype, diff --git a/python/cudnn/native_sparse_attention/jax_utils.py b/python/cudnn/native_sparse_attention/jax_utils.py index 9d9a25df0..e4dd4c367 100644 --- a/python/cudnn/native_sparse_attention/jax_utils.py +++ b/python/cudnn/native_sparse_attention/jax_utils.py @@ -7,7 +7,6 @@ from typing import Any -import jax import jax.numpy as jnp from .. import data_type @@ -108,12 +107,18 @@ def make_fixed_output( kernel_stride_order: tuple[int, ...] | None = None, init_value: bool | int | float | None = None, ) -> JaxTensorDesc: - return describe_fixed_data( - jax.ShapeDtypeStruct(public_shape, dtype), - name, - layout=layout, - kernel_axes=kernel_axes, - kernel_stride_order=kernel_stride_order, + mode = fixed_data_mode(layout, kernel_axes=kernel_axes) + public_stride_order = ( + None + if kernel_stride_order is None + else stride_order_to_public(kernel_stride_order, mode) + ) + return JaxTensorDesc.from_shape( + public_shape, + dtype, + name=name, + mode=mode, + public_stride_order=public_stride_order, init_value=init_value, ) diff --git a/python/cudnn/native_sparse_attention/selection/__init__.py b/python/cudnn/native_sparse_attention/selection/__init__.py index 231338d4e..e61451af3 100644 --- a/python/cudnn/native_sparse_attention/selection/__init__.py +++ b/python/cudnn/native_sparse_attention/selection/__init__.py @@ -1,6 +1,6 @@ """Lazy Torch selection-attention API exports.""" -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/native_sparse_attention/selection/jax.py b/python/cudnn/native_sparse_attention/selection/jax.py index d7082a0a9..381a50961 100644 --- a/python/cudnn/native_sparse_attention/selection/jax.py +++ b/python/cudnn/native_sparse_attention/selection/jax.py @@ -13,8 +13,8 @@ import jax.numpy as jnp from ... import data_type -from ..._cute_compiler import compile_options_for_target -from ..._jax import JaxApiBase, TupleDict +from ..._jax.compiler import compile_options_for_target +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict from ..jax_utils import normalize_supported_dtype SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) @@ -148,38 +148,37 @@ def __init__( self.target_compute_capability = target_compute_capability self.compute_capability: int | None = None - self.q_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *self.q_desc.shape), self.q_desc.dtype), - "q_tensor", + self.q_kernel_desc = JaxTensorDesc.from_shape( + (1, *self.q_desc.shape), + self.q_desc.dtype, + name="q_tensor", ) - self.k_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *self.k_desc.shape), self.k_desc.dtype), - "k_tensor", + self.k_kernel_desc = JaxTensorDesc.from_shape( + (1, *self.k_desc.shape), + self.k_desc.dtype, + name="k_tensor", ) - self.v_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *self.v_desc.shape), self.v_desc.dtype), - "v_tensor", + self.v_kernel_desc = JaxTensorDesc.from_shape( + (1, *self.v_desc.shape), + self.v_desc.dtype, + name="v_tensor", ) - self.o_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct( - (1, self.total_tokens, self.num_query_heads, self.value_dim), - self.output_dtype, - ), - "o_tensor", + self.o_desc = JaxTensorDesc.from_shape( + (1, self.total_tokens, self.num_query_heads, self.value_dim), + self.output_dtype, + name="o_tensor", init_value=0, ) - self.l_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct( - (1, self.total_tokens, self.num_query_heads), jnp.float32 - ), - "l_tensor", + self.l_desc = JaxTensorDesc.from_shape( + (1, self.total_tokens, self.num_query_heads), + jnp.float32, + name="l_tensor", init_value=0.0, ) - self.m_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct( - (1, self.total_tokens, self.num_query_heads), jnp.float32 - ), - "m_tensor", + self.m_desc = JaxTensorDesc.from_shape( + (1, self.total_tokens, self.num_query_heads), + jnp.float32, + name="m_tensor", init_value=float("-inf"), ) @@ -205,9 +204,6 @@ def __call__( (q_tensor, self.q_desc), (k_tensor, self.k_desc), (v_tensor, self.v_desc), - (block_indices_tensor, self.block_indices_desc), - (block_counts_tensor, self.block_counts_desc), - (cum_seqlen_tensor, self.cum_seqlen_desc), ): self._check_tensor_signature(value, desc) @@ -225,13 +221,13 @@ def __call__( ), launch=self._launch_kernel, output_descs=(self.o_desc, self.l_desc, self.m_desc), - input_spec=( - self._to_tensor_spec(self.q_kernel_desc), - self._to_tensor_spec(self.k_kernel_desc), - self._to_tensor_spec(self.v_kernel_desc), - self._to_tensor_spec(self.block_indices_desc), - self._to_tensor_spec(self.block_counts_desc), - self._to_tensor_spec(self.cum_seqlen_desc), + input_descs=( + self.q_kernel_desc, + self.k_kernel_desc, + self.v_kernel_desc, + self.block_indices_desc, + self.block_counts_desc, + self.cum_seqlen_desc, ), compile_options=compile_options_for_target(self.compute_capability), ) @@ -323,9 +319,8 @@ def selection_attention_wrapper( block_counts_tensor, cum_seqlen_tensor, ) - samples = tuple(jax.ShapeDtypeStruct(value.shape, value.dtype) for value in values) return SelectionAttention( - *samples, + *values, max_s_q=max_s_q, max_s_k=max_s_k, block_size=block_size, diff --git a/python/cudnn/native_sparse_attention/sliding_window_attention/__init__.py b/python/cudnn/native_sparse_attention/sliding_window_attention/__init__.py index 503009633..980cc3fcb 100644 --- a/python/cudnn/native_sparse_attention/sliding_window_attention/__init__.py +++ b/python/cudnn/native_sparse_attention/sliding_window_attention/__init__.py @@ -1,6 +1,6 @@ """Lazy Torch sliding-window-attention API exports.""" -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/native_sparse_attention/sliding_window_attention/jax.py b/python/cudnn/native_sparse_attention/sliding_window_attention/jax.py index 5125b7f8d..8b24bf21f 100644 --- a/python/cudnn/native_sparse_attention/sliding_window_attention/jax.py +++ b/python/cudnn/native_sparse_attention/sliding_window_attention/jax.py @@ -215,7 +215,7 @@ def __call__( (k_tensor, self.k_desc), (v_tensor, self.v_desc), ): - self._check_tensor_signature(value, desc, mode=self.data_mode) + self._check_tensor_signature(value, desc) if self.input_layout == "BHSD": if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: @@ -331,24 +331,12 @@ def sliding_window_attention_wrapper( intentionally outside this functional JAX API. """ - samples = tuple( - jax.ShapeDtypeStruct(value.shape, value.dtype) - for value in (q_tensor, k_tensor, v_tensor) - ) - sample_cum_seqlen_q = ( - None - if cum_seqlen_q_tensor is None - else jax.ShapeDtypeStruct(cum_seqlen_q_tensor.shape, cum_seqlen_q_tensor.dtype) - ) - sample_cum_seqlen_k = ( - None - if cum_seqlen_k_tensor is None - else jax.ShapeDtypeStruct(cum_seqlen_k_tensor.shape, cum_seqlen_k_tensor.dtype) - ) return SlidingWindowAttention( - *samples, - sample_cum_seqlen_q, - sample_cum_seqlen_k, + q_tensor, + k_tensor, + v_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, max_s_q=max_s_q, max_s_k=max_s_k, left_bound=left_bound, diff --git a/python/cudnn/native_sparse_attention/top_k/__init__.py b/python/cudnn/native_sparse_attention/top_k/__init__.py index d60e6bd78..705ddb1a5 100644 --- a/python/cudnn/native_sparse_attention/top_k/__init__.py +++ b/python/cudnn/native_sparse_attention/top_k/__init__.py @@ -1,6 +1,6 @@ """Lazy Torch top-K API exports.""" -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/native_sparse_attention/top_k/jax.py b/python/cudnn/native_sparse_attention/top_k/jax.py index b0e1062c6..73efbec05 100644 --- a/python/cudnn/native_sparse_attention/top_k/jax.py +++ b/python/cudnn/native_sparse_attention/top_k/jax.py @@ -13,8 +13,8 @@ import jax.numpy as jnp from ... import data_type -from ..._cute_compiler import compile_options_for_target -from ..._jax import JaxApiBase, TupleDict +from ..._jax.compiler import compile_options_for_target +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict from ..._jax.layout import to_public_axes from ..jax_utils import ( FIXED_LAYOUTS, @@ -135,14 +135,16 @@ def __init__( ) else: output_shape = (self.total_q_tokens, self.num_kv_heads, self.k_value) - self.scores_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct(output_shape, jnp.float32), - "topk_scores_tensor", + self.scores_desc = JaxTensorDesc.from_shape( + output_shape, + jnp.float32, + name="topk_scores_tensor", init_value=float("-inf"), ) - self.indices_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct(output_shape, jnp.int32), - "topk_indices_tensor", + self.indices_desc = JaxTensorDesc.from_shape( + output_shape, + jnp.int32, + name="topk_indices_tensor", init_value=-1, ) @@ -315,27 +317,22 @@ def _init_packed_shape( # These virtual rank-4 views preserve the packed THD byte order while # exposing the B,H,T,D shapes read by the CuTe kernel. LSE is likewise # presented as B,H,T, matching the Torch wrapper's transpose. - self.q_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct( - (1, self.num_query_heads, self.total_q_tokens, self.head_dim), - self.q_desc.dtype, - ), - "q_tensor", + self.q_kernel_desc = JaxTensorDesc.from_shape( + (1, self.num_query_heads, self.total_q_tokens, self.head_dim), + self.q_desc.dtype, + name="q_tensor", public_stride_order=_KERNEL_DATA_STRIDE_ORDER, ) - self.k_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct( - (1, self.num_kv_heads, self.total_k_tokens, self.head_dim), - self.k_desc.dtype, - ), - "k_tensor", + self.k_kernel_desc = JaxTensorDesc.from_shape( + (1, self.num_kv_heads, self.total_k_tokens, self.head_dim), + self.k_desc.dtype, + name="k_tensor", public_stride_order=_KERNEL_DATA_STRIDE_ORDER, ) - self.lse_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct( - (1, self.num_query_heads, self.total_q_tokens), jnp.float32 - ), - "lse_tensor", + self.lse_kernel_desc = JaxTensorDesc.from_shape( + (1, self.num_query_heads, self.total_q_tokens), + jnp.float32, + name="lse_tensor", ) def _check_configuration(self) -> None: @@ -383,20 +380,16 @@ def __call__( cum_seqlen_k_tensor: Any | None = None, ) -> TupleDict: self.check_support() - self._check_tensor_signature(q_tensor, self.q_desc, mode=self.data_mode) - self._check_tensor_signature(k_tensor, self.k_desc, mode=self.data_mode) - self._check_tensor_signature(lse_tensor, self.lse_desc) - if self.input_layout in FIXED_LAYOUTS: if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: raise ValueError( "cum_seqlen_q and cum_seqlen_k must be omitted for fixed layout" ) inputs = (q_tensor, k_tensor, lse_tensor) - input_spec = ( - self._to_tensor_spec(self.q_kernel_desc, mode=self.data_mode), - self._to_tensor_spec(self.k_kernel_desc, mode=self.data_mode), - self._to_tensor_spec(self.lse_kernel_desc), + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.lse_kernel_desc, ) launch = self._launch_kernel else: @@ -404,8 +397,12 @@ def __call__( raise ValueError( "cum_seqlen_q and cum_seqlen_k are both required for THD layout" ) - self._check_tensor_signature(cum_seqlen_q_tensor, self.cum_seqlen_q_desc) - self._check_tensor_signature(cum_seqlen_k_tensor, self.cum_seqlen_k_desc) + for value, desc in ( + (q_tensor, self.q_desc), + (k_tensor, self.k_desc), + (lse_tensor, self.lse_desc), + ): + self._check_tensor_signature(value, desc) q_storage = jnp.transpose( jnp.reshape( q_tensor, @@ -431,25 +428,20 @@ def __call__( cum_seqlen_q_tensor, cum_seqlen_k_tensor, ) - input_spec = ( - self._to_tensor_spec(self.q_kernel_desc), - self._to_tensor_spec(self.k_kernel_desc), - self._to_tensor_spec(self.lse_kernel_desc), - self._to_tensor_spec(self.cum_seqlen_q_desc), - self._to_tensor_spec(self.cum_seqlen_k_desc), + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.lse_kernel_desc, + self.cum_seqlen_q_desc, + self.cum_seqlen_k_desc, ) launch = self._launch_packed_kernel - output_mode = self.data_mode if self.input_layout in FIXED_LAYOUTS else None scores, indices = self._call_kernel( inputs, launch=launch, output_descs=(self.scores_desc, self.indices_desc), - input_spec=input_spec, - output_spec=( - self._to_tensor_spec(self.scores_desc, mode=output_mode), - self._to_tensor_spec(self.indices_desc, mode=output_mode), - ), + input_descs=input_descs, compile_options=compile_options_for_target(self.compute_capability), ) return TupleDict(topk_scores_tensor=scores, topk_indices_tensor=indices) @@ -585,24 +577,12 @@ def topk_reduction_wrapper( Fixed outputs follow ``layout``; fixed LSE remains ``(B, H, S)``. """ - samples = tuple( - jax.ShapeDtypeStruct(value.shape, value.dtype) - for value in (q_tensor, k_tensor, lse_tensor) - ) - sample_cum_seqlen_q = ( - None - if cum_seqlen_q_tensor is None - else jax.ShapeDtypeStruct(cum_seqlen_q_tensor.shape, cum_seqlen_q_tensor.dtype) - ) - sample_cum_seqlen_k = ( - None - if cum_seqlen_k_tensor is None - else jax.ShapeDtypeStruct(cum_seqlen_k_tensor.shape, cum_seqlen_k_tensor.dtype) - ) return TopKReduction( - *samples, - sample_cum_seqlen_q, - sample_cum_seqlen_k, + q_tensor, + k_tensor, + lse_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, max_s_q=max_s_q, max_s_k=max_s_k, acc_dtype=acc_dtype, diff --git a/python/cudnn/rmsnorm_rht_amax/__init__.py b/python/cudnn/rmsnorm_rht_amax/__init__.py index 71ee075c9..de543a280 100644 --- a/python/cudnn/rmsnorm_rht_amax/__init__.py +++ b/python/cudnn/rmsnorm_rht_amax/__init__.py @@ -6,7 +6,7 @@ JAX APIs are exported through :mod:`cudnn.jax`. """ -from .._operation_api import make_operation_api +from ..common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/rmsnorm_rht_amax/jax.py b/python/cudnn/rmsnorm_rht_amax/jax.py index 046b973ef..0d94335cc 100644 --- a/python/cudnn/rmsnorm_rht_amax/jax.py +++ b/python/cudnn/rmsnorm_rht_amax/jax.py @@ -81,17 +81,13 @@ def check_support(self) -> bool: def __call__(self, x: Any, weight: Any) -> tuple[Any, Any]: self.check_support() - self._check_tensor_signature(x, self.x_desc) - self._check_tensor_signature(weight, self.w_desc) - x_spec = self._to_tensor_spec( - self.x_desc, - divisibility=(self._op.rows_per_cta, 16), + x_desc = self.x_desc.with_divisibility( + (self._op.rows_per_cta, 16) ) - weight_spec = self._to_tensor_spec(self.w_desc, divisibility=(16,)) - output_spec = self._to_tensor_spec( - self.o_desc, - divisibility=(self._op.rows_per_cta, 16), + weight_desc = self.w_desc.with_divisibility((16,)) + output_desc = self.o_desc.with_divisibility( + (self._op.rows_per_cta, 16) ) def launch( @@ -112,9 +108,8 @@ def launch( return self._call_kernel( (x, weight), launch=launch, - output_descs=(self.o_desc, self.amax_desc), - input_spec=(x_spec, weight_spec), - output_spec=(output_spec, None), + output_descs=(output_desc, self.amax_desc), + input_descs=(x_desc, weight_desc), ) @@ -129,11 +124,9 @@ def rmsnorm_rht_amax_sm100( ) -> tuple[Any, Any]: """Apply fused RMSNorm, 16-wide RHT, and per-CTA amax from JAX.""" - sample_x = jax.ShapeDtypeStruct(x.shape, x.dtype) - sample_weight = jax.ShapeDtypeStruct(weight.shape, weight.dtype) return RmsNormRhtAmaxSm100( - sample_x, - sample_weight, + x, + weight, eps=eps, num_threads=num_threads, rows_per_cta=rows_per_cta, diff --git a/python/cudnn/rmsnorm_rht_amax/kernel.py b/python/cudnn/rmsnorm_rht_amax/kernel.py index b0ca5c3bb..01715d98b 100644 --- a/python/cudnn/rmsnorm_rht_amax/kernel.py +++ b/python/cudnn/rmsnorm_rht_amax/kernel.py @@ -70,6 +70,35 @@ def redux_sync_max_f32(val, *, loc=None, ip=None): return Float32(result) +@dsl_user_op +def _as_global_tensor(tensor: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor: + """Preserve a generic kernel argument's global address space for cp.async.""" + + if tensor.memspace == cute.AddressSpace.gmem: + return tensor + if tensor.memspace != cute.AddressSpace.generic: + raise ValueError(f"Expected a global or generic tensor, got {tensor.memspace}") + + pointer = tensor.iterator + global_pointer = llvm.addrspacecast( + llvm.PointerType.get(cute.AddressSpace.gmem), + pointer.llvm_ptr, + loc=loc, + ip=ip, + ) + return cute.make_tensor( + cute.make_ptr( + tensor.element_type, + global_pointer, + cute.AddressSpace.gmem, + assumed_align=pointer.alignment, + loc=loc, + ip=ip, + ), + tensor.layout, + ) + + class RMSNormRHTAmaxKernel: """CuTe implementation of fused RMSNorm + RHT + per-CTA amax.""" @@ -119,6 +148,7 @@ def __init__( @cute.kernel def kernel(self, m_x: cute.Tensor, m_w: cute.Tensor, m_o: cute.Tensor, m_amax: cute.Tensor, eps: Float32, tv_layout: cute.Layout, tiler_mn: cute.Shape): cfg = self + m_x = _as_global_tensor(m_x) tid = cute.arch.thread_idx()[0] bid = cute.arch.block_idx()[0] inv_sqrt_had = cutlass.Float32(cfg.inv_sqrt_had) diff --git a/python/cudnn/rmsnorm_rht_amax/op.py b/python/cudnn/rmsnorm_rht_amax/op.py index 2c342cd75..b3b967a8b 100644 --- a/python/cudnn/rmsnorm_rht_amax/op.py +++ b/python/cudnn/rmsnorm_rht_amax/op.py @@ -8,8 +8,8 @@ from typing import Any, Optional from .. import data_type -from .._op import Op -from .._tensor_desc import TensorDesc +from ..common.op import Op +from ..common.tensor_desc import TensorDesc HAD_BLOCK = 16 diff --git a/python/cudnn/sdpa/__init__.py b/python/cudnn/sdpa/__init__.py index f743f26dc..a77765b5f 100644 --- a/python/cudnn/sdpa/__init__.py +++ b/python/cudnn/sdpa/__init__.py @@ -3,7 +3,7 @@ """Lazy Torch API exports for the CuTe SDPA operations.""" -from .._operation_api import make_operation_api +from ..common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/sdpa/bwd/__init__.py b/python/cudnn/sdpa/bwd/__init__.py index 9568c8ca7..55775acd8 100644 --- a/python/cudnn/sdpa/bwd/__init__.py +++ b/python/cudnn/sdpa/bwd/__init__.py @@ -3,7 +3,7 @@ """Lazy Torch SDPA-backward API exports.""" -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/sdpa/bwd/jax.py b/python/cudnn/sdpa/bwd/jax.py index bad2dc364..3af81d368 100644 --- a/python/cudnn/sdpa/bwd/jax.py +++ b/python/cudnn/sdpa/bwd/jax.py @@ -12,7 +12,7 @@ import jax.numpy as jnp from ... import data_type -from ..._cute_compiler import compile_options_for_target +from ..._jax.compiler import compile_options_for_target from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict from ..jax_utils import ( FIXED_LAYOUTS, @@ -147,29 +147,23 @@ def __init__( self.workspace_seqlen_q = self.seqlen_q self.workspace_batch = self.batch else: - self.dq_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct(tuple(sample_q.shape), sample_q.dtype), - "dq_tensor", + self.dq_desc = JaxTensorDesc.from_shape( + tuple(sample_q.shape), sample_q.dtype, name="dq_tensor" ) - self.dk_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct(tuple(sample_k.shape), sample_k.dtype), - "dk_tensor", + self.dk_desc = JaxTensorDesc.from_shape( + tuple(sample_k.shape), sample_k.dtype, name="dk_tensor" ) - self.dv_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct(tuple(sample_v.shape), sample_v.dtype), - "dv_tensor", + self.dv_desc = JaxTensorDesc.from_shape( + tuple(sample_v.shape), sample_v.dtype, name="dv_tensor" ) - self.dq_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *sample_q.shape), sample_q.dtype), - "dq_kernel_tensor", + self.dq_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_q.shape), sample_q.dtype, name="dq_kernel_tensor" ) - self.dk_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *sample_k.shape), sample_k.dtype), - "dk_kernel_tensor", + self.dk_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_k.shape), sample_k.dtype, name="dk_kernel_tensor" ) - self.dv_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *sample_v.shape), sample_v.dtype), - "dv_kernel_tensor", + self.dv_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_v.shape), sample_v.dtype, name="dv_kernel_tensor" ) # Workspace size follows the promoted (1, T, H, D) ABI, matching # the Torch adapter's use of s_qo=T and b=1. @@ -353,27 +347,26 @@ def _init_packed_shape( "max_s_q and max_s_k must be positive and no larger than their packed token counts" ) - self.q_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *sample_q.shape), sample_q.dtype), "q_tensor" + self.q_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_q.shape), sample_q.dtype, name="q_tensor" ) - self.k_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *sample_k.shape), sample_k.dtype), "k_tensor" + self.k_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_k.shape), sample_k.dtype, name="k_tensor" ) - self.v_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *sample_v.shape), sample_v.dtype), "v_tensor" + self.v_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_v.shape), sample_v.dtype, name="v_tensor" ) - self.o_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *sample_o.shape), sample_o.dtype), "o_tensor" + self.o_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_o.shape), sample_o.dtype, name="o_tensor" ) - self.do_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *sample_do.shape), sample_do.dtype), "do_tensor" + self.do_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_do.shape), sample_do.dtype, name="do_tensor" ) # The kernel addresses packed LSE as (1, H, T) with T contiguous. - self.lse_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct( - (1, self.num_query_heads, self.total_q_tokens), jnp.float32 - ), - "lse_tensor", + self.lse_kernel_desc = JaxTensorDesc.from_shape( + (1, self.num_query_heads, self.total_q_tokens), + jnp.float32, + name="lse_tensor", ) def check_support(self) -> bool: @@ -396,36 +389,19 @@ def __call__( cum_seqlen_k_tensor: Any | None = None, ) -> TupleDict: self.check_support() - signature_mode = self.data_mode - for value, desc in ( - (q_tensor, self.q_desc), - (k_tensor, self.k_desc), - (v_tensor, self.v_desc), - (o_tensor, self.o_desc), - (do_tensor, self.do_desc), - (lse_tensor, self.lse_desc), - ): - mode = signature_mode if desc is not self.lse_desc else None - self._check_tensor_signature(value, desc, mode=mode) - if self.input_layout in FIXED_LAYOUTS: if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: raise ValueError( "cum_seqlen_q and cum_seqlen_k must be omitted for fixed layout" ) inputs = (q_tensor, k_tensor, v_tensor, o_tensor, do_tensor, lse_tensor) - input_spec = ( - self._to_tensor_spec(self.q_kernel_desc, mode=self.data_mode), - self._to_tensor_spec(self.k_kernel_desc, mode=self.data_mode), - self._to_tensor_spec(self.v_kernel_desc, mode=self.data_mode), - self._to_tensor_spec(self.o_kernel_desc, mode=self.data_mode), - self._to_tensor_spec(self.do_kernel_desc, mode=self.data_mode), - self._to_tensor_spec(self.lse_kernel_desc), - ) - output_spec = ( - self._to_tensor_spec(self.dq_kernel_desc, mode=self.data_mode), - self._to_tensor_spec(self.dk_kernel_desc, mode=self.data_mode), - self._to_tensor_spec(self.dv_kernel_desc, mode=self.data_mode), + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.v_kernel_desc, + self.o_kernel_desc, + self.do_kernel_desc, + self.lse_kernel_desc, ) launch = self._launch_kernel else: @@ -433,8 +409,15 @@ def __call__( raise ValueError( "cum_seqlen_q and cum_seqlen_k are both required for THD layout" ) - self._check_tensor_signature(cum_seqlen_q_tensor, self.cum_q_desc) - self._check_tensor_signature(cum_seqlen_k_tensor, self.cum_k_desc) + for value, desc in ( + (q_tensor, self.q_desc), + (k_tensor, self.k_desc), + (v_tensor, self.v_desc), + (o_tensor, self.o_desc), + (do_tensor, self.do_desc), + (lse_tensor, self.lse_desc), + ): + self._check_tensor_signature(value, desc) lse_storage = jnp.transpose( jnp.reshape( lse_tensor, @@ -452,20 +435,15 @@ def __call__( cum_seqlen_q_tensor, cum_seqlen_k_tensor, ) - input_spec = ( - self._to_tensor_spec(self.q_kernel_desc), - self._to_tensor_spec(self.k_kernel_desc), - self._to_tensor_spec(self.v_kernel_desc), - self._to_tensor_spec(self.o_kernel_desc), - self._to_tensor_spec(self.do_kernel_desc), - self._to_tensor_spec(self.lse_kernel_desc), - self._to_tensor_spec(self.cum_q_desc), - self._to_tensor_spec(self.cum_k_desc), - ) - output_spec = ( - self._to_tensor_spec(self.dq_kernel_desc), - self._to_tensor_spec(self.dk_kernel_desc), - self._to_tensor_spec(self.dv_kernel_desc), + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.v_kernel_desc, + self.o_kernel_desc, + self.do_kernel_desc, + self.lse_kernel_desc, + self.cum_q_desc, + self.cum_k_desc, ) launch = self._launch_varlen_kernel @@ -477,9 +455,8 @@ def __call__( self.dk_kernel_desc, self.dv_kernel_desc, ), + input_descs=input_descs, workspace_descs=(self._workspace_desc(),), - input_spec=input_spec, - output_spec=output_spec, compile_options=compile_options_for_target(self.compute_capability), ) if self.input_layout == _PACKED_LAYOUT: @@ -683,22 +660,15 @@ def sdpa_bwd_wrapper_sm100_d256( remains ``(B, H, S)``. """ - values = (q_tensor, k_tensor, v_tensor, o_tensor, do_tensor, lse_tensor) - samples = tuple(jax.ShapeDtypeStruct(value.shape, value.dtype) for value in values) - sample_cum_seqlen_q = ( - None - if cum_seqlen_q_tensor is None - else jax.ShapeDtypeStruct(cum_seqlen_q_tensor.shape, cum_seqlen_q_tensor.dtype) - ) - sample_cum_seqlen_k = ( - None - if cum_seqlen_k_tensor is None - else jax.ShapeDtypeStruct(cum_seqlen_k_tensor.shape, cum_seqlen_k_tensor.dtype) - ) return SdpabwdSm100D256( - *samples, - sample_cum_seqlen_q, - sample_cum_seqlen_k, + q_tensor, + k_tensor, + v_tensor, + o_tensor, + do_tensor, + lse_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, max_s_q=max_s_q, max_s_k=max_s_k, acc_dtype=acc_dtype, diff --git a/python/cudnn/sdpa/fwd/__init__.py b/python/cudnn/sdpa/fwd/__init__.py index 44f91d48a..c2cfef51f 100644 --- a/python/cudnn/sdpa/fwd/__init__.py +++ b/python/cudnn/sdpa/fwd/__init__.py @@ -3,7 +3,7 @@ """Lazy Torch SDPA-forward API exports.""" -from ..._operation_api import make_operation_api +from ...common.operation_api import make_operation_api __all__, __getattr__, __dir__ = make_operation_api( globals(), diff --git a/python/cudnn/sdpa/fwd/jax.py b/python/cudnn/sdpa/fwd/jax.py index 254497d1f..cedfaab03 100644 --- a/python/cudnn/sdpa/fwd/jax.py +++ b/python/cudnn/sdpa/fwd/jax.py @@ -13,8 +13,8 @@ import jax.numpy as jnp from ... import data_type -from ..._cute_compiler import compile_options_for_target -from ..._jax import JaxApiBase, TupleDict +from ..._jax.compiler import compile_options_for_target +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict from ..jax_utils import ( FIXED_LAYOUTS, describe_fixed_data, @@ -119,30 +119,29 @@ def __init__( layout=self.input_layout, ) self.o_kernel_desc = self.o_desc - self.lse_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct( - (self.batch, self.num_query_heads, self.seqlen_q), - jnp.float32, - ), - "lse_tensor", + self.lse_desc = JaxTensorDesc.from_shape( + (self.batch, self.num_query_heads, self.seqlen_q), + jnp.float32, + name="lse_tensor", ) self.lse_kernel_desc = self.lse_desc else: - self.o_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct(tuple(sample_q.shape), sample_q.dtype), - "o_tensor", + self.o_desc = JaxTensorDesc.from_shape( + tuple(sample_q.shape), + sample_q.dtype, + name="o_tensor", ) - self.o_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *sample_q.shape), sample_q.dtype), - "o_kernel_tensor", + self.o_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_q.shape), + sample_q.dtype, + name="o_kernel_tensor", ) - self.lse_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct( - (self.total_q_tokens, self.num_query_heads), jnp.float32 - ), - "lse_tensor", - # The kernel writes sequence-contiguous LSE. TensorSpec asks - # XLA for that physical order while preserving public (T, H). + self.lse_desc = JaxTensorDesc.from_shape( + (self.total_q_tokens, self.num_query_heads), + jnp.float32, + name="lse_tensor", + # The kernel writes sequence-contiguous LSE. The descriptor + # records that physical order while preserving public (T, H). public_stride_order=(0, 1), ) # Forward consumes packed LSE through its iterator, just as the @@ -281,14 +280,14 @@ def _init_packed_shape( # Public THD remains row-major; adding the singleton batch dimension # yields the exact (1, T, H, D) ABI used by the Torch wrapper. - self.q_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *sample_q.shape), sample_q.dtype), "q_tensor" + self.q_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_q.shape), sample_q.dtype, name="q_tensor" ) - self.k_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *sample_k.shape), sample_k.dtype), "k_tensor" + self.k_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_k.shape), sample_k.dtype, name="k_tensor" ) - self.v_kernel_desc = self._to_tensor_desc( - jax.ShapeDtypeStruct((1, *sample_v.shape), sample_v.dtype), "v_tensor" + self.v_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_v.shape), sample_v.dtype, name="v_tensor" ) def check_support(self) -> bool: @@ -308,28 +307,16 @@ def __call__( cum_seqlen_k_tensor: Any | None = None, ) -> TupleDict: self.check_support() - signature_mode = self.data_mode - for value, desc in ( - (q_tensor, self.q_desc), - (k_tensor, self.k_desc), - (v_tensor, self.v_desc), - ): - self._check_tensor_signature(value, desc, mode=signature_mode) - if self.input_layout in FIXED_LAYOUTS: if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: raise ValueError( "cum_seqlen_q and cum_seqlen_k must be omitted for fixed layout" ) inputs = (q_tensor, k_tensor, v_tensor) - input_spec = ( - self._to_tensor_spec(self.q_kernel_desc, mode=self.data_mode), - self._to_tensor_spec(self.k_kernel_desc, mode=self.data_mode), - self._to_tensor_spec(self.v_kernel_desc, mode=self.data_mode), - ) - output_spec = ( - self._to_tensor_spec(self.o_kernel_desc, mode=self.data_mode), - self._to_tensor_spec(self.lse_kernel_desc), + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.v_kernel_desc, ) launch = self._launch_kernel else: @@ -337,8 +324,12 @@ def __call__( raise ValueError( "cum_seqlen_q and cum_seqlen_k are both required for THD layout" ) - self._check_tensor_signature(cum_seqlen_q_tensor, self.cum_q_desc) - self._check_tensor_signature(cum_seqlen_k_tensor, self.cum_k_desc) + for value, desc in ( + (q_tensor, self.q_desc), + (k_tensor, self.k_desc), + (v_tensor, self.v_desc), + ): + self._check_tensor_signature(value, desc) inputs = ( jnp.reshape(q_tensor, self.q_kernel_desc.shape), jnp.reshape(k_tensor, self.k_kernel_desc.shape), @@ -346,16 +337,12 @@ def __call__( cum_seqlen_q_tensor, cum_seqlen_k_tensor, ) - input_spec = ( - self._to_tensor_spec(self.q_kernel_desc), - self._to_tensor_spec(self.k_kernel_desc), - self._to_tensor_spec(self.v_kernel_desc), - self._to_tensor_spec(self.cum_q_desc), - self._to_tensor_spec(self.cum_k_desc), - ) - output_spec = ( - self._to_tensor_spec(self.o_kernel_desc), - self._to_tensor_spec(self.lse_kernel_desc), + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.v_kernel_desc, + self.cum_q_desc, + self.cum_k_desc, ) launch = self._launch_varlen_kernel @@ -363,8 +350,7 @@ def __call__( inputs, launch=launch, output_descs=(self.o_kernel_desc, self.lse_kernel_desc), - input_spec=input_spec, - output_spec=output_spec, + input_descs=input_descs, compile_options=compile_options_for_target(self.compute_capability), ) if self.input_layout == _PACKED_LAYOUT: @@ -503,24 +489,12 @@ def sdpa_fwd_wrapper_sm100_d256( ``(B, H, S)``. """ - samples = tuple( - jax.ShapeDtypeStruct(value.shape, value.dtype) - for value in (q_tensor, k_tensor, v_tensor) - ) - sample_cum_seqlen_q = ( - None - if cum_seqlen_q_tensor is None - else jax.ShapeDtypeStruct(cum_seqlen_q_tensor.shape, cum_seqlen_q_tensor.dtype) - ) - sample_cum_seqlen_k = ( - None - if cum_seqlen_k_tensor is None - else jax.ShapeDtypeStruct(cum_seqlen_k_tensor.shape, cum_seqlen_k_tensor.dtype) - ) return SdpafwdSm100D256( - *samples, - sample_cum_seqlen_q, - sample_cum_seqlen_k, + q_tensor, + k_tensor, + v_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, max_s_q=max_s_q, max_s_k=max_s_k, qk_acc_dtype=qk_acc_dtype, diff --git a/python/cudnn/sdpa/jax_utils.py b/python/cudnn/sdpa/jax_utils.py index c54fa7fd5..64547de05 100644 --- a/python/cudnn/sdpa/jax_utils.py +++ b/python/cudnn/sdpa/jax_utils.py @@ -8,7 +8,6 @@ import math from typing import Any -import jax import jax.numpy as jnp from .. import data_type @@ -79,10 +78,13 @@ def make_fixed_output( layout: str, init_value: bool | int | float | None = None, ) -> JaxTensorDesc: - return describe_fixed_data( - jax.ShapeDtypeStruct(public_shape, dtype), - name, - layout=layout, + mode = fixed_data_mode(layout) + return JaxTensorDesc.from_shape( + public_shape, + dtype, + name=name, + mode=mode, + public_stride_order=stride_order_to_public(KERNEL_STRIDE_ORDER, mode), init_value=init_value, ) diff --git a/test/python/fe_api/test_dense_gemm_helpers.py b/test/python/fe_api/test_dense_gemm_helpers.py index 2db698500..81d66bd9c 100644 --- a/test/python/fe_api/test_dense_gemm_helpers.py +++ b/test/python/fe_api/test_dense_gemm_helpers.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for framework-neutral dense GEMM helpers.""" +"""Contracts for framework-neutral dense GEMM helpers.""" import ast from enum import Enum, auto @@ -45,8 +45,8 @@ def setUpClass(cls) -> None: package.data_type = _DataType sys.modules[_PACKAGE] = package - cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") - cls.helpers = importlib.import_module(f"{_PACKAGE}._dense_gemm") + cls.tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + cls.helpers = importlib.import_module(f"{_PACKAGE}.gemm.helpers") @classmethod def tearDownClass(cls) -> None: @@ -180,7 +180,7 @@ def test_cluster_shape_validation_handles_cta_groups(self): self.helpers.require_cluster_shape(shape, **kwargs) def test_module_has_no_framework_or_kernel_dependencies(self): - tree = ast.parse((_CUDNN_ROOT / "_dense_gemm.py").read_text()) + tree = ast.parse((_CUDNN_ROOT / "gemm" / "helpers.py").read_text()) imports = [] for node in tree.body: if isinstance(node, ast.Import): diff --git a/test/python/fe_api/test_dsa_compiler.py b/test/python/fe_api/test_dsa_compiler.py deleted file mode 100644 index 54ad206f7..000000000 --- a/test/python/fe_api/test_dsa_compiler.py +++ /dev/null @@ -1,52 +0,0 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: MIT - -"""Dependency-free contracts for JAX CuTe compilation targets.""" - -from importlib.util import module_from_spec, spec_from_file_location -from pathlib import Path -import sys -import unittest - -try: - import pytest -except ImportError: - pass -else: - pytestmark = pytest.mark.L0 - - -_MODULE_PATH = Path(__file__).resolve().parents[3] / "python" / "cudnn" / "_cute_compiler.py" - - -class JaxCompilerTest(unittest.TestCase): - @classmethod - def setUpClass(cls) -> None: - spec = spec_from_file_location("dsa_compiler_contract", _MODULE_PATH) - assert spec is not None and spec.loader is not None - cls.module = module_from_spec(spec) - spec.loader.exec_module(cls.module) - - def test_explicit_targets_do_not_import_torch(self): - torch_before = sys.modules.get("torch") - self.assertEqual(self.module.gpu_arch_flag_for_compute_capability(90), "sm_90a") - self.assertEqual(self.module.gpu_arch_flag_for_compute_capability(100), "sm_100a") - self.assertEqual(self.module.gpu_arch_flag_for_compute_capability(103), "sm_103a") - self.assertEqual(self.module.gpu_arch_flag_for_compute_capability(107), "sm_100f") - self.assertIs(sys.modules.get("torch"), torch_before) - - def test_builds_jax_compile_options_from_explicit_target(self): - self.assertEqual( - self.module.compile_options_for_target(103, "--opt-level 3"), - "--enable-tvm-ffi --gpu-arch sm_103a --opt-level 3", - ) - - def test_rejects_unknown_or_non_integer_targets(self): - with self.assertRaisesRegex(RuntimeError, "Unsupported GPU compute capability SM101"): - self.module.compile_options_for_target(101) - with self.assertRaisesRegex(TypeError, "must be an int"): - self.module.compile_options_for_target(True) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/python/fe_api/test_dsa_indexer_backward_op.py b/test/python/fe_api/test_dsa_indexer_backward_op.py index e861da4c0..17e8fba78 100644 --- a/test/python/fe_api/test_dsa_indexer_backward_op.py +++ b/test/python/fe_api/test_dsa_indexer_backward_op.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for DSA indexer-backward metadata.""" +"""Contracts for DSA indexer-backward metadata.""" from enum import Enum, auto import importlib @@ -56,8 +56,8 @@ def setUpClass(cls) -> None: sys.modules[operation_name] = operation try: - cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") - cls.base_module = importlib.import_module(f"{_PACKAGE}._op") + cls.tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}.common.op") cls.op_module = importlib.import_module(f"{operation_name}.op") except Exception: cls.tearDownClass() diff --git a/test/python/fe_api/test_dsa_indexer_ops.py b/test/python/fe_api/test_dsa_indexer_ops.py index 736aad8cb..b38ce52b1 100644 --- a/test/python/fe_api/test_dsa_indexer_ops.py +++ b/test/python/fe_api/test_dsa_indexer_ops.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for the DSA indexer logical operations.""" +"""Contracts for the DSA indexer logical operations.""" import ast from enum import Enum, auto @@ -98,8 +98,8 @@ def setUpClass(cls) -> None: module.__spec__ = ModuleSpec(name, loader=None, is_package=True) sys.modules[name] = module - cls.tensor = importlib.import_module(f"{_PACKAGE}._tensor_desc") - importlib.import_module(f"{_PACKAGE}._op") + cls.tensor = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + importlib.import_module(f"{_PACKAGE}.common.op") cls.forward = importlib.import_module(f"{dsa_name}.indexer_forward.op") cls.topk = importlib.import_module(f"{dsa_name}.indexer_top_k.op") @@ -192,6 +192,21 @@ def test_forward_requires_the_kernel_vector_axis_to_be_contiguous(self): with self.assertRaisesRegex(ValueError, "canonical D axis contiguous"): operation.check_support() + def test_sm100_forward_requires_a_complete_packed_query_tile(self): + operation = self.forward.IndexerForwardOp( + q=self.desc((1, 4, 32, 128), _DataType.BFLOAT16), + k=self.desc((1, 5, 1, 128), _DataType.BFLOAT16), + weight=self.desc((1, 4, 32), _DataType.BFLOAT16), + output=self.desc((1, 4, 8), _DataType.FLOAT), + ratio=1, + target_compute_capability=100, + ) + with self.assertRaisesRegex(ValueError, "complete packed query tile"): + operation.check_support() + + operation.target_compute_capability = 90 + self.assertTrue(operation.check_support()) + def make_topk(self, *, return_val=True, **overrides): input_values = self.desc((4, 64), _DataType.FLOAT) arguments = { @@ -388,10 +403,47 @@ def tearDownClass(cls) -> None: @staticmethod def fake_call(captured): def call(_inputs, *, output_descs, **options): + input_descs = options.get("input_descs") + if input_descs is not None: + derived_input_specs = tuple( + DsaIndexerJaxAdapterContractTest.jax_base.JaxApiBase._to_tensor_spec( + desc + ) + for desc in input_descs + ) + supplied_input_specs = options.get("input_spec") + options["input_spec"] = ( + derived_input_specs + if supplied_input_specs is None + else tuple( + derived if supplied is None else supplied + for derived, supplied in zip( + derived_input_specs, + supplied_input_specs, + ) + ) + ) + + derived_output_specs = tuple( + DsaIndexerJaxAdapterContractTest.jax_base.JaxApiBase._to_tensor_spec( + desc + ) + for desc in output_descs + ) + supplied_output_specs = options.get("output_spec") + options["output_spec"] = ( + derived_output_specs + if supplied_output_specs is None + else tuple( + derived if supplied is None else supplied + for derived, supplied in zip( + derived_output_specs, + supplied_output_specs, + ) + ) + ) captured.update(output_descs=output_descs, **options) - output_specs = options.get("output_spec") - if output_specs is None: - output_specs = (None,) * len(output_descs) + output_specs = options["output_spec"] def public_shape(desc, spec): if spec is None or spec.mode is None: diff --git a/test/python/fe_api/test_dsa_score_recompute_op.py b/test/python/fe_api/test_dsa_score_recompute_op.py index c6d02f919..31cd2a4b2 100644 --- a/test/python/fe_api/test_dsa_score_recompute_op.py +++ b/test/python/fe_api/test_dsa_score_recompute_op.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for DSA score-recompute operations.""" +"""Contracts for DSA score-recompute operations.""" from enum import Enum, auto import importlib @@ -57,7 +57,7 @@ def setUpClass(cls) -> None: sys.modules[score_name] = score try: - cls.tensor = importlib.import_module(f"{_PACKAGE}._tensor_desc") + cls.tensor = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") cls.op = importlib.import_module(f"{score_name}.op") cls.config = importlib.import_module(f"{score_name}.config") except Exception: diff --git a/test/python/fe_api/test_dsa_sparse_attention_backward_op.py b/test/python/fe_api/test_dsa_sparse_attention_backward_op.py index da72104a3..f81acc026 100644 --- a/test/python/fe_api/test_dsa_sparse_attention_backward_op.py +++ b/test/python/fe_api/test_dsa_sparse_attention_backward_op.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for DSA sparse-attention backward metadata.""" +"""Contracts for DSA sparse-attention backward metadata.""" from enum import Enum, auto import importlib @@ -58,8 +58,8 @@ def setUpClass(cls) -> None: sys.modules[operation_name] = operation try: - cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") - cls.base_module = importlib.import_module(f"{_PACKAGE}._op") + cls.tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}.common.op") cls.op_module = importlib.import_module(f"{operation_name}.op") except Exception: cls.tearDownClass() diff --git a/test/python/fe_api/test_gemm_amax_op.py b/test/python/fe_api/test_gemm_amax_op.py index b61e6ab57..942a86b0e 100644 --- a/test/python/fe_api/test_gemm_amax_op.py +++ b/test/python/fe_api/test_gemm_amax_op.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for the dense GEMM + amax operation.""" +"""Contracts for the dense GEMM + amax operation.""" from enum import Enum, auto import importlib @@ -55,8 +55,8 @@ def setUpClass(cls) -> None: sys.modules[operation_name] = operation try: - cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") - cls.base_module = importlib.import_module(f"{_PACKAGE}._op") + cls.tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}.common.op") cls.op_module = importlib.import_module(f"{operation_name}.op") except Exception: cls.tearDownClass() diff --git a/test/python/fe_api/test_gemm_relu_imports.py b/test/python/fe_api/test_gemm_relu_imports.py index d336471d6..10dc35344 100644 --- a/test/python/fe_api/test_gemm_relu_imports.py +++ b/test/python/fe_api/test_gemm_relu_imports.py @@ -105,7 +105,9 @@ def test_exports_route_to_op_and_torch_modules_independently(self): def test_framework_neutral_op_modules_have_no_framework_imports(self): paths = [ - _CUDNN_ROOT / "_gemm_relu.py", + _CUDNN_ROOT / "gemm" / "__init__.py", + _CUDNN_ROOT / "gemm" / "helpers.py", + _CUDNN_ROOT / "gemm" / "srelu.py", _CUDNN_ROOT / "gemm_srelu" / "op.py", _CUDNN_ROOT / "gemm_dsrelu" / "op.py", ] diff --git a/test/python/fe_api/test_gemm_swiglu_op.py b/test/python/fe_api/test_gemm_swiglu_op.py index fba11954f..a2b5bb611 100644 --- a/test/python/fe_api/test_gemm_swiglu_op.py +++ b/test/python/fe_api/test_gemm_swiglu_op.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for the dense GEMM + SwiGLU operation.""" +"""Contracts for the dense GEMM + SwiGLU operation.""" from enum import Enum, auto import importlib @@ -53,8 +53,8 @@ def setUpClass(cls) -> None: sys.modules[operation_name] = operation try: - cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") - cls.base_module = importlib.import_module(f"{_PACKAGE}._op") + cls.tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}.common.op") cls.op_module = importlib.import_module(f"{operation_name}.op") except Exception: cls.tearDownClass() diff --git a/test/python/fe_api/test_jax_api_base.py b/test/python/fe_api/test_jax_api_base.py index dbb7bf26e..0a9d1f764 100644 --- a/test/python/fe_api/test_jax_api_base.py +++ b/test/python/fe_api/test_jax_api_base.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for the JAX API base.""" +"""Contracts for the JAX API base.""" from enum import Enum, auto import importlib @@ -38,10 +38,18 @@ def __init__(self, shape, dtype, label=None): class _TensorSpec: - def __init__(self, *, layout=None, mode=None, divisibility=None): + def __init__( + self, + *, + layout=None, + mode=None, + divisibility=None, + ptr_assumed_align=None, + ): self.layout = layout self.mode = mode self.divisibility = divisibility + self.ptr_assumed_align = ptr_assumed_align class _ShapeDtypeStruct: @@ -100,7 +108,7 @@ def setUpClass(cls) -> None: "cutlass.jax": None, }, ): - cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + cls.tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") cls.module = importlib.import_module(f"{internal_name}.api_base") except Exception: cls.tearDownClass() @@ -275,15 +283,31 @@ def test_converts_array_metadata_to_shared_tensor_desc(self): self.assertEqual(desc.cudnn_dtype, _DataType.BFLOAT16) def test_converts_public_array_axes_to_canonical_descriptor_axes(self): + mode = (2, 0, 1) desc = self.module.JaxApiBase._to_tensor_desc( _ArrayMetadata((2, 3, 4), "bfloat16"), "sample", - mode=(2, 0, 1), + mode=mode, ) self.assertEqual(desc.shape, (4, 2, 3)) self.assertEqual(desc.stride, (1, 12, 4)) self.assertEqual(desc.stride_order, (0, 2, 1)) + self.assertEqual(desc.mode, mode) + self.assertEqual(desc.array_shape, (2, 3, 4)) + + def test_constructs_descriptor_directly_from_shape(self): + desc = self.module.JaxTensorDesc.from_shape( + (2, 3, 4), + "bfloat16", + name="sample", + mode=(2, 0, 1), + ) + + self.assertEqual(desc.shape, (4, 2, 3)) + self.assertEqual(desc.array_shape, (2, 3, 4)) + self.assertEqual(desc.mode, (2, 0, 1)) + self.assertEqual(desc.cudnn_dtype, _DataType.BFLOAT16) def test_converts_public_physical_layout_to_canonical_descriptor_axes(self): desc = self.module.JaxApiBase._to_tensor_desc( @@ -329,6 +353,7 @@ def test_jax_descriptor_derives_compact_output_metadata(self): self.assertEqual(output.stride_order, (0, 1)) self.assertEqual(output.name, "output") self.assertEqual(output.init_value, -2.0) + self.assertEqual(output.mode, (0, 1)) def test_checks_invocation_signature(self): desc = self.module.JaxApiBase._to_tensor_desc(_ArrayMetadata((2, 3), "bfloat16"), "sample") @@ -350,13 +375,11 @@ def test_checks_public_invocation_against_canonical_descriptor(self): self.module.JaxApiBase._check_tensor_signature( _ArrayMetadata((2, 3, 4), "bfloat16"), desc, - mode=mode, ) with self.assertRaisesRegex(ValueError, "sample tensor shape mismatch"): self.module.JaxApiBase._check_tensor_signature( _ArrayMetadata((1, 3, 4), "bfloat16"), desc, - mode=mode, ) def test_builds_public_tensor_spec_from_canonical_metadata(self): @@ -365,6 +388,8 @@ def test_builds_public_tensor_spec_from_canonical_metadata(self): _ArrayMetadata((2, 3, 4), "bfloat16"), "sample", mode=mode, + divisibility=(8, 2, 4), + ptr_assumed_align=64, ) cutlass = types.ModuleType("cutlass") cutlass.__path__ = [] @@ -373,15 +398,12 @@ def test_builds_public_tensor_spec_from_canonical_metadata(self): cutlass.jax = cutlass_jax with mock.patch.dict(sys.modules, {"cutlass": cutlass, "cutlass.jax": cutlass_jax}): - spec = self.module.JaxApiBase._to_tensor_spec( - desc, - mode=mode, - divisibility=(8, 2, 4), - ) + spec = self.module.JaxApiBase._to_tensor_spec(desc) self.assertEqual(spec.layout, (2, 1, 0)) self.assertEqual(spec.mode, mode) self.assertEqual(spec.divisibility, (2, 4, 8)) + self.assertEqual(spec.ptr_assumed_align, 64) def test_tensor_spec_preserves_explicit_public_physical_layout(self): mode = (1, 2, 0) @@ -398,12 +420,12 @@ def test_tensor_spec_preserves_explicit_public_physical_layout(self): cutlass.jax = cutlass_jax with mock.patch.dict(sys.modules, {"cutlass": cutlass, "cutlass.jax": cutlass_jax}): - spec = self.module.JaxApiBase._to_tensor_spec(desc, mode=mode) + spec = self.module.JaxApiBase._to_tensor_spec(desc) self.assertEqual(spec.layout, (0, 2, 1)) self.assertEqual(spec.mode, mode) - def test_tensor_spec_defaults_to_identity_axis_binding(self): + def test_default_descriptor_uses_cutlass_default_tensor_spec(self): desc = self.module.JaxApiBase._to_tensor_desc( _ArrayMetadata((2, 3), "bfloat16"), "sample", @@ -415,11 +437,7 @@ def test_tensor_spec_defaults_to_identity_axis_binding(self): cutlass.jax = cutlass_jax with mock.patch.dict(sys.modules, {"cutlass": cutlass, "cutlass.jax": cutlass_jax}): - spec = self.module.JaxApiBase._to_tensor_spec(desc) - - self.assertEqual(spec.layout, (1, 0)) - self.assertEqual(spec.mode, (0, 1)) - self.assertIsNone(spec.divisibility) + self.assertIsNone(self.module.JaxApiBase._to_tensor_spec(desc)) def test_tensor_spec_rejects_divisibility_with_the_wrong_rank(self): desc = self.module.JaxApiBase._to_tensor_desc( @@ -428,54 +446,54 @@ def test_tensor_spec_rejects_divisibility_with_the_wrong_rank(self): ) with self.assertRaisesRegex(ValueError, "divisibility rank mismatch"): - self.module.JaxApiBase._to_tensor_spec(desc, divisibility=(2,)) - - def test_materializes_canonical_output_in_public_axis_order(self): - desc = self.tensor_module.TensorDesc( - dtype=_DataType.BFLOAT16, - shape=(4, 2, 3), - stride=(1, 12, 4), - stride_order=(0, 2, 1), + self.module.JaxTensorDesc.from_array( + _ArrayMetadata((2, 3), "bfloat16"), + name="sample", + divisibility=(2,), + ) + + def test_builds_shape_dtype_struct_in_array_axis_order(self): + desc = self.module.JaxTensorDesc.from_shape( + (2, 3, 4), + "bfloat16", + mode=(2, 0, 1), ) jax = types.ModuleType("jax") jax.ShapeDtypeStruct = _ShapeDtypeStruct with mock.patch.dict(sys.modules, {"jax": jax}): - output = self.module.JaxApiBase._materialize_tensor_desc( - desc, - mode=(2, 0, 1), - ) + output = self.module.JaxApiBase._to_shape_dtype_struct(desc) self.assertEqual(output.shape, (2, 3, 4)) self.assertEqual(output.dtype, "bfloat16") - def test_calls_explicit_launcher_with_initialized_outputs_and_hidden_workspaces(self): - tensor_module = self.tensor_module + def test_calls_explicit_launcher_with_initialized_and_fresh_outputs(self): seen = {} output_descs = ( - tensor_module.make_compact_tensor_desc( - dtype=_DataType.FLOAT, - shape=(2,), + self.module.JaxTensorDesc.from_shape( + (2,), + "float32", name="output", ), - tensor_module.make_compact_tensor_desc( - dtype=_DataType.FLOAT, - shape=(4, 2, 3), + self.module.JaxTensorDesc.from_shape( + (2, 3, 4), + "float32", name="amax", init_value=float("-inf"), + mode=(2, 0, 1), ), ) workspace_descs = ( - tensor_module.make_compact_tensor_desc( - dtype=_DataType.FLOAT, - shape=(3,), + self.module.JaxTensorDesc.from_shape( + (3,), + "float32", name="counter", init_value=0, ), - tensor_module.make_compact_tensor_desc( - dtype=_DataType.BFLOAT16, - shape=(5,), + self.module.JaxTensorDesc.from_shape( + (5,), + "bfloat16", name="scratch", ), ) @@ -491,6 +509,7 @@ def __call__(self, value): return self._call_kernel( (value,), launch=launch, + input_descs=(input_desc,), output_descs=output_descs, workspace_descs=workspace_descs, ) @@ -511,10 +530,23 @@ def full(shape, value, dtype): cutlass = types.ModuleType("cutlass") cutlass.__path__ = [] + cutlass_cute = types.ModuleType("cutlass.cute") + + def cute_jit(*, preprocess): + self.assertFalse(preprocess) + + def decorate(function): + seen["cute_launcher"] = function + return function + + return decorate + + cutlass_cute.jit = cute_jit cutlass_jax = types.ModuleType("cutlass.jax") cutlass_jax.TensorSpec = _TensorSpec def cutlass_call(launcher, **options): + seen["cutlass_launcher"] = launcher seen["call_options"] = options def invoke(*args): @@ -538,12 +570,14 @@ def invoke(*args): return invoke cutlass_jax.cutlass_call = cutlass_call + cutlass.cute = cutlass_cute cutlass.jax = cutlass_jax input_value = _ArrayMetadata((2,), "float32", label="input") - output_seed = _ArrayMetadata((2,), "float32", label="seed") - input_tensor_spec = _TensorSpec(layout=(0,), mode=(0,)) - initialized_output_spec = _TensorSpec(layout=(2, 1, 0), mode=(2, 0, 1)) + input_desc = self.module.JaxTensorDesc.from_array( + input_value, + name="input", + ) api = Adapter() with mock.patch.dict( sys.modules, @@ -551,64 +585,62 @@ def invoke(*args): "jax": jax, "jax.numpy": jnp, "cutlass": cutlass, + "cutlass.cute": cutlass_cute, "cutlass.jax": cutlass_jax, }, ): results = api._call_kernel( (input_value,), launch=launch, + input_descs=(input_desc,), output_descs=output_descs, - output_seeds=(output_seed, None), workspace_descs=workspace_descs, - input_spec=(input_tensor_spec,), - output_spec=(None, initialized_output_spec), ) - with self.assertRaisesRegex(ValueError, "Expected 2 output seeds, got 1"): + with self.assertRaisesRegex(TypeError, "input #0 must have shape and dtype metadata"): api._call_kernel( - (input_value,), + ([input_value],), launch=launch, + input_descs=(input_desc,), output_descs=output_descs, - output_seeds=(output_seed,), + workspace_descs=workspace_descs, ) - with self.assertRaisesRegex(ValueError, "output seed shape mismatch"): + with self.assertRaisesRegex(ValueError, "Expected 1 input descriptors, got 0"): api._call_kernel( (input_value,), launch=launch, + input_descs=(), output_descs=output_descs, - output_seeds=( - _ArrayMetadata((3,), "float32", label="wrong shape"), - None, - ), ) - with self.assertRaisesRegex(ValueError, "output seed dtype mismatch"): + with self.assertRaisesRegex(TypeError, "input_descs must contain JaxTensorDesc"): api._call_kernel( (input_value,), launch=launch, - output_descs=output_descs, - output_seeds=( - _ArrayMetadata((2,), "bfloat16", label="wrong dtype"), - None, + input_descs=( + self.tensor_module.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2,), + ), ), - ) - with self.assertRaisesRegex(TypeError, "input #0 must have shape and dtype metadata"): - api._call_kernel( - ([input_value],), - launch=launch, output_descs=output_descs, - workspace_descs=workspace_descs, ) - with self.assertRaisesRegex(TypeError, "output_spec must contain only TensorSpec or None"): + with self.assertRaisesRegex(TypeError, "output_descs must contain JaxTensorDesc"): api._call_kernel( (input_value,), launch=launch, - output_descs=output_descs, + input_descs=(input_desc,), + output_descs=( + self.tensor_module.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2,), + ), + ), workspace_descs=workspace_descs, - output_spec=((0,), initialized_output_spec), ) with self.assertRaisesRegex(TypeError, "launch must be callable"): api._call_kernel( (input_value,), launch=None, + input_descs=(input_desc,), output_descs=output_descs, ) @@ -621,24 +653,28 @@ def invoke(*args): ) self.assertEqual( seen["call_options"]["input_output_aliases"], - {1: 0, 2: 1, 3: 2}, + {1: 1, 2: 2}, ) - self.assertEqual(len(seen["call_options"]["input_spec"]), 4) + self.assertEqual(len(seen["call_options"]["input_spec"]), 3) + self.assertIsNone(seen["call_options"]["input_spec"][0]) self.assertIs( seen["call_options"]["input_spec"][1], - seen["call_options"]["output_spec"][0], + seen["call_options"]["output_spec"][1], + ) + self.assertIs( + seen["call_options"]["input_spec"][2], + seen["call_options"]["output_spec"][2], ) - self.assertIs(seen["call_options"]["input_spec"][2], initialized_output_spec) - self.assertIs(seen["call_options"]["input_spec"][3], seen["call_options"]["output_spec"][2]) stream, launch_input, output, amax, counter, scratch = seen["launch_args"] self.assertEqual(stream, "stream") + self.assertIs(seen["cute_launcher"], seen["cutlass_launcher"]) self.assertIs(launch_input, input_value) - self.assertIs(output, output_seed) + self.assertEqual(output.label, "allocated(0)") self.assertIs(amax, full_calls[0][3]) self.assertIs(counter, full_calls[1][3]) self.assertEqual(scratch.label, "allocated(3)") - self.assertEqual([result.label for result in results], ["seed", "full(-inf)"]) + self.assertEqual([result.label for result in results], ["allocated(0)", "full(-inf)"]) self.assertNotIn(full_calls[0][3], vars(api).values()) self.assertNotIn(full_calls[1][3], vars(api).values()) diff --git a/test/python/fe_api/test_jax_bsa_contract.py b/test/python/fe_api/test_jax_bsa_contract.py index 1a2e77345..c648bf2ec 100644 --- a/test/python/fe_api/test_jax_bsa_contract.py +++ b/test/python/fe_api/test_jax_bsa_contract.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for the JAX block-sparse attention adapters.""" +"""Contracts for the JAX block-sparse attention adapters.""" from __future__ import annotations @@ -81,22 +81,48 @@ def setUpClass(cls) -> None: internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) sys.modules[internal_name] = internal - tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") layout_module = importlib.import_module(f"{internal_name}.layout") - result_module = importlib.import_module(f"{_PACKAGE}._result") + result_module = importlib.import_module(f"{_PACKAGE}.common.result") class JaxTensorDesc(tensor_module.TensorDesc): + @classmethod + def from_shape( + cls, + shape, + dtype, + *, + name="", + mode=None, + public_stride_order=None, + init_value=None, + ): + return JaxApiBase._to_tensor_desc( + _Array(shape, dtype), + name, + mode=mode, + public_stride_order=public_stride_order, + init_value=init_value, + ) + @property def cudnn_dtype(self): return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) def compact_like( - self, *, cudnn_dtype, shape, stride_order=None, name="", init_value=None + self, + *, + cudnn_dtype, + shape, + stride_order=None, + name="", + init_value=None, + mode=None, ): if stride_order is None: stride_order = tuple(reversed(range(len(shape)))) stride = layout_module.compact_stride(tuple(shape), tuple(stride_order)) - return JaxTensorDesc( + desc = JaxTensorDesc( dtype=_CUDNN_TO_DTYPE[cudnn_dtype], shape=tuple(shape), stride=stride, @@ -104,6 +130,12 @@ def compact_like( name=name, init_value=init_value, ) + object.__setattr__( + desc, + "mode", + layout_module.normalize_mode(len(shape), mode), + ) + return desc class JaxApiBase: @staticmethod @@ -123,7 +155,7 @@ def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): canonical_axis_by_public_axis = layout_module.to_public_axes( tuple(range(len(public_shape))), mode ) - return JaxTensorDesc( + desc = JaxTensorDesc( dtype=value.dtype, shape=layout_module.to_canonical_axes(public_shape, mode), stride=layout_module.to_canonical_axes(public_stride, mode), @@ -133,9 +165,13 @@ def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): name=name, init_value=init_value, ) + object.__setattr__(desc, "mode", mode) + return desc @staticmethod def _check_tensor_signature(value, expected, *, mode=None): + if mode is None: + mode = expected.mode actual_shape = layout_module.to_canonical_axes(tuple(value.shape), mode) if actual_shape != expected.shape: raise ValueError(f"{expected.name} shape mismatch") @@ -147,6 +183,8 @@ def _check_tensor_signature(value, expected, *, mode=None): @staticmethod def _to_tensor_spec(desc, *, mode=None, divisibility=None): + if mode is None: + mode = desc.mode return _TensorSpec( layout=layout_module.to_cutlass_layout( desc.shape, @@ -159,7 +197,22 @@ def _to_tensor_spec(desc, *, mode=None, divisibility=None): divisibility=divisibility, ) - def _call_kernel(self, inputs, *, output_descs, output_spec, **options): + def _call_kernel( + self, + inputs, + *, + output_descs, + input_descs=None, + output_spec=None, + **options, + ): + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc) + if output_spec is None: + output_spec = tuple( + self._to_tensor_spec(desc) for desc in output_descs + ) self.captured_call = { "inputs": tuple(inputs), "output_descs": tuple(output_descs), diff --git a/test/python/fe_api/test_jax_compiler.py b/test/python/fe_api/test_jax_compiler.py new file mode 100644 index 000000000..7ba5efc0d --- /dev/null +++ b/test/python/fe_api/test_jax_compiler.py @@ -0,0 +1,74 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for CUTLASS JAX compilation targets.""" + +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_PACKAGE = "cudnn_jax_compiler_test" + + +class JaxCompilerTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + sys.modules[_PACKAGE] = root + + jax_package_name = f"{_PACKAGE}._jax" + jax_package = types.ModuleType(jax_package_name) + jax_package.__path__ = [str(_CUDNN_ROOT / "_jax")] + jax_package.__package__ = jax_package_name + jax_package.__spec__ = ModuleSpec(jax_package_name, loader=None, is_package=True) + sys.modules[jax_package_name] = jax_package + + cls.arch = importlib.import_module(f"{_PACKAGE}.common.cute_arch") + cls.compiler = importlib.import_module(f"{jax_package_name}.compiler") + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def test_explicit_targets_do_not_import_frameworks(self): + torch_before = sys.modules.get("torch") + jax_before = sys.modules.get("jax") + self.assertEqual(self.arch.gpu_arch_flag_for_compute_capability(90), "sm_90a") + self.assertEqual(self.arch.gpu_arch_flag_for_compute_capability(100), "sm_100a") + self.assertEqual(self.arch.gpu_arch_flag_for_compute_capability(103), "sm_103a") + self.assertEqual(self.arch.gpu_arch_flag_for_compute_capability(107), "sm_100f") + self.assertIs(sys.modules.get("torch"), torch_before) + self.assertIs(sys.modules.get("jax"), jax_before) + + def test_builds_jax_compile_options_from_explicit_target(self): + self.assertEqual( + self.compiler.compile_options_for_target(103, "--opt-level 3"), + "--gpu-arch sm_103a --opt-level 3", + ) + + def test_rejects_unknown_or_non_integer_targets(self): + with self.assertRaisesRegex(RuntimeError, "Unsupported GPU compute capability SM101"): + self.compiler.compile_options_for_target(101) + with self.assertRaisesRegex(TypeError, "must be an int"): + self.compiler.compile_options_for_target(True) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_datatypes.py b/test/python/fe_api/test_jax_datatypes.py index 40d820dea..d6f554274 100644 --- a/test/python/fe_api/test_jax_datatypes.py +++ b/test/python/fe_api/test_jax_datatypes.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free tests for common JAX/cuDNN dtype conversion.""" +"""Tests for common JAX/cuDNN dtype conversion.""" from enum import Enum, auto import importlib diff --git a/test/python/fe_api/test_jax_discrete_grouped_gemm_contract.py b/test/python/fe_api/test_jax_discrete_grouped_gemm_contract.py index a4101532f..1c344c63b 100644 --- a/test/python/fe_api/test_jax_discrete_grouped_gemm_contract.py +++ b/test/python/fe_api/test_jax_discrete_grouped_gemm_contract.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for the discrete grouped GEMM JAX adapters.""" +"""Contracts for the discrete grouped GEMM JAX adapters.""" from __future__ import annotations @@ -101,11 +101,34 @@ def setUpClass(cls) -> None: internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) sys.modules[internal_name] = internal - tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") layout_module = importlib.import_module(f"{internal_name}.layout") - result_module = importlib.import_module(f"{_PACKAGE}._result") + result_module = importlib.import_module(f"{_PACKAGE}.common.result") class JaxTensorDesc(tensor_module.TensorDesc): + @classmethod + def from_shape( + cls, + shape, + dtype, + *, + name="", + mode=None, + public_stride_order=None, + init_value=None, + divisibility=None, + ptr_assumed_align=None, + ): + return JaxApiBase._to_tensor_desc( + _Array(shape, dtype), + name, + mode=mode, + public_stride_order=public_stride_order, + init_value=init_value, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + ) + @property def cudnn_dtype(self): return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) @@ -116,7 +139,16 @@ def _resolve_compute_capability(_target, _supported, _name): return 100 @staticmethod - def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): + def _to_tensor_desc( + value, + name, + *, + mode=None, + init_value=None, + divisibility=None, + ptr_assumed_align=None, + **_unused, + ): public_shape = tuple(value.shape) mode = layout_module.normalize_mode(len(public_shape), mode) public_order = tuple(reversed(range(len(public_shape)))) @@ -124,7 +156,7 @@ def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): canonical_axis_by_public_axis = layout_module.to_public_axes( tuple(range(len(public_shape))), mode ) - return JaxTensorDesc( + desc = JaxTensorDesc( dtype=value.dtype, shape=layout_module.to_canonical_axes(public_shape, mode), stride=layout_module.to_canonical_axes(public_stride, mode), @@ -134,9 +166,15 @@ def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): name=name, init_value=init_value, ) + object.__setattr__(desc, "mode", mode) + object.__setattr__(desc, "divisibility", divisibility) + object.__setattr__(desc, "ptr_assumed_align", ptr_assumed_align) + return desc @staticmethod def _check_tensor_signature(value, expected, *, mode=None): + if mode is None: + mode = expected.mode actual = layout_module.to_canonical_axes(tuple(value.shape), mode) if actual != expected.shape: raise ValueError( @@ -150,10 +188,19 @@ def _check_tensor_signature(value, expected, *, mode=None): @staticmethod def _to_tensor_spec(_desc, *, mode=None, **_unused): - return _TensorSpec(mode) + if mode is None: + mode = _desc.mode + return _TensorSpec( + mode, + ptr_assumed_align=getattr( + _desc, "ptr_assumed_align", None + ), + ) @staticmethod - def _materialize_tensor_desc(desc, *, mode=None): + def _to_shape_dtype_struct(desc, *, mode=None): + if mode is None: + mode = desc.mode return _Array( layout_module.to_public_axes(desc.shape, mode), desc.dtype ) @@ -168,11 +215,23 @@ def _call_kernel( *, launch, output_descs, + input_descs=None, workspace_descs=(), - output_spec=(), + output_spec=None, workspace_spec=None, **options, ): + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc) + if output_spec is None: + output_spec = tuple( + self._to_tensor_spec(desc) for desc in output_descs + ) + if workspace_spec is None: + workspace_spec = tuple( + self._to_tensor_spec(desc) for desc in workspace_descs + ) outputs = tuple( _Array( layout_module.to_public_axes(desc.shape, spec.mode), desc.dtype diff --git a/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py b/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py index 8454d0ee7..57ec4bae7 100644 --- a/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py +++ b/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free JAX contracts for DSA indexer backward.""" +"""JAX contracts for DSA indexer backward.""" from __future__ import annotations @@ -121,7 +121,7 @@ def setUpClass(cls) -> None: operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) sys.modules[operation_name] = operation - tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") class JaxTensorDesc(tensor_module.TensorDesc): @property @@ -129,7 +129,14 @@ def cudnn_dtype(self): return cls.dtype_to_cudnn.get(self.dtype, _DataType.NOT_SET) def compact_like( - self, *, cudnn_dtype, shape, stride_order=None, name="", init_value=None + self, + *, + cudnn_dtype, + shape, + stride_order=None, + name="", + init_value=None, + mode=None, ): canonical = tensor_module.make_compact_tensor_desc( dtype=cudnn_dtype, @@ -138,7 +145,7 @@ def compact_like( name=name, init_value=init_value, ) - return JaxTensorDesc( + desc = JaxTensorDesc( dtype=cls.cudnn_to_dtype[cudnn_dtype], shape=canonical.shape, stride=canonical.stride, @@ -146,6 +153,25 @@ def compact_like( name=name, init_value=init_value, ) + object.__setattr__( + desc, + "mode", + tuple(range(len(shape))) if mode is None else tuple(mode), + ) + return desc + + def with_divisibility(self, divisibility): + desc = JaxTensorDesc( + dtype=self.dtype, + shape=self.shape, + stride=self.stride, + stride_order=self.stride_order, + name=self.name, + init_value=self.init_value, + ) + object.__setattr__(desc, "mode", self.mode) + object.__setattr__(desc, "divisibility", tuple(divisibility)) + return desc class JaxApiBase: captured_call = None @@ -162,7 +188,7 @@ def _to_tensor_desc(value, name, *, mode=None, init_value=None): canonical_axis_by_public = [0] * len(mode) for canonical_axis, public_axis in enumerate(mode): canonical_axis_by_public[public_axis] = canonical_axis - return JaxTensorDesc( + desc = JaxTensorDesc( dtype=value.dtype, shape=tuple(public_shape[axis] for axis in mode), stride=tuple(public_stride[axis] for axis in mode), @@ -173,12 +199,14 @@ def _to_tensor_desc(value, name, *, mode=None, init_value=None): name=name, init_value=init_value, ) + object.__setattr__(desc, "mode", mode) + return desc @staticmethod def _resolve_compute_capability(target, supported, operation_name): del supported, operation_name if target is None: - raise RuntimeError("target required by dependency-free test") + raise RuntimeError("target required by test") return target @staticmethod @@ -189,7 +217,7 @@ def _compute_capability_family(target, supported): @staticmethod def _check_tensor_signature(value, expected, *, mode=None): - mode = tuple(range(len(value.shape))) if mode is None else tuple(mode) + mode = expected.mode if mode is None else tuple(mode) actual_dtype = cls.dtype_to_cudnn.get(value.dtype, _DataType.NOT_SET) actual_shape = tuple(value.shape[axis] for axis in mode) if ( @@ -200,11 +228,51 @@ def _check_tensor_signature(value, expected, *, mode=None): @staticmethod def _to_tensor_spec(desc, *, mode=None, divisibility=None): - del desc + if mode is None: + mode = desc.mode + if divisibility is None: + divisibility = getattr(desc, "divisibility", None) return _TensorSpec(mode=mode, divisibility=divisibility) def _call_kernel(self, inputs, *, launch, **options): options["launch"] = launch + input_descs = options.get("input_descs") + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc) + derived_specs = tuple( + self._to_tensor_spec(desc) for desc in input_descs + ) + supplied_specs = options.get("input_spec") + options["input_spec"] = ( + derived_specs + if supplied_specs is None + else tuple( + derived if supplied is None else supplied + for derived, supplied in zip( + derived_specs, supplied_specs + ) + ) + ) + output_descs = options["output_descs"] + derived_output_specs = tuple( + self._to_tensor_spec(desc) for desc in output_descs + ) + supplied_output_specs = options.get("output_spec") + options["output_spec"] = ( + derived_output_specs + if supplied_output_specs is None + else tuple( + derived if supplied is None else supplied + for derived, supplied in zip( + derived_output_specs, supplied_output_specs + ) + ) + ) + options["workspace_spec"] = tuple( + self._to_tensor_spec(desc) + for desc in options.get("workspace_descs", ()) + ) JaxApiBase.captured_call = (inputs, options) specs = options.get("output_spec") or (None,) * len( options["output_descs"] @@ -212,7 +280,7 @@ def _call_kernel(self, inputs, *, launch, **options): results = [] for desc, spec in zip(options["output_descs"], specs): mode = ( - tuple(range(desc.ndim)) + desc.mode if spec is None or spec.mode is None else tuple(spec.mode) ) diff --git a/test/python/fe_api/test_jax_dsa_indexers.py b/test/python/fe_api/test_jax_dsa_indexers.py index b18b7774a..d41bf4e83 100644 --- a/test/python/fe_api/test_jax_dsa_indexers.py +++ b/test/python/fe_api/test_jax_dsa_indexers.py @@ -58,21 +58,21 @@ def test_jax_dsa_indexer_abstract_shapes(monkeypatch): monkeypatch.setattr(JaxApiBase, "_local_gpu_capabilities", staticmethod(lambda _operation_name: ((object(), 100),))) - q = jax.ShapeDtypeStruct((1, 4, 32, 128), jnp.bfloat16) + q = jax.ShapeDtypeStruct((1, 8, 32, 128), jnp.bfloat16) k = jax.ShapeDtypeStruct((1, 5, 1, 128), jnp.bfloat16) - w = jax.ShapeDtypeStruct((1, 4, 32), jnp.bfloat16) + w = jax.ShapeDtypeStruct((1, 8, 32), jnp.bfloat16) fixed = jax.eval_shape( lambda q, k, w: indexer_forward_wrapper(q, k, w, ratio=1, target_compute_capability=100), q, k, w, )["scores"] - assert fixed.shape == (1, 4, 5) + assert fixed.shape == (1, 8, 5) assert fixed.dtype == jnp.float32 - q_sequence_major = jax.ShapeDtypeStruct((4, 2, 32, 128), jnp.bfloat16) + q_sequence_major = jax.ShapeDtypeStruct((8, 2, 32, 128), jnp.bfloat16) k_sequence_major = jax.ShapeDtypeStruct((5, 2, 1, 128), jnp.bfloat16) - w_sequence_major = jax.ShapeDtypeStruct((4, 2, 32), jnp.bfloat16) + w_sequence_major = jax.ShapeDtypeStruct((8, 2, 32), jnp.bfloat16) sequence_major = jax.eval_shape( lambda q, k, w: indexer_forward_wrapper( q, @@ -89,12 +89,12 @@ def test_jax_dsa_indexer_abstract_shapes(monkeypatch): k_sequence_major, w_sequence_major, )["scores"] - assert sequence_major.shape == (4, 2, 5) + assert sequence_major.shape == (8, 2, 5) assert sequence_major.dtype == jnp.float32 - q_thd = jax.ShapeDtypeStruct((4, 32, 128), jnp.bfloat16) + q_thd = jax.ShapeDtypeStruct((8, 32, 128), jnp.bfloat16) k_thd = jax.ShapeDtypeStruct((5, 1, 128), jnp.bfloat16) - w_thd = jax.ShapeDtypeStruct((4, 32), jnp.bfloat16) + w_thd = jax.ShapeDtypeStruct((8, 32), jnp.bfloat16) cu = jax.ShapeDtypeStruct((2,), jnp.int32) packed = jax.eval_shape( lambda q, k, w, cu_q, cu_k: indexer_forward_wrapper( @@ -104,7 +104,7 @@ def test_jax_dsa_indexer_abstract_shapes(monkeypatch): ratio=1, cu_seqlens_q=cu_q, cu_seqlens_k=cu_k, - max_seqlen_q=4, + max_seqlen_q=8, max_seqlen_k=5, target_compute_capability=100, ), @@ -114,7 +114,7 @@ def test_jax_dsa_indexer_abstract_shapes(monkeypatch): cu, cu, )["scores"] - assert packed.shape == (4, 5) + assert packed.shape == (8, 5) values = jax.ShapeDtypeStruct((2, 64), jnp.float32) lengths = jax.ShapeDtypeStruct((2,), jnp.int32) @@ -180,7 +180,7 @@ def test_jax_indexer_forward_jit_and_numerics(): device, capability = _supported_gpu(jax, 90) from cudnn.jax import indexer_forward_wrapper - batch, seqlen_q, seqlen_k = 2, 4, 5 + batch, seqlen_q, seqlen_k = 2, 8, 5 num_query_heads, head_dim = 32, 128 q = jax.device_put( jax.random.normal(jax.random.key(0), (batch, seqlen_q, num_query_heads, head_dim), dtype=jnp.bfloat16), diff --git a/test/python/fe_api/test_jax_dsa_score_recompute_contract.py b/test/python/fe_api/test_jax_dsa_score_recompute_contract.py index c99f20c57..e2b8a743a 100644 --- a/test/python/fe_api/test_jax_dsa_score_recompute_contract.py +++ b/test/python/fe_api/test_jax_dsa_score_recompute_contract.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for JAX DSA score-recompute adapters.""" +"""Contracts for JAX DSA score-recompute adapters.""" from enum import Enum, auto import ast @@ -68,7 +68,7 @@ def setUpClass(cls) -> None: root.data_type = _DataType sys.modules[_PACKAGE] = root - cls.tensor = importlib.import_module(f"{_PACKAGE}._tensor_desc") + cls.tensor = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") fake_jnp = types.ModuleType("jax.numpy") fake_jnp.bfloat16 = _DataType.BFLOAT16 @@ -90,6 +90,29 @@ def setUpClass(cls) -> None: tensor_module = cls.tensor + class JaxTensorDesc(tensor_module.TensorDesc): + @classmethod + def from_shape( + cls, + shape, + dtype, + *, + name="", + mode=None, + public_stride_order=None, + init_value=None, + ): + return FakeJaxApiBase._to_tensor_desc( + _Array(shape, dtype), + name, + mode=mode, + init_value=init_value, + ) + + @property + def cudnn_dtype(self): + return self.dtype + class FakeJaxApiBase: @staticmethod def _to_tensor_desc(value, name, *, mode=None, init_value=None): @@ -103,7 +126,7 @@ def _to_tensor_desc(value, name, *, mode=None, init_value=None): canonical_axis_by_public = [0] * len(mode) for canonical_axis, public_axis in enumerate(mode): canonical_axis_by_public[public_axis] = canonical_axis - return tensor_module.TensorDesc( + desc = JaxTensorDesc( dtype=value.dtype, shape=tuple(public_shape[axis] for axis in mode), stride=tuple(public_stride[axis] for axis in mode), @@ -114,6 +137,8 @@ def _to_tensor_desc(value, name, *, mode=None, init_value=None): name=name, init_value=init_value, ) + object.__setattr__(desc, "mode", mode) + return desc @staticmethod def _resolve_compute_capability(target, supported, operation_name): @@ -125,7 +150,7 @@ def _resolve_compute_capability(target, supported, operation_name): @staticmethod def _check_tensor_signature(value, expected, *, mode=None): - mode = tuple(range(len(value.shape))) if mode is None else tuple(mode) + mode = expected.mode if mode is None else tuple(mode) if tuple(value.shape[axis] for axis in mode) != expected.shape: raise ValueError(f"{expected.name} tensor shape mismatch") if value.dtype != expected.cudnn_dtype: @@ -133,12 +158,23 @@ def _check_tensor_signature(value, expected, *, mode=None): @staticmethod def _to_tensor_spec(desc, *, mode=None, divisibility=None): - del desc + if mode is None: + mode = desc.mode return _TensorSpec(mode=mode, divisibility=divisibility) def _call_kernel( self, inputs, *, launch, output_descs, workspace_descs=(), **options ): + input_descs = options.get("input_descs") + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc) + options["input_spec"] = tuple( + self._to_tensor_spec(desc) for desc in input_descs + ) + options["output_spec"] = tuple( + self._to_tensor_spec(desc) for desc in output_descs + ) self.captured_call = { "inputs": tuple(inputs), "outputs": tuple(output_descs), @@ -150,7 +186,7 @@ def _call_kernel( outputs = [] for desc, spec in zip(output_descs, specs): mode = ( - tuple(range(desc.ndim)) + desc.mode if spec is None or spec.mode is None else tuple(spec.mode) ) @@ -175,7 +211,7 @@ def _call_kernel( fake_internal_jax.__name__, loader=None, is_package=True ) fake_internal_jax.JaxApiBase = FakeJaxApiBase - fake_internal_jax.JaxTensorDesc = tensor_module.TensorDesc + fake_internal_jax.JaxTensorDesc = JaxTensorDesc fake_internal_jax.TupleDict = _TupleDict dsa_name = f"{_PACKAGE}.deepseek_sparse_attention" @@ -193,9 +229,6 @@ def _call_kernel( utils_name = f"{dsa_name}.utils" utils = types.ModuleType(utils_name) utils.__path__ = [] - compiler = types.ModuleType(f"{utils_name}.compiler") - compiler.compile_options_for_target = lambda target: f"--gpu-arch sm_{target}" - cls.optional_modules = { "jax": fake_jax, "jax.numpy": fake_jnp, @@ -203,7 +236,6 @@ def _call_kernel( dsa_name: dsa, score_name: score, utils_name: utils, - compiler.__name__: compiler, } sys.modules.update(cls.optional_modules) try: @@ -274,7 +306,7 @@ def test_sm100_sparse_infers_output_and_hidden_length_workspace(self): self.assertEqual( (workspace.shape, workspace.cudnn_dtype), ((1, 1), _DataType.INT32) ) - self.assertEqual(api.captured_call["compile_options"], "--gpu-arch sm_100") + self.assertEqual(api.captured_call["compile_options"], "--gpu-arch sm_100a") def test_explicit_topk_length_remains_a_real_input(self): q, k, _, indices = self.sparse_samples() @@ -330,7 +362,7 @@ def test_sparse_sequence_major_layouts_preserve_public_shapes(self): self.assertEqual(kernel_weights.shape, (2, 32, 4)) self.assertEqual( tuple(spec.mode for spec in api.captured_call["input_spec"]), - ((1, 0, 2, 3), (1, 0, 2, 3), None, (1, 0, 2)), + ((1, 0, 2, 3), (1, 0, 2, 3), (0, 1, 2), (1, 0, 2)), ) self.assertEqual( api.captured_call["output_spec"][0].mode, diff --git a/test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py b/test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py index 26ee1d383..8ba4f7beb 100644 --- a/test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py +++ b/test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free JAX contracts for DSA sparse-attention backward.""" +"""JAX contracts for DSA sparse-attention backward.""" from __future__ import annotations @@ -88,7 +88,7 @@ def setUpClass(cls) -> None: operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) sys.modules[operation_name] = operation - tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") class JaxTensorDesc(tensor_module.TensorDesc): @property @@ -112,6 +112,18 @@ def compact_like(self, *, cudnn_dtype, shape, stride_order=None, name="", init_v init_value=init_value, ) + def with_divisibility(self, divisibility): + desc = JaxTensorDesc( + dtype=self.dtype, + shape=self.shape, + stride=self.stride, + stride_order=self.stride_order, + name=self.name, + init_value=self.init_value, + ) + object.__setattr__(desc, "divisibility", tuple(divisibility)) + return desc + class JaxApiBase: captured_call = None @@ -138,7 +150,7 @@ def _to_tensor_desc(value, name, *, mode=None, init_value=None): def _resolve_compute_capability(target, supported, operation_name): del supported, operation_name if target is None: - raise RuntimeError("target required by dependency-free test") + raise RuntimeError("target required by test") return target @staticmethod @@ -162,6 +174,9 @@ def _call_kernel(self, inputs, *, launch, **options): return tuple(_Array(desc.shape, _CUDNN_TO_DTYPE[desc.cudnn_dtype]) for desc in options["output_descs"]) internal = types.ModuleType(f"{_PACKAGE}._jax") + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal.__name__ + internal.__spec__ = ModuleSpec(internal.__name__, loader=None, is_package=True) internal.JaxApiBase = JaxApiBase internal.JaxTensorDesc = JaxTensorDesc internal.TupleDict = _TupleDict diff --git a/test/python/fe_api/test_jax_gemm_amax.py b/test/python/fe_api/test_jax_gemm_amax.py index 09a563b51..bb21cc7be 100644 --- a/test/python/fe_api/test_jax_gemm_amax.py +++ b/test/python/fe_api/test_jax_gemm_amax.py @@ -65,10 +65,12 @@ def test_jax_gemm_amax_abstract_contract(monkeypatch): from cudnn._jax import JaxApiBase from cudnn.jax import GemmAmaxSm100, gemm_amax_wrapper_sm100 - def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): + def abstract_call(self, _inputs, *, output_descs, output_spec=None, **_options): + if output_spec is None: + output_spec = tuple(self._to_tensor_spec(desc) for desc in output_descs) values = [] for desc, spec in zip(output_descs, output_spec): - shape = self._materialize_tensor_desc(desc, mode=spec.mode).shape + shape = self._to_shape_dtype_struct(desc, mode=spec.mode).shape fill = 0 if desc.init_value is None else desc.init_value values.append(jnp.full(shape, fill, dtype=desc.dtype)) return tuple(values) diff --git a/test/python/fe_api/test_jax_gemm_amax_contract.py b/test/python/fe_api/test_jax_gemm_amax_contract.py index e510c6d5c..9cdd1cf30 100644 --- a/test/python/fe_api/test_jax_gemm_amax_contract.py +++ b/test/python/fe_api/test_jax_gemm_amax_contract.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for the JAX GEMM + amax adapter.""" +"""Contracts for the JAX GEMM + amax adapter.""" from __future__ import annotations @@ -90,11 +90,29 @@ def setUpClass(cls) -> None: internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) sys.modules[internal_name] = internal - tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") layout_module = importlib.import_module(f"{internal_name}.layout") - result_module = importlib.import_module(f"{_PACKAGE}._result") + result_module = importlib.import_module(f"{_PACKAGE}.common.result") class JaxTensorDesc(tensor_module.TensorDesc): + @classmethod + def from_shape(cls, shape, dtype, *, name, mode=None, init_value=None): + public_shape = tuple(shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + public_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride(public_shape, public_order) + canonical_axis_by_public_axis = layout_module.to_public_axes(tuple(range(len(public_shape))), mode) + desc = cls( + dtype=dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple(canonical_axis_by_public_axis[axis] for axis in public_order), + name=name, + init_value=init_value, + ) + object.__setattr__(desc, "mode", mode) + return desc + @property def cudnn_dtype(self): return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) @@ -121,7 +139,7 @@ def _to_tensor_desc( tuple(range(len(public_shape))), mode, ) - return JaxTensorDesc( + desc = JaxTensorDesc( dtype=value.dtype, shape=layout_module.to_canonical_axes(public_shape, mode), stride=layout_module.to_canonical_axes(public_stride, mode), @@ -129,6 +147,8 @@ def _to_tensor_desc( name=name, init_value=init_value, ) + object.__setattr__(desc, "mode", mode) + return desc @staticmethod def _resolve_compute_capability(target, supported, operation_name): @@ -163,6 +183,15 @@ def _to_tensor_spec(desc, *, mode=None, divisibility=None): ) def _call_kernel(self, inputs, **options): + input_descs = options.get("input_descs") + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc, mode=desc.mode) + options["input_spec"] = tuple(self._to_tensor_spec(desc, mode=desc.mode) for desc in input_descs) + if "output_spec" not in options: + options["output_spec"] = tuple( + self._to_tensor_spec(desc, mode=getattr(desc, "mode", None)) for desc in options["output_descs"] + ) self.captured_call = (tuple(inputs), options) return tuple( _Array( @@ -261,10 +290,10 @@ def test_default_layouts_bind_row_major_scales_and_infer_outputs(self): samples = self._samples() api = self.module.GemmAmaxSm100(*samples) - self.assertEqual(api.a_mode, (1, 2, 0)) - self.assertEqual(api.b_mode, (1, 2, 0)) - self.assertEqual(api.output_mode, (1, 2, 0)) - self.assertEqual(api.scale_mode, (3, 4, 1, 5, 2, 0)) + self.assertEqual(api.a_desc.mode, (1, 2, 0)) + self.assertEqual(api.b_desc.mode, (1, 2, 0)) + self.assertEqual(api.c_desc.mode, (1, 2, 0)) + self.assertEqual(api.sfa_desc.mode, (3, 4, 1, 5, 2, 0)) self.assertEqual(api.sfa_desc.shape, (32, 4, 1, 4, 1, 2)) self.assertEqual(api.sfa_desc.stride_order, (3, 1, 0, 4, 2, 5)) self.assertEqual(api.amax_desc.init_value, float("-inf")) @@ -280,7 +309,7 @@ def test_default_layouts_bind_row_major_scales_and_infer_outputs(self): self.assertTrue(callable(options["launch"])) self.assertEqual( tuple(spec.mode for spec in options["input_spec"]), - (api.a_mode, api.b_mode, api.scale_mode, api.scale_mode), + (api.a_desc.mode, api.b_desc.mode, api.sfa_desc.mode, api.sfb_desc.mode), ) self.assertEqual(options["input_spec"][2].layout, (5, 4, 3, 2, 1, 0)) self.assertEqual(options["input_spec"][3].layout, (5, 4, 3, 2, 1, 0)) @@ -333,7 +362,6 @@ def test_explicit_outputs_and_runtime_signatures_are_checked(self): *samples, sample_c=_Array((2, 128, 256), "bfloat16"), sample_amax=_Array((1, 1, 1), "float32"), - c_dtype="bfloat16", ) result = api(*samples) self.assertEqual(result["c_tensor"].dtype, "bfloat16") @@ -346,6 +374,13 @@ def test_explicit_outputs_and_runtime_signatures_are_checked(self): *samples, sample_c=_Array((2, 128, 256), "float32"), ) + with self.assertRaisesRegex(ValueError, "c_dtype cannot be specified"): + self.module.GemmAmaxSm100( + *samples, + sample_c=_Array((2, 128, 256), "bfloat16"), + sample_amax=_Array((1, 1, 1), "float32"), + c_dtype="bfloat16", + ) with self.assertRaisesRegex(ValueError, "sample_a tensor shape mismatch"): api(_Array((2, 64, 128), samples[0].dtype), *samples[1:]) diff --git a/test/python/fe_api/test_jax_gemm_layout.py b/test/python/fe_api/test_jax_gemm_layout.py index 5ce45a971..fab222835 100644 --- a/test/python/fe_api/test_jax_gemm_layout.py +++ b/test/python/fe_api/test_jax_gemm_layout.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for dense JAX GEMM axis bindings.""" +"""Contracts for dense JAX GEMM axis bindings.""" import importlib from importlib.machinery import ModuleSpec diff --git a/test/python/fe_api/test_jax_gemm_relu.py b/test/python/fe_api/test_jax_gemm_relu.py index ef12440ba..94cdf8469 100644 --- a/test/python/fe_api/test_jax_gemm_relu.py +++ b/test/python/fe_api/test_jax_gemm_relu.py @@ -50,10 +50,12 @@ def test_jax_gemm_relu_abstract_contract(monkeypatch): from cudnn._jax import JaxApiBase from cudnn.jax import GemmDsreluSm100, GemmSreluSm100, gemm_dsrelu_wrapper_sm100, gemm_srelu_wrapper_sm100 - def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): + def abstract_call(self, _inputs, *, output_descs, output_spec=None, **_options): + if output_spec is None: + output_spec = tuple(self._to_tensor_spec(desc) for desc in output_descs) return tuple( jnp.empty( - self._materialize_tensor_desc(desc, mode=spec.mode).shape, + self._to_shape_dtype_struct(desc, mode=spec.mode).shape, dtype=desc.dtype, ) for desc, spec in zip(output_descs, output_spec) @@ -107,6 +109,8 @@ def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): probability, c_layout="LNM", sf_vec_size=32, + a_layout="LKM", + b_layout="LKN", ), jax.ShapeDtypeStruct((batch, k, m), jnp.float8_e4m3fn), jax.ShapeDtypeStruct((batch, k, n), jnp.float8_e4m3fn), @@ -116,13 +120,11 @@ def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): ) assert alternate["c_tensor"].shape == (batch, n, m) - norm_const = jax.ShapeDtypeStruct((1,), jnp.float32) - with pytest.raises(NotImplementedError, match="does not implement SFD generation"): + with pytest.raises(NotImplementedError, match="FP8 D output is unavailable"): jax.eval_shape( lambda *args: gemm_dsrelu_wrapper_sm100( - *args[:-1], + *args, d_dtype=jnp.float8_e4m3fn, - norm_const_tensor=args[-1], sf_vec_size=32, ), a, @@ -131,7 +133,6 @@ def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): scales, scales, prob, - norm_const, ) fp4_dtype = getattr(jnp, "float4_e2m1fn", None) diff --git a/test/python/fe_api/test_jax_gemm_relu_contract.py b/test/python/fe_api/test_jax_gemm_relu_contract.py index 6f6f3feff..fc690e3fb 100644 --- a/test/python/fe_api/test_jax_gemm_relu_contract.py +++ b/test/python/fe_api/test_jax_gemm_relu_contract.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for the JAX sReLU and dsReLU adapters.""" +"""Contracts for the JAX sReLU and dsReLU adapters.""" from __future__ import annotations @@ -89,11 +89,29 @@ def setUpClass(cls) -> None: internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) sys.modules[internal_name] = internal - tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") layout_module = importlib.import_module(f"{internal_name}.layout") - result_module = importlib.import_module(f"{_PACKAGE}._result") + result_module = importlib.import_module(f"{_PACKAGE}.common.result") class JaxTensorDesc(tensor_module.TensorDesc): + @classmethod + def from_shape(cls, shape, dtype, *, name, mode=None, init_value=None): + public_shape = tuple(shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + public_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride(public_shape, public_order) + canonical_axis_by_public_axis = layout_module.to_public_axes(tuple(range(len(public_shape))), mode) + desc = cls( + dtype=dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple(canonical_axis_by_public_axis[axis] for axis in public_order), + name=name, + init_value=init_value, + ) + object.__setattr__(desc, "mode", mode) + return desc + @property def cudnn_dtype(self): return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) @@ -123,7 +141,7 @@ def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): public_order = tuple(reversed(range(len(public_shape)))) public_stride = layout_module.compact_stride(public_shape, public_order) canonical_axis_by_public_axis = layout_module.to_public_axes(tuple(range(len(public_shape))), mode) - return JaxTensorDesc( + desc = JaxTensorDesc( dtype=value.dtype, shape=layout_module.to_canonical_axes(public_shape, mode), stride=layout_module.to_canonical_axes(public_stride, mode), @@ -131,6 +149,8 @@ def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): name=name, init_value=init_value, ) + object.__setattr__(desc, "mode", mode) + return desc @staticmethod def _resolve_compute_capability(target, supported, operation_name): @@ -164,6 +184,15 @@ def _to_tensor_spec(desc, *, mode=None, divisibility=None): ) def _call_kernel(self, inputs, **options): + input_descs = options.get("input_descs") + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc, mode=desc.mode) + options["input_spec"] = tuple(self._to_tensor_spec(desc, mode=desc.mode) for desc in input_descs) + if "output_spec" not in options: + options["output_spec"] = tuple( + self._to_tensor_spec(desc, mode=getattr(desc, "mode", None)) for desc in options["output_descs"] + ) self.captured_call = (tuple(inputs), options) return tuple( _Array( @@ -271,7 +300,13 @@ def test_forward_and_backward_use_row_major_public_auxiliary_layouts(self): self.assertTrue(callable(forward_options["launch"])) self.assertEqual( tuple(spec.mode for spec in forward_options["input_spec"]), - (forward.a_mode, forward.b_mode, forward.scale_mode, forward.scale_mode, forward.probability_mode), + ( + forward.a_desc.mode, + forward.b_desc.mode, + forward.sfa_desc.mode, + forward.sfb_desc.mode, + forward.prob_desc.mode, + ), ) self.assertEqual(forward_options["input_spec"][2].layout, (5, 4, 3, 2, 1, 0)) self.assertEqual(forward_options["input_spec"][4].layout, (2, 1, 0)) @@ -294,7 +329,6 @@ def test_native_fp4_returns_initialized_amax(self): def test_fp8_d_output_is_rejected_until_sfd_generation_is_implemented(self): a, b, c, sfa, sfb, prob = self._samples() - norm_const = _Array((1,), "float32") constructors = ( lambda: self.srelu.GemmSreluSm100( a, @@ -302,7 +336,6 @@ def test_fp8_d_output_is_rejected_until_sfd_generation_is_implemented(self): sfa, sfb, prob, - sample_norm_const=norm_const, d_dtype="float8_e4m3fn", sf_vec_size=32, ), @@ -313,14 +346,68 @@ def test_fp8_d_output_is_rejected_until_sfd_generation_is_implemented(self): sfa, sfb, prob, - sample_norm_const=norm_const, d_dtype="float8_e4m3fn", sf_vec_size=32, ), ) for construct in constructors: - with self.subTest(construct=construct), self.assertRaisesRegex(NotImplementedError, "does not implement SFD generation"): - construct() + with self.subTest(construct=construct), self.assertRaisesRegex(NotImplementedError, "FP8 D output is unavailable"): + construct().check_support() + + def test_explicit_output_samples_are_alternative_to_dtype_arguments(self): + a, b, c, sfa, sfb, prob = self._samples() + sample_d = _Array(c.shape, "bfloat16") + sample_dprob = _Array(prob.shape, "float32") + + forward = self.srelu.GemmSreluSm100( + a, + b, + sfa, + sfb, + prob, + sample_c=c, + sample_d=sample_d, + sf_vec_size=32, + ) + backward = self.dsrelu.GemmDsreluSm100( + a, + b, + c, + sfa, + sfb, + prob, + sample_d=sample_d, + sample_dprob=sample_dprob, + sf_vec_size=32, + ) + self.assertTrue(forward.check_support()) + self.assertTrue(backward.check_support()) + + with self.assertRaisesRegex(ValueError, "c_dtype and d_dtype cannot be specified"): + self.srelu.GemmSreluSm100( + a, + b, + sfa, + sfb, + prob, + sample_c=c, + sample_d=sample_d, + c_dtype="bfloat16", + sf_vec_size=32, + ) + with self.assertRaisesRegex(ValueError, "d_dtype cannot be specified"): + self.dsrelu.GemmDsreluSm100( + a, + b, + c, + sfa, + sfb, + prob, + sample_d=sample_d, + sample_dprob=sample_dprob, + d_dtype="bfloat16", + sf_vec_size=32, + ) def test_unit_scale_tiles_and_batch_preserve_packed_layout(self): a, b, c, sfa, sfb, prob = self._samples(batch=1, m=128, n=128, k=128) diff --git a/test/python/fe_api/test_jax_gemm_swiglu.py b/test/python/fe_api/test_jax_gemm_swiglu.py index 6f6737eb2..28d3a18d2 100644 --- a/test/python/fe_api/test_jax_gemm_swiglu.py +++ b/test/python/fe_api/test_jax_gemm_swiglu.py @@ -52,10 +52,12 @@ def test_jax_gemm_swiglu_abstract_contract( from cudnn._jax import JaxApiBase from cudnn.jax import GemmSwigluSm100, gemm_swiglu_wrapper_sm100 - def abstract_call(self, _inputs, *, output_descs, output_spec, **_options): + def abstract_call(self, _inputs, *, output_descs, output_spec=None, **_options): + if output_spec is None: + output_spec = tuple(self._to_tensor_spec(desc) for desc in output_descs) return tuple( jnp.empty( - self._materialize_tensor_desc(desc, mode=spec.mode).shape, + self._to_shape_dtype_struct(desc, mode=spec.mode).shape, dtype=desc.dtype, ) for desc, spec in zip(output_descs, output_spec) diff --git a/test/python/fe_api/test_jax_gemm_swiglu_contract.py b/test/python/fe_api/test_jax_gemm_swiglu_contract.py index 32b368528..8d3e8611f 100644 --- a/test/python/fe_api/test_jax_gemm_swiglu_contract.py +++ b/test/python/fe_api/test_jax_gemm_swiglu_contract.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for the standard JAX GEMM + SwiGLU adapter.""" +"""Contracts for the standard JAX GEMM + SwiGLU adapter.""" from __future__ import annotations @@ -88,11 +88,29 @@ def setUpClass(cls) -> None: internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) sys.modules[internal_name] = internal - tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") layout_module = importlib.import_module(f"{internal_name}.layout") - result_module = importlib.import_module(f"{_PACKAGE}._result") + result_module = importlib.import_module(f"{_PACKAGE}.common.result") class JaxTensorDesc(tensor_module.TensorDesc): + @classmethod + def from_shape(cls, shape, dtype, *, name, mode=None, init_value=None): + public_shape = tuple(shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + public_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride(public_shape, public_order) + canonical_axis_by_public_axis = layout_module.to_public_axes(tuple(range(len(public_shape))), mode) + desc = cls( + dtype=dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple(canonical_axis_by_public_axis[axis] for axis in public_order), + name=name, + init_value=init_value, + ) + object.__setattr__(desc, "mode", mode) + return desc + @property def cudnn_dtype(self): return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) @@ -105,7 +123,7 @@ def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): public_order = tuple(reversed(range(len(public_shape)))) public_stride = layout_module.compact_stride(public_shape, public_order) canonical_axis_by_public_axis = layout_module.to_public_axes(tuple(range(len(public_shape))), mode) - return JaxTensorDesc( + desc = JaxTensorDesc( dtype=value.dtype, shape=layout_module.to_canonical_axes(public_shape, mode), stride=layout_module.to_canonical_axes(public_stride, mode), @@ -113,6 +131,8 @@ def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): name=name, init_value=init_value, ) + object.__setattr__(desc, "mode", mode) + return desc @staticmethod def _resolve_compute_capability(target, supported, operation_name): @@ -146,6 +166,15 @@ def _to_tensor_spec(desc, *, mode=None, divisibility=None): ) def _call_kernel(self, inputs, *, launch, **options): + input_descs = options.get("input_descs") + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc, mode=desc.mode) + options["input_spec"] = tuple(self._to_tensor_spec(desc, mode=desc.mode) for desc in input_descs) + if "output_spec" not in options: + options["output_spec"] = tuple( + self._to_tensor_spec(desc, mode=getattr(desc, "mode", None)) for desc in options["output_descs"] + ) options["launch"] = launch self.captured_call = (tuple(inputs), options) return tuple( @@ -236,9 +265,9 @@ def test_default_layouts_infer_outputs_and_declare_cutlass_call(self): sample_b, ) - self.assertEqual(api.a_mode, (1, 2, 0)) - self.assertEqual(api.b_mode, (1, 2, 0)) - self.assertEqual(api.output_mode, (1, 2, 0)) + self.assertEqual(api.a_desc.mode, (1, 2, 0)) + self.assertEqual(api.b_desc.mode, (1, 2, 0)) + self.assertEqual(api.c_desc.mode, (1, 2, 0)) self.assertEqual(api.a_desc.shape, (128, 64, 3)) self.assertEqual(api.b_desc.shape, (192, 64, 3)) self.assertEqual(api.ab12_desc.shape, (128, 192, 3)) @@ -266,11 +295,11 @@ def test_default_layouts_infer_outputs_and_declare_cutlass_call(self): self.assertEqual(options["output_descs"], (api.ab12_desc, api.c_desc)) self.assertEqual( tuple(spec.mode for spec in options["input_spec"]), - (api.a_mode, api.b_mode), + (api.a_desc.mode, api.b_desc.mode), ) self.assertEqual( tuple(spec.mode for spec in options["output_spec"]), - (api.output_mode, api.output_mode), + (api.ab12_desc.mode, api.c_desc.mode), ) self.assertEqual( tuple(spec.layout for spec in options["input_spec"]), @@ -293,9 +322,9 @@ def test_alternate_layouts_map_public_shapes_to_the_same_canonical_op(self): ) result = api(sample_a, sample_b) - self.assertEqual(api.a_mode, (2, 1, 0)) - self.assertEqual(api.b_mode, (2, 1, 0)) - self.assertEqual(api.output_mode, (2, 1, 0)) + self.assertEqual(api.a_desc.mode, (2, 1, 0)) + self.assertEqual(api.b_desc.mode, (2, 1, 0)) + self.assertEqual(api.c_desc.mode, (2, 1, 0)) self.assertEqual(api.a_desc.shape, (128, 64, 3)) self.assertEqual(api.b_desc.shape, (192, 64, 3)) self.assertEqual(api.a_desc.stride_order, (0, 1, 2)) @@ -315,8 +344,6 @@ def test_explicit_output_exemplars_are_checked_and_preserved(self): sample_b, sample_ab12=sample_ab12, sample_c=sample_c, - ab12_dtype="bfloat16", - c_dtype="bfloat16", ) result = api(sample_a, sample_b) self.assertEqual(result["ab12_tensor"].dtype, "bfloat16") @@ -328,7 +355,7 @@ def test_explicit_output_exemplars_are_checked_and_preserved(self): sample_b, sample_ab12=sample_ab12, ) - with self.assertRaisesRegex(ValueError, "c_dtype=float16 does not match the explicit sample dtype"): + with self.assertRaisesRegex(ValueError, "ab12_dtype and c_dtype cannot be specified"): self.module.GemmSwigluSm100( sample_a, sample_b, @@ -348,24 +375,17 @@ def test_runtime_inputs_must_match_the_specialized_signature(self): with self.assertRaisesRegex(ValueError, "sample_b tensor dtype mismatch"): api(sample_a, _Array(sample_b.shape, "float16")) - def test_rejects_unreachable_fp8_c_and_block_only_standard_options(self): + def test_rejects_fp8_c_and_block_only_standard_options(self): sample_a, sample_b = self._samples() sample_ab12 = _Array((3, 128, 192), "bfloat16") sample_c_fp8 = _Array((3, 128, 96), "float8_e4m3fn") - with self.assertRaisesRegex(NotImplementedError, "does not support FP8 C"): + with self.assertRaisesRegex(ValueError, "C dtype must be float16 or bfloat16"): self.module.GemmSwigluSm100( sample_a, sample_b, sample_ab12=sample_ab12, sample_c=sample_c_fp8, - ) - - with self.assertRaisesRegex(NotImplementedError, "does not expose FP8 C or SFC"): - self.module.GemmSwigluSm100( - sample_a, - sample_b, - sample_sfc=_Array((1,), "float8_e8m0fnu"), - ) + ).check_support() for option, value in (("sf_vec_size", 32), ("vector_f32", True), ("ab12_stages", 3)): with self.subTest(option=option), self.assertRaisesRegex(ValueError, "only applies to block-scaled"): @@ -425,7 +445,7 @@ def test_block_scaled_mxfp8_uses_packed_scale_modes_and_kernel_abi(self): ) self.assertTrue(api.is_block_scaled) - self.assertEqual(api.scale_mode, (3, 4, 1, 5, 2, 0)) + self.assertEqual(api.sfa_desc.mode, (3, 4, 1, 5, 2, 0)) self.assertEqual(api.sfa_desc.shape, (32, 4, 1, 4, 1, batch)) self.assertEqual(api.sfa_desc.stride_order, (3, 1, 0, 4, 2, 5)) result = api(sample_a, sample_b, sample_sfa, sample_sfb) @@ -440,9 +460,12 @@ def test_block_scaled_mxfp8_uses_packed_scale_modes_and_kernel_abi(self): self.assertEqual(inputs, (sample_a, sample_b, sample_sfa, sample_sfb)) self.assertEqual( tuple(spec.mode for spec in options["input_spec"]), - (api.a_mode, api.b_mode, api.scale_mode, api.scale_mode), + (api.a_desc.mode, api.b_desc.mode, api.sfa_desc.mode, api.sfb_desc.mode), + ) + self.assertEqual( + tuple(spec.mode for spec in options["output_spec"]), + (api.ab12_desc.mode, api.c_desc.mode), ) - self.assertEqual(tuple(spec.mode for spec in options["output_spec"]), (api.output_mode, api.output_mode)) self.assertEqual(api.captured_active_cluster_queries, [(1, 0)]) seen = {} diff --git a/test/python/fe_api/test_jax_grouped_gemm_contract.py b/test/python/fe_api/test_jax_grouped_gemm_contract.py index 854772c06..97cd8c8b1 100644 --- a/test/python/fe_api/test_jax_grouped_gemm_contract.py +++ b/test/python/fe_api/test_jax_grouped_gemm_contract.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free source contracts for JAX grouped GEMM adapters.""" +"""Source contracts for JAX grouped GEMM adapters.""" from __future__ import annotations @@ -120,9 +120,9 @@ def test_shared_adapter_uses_current_jax_api_base_and_call_kernel(): def test_grouped_layouts_are_explicit_row_major_public_mappings(): source = (_GROUPED_ROOT / "_jax_api.py").read_text() - assert "return _row_major_spec(6, BLOCK_SCALE_MODE)" in source - assert "return _row_major_spec(3, (2, 1, 0))" in source - assert "return _row_major_spec(2, (1, 0))" in source + assert "PROBABILITY_MODE = (2, 1, 0)" in source + assert "GROUPED_BIAS_MODE = (1, 0)" in source + assert "GROUPED_WORKSPACE_ALIGNMENT = 128" in source assert "_canonical_block_scale_shape(rows, k, batch, sf_vec_size)" in source assert "BLOCK_SCALE_MODE" in source @@ -149,10 +149,36 @@ def test_grouped_buffers_use_tensor_descriptors_directly(): source = (_GROUPED_ROOT / "_jax_api.py").read_text() assert "class BufferSpec" not in source + assert "class _SampleDesc" not in source assert "def make_buffer_desc(" in source + assert "JaxTensorDesc.from_shape(" in source + assert "JaxTensorDesc.from_array(" in source assert "init_value=init_value" in source +@pytest.mark.parametrize("family", _FAMILIES) +def test_grouped_lowerings_pass_descriptors_without_parallel_specs(family: str): + tree = ast.parse(_jax_path(family).read_text()) + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "call_cutedsl" + ] + + assert calls + for call in calls: + keywords = {keyword.arg for keyword in call.keywords} + assert "input_descs" in keywords + assert "outputs" in keywords + assert not { + "input_specs", + "output_specs", + "workspace_specs", + } & keywords + + @pytest.mark.parametrize("family", _NATIVE_FP4_FAMILIES) def test_dense_grouped_adapters_use_native_fp4_recipe_validation(family: str): source = _jax_path(family).read_text() @@ -338,20 +364,22 @@ def __init__(self, shape, dtype, fill_value=None): with mock.patch.dict(sys.modules, modules): adapter = importlib.import_module(f"{package_name}.grouped_gemm._jax_api") - assert adapter.gemm_a_tensor_spec("LMK").mode == (1, 2, 0) - assert adapter.gemm_b_tensor_spec("LKN").mode == (2, 1, 0) - assert adapter.probability_tensor_spec().mode == (2, 1, 0) - assert adapter.grouped_bias_tensor_spec().mode == (1, 0) + assert adapter.gemm_a_mode("LMK") == (1, 2, 0) + assert adapter.gemm_b_mode("LKN") == (2, 1, 0) + assert adapter.PROBABILITY_MODE == (2, 1, 0) + assert adapter.GROUPED_BIAS_MODE == (1, 0) assert adapter.block_scale_shape(256, 64, 4, 32) == (4, 2, 1, 32, 4, 4) sample = ShapeDtypeStruct((1, 256, 64), "float8") desc = adapter.ApiBaseJax().make_tensor_desc( sample, - tensor_spec=adapter.gemm_a_tensor_spec("LMK"), + mode=adapter.gemm_a_mode("LMK"), name="sample_a", ) + assert type(desc) is adapter.JaxTensorDesc assert desc.shape == (256, 64, 1) assert desc.stride == (64, 1, 16384) + assert desc.mode == (1, 2, 0) assert desc.array_shape == sample.shape fp4_sfa = ShapeDtypeStruct( @@ -464,6 +492,12 @@ def fake_call_kernel(inputs, **options): options["launch"]("stream", *inputs, "output", "workspace") return ("result",) + input_desc = adapter.as_gemm_tensor_desc( + "input", + sample, + mode=adapter.gemm_a_mode("LMK"), + ) + output_mode = adapter.gemm_output_mode("LMN", name="output_layout") with ( mock.patch.object( adapter._CALLER, "_get_max_active_clusters", return_value=7 @@ -478,16 +512,23 @@ def fake_call_kernel(inputs, **options): result = adapter.call_cutedsl( launch, (sample,), + input_descs=(input_desc,), outputs=( adapter.make_buffer_desc( "output", (1, 256, 128), "bfloat16", - tensor_spec=adapter.gemm_c_tensor_spec("LMN"), + mode=output_mode, + ), + ), + workspaces=( + adapter.make_buffer_desc( + "workspace", + (4,), + "uint8", + ptr_assumed_align=adapter.GROUPED_WORKSPACE_ALIGNMENT, ), ), - workspaces=(adapter.make_buffer_desc("workspace", (4,), "uint8"),), - input_specs=(adapter.gemm_a_tensor_spec("LMK"),), static_args={ "cluster_shape_mn": (2, 1), "cluster_overlap_margin": 1, @@ -511,11 +552,26 @@ def fake_call_kernel(inputs, **options): ) ] assert lowered["output_descs"][0].shape == (256, 128, 1) + assert type(lowered["output_descs"][0]) is adapter.JaxTensorDesc + assert lowered["output_descs"][0].mode == output_mode + assert lowered["input_descs"] == (input_desc,) assert lowered["workspace_descs"][0].shape == (4,) + assert ( + lowered["workspace_descs"][0].ptr_assumed_align + == adapter.GROUPED_WORKSPACE_ALIGNMENT + ) + + with pytest.raises(ValueError, match="Expected 1 input descriptors, got 0"): + adapter.call_cutedsl( + launch, + (sample,), + input_descs=(), + outputs=(adapter.make_buffer_desc("output", (1, 0, 128), "bfloat16"),), + ) - def materialize(desc, *, mode=None): + def materialize(desc): return ShapeDtypeStruct( - adapter.to_public_axes(desc.shape, mode), + desc.array_shape, desc.dtype, ) @@ -534,19 +590,20 @@ def materialize(desc, *, mode=None): ), mock.patch.object( adapter._CALLER, - "_materialize_tensor_desc", + "_to_shape_dtype_struct", side_effect=materialize, ), ): empty, initialized = adapter.call_cutedsl( launch, (sample,), + input_descs=(input_desc,), outputs=( adapter.make_buffer_desc( "empty", (1, 0, 128), "bfloat16", - tensor_spec=adapter.gemm_c_tensor_spec("LMN"), + mode=output_mode, ), adapter.make_buffer_desc( "amax", diff --git a/test/python/fe_api/test_jax_grouped_wgrad_accumulation_contract.py b/test/python/fe_api/test_jax_grouped_wgrad_accumulation_contract.py index f8d7a6674..ed88fc600 100644 --- a/test/python/fe_api/test_jax_grouped_wgrad_accumulation_contract.py +++ b/test/python/fe_api/test_jax_grouped_wgrad_accumulation_contract.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for functional JAX grouped-wgrad accumulation.""" +"""Contracts for JAX-owned grouped-wgrad outputs.""" from __future__ import annotations @@ -16,6 +16,7 @@ _WGRAD_PATH = ( _ROOT / "python" / "cudnn" / "grouped_gemm" / "grouped_gemm_wgrad" / "jax.py" ) +_SHARED_PATH = _ROOT / "python" / "cudnn" / "grouped_gemm" / "_jax_api.py" def _definitions() -> dict[str, ast.AST]: @@ -41,20 +42,25 @@ def _argument_names(function: ast.FunctionDef) -> tuple[str, ...]: ) -def test_wgrad_seed_is_part_of_the_class_and_wrapper_signatures(): +def test_wgrad_output_is_not_part_of_the_class_or_wrapper_signatures(): definitions = _definitions() operation = definitions["GroupedGemmWgradSm100"] assert isinstance(operation, ast.ClassDef) wrapper = definitions["grouped_gemm_wgrad_wrapper_sm100"] assert isinstance(wrapper, ast.FunctionDef) + implementation = definitions["_grouped_gemm_wgrad_impl"] + assert isinstance(implementation, ast.FunctionDef) - assert "sample_wgrad_tensor" in _argument_names(_method(operation, "__init__")) - assert "wgrad_tensor" in _argument_names(_method(operation, "__call__")) - assert "wgrad_tensor" in _argument_names(_method(operation, "_call_impl")) - assert "wgrad_tensor" in _argument_names(wrapper) + assert "sample_wgrad_tensor" not in _argument_names( + _method(operation, "__init__") + ) + assert "wgrad_tensor" not in _argument_names(_method(operation, "__call__")) + assert "wgrad_tensor" not in _argument_names(_method(operation, "_call_impl")) + assert "wgrad_tensor" not in _argument_names(implementation) + assert "wgrad_tensor" not in _argument_names(wrapper) -def test_wgrad_seed_is_aliased_as_the_custom_call_result(): +def test_accumulating_wgrad_is_zero_initialized_by_its_descriptor(): implementation = _definitions()["_grouped_gemm_wgrad_impl"] assert isinstance(implementation, ast.FunctionDef) calls = [ @@ -67,18 +73,66 @@ def test_wgrad_seed_is_aliased_as_the_custom_call_result(): assert len(calls) == 1 keywords = {keyword.arg: keyword.value for keyword in calls[0].keywords} - assert ast.unparse(keywords["output_seeds"]) == "(wgrad_tensor,)" + assert "output_seeds" not in keywords + + output_descs = [ + node + for node in ast.walk(implementation) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "make_buffer_desc" + and node.args + and isinstance(node.args[0], ast.Constant) + and node.args[0].value == "wgrad_tensor" + ] + assert len(output_descs) == 1 + output_keywords = { + keyword.arg: keyword.value for keyword in output_descs[0].keywords + } + assert ast.unparse(output_keywords["init_value"]) == ( + "0.0 if accumulate_on_output else None" + ) source = ast.get_source_segment(_WGRAD_PATH.read_text(), implementation) assert source is not None - assert "wgrad_tensor is only valid when accumulate_on_output=True" in source - assert "wgrad_dtype must match wgrad_tensor.dtype" in source - assert "accumulate_on_output and wgrad_tensor is None" in source + assert "default=jnp.bfloat16" in source + + +def test_grouped_lowering_does_not_accept_preallocated_outputs(): + tree = ast.parse(_SHARED_PATH.read_text()) + call_cutedsl = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "call_cutedsl" + ) + + assert "output_seeds" not in _argument_names(call_cutedsl) + assert "output_seeds" not in ast.unparse(call_cutedsl) + + +def test_wgrad_layout_and_alignment_constraints_live_on_descriptors(): + source = _WGRAD_PATH.read_text() + implementation = _definitions()["_grouped_gemm_wgrad_impl"] + assert isinstance(implementation, ast.FunctionDef) + call = next( + node + for node in ast.walk(implementation) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "call_cutedsl" + ) + keywords = {keyword.arg for keyword in call.keywords} + + assert "WGRAD_B_STRIDE_ORDER = (0, 1)" in source + assert "WGRAD_ALIGNMENT = 16" in source + assert "ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT" in source + assert "input_descs" in keywords + assert not {"input_specs", "output_specs", "workspace_specs"} & keywords -def test_fresh_output_and_seeded_accumulation_are_documented_separately(): +def test_jax_owned_output_and_zero_initialization_are_documented(): source = _WGRAD_PATH.read_text() - assert "Omitting the seed preserves the" in source - assert "simple fresh-output behavior by starting from zero" in source + assert "The output is inferred and allocated by JAX" in source + assert "zero-initialized output" in source assert "Pointer-table outputs are not supported" in source diff --git a/test/python/fe_api/test_jax_layout.py b/test/python/fe_api/test_jax_layout.py index 1e605a8d8..068af5cf9 100644 --- a/test/python/fe_api/test_jax_layout.py +++ b/test/python/fe_api/test_jax_layout.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free tests for CUTLASS JAX layout conversion.""" +"""Tests for CUTLASS JAX layout conversion.""" import importlib.util from pathlib import Path diff --git a/test/python/fe_api/test_jax_nsa_topk_thd_contract.py b/test/python/fe_api/test_jax_nsa_topk_thd_contract.py index 8d7ed478e..db9e98ea4 100644 --- a/test/python/fe_api/test_jax_nsa_topk_thd_contract.py +++ b/test/python/fe_api/test_jax_nsa_topk_thd_contract.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for packed JAX NSA top-K reduction.""" +"""Contracts for packed JAX NSA top-K reduction.""" from __future__ import annotations @@ -97,11 +97,30 @@ def setUpClass(cls) -> None: internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) sys.modules[internal_name] = internal - tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") layout_module = importlib.import_module(f"{internal_name}.layout") - result_module = importlib.import_module(f"{_PACKAGE}._result") + result_module = importlib.import_module(f"{_PACKAGE}.common.result") class JaxTensorDesc(tensor_module.TensorDesc): + @classmethod + def from_shape( + cls, + shape, + dtype, + *, + name="", + mode=None, + public_stride_order=None, + init_value=None, + ): + return JaxApiBase._to_tensor_desc( + _Array(shape, dtype), + name, + mode=mode, + public_stride_order=public_stride_order, + init_value=init_value, + ) + @property def cudnn_dtype(self): return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) @@ -134,7 +153,7 @@ def _to_tensor_desc( canonical_axis_by_public_axis = layout_module.to_public_axes( tuple(range(len(public_shape))), mode ) - return JaxTensorDesc( + desc = JaxTensorDesc( dtype=value.dtype, shape=layout_module.to_canonical_axes(public_shape, mode), stride=layout_module.to_canonical_axes(public_stride, mode), @@ -145,9 +164,13 @@ def _to_tensor_desc( name=name, init_value=init_value, ) + object.__setattr__(desc, "mode", mode) + return desc @staticmethod def _check_tensor_signature(value, expected, *, mode=None): + if mode is None: + mode = expected.mode actual_shape = layout_module.to_canonical_axes(tuple(value.shape), mode) if actual_shape != expected.shape: raise ValueError(f"{expected.name} shape mismatch") @@ -159,6 +182,8 @@ def _check_tensor_signature(value, expected, *, mode=None): @staticmethod def _to_tensor_spec(desc, *, mode=None, divisibility=None): + if mode is None: + mode = desc.mode mode = layout_module.normalize_mode(desc.ndim, mode) return _TensorSpec( layout=layout_module.to_cutlass_layout( @@ -178,10 +203,22 @@ def _call_kernel( *, launch, output_descs, - input_spec, - output_spec, + input_descs=None, + input_spec=None, + output_spec=None, **options, ): + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc) + if input_spec is None: + input_spec = tuple( + self._to_tensor_spec(desc) for desc in input_descs or () + ) + if output_spec is None: + output_spec = tuple( + self._to_tensor_spec(desc) for desc in output_descs + ) self.captured_call = { "inputs": tuple(inputs), "launch": launch, @@ -191,8 +228,11 @@ def _call_kernel( **options, } return tuple( - _Array(desc.shape, _CUDNN_TO_DTYPE[desc.cudnn_dtype]) - for desc in output_descs + _Array( + layout_module.to_public_axes(desc.shape, spec.mode), + _CUDNN_TO_DTYPE[desc.cudnn_dtype], + ) + for desc, spec in zip(output_descs, output_spec) ) internal.JaxApiBase = JaxApiBase diff --git a/test/python/fe_api/test_jax_sdpa_thd_contract.py b/test/python/fe_api/test_jax_sdpa_thd_contract.py index 417299148..050edce56 100644 --- a/test/python/fe_api/test_jax_sdpa_thd_contract.py +++ b/test/python/fe_api/test_jax_sdpa_thd_contract.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free source contracts for packed-THD JAX SDPA.""" +"""Source contracts for packed-THD JAX SDPA.""" from __future__ import annotations @@ -135,7 +135,8 @@ def test_thd_is_promoted_without_changing_public_result_ranks(self): backward = _BACKWARD.read_text() for source in (forward, backward): - self.assertIn("jax.ShapeDtypeStruct((1, *sample_q.shape)", source) + self.assertIn("JaxTensorDesc.from_shape(", source) + self.assertIn("(1, *sample_q.shape)", source) self.assertIn("jnp.reshape(q_tensor, self.q_kernel_desc.shape)", source) self.assertIn("max_s_q and max_s_k are both required for THD", source) diff --git a/test/python/fe_api/test_result.py b/test/python/fe_api/test_result.py index 5a6918037..8a7376fcc 100644 --- a/test/python/fe_api/test_result.py +++ b/test/python/fe_api/test_result.py @@ -11,7 +11,10 @@ import pytest _CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" -_COMMON_SPEC = importlib.util.spec_from_file_location("_cudnn_result_contract", _CUDNN_ROOT / "_result.py") +_COMMON_SPEC = importlib.util.spec_from_file_location( + "_cudnn_result_contract", + _CUDNN_ROOT / "common" / "result.py", +) assert _COMMON_SPEC is not None and _COMMON_SPEC.loader is not None _COMMON_MODULE = importlib.util.module_from_spec(_COMMON_SPEC) _COMMON_SPEC.loader.exec_module(_COMMON_MODULE) @@ -37,11 +40,14 @@ def test_tuple_dict_is_a_jax_pytree(monkeypatch): package_name = "_cudnn_result_jax_contract" package = types.ModuleType(package_name) package.__path__ = [str(_CUDNN_ROOT)] + common_package = types.ModuleType(f"{package_name}.common") + common_package.__path__ = [str(_CUDNN_ROOT / "common")] jax_package = types.ModuleType(f"{package_name}._jax") jax_package.__path__ = [str(_CUDNN_ROOT / "_jax")] monkeypatch.setitem(sys.modules, package_name, package) + monkeypatch.setitem(sys.modules, f"{package_name}.common", common_package) monkeypatch.setitem(sys.modules, f"{package_name}._jax", jax_package) - monkeypatch.setitem(sys.modules, f"{package_name}._result", _COMMON_MODULE) + monkeypatch.setitem(sys.modules, f"{package_name}.common.result", _COMMON_MODULE) spec = importlib.util.spec_from_file_location( f"{package_name}._jax.result", diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_imports.py b/test/python/fe_api/test_rmsnorm_rht_amax_imports.py index 0f576f44e..a4d13faa2 100644 --- a/test/python/fe_api/test_rmsnorm_rht_amax_imports.py +++ b/test/python/fe_api/test_rmsnorm_rht_amax_imports.py @@ -189,7 +189,7 @@ def imports_framework(node, framework: str): return any(alias.name == framework or alias.name.startswith(f"{framework}.") for alias in node.names) return node.level == 0 and (node.module == framework or (node.module or "").startswith(f"{framework}.")) - base_op_imports = all_imports(_CUDNN_ROOT / "_op.py") + base_op_imports = all_imports(_CUDNN_ROOT / "common" / "op.py") self.assertFalse(any(imports_framework(node, framework) for node in base_op_imports for framework in ("torch", "jax", "cutlass", "cuda"))) operation_imports = all_imports("op.py") @@ -367,8 +367,8 @@ def imports_framework(node, framework: str): self.assertNotIn("_trace_lock", {node.attr for node in ast.walk(jax_adapter) if isinstance(node, ast.Attribute)}) root_source = (_CUDNN_ROOT / "__init__.py").read_text() - self.assertIn("from ._op import Op", root_source) - self.assertIn("from ._tensor_desc import TensorDesc", root_source) + self.assertIn("from .common.op import Op", root_source) + self.assertIn("from .common.tensor_desc import TensorDesc", root_source) self.assertIn('if name == "jax":', root_source) self.assertIn('importlib.import_module(".jax", __name__)', root_source) diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py b/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py index 3eca62eb1..bf8a36a4c 100644 --- a/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py +++ b/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for the RMSNorm CuTe kernel.""" +"""Contracts for the RMSNorm CuTe kernel.""" import ast import importlib diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_op.py b/test/python/fe_api/test_rmsnorm_rht_amax_op.py index 64815b5d5..5db57fa78 100644 --- a/test/python/fe_api/test_rmsnorm_rht_amax_op.py +++ b/test/python/fe_api/test_rmsnorm_rht_amax_op.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free contracts for the RMSNorm logical operation.""" +"""Contracts for the RMSNorm logical operation.""" from enum import Enum, auto import importlib @@ -48,8 +48,8 @@ def setUpClass(cls) -> None: sys.modules[operation_name] = operation try: - cls.tensor_module = importlib.import_module(f"{_PACKAGE}._tensor_desc") - cls.base_module = importlib.import_module(f"{_PACKAGE}._op") + cls.tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}.common.op") cls.op_module = importlib.import_module(f"{operation_name}.op") except Exception: cls.tearDownClass() diff --git a/test/python/fe_api/test_shared_tensor_desc.py b/test/python/fe_api/test_shared_tensor_desc.py index 218ad44b4..cd48a097d 100644 --- a/test/python/fe_api/test_shared_tensor_desc.py +++ b/test/python/fe_api/test_shared_tensor_desc.py @@ -1,7 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""Dependency-free tests for framework-neutral tensor descriptors.""" +"""Tests for framework-neutral tensor descriptors.""" from enum import Enum, auto import importlib @@ -35,7 +35,7 @@ def _load_tensor_desc_module(): package.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) package.data_type = _DataType sys.modules[_PACKAGE] = package - return importlib.import_module(f"{_PACKAGE}._tensor_desc") + return importlib.import_module(f"{_PACKAGE}.common.tensor_desc") class SharedTensorDescTest(unittest.TestCase):