diff --git a/docs/fe-oss-apis/attention/sdpa_bwd_d256.md b/docs/fe-oss-apis/attention/sdpa_bwd_d256.md index e183af066..07550e606 100644 --- a/docs/fe-oss-apis/attention/sdpa_bwd_d256.md +++ b/docs/fe-oss-apis/attention/sdpa_bwd_d256.md @@ -45,6 +45,13 @@ Here, the `sum(dP \odot P)` term is reduced row-wise over the key/softmax dimens ### High-level wrapper +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python import torch from cudnn import sdpa_bwd_wrapper_sm100_d256 @@ -73,6 +80,50 @@ dq, dk, dv = result # Key access: result["dq_tensor"], result["dk_tensor"], result["dv_tensor"] ``` +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +from cudnn.jax import sdpa_bwd_wrapper_sm100_d256 + +@jax.jit +def sdpa_bwd(q, k, v, o, do, lse): + return sdpa_bwd_wrapper_sm100_d256( + q, + k, + v, + o, + do, + lse, + mma_tiler_mn=(128, 128), + dkdv_mma_tiler_mn=(128, 64), + is_causal=False, + window_size=(-1, -1), + scale_softmax=None, + ) + +result = sdpa_bwd(q, k, v, o, do, lse) +dq, dk, dv = result["dq"]_tensor, result["dk_tensor"], result["dv_tensor"] +``` + +The JAX wrapper returns +`TupleDict(dq_tensor=..., dk_tensor=..., dv_tensor=...)`. It supports +fixed BHSD arrays only: Q, K, V, O, and dO have logical shape +`(B,H,S,256)`, all floating inputs except FP32 LSE share FP16 or BF16 dtype, +and `H_q % H_kv == 0`. Packed THD inputs, cumulative lengths, `max_s_*`, +explicit output buffers, and stream arguments remain PyTorch-only. Tensor +inputs are runtime operands; shapes, dtypes, tilers, mask/window selection, +and `scale_softmax` are static compilation state. XLA owns dQ/dK/dV and the +zero-initialized hidden workspace used by the two-kernel implementation. This +is an explicit backward operation and does not register a JAX autodiff rule. + +````` + +`````` + ### Class API ```python @@ -201,7 +252,7 @@ sdpa_bwd.execute( ### Wrapper return values -Returns a `TupleDict` with keys: +The PyTorch wrapper returns a `TupleDict` with keys: - `dq_tensor` - `dk_tensor` @@ -209,6 +260,9 @@ Returns a `TupleDict` with keys: Tuple unpacking order is: `(dq_tensor, dk_tensor, dv_tensor)`. +The JAX wrapper returns `TupleDict` with the same field names and tuple +order. + ### Class-specific parameters: `SdpabwdSm100D256` #### `SdpabwdSm100D256` (constructor) @@ -277,4 +331,3 @@ For runnable examples and reference-comparison checks, see: - `test/python/fe_api/test_sdpa_bwd.py` - `test/python/fe_api/test_sdpa_bwd_utils.py` - diff --git a/docs/fe-oss-apis/attention/sdpa_fwd_d256.md b/docs/fe-oss-apis/attention/sdpa_fwd_d256.md index 56a03fae3..ceafaa736 100644 --- a/docs/fe-oss-apis/attention/sdpa_fwd_d256.md +++ b/docs/fe-oss-apis/attention/sdpa_fwd_d256.md @@ -13,6 +13,13 @@ This is available through a standalone API (documented below) and is also used b ### High-level wrapper +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python import torch from cudnn import sdpa_fwd_wrapper_sm100_d256 @@ -39,6 +46,45 @@ o_tensor, lse_tensor = result # Key access: result["o_tensor"], result["lse_tensor"] ``` +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +from cudnn.jax import sdpa_fwd_wrapper_sm100_d256 + +@jax.jit +def sdpa_fwd(q, k, v): + return sdpa_fwd_wrapper_sm100_d256( + q, + k, + v, + mma_tiler_mn=(128, 128), + is_causal=False, + window_size=(-1, -1), + scale_softmax=None, + scale_output=1.0, + ) + +result = sdpa_fwd(q, k, v) +o_tensor, lse_tensor = result["o_tensor"], result["lse_tensor"] +``` + +The JAX wrapper returns `TupleDict(o_tensor=..., lse_tensor=...)`. It +supports fixed BHSD arrays only: Q, K, and V use logical shape +`(B,H,S,256)`, share FP16 or BF16 dtype, and satisfy `H_q % H_kv == 0`. +Packed THD inputs, cumulative lengths, `max_s_*`, explicit outputs, and stream +arguments remain PyTorch-only. Tensor data is runtime; shapes, dtypes, +accumulator types, tiling, mask/window configuration, and both scale values +are static compilation state. XLA supplies the runtime stream and owns both +results. + +````` + +`````` + ### Class API ```python @@ -124,13 +170,16 @@ sdpa_fwd.execute( ## Wrapper return values -Returns a `TupleDict` with keys: +The PyTorch wrapper returns a `TupleDict` with keys: - `o_tensor` - `lse_tensor` Tuple unpacking order is `(o_tensor, lse_tensor)`. +The JAX wrapper returns `TupleDict` with the same `o_tensor` and +`lse_tensor` field names and tuple order. + ## Support surface and constraints - `head_dim` must be exactly `256` diff --git a/docs/fe-oss-apis/cutedsl-jax-design.md b/docs/fe-oss-apis/cutedsl-jax-design.md new file mode 100644 index 000000000..c9af27a3b --- /dev/null +++ b/docs/fe-oss-apis/cutedsl-jax-design.md @@ -0,0 +1,280 @@ +# JAX support for cuDNN Frontend OSS APIs + +Status: 31 of 33 public FE-OSS entry points have source-level JAX bindings; +representative hardware tests currently cover 11 of those 31 entry points. + +## Decision + +JAX is an optional, explicit API surface under `cudnn.jax`. The existing +PyTorch operation API remains first class: unqualified `cudnn` and +operation-package symbols continue to mean PyTorch, regardless of which +frameworks are installed. Compatibility of lower-level public classes still +needs one correction before release, described below. + +Both bindings use the same CuTe DSL kernels and share framework-neutral tensor +metadata and validation where the contracts genuinely match. They do not share +an execution lifecycle. PyTorch retains its explicit compile/execute flow; JAX +returns ordinary, un-jitted callables that lower through +`cutlass.jax.cutlass_call` when traced by JAX. + +This is the right boundary because a JAX tracer describes an abstract value, +not a concrete device allocation. XLA must own output allocation, workspace +lifetime, stream ordering, and buffer aliasing. + +## User and package model + +JAX users install the optional dependency and opt into its namespace: + +```bash +pip install 'nvidia-cudnn-frontend[jax]' +``` + +```python +import jax +import cudnn +import cudnn.jax + +# Existing PyTorch API; its meaning never depends on whether JAX is installed. +torch_result = cudnn.rmsnorm_rht_amax_sm100(torch_x, torch_weight) + +@jax.jit +def f(x, weight): + return cudnn.jax.rmsnorm_rht_amax_sm100(x, weight) +``` + +Each operation co-locates its bindings: + +```text +cudnn// + api.py # PyTorch API + jax.py # JAX API + kernel.py # shared CuTe DSL implementation +``` + +Operation packages lazily expose `.jax`, while their unqualified exports come +from `api.py`. Importing `cudnn` does not import JAX, CUTLASS, or PyTorch +eagerly. Importing `cudnn.jax` is the JAX dependency boundary: it checks for +JAX and CUTLASS, verifies `cutlass.jax.is_available()`, and gives an install +hint when the integration is unavailable. Before release, an operation-local +`.jax` import must route through the same check; today it can bypass the +friendly compatibility diagnostic. + +Comparable operation names and concepts are preferred, but exact signatures +are not required. Defaults, supported layouts, result containers, and +framework-specific controls may differ. A new PyTorch wrapper is not +automatically blocked on a JAX implementation, but review should identify the +gap and document whether it is intentional. + +## Execution architecture + +| Concern | PyTorch | JAX | +| --- | --- | --- | +| Inputs | Concrete tensors with observable strides and devices | Arrays or tracers with abstract shape and dtype | +| Compilation | Explicit `cute.compile()` lifecycle | CuTe compilation during XLA lowering | +| Execution | Python invokes the compiled callable | XLA invokes a typed custom call | +| Outputs | Allocated by the wrapper | Declared by shape/dtype; allocated by XLA | +| Stream | Torch/CUTLASS stream | Supplied by XLA/PJRT | +| Workspace | May be cached on the API object | XLA-owned, per-invocation hidden result | +| Mutation | Preallocated outputs are written | Public API is functional; aliases are declared | + +`cutlass.jax.cutlass_call` is the integration seam. During lowering it builds +a CUTLASS function specification from abstract inputs, output metadata, tensor +layouts, aliases, compile options, and static arguments. It compiles the CuTe +launcher and embeds the resulting object in a typed XLA FFI custom call. At +runtime, the custom call receives XLA-owned buffers and XLA's CUDA stream; +Python is not on the execution path. + +NSA sliding-window attention is the one non-CuTe exception. Its fixed-shape +inference binding delegates to JAX's public `dot_product_attention` cuDNN +implementation, which lowers to JAX's registered cuDNN custom call. The +frontend does not invoke its PyTorch/pygraph execution path while tracing. + +`call_cutedsl` preserves CUTLASS's `allow_cuda_graph=True` default. This permits +the graph-capable FFI target but does not force capture; XLA decides whether to +use it. The frontend only needs to opt out if an operation introduces a known +non-capturable runtime effect outside the ordinary stream-ordered launch path. + +The frontend's `call_cutedsl` adapter gives every kernel one canonical ABI: + +```text +launcher(stream, *inputs, *outputs, *workspaces, **static_args) +``` + +Public outputs and hidden workspaces are described by `BufferSpec`. A +`fill_value` requests XLA-visible initialization and an input/output alias; +`None` means the kernel fully overwrites the buffer. Workspace shapes must be +derivable from abstract metadata. No tensor address, stream, output buffer, or +workspace is retained by a JAX API object. + +## Shared metadata and validation + +The base classes deliberately separate common contracts from framework policy: + +- `TensorDesc` contains logical shape, dtype, packing, and kernel-visible + layout metadata without importing a framework. +- `TorchTensorDesc` adds observed strides, device, view behavior, and packed + storage rules. +- `JaxTensorDesc` adds the native `cutlass.jax.TensorSpec` used to declare + compact physical layout and logical mode order to XLA. It never reads a + physical JAX buffer stride. +- `ApiBase` contains framework-neutral validation helpers. +- `ApiBaseTorch` owns support checks, compilation, execution, allocation, and + streams. The historical `APIBase` name remains its compatibility alias. +- `ApiBaseJax` owns sample-signature validation and is itself a stable, + un-jitted callable. `get_jax_callable()` returns that same object. + +JAX operation classes accept sample array-like values in their constructors, +immediately convert them to descriptors, and retain no arrays or tracers. Their +calls accept the real operands and verify that shape, dtype, and optional +presence match the validated sample signature. Applications retain control of +`jax.jit`, donation, sharding, and device placement. + +Rank, logical shape and dtype, packing, divisibility, tile domains, and output +inference should be shared when they describe the same kernel path. Physical +Torch strides and devices, JAX `TensorSpec` declarations, framework streams, +allocation, and lowering remain adapter-specific. Validators should return +immutable, kernel-specific plans rather than branch on a framework flag. + +For GEMM, a Torch tensor's observed strides determine M-major, N-major, or +K-major interpretation. JAX has no equivalent user-visible stride contract. +JAX GEMM APIs instead use compact row-major public axis-order strings with the +batch or expert mode `L` outermost: A uses `LMK` or `LKM`, B uses `LNK` or +`LKN`, and C/D use `LMN` or `LNM` where the selected kernel supports those +orders. `TensorSpec.mode` maps the public axes to the kernel's canonical +`(M,K,L)`, `(N,K,L)`, and `(M,N,L)` modes, while `TensorSpec.layout` declares +the physical XLA layout at the custom-call boundary. XLA may insert a layout +conversion to satisfy that constraint. + +| Operand | Public order and shape | `TensorSpec.mode` | Contiguous kernel mode | +| --- | --- | --- | --- | +| A | `LMK`: `(L, M, K)` | `(1, 2, 0)` | K | +| A | `LKM`: `(L, K, M)` | `(2, 1, 0)` | M | +| B | `LNK`: `(L, N, K)` | `(1, 2, 0)` | K | +| B | `LKN`: `(L, K, N)` | `(2, 1, 0)` | N | +| C/D | `LMN`: `(L, M, N)` | `(1, 2, 0)` | N | +| C/D | `LNM`: `(L, N, M)` | `(2, 1, 0)` | M | + +The current grouped-GEMM bindings fix A to `LMK` and matrix outputs to `LMN`. +Only quant, sReLU, dsReLU, and dGLU accept both `LNK` and `LKN` for B; the +SwiGLU, dSwiGLU, GLU, and GLU + Hadamard paths use `LNK`. Grouped wgrad keeps +its separate two-dimensional input contract. These layout strings apply to +matrix operands and results. Packed scale tensors, probabilities, reductions, +and other auxiliary buffers retain their operator-specific layouts. + +## Coverage and limits + +| Family | PyTorch wrappers | JAX wrappers | +| --- | ---: | ---: | +| RMSNorm + RHT + amax | 1 | 1 | +| Dense GEMM fusions | 4 | 4 | +| Grouped GEMM | 9 | 9 | +| Discrete grouped GEMM | 2 | 0 | +| SDPA | 2 | 2 | +| Native Sparse Attention | 4 | 4 | +| DeepSeek Sparse Attention | 11 | 11 | +| **Total** | **33** | **31** | + +The two current gaps need separate ABI work, not another generic adapter: + +1. Two discrete grouped-GEMM wrappers consume device-resident tables of raw + addresses. JAX must receive the expert buffers as real operands and build + pointer tables in XLA-owned device workspace. + +The sliding-window binding deliberately covers fixed-shape BHSD inference with +`right_bound=0`. Training statistics and packed THD/ragged inputs still need a +native frontend XLA FFI adapter; depending on JAX's private cuDNN APIs would be +an unstable OSS contract. + +Dense DSA indexer backward now uses a runtime `grad_loss` operand, an +XLA-owned score-gradient workspace, and a kernel-cleared FP32 accumulation +buffer. Its first JAX binding deliberately covers fixed-shape SM100 BSHD +inputs; the broader PyTorch SM90 and packed-THD paths remain unchanged. + +Only RMSNorm, the four dense GEMMs, five DSA entry points, and NSA +sliding-window inference currently have real GPU execution tests. These tests +include numerical reference coverage for dense indexer backward across multiple +query and key tiles, runtime loss scaling, and the registered cuDNN +sliding-window lowering. The other families have source and CPU contract +coverage but have not been qualified through real CUTLASS lowering and SM100 +execution. + +Current JAX bindings require concrete shapes during tracing. They do not define +autodiff, `vmap`, or automatic SPMD partitioning rules. SDPA backward is an +explicit operation, not an autodiff rule. Operator-specific `shard_map` or +custom partitioning can be added only after local shapes, communication, and +workspace semantics are defined. Heterogeneous-GPU processes also remain to be +qualified. + +## OSS readiness review + +The architecture is appropriate for an open-source API, but four items should +block release: + +1. **Publish and enforce a compatibility matrix.** The base package advertises + Python 3.9+, but the new base descriptors use Python 3.10-only + `dataclass(kw_only=True)`, and the JAX dependency requires Python 3.11+. + The normal test requirements also pin NumPy below 2 while JAX 0.9.1 requires + NumPy 2. Split the JAX environment, document Linux/CUDA 13 support, and test + the minimum and current qualified JAX versions with CUTLASS 4.5.0. + Installation tests must build a wheel, install only `[jax]` in a Torch-free + environment, import `cudnn.jax`, and lower a representative operation. +2. **Add required SM100 GPU CI.** CPU contract tests and fake CUTLASS modules + are useful but cannot validate generated code, workspace lifetime, layout + conversion, stream use, or numerical correctness. CI should cover every + family and representative initialized-output, hidden-workspace, + multi-launch, and repeated-call paths. The checked-in wrapper count should + be generated or tested rather than maintained only in prose. +3. **Make runtime metadata safe.** Grouped-GEMM offsets and sparse-attention + counts/indices are documented as trusted values, but kernels use them for + device indexing without complete bounds checks. Add device-side validation + or guards, or define an explicit unsafe precondition with precise behavior + and negative tests. Invalid public inputs must not silently permit + out-of-bounds GPU access. +4. **Resolve the `TensorDesc` compatibility break.** This branch repurposes the + previously Torch-specific public name as a keyword-only neutral descriptor + and moves its old constructor and tensor-like behavior to + `TorchTensorDesc`. Keep a deprecated Torch-compatible spelling or make the + break explicit through versioning and migration documentation. Preserving + only the historical `APIBase` alias is not sufficient. + +Release and legal owners should also confirm the distribution terms for the +optional proprietary CuTe DSL dependency; the frontend source is MIT, but that +does not make the complete runtime stack MIT. + +The following are important follow-ups, but need not block the first release if +their limits are explicit: + +- Extract the remaining dual-purpose `_validate_only` functions into pure + validators that return frozen plans. Revalidating in constructors and calls + is a drift risk across 19 operation modules. +- Ensure every `JaxTensorDesc` or frozen plan contains the exact `TensorSpec` + used for lowering. Some SDPA sample descriptors currently use defaults while + the call path supplies a nontrivial layout. +- Define whether public `JaxTensorDesc.tensor_spec` is a stable API or an + advanced escape hatch coupled to CUTLASS. Its compatibility policy should be + tested. +- Coordinate with CUTLASS on its unbounded compilation cache and whether the + target architecture participates in the key. Until then, document + heterogeneous-architecture processes as unsupported. +- Add HLO and performance tests for declared GEMM layouts so hidden XLA + transposes do not erase kernel gains. +- Publish a per-operation support matrix for shapes, dtypes, layouts, + transforms, and architecture restrictions. Name parity alone does not imply + domain parity. +- Make `TupleDict` immutable or keep its cached key order synchronized after + dictionary mutation; its current tuple iteration can become stale. + +## Recommendation + +Keep the explicit `cudnn.jax` namespace, co-located `api.py`/`jax.py` +bindings, framework-neutral descriptors, and separate runtime adapters. Do not +unify PyTorch and JAX compile/execute behavior or dispatch implicitly on array +type. With the packaging matrix and GPU CI gates above, this design is a clear, +maintainable way to expose the same FE-OSS kernels to both frameworks. + +## References + +- [cuDNN Frontend OSS APIs](https://docs.nvidia.com/deeplearning/cudnn/latest/fe-oss-apis/fe-oss-apis.html) +- [CUTLASS `cutlass.jax.cutlass_call`](https://github.com/NVIDIA/cutlass/blob/v4.5.0/python/CuTeDSL/cutlass/jax/primitive.py) +- [JAX FFI API](https://docs.jax.dev/en/latest/_autosummary/jax.ffi.ffi_call.html) +- [JAX package metadata](https://pypi.org/project/jax/) diff --git a/docs/fe-oss-apis/dsa.md b/docs/fe-oss-apis/dsa.md index 9ec5f3400..9e921dc87 100644 --- a/docs/fe-oss-apis/dsa.md +++ b/docs/fe-oss-apis/dsa.md @@ -8,7 +8,7 @@ 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` +delivered as Python classes / wrappers that follow the same `ApiBaseTorch` pattern as other cuDNN Frontend operations. **Scope:** this module ships CuTe-DSL kernels for DSA backward, indexer @@ -59,16 +59,42 @@ Training-score loss path: ## Installation +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```bash pip install nvidia-cudnn-frontend[cutedsl] ``` +````` + +`````{tab-item} JAX +:sync: jax + +```bash +pip install nvidia-cudnn-frontend[jax] +``` + +The JAX optional dependency set requires Python 3.11 or newer. + +````` + +`````` + --- ## API Usage ### DSA Namespace +PyTorch remains the default and exposes every DSA class and wrapper below. +The optional JAX namespace exposes functional wrappers for the supported +fixed-shape subsets; the PyTorch class APIs remain unchanged. + ```python from cudnn import DSA @@ -100,6 +126,24 @@ DSA.DenseIndexerBackward DSA.dense_indexer_backward_wrapper ``` +The explicit JAX facade exposes its implemented functional subset: + +```python +from cudnn.jax import DSA + +DSA.indexer_forward_wrapper +DSA.indexer_top_k_wrapper +DSA.local_to_global_wrapper +DSA.compactify_wrapper +DSA.sparse_indexer_score_recompute_wrapper +DSA.sparse_attn_score_recompute_wrapper +DSA.dense_indexer_score_recompute_wrapper +DSA.dense_attn_score_recompute_wrapper +DSA.indexer_backward_wrapper +DSA.dense_indexer_backward_wrapper +DSA.sparse_attention_backward_wrapper +``` + --- ## Components @@ -120,7 +164,18 @@ Backward pass for DeepSeek Sparse Attention. Expects the forward outputs - **Outputs** — tuple `(dq, dkv, d_sink)` - **Constraints** — SM90 or SM100; SM90 supports the FlashMLA DSA shape with `head_dim ∈ {512, 576}` +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python +import math + +from cudnn import DSA + result = DSA.sparse_attention_backward_wrapper( q, kv, out, dout, lse, attn_sink, topk_idxs, softmax_scale=1.0 / math.sqrt(D), @@ -129,6 +184,55 @@ result = DSA.sparse_attention_backward_wrapper( dq, dkv, d_sink = result["dq"], result["dkv"], result["d_sink"] ``` +````` + +`````{tab-item} JAX +:sync: jax + +```python +import math + +import jax +from cudnn.jax import DSA + +@jax.jit +def sparse_attention_bwd(q, kv, out, dout, lse, attn_sink, topk_idxs, topk_length): + return DSA.sparse_attention_backward_wrapper( + q, + kv, + out, + dout, + lse, + attn_sink, + topk_idxs, + softmax_scale=1.0 / math.sqrt(512), + topk_length=topk_length, + ) + +result = sparse_attention_bwd( + q, kv, out, dout, lse, attn_sink, topk_idxs, topk_length +) +dq, dkv, d_sink = result["dq"], result["dkv"], result["d_sink"] +``` + +The JAX wrapper returns +`TupleDict(dq=..., dkv=..., d_sink=...)`. It supports the +fixed SM100 flat-MQA domain only: BF16 `q`, `out`, and `dout` have shape +`(S_q, H, 512)`, BF16 `kv` has shape `(S_kv, 512)`, `lse` and `attn_sink` are +FP32, and `H` is divisible by 64. Indices are global INT32 values and the +optional runtime `topk_length` has shape `(S_q,)`. + +`softmax_scale`, `block_tile=64`, and whether `topk_length` is present are +static compilation choices. Tensor data, indices, and lengths remain runtime +operands. XLA owns the zero-initialized reduction workspaces and output +accumulators. This is an explicit backward operation, not a registered JAX +autodiff rule; SM90 and broader packed/variable-length layouts remain +PyTorch-only. + +````` + +`````` + ### 2. Indexer Forward (score-only) Computes dense indexer scores: @@ -153,13 +257,57 @@ compressed column 0. - **Output** — `scores`: `(B, S_q, S_k)` FP32 - **Constraints** — SM100+, `head_dim == 128`, `qhead_per_kv_head ∈ {32, 64}` +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python +from cudnn import DSA + result = DSA.indexer_forward_wrapper( q, k, w, ratio=4, q_causal_offsets=q_causal_offsets, ) scores = result["scores"] ``` +The existing wrapper returns a `TupleDict` and retains its optional stream, +THD inputs, and explicit sequence-bound controls. + +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +from cudnn.jax import DSA + +@jax.jit +def indexer_scores(q, k, w): + return DSA.indexer_forward_wrapper(q, k, w, ratio=4) + +scores = indexer_scores(q, k, w)["scores"] +``` + +The JAX wrapper returns `TupleDict(scores=...)`. Its public result +has shape `(B, S_q, S_k)`. Internally, the custom call uses an FP32 result whose +last dimension is padded to a multiple of four for TMA. That physical result is +initialized to `-inf`, aliased into the custom call, and sliced back to `S_k`; +the initialization preserves positions that the causal kernel deliberately +does not write. + +The JAX wrapper supports concrete, compact BSHD inputs only. THD inputs, +`cu_seqlens_*`, `q_causal_offsets`, and shape-polymorphic export are +unsupported. Tuning options and `sm_scale` are Python-static compilation +state, and XLA supplies the runtime stream. + +````` + +`````` + ### 3. Indexer Top-K Radix top-K kernel for selecting candidate KV indices from indexer scores, @@ -172,7 +320,16 @@ with variable per-row effective length. `return_val=False`) - **Constraints** — SM100+, `top_k ≤ 2048` +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python +from cudnn import DSA + result = DSA.indexer_top_k_wrapper( scores.reshape(-1, scores.shape[-1]), seq_lens, top_k=512, @@ -180,6 +337,59 @@ result = DSA.indexer_top_k_wrapper( indices, values = result["indices"], result["values"] ``` +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +from cudnn.jax import DSA + +@jax.jit +def select_scores(scores, seq_lens): + return DSA.indexer_top_k_wrapper( + scores.reshape(-1, scores.shape[-1]), + seq_lens, + top_k=512, + next_n=1, + return_val=True, + ) + +result = select_scores(scores, seq_lens) +indices, values = result["indices"], result["values"] +``` + +The JAX wrapper returns `TupleDict(indices=..., values=...)`, with both +arrays shaped `(n_rows, top_k)`. The JAX API requires `return_val=True`; the +PyTorch wrapper retains its `return_val=False` mode. The +radix kernel needs an INT32 temporary of shape +`(n_rows, buffer_count, num_cols)`, where `buffer_count` is two for FP32 input +and one for FP16/BF16. The adapter declares that temporary as an uninitialized +custom-call result and drops it from the public result, allowing XLA to own its +lifetime instead of caching mutable workspace in Python. The kernel consumes +the global scratch path when its bucketed column count exceeds the in-CTA +candidate capacity; the launcher retains the same buffer ABI for smaller +problems. + +`top_k`, `next_n`, `return_val`, and `num_copy_bits` are Python-static. Shapes +must be concrete, `n_rows` must equal `seq_lens.shape[0] * next_n`, and +`0 < top_k <= min(num_cols, 2048)`. The kernel has no scalar tail when it +selects a vectorized output path, so the wrapper rejects configurations where +`top_k` is not divisible by the selected output vector width. Runtime lengths +are trusted kernel inputs: +every `seq_lens[b]` must satisfy +`top_k + next_n - 1 <= seq_lens[b] <= num_cols`, which keeps every staggered +row at least `top_k` elements long and prevents out-of-range reads. JAX tracing +does not copy those values to the host for validation. The JAX wrapper rejects +workspace sizes above the INT32 indexing limit instead of implementing the +PyTorch path's row-chunked fallback, and it does not expose the persistent +global-counter path. + +````` + +`````` + ### 4. Sparse Indexer Score Recompute Computes softmax over top-K entries of the indexer score: @@ -191,6 +401,68 @@ Computes softmax over top-K entries of the indexer score: `batch_idx * S_k + local_idx`. - **Output** — `predict`: `(B, S_q, topk)` FP32. +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + +```python +from cudnn import DSA + +result = DSA.sparse_indexer_score_recompute_wrapper( + q_indexer, + k_indexer, + weights, + topk_indices, + topk_length=topk_length, +) +predict = result["predict"] +``` + +The PyTorch wrapper returns a `TupleDict`, supports SM90 and SM100+, and retains +optional output-buffer and stream controls. + +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +from cudnn.jax import DSA + +@jax.jit +def recompute_predict(q_indexer, k_indexer, weights, topk_indices, topk_length): + return DSA.sparse_indexer_score_recompute_wrapper( + q_indexer, + k_indexer, + weights, + topk_indices, + topk_length=topk_length, + ) + +predict = recompute_predict( + q_indexer, k_indexer, weights, topk_indices, topk_length +)["predict"] +``` + +The JAX wrapper returns +`TupleDict(predict=...)`. It supports the SM100 fixed +batched MQA layout: BF16 `q_indexer` has shape `(B, S_q, H_q, D)`, BF16 +`k_indexer` has shape `(B, S_k, D)`, BF16 `weights` has shape +`(B, S_q, H_q)`, and INT32 `topk_indices` has shape `(B, S_q, topk)`. The +optional INT32 `topk_length` has shape `(B, S_q)`. The FP32 result has the same +shape as `topk_indices` and is fully overwritten by the kernel, so it needs no +initialized alias or algorithmic scratch workspace. When `topk_length` is +omitted, the adapter supplies one hidden `(1, 1)` INT32 placeholder required by +the native launch ABI; the non-compact kernel does not read it. + +````` + +`````` + ### 5. Sparse Attn Score Recompute L1-normalised head-summed softmax over top-K entries: @@ -203,12 +475,89 @@ L1-normalised head-summed softmax over top-K entries: - **Output** — `target`: `(B, S_q, topk)` FP32. - Note: the wrapper handles the `-log2(e) * lse` preprocessing internally. +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + +```python +from cudnn import DSA + +result = DSA.sparse_attn_score_recompute_wrapper( + q_attn, + k_attn, + lse, + topk_indices, + softmax_scale, + topk_length=topk_length, +) +target = result["target"] +``` + +The PyTorch wrapper returns a `TupleDict`, supports SM90 and SM100+, and retains +optional output-buffer and stream controls. + +````` + +`````{tab-item} JAX +:sync: jax + +```python +import math + +import jax +from cudnn.jax import DSA + +softmax_scale = 1.0 / math.sqrt(D) + +@jax.jit +def recompute_target(q_attn, k_attn, lse, topk_indices, topk_length): + return DSA.sparse_attn_score_recompute_wrapper( + q_attn, + k_attn, + lse, + topk_indices, + softmax_scale, + topk_length=topk_length, + ) + +target = recompute_target( + q_attn, k_attn, lse, topk_indices, topk_length +)["target"] +``` + +The JAX wrapper returns `TupleDict(target=...)`. It uses +the same SM100 fixed batched MQA shapes as the indexer variant, with BF16 +`q_attn` and `k_attn`, FP32 `lse` shaped `(B, S_q, H_q)`, INT32 indices and +optional lengths, and an FP32 `(B, S_q, topk)` result. The kernel fully +overwrites the result and needs no algorithmic scratch workspace. As in the +indexer variant, omitting `topk_length` creates only the unused hidden ABI +placeholder. + +````` + +`````` + +For both JAX sparse score wrappers, shapes must be concrete and compact. THD +inputs and the SM90 kernels remain PyTorch-only. The sparse MQA kernel requires +`qhead_per_kv_head == H_q`, with `H_q` in `{32, 64, 128}`, and a positive head +dimension divisible by 64. The indexer variant requires `topk` to be divisible +by 128; the attention variant requires divisibility by its selected 64- or +128-column tile. `qhead_per_kv_head`, `topk_indices_global`, and +`softmax_scale` (for the attention variant) are Python-static compilation +state; whether `topk_length` is present selects a static kernel variant. Runtime +index and length values are trusted: valid entries must identify KV rows in the +selected local or global index space, and each `topk_length` must be between +zero and `topk`. Positions at or beyond that length are returned as zero. + ### 6. Dense Indexer / Dense Attn Score Recompute 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. +written as zero and excluded from `denom`. Pass the same `q_causal_offsets` +to all dense score tensors that feed the same loss path. ### 7. Indexer Backward @@ -227,7 +576,16 @@ indexer tower: (CuTe-DSL only).** If the CuTe-DSL path fails the wrapper raises `RuntimeError` rather than silently falling back. +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python +from cudnn import DSA + result = DSA.indexer_backward_wrapper( index_q, weights, index_k, attn_score, index_score, topk_indices, @@ -238,6 +596,75 @@ d_index_q, d_weights, d_index_k = ( ) ``` +The PyTorch score-gradient stage overwrites `attn_score` and `index_score` +in place, preserving the existing eager pipeline and class API. + +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +from cudnn.jax import DSA + +@jax.jit +def indexer_bwd( + index_q, + weights, + index_k, + attn_score, + index_score, + topk_indices, + grad_loss, +): + return DSA.indexer_backward_wrapper( + index_q, + weights, + index_k, + attn_score, + index_score, + topk_indices, + sm_scale=1.0, + loss_coeff=1.0, + grad_loss=grad_loss, + block_I=128, + topk_indices_global=False, + ) + +result = indexer_bwd( + index_q, + weights, + index_k, + attn_score, + index_score, + topk_indices, + grad_loss, +) +d_index_q = result["d_index_q"] +d_weights = result["d_weights"] +d_index_k = result["d_index_k"] +``` + +The JAX wrapper returns +`TupleDict(d_index_q=..., d_weights=..., d_index_k=...)` for the +fixed SM100 BSHD subset. It requires BF16 `index_q=(B,S_q,64,128)`, BF16 +`weights=(B,S_q,64)`, BF16 `index_k=(B,S_k,128)`, FP32 score tensors, and +INT32 indices with common shape `(B,S_q,topk)`. `topk` must be divisible by +`block_I`, and the JAX path currently requires `block_I=128`. + +`grad_loss` is a runtime FP32 scalar or one-element array. `sm_scale`, +`loss_coeff`, `block_I`, and `topk_indices_global` are static compilation +state. Unlike the PyTorch API, JAX leaves both score inputs unchanged: an +XLA-owned hidden buffer carries `grad_signal` between the score-gradient and +GEMM stages, and an XLA-owned zeroed FP32 accumulator is cast to the returned +BF16 `d_index_k`. This is a standalone backward wrapper and does not register +a custom VJP. + +````` + +`````` + ### 8. Dense Indexer Backward Full-KV counterpart to Indexer Backward. It consumes raw dense score tensors @@ -254,6 +681,13 @@ and denominators produced by Dense Indexer / Dense Attn Score Recompute. - **Outputs** — `d_index_q`, `d_weights`, `d_index_k` - **Constraints** — SM90 or SM100+, `H >= 64`, `ratio >= 1` +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python dense_index = DSA.dense_indexer_score_recompute_wrapper( index_q, index_k.unsqueeze(2), weights, @@ -273,6 +707,58 @@ result = DSA.dense_indexer_backward_wrapper( ) ``` +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +from cudnn.jax import DSA + +@jax.jit +def dense_indexer_bwd( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + grad_loss, +): + return DSA.dense_indexer_backward_wrapper( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + sm_scale=1.0, + loss_coeff=1.0, + grad_loss=grad_loss, + block_I=128, + ratio=1, + ) +``` + +The JAX binding covers fixed-shape SM100 BSHD inputs with `H=64`, `D=128`, +and `block_I=128`. Inputs use BF16 model tensors, FP32 scores and +denominators, and a runtime FP32 scalar or one-element `grad_loss` array. +`sm_scale`, `loss_coeff`, `block_I`, and `ratio` are static compilation state. +Packed THD inputs and `q_causal_offsets` remain PyTorch-only. + +Both score inputs remain immutable. The score-gradient and GEMM kernels run in +one custom call on XLA's stream, using an XLA-owned FP32 `grad_signal` +workspace. The GEMM launch clears its XLA-owned FP32 `d_index_k` accumulator +before the bulk reductions. The returned +`d_index_q`, `d_weights`, and `d_index_k` arrays are BF16. + +````` + +`````` + --- ## Limitations @@ -286,3 +772,9 @@ result = DSA.dense_indexer_backward_wrapper( `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 coverage includes Indexer Forward, Indexer Top-K, sparse and dense + score recompute, sparse and dense Indexer Backward, and Sparse Attention + Backward.** + The wrappers require concrete compact shapes and do not define autodiff, + `vmap`, or automatic sharding rules. Their broader SM90/THD PyTorch paths + remain unchanged. diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_amax.md b/docs/fe-oss-apis/gemm_fusions/gemm_amax.md index 12f14f11e..9b03761d4 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_amax.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_amax.md @@ -64,6 +64,14 @@ A (MxKxL), SFA B (NxKxL), SFB ## API Usage ### High-level wrapper + +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python result = gemm_amax_wrapper_sm100( a_tensor, @@ -82,7 +90,66 @@ c, amax = result # Key access: result["c_tensor"], result["amax_tensor"] ``` +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX integration with `pip install nvidia-cudnn-frontend[jax]`. + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import gemm_amax_wrapper_sm100 + +@jax.jit +def run(a, b, sfa, sfb): + # A is (L, M, K); B is (L, N, K). + return gemm_amax_wrapper_sm100( + a, + b, + sfa, + sfb, + c_layout="LMN", + c_dtype=jnp.float32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + sf_vec_size=32, + a_layout="LMK", + b_layout="LNK", + ) + +c, amax = run(a, b, sfa, sfb) +# c.shape == (L, M, N) +``` + +The JAX API supports FP8 A/B, E8M0 scale factors, `sf_vec_size=32`, and +float32/float16/bfloat16 C. Layout strings describe the public, batch-first +axis order: `a_layout` accepts `"LMK"` or `"LKM"`, `b_layout` accepts +`"LNK"` or `"LKN"`, and `c_layout` accepts `"LMN"` or `"LNM"`. All accepted +layouts keep `L` outermost; the final letter is the contiguous mode of the +compact row-major JAX array. The defaults are `"LMK"`, `"LNK"`, and `"LMN"`. +The wrapper maps these arrays to the kernel's canonical `(M, K, L)`, +`(N, K, L)`, and `(M, N, L)` tensor views. + +The specialized `SFA` and `SFB` shapes remain unchanged with `L` in their +final dimension. Packed FP4x2 is not exposed because JAX's scalar FP4 type has +a different storage ABI. JAX initializes the amax reduction on every +invocation. + +````` + +`````` + ### Class API + +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cuda.bindings import driver as cuda @@ -103,9 +170,48 @@ op.compile() op.execute(a, b, sfa, sfb, c, amax, current_stream=None) ``` +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import GemmAmaxSm100 + +op = GemmAmaxSm100( + sample_a=jax.ShapeDtypeStruct(a.shape, a.dtype), + sample_b=jax.ShapeDtypeStruct(b.shape, b.dtype), + sample_sfa=jax.ShapeDtypeStruct(sfa.shape, sfa.dtype), + sample_sfb=jax.ShapeDtypeStruct(sfb.shape, sfb.dtype), + c_layout="LMN", + c_dtype=jnp.float32, + acc_dtype=jnp.float32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + sf_vec_size=32, + a_layout="LMK", + b_layout="LNK", +) +assert op.check_support() +c, amax = jax.jit(op)(a, b, sfa, sfb) +``` + +The constructor retains only the input descriptors. Outputs are allocated by +the JAX call. + +````` + +`````` + --- -## Parameters +## PyTorch parameter reference + +This section documents the existing PyTorch tensor shapes, strides, and +major-mode arguments. The JAX public shapes and layout strings are described +in the JAX usage tabs above. ### Input/Output tensors - Input tensor **A**: `a_tensor` (wrapper) or `sample_a`/`a_tensor` (class) @@ -158,6 +264,8 @@ Returns a `TupleDict` with keys: Tuple unpacking order is: `(c_tensor, amax_tensor)`. +The JAX API returns the same `TupleDict` keys and unpacking order. + ### Class-specific parameters: `GemmAmaxSm100` #### `GemmAmaxSm100` (constructor) diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md b/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md index a18656744..ddeadf4d6 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md @@ -75,6 +75,13 @@ A (MxKxL), SFA B (NxKxL), SFB ### High-level wrapper +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import gemm_dsrelu_wrapper_sm100 @@ -100,8 +107,71 @@ result = gemm_dsrelu_wrapper_sm100( d, dprob, amax, sfd = result ``` +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX integration with `pip install nvidia-cudnn-frontend[jax]`. + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import gemm_dsrelu_wrapper_sm100 + +@jax.jit +def run(a, b, c, sfa, sfb, prob): + # A is (L, M, K); B is (L, N, K); C is (L, M, N). + # prob retains its specialized (M, 1, L) ABI. + return gemm_dsrelu_wrapper_sm100( + a, + b, + c, + sfa, + sfb, + prob, + alpha=1.0, + d_layout="LMN", + d_dtype=jnp.bfloat16, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + a_layout="LMK", + b_layout="LNK", + ) + +d, dprob, amax, sfd = run(a, b, c, sfa, sfb, prob) +# d.shape == (L, M, N); dprob.shape == (M, 1, L) +``` + +The JAX API supports FP8 A/B with E8M0 scale factors and +`sf_vec_size=32`. C and D may be float32, float16, or bfloat16. The API +creates and zero-initializes `dprob` for each invocation. Output quantization +is not exposed, so `amax` and `sfd` are `None`. `M` must be divisible by the +per-CTA M extent because the native probability load is not predicated for a +partial final CTA tile. That extent is 128 rows for both `TILE_M=128` and the +two-CTA `TILE_M=256` mode. + +Layout strings describe the public, batch-first axis order. `a_layout` accepts +`"LMK"` or `"LKM"`, `b_layout` accepts `"LNK"` or `"LKN"`, and `d_layout` +accepts `"LMN"` or `"LNM"` for both the C input and D output. All accepted +layouts keep `L` outermost; the defaults are `"LMK"`, `"LNK"`, and `"LMN"`. +The specialized scale-factor tensors retain the shapes above, while `prob` and +`dprob` remain `(M, 1, L)`. + +````` + +`````` + ### Class API +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import GemmDsreluSm100 @@ -142,9 +212,51 @@ op.execute( ) ``` +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import GemmDsreluSm100 + +op = GemmDsreluSm100( + sample_a=jax.ShapeDtypeStruct(a.shape, a.dtype), + sample_b=jax.ShapeDtypeStruct(b.shape, b.dtype), + sample_c=jax.ShapeDtypeStruct(c.shape, c.dtype), + sample_sfa=jax.ShapeDtypeStruct(sfa.shape, sfa.dtype), + sample_sfb=jax.ShapeDtypeStruct(sfb.shape, sfb.dtype), + sample_prob=jax.ShapeDtypeStruct(prob.shape, prob.dtype), + alpha=1.0, + d_layout="LMN", + d_dtype=jnp.bfloat16, + acc_dtype=jnp.float32, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + a_layout="LMK", + b_layout="LNK", +) +assert op.check_support() +d, dprob, amax, sfd = jax.jit(op)(a, b, c, sfa, sfb, prob) +``` + +The constructor retains only the input descriptors. Outputs are allocated by +the JAX call; `amax` and `sfd` are `None` in the supported output modes. + +````` + +`````` + --- -## Parameters +## PyTorch parameter reference + +This section documents the existing PyTorch tensor shapes and major-mode +arguments. The JAX public shapes and layout strings are described in the JAX +usage tabs above. ### Input/Output tensors @@ -219,6 +331,8 @@ Returns a `TupleDict` with keys: Tuple unpacking order is: `(d_tensor, dprob_tensor, amax_tensor, sfd_tensor)`. +The JAX API returns the same `TupleDict` keys and unpacking order. + --- ## Support surface and constraints diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md b/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md index 963165e3f..d0e3fc118 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md @@ -68,6 +68,13 @@ A (MxKxL), SFA B (NxKxL), SFB ### High-level wrapper +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import gemm_srelu_wrapper_sm100 @@ -93,8 +100,70 @@ result = gemm_srelu_wrapper_sm100( c, d, amax, sfd = result ``` +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX integration with `pip install nvidia-cudnn-frontend[jax]`. + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import gemm_srelu_wrapper_sm100 + +@jax.jit +def run(a, b, sfa, sfb, prob): + # A is (L, M, K); B is (L, N, K). + # prob retains its specialized (M, 1, L) ABI. + return gemm_srelu_wrapper_sm100( + a, + b, + sfa, + sfb, + prob, + alpha=1.0, + c_layout="LMN", + c_dtype=jnp.bfloat16, + d_dtype=jnp.bfloat16, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + a_layout="LMK", + b_layout="LNK", + ) + +c, d, amax, sfd = run(a, b, sfa, sfb, prob) +# c.shape == d.shape == (L, M, N) +``` + +The JAX API supports FP8 A/B with E8M0 scale factors and +`sf_vec_size=32`. C and D may be float32, float16, or bfloat16; output +quantization (`sfd`, `amax`, and `norm_const_tensor`) is not exposed in this +surface, so `amax` and `sfd` are `None`. `M` must be divisible by the per-CTA +M extent because the native probability load is not predicated for a partial +final CTA tile. That extent is 128 rows for both `TILE_M=128` and the two-CTA +`TILE_M=256` mode. + +Layout strings describe the public, batch-first axis order. `a_layout` accepts +`"LMK"` or `"LKM"`, `b_layout` accepts `"LNK"` or `"LKN"`, and `c_layout` +accepts `"LMN"` or `"LNM"` for both C and D. All accepted layouts keep `L` +outermost; the defaults are `"LMK"`, `"LNK"`, and `"LMN"`. The specialized +scale-factor tensors retain the shapes above, and `prob` remains `(M, 1, L)`. + +````` + +`````` + ### Class API +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import GemmSreluSm100 @@ -133,9 +202,51 @@ op.execute( ) ``` +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import GemmSreluSm100 + +op = GemmSreluSm100( + sample_a=jax.ShapeDtypeStruct(a.shape, a.dtype), + sample_b=jax.ShapeDtypeStruct(b.shape, b.dtype), + sample_sfa=jax.ShapeDtypeStruct(sfa.shape, sfa.dtype), + sample_sfb=jax.ShapeDtypeStruct(sfb.shape, sfb.dtype), + sample_prob=jax.ShapeDtypeStruct(prob.shape, prob.dtype), + alpha=1.0, + c_layout="LMN", + c_dtype=jnp.bfloat16, + d_dtype=jnp.bfloat16, + acc_dtype=jnp.float32, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + a_layout="LMK", + b_layout="LNK", +) +assert op.check_support() +c, d, amax, sfd = jax.jit(op)(a, b, sfa, sfb, prob) +``` + +The constructor retains only the input descriptors. Outputs are allocated by +the JAX call; `amax` and `sfd` are `None` in the supported output modes. + +````` + +`````` + --- -## Parameters +## PyTorch parameter reference + +This section documents the existing PyTorch tensor shapes and major-mode +arguments. The JAX public shapes and layout strings are described in the JAX +usage tabs above. ### Input/Output tensors @@ -204,6 +315,8 @@ Returns a `TupleDict` with keys: Tuple unpacking order is: `(c_tensor, d_tensor, amax_tensor, sfd_tensor)`. +The JAX API returns the same `TupleDict` keys and unpacking order. + --- ## Support surface and constraints diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md index e03c041e4..8ddbdd0b8 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md @@ -68,6 +68,13 @@ Notes: ### High-level wrapper (Standard Mode) +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python result = gemm_swiglu_wrapper_sm100( a_tensor, @@ -86,10 +93,64 @@ ab12, c, sfc, amax = result # Key access: result["ab12_tensor"], result["c_tensor"] ``` +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX integration with `pip install nvidia-cudnn-frontend[jax]`. + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import gemm_swiglu_wrapper_sm100 + +@jax.jit +def run(a, b): + # A is (L, M, K); B is (L, N, K). + return gemm_swiglu_wrapper_sm100( + a, + b, + alpha=1.0, + c_layout="LNM", + ab12_dtype=jnp.float32, + c_dtype=jnp.float16, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + a_layout="LMK", + b_layout="LNK", + ) + +ab12, c, sfc, amax = run(a, b) +# ab12.shape == (L, N, M); c.shape == (L, N // 2, M) +``` + +The JAX API supports the standard mode. Layout strings describe the public, +batch-first axis order: `a_layout` accepts `"LMK"` or `"LKM"`, `b_layout` +accepts `"LNK"` or `"LKN"`, and `c_layout` accepts `"LMN"` or `"LNM"` for +both AB12 and C. All accepted layouts keep `L` outermost; the defaults are +`"LMK"`, `"LNK"`, and `"LMN"`. The wrapper maps these arrays to the kernel's +canonical `(M, K, L)`, `(N, K, L)`, and `(M, N, L)` tensor views. + +`TILE_N` must be one of `{64, 128, 192, 256}` because the epilogue consumes +32-column subtiles in pairs. With `TILE_M=256`, `M` must be divisible by 256 so +both CTAs in each MMA pair are present. `sfc` and `amax` are `None`. + +````` + +`````` + ### High-level wrapper (Quantized Mode) When scale factor tensors are provided, the wrapper uses the block-scaled quantized kernel. +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python result = gemm_swiglu_wrapper_sm100( a_tensor, @@ -114,8 +175,66 @@ ab12, c, sfc, amax = result # Key access: result["ab12_tensor"], result["c_tensor"], result["sfc_tensor"], result["amax_tensor"] ``` +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import gemm_swiglu_wrapper_sm100 + +@jax.jit +def run_mxfp8(a, b, sfa, sfb): + return gemm_swiglu_wrapper_sm100( + a, + b, + sfa_tensor=sfa, + sfb_tensor=sfb, + c_layout="LMN", + ab12_dtype=jnp.bfloat16, + c_dtype=jnp.bfloat16, + acc_dtype=jnp.float32, + sf_vec_size=32, + vector_f32=False, + ab12_stages=4, + a_layout="LMK", + b_layout="LNK", + ) + +result = run_mxfp8(a, b, sfa, sfb) +ab12, c, sfc, amax = result +# sfc is None; amax is None for MXFP8 inputs. +``` + +The JAX block-scaled path supports MXFP8 inputs and native +`jnp.float4_e2m1fn` inputs. Pass SFA and SFB together with shapes +`(32, 4, ceil(M/128), 4, ceil(ceil(K/sf_vec_size)/4), L)` and +`(32, 4, ceil(N/128), 4, ceil(ceil(K/sf_vec_size)/4), L)`, respectively. +Native FP4 requires the default K-major input and N-major output layouts: +`a_layout="LMK"`, `b_layout="LNK"`, and `c_layout="LMN"`. With native FP4 +inputs and `c_dtype=jnp.bfloat16`, the wrapper allocates and returns an +initialized `(1, 1, 1)` FP32 AMAX tensor. `vector_f32` and `ab12_stages` are +static block-scaled kernel controls and may be configured before tracing. + +Raw `uint8` FP4 payloads are specific to the PyTorch API and are not accepted +by JAX. FP8 `C` output, SFC generation, and `norm_const_tensor` are not yet +supported in the JAX API. + +````` + +`````` + ### Class API (Standard Mode) +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python gemm = GemmSwigluSm100( sample_a, @@ -139,8 +258,49 @@ gemm.execute( ) ``` +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import GemmSwigluSm100 + +gemm = GemmSwigluSm100( + sample_a=jax.ShapeDtypeStruct(a.shape, a.dtype), + sample_b=jax.ShapeDtypeStruct(b.shape, b.dtype), + alpha=1.0, + c_layout="LNM", + ab12_dtype=jnp.float32, + c_dtype=jnp.float16, + acc_dtype=jnp.float32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + a_layout="LMK", + b_layout="LNK", +) +assert gemm.check_support() +ab12, c, sfc, amax = jax.jit(gemm)(a, b) +``` + +The constructor retains only the input descriptors. Outputs are allocated by +the JAX call; `sfc` and `amax` are `None` in standard mode. + +````` + +`````` + ### Class API (Quantized Mode) +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python gemm = GemmSwigluSm100( sample_a, @@ -178,9 +338,57 @@ gemm.execute( ) ``` +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import GemmSwigluSm100 + +# a and b have dtype jnp.float4_e2m1fn; sfa and sfb use a supported FP8 scale dtype. +gemm = GemmSwigluSm100( + sample_a=jax.ShapeDtypeStruct(a.shape, a.dtype), + sample_b=jax.ShapeDtypeStruct(b.shape, b.dtype), + sample_sfa=jax.ShapeDtypeStruct(sfa.shape, sfa.dtype), + sample_sfb=jax.ShapeDtypeStruct(sfb.shape, sfb.dtype), + alpha=1.0, + c_layout="LMN", + ab12_dtype=jnp.bfloat16, + c_dtype=jnp.bfloat16, + acc_dtype=jnp.float32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + sf_vec_size=16, + vector_f32=False, + ab12_stages=4, + a_layout="LMK", + b_layout="LNK", +) +assert gemm.check_support() +ab12, c, sfc, amax = jax.jit(gemm)(a, b, sfa, sfb) +# sfc is None; amax is an initialized FP32 result for FP4 inputs with BF16 C. +``` + +The JAX constructor binds input metadata and static kernel configuration. It +derives and functionally allocates AB12, C, and any applicable AMAX output; +there are no `sample_ab12`, `sample_c`, or `sample_amax` arguments. The +application owns `jax.jit` and passes the invocation-time arrays to the +callable object. + +````` + +`````` + --- -## Parameters +## PyTorch parameter reference + +This section documents the existing PyTorch tensor shapes, strides, and +major-mode arguments. The JAX public shapes and layout strings are described +in the JAX usage tabs above. ### Input/Output tensors @@ -241,7 +449,7 @@ gemm.execute( - `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)` @@ -286,6 +494,10 @@ Tuple unpacking order is always: - **Standard mode**: `sfc_tensor is None` and `amax_tensor is None` - **Quantized mode**: `sfc_tensor` and/or `amax_tensor` are populated based on dtype/configuration +The JAX API returns the same four `TupleDict` keys and unpacking order. SFC is +currently always `None`. AMAX is populated for native FP4 inputs with BF16 C; +it is `None` in standard mode and for MXFP8 inputs. + ### Class-specific parameters #### `GemmSwigluSm100` (constructor) @@ -319,7 +531,7 @@ Tuple unpacking order is always: #### Quantized mode -The quantized kernel supports the following configurations: +The PyTorch quantized API supports the following configurations: | Format | ab_dtype | sf_dtype | sf_vec_size | Notes | |--------|-----------|----------|-------------|-------| @@ -327,6 +539,19 @@ 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 | +The JAX quantized API uses native JAX dtypes: + +| Format | A/B dtype | SFA/SFB dtype | sf_vec_size | +|--------|-----------|---------------|-------------| +| **NVFP4** | `jnp.float4_e2m1fn` | `jnp.float8_e4m3fn` or `jnp.float8_e8m0fnu` | 16 | +| **MXFP4** | `jnp.float4_e2m1fn` | `jnp.float8_e8m0fnu` | 32 | +| **MXFP8** | `jnp.float8_e4m3fn` or `jnp.float8_e5m2` | `jnp.float8_e8m0fnu` | 32 | + +JAX does not reinterpret `uint8` as packed FP4. Its FP4 arrays contain native +logical FP4 elements. JAX currently supports FP16, BF16, or FP32 C output; +SFC, FP8 C output, and `norm_const_tensor` are deferred. MXFP8 requires FP16 +or BF16 AB12. Native FP4 with BF16 C returns AMAX. + Additional constraints: - `acc_dtype` must be `float32` - Not compatible with FP8 c_dtype. BF16 `c_dtype` is expected. diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md index 43cf67715..5cc853346 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md @@ -116,6 +116,13 @@ $$ **Dense mode:** +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import grouped_gemm_dglu_wrapper_sm100 from cuda.bindings import driver as cuda @@ -161,6 +168,62 @@ sfd_col = outputs["sfd_col_tensor"] d_row, d_col, dprob, dbias, amax, sfd_row, sfd_col = outputs ``` +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX integration with `pip install nvidia-cudnn-frontend[jax]`. + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import grouped_gemm_dglu_wrapper_sm100 + +@jax.jit +def run(a, b, c, sfa, sfb, padded_offsets, alpha, beta, prob, norm_const): + return grouped_gemm_dglu_wrapper_sm100( + a_tensor=a, + c_tensor=c, + sfa_tensor=sfa, + padded_offsets=padded_offsets, + alpha_tensor=alpha, + beta_tensor=beta, + prob_tensor=prob, + b_tensor=b, + sfb_tensor=sfb, + generate_dbias=True, + norm_const_tensor=norm_const, + d_dtype=jnp.float8_e4m3fn, + output_layout="LMN", + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + act_func="dswiglu", + b_layout="LNK", + ) + +result = run(a, b, c, sfa, sfb, padded_offsets, alpha, beta, prob, norm_const) +d_row, d_col, dprob, dbias, amax, sfd_row, sfd_col = result +``` + +The JAX layout strings describe public axis order, not physical strides. `A` +is fixed to `LMK` with shape `(1, valid_m, K)`. `b_layout="LNK"` uses shape +`(L, N, K)` and `b_layout="LKN"` uses `(L, K, N)`. The input `C` and matrix +outputs are fixed to `output_layout="LMN"` and use shape +`(1, valid_m, 2N)`. The expert mode `L` must remain outermost in these strings. + +Packed scale tensors, `prob`, `dprob`, and `dbias` retain their +specialized shapes documented below. The JAX API returns `TupleDict`; `dprob` +and optional `dbias` are fresh zero-initialized functional results. The current +surface supports dense FP8 A/B and FP8 D with E8M0 scales and +`sf_vec_size=32`; `amax` is `None`. Shapes, dtypes, layouts, and configuration +arguments are static under `jax.jit`. + +````` + +`````` + **Discrete mode:** ```python @@ -347,10 +410,13 @@ Returns a `TupleDict` (dictionary + tuple unpacking): - `d_row_tensor`: Row-quantized dGLU output - `d_col_tensor`: Column-quantized dGLU output - `dprob_tensor`: Gradient of prob +- `dbias_tensor`: Optional bias gradient - `amax_tensor`: Per-group amax (when `d_dtype` is bf16/fp16) - `sfd_row_tensor`: Row scale factors (when SFD enabled) - `sfd_col_tensor`: Column scale factors (when SFD enabled) +Tuple unpacking order is: `(d_row_tensor, d_col_tensor, dprob_tensor, dbias_tensor, amax_tensor, sfd_row_tensor, sfd_col_tensor)`. + --- ## Support Surface and Constraints diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.md index 128276284..a18b53207 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dsrelu.md @@ -83,6 +83,13 @@ A (valid_m×K×1), SFA B (N×K×L), SFB padded_offsets ### High-level wrapper +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import grouped_gemm_dsrelu_wrapper_sm100 @@ -111,6 +118,60 @@ result = grouped_gemm_dsrelu_wrapper_sm100( d_row, d_col, dprob, dbias, amax, sfd_row, sfd_col = result ``` +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX integration with `pip install nvidia-cudnn-frontend[jax]`. + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import grouped_gemm_dsrelu_wrapper_sm100 + +@jax.jit +def run(a, b, c, sfa, sfb, padded_offsets, alpha, prob, norm_const): + return grouped_gemm_dsrelu_wrapper_sm100( + a, + b, + c, + sfa, + sfb, + padded_offsets, + alpha, + prob, + generate_dbias=True, + norm_const_tensor=norm_const, + d_dtype=jnp.float8_e4m3fn, + output_layout="LMN", + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + b_layout="LNK", + ) + +result = run(a, b, c, sfa, sfb, padded_offsets, alpha, prob, norm_const) +d_row, d_col, d_srelu, dprob, dbias, amax, sfd_row, sfd_col, sfd_col_d_srelu = result +``` + +The JAX layout strings describe public axis order, not physical strides. `A` +is fixed to `LMK` with shape `(1, valid_m, K)`. `b_layout="LNK"` uses shape +`(L, N, K)` and `b_layout="LKN"` uses `(L, K, N)`. The input `C` and matrix +outputs are fixed to `output_layout="LMN"` with shape `(1, valid_m, N)`. The +expert mode `L` must remain outermost in these strings. + +Packed scale tensors, `prob`, `dprob`, and `dbias` retain their specialized +shapes documented below. The JAX API returns `TupleDict`; `dprob` and optional +`dbias` are fresh zero-initialized functional results. It supports dense FP8 +A/B and FP8 outputs with E8M0 scales and `sf_vec_size=32`; `amax` is `None`. +Shapes, dtypes, layouts, and configuration arguments are static under +`jax.jit`. + +````` + +`````` + ### Discrete-weight wrapper ```python @@ -274,13 +335,15 @@ Returns a `TupleDict` with keys: - `d_row_tensor` - `d_col_tensor` +- `d_srelu_tensor` - `dprob_tensor` - `dbias_tensor` - `amax_tensor` - `sfd_row_tensor` - `sfd_col_tensor` +- `sfd_col_d_srelu_tensor` -Tuple unpacking order is: `(d_row_tensor, d_col_tensor, dprob_tensor, dbias_tensor, amax_tensor, sfd_row_tensor, sfd_col_tensor)`. +Tuple unpacking order is: `(d_row_tensor, d_col_tensor, d_srelu_tensor, dprob_tensor, dbias_tensor, amax_tensor, sfd_row_tensor, sfd_col_tensor, sfd_col_d_srelu_tensor)`. --- diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dswiglu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dswiglu.md index c7f1de2e1..5bd7e233a 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dswiglu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dswiglu.md @@ -100,6 +100,13 @@ $$ ### High-level Wrapper +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import grouped_gemm_dswiglu_wrapper_sm100 from cuda.bindings import driver as cuda @@ -142,6 +149,58 @@ sfd_col = outputs["sfd_col_tensor"] # column scale factors (when d_dtype is FP8) d_row, d_col, dprob, amax, sfd_row, sfd_col = outputs ``` +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX integration with `pip install nvidia-cudnn-frontend[jax]`. + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import grouped_gemm_dswiglu_wrapper_sm100 + +@jax.jit +def run(a, b, c, sfa, sfb, padded_offsets, alpha, beta, prob, norm_const): + return grouped_gemm_dswiglu_wrapper_sm100( + a, + b, + c, + sfa, + sfb, + padded_offsets, + alpha, + beta, + prob, + norm_const, + d_dtype=jnp.float8_e4m3fn, + output_layout="LMN", + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + ) + +result = run(a, b, c, sfa, sfb, padded_offsets, alpha, beta, prob, norm_const) +d_row, d_col, dprob, amax, sfd_row, sfd_col = result +``` + +The JAX matrix inputs have fixed public axis orders: `A` is `LMK` with shape +`(1, valid_m, K)`, and `B` is `LNK` with shape `(L, N, K)`. The input `C` and +matrix outputs are fixed to `output_layout="LMN"`; for this backward fusion +they use shape `(1, valid_m, 2N)`. These strings describe public axis order +rather than physical strides, and the batch/expert mode must remain outermost. + +Packed scale tensors, `prob`, and `dprob` retain their specialized shapes +documented below. The JAX API returns `TupleDict` and creates `dprob` as a +fresh zero-initialized functional result. It supports dense FP8 A/B and FP8 D +with E8M0 scales and `sf_vec_size=32`; `amax` is `None`. Shapes, dtypes, +layouts, and configuration arguments are static under `jax.jit`. + +````` + +`````` + ### Class API ```python diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md index 2061c8be4..4da14243b 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md @@ -119,6 +119,13 @@ $$ **Dense mode:** +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import grouped_gemm_glu_wrapper_sm100 from cuda.bindings import driver as cuda @@ -163,6 +170,62 @@ sfd_col = outputs["sfd_col_tensor"] c, d, d_col, amax, sfd_row, sfd_col = outputs ``` +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX integration with `pip install nvidia-cudnn-frontend[jax]`. + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import grouped_gemm_glu_wrapper_sm100 + +@jax.jit +def run(a, b, sfa, sfb, padded_offsets, alpha, prob, norm_const, bias): + return grouped_gemm_glu_wrapper_sm100( + a_tensor=a, + sfa_tensor=sfa, + padded_offsets=padded_offsets, + alpha_tensor=alpha, + b_tensor=b, + sfb_tensor=sfb, + bias_tensor=bias, + norm_const_tensor=norm_const, + prob_tensor=prob, + c_dtype=jnp.bfloat16, + d_dtype=jnp.float8_e4m3fn, + output_layout="LMN", + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + act_func="swiglu", + b_layout="LNK", + ) + +result = run(a, b, sfa, sfb, padded_offsets, alpha, prob, norm_const, bias) +c, d, d_col, amax, sfd_row, sfd_col = result +``` + +The JAX matrix inputs have fixed public axis orders: `A` is `LMK` with shape +`(1, valid_m, K)`, and the supported `b_layout="LNK"` uses shape `(L, N, K)`. +The `LKN` B order is not supported by this kernel. Matrix outputs are fixed to +`output_layout="LMN"`: `C` has shape `(1, valid_m, N)`, while `D` and `D_col` +have shape `(1, valid_m, N/2)` for GLU activations. The expert mode must remain +outermost in layout strings. + +Packed scale tensors, `prob`, `bias`, and reduction outputs retain their +specialized shapes documented below. The JAX API returns `TupleDict` and +supports dense FP8 A/B with E8M0 scales and `sf_vec_size=32`; the FP8 D +configuration above returns `d_col`, `sfd_row`, and `sfd_col`, with +`amax=None`. Discrete weight pointers and packed FP4 are not exposed. Shapes, +dtypes, layouts, and configuration arguments are static under `jax.jit`. + +````` + +`````` + **Discrete mode:** ```python diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard.md index 1d41af336..152bafbbc 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu_hadamard.md @@ -101,6 +101,13 @@ A (valid_m×K×1), SFA B (N×K×L), SFB padded_offsets ### High-level wrapper +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import grouped_gemm_glu_hadamard_wrapper_sm100 @@ -131,6 +138,61 @@ c_tensor, d_tensor, amax_tensor = result The wrapper constructs the fixed Hadamard matrix internally. +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX integration with `pip install nvidia-cudnn-frontend[jax]`. + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import grouped_gemm_glu_hadamard_wrapper_sm100 + +@jax.jit +def run(a, b, sfa, sfb, padded_offsets, alpha, prob, bias): + return grouped_gemm_glu_hadamard_wrapper_sm100( + a, + b, + sfa, + sfb, + padded_offsets, + alpha, + prob, + bias_tensor=bias, + c_dtype=jnp.bfloat16, + d_dtype=jnp.bfloat16, + output_layout="LMN", + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=16, + act_func="swiglu", + b_layout="LNK", + ) + +result = run(a, b, sfa, sfb, padded_offsets, alpha, prob, bias) +c, d, amax, post_rht_amax = result +``` + +The JAX matrix inputs have fixed public axis orders: `A` is `LMK` with shape +`(1, valid_m, K)`, and the supported `b_layout="LNK"` uses shape `(L, N, K)`. +The `LKN` B order is not supported by this kernel. Matrix outputs are fixed to +`output_layout="LMN"`: `C` has shape `(1, valid_m, N)`, while `D` has shape +`(1, valid_m, N/2)` for SwiGLU and GEGLU. The expert mode must remain outermost +in layout strings. + +Packed scale tensors, `prob`, `bias`, and reduction outputs retain their +specialized shapes documented below. Unlike the other grouped JAX wrappers, +this kernel's supported input path is dense native FP4 +(`jnp.float4_e2m1fn`), with E8M0 or E4M3 scale factors. Raw packed `uint8` FP4 +and discrete weight pointers are not exposed. Shapes, dtypes, layouts, and +configuration arguments are static under `jax.jit`. + +````` + +`````` + ### Class API ```python @@ -250,8 +312,9 @@ Returns a `TupleDict` with keys: - `c_tensor` - `d_tensor` - `amax_tensor` +- `post_rht_amax_tensor` -Tuple unpacking order is: `(c_tensor, d_tensor, amax_tensor)`. +Tuple unpacking order is: `(c_tensor, d_tensor, amax_tensor, post_rht_amax_tensor)`. --- diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant.md index 3d987fae4..8bc88c9b4 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant.md @@ -93,6 +93,13 @@ $$ ### High-level Wrapper +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import grouped_gemm_quant_wrapper_sm100 from cuda.bindings import driver as cuda @@ -131,6 +138,60 @@ sfd_col = outputs["sfd_col_tensor"] # column scale factors (when SFD outputs are d, d_col, amax, sfd_row, sfd_col = outputs # d_col is None for bf16/fp16/fp32 outputs ``` +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX integration with `pip install nvidia-cudnn-frontend[jax]`. + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import grouped_gemm_quant_wrapper_sm100 + +@jax.jit +def run(a, b, sfa, sfb, padded_offsets, alpha, prob, norm_const): + return grouped_gemm_quant_wrapper_sm100( + a_tensor=a, + sfa_tensor=sfa, + padded_offsets=padded_offsets, + alpha_tensor=alpha, + b_tensor=b, + sfb_tensor=sfb, + prob_tensor=prob, + norm_const_tensor=norm_const, + d_dtype=jnp.float8_e4m3fn, + output_layout="LMN", + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + b_layout="LNK", + ) + +result = run(a, b, sfa, sfb, padded_offsets, alpha, prob, norm_const) +d, d_col, amax, sfd_row, sfd_col = result +``` + +The JAX API returns `TupleDict`. Its layout strings describe public axis order, +not physical strides. `A` is fixed to `LMK` with shape `(1, valid_m, K)`. +`b_layout="LNK"` uses shape `(L, N, K)` and `b_layout="LKN"` uses +`(L, K, N)`, where `L` is the expert count. Matrix outputs are fixed to +`output_layout="LMN"` and have shape `(1, valid_m, N)`. The batch/expert mode +must remain outermost in these strings. + +The packed scale tensors, `prob`, and reduction outputs retain their +operator-specific shapes documented below; the matrix layout strings do not +reorder them. The JAX path supports dense FP8 A/B with E8M0 scales and +`sf_vec_size=32`; FP8 D produces `d_col`, `sfd_row`, and `sfd_col`, while +`amax` is `None`. Discrete weight pointers and packed FP4 are not exposed. +Shapes, dtypes, layouts, and configuration arguments are static under +`jax.jit`. + +````` + +`````` + ### Class API ```python diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant_unified.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant_unified.md index daec4e61d..ba54404af 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant_unified.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant_unified.md @@ -97,6 +97,13 @@ $$ ### High-level Wrapper +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import grouped_gemm_quant_wrapper_sm100 from cuda.bindings import driver as cuda @@ -135,6 +142,58 @@ sfd_col = outputs["sfd_col_tensor"] d, d_col, amax, sfd_row, sfd_col = outputs # d_col is None for bf16/fp16/fp32 outputs ``` +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX integration with `pip install nvidia-cudnn-frontend[jax]`. + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import grouped_gemm_quant_wrapper_sm100 + +@jax.jit +def run(a, b, sfa, sfb, padded_offsets, alpha, prob, norm_const): + return grouped_gemm_quant_wrapper_sm100( + a_tensor=a, + sfa_tensor=sfa, + padded_offsets=padded_offsets, + alpha_tensor=alpha, + b_tensor=b, + sfb_tensor=sfb, + prob_tensor=prob, + norm_const_tensor=norm_const, + d_dtype=jnp.float8_e4m3fn, + output_layout="LMN", + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + b_layout="LNK", + ) + +result = run(a, b, sfa, sfb, padded_offsets, alpha, prob, norm_const) +d, d_col, amax, sfd_row, sfd_col = result +``` + +The JAX namespace currently exposes the dense-weight subset of this unified +API. `A` has fixed public order `LMK` and shape `(1, valid_m, K)`. +`b_layout="LNK"` uses shape `(L, N, K)`, while `b_layout="LKN"` uses +`(L, K, N)`. Matrix outputs have fixed `output_layout="LMN"` and shape +`(1, valid_m, N)`. The expert mode must remain outermost in these compact +row-major public layouts. + +The packed scale tensors, `prob`, and reductions keep their operator-specific +layouts. The wrapper returns `TupleDict` and supports FP8 A/B with E8M0 scales +and `sf_vec_size=32`. FP8 D produces `d_col`, `sfd_row`, and `sfd_col`, while +`amax` is `None`. Shapes, dtypes, layouts, and configuration arguments are +static under `jax.jit`. + +````` + +`````` + ### Class API ```python @@ -179,7 +238,11 @@ api.execute( --- -## Parameters +## PyTorch parameter reference + +This section documents the existing PyTorch tensor shapes and major-mode +arguments. The JAX public shapes and layout strings are described in the JAX +usage tab above. ### Input/Output Tensors diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_srelu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_srelu.md index d70601ac4..802de0c59 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_srelu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_srelu.md @@ -76,6 +76,13 @@ A (valid_m×K×1), SFA B (N×K×L), SFB padded_offsets ### High-level wrapper +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import grouped_gemm_srelu_wrapper_sm100 @@ -104,6 +111,60 @@ result = grouped_gemm_srelu_wrapper_sm100( c, d, d_col, amax, sfd_row, sfd_col = result ``` +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX integration with `pip install nvidia-cudnn-frontend[jax]`. + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import grouped_gemm_srelu_wrapper_sm100 + +@jax.jit +def run(a, b, sfa, sfb, padded_offsets, alpha, prob, norm_const): + return grouped_gemm_srelu_wrapper_sm100( + a, + b, + sfa, + sfb, + padded_offsets, + alpha, + norm_const_tensor=norm_const, + prob_tensor=prob, + c_dtype=jnp.bfloat16, + d_dtype=jnp.float8_e4m3fn, + output_layout="LMN", + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + b_layout="LNK", + ) + +result = run(a, b, sfa, sfb, padded_offsets, alpha, prob, norm_const) +c, d, d_col, amax, sfd_row, sfd_col = result +``` + +The JAX API returns `TupleDict`. Its layout strings describe public axis order, +not physical strides. `A` is fixed to `LMK` with shape `(1, valid_m, K)`. +`b_layout="LNK"` uses shape `(L, N, K)` and `b_layout="LKN"` uses +`(L, K, N)`. The `C`, `D`, and `D_col` matrix results are fixed to +`output_layout="LMN"` and use shape `(1, valid_m, N)`. The expert mode `L` +must remain outermost in these strings. + +Packed scale tensors, `prob`, and reduction outputs retain their specialized +shapes documented below. The JAX path supports dense FP8 A/B with E8M0 scales +and `sf_vec_size=32`; the FP8 D configuration above returns both quantized +layouts and their scale factors, with `amax=None`. Discrete weight pointers and +packed FP4 are not exposed. Shapes, dtypes, layouts, and configuration +arguments are static under `jax.jit`. + +````` + +`````` + ### Discrete-weight wrapper ```python diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_swiglu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_swiglu.md index 6bb8a1a34..e8e335ed1 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_swiglu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_swiglu.md @@ -98,6 +98,13 @@ $$ ### High-level Wrapper +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import grouped_gemm_swiglu_wrapper_sm100 from cuda.bindings import driver as cuda @@ -138,6 +145,60 @@ sfd_col = outputs["sfd_col_tensor"] # column scale factors (when d_dtype is FP8) c, d, d_col, amax, sfd_row, sfd_col = outputs ``` +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX integration with `pip install nvidia-cudnn-frontend[jax]`. + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import grouped_gemm_swiglu_wrapper_sm100 + +@jax.jit +def run(a, b, sfa, sfb, padded_offsets, alpha, norm_const, prob): + return grouped_gemm_swiglu_wrapper_sm100( + a, + b, + sfa, + sfb, + padded_offsets, + alpha, + norm_const, + prob, + acc_dtype=jnp.float32, + c_dtype=jnp.bfloat16, + d_dtype=jnp.float8_e4m3fn, + output_layout="LMN", + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + ) + +result = run(a, b, sfa, sfb, padded_offsets, alpha, norm_const, prob) +c, d, d_col, amax, sfd_row, sfd_col = result +``` + +The JAX matrix inputs have fixed public axis orders: `A` is `LMK` with shape +`(1, valid_m, K)`, and `B` is `LNK` with shape `(L, N, K)`. Matrix outputs are +fixed to `output_layout="LMN"`: `C` has shape `(1, valid_m, N)`, while `D` and +`D_col` have shape `(1, valid_m, N/2)`. These strings describe public axis +order rather than physical strides, and the batch/expert mode must remain +outermost. + +Packed scale tensors, `prob`, and reduction outputs retain their specialized +shapes documented below. The JAX API returns `TupleDict` and supports dense FP8 +A/B with E8M0 scale factors and `sf_vec_size=32`; discrete weight pointers and +packed FP4 are not exposed. With the FP8 D type above, `amax` is `None`. +Shapes, dtypes, layouts, and configuration arguments are static under +`jax.jit`. + +````` + +`````` + ### Class API ```python diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md index a71de62ef..9674b6f32 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md @@ -43,6 +43,13 @@ The API supports two output modes through one public surface: ## Wrapper Example +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python import cudnn import torch @@ -61,6 +68,47 @@ result = cudnn.grouped_gemm_wgrad_wrapper_sm100( wgrad_tensor = result["wgrad_tensor"] ``` +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX integration with `pip install nvidia-cudnn-frontend[jax]`. + +```python +import jax +import jax.numpy as jnp +from cudnn.jax import grouped_gemm_wgrad_wrapper_sm100 + +@jax.jit +def run(a, b, sfa, sfb, offsets): + return grouped_gemm_wgrad_wrapper_sm100( + a, + b, + sfa, + sfb, + offsets, + wgrad_dtype=jnp.bfloat16, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + input_order="tensor2d", + ) + +result = run(a_tensor, b_tensor, sfa_tensor, sfb_tensor, offsets_tensor) +wgrad_tensor = result["wgrad_tensor"] +``` + +The JAX API returns `TupleDict` and exposes only the dense-output +mode. A/B are FP8 with E8M0 scale factors and `sf_vec_size=32`; raw discrete +output pointers and packed FP4 are not exposed. Tensor shapes and dtypes, +expert count, `input_order`, and other configuration arguments are static +under `jax.jit`; the offset values remain runtime data. + +````` + +`````` + ## Class API Example ```python diff --git a/docs/fe-oss-apis/nsa.md b/docs/fe-oss-apis/nsa.md index 10a2458a6..a0cfb894e 100644 --- a/docs/fe-oss-apis/nsa.md +++ b/docs/fe-oss-apis/nsa.md @@ -43,19 +43,39 @@ Each component can be used independently or combined for the full NSA pipeline. ## Installation -Install the cuDNN Frontend package with the CuteDSL optional dependencies: +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: ```bash pip install nvidia-cudnn-frontend[cutedsl] ``` +````` + +`````{tab-item} JAX +:sync: jax + +```bash +pip install nvidia-cudnn-frontend[jax] +``` + +The JAX optional dependency set requires Python 3.11 or newer. + +````` + +`````` + --- ## API Usage ### NSA Namespace -All NSA components are accessible through the `NSA` namespace: +PyTorch exposes all NSA components through the `NSA` namespace: ```python from cudnn import NSA @@ -74,6 +94,20 @@ NSA.TopKReduction NSA.topk_reduction_wrapper ``` +The explicit JAX namespace exposes the functional JAX subset: + +```python +from cudnn.jax import NSA + +NSA.selection_attention_wrapper +NSA.compression_attention_wrapper +NSA.topk_reduction_wrapper +NSA.sliding_window_attention_wrapper +``` + +Selection, compression, and top-K use CuTe DSL. Sliding-window inference uses +JAX's registered cuDNN attention lowering. + --- ## Components @@ -108,6 +142,13 @@ $$ #### High-level Wrapper +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import NSA @@ -131,6 +172,48 @@ o, l, m = result # Key access: result["o_tensor"], result["l_tensor"], result["m_tensor"] ``` +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +from cudnn.jax import NSA + +@jax.jit +def selection(q, k, v, block_indices, block_counts, cum_q, cum_k): + return NSA.selection_attention_wrapper( + q, + k, + v, + block_indices, + block_counts, + cum_q, + cum_k, + max_s_q=1024, + max_s_k=1024, + block_size=64, + ) + +result = selection(q, k, v, block_indices, block_counts, cum_seqlen_q, cum_seqlen_k) +o, l, m = result["o_tensor"], result["l_tensor"], result["m_tensor"] +``` + +The JAX wrapper returns +`TupleDict(o_tensor=..., l_tensor=..., m_tensor=...)`. It +supports packed, compact THD self-attention on SM90: FP16/BF16 Q, K, and V, +runtime cumulative lengths and block metadata, equal Q/KV token counts, and +identical Q/K runtime cumulative-length values. `max_s_q` and `max_s_k` are +required equal static integers. `block_size`, `scale_softmax`, output and +accumulator dtypes, and the GQA/head geometry are also static at trace time. +The current GQA ratio is one of `{1, 2, 4, 8}`; QK and V dimensions are +positive multiples of 16. XLA owns and initializes all three outputs. + +````` + +`````` + #### Class API ```python @@ -221,6 +304,13 @@ where $\alpha_q$, $\alpha_k$, $\alpha_v$, $\alpha_o$ are optional scaling factor #### High-level Wrapper +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import NSA @@ -247,6 +337,42 @@ o, lse = result # Key access: result["o_tensor"], result["lse_tensor"] ``` +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +from cudnn.jax import NSA + +@jax.jit +def compression(q, k, v): + return NSA.compression_attention_wrapper( + q, + k, + v, + enable_lse=True, + mma_tiler_mn=(128, 128), + is_persistent=False, + ) + +result = compression(q, k, v) +o, lse = result["o_tensor"], result["lse_tensor"] +``` + +The JAX wrapper returns +`TupleDict(o_tensor=..., lse_tensor=...)`, where +`lse_tensor` is `None` when `enable_lse=False`. It supports fixed BHSD inputs +on SM100 with matching FP16/BF16 Q, K, and V, `D` in `{32, 64, 128}`, and +`S_q` an integer multiple of `S_k`. THD inputs and cumulative-length operands +remain PyTorch-only. Tensor data is runtime; `enable_lse`, dtypes, tiling, +persistence, all scale values, and the fixed shapes are compilation state. + +````` + +`````` + #### Class API ```python @@ -332,13 +458,20 @@ This implementation is a wrapper around cudnn backend (and is not strictly open For each query position $q$, attention is restricted to key positions within the window: $$ -O[q] = \sum_{k : q - L \leq k \leq q + R} \text{softmax}\left(\frac{Q[q] \cdot K[k]^T}{\sqrt{D}}\right) V[k] +O[q] = \sum_{k : q - L < k \leq q + R} \text{softmax}\left(\frac{Q[q] \cdot K[k]^T}{\sqrt{D}}\right) V[k] $$ where $L$ is `left_bound` and $R$ is `right_bound`. #### High-level Wrapper +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import NSA @@ -362,6 +495,43 @@ o, stats = result # Key access: result["o_tensor"], result["stats_tensor"] ``` +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +from cudnn.jax import NSA + +@jax.jit +def sliding_window_attention(q, k, v): + return NSA.sliding_window_attention_wrapper( + q, + k, + v, + left_bound=512, + right_bound=0, + is_infer=True, + attn_scale=None, + ) +``` + +The JAX binding covers fixed-shape FP16/BF16 `(B, H, S, D)` self-attention +inference with `S_q == S_k`, `left_bound >= 1`, and `right_bound=0`. `D` must +be 32, 64, or 128, `D_v == D`, and `H_q` must be divisible by `H_kv`. It +lowers through JAX's registered cuDNN attention custom call, so XLA owns the +stream, output, and workspace. `o_dtype` may be FP16 or BF16 and defaults to +the input dtype. The JAX default `left_bound=1` selects only the current token; +pass the intended window explicitly when porting a PyTorch call. Packed THD +inputs, ragged offsets, and FP32 training statistics remain PyTorch-only. The +result uses the same `o_tensor` and `stats_tensor` keys; `stats_tensor` is +`None` for this inference-only subset. + +````` + +`````` + #### Class API ```python @@ -454,6 +624,13 @@ $$ #### High-level Wrapper +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + ```python from cudnn import NSA @@ -478,6 +655,46 @@ topk_scores, topk_indices = result # Key access: result["topk_scores_tensor"], result["topk_indices_tensor"] ``` +````` + +`````{tab-item} JAX +:sync: jax + +```python +import jax +from cudnn.jax import NSA + +@jax.jit +def reduce_topk(q, k, lse): + return NSA.topk_reduction_wrapper( + q, + k, + lse, + k_value=16, + selection_block_size=64, + compress_stride=32, + is_causal=True, + ) + +result = reduce_topk(q, k, lse) +topk_scores = result["topk_scores_tensor"] +topk_indices = result["topk_indices_tensor"] +``` + +The JAX wrapper returns +`TupleDict(topk_scores_tensor=..., topk_indices_tensor=...)` with +shape `(B,H_kv,S_q,k_value)`, FP32 scores, and INT32 indices. It supports +fixed FP16/BF16 BHSD inputs on SM100; THD inputs, cumulative lengths, and +`max_s_*` controls remain PyTorch-only. Q/K/LSE are runtime operands. Shapes, +`k_value`, block/stride geometry, causality, tiling, accumulator dtype, and +`scale_softmax` are static. `k_value` must be a positive multiple of four no +larger than the per-tile candidate count. XLA initializes invalid scores to +`-inf` and indices to `-1` before the kernel writes valid selections. + +````` + +`````` + #### Class API ```python @@ -551,6 +768,10 @@ topk.execute( | Sliding Window Attention | ✅ | ✅ | | Top-K Reduction | ✅ | ✅ | +The table above describes the PyTorch surface. In JAX, Compression Attention +and Top-K Reduction support fixed BHSD only, Selection Attention supports +packed THD self-attention only, and Sliding Window Attention is unavailable. + ### T,H,D Format (Variable-Length Batched) Used for sequences of varying lengths packed into a single tensor: @@ -595,6 +816,10 @@ All components require `float32` accumulator dtype for numerical stability. ## Usage Examples +The end-to-end example below uses the PyTorch surface because it includes +Sliding Window Attention. The three CuTe-DSL stages can instead use the JAX +wrappers shown in their component tabs. + For complete usage examples and tests, see: - `test/python/fe_api/nsa/test_NSA_selection_attention.py` diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 799f406c4..8c9946026 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -2,6 +2,9 @@ **FE-OSS APIs are experimental and subject to change.** +Design work for tracing FE-OSS CuTe DSL kernels from JAX is documented in +[CuTe DSL + JAX support for FE-OSS APIs](cutedsl-jax-design.md). + This folder documents the Python FE APIs implemented under `python/cudnn`. For details on currently implemented operations, see: - [GEMM + Amax](gemm_fusions/gemm_amax.md) - [GEMM + SwiGLU](gemm_fusions/gemm_swiglu.md) @@ -35,9 +38,69 @@ pip install nvidia-cudnn-frontend[cutedsl] After installation, you can import the APIs directly from the `cudnn` package, i.e. `from cudnn import {your_operation}` +The JAX integration uses a separate optional dependency set: + +```bash +pip install nvidia-cudnn-frontend[jax] +``` + +This optional dependency set requires Python 3.11 or newer, matching the JAX +CUDA package requirement; the base frontend package retains its broader Python +support. + +Importing `cudnn` does not load JAX or CuTe DSL. Importing `cudnn.jax`, or +accessing `cudnn.jax` after importing `cudnn`, is the explicit JAX opt-in. It +validates JAX and CuTe DSL availability and points missing installations to the +`jax` extra, including checking `cutlass.jax.is_available()` and reporting +CUTLASS's minimum supported JAX version when unavailable. It then loads the JAX +operation wrappers and shared CuTe DSL bridge. Architecture-specific kernel +modules remain deferred until an operation is traced. Each implemented +operation keeps `api.py` as its Torch binding and a sibling `jax.py` as its JAX +binding. + ## API Usage -Each operation exposes two APIs: +PyTorch remains the default FE-OSS interface and preserves its existing +wrappers and classes. Some operations also provide a functional JAX API under +`cudnn.jax`. Each operation-backed JAX wrapper has a callable class with the +aligned Torch class name; the two DSA layout helpers remain functional-only. +Supported bindings also retain aligned functional names across the two +namespaces: + +```python +from cudnn import rmsnorm_rht_amax_sm100 +from cudnn.jax import rmsnorm_rht_amax_sm100 + +from cudnn import gemm_swiglu_wrapper_sm100 +from cudnn.jax import gemm_swiglu_wrapper_sm100 +``` + +JAX class constructors accept array-like samples, immediately reduce them to +shape/dtype descriptors, and do not retain the sample arrays. Actual arrays are +passed when the object is called. The object is intentionally not pre-jitted, +so applications retain control over JIT, sharding, donation, and placement: + +```python +import jax +from cudnn.jax import RmsNormRhtAmaxSm100 + +op = RmsNormRhtAmaxSm100( + jax.ShapeDtypeStruct(x.shape, x.dtype), + jax.ShapeDtypeStruct(weight.shape, weight.dtype), +) +op.check_support() +output, amax = jax.jit(op)(x, weight) +``` + +Compile-time configuration becomes immutable on the first call because JAX +caches by callable identity. Construct a new operation object to change static +options after tracing. + +Where both bindings exist, they should offer comparable operation semantics and +recognizable inputs, options, and results. Exact names, signatures, defaults, +layouts, result containers, lifecycle controls, and supported domains may +differ by framework. JAX does not replace or narrow the PyTorch APIs, and +backend selection is never inferred from array types. ### 1. High-level wrapper diff --git a/docs/fe-oss-apis/rmsnorm_rht_amax.md b/docs/fe-oss-apis/rmsnorm_rht_amax.md index ba875fef8..49b6f58be 100644 --- a/docs/fe-oss-apis/rmsnorm_rht_amax.md +++ b/docs/fe-oss-apis/rmsnorm_rht_amax.md @@ -8,7 +8,10 @@ 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`) +- a PyTorch functional API (`cudnn.rmsnorm_rht_amax_sm100`) +- the existing wrapper API (`rmsnorm_rht_amax_wrapper_sm100`) +- a JAX class API (`cudnn.jax.RmsNormRhtAmaxSm100`) +- a JAX functional API (`cudnn.jax.rmsnorm_rht_amax_sm100`) - grouped-gemm-style regression coverage for compile/execute, wrapper use, and cache reuse ## Shapes @@ -53,7 +56,120 @@ over every element produced by that CTA. ## API Usage -### High-level wrapper +### Comparable framework APIs + +Both bindings implement the same forward computation and expose comparable +inputs, options, and result roles. This operation currently uses the same +function name in both namespaces; return types, lifecycle, supported layouts, +and stream handling remain framework-specific. + +``````{tab-set} +:sync-group: frontend-framework + +`````{tab-item} PyTorch +:sync: torch +:selected: + +Install the CuTe DSL dependencies: + +```bash +pip install nvidia-cudnn-frontend[cutedsl] +``` + +```python +from cudnn import rmsnorm_rht_amax_sm100 + +result = rmsnorm_rht_amax_sm100( + x, + weight, + eps=1e-5, + num_threads=None, + rows_per_cta=None, + current_stream=None, +) + +output, amax = result +``` + +The PyTorch functional API returns `TupleDict(output=..., amax=...)`. + +````` + +`````{tab-item} JAX +:sync: jax + +Install the JAX-specific optional dependencies: + +```bash +pip install nvidia-cudnn-frontend[jax] +``` + +The JAX optional dependency set requires Python 3.11 or newer. + +```python +import jax +from cudnn.jax import rmsnorm_rht_amax_sm100 + +eps = 1e-5 +num_threads = 128 +rows_per_cta = 2 + +@jax.jit +def run(x, weight): + return rmsnorm_rht_amax_sm100( + x, + weight, + eps=eps, + num_threads=num_threads, + rows_per_cta=rows_per_cta, + ) + +output, amax = run(x, weight) +``` + +The JAX function returns `TupleDict(output=..., amax=...)`, registered as a +JAX pytree with the same key and unpacking behavior as the PyTorch result. + +The JAX class API mirrors the sample-signature style without retaining sample +arrays: + +```python +from cudnn.jax import RmsNormRhtAmaxSm100 + +op = RmsNormRhtAmaxSm100( + jax.ShapeDtypeStruct(x.shape, x.dtype), + jax.ShapeDtypeStruct(weight.shape, weight.dtype), + eps=eps, + num_threads=num_threads, + rows_per_cta=rows_per_cta, +) +assert op.check_support() +run = jax.jit(op) +output, amax = run(x, weight) +``` + +JAX `check_support()` validates sample shape, dtype, and static launch +configuration. The final SM100 device check belongs to lowering/runtime because +abstract samples such as `ShapeDtypeStruct` do not have a device placement. + +````` + +`````` + +The JAX API requires concrete `M` and `N` during tracing and does +not yet define autodiff, `vmap`, or automatic partitioning rules. `eps`, +`num_threads`, and `rows_per_cta` are static compilation state; close them over +as above or list them in `jax.jit(static_argnames=...)`. JAX always uses XLA's +runtime stream; Torch optionally accepts `current_stream`. See +[CuTe DSL + JAX support for FE-OSS APIs](cutedsl-jax-design.md) for the design, +workspace model, rollout plan, and current limitations. + +### Additional PyTorch APIs + +The functional API above is additive. The existing PyTorch wrapper and class +APIs remain first class and retain their original contracts. + +#### Existing high-level wrapper ```python from cudnn import rmsnorm_rht_amax_wrapper_sm100 @@ -72,7 +188,7 @@ o_tensor, amax_tensor = result When no overrides are supplied, the wrapper uses the upstream-tuned thread table when available and an upstream-style `rows_per_cta` heuristic. -### Class API +#### Class API ```python from cudnn import RmsNormRhtAmaxSm100 @@ -101,21 +217,21 @@ op.execute( ### Input and output tensors -- `x_tensor` / `sample_x` +- `x` / `x_tensor` / `sample_x` - Shape: `(M, N)` - Layout: row-major contiguous - - Dtype: `torch.bfloat16` -- `w_tensor` / `sample_w` + - Dtype: bfloat16 (`torch.bfloat16` or `jax.numpy.bfloat16`) +- `weight` / `w_tensor` / `sample_w` - Shape: `(N,)` - Layout: contiguous - - Dtype: `torch.bfloat16` -- `o_tensor` / `sample_o` + - Dtype: bfloat16 (`torch.bfloat16` or `jax.numpy.bfloat16`) +- `output` / `o_tensor` / `sample_o` - Shape: `(M, N)` - Layout: row-major contiguous - - Dtype: `torch.bfloat16` -- `amax_tensor` / `sample_amax` + - Dtype: bfloat16 (`torch.bfloat16` or `jax.numpy.bfloat16`) +- `amax` / `amax_tensor` / `sample_amax` - Shape: `(M / rows_per_cta,)` - - Dtype: `torch.float32` + - Dtype: float32 (`torch.float32` or `jax.numpy.float32`) ### Common parameters @@ -125,15 +241,21 @@ op.execute( - 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`) +- Torch-only CUDA stream (`current_stream`); JAX always uses XLA's stream + +### Return values + +The PyTorch functional API returns a `TupleDict` with keys: +- `output` +- `amax` -### Wrapper return values +The JAX functional API returns a `TupleDict` with keys `output` and `amax`. -Returns a `TupleDict` with keys: +The legacy wrapper returns a `TupleDict` with keys: - `o_tensor` - `amax_tensor` -Tuple unpacking order is `(o_tensor, amax_tensor)`. +Tuple unpacking follows each key order. ## Support surface and constraints @@ -143,6 +265,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`. - Inputs and output are currently bf16 only. +- The JAX API requires concrete shapes and compact row-major inputs. - The frontend integration matches the upstream RMSNorm kernel semantics; it does not expose full LayerNorm mean/bias behavior. ## Verification 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/README.md b/python/cudnn/README.md index 61b5b3c94..95f024c6a 100644 --- a/python/cudnn/README.md +++ b/python/cudnn/README.md @@ -14,27 +14,72 @@ A simplified view of package structure: pyproject.toml # Project metadata and dependencies. Optional dependencies for frontend-only APIs are registered here. python/cudnn/ ├── __init__.py # Top-level exports (Graph, graph, jit, wrappers, kernels) +├── _operation_api.py # Shared lazy operation API exports +├── jax/ # Explicit JAX facade and shared CuTe/JAX adapter ├── graph.py # Low-level graph helpers (graph, jit, graph_cache) ├── wrapper.py # High-level Graph wrapper class ├── datatypes.py # Data type conversions and helpers ├── api_base.py # Abstract API base class for frontend-only APIs ├── {frontend-only-api-name}/ -│ ├── __init__.py # Frontend-only API class -│ └── api.py # High-level API implementation +│ ├── __init__.py # Lazy api.py exports and explicit .jax namespace +│ ├── api.py # PyTorch API implementation +│ ├── jax.py # Optional JAX functional API +│ ├── config.py # Framework-neutral operator configuration │ └── {kernel_name}.py # Kernel implementation, i.e CuteDSL test/python/ # Test files └── fe_api/ # Test files for frontend-only APIs ``` -## - ## Adding new frontend-only APIs -To add a new frontend-only API, follow these steps: -1. Create a new directory in the `python/cudnn` directory with the name of the API. -2. Add your kernel implementation and implement the high level API implementation in `api.py`, extending the `APIBase` class in `api_base.py`. -3. Expose the API import in `python/cudnn/__init__.py` and register the folder in `pyproject.toml`. Register any optional dependences if required. -4. Add a sample usage/test file in `test/python/fe_api/`. +The review unit is a user-visible operation, which may use one or more +main/helper CuTe kernels. Existing unqualified names retain their PyTorch +meaning; JAX is an explicit optional functional namespace. When both bindings +exist, they should implement comparable functionality on their documented +overlapping domain and use familiar terminology where practical: + +```python +from cudnn import my_operation_wrapper +torch_result = my_operation_wrapper(torch_inputs, ...) + +from cudnn.jax import my_operation +jax_result = my_operation(jax_inputs, ...) +``` + +To add a new frontend-only API: + +1. Document logical inputs and outputs, supported shapes/dtypes/layouts, + workspace, aliasing, initialization, and transformation behavior. +2. Add the CuTe implementation and preserve the existing PyTorch class/wrapper + conventions and compatibility behavior. +3. If JAX support is in scope, add a functional adapter using + `cutlass.jax.cutlass_call`. It must infer outputs/workspace from abstract + metadata, accept XLA's stream, and avoid Torch imports or host reads during + tracing. +4. Prefer recognizable operation, operand, option, and result names across + frameworks, but document intentional differences. Exact Python names, + signatures, defaults, layouts, result containers, and supported domains are + not required to match. +5. Test each framework's lifecycle and compare numerical behavior on the domain + they share. JAX coverage should include `eval_shape`, `jit`, lowering, and + execution on supported hardware. +6. During review, consider whether a new or modified PyTorch operation should + update JAX. Static lint or LLM review may report likely gaps or drift, but is + advisory rather than a public API contract or merge gate. +7. Keep `api.py` as the PyTorch implementation and place an optional JAX + implementation in `jax.py` beside it. Unqualified operation symbols always + resolve from `api.py`; use `.jax` or `cudnn.jax` explicitly for + JAX. Dependency availability must not change a symbol's framework. Register + JAX dependencies only in the optional extra. + +Do not use tensor-type dispatch or a traced `target=` argument. JAX does not +emulate the mutable `APIBase.compile()` / `execute()` lifecycle, and adding or +updating a PyTorch operation does not require a JAX binding. + +The JAX import-boundary test discovers every `jax.py` automatically and follows +its local kernel dependencies. New JAX operations do not require a per-kernel +import allowlist; keep shared modules free of required module-scope Torch +dependencies. **Currently implemented frontend-only APIs**: - `GEMM + Amax` @@ -57,6 +102,7 @@ To add a new frontend-only API, follow these steps: - `Block Sparse Attention (BSA)` - `SDPA Forward (SM100, D=256)` - `SDPA Backward (SM100, D=256)` +- `DeepSeek Sparse Attention (DSA) stages` **In progress frontend-only APIs**: - GEMM + Dswiglu diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index af48c763f..c2b46ecf7 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -277,6 +277,7 @@ def _dlopen_cudnn(): "GemmAmaxSm100": (".gemm_amax", "GemmAmaxSm100"), "gemm_amax_wrapper_sm100": (".gemm_amax", "gemm_amax_wrapper_sm100"), "RmsNormRhtAmaxSm100": (".rmsnorm_rht_amax", "RmsNormRhtAmaxSm100"), + "rmsnorm_rht_amax_sm100": (".rmsnorm_rht_amax", "rmsnorm_rht_amax_sm100"), "rmsnorm_rht_amax_wrapper_sm100": (".rmsnorm_rht_amax", "rmsnorm_rht_amax_wrapper_sm100"), "grouped_gemm": (".grouped_gemm", None), "GroupedGemmSwigluSm100": (".grouped_gemm", "GroupedGemmSwigluSm100"), @@ -340,6 +341,11 @@ def __getattr__(name: str) -> Any: globals()["experimental"] = _experimental return _experimental + if name == "jax": + _jax = importlib.import_module(".jax", __name__) + globals()["jax"] = _jax + return _jax + if name in _LAZY_OPTIONAL_IMPORTS: return _load_optional_symbol(name) diff --git a/python/cudnn/_experimental_warnings.py b/python/cudnn/_experimental_warnings.py index 762e3c608..e18f0b7b3 100644 --- a/python/cudnn/_experimental_warnings.py +++ b/python/cudnn/_experimental_warnings.py @@ -1,18 +1,21 @@ import logging +import threading _experimental_api_warnings_emitted = set() +_experimental_api_warnings_lock = threading.Lock() def warn_experimental_api_once(logger: logging.Logger, api_name: str) -> None: """Emit the experimental API warning once per API class per process.""" - if api_name in _experimental_api_warnings_emitted: - return - - _experimental_api_warnings_emitted.add(api_name) + with _experimental_api_warnings_lock: + if api_name in _experimental_api_warnings_emitted: + return + _experimental_api_warnings_emitted.add(api_name) logger.warning("%s is an experimental API", api_name) def _reset_experimental_api_warning_registry() -> None: """Reset experimental API warning state for tests.""" - _experimental_api_warnings_emitted.clear() + with _experimental_api_warnings_lock: + _experimental_api_warnings_emitted.clear() diff --git a/python/cudnn/_jax/__init__.py b/python/cudnn/_jax/__init__.py new file mode 100644 index 000000000..fbeca165a --- /dev/null +++ b/python/cudnn/_jax/__init__.py @@ -0,0 +1 @@ +"""Internal implementation helpers for the optional JAX API.""" diff --git a/python/cudnn/_jax/api_base.py b/python/cudnn/_jax/api_base.py new file mode 100644 index 000000000..2c5b3973c --- /dev/null +++ b/python/cudnn/_jax/api_base.py @@ -0,0 +1,662 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API base, tensor metadata, and CuTe DSL call adapter.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Any, Optional, Tuple, final + +import cutlass +import cutlass.cute as cute +import cutlass.jax as cutlass_jax +import jax +import jax.numpy as jnp +from cutlass.jax import TensorSpec + +from ..api_base import ApiBase, TensorDesc, TupleDict + +_NO_DEFAULT = object() + + +def as_dtype(value: Any) -> Any: + """Return a JAX dtype without retaining a dtype-bearing value.""" + + # Scalar dtype classes such as numpy.float32 expose an instance-level + # dtype descriptor, so only unwrap dtype-bearing values, not classes. + if not isinstance(value, type) and hasattr(value, "dtype"): + value = value.dtype + return jnp.dtype(value) + + +def as_optional_dtype(value: Any | None) -> Any | None: + """Return ``None`` or a JAX dtype without retaining the source value.""" + + return None if value is None else as_dtype(value) + + +def require_dtype( + value: Any, + valid_dtypes: Iterable[Any], + *, + name: str | None = None, + default: Any = _NO_DEFAULT, +) -> Any: + """Return a supported dtype from a dtype-like value or object with ``dtype``. + + The diagnostic name defaults to ``".dtype"`` for named + dtype-bearing values, then ``"dtype"``. + """ + + if not name: + value_name = getattr(value, "name", None) if not isinstance(value, type) and hasattr(value, "dtype") else None + name = f"{value_name}.dtype" if value_name else "dtype" + + if value is None: + if default is _NO_DEFAULT: + raise ValueError(f"{name} must not be None") + value = default + + dtype = as_dtype(value) + valid_dtypes = tuple(as_dtype(item) for item in valid_dtypes) + if dtype not in valid_dtypes: + supported = ", ".join(item.name for item in valid_dtypes) + raise ValueError(f"{name} must be one of {{{supported}}}, got {dtype}") + return dtype + + +def require_array( + value: Any, + *, + name: str | None = None, + rank: int | Iterable[int] | None = None, + shape: Sequence[Any] | None = None, + dtype: Any | Iterable[Any] | None = None, +) -> tuple[Any, ...]: + """Validate array metadata and return its shape. + + The diagnostic name defaults to a non-empty ``value.name``, then + ``"value"``. ``dtype`` accepts either one dtype or an iterable of supported + dtypes. Validation applies to the value's exposed shape. For a + :class:`JaxTensorDesc`, that is the kernel-visible shape; use + :attr:`JaxTensorDesc.array_shape` when validating a public JAX array against + a descriptor. + """ + + name = name or getattr(value, "name", None) or "value" + if not hasattr(value, "shape") or not hasattr(value, "dtype"): + raise TypeError(f"{name} must have shape and dtype metadata") + + actual_shape = tuple(value.shape) + if rank is not None: + valid_ranks = (rank,) if isinstance(rank, int) else tuple(rank) + if len(actual_shape) not in valid_ranks: + if len(valid_ranks) == 1: + raise ValueError(f"{name} must have rank {valid_ranks[0]}, got shape {actual_shape}") + expected = ", ".join(str(item) for item in valid_ranks) + raise ValueError(f"{name} must have one of ranks {{{expected}}}, got shape {actual_shape}") + + if shape is not None: + expected_shape = tuple(shape) + if actual_shape != expected_shape: + raise ValueError(f"{name} must have shape {expected_shape}, got {actual_shape}") + + if dtype is not None: + is_dtype_collection = isinstance(dtype, Iterable) and not isinstance(dtype, (str, bytes, type)) and not hasattr(dtype, "dtype") + valid_dtypes = tuple(dtype) if is_dtype_collection else (dtype,) + require_dtype(value, valid_dtypes, name=f"{name}.dtype") + + return actual_shape + + +def _flatten_tuple_dict(value: TupleDict) -> tuple[tuple[Any, ...], tuple[Any, ...]]: + """Flatten the current mapping values in insertion order.""" + + keys = tuple(dict.keys(value)) + children = tuple(dict.__getitem__(value, key) for key in keys) + return children, keys + + +def _flatten_tuple_dict_with_keys(value: TupleDict): + children, keys = _flatten_tuple_dict(value) + return tuple((jax.tree_util.DictKey(key), child) for key, child in zip(keys, children)), keys + + +def _unflatten_tuple_dict(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_tuple_dict_with_keys, + _unflatten_tuple_dict, + _flatten_tuple_dict, + ) + TupleDict._jax_pytree_registered = True + + +def _resolve_layout_mode( + rank: int, + *, + tensor_spec: TensorSpec | None, +) -> tuple[tuple[int, ...], tuple[int, ...]]: + layout = None if tensor_spec is None else tensor_spec.layout + mode = None if tensor_spec is None else tensor_spec.mode + + dimensions = tuple(range(rank)) + layout = tuple(reversed(dimensions)) if layout is None else tuple(layout) + mode = dimensions if mode is None else tuple(mode) + if tuple(sorted(layout)) != dimensions: + raise ValueError(f"layout must be a permutation for rank {rank}, got {layout}") + if tuple(sorted(mode)) != dimensions: + raise ValueError(f"mode must be a permutation for rank {rank}, got {mode}") + return layout, mode + + +@dataclass(frozen=True, kw_only=True) +class JaxTensorDesc(TensorDesc): + """Abstract JAX tensor metadata and its CUTLASS lowering specification. + + The descriptor reads only shape and dtype from an array-like value. Its + shape and stride describe the modes presented to the kernel after the + declared ``TensorSpec.mode`` permutation. The stride is derived from the + compact layout requested from XLA; it never inspects a physical JAX buffer + or device. ``tensor_spec`` is retained unchanged for use by + :func:`call_cutedsl`; ``None`` preserves CUTLASS's inferred-default + behavior. The ``layout`` and ``mode`` properties expose normalized + defaults without rewriting that native object. Prefer :meth:`from_value` + for input metadata because it applies ``TensorSpec.mode`` to the public + array shape. A shape passed directly to the constructor must already be in + kernel-visible mode order. + """ + + tensor_spec: TensorSpec | None = None + jax_layout: tuple[int, ...] = field(init=False) + jax_mode: tuple[int, ...] = field(init=False) + + def __post_init__(self) -> None: + rank = len(tuple(self.shape)) + layout, mode = _resolve_layout_mode( + rank, + tensor_spec=self.tensor_spec, + ) + dimensions = tuple(range(rank)) + + expected_order = tuple(sorted(dimensions, key=lambda dim: layout[mode[dim]])) + if self.stride_order is None: + object.__setattr__(self, "stride_order", expected_order) + elif tuple(self.stride_order) != expected_order: + raise ValueError(f"stride_order must agree with jax_layout and jax_mode: expected {expected_order}, got {tuple(self.stride_order)}") + + object.__setattr__(self, "dtype", jnp.dtype(self.dtype)) + object.__setattr__(self, "jax_layout", layout) + object.__setattr__(self, "jax_mode", mode) + super().__post_init__() + + expected_stride: list[Any] = [None] * rank + running: Any = 1 + for dim in expected_order: + expected_stride[dim] = running + running *= self.shape[dim] + if self.stride != tuple(expected_stride): + raise ValueError( + "stride must describe the compact layout declared by jax_layout and jax_mode: " f"expected {tuple(expected_stride)}, got {self.stride}" + ) + + @classmethod + def from_value( + cls, + value: Any, + *, + tensor_spec: TensorSpec | None = None, + layout: Sequence[int] | None = None, + mode: Sequence[int] | None = None, + name: str = "", + ) -> "JaxTensorDesc": + """Capture array metadata together with its native lowering spec. + + ``layout`` and ``mode`` are convenience forms that construct a native + ``TensorSpec``. Supplying neither keeps ``tensor_spec=None`` so CUTLASS + can infer its complete default, including divisibility. + """ + + input_shape = require_array(value, name=name or None) + if tensor_spec is not None: + if layout is not None or mode is not None: + raise ValueError("tensor_spec cannot be combined with explicit layout or mode") + elif layout is not None or mode is not None: + tensor_spec = TensorSpec( + layout=None if layout is None else tuple(layout), + mode=None if mode is None else tuple(mode), + ) + + _, mode = _resolve_layout_mode(len(input_shape), tensor_spec=tensor_spec) + shape = tuple(input_shape[dim] for dim in mode) + return cls( + dtype=value.dtype, + shape=shape, + tensor_spec=tensor_spec, + name=name, + ) + + @property + def layout(self) -> tuple[int, ...]: + return self.jax_layout + + @property + def mode(self) -> tuple[int, ...]: + assert self.jax_mode is not None + return self.jax_mode + + @property + def array_shape(self) -> tuple[Any, ...]: + """Return the public JAX shape before ``TensorSpec.mode`` is applied.""" + + shape: list[Any] = [None] * self.ndim + for kernel_dim, array_dim in enumerate(self.mode): + shape[array_dim] = self.shape[kernel_dim] + return tuple(shape) + + +class ApiBaseJax(ApiBase, ABC): + """Base for sample-signature-bound, traceable JAX callable objects. + + Instances are intentionally not wrapped in ``jax.jit``. The application + owns JIT, sharding, donation, and device-placement policy. + """ + + def __init__(self) -> None: + super().__init__() + self._configuration_frozen = False + + def __setattr__(self, name: str, value: Any) -> None: + """Invalidate pre-call support state and reject post-call mutation. + + JAX caches an executable by callable identity, not by the current + contents of a callable object's attributes. Compile-affecting state may + therefore change before the first invocation, but becomes immutable as + soon as this object has participated in tracing or execution. + """ + + if self.__dict__.get("_configuration_frozen", False): + raise AttributeError( + f"{self.__class__.__name__} configuration is immutable after its first call; " "construct a new instance for different static options" + ) + + if name not in {"_configuration_frozen", "_is_supported"} and self.__dict__.get("_is_supported", False): + object.__setattr__(self, "_is_supported", False) + object.__setattr__(self, name, value) + + @final + def check_support(self) -> bool: + """Validate and cache sample metadata and static configuration.""" + + if self._is_supported: + return True + self._check_support() + object.__setattr__(self, "_is_supported", True) + return True + + @abstractmethod + def _check_support(self) -> None: + """Validate and resolve operation-specific configuration or raise.""" + + @abstractmethod + def _call_impl(self, *args: Any, **kwargs: Any) -> Any: + """Lower the operation using invocation-time JAX arrays.""" + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """Invoke this object with JAX arrays matching its sample signature.""" + + self._ensure_support_checked() + object.__setattr__(self, "_configuration_frozen", True) + return self._call_impl(*args, **kwargs) + + def make_tensor_desc( + self, + value: Any, + *, + tensor_spec: TensorSpec | None = None, + layout: Sequence[int] | None = None, + mode: Sequence[int] | None = None, + name: str = "", + ) -> JaxTensorDesc: + """Return abstract JAX metadata using a declared CUTLASS tensor spec.""" + + return JaxTensorDesc.from_value( + value, + tensor_spec=tensor_spec, + layout=layout, + mode=mode, + name=name, + ) + + def as_dtype(self, value: Any) -> Any: + """Return a JAX dtype without retaining a dtype-bearing value.""" + + return as_dtype(value) + + def as_optional_dtype(self, value: Any | None) -> Any | None: + """Return ``None`` or a JAX dtype without retaining the source value.""" + + return as_optional_dtype(value) + + def require_dtype( + self, + value: Any, + valid_dtypes: Iterable[Any], + *, + name: str | None = None, + default: Any = _NO_DEFAULT, + ) -> Any: + """Return a supported dtype from a dtype-like value or descriptor.""" + + return require_dtype(value, valid_dtypes, name=name, default=default) + + def make_optional_tensor_desc( + self, + value: Any | None, + *, + tensor_spec: TensorSpec | None = None, + layout: Sequence[int] | None = None, + mode: Sequence[int] | None = None, + name: str = "", + ) -> JaxTensorDesc | None: + """Return metadata for an optional sample value.""" + + if value is None: + return None + return self.make_tensor_desc( + value, + tensor_spec=tensor_spec, + layout=layout, + mode=mode, + name=name, + ) + + def check_tensor_signature( + self, + value: Any, + expected: JaxTensorDesc, + *, + name: str = "", + ) -> JaxTensorDesc: + """Validate an invocation-time value against a sample descriptor.""" + + actual_shape = require_array(value, name=name or None) + if actual_shape != expected.array_shape: + raise ValueError(f"{name} tensor shape mismatch: expected {expected.array_shape}, got {actual_shape}") + actual = self.make_tensor_desc( + value, + tensor_spec=expected.tensor_spec, + name=name, + ) + self.check_dtype(actual, expected.dtype_name, name) + return actual + + def check_optional_tensor_signature( + self, + value: Any | None, + expected: JaxTensorDesc | None, + *, + name: str = "", + ) -> JaxTensorDesc | None: + """Validate optional-operand presence and, when present, its signature.""" + + if value is None and expected is None: + return None + if value is None or expected is None: + expected_presence = "present" if expected is not None else "absent" + actual_presence = "present" if value is not None else "absent" + raise ValueError(f"{name} presence mismatch: expected {expected_presence}, got {actual_presence}") + return self.check_tensor_signature(value, expected, name=name) + + @staticmethod + def freeze_mapping(values: Mapping[str, Any]) -> Mapping[str, Any]: + """Return an immutable copy suitable for persistent callable state.""" + + return MappingProxyType(dict(values)) + + def check_tensor_signatures( + self, + expected: Mapping[str, JaxTensorDesc | None], + values: Mapping[str, Any], + ) -> None: + """Validate invocation metadata against named sample descriptors.""" + + for name, expected_desc in expected.items(): + self.check_optional_tensor_signature(values[name], expected_desc, name=name) + + def get_jax_callable(self) -> Callable[..., Any]: + """Return this stable, un-jitted callable object.""" + + return self + + +@dataclass(frozen=True) +class BufferSpec: + """Shape, dtype, tensor metadata, and optional fill value for a result. + + A spec passed through ``outputs`` is returned to the caller. A spec passed + through ``workspaces`` is supplied to the CuTe launcher after the public + outputs and then omitted from the JAX-visible result. ``tensor_spec`` is a + native :class:`cutlass.jax.TensorSpec`; ``None`` asks CUTLASS to infer its + default from the shape and dtype. ``fill_value=None`` leaves the result + uninitialized; any other value initializes it before the kernel launch. + """ + + name: str + shape: Tuple[Any, ...] + dtype: Any + tensor_spec: TensorSpec | None = None + fill_value: Any = None + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("BufferSpec.name must not be empty") + object.__setattr__(self, "shape", tuple(self.shape)) + + +@dataclass(frozen=True) +class _CallPlan: + num_user_inputs: int + all_results: Tuple[BufferSpec, ...] + num_public_results: int + result_input_sources: Tuple[Optional[int], ...] + input_output_aliases: Tuple[Tuple[int, int], ...] + initialized_result_indices: Tuple[int, ...] + + @property + def num_total_inputs(self) -> int: + return self.num_user_inputs + len(self.initialized_result_indices) + + +@dataclass(frozen=True) +class _LaunchConfig: + """Immutable compile-time state for the shared CuTe launch adapter.""" + + fn: Callable[..., None] + num_user_inputs: int + num_total_inputs: int + result_input_sources: Tuple[Optional[int], ...] + static_args: Tuple[Tuple[str, Any], ...] + + +def _build_call_plan( + *, + num_user_inputs: int, + outputs: Sequence[BufferSpec], + workspaces: Sequence[BufferSpec], +) -> _CallPlan: + outputs = tuple(outputs) + workspaces = tuple(workspaces) + if not outputs: + raise ValueError("call_cutedsl requires at least one public output") + + all_results = outputs + workspaces + if not all(isinstance(spec, BufferSpec) for spec in all_results): + raise TypeError("outputs and workspaces must contain only BufferSpec values") + names = [spec.name for spec in all_results] + if len(set(names)) != len(names): + raise ValueError(f"Buffer names must be unique, got {names}") + + aliases = [] + result_sources: list[Optional[int]] = [None] * len(all_results) + initialized_indices = [] + next_input_idx = num_user_inputs + for result_idx, spec in enumerate(all_results): + if spec.fill_value is None: + continue + result_sources[result_idx] = next_input_idx + aliases.append((next_input_idx, result_idx)) + initialized_indices.append(result_idx) + next_input_idx += 1 + + return _CallPlan( + num_user_inputs=num_user_inputs, + all_results=all_results, + num_public_results=len(outputs), + result_input_sources=tuple(result_sources), + input_output_aliases=tuple(aliases), + initialized_result_indices=tuple(initialized_indices), + ) + + +@cute.jit(preprocess=False) +def _launch_adapter(stream, *args, config: cutlass.Constexpr): + """Reconstruct canonical ``inputs, outputs, workspaces`` launcher order.""" + + input_args = args[: config.num_total_inputs] + fresh_result_idx = config.num_total_inputs + canonical_results = [] + for source in config.result_input_sources: + if source is None: + canonical_results.append(args[fresh_result_idx]) + fresh_result_idx += 1 + else: + canonical_results.append(input_args[source]) + + if fresh_result_idx != len(args): + raise RuntimeError("CuTe launcher received more result buffers than expected") + + config.fn( + stream, + *input_args[: config.num_user_inputs], + *canonical_results, + **dict(config.static_args), + ) + + +def call_cutedsl( + fn: Callable[..., None], + inputs: Sequence[Any], + *, + outputs: Sequence[BufferSpec], + workspaces: Sequence[BufferSpec] = (), + input_specs: Optional[Sequence[Optional[TensorSpec]]] = None, + static_args: Optional[Mapping[str, Any]] = None, + allow_cuda_graph: bool = True, + compile_options: Any = None, + use_static_tensors: bool = True, +) -> Tuple[Any, ...]: + """Invoke a CuTe DSL launcher as a functional JAX operation. + + The launcher's canonical signature is:: + + fn(stream, *inputs, *outputs, *workspaces, **static_args) -> None + + Call this function while tracing a JAX function (normally inside + ``jax.jit``). Public results are returned as a tuple in ``outputs`` order; + workspace results are hidden. All output and workspace sizes must be + derivable from abstract input metadata, not from runtime tensor values. + + Inputs must be a flat sequence of array-like values; nested pytrees are + intentionally outside the FE ABI. ``input_specs`` and + ``BufferSpec.tensor_spec`` accept native ``cutlass.jax.TensorSpec`` objects; + validation is delegated to ``cutlass.jax.cutlass_call``. Tensor shapes and + strides are compile-time constants by default; pass + ``use_static_tensors=False`` for a future symbolic-shape path. Filled + buffers are passed as internal aliased inputs so the public JAX function + remains functional. ``fn`` and every value in ``static_args`` must be + immutable and hashable because they participate in JAX and CUTLASS + compilation cache keys. + """ + + inputs = tuple(inputs) + static_args = dict(static_args or {}) + plan = _build_call_plan( + num_user_inputs=len(inputs), + outputs=outputs, + workspaces=workspaces, + ) + + if input_specs is None: + normalized_input_specs: Tuple[Optional[TensorSpec], ...] = (None,) * len(inputs) + else: + normalized_input_specs = tuple(input_specs) + if len(normalized_input_specs) != len(inputs): + spec_label = "spec" if len(inputs) == 1 else "specs" + raise ValueError(f"Expected {len(inputs)} input tensor {spec_label}, got " f"{len(normalized_input_specs)}") + + for input_idx, value in enumerate(inputs): + if not hasattr(value, "shape") or not hasattr(value, "dtype"): + raise TypeError( + "call_cutedsl inputs must be a flat sequence of array-like " + f"values with shape and dtype metadata; input #{input_idx} " + f"is {type(value).__name__}" + ) + + initialized_buffers = [] + initialized_specs = [] + for result_idx in plan.initialized_result_indices: + spec = plan.all_results[result_idx] + initialized_buffers.append(jnp.full(spec.shape, spec.fill_value, dtype=spec.dtype)) + initialized_specs.append(spec.tensor_spec) + + cutlass_input_specs = normalized_input_specs + tuple(initialized_specs) + cutlass_output_specs = tuple(x.tensor_spec for x in plan.all_results) + result_shape_dtypes = tuple(jax.ShapeDtypeStruct(x.shape, x.dtype) for x in plan.all_results) + + launch_config = _LaunchConfig( + fn=fn, + num_user_inputs=plan.num_user_inputs, + num_total_inputs=plan.num_total_inputs, + result_input_sources=plan.result_input_sources, + static_args=tuple(sorted(static_args.items())), + ) + call = cutlass_jax.cutlass_call( + _launch_adapter, + output_shape_dtype=result_shape_dtypes, + input_spec=cutlass_input_specs, + output_spec=cutlass_output_specs, + input_output_aliases=dict(plan.input_output_aliases), + allow_cuda_graph=allow_cuda_graph, + compile_options=compile_options, + use_static_tensors=use_static_tensors, + config=launch_config, + ) + all_results = call(*inputs, *initialized_buffers) + + if not isinstance(all_results, (tuple, list)): + all_results = (all_results,) + if len(all_results) != len(plan.all_results): + raise RuntimeError(f"CuTe call returned {len(all_results)} buffers; expected " f"{len(plan.all_results)}") + return tuple(all_results[: plan.num_public_results]) + + +__all__ = [ + "ApiBaseJax", + "BufferSpec", + "JaxTensorDesc", + "TupleDict", + "as_dtype", + "as_optional_dtype", + "call_cutedsl", + "require_array", + "require_dtype", +] diff --git a/python/cudnn/_jax/gemm.py b/python/cudnn/_jax/gemm.py new file mode 100644 index 000000000..cf2d9c4b2 --- /dev/null +++ b/python/cudnn/_jax/gemm.py @@ -0,0 +1,197 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared JAX tensor metadata for FE-OSS GEMM operations.""" + +from __future__ import annotations + +from typing import Any + +import jax.numpy as jnp +from cutlass.jax import TensorSpec + +from ..gemm_validation import ( + require_block_scale_shapes, + require_contiguous_alignment, + require_gemm_shapes, +) +from .api_base import JaxTensorDesc, as_dtype, require_array + +GEMM_A_LAYOUTS = ("LMK", "LKM") +GEMM_B_LAYOUTS = ("LNK", "LKN") +GEMM_C_LAYOUTS = ("LMN", "LNM") + + +def require_layout(name: str, value: str, supported: tuple[str, ...]) -> str: + """Return a canonical public axis-order string from a supported set.""" + + if not isinstance(value, str): + raise TypeError(f"{name} must be a string, got {type(value).__name__}") + value = value.upper() + 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 _gemm_tensor_spec( + name: str, + layout: str, + *, + kernel_modes: str, + supported: tuple[str, ...], +) -> TensorSpec: + """Map a compact row-major public layout to canonical kernel modes.""" + + layout = require_layout(name, layout, supported) + mode = tuple(layout.index(dim) for dim in kernel_modes) + return TensorSpec( + layout=(2, 1, 0), + mode=mode, + ) + + +def as_gemm_tensor_desc( + name: str, + value: Any, + tensor_spec: TensorSpec, +) -> JaxTensorDesc: + """Return kernel-visible GEMM metadata for an array or existing descriptor.""" + + if isinstance(value, JaxTensorDesc): + expected_layout = tuple(tensor_spec.layout) + expected_mode = tuple(tensor_spec.mode) + if value.layout != expected_layout or value.mode != expected_mode: + raise ValueError( + f"{name} descriptor layout does not match the requested GEMM layout: " + f"expected layout={expected_layout}, mode={expected_mode}; " + f"got layout={value.layout}, mode={value.mode}" + ) + return value + return JaxTensorDesc.from_value( + value, + tensor_spec=tensor_spec, + name=name, + ) + + +def require_gemm_inputs( + a_tensor: Any, + b_tensor: Any, +) -> tuple[int, int, int, int, Any]: + """Validate common dense GEMM metadata and return ``M, N, K, L, dtype``.""" + + a_shape = require_array(a_tensor, name="a_tensor", rank=3) + a_dtype = as_dtype(a_tensor) + b_shape = require_array(b_tensor, name="b_tensor", rank=3, dtype=a_dtype) + m, n, k, batch = require_gemm_shapes(a_shape, b_shape) + return m, n, k, batch, a_dtype + + +def require_fp8_block_scales( + sfa_tensor: Any, + sfb_tensor: Any, + *, + m: int, + n: int, + k: int, + batch: int, + sf_vec_size: int, +) -> None: + """Validate the native MXFP8 scale-factor ABI used by dense kernels.""" + + if sf_vec_size != 32: + raise NotImplementedError("The JAX MXFP8 path requires sf_vec_size=32, " f"got {sf_vec_size}") + sfa_shape = require_array( + sfa_tensor, + name="sfa_tensor", + rank=6, + dtype=jnp.float8_e8m0fnu, + ) + sfb_shape = require_array( + sfb_tensor, + name="sfb_tensor", + rank=6, + dtype=jnp.float8_e8m0fnu, + ) + require_block_scale_shapes( + sfa_shape, + sfb_shape, + m=m, + n=n, + k=k, + batch=batch, + sf_vec_size=sf_vec_size, + ) + + +def require_16_byte_extent(name: str, elements: int, dtype: Any) -> None: + """Require the kernel's contiguous mode to span a multiple of 16 bytes.""" + + dtype = jnp.dtype(dtype) + require_contiguous_alignment(name, elements, dtype.itemsize * 8) + + +def gemm_a_tensor_spec(layout: str) -> TensorSpec: + """Describe public A layout ``LMK`` or ``LKM`` as kernel ``(M,K,L)``.""" + + return _gemm_tensor_spec( + "a_layout", + layout, + kernel_modes="MKL", + supported=GEMM_A_LAYOUTS, + ) + + +def gemm_b_tensor_spec(layout: str) -> TensorSpec: + """Describe public B layout ``LNK`` or ``LKN`` as kernel ``(N,K,L)``.""" + + return _gemm_tensor_spec( + "b_layout", + layout, + kernel_modes="NKL", + supported=GEMM_B_LAYOUTS, + ) + + +def gemm_c_tensor_spec(layout: str, *, name: str = "c_layout") -> TensorSpec: + """Describe public C/D layout ``LMN`` or ``LNM`` as kernel ``(M,N,L)``.""" + + return _gemm_tensor_spec( + name, + layout, + kernel_modes="MNL", + supported=GEMM_C_LAYOUTS, + ) + + +def block_scale_tensor_spec() -> TensorSpec: + """Describe the compact six-dimensional block-scale atom layout.""" + + return TensorSpec( + layout=(2, 1, 4, 0, 3, 5), + mode=(0, 1, 2, 3, 4, 5), + ) + + +def probability_tensor_spec() -> TensorSpec: + """Describe a logical ``(M, 1, L)`` tensor with contiguous ``M`` mode.""" + + return TensorSpec(layout=(0, 1, 2), mode=(0, 1, 2)) + + +__all__ = [ + "GEMM_A_LAYOUTS", + "GEMM_B_LAYOUTS", + "GEMM_C_LAYOUTS", + "as_gemm_tensor_desc", + "block_scale_tensor_spec", + "gemm_a_tensor_spec", + "gemm_b_tensor_spec", + "gemm_c_tensor_spec", + "probability_tensor_spec", + "require_fp8_block_scales", + "require_gemm_inputs", + "require_16_byte_extent", + "require_layout", +] diff --git a/python/cudnn/_jax/grouped_gemm.py b/python/cudnn/_jax/grouped_gemm.py new file mode 100644 index 000000000..608a2a3fb --- /dev/null +++ b/python/cudnn/_jax/grouped_gemm.py @@ -0,0 +1,171 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared JAX metadata validation for contiguous grouped GEMMs.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import jax.numpy as jnp +from cutlass.jax import TensorSpec + +from ..gemm_validation import block_scale_shape +from .api_base import as_dtype, require_array + + +def require_grouped_gemm_inputs( + a_tensor: Any, + b_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + *, + max_experts: int, + valid_ab_dtypes: Iterable[Any] | None = None, +) -> tuple[int, int, int, int, Any]: + """Validate the common contiguous grouped-GEMM inputs. + + A has logical shape ``(M, K, 1)`` while B has ``(N, K, L)``. The + grouped dimension is flattened into A's M dimension and described at + runtime by ``padded_offsets``. + """ + + if valid_ab_dtypes is None: + valid_ab_dtypes = (jnp.float8_e4m3fn, jnp.float8_e5m2) + a_shape = require_array( + a_tensor, + name="a_tensor", + rank=3, + dtype=valid_ab_dtypes, + ) + a_dtype = as_dtype(a_tensor) + b_shape = require_array(b_tensor, name="b_tensor", rank=3, dtype=a_dtype) + m, k, a_batch = a_shape + n, b_k, experts = b_shape + if a_batch != 1: + raise ValueError(f"a_tensor must have shape (M, K, 1), got {a_shape}") + if b_k != k: + raise ValueError(f"a_tensor and b_tensor must have matching K dimensions, got {a_shape} and {b_shape}") + dimensions = {"M": m, "N": n, "K": k, "L": experts} + nonpositive = [f"{name}={value}" for name, value in dimensions.items() if value <= 0] + if nonpositive: + raise ValueError("Grouped GEMM dimensions must be positive, got " + ", ".join(nonpositive)) + 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, a_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: + """Validate the E8M0 scale-factor ABI for contiguous grouped GEMMs.""" + + 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_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: + """Validate native block-scale shapes and a caller-supplied dtype set.""" + + 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_vector( + name: str, + tensor: Any, + *, + length: int, + dtype: Any = None, +) -> Any: + """Validate a contiguous one-dimensional grouped-GEMM tensor.""" + + if dtype is None: + dtype = jnp.float32 + require_array(tensor, name=name, shape=(length,), dtype=dtype) + return as_dtype(tensor) + + +def require_grouped_probability(name: str, tensor: Any, *, m: int) -> None: + """Validate a per-row FP32 probability tensor.""" + + require_array( + tensor, + name=name, + shape=(m, 1, 1), + dtype=jnp.float32, + ) + + +def grouped_bias_tensor_spec() -> TensorSpec: + """Describe a logical ``(N, L)`` bias with contiguous N mode.""" + + return TensorSpec(layout=(0, 1), mode=(0, 1)) + + +def grouped_workspace_tensor_spec() -> TensorSpec: + """Describe the 128-byte-aligned grouped-kernel workspace.""" + + return TensorSpec(ptr_assumed_align=128) + + +__all__ = [ + "grouped_bias_tensor_spec", + "grouped_workspace_tensor_spec", + "require_grouped_block_scales", + "require_grouped_fp8_scales", + "require_grouped_gemm_inputs", + "require_grouped_probability", + "require_grouped_vector", +] diff --git a/python/cudnn/_operation_api.py b/python/cudnn/_operation_api.py new file mode 100644 index 000000000..7a0a307d4 --- /dev/null +++ b/python/cudnn/_operation_api.py @@ -0,0 +1,41 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy exports for frontend-only operation packages.""" + +from __future__ import annotations + +from importlib import import_module +from typing import Any, Callable, MutableMapping, Sequence, Tuple + + +def make_operation_api( + module_globals: MutableMapping[str, Any], + *, + exports: Sequence[str], +) -> Tuple[list[str], Callable[[str], Any]]: + """Create ``__all__`` and ``__getattr__`` for an operation package. + + Unqualified symbols always come from the sibling ``api.py`` module. The + sibling ``jax.py`` module is available only through the explicit ``jax`` + attribute. Dependency availability does not change how a name is routed. + """ + + package_name = module_globals["__name__"] + exported_names = tuple(exports) + + def get_attribute(name: str) -> Any: + if name in {"api", "jax"}: + value = import_module(f".{name}", package_name) + elif name in exported_names: + value = getattr(import_module(".api", package_name), name) + else: + raise AttributeError(f"module {package_name!r} has no attribute {name!r}") + + module_globals[name] = value + return value + + return list(exported_names), get_attribute + + +__all__ = ["make_operation_api"] diff --git a/python/cudnn/api_base.py b/python/cudnn/api_base.py index 2a2a1741a..6e31ab5eb 100644 --- a/python/cudnn/api_base.py +++ b/python/cudnn/api_base.py @@ -1,27 +1,25 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -""" -Base classes for cuDNN API wrappers. +"""Framework-neutral metadata and validation support for FE-OSS APIs. -This module provides abstract base classes that define common interfaces -for cuDNN API wrapper classes, including validation, compilation, and execution patterns. +The eager Torch lifecycle is loaded only when one of its compatibility names +is requested. JAX wrappers can therefore reuse :class:`TensorDesc` and +:class:`ApiBase` without importing Torch, CUDA bindings, or CuTe DSL. """ from __future__ import annotations -from abc import ABC, abstractmethod +from collections.abc import Collection, Sequence from dataclasses import dataclass, field -from typing import Any, List, Tuple, Optional +from importlib import import_module import logging -import threading -import cuda.bindings.driver as cuda -import cutlass -import torch +from typing import Any, TYPE_CHECKING + +from ._experimental_warnings import warn_experimental_api_once -import cutlass.cute as cute -from cudnn._experimental_warnings import warn_experimental_api_once -from cudnn.datatypes import _convert_to_cutlass_data_type +if TYPE_CHECKING: + from .api_base_torch import APIBase, ApiBaseTorch, TorchTensorDesc def ceil_div(a: int, b: int) -> int: @@ -29,1025 +27,355 @@ def ceil_div(a: int, b: int) -> int: def is_power_of_2(n: int) -> bool: - """Check if n is a power of 2.""" - return n > 0 and (n & (n - 1)) == 0 - - -_experimental_api_warnings_emitted = set() -_experimental_api_warnings_lock = threading.Lock() - + """Return whether ``n`` is a positive power of two.""" -def warn_experimental_api_once(logger: logging.Logger, api_name: str) -> None: - """Emit the experimental API warning once per API class per process.""" - with _experimental_api_warnings_lock: - if api_name in _experimental_api_warnings_emitted: - return - _experimental_api_warnings_emitted.add(api_name) - - logger.warning("%s is an experimental API", api_name) - - -def _reset_experimental_api_warning_registry() -> None: - """Reset experimental API warning state for tests.""" - with _experimental_api_warnings_lock: - _experimental_api_warnings_emitted.clear() + return n > 0 and (n & (n - 1)) == 0 -@dataclass(frozen=True) +_DTYPE_ALIASES = { + "half": "float16", + "float": "float32", + "double": "float64", + "long": "int64", + # Torch names the storage type; the logical element remains E2M1. + "float4_e2m1fn_x2": "float4_e2m1fn", + # CUTLASS numeric classes omit separators used by framework dtypes. + "float4e2m1fn": "float4_e2m1fn", + "float8e4m3fn": "float8_e4m3fn", + "float8e5m2": "float8_e5m2", + "float8e8m0fnu": "float8_e8m0fnu", +} + +_DTYPE_BITS = { + "bool": 1, + "float4_e2m1fn": 4, + "float8_e4m3fn": 8, + "float8_e5m2": 8, + "float8_e8m0fnu": 8, + "int8": 8, + "uint8": 8, + "bfloat16": 16, + "float16": 16, + "int16": 16, + "uint16": 16, + "float32": 32, + "int32": 32, + "uint32": 32, + "float64": 64, + "int64": 64, + "uint64": 64, +} + + +def _storage_dtype_name(dtype: Any) -> str: + name = getattr(dtype, "name", None) + if not isinstance(name, str): + name = getattr(dtype, "__name__", None) + if not isinstance(name, str): + name = str(dtype) + return name.rsplit(".", 1)[-1].lower() + + +def canonical_dtype_name(dtype: Any) -> str: + """Return a stable logical dtype name without importing a framework.""" + + name = _storage_dtype_name(dtype) + return _DTYPE_ALIASES.get(name, name) + + +def dtype_bits(dtype: Any) -> int | None: + """Return the logical element width when it can be determined safely.""" + + name = canonical_dtype_name(dtype) + if name in _DTYPE_BITS: + return _DTYPE_BITS[name] + itemsize = getattr(dtype, "itemsize", None) + if isinstance(itemsize, int) and itemsize > 0: + return itemsize * 8 + return None + + +def _compact_stride(shape: tuple[Any, ...], order: tuple[int, ...]) -> tuple[Any, ...]: + stride: list[Any] = [None] * len(shape) + running: Any = 1 + for dim in order: + stride[dim] = running + running *= shape[dim] + return tuple(stride) + + +@dataclass(frozen=True, kw_only=True) class TensorDesc: - """Metadata needed to validate/compile tensor signatures without storage.""" - - 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) + """Logical tensor metadata shared by framework adapters. + + ``shape`` and ``stride`` are expressed in the mode order presented to the + kernel. For Torch that is the observed tensor order unless a wrapper + explicitly transforms the descriptor. For JAX the adapter applies the + declared ``TensorSpec.mode`` permutation. JAX strides describe the compact + layout requested from XLA; they are not inspected from a physical buffer. + ``dtype_name`` is logical; ``storage_dtype_name`` and ``packing`` retain + adapter-specific storage details such as Torch's packed-uint8 FP4 ABI. + """ + + dtype: Any + shape: tuple[Any, ...] + stride: tuple[Any, ...] | None = None + stride_order: tuple[int, ...] | None = None + packing: str = "native" name: str = "" + ndim: int = field(init=False) + storage_dtype_name: str = field(init=False) + dtype_name: str = field(init=False) + element_bits: int | None = field(init=False) - def __post_init__(self): + def __post_init__(self) -> None: shape = tuple(self.shape) - stride = tuple(self.stride) - stride_order = tuple(self.stride_order) - device = self.device - if not isinstance(device, torch.device): - try: - device = torch.device(device) - except (TypeError, ValueError, RuntimeError) as exc: - raise TypeError(f"Invalid device for TensorDesc: {self.device!r}") from exc - + stride = None if self.stride is None else tuple(self.stride) + stride_order = None if self.stride_order is None else tuple(self.stride_order) ndim = len(shape) - if len(stride) != ndim: + packing = self.packing + storage_dtype_name = _storage_dtype_name(self.dtype) + if storage_dtype_name == "float4_e2m1fn_x2" and packing == "native": + packing = "fp4x2" + if packing not in {"native", "fp4x2"}: + raise ValueError(f"Unsupported tensor packing {packing!r}") + + if stride is not None and 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}") + if stride_order is not None: + 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}") + if stride is None and stride_order is not None: + stride = _compact_stride(shape, 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) + object.__setattr__(self, "packing", packing) + if packing == "fp4x2" and storage_dtype_name not in { + "uint8", + "float4_e2m1fn", + "float4_e2m1fn_x2", + }: + raise ValueError(f"fp4x2 packing requires uint8 or float4 storage, got {storage_dtype_name}") + dtype_name = "float4_e2m1fn" if packing == "fp4x2" else canonical_dtype_name(self.dtype) + object.__setattr__(self, "storage_dtype_name", storage_dtype_name) + object.__setattr__(self, "dtype_name", dtype_name) + object.__setattr__(self, "element_bits", _DTYPE_BITS.get(dtype_name, dtype_bits(self.dtype))) + + def size(self, dim: int | None = None) -> tuple[Any, ...] | Any: + """Return the logical shape or one dimension.""" - @staticmethod - def _normalize_dim(dim: int, ndim: int, *, allow_new_dim: bool = False) -> int: - min_dim = -ndim - (1 if allow_new_dim else 0) - max_dim = ndim if allow_new_dim else ndim - 1 - if dim < min_dim or dim > max_dim: - raise IndexError(f"Dimension out of range (expected to be in range of [{min_dim}, {max_dim}], but got {dim})") - if dim < 0: - dim += ndim + (1 if allow_new_dim else 0) - return dim + if dim is None: + return self.shape + return self.shape[dim] - @staticmethod - def _compute_stride_order(shape: Tuple[int, ...], stride: Tuple[int, ...]) -> Tuple[int, ...]: - return tuple(i for i, _ in sorted(enumerate(stride), key=lambda x: (x[1], shape[x[0]]))) + @property + def is_fp4(self) -> bool: + return self.dtype_name.startswith("float4_") - @staticmethod - def _numel(shape: Tuple[int, ...]) -> int: - numel = 1 - for size in shape: - numel *= size - return numel + @property + def is_fp8(self) -> bool: + return self.dtype_name.startswith("float8_") - @staticmethod - def _compute_contiguous_stride(shape: Tuple[int, ...]) -> Tuple[int, ...]: - if not shape: - return () - strides = [0] * len(shape) - running = 1 - for i in range(len(shape) - 1, -1, -1): - strides[i] = running - running *= max(shape[i], 1) - return tuple(strides) + @property + def is_f16(self) -> bool: + return self.dtype_name in {"float16", "bfloat16"} - @staticmethod - def _is_contiguous_with_order(shape: Tuple[int, ...], stride: Tuple[int, ...], order: Tuple[int, ...]) -> bool: - expected_stride = 1 - for dim in order: - size = shape[dim] - if size == 1: - continue - if stride[dim] != expected_stride: - return False - expected_stride *= size - return True - @staticmethod - def _compute_view_stride( - old_shape: Tuple[int, ...], - old_stride: Tuple[int, ...], - new_shape: Tuple[int, ...], - ) -> Optional[Tuple[int, ...]]: - old_numel = TensorDesc._numel(old_shape) - if old_numel == 0: - return TensorDesc._compute_contiguous_stride(new_shape) - - new_stride = [0] * len(new_shape) - view_dim = len(new_shape) - 1 - tensor_numel = 1 - view_numel = 1 - - for tensor_dim in range(len(old_shape) - 1, -1, -1): - tensor_numel *= old_shape[tensor_dim] - is_contiguous_chunk_end = tensor_dim == 0 or ( - old_shape[tensor_dim - 1] != 1 and old_stride[tensor_dim - 1] != tensor_numel * old_stride[tensor_dim] - ) - if is_contiguous_chunk_end: - while view_dim >= 0 and (view_numel < tensor_numel or new_shape[view_dim] == 1): - new_stride[view_dim] = view_numel * old_stride[tensor_dim] - view_numel *= new_shape[view_dim] - view_dim -= 1 - - if view_numel != tensor_numel: - return None - - if tensor_dim > 0: - tensor_numel = 1 - view_numel = 1 - - if view_dim != -1: - return None - return tuple(new_stride) - - def _with_layout(self, shape: Tuple[int, ...], stride: Tuple[int, ...]) -> "TensorDesc": - return TensorDesc( - dtype=self.dtype, - shape=shape, - stride=stride, - stride_order=self._compute_stride_order(shape, stride), - device=self.device, - interpret_uint8_as_fp4x2=self.interpret_uint8_as_fp4x2, - name=self.name, - ) - - def __len__(self) -> int: - if self.ndim == 0: - raise TypeError("len() of a 0-d tensor") - return self.shape[0] - - def size(self, dim: Optional[int] = None) -> torch.Size | int: - if dim is None: - return torch.Size(self.shape) - dim = self._normalize_dim(int(dim), self.ndim) - return self.shape[dim] +def _shape_of(value: TensorDesc | Sequence[Any] | Any) -> tuple[Any, ...]: + if isinstance(value, TensorDesc) or hasattr(value, "shape"): + return tuple(value.shape) + return tuple(value) - def permute(self, *dims: int | Tuple[int, ...] | List[int]) -> "TensorDesc": - if len(dims) == 1 and isinstance(dims[0], (tuple, list)): - dims = tuple(dims[0]) - dims = tuple(int(d) for d in dims) - if len(dims) != self.ndim: - raise RuntimeError(f"permute(): expected {self.ndim} dims, got {len(dims)}") - dims = tuple(self._normalize_dim(d, self.ndim) for d in dims) - if len(set(dims)) != self.ndim: - raise RuntimeError(f"permute(): dims must be unique, got {dims}") - - new_shape = tuple(self.shape[d] for d in dims) - new_stride = tuple(self.stride[d] for d in dims) - return self._with_layout(new_shape, new_stride) - - def transpose(self, dim0: int, dim1: int) -> "TensorDesc": - dim0 = self._normalize_dim(dim0, self.ndim) - dim1 = self._normalize_dim(dim1, self.ndim) - if dim0 == dim1: - return self - dims = list(range(self.ndim)) - dims[dim0], dims[dim1] = dims[dim1], dims[dim0] - return self.permute(dims) - - def squeeze(self, dim: Optional[int | Tuple[int, ...] | List[int]] = None) -> "TensorDesc": - if dim is None: - keep_dims = [i for i, size in enumerate(self.shape) if size != 1] - elif isinstance(dim, (tuple, list)): - squeeze_dims = tuple(self._normalize_dim(int(d), self.ndim) for d in dim) - if len(set(squeeze_dims)) != len(squeeze_dims): - raise RuntimeError(f"squeeze(): dims must be unique, got {squeeze_dims}") - squeeze_dims = {d for d in squeeze_dims if self.shape[d] == 1} - keep_dims = [i for i in range(self.ndim) if i not in squeeze_dims] - else: - squeeze_dim = self._normalize_dim(int(dim), self.ndim) - if self.shape[squeeze_dim] != 1: - return self - keep_dims = [i for i in range(self.ndim) if i != squeeze_dim] - - new_shape = tuple(self.shape[i] for i in keep_dims) - new_stride = tuple(self.stride[i] for i in keep_dims) - if new_shape == self.shape and new_stride == self.stride: - return self - return self._with_layout(new_shape, new_stride) - - def unsqueeze(self, dim: int) -> "TensorDesc": - dim = self._normalize_dim(dim, self.ndim, allow_new_dim=True) - - if dim >= self.ndim: - inserted_stride = 1 - else: - inserted_stride = self.stride[dim] * self.shape[dim] - - new_shape = self.shape[:dim] + (1,) + self.shape[dim:] - new_stride = self.stride[:dim] + (inserted_stride,) + self.stride[dim:] - return self._with_layout(new_shape, new_stride) - - def is_contiguous(self, memory_format: torch.memory_format = torch.contiguous_format) -> bool: - if memory_format in {torch.contiguous_format, torch.preserve_format}: - if self._numel(self.shape) == 0: - return True - return self._is_contiguous_with_order(self.shape, self.stride, tuple(range(self.ndim - 1, -1, -1))) - if memory_format == torch.channels_last: - if self.ndim != 4: - return False - return self._is_contiguous_with_order(self.shape, self.stride, (1, 3, 2, 0)) - if memory_format == torch.channels_last_3d: - if self.ndim != 5: - return False - return self._is_contiguous_with_order(self.shape, self.stride, (1, 4, 3, 2, 0)) - - raise ValueError(f"Unsupported memory format: {memory_format}") - - def contiguous(self) -> "TensorDesc": - if self.is_contiguous(): - return self - contiguous_stride = self._compute_contiguous_stride(self.shape) - return self._with_layout(self.shape, contiguous_stride) - - def view(self, *shape: int | Tuple[int, ...] | List[int]) -> "TensorDesc": - if len(shape) == 1 and isinstance(shape[0], (tuple, list)): - shape = tuple(shape[0]) - new_shape = tuple(int(s) for s in shape) - - old_numel = self._numel(self.shape) - infer_dim = None - known_numel = 1 - for i, size in enumerate(new_shape): - if size == -1: - if infer_dim is not None: - raise RuntimeError("only one dimension can be inferred") - infer_dim = i - continue - if size < 0: - raise RuntimeError(f"invalid shape dimension {size}") - known_numel *= size - - if infer_dim is not None: - if known_numel == 0 or old_numel % known_numel != 0: - raise RuntimeError(f"shape '{new_shape}' is invalid for input of size {old_numel}") - inferred_size = old_numel // known_numel - new_shape = new_shape[:infer_dim] + (inferred_size,) + new_shape[infer_dim + 1 :] - known_numel *= inferred_size - - if known_numel != old_numel: - raise RuntimeError(f"shape '{new_shape}' is invalid for input of size {old_numel}") - - new_stride = self._compute_view_stride(self.shape, self.stride, new_shape) - if new_stride is None: - raise RuntimeError( - "view size is not compatible with input tensor's size and stride " "(at least one dimension spans across two contiguous subspaces)" - ) - - return self._with_layout(new_shape, new_stride) - - def as_strided( - self, - size: Tuple[int, ...] | List[int], - stride: Tuple[int, ...] | List[int], - storage_offset: int = 0, - ) -> "TensorDesc": - if storage_offset != 0: - raise RuntimeError("TensorDesc.as_strided(): non-zero storage_offset is unsupported") - if not isinstance(size, (tuple, list)) or not isinstance(stride, (tuple, list)): - raise TypeError("TensorDesc.as_strided(): size and stride must be tuple/list") - - size = tuple(int(s) for s in size) - stride = tuple(int(s) for s in stride) - if len(size) != len(stride): - raise RuntimeError(f"TensorDesc.as_strided(): mismatch in length of size ({len(size)}) and stride ({len(stride)})") - if any(s < 0 for s in size): - raise RuntimeError(f"TensorDesc.as_strided(): invalid size, got {size}") - if any(s < 0 for s in stride): - raise RuntimeError(f"TensorDesc.as_strided(): invalid stride, got {stride}") - - return self._with_layout(size, stride) - - -class APIBase(ABC): - """Abstract base class for cuDNN API wrappers. - - This class defines the common interface that all API wrapper implementations - should follow, including configuration validation, compilation, and execution. - - Provides common functionality: - - Logging via self._logger - - Support validation tracking via self._is_supported - - Compiled kernel caching via self._compiled_kernel - - Stream management helpers - - Subclasses should implement the abstract methods to provide - API-specific validation logic and execution behavior. - - Example: - >>> class MyKernelAPI(APIBase): - ... def __init__(self, sample_input, sample_output, config): - ... super().__init__() - ... self.sample_input = sample_input - ... self.sample_output = sample_output - ... self.config = config - ... self._kernel = MyKernel - ... - ... def check_support(self) -> bool: - ... # Validate inputs and configuration - ... assert self.sample_input.dtype == torch.float32 - ... self._is_supported = True - ... return True - ... - ... def compile(self): - ... self._ensure_support_checked() - ... # Create and compile kernel - ... kernel = self._kernel(self.config) - ... self._compiled_kernel = cute.compile(kernel, ...) - ... - ... def execute(self, input_tensor, output_tensor, current_stream=None): - ... current_stream = self._get_default_stream(current_stream) - ... self._compiled_kernel(input_tensor, output_tensor, current_stream) - """ - def __init__(self): - """Initialize the API base. +def _dtype_name_of(value: TensorDesc | Any) -> str: + if isinstance(value, TensorDesc): + return value.dtype_name + if not isinstance(value, type) and hasattr(value, "dtype"): + value = value.dtype + return canonical_dtype_name(value) - Sets up: - - self._is_supported: Flag indicating if configuration is validated - - self._kernel: Kernel instance - - self._compiled_kernel: Cache for compiled kernel - - self._logger: Logger instance for this class - """ - self._is_supported = False - self._kernel = None - self._compiled_kernel = None - self._interpret_uint8_as_fp4x2 = False + +class ApiBase: + """Framework-neutral base for descriptor validation. + + This class deliberately has no compile or execute lifecycle. Framework + adapters own those policies while operation validators share descriptors. + """ + + def __init__(self) -> None: self._logger = logging.getLogger(self.__class__.__name__) + self._is_supported = False def _warn_experimental_api(self) -> None: warn_experimental_api_once(self._logger, self.__class__.__name__) - @abstractmethod - def check_support(self) -> bool: - """Check if the current configuration is supported by the kernel. - - This method should validate: - - Input/output tensor shapes and strides - - Data types compatibility - - Hardware capabilities (compute capability, memory, etc.) - - Configuration parameters (tile sizes, cluster shapes, etc.) - - Implementations should set self._is_supported = True if valid. - - :return: True if the configuration is supported - :rtype: bool - :raises AssertionError: If a configuration requirement is not met - - Example: - >>> def check_support(self) -> bool: - ... self._logger.debug("Checking support") - ... assert self.input.dtype in {torch.float16, torch.float32} - ... assert self.input.shape[0] % 16 == 0, "Shape must be 16-aligned" - ... self._is_supported = True - ... return True - """ - pass - - @abstractmethod - def compile(self) -> None: - """Compile the kernel with the current configuration. - - This method should: - 1. Ensure support has been checked (use self._ensure_support_checked()) - 2. Create the underlying kernel implementation and fake cute tensors from the sample tensor descriptors - 3. Compile the kernel using cute.compile() - 4. Cache the compiled kernel in self._compiled_kernel - - :raises AssertionError: If the configuration is not supported - - Example: - >>> def compile(self): - ... self._ensure_support_checked() - ... - ... kernel = self._kernel(self.config) - ... sample_input_cute = self._make_fake_cute_tensor_from_desc(self.sample_input) - ... sample_output_cute = self._make_fake_cute_tensor_from_desc(self.sample_output) - ... - ... self._compiled_kernel = cute.compile( - ... kernel, - ... sample_input_cute, - ... sample_output_cute - ... ) - """ - pass - - @abstractmethod - def execute( - self, - *args, - current_stream: Optional[cuda.CUstream] = None, - **kwargs, - ) -> Any: - """Execute the kernel with the provided inputs. - - This method should execute using the cached compiled kernel. - - :param args: Positional arguments (typically input/output tensors) - :param current_stream: CUDA stream for execution (optional) - :type current_stream: cuda.CUstream or None - :param kwargs: Additional keyword arguments for execution - :return: Execution result (if any) - :raises AssertionError: If compiled kernel is not available - - Example: - >>> def execute(self, input_tensor, output_tensor, current_stream=None): - ... current_stream = self._get_default_stream(current_stream) - ... assert self._compiled_kernel is not None, "Kernel not compiled" - ... self._logger.debug("Executing with compiled kernel") - ... self._compiled_kernel(input_tensor, output_tensor, current_stream) - """ - pass - - def __call__(self, *args, **kwargs) -> Any: - """Convenience method to execute the kernel. - - This is a shorthand for compiling (if needed) and then executing. - - :param args: Positional arguments passed to execute() - :param kwargs: Keyword arguments passed to execute() - :return: Result from execute() - - Example: - >>> api = MyKernelAPI(...) - >>> api.check_support() - >>> api.compile() - >>> api(input_tensor, output_tensor) - """ - if self._compiled_kernel is None: - self.compile() - return self.execute(*args, **kwargs) - - def _ensure_support_checked(self) -> None: - """Helper to ensure check_support() was called before compilation. - - If check_support() has not been called yet (self._is_supported is False), - this method will automatically call it. This prevents compilation - with invalid configurations. - - :raises AssertionError: If check_support() returns False or raises - - Example: - >>> def compile(self): - ... self._ensure_support_checked() # Automatic validation - ... # ... rest of compilation - """ - if not self._is_supported: - self._logger.info(f"{self.__class__.__name__}: check_support not previously called, calling now") - assert self.check_support(), "Unsupported configuration" - - def _get_default_stream(self, stream: Optional[cuda.CUstream]) -> cuda.CUstream: - """Get default CUDA stream if none provided. - - This is a convenience helper to handle optional stream parameters. - If a stream is provided, it is returned as-is. If None, the default - CUDA stream is returned. - - :param stream: CUDA stream or None - :type stream: cuda.CUstream or None - :return: CUDA stream (either provided or default) - :rtype: cuda.CUstream - - Example: - >>> def execute(self, input_tensor, output_tensor, current_stream=None): - ... current_stream = self._get_default_stream(current_stream) - ... # Now current_stream is guaranteed to be a valid stream - """ - if stream is None: - self._logger.debug(f"{self.__class__.__name__}: No CUDA stream provided, using default stream") - return cutlass.cuda.default_stream() - return stream - - def _pad_tensor_to_ndim( - self, - tensor: Optional[torch.Tensor | TensorDesc], - ndim: int, - name: str, - ) -> Optional[torch.Tensor | TensorDesc]: - """Pad a tensor/descriptor by unsqueezing at dim -1 until it reaches ndim rank. - - - If tensor is None, returns None. - - Unsqueezes at dim -1 until tensor/descriptor rank == ndim. - - Logs final reshape for traceability. - - :param tensor: The tensor/descriptor to pad (or None) - :param ndim: Target rank (pad trailing dims until reached) - :param name: Logical tensor name for logging - :return: The padded tensor/descriptor (or None) - """ - if tensor is None: - return None - - if tensor.ndim < ndim: - self._logger.info(f"Padding {name} to {ndim}D from {tensor.shape}") - for _ in range(ndim - tensor.ndim): - tensor = tensor.unsqueeze(-1) - return tensor - - def _unpad_tensor_to_ndim( - self, - tensor: Optional[torch.Tensor | TensorDesc], - ndim: int, - name: str, - ) -> Optional[torch.Tensor | TensorDesc]: - """Unpad a tensor/descriptor by squeezing at dim -1 until it reaches ndim rank. - - - If tensor is None, returns None. - - Squeezes at dim -1 until tensor/descriptor rank == ndim. - - Logs final reshape for traceability. - - :param tensor: The tensor/descriptor to unpad (or None) - :param ndim: Target rank (squeeze trailing dims until reached) - :param name: Logical tensor name for logging - :return: The unpadded tensor/descriptor (or None) - """ - if tensor is None: - return None - - if tensor.ndim > ndim: - self._logger.info(f"Unpadding {name} from {tensor.shape} to {ndim}D") - for _ in range(tensor.ndim - ndim): - if tensor.shape and tensor.shape[-1] == 1: - tensor = tensor.squeeze(-1) - else: - break - - if tensor.ndim != ndim: - self._logger.critical(f"Unpadding {name} resulted in shape {tensor.shape}, expected {ndim}D") - return tensor - - def _is_fp4x2(self, tensor_or_dtype: torch.Tensor | torch.dtype | TensorDesc) -> bool: - """Check if tensor or dtype is an FP4x2 packed datatype. - - :param tensor_or_dtype: The torch tensor or dtype to check - :type tensor_or_dtype: torch.Tensor | torch.dtype - :return: True if tensor/dtype is an FP4x2 packed type - :rtype: bool - """ - if tensor_or_dtype is None: - return False - if isinstance(tensor_or_dtype, TensorDesc): - dtype = tensor_or_dtype.dtype - interpret_uint8_as_fp4x2 = tensor_or_dtype.interpret_uint8_as_fp4x2 - elif isinstance(tensor_or_dtype, torch.Tensor): - dtype = tensor_or_dtype.dtype - interpret_uint8_as_fp4x2 = self._interpret_uint8_as_fp4x2 - else: - dtype = tensor_or_dtype - interpret_uint8_as_fp4x2 = self._interpret_uint8_as_fp4x2 - return (dtype == torch.float4_e2m1fn_x2) or (interpret_uint8_as_fp4x2 and dtype == torch.uint8) - - def _is_fp8(self, tensor_or_dtype: torch.Tensor | torch.dtype | TensorDesc) -> bool: - """Check if tensor or dtype is an FP8 datatype. - - :param tensor_or_dtype: The torch tensor or dtype to check - :type tensor_or_dtype: torch.Tensor | torch.dtype - :return: True if tensor/dtype is an FP8 type - :rtype: bool - """ - if tensor_or_dtype is None: - return False - dtype = tensor_or_dtype.dtype if isinstance(tensor_or_dtype, (torch.Tensor, TensorDesc)) else tensor_or_dtype - return dtype in {torch.float8_e5m2, torch.float8_e4m3fn} - - def _is_f16(self, tensor_or_dtype: torch.Tensor | torch.dtype | TensorDesc) -> bool: - """Check if tensor or dtype is an fp16 or bf16 datatype. - - :param tensor_or_dtype: The torch tensor or dtype to check - :type tensor_or_dtype: torch.Tensor | torch.dtype - :return: True if tensor/dtype is an fp16 or bf16 type - :rtype: bool - """ - if tensor_or_dtype is None: - return False - dtype = tensor_or_dtype.dtype if isinstance(tensor_or_dtype, (torch.Tensor, TensorDesc)) else tensor_or_dtype - return dtype in {torch.float16, torch.bfloat16} - - def _get_innermost_stride_dim(self, tensor: torch.Tensor, name: str = "") -> int: - """Return index of innermost contiguous dimension (stride == 1). - - :raises RuntimeError: If no dimension with stride 1 is found. - """ - idx = next((i for i, s in enumerate(tensor.stride()) if s == 1), None) - if idx is None: - self._logger.critical( - f"tensor {name} has shape: {tensor.shape} stride {tensor.stride()} – innermost contiguous (stride == 1) dimension not found. " - ) - raise RuntimeError(f"tensor {name} has shape: {tensor.shape} stride {tensor.stride()} – innermost contiguous (stride == 1) dimension not found. ") - return idx - - def _tensor_shape( - self, - tensor: Optional[torch.Tensor | TensorDesc], - name: str = "", - ) -> Optional[Tuple[int, ...]]: - """Get the logical shape of a tensor, handling FP4x2 packed datatypes. - - For FP4x2 datatypes, two values are packed per byte. The innermost - contiguous dimension (with stride 1) contains packed values, so the - logical shape for that dimension is 2x the physical shape. - - :param tensor: The tensor to get shape from (or None) - :type tensor: torch.Tensor or None - :param name: Logical tensor name for logging - :type name: str - :return: The logical shape tuple (or None if tensor is None) - :rtype: Tuple[int, ...] or None - """ - if tensor is None: - return None - if isinstance(tensor, TensorDesc): - return tensor.shape - - if self._is_fp4x2(tensor): - innermost_dim_index = self._get_innermost_stride_dim(tensor, name=name) - shape = tuple(dim * 2 if i == innermost_dim_index else dim for i, dim in enumerate(tensor.shape)) - self._logger.debug(f"FP4x2 tensor {name}: physical shape {tensor.shape} -> logical shape {shape}") - return shape - else: - return tensor.shape - - def _tensor_stride( + def check_tensor_shape( self, - tensor: Optional[torch.Tensor | TensorDesc], + tensor_or_shape: TensorDesc | Sequence[Any] | Any, + expected: Sequence[Any] | Collection[Sequence[Any]], name: str = "", - ) -> Optional[Tuple[int, ...]]: - """Get the logical stride of a tensor, handling FP4x2 packed datatypes. - - For FP4x2 datatypes, two values are packed per byte. The strides must - be adjusted to reflect logical element spacing. All strides are - multiplied by 2 since each physical element contains 2 logical elements. - - :param tensor: The tensor to get stride from (or None) - :type tensor: torch.Tensor or None - :param name: Logical tensor name for logging - :type name: str - :return: The logical stride tuple (or None if tensor is None) - :rtype: Tuple[int, ...] or None - """ - if tensor is None: - return None - if isinstance(tensor, TensorDesc): - return tensor.stride - - if self._is_fp4x2(tensor): - innermost_dim_index = self._get_innermost_stride_dim(tensor, name=name) - strides = tuple(s * 2 if i != innermost_dim_index else s for i, s in enumerate(tensor.stride())) - self._logger.debug(f"FP4x2 tensor {name}: physical stride {tensor.stride()} -> logical stride {strides}") - return strides - else: - return tensor.stride() + ) -> tuple[Any, ...]: + """Validate a logical tensor shape and return it as a tuple.""" - def _check_tensor_shape( - self, - tensor_or_shape: torch.Tensor | TensorDesc | Tuple[int, ...], - shape: Tuple[int, ...] | List[Tuple[int, ...]], - name: str = "", - ) -> Optional[Tuple[int, ...]]: - """Check if the shape of a tensor matches the expected shape(s). - - :param tensor_or_shape: The tensor to get shape from or the shape to check - :type tensor_or_shape: torch.Tensor | TensorDesc | Tuple[int, ...] - :param shape: expected shape or list of expected shapes - :type shape: Tuple[int, ...] | List[Tuple[int, ...]] - :param name: Logical tensor name for logging - :type name: str - :raises ValueError: If the shape of the tensor does not match the expected shape(s) - :return: The logical shape of the tensor - :rtype: Optional[Tuple[int, ...]] - """ - if tensor_or_shape is None: - return None - tensor_shape = self._tensor_shape(tensor_or_shape, name=name) if isinstance(tensor_or_shape, (torch.Tensor, TensorDesc)) else tensor_or_shape - if isinstance(shape, tuple): - if tensor_shape != shape: - raise ValueError(f"{name} tensor shape mismatch: expected {shape}, got {tensor_shape}") - elif isinstance(shape, list): - if tensor_shape not in shape: - raise ValueError(f"{name} tensor shape mismatch: expected one of {shape}, got {tensor_shape}") + actual = _shape_of(tensor_or_shape) + if isinstance(expected, tuple): + if actual != expected: + raise ValueError(f"{name} tensor shape mismatch: expected {expected}, got {actual}") else: - raise ValueError(f"Expected shape to be a tuple or list, got {type(shape)}") - return tensor_shape + choices = [tuple(shape) for shape in expected] + if actual not in choices: + raise ValueError(f"{name} tensor shape mismatch: expected one of {choices}, got {actual}") + return actual - def _check_tensor_stride( + def check_tensor_stride( self, - tensor_or_stride: torch.Tensor | TensorDesc | Tuple[int, ...], - stride: Optional[Tuple[int, ...] | List[Tuple[int, ...]]] = None, - stride_order: Optional[Tuple[int, ...] | List[Tuple[int, ...]]] = None, + tensor_or_stride: TensorDesc | Sequence[Any] | None, + stride: Sequence[Any] | Collection[Sequence[Any]] | None = None, + stride_order: Sequence[int] | Collection[Sequence[int]] | None = None, name: str = "", extra_error_msg: str = "", - ) -> Optional[Tuple[Tuple[int, ...], Tuple[int, ...]]]: - """Check if the stride of a tensor matches the expected stride(s) or stride order(s). - - :param tensor_or_stride: The tensor to get stride from or the stride to check - :type tensor_or_stride: torch.Tensor | TensorDesc | Tuple[int, ...] - :param stride: The expected stride(s) - :type stride: Tuple[int, ...] | List[Tuple[int, ...]] - :param stride_order: The expected stride order(s) - :type stride_order: Tuple[int, ...] | List[Tuple[int, ...]] - :param name: Logical tensor name for logging - :type name: str - :param extra_error_msg: Extra error message to add to the error - :type extra_error_msg: str - :raises ValueError: If the stride of the tensor does not match the expected stride order - :return: The stride and stride order of the tensor - :rtype: Optional[Tuple[Tuple[int, ...], Tuple[int, ...]]] - """ + ) -> tuple[tuple[Any, ...] | None, tuple[int, ...] | None]: + """Validate a kernel-visible layout contract.""" + if tensor_or_stride is None: return None, None - if isinstance(tensor_or_stride, TensorDesc): - tensor_stride = tensor_or_stride.stride - tensor_stride_order = tensor_or_stride.stride_order - elif isinstance(tensor_or_stride, torch.Tensor): - tensor_stride = self._tensor_stride(tensor_or_stride, name=name) - tensor_stride_order = tuple(i for i, s in sorted(enumerate(tensor_stride), key=lambda x: x[1])) + elif isinstance(tensor_or_stride, TensorDesc): + actual_stride = tensor_or_stride.stride + actual_order = tensor_or_stride.stride_order else: - tensor_stride = tensor_or_stride - tensor_stride_order = tuple(i for i, s in sorted(enumerate(tensor_stride), key=lambda x: x[1])) + actual_stride = tuple(tensor_or_stride) + actual_order = tuple(dim for dim, _ in sorted(enumerate(actual_stride), key=lambda item: item[1])) + + def with_context(message: str) -> str: + if extra_error_msg: + return f"{message}: {extra_error_msg}" + return message if stride is not None: - if isinstance(stride, tuple): - if tensor_stride != stride: - error_msg = f"{name} tensor stride mismatch: expected {stride}, got {tensor_stride}" - if extra_error_msg: - error_msg += f": {extra_error_msg}" - raise ValueError(error_msg) - elif isinstance(stride, list): - if tensor_stride not in stride: - error_msg = f"{name} tensor stride mismatch: expected one of {stride}, got {tensor_stride}" - if extra_error_msg: - error_msg += f": {extra_error_msg}" - raise ValueError(error_msg) - else: - error_msg = f"Expected stride to be a tuple or list, got {type(stride)}" - if extra_error_msg: - error_msg += f": {extra_error_msg}" - raise ValueError(error_msg) + if not isinstance(stride, (tuple, list)): + raise ValueError(with_context(f"Expected stride to be a tuple or list, got {type(stride)}")) + expected = tuple(stride) if isinstance(stride, tuple) else [tuple(item) for item in stride] + stride_matches = actual_stride in expected if isinstance(expected, list) else actual_stride == expected + if not stride_matches: + qualifier = "one of " if isinstance(expected, list) else "" + raise ValueError(with_context(f"{name} tensor stride mismatch: expected {qualifier}{expected}, got {actual_stride}")) if stride_order is not None: - if isinstance(stride_order, tuple): - if tensor_stride_order != stride_order: - error_msg = f"{name} tensor stride order mismatch: expected {stride_order}, got {tensor_stride_order}" - if extra_error_msg: - error_msg += f": {extra_error_msg}" - raise ValueError(error_msg) - elif isinstance(stride_order, list): - if tensor_stride_order not in stride_order: - error_msg = f"{name} tensor stride order mismatch: expected one of {stride_order}, got {tensor_stride_order}" - if extra_error_msg: - error_msg += f": {extra_error_msg}" - raise ValueError(error_msg) - else: - error_msg = f"Expected stride order to be a tuple or list, got {type(stride_order)}" - if extra_error_msg: - error_msg += f": {extra_error_msg}" - raise ValueError(error_msg) - return tensor_stride, tensor_stride_order - - def _check_dtype( + if not isinstance(stride_order, (tuple, list)): + raise ValueError(with_context(f"Expected stride order to be a tuple or list, got {type(stride_order)}")) + expected_order = tuple(stride_order) if isinstance(stride_order, tuple) else [tuple(item) for item in stride_order] + order_matches = actual_order in expected_order if isinstance(expected_order, list) else actual_order == expected_order + if not order_matches: + qualifier = "one of " if isinstance(expected_order, list) else "" + raise ValueError(with_context(f"{name} tensor stride order mismatch: expected {qualifier}{expected_order}, got {actual_order}")) + return actual_stride, actual_order + + def check_dtype( self, - tensor_or_dtype: torch.Tensor | TensorDesc | torch.dtype, - dtype: torch.dtype | List[torch.dtype], + tensor_or_dtype: TensorDesc | Any, + expected: Any | Collection[Any], name: str = "", extra_error_msg: str = "", - ) -> Optional[torch.dtype]: - """Check if the dtype of a tensor or dtype matches the expected dtype(s). - - :param tensor_or_dtype: The tensor to get dtype from or the dtype to check - :type tensor_or_dtype: torch.Tensor | TensorDesc | torch.dtype - :param dtype: The expected dtype(s) - :type dtype: torch.dtype | List[torch.dtype] - :param name: Logical tensor name for logging - :type name: str - :raises ValueError: If the dtype of the tensor does not match the expected dtype(s) - :return: The dtype of the tensor - :rtype: Optional[torch.dtype] - """ - if tensor_or_dtype is None: - return None - tensor_dtype = tensor_or_dtype.dtype if isinstance(tensor_or_dtype, (torch.Tensor, TensorDesc)) else tensor_or_dtype - if isinstance(dtype, torch.dtype): - if tensor_dtype != dtype: - error_msg = f"{name} dtype mismatch: expected {dtype}, got {tensor_dtype}" - if extra_error_msg: - error_msg += f": {extra_error_msg}" - raise ValueError(error_msg) - elif isinstance(dtype, list): - if tensor_dtype not in dtype: - error_msg = f"{name} dtype mismatch: expected one of {dtype}, got {tensor_dtype}" - if extra_error_msg: - error_msg += f": {extra_error_msg}" - raise ValueError(error_msg) + ) -> str: + """Validate a dtype through its framework-neutral logical name.""" + + actual = _dtype_name_of(tensor_or_dtype) + if isinstance(expected, Collection) and not isinstance(expected, (str, bytes)): + choices = tuple(_dtype_name_of(item) for item in expected) else: - raise ValueError(f"Expected dtype to be a torch.dtype or list, got {type(dtype)}") - return tensor_dtype - - def _value_error_if(self, condition: bool, error_msg: str) -> None: - """Raise a ValueError if the condition is true. - - :param condition: The condition to check - :type condition: bool - :param error_msg: The error message to raise - :type error_msg: str - :raises ValueError: If the condition is true - """ - if condition: + choices = (_dtype_name_of(expected),) + if actual not in choices: + supported = ", ".join(choices) + error_msg = f"{name} dtype mismatch: expected one of {{{supported}}}, got {actual}" + if extra_error_msg: + error_msg += f": {extra_error_msg}" raise ValueError(error_msg) + return actual - def _not_implemented_error_if(self, condition: bool, error_msg: str) -> None: - """Raise a NotImplementedError if the condition is true. + # Compatibility spellings used by the existing Torch wrappers. + _check_tensor_shape = check_tensor_shape + _check_tensor_stride = check_tensor_stride + _check_dtype = check_dtype - :param condition: The condition to check - :type condition: bool - :param error_msg: The error message to raise - :type error_msg: str - :raises NotImplementedError: If the condition is true - """ + @staticmethod + def value_error_if(condition: bool, error_msg: str) -> None: if condition: - raise NotImplementedError(error_msg) + raise ValueError(error_msg) - def _runtime_error_if(self, condition: bool, error_msg: str) -> None: - """Raise a RuntimeError if the condition is true. + @staticmethod + def not_implemented_error_if(condition: bool, error_msg: str) -> None: + if condition: + raise NotImplementedError(error_msg) - :param condition: The condition to check - :type condition: bool - :param error_msg: The error message to raise - :type error_msg: str - :raises RuntimeError: If the condition is true - """ + @staticmethod + def runtime_error_if(condition: bool, error_msg: str) -> None: if condition: raise RuntimeError(error_msg) - def _make_fake_cute_tensor_like( - self, - tensor: torch.Tensor, - assumed_align: int = 16, - name: str = "", - ) -> cute.Pointer: - """Make a fake tensor like the provided tensor. - :param tensor: The tensor to make a fake tensor like - :type tensor: torch.Tensor - :param assumed_align: The assumed alignment of the tensor - :type assumed_align: int - :param name: Logical tensor name for logging - :type name: str - :return: A fake tensor like the provided tensor - :rtype: cute.Pointer - """ - return self._make_fake_cute_tensor_from_desc( - self._make_tensor_desc(tensor, name=name), - assumed_align=assumed_align, - ) - - def _make_tensor_desc( - self, - tensor: Optional[torch.Tensor], - name: str = "", - interpret_uint8_as_fp4x2: Optional[bool] = None, - ) -> Optional[TensorDesc]: - """Capture logical tensor metadata that is sufficient for validation/compile.""" - if tensor is None: - return None - if interpret_uint8_as_fp4x2 is None: - interpret_uint8_as_fp4x2 = self._interpret_uint8_as_fp4x2 - prev_interpret = self._interpret_uint8_as_fp4x2 - self._interpret_uint8_as_fp4x2 = interpret_uint8_as_fp4x2 - try: - tensor_shape = self._tensor_shape(tensor, name=name) - 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, - 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( - self, - tensor_desc: Optional[TensorDesc], - assumed_align: int = 16, - ) -> Optional[cute.Pointer]: - """Build a fake cute tensor from a descriptor.""" - if tensor_desc is None: - return None - return self._make_fake_cute_tensor( - dtype=tensor_desc.dtype, - shape=tensor_desc.shape, - stride=tensor_desc.stride, - assumed_align=assumed_align, - interpret_uint8_as_fp4x2=tensor_desc.interpret_uint8_as_fp4x2, - ) - - def _make_fake_cute_tensor( - self, - dtype: torch.dtype, - shape: Tuple[int, ...], - stride: Tuple[int, ...], - assumed_align: int = 16, - interpret_uint8_as_fp4x2: Optional[bool] = None, - ) -> cute.Pointer: - """Make a fake tensor. - - :param dtype: The dtype of the tensor - :type dtype: torch.dtype - :param shape: The shape of the tensor - :type shape: Tuple[int, ...] - :param stride: The stride of the tensor - :type stride: Tuple[int, ...] - :param assumed_align: The assumed alignment of the tensor - :type assumed_align: int - :return: A fake tensor - :rtype: cute.Pointer - """ - if interpret_uint8_as_fp4x2 is None: - interpret_uint8_as_fp4x2 = self._interpret_uint8_as_fp4x2 - return cute.runtime.make_fake_tensor( - dtype=_convert_to_cutlass_data_type(dtype, interpret_uint8_as_fp4x2=interpret_uint8_as_fp4x2), - shape=shape, - stride=stride, - assumed_align=assumed_align, - ) - - def _make_fake_cute_compact_tensor( - self, - dtype: torch.dtype, - shape: Tuple[int, ...], - stride_order: Tuple[int, ...], - assumed_align: int = 16, - dynamic_mode: Optional[int] = None, - divisibility: int = 16, - interpret_uint8_as_fp4x2: Optional[bool] = None, - ) -> cute.Pointer: - """Make a fake compact tensor. - :param dtype: The dtype of the tensor - :type dtype: torch.dtype - :param shape: The shape of the tensor - :type shape: Tuple[int, ...] - :param stride_order: The stride order of the tensor - :type stride_order: Tuple[int, ...] - :param assumed_align: The assumed alignment of the tensor - :type assumed_align: int - :return: A fake compact tensor - :rtype: cute.Pointer - """ - if dynamic_mode is not None: - dynamic_dim = cute.sym_int(divisibility=divisibility) - shape = shape[:dynamic_mode] + (dynamic_dim,) + shape[dynamic_mode + 1 :] - - if interpret_uint8_as_fp4x2 is None: - interpret_uint8_as_fp4x2 = self._interpret_uint8_as_fp4x2 - return cute.runtime.make_fake_compact_tensor( - dtype=_convert_to_cutlass_data_type(dtype, interpret_uint8_as_fp4x2=interpret_uint8_as_fp4x2), - shape=shape, - stride_order=stride_order, - assumed_align=assumed_align, - ) + _value_error_if = value_error_if + _not_implemented_error_if = not_implemented_error_if + _runtime_error_if = runtime_error_if + def _ensure_support_checked(self) -> None: + """Run an adapter's support check once before lowering or compilation.""" -class TupleDict(dict): - """A dictionary that supports tuple unpacking. + if self._is_supported: + return + check_support = getattr(self, "check_support", None) + if check_support is None: + raise NotImplementedError(f"{self.__class__.__name__} does not define check_support()") + self._logger.info(f"{self.__class__.__name__}: check_support not previously called, calling now") + if not check_support(): + raise AssertionError("Unsupported configuration") - 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) - """ +class TupleDict(dict): + """Dictionary with value-order iteration and integer indexing.""" - def __init__(self, *args, **kwargs): + def __init__(self, *args: Any, **kwargs: Any) -> None: 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) + return (self[key] for key in self._keys) - def __getitem__(self, key): - """Support both string keys and integer indices.""" + 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") - return super().__getitem__(self._keys[key]) + key = self._keys[key] return super().__getitem__(key) + + +_TORCH_EXPORTS = {"TorchTensorDesc", "ApiBaseTorch", "APIBase"} + + +def __getattr__(name: str) -> Any: + if name not in _TORCH_EXPORTS: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + module = import_module(".api_base_torch", __package__) + value = getattr(module, name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted((*globals(), *_TORCH_EXPORTS)) + + +__all__ = [ + "APIBase", + "ApiBase", + "ApiBaseTorch", + "TensorDesc", + "TorchTensorDesc", + "TupleDict", + "canonical_dtype_name", + "ceil_div", + "dtype_bits", + "is_power_of_2", +] diff --git a/python/cudnn/api_base_torch.py b/python/cudnn/api_base_torch.py new file mode 100644 index 000000000..de5d107ff --- /dev/null +++ b/python/cudnn/api_base_torch.py @@ -0,0 +1,910 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Torch implementation of the FE-OSS API lifecycle.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, List, Tuple, Optional +import cuda.bindings.driver as cuda +import cutlass +import torch + +import cutlass.cute as cute +from .api_base import ApiBase, TensorDesc, TupleDict, ceil_div, is_power_of_2 +from .datatypes import _convert_to_cutlass_data_type + + +@dataclass(frozen=True, kw_only=True) +class TorchTensorDesc(TensorDesc): + """Torch tensor metadata, including observed storage layout and device.""" + + dtype: torch.dtype + shape: Tuple[int, ...] + stride: Tuple[int, ...] + stride_order: Tuple[int, ...] + device: torch.device + interpret_uint8_as_fp4x2: bool = False + name: str = "" + + def __post_init__(self): + raw_dtype_name = str(self.dtype).rsplit(".", 1)[-1].lower() + is_uint8 = self.dtype == getattr(torch, "uint8", None) or raw_dtype_name == "uint8" + is_native_fp4x2 = self.dtype == getattr(torch, "float4_e2m1fn_x2", None) or raw_dtype_name == "float4_e2m1fn_x2" + if (self.interpret_uint8_as_fp4x2 and is_uint8) or is_native_fp4x2: + object.__setattr__(self, "packing", "fp4x2") + super().__post_init__() + device = self.device + if not isinstance(device, torch.device): + try: + device = torch.device(device) + except (TypeError, ValueError, RuntimeError) as exc: + raise TypeError(f"Invalid device for TorchTensorDesc: {self.device!r}") from exc + object.__setattr__(self, "device", device) + + @classmethod + def from_tensor( + cls, + tensor: torch.Tensor, + *, + name: str = "", + interpret_uint8_as_fp4x2: bool = False, + ) -> "TorchTensorDesc": + """Capture logical shape and observed layout from a Torch tensor.""" + + shape = tuple(tensor.shape) + stride = tuple(tensor.stride()) + is_packed_fp4 = tensor.dtype == getattr(torch, "float4_e2m1fn_x2", None) or (interpret_uint8_as_fp4x2 and tensor.dtype == torch.uint8) + if is_packed_fp4: + innermost_dim = next((dim for dim, value in enumerate(stride) if value == 1), None) + if innermost_dim is None: + raise RuntimeError(f"tensor {name} has shape: {shape} stride {stride} – innermost contiguous (stride == 1) dimension not found. ") + shape = tuple(size * 2 if dim == innermost_dim else size for dim, size in enumerate(shape)) + stride = tuple(value if dim == innermost_dim else value * 2 for dim, value in enumerate(stride)) + + stride_order = tuple(dim for dim, _ in sorted(enumerate(stride), key=lambda item: (item[1], shape[item[0]]))) + return cls( + dtype=tensor.dtype, + shape=shape, + stride=stride, + stride_order=stride_order, + device=tensor.device, + interpret_uint8_as_fp4x2=interpret_uint8_as_fp4x2, + name=name, + ) + + @staticmethod + def _normalize_dim(dim: int, ndim: int, *, allow_new_dim: bool = False) -> int: + min_dim = -ndim - (1 if allow_new_dim else 0) + max_dim = ndim if allow_new_dim else ndim - 1 + if dim < min_dim or dim > max_dim: + raise IndexError(f"Dimension out of range (expected to be in range of [{min_dim}, {max_dim}], but got {dim})") + if dim < 0: + dim += ndim + (1 if allow_new_dim else 0) + return dim + + @staticmethod + def _compute_stride_order(shape: Tuple[int, ...], stride: Tuple[int, ...]) -> Tuple[int, ...]: + return tuple(i for i, _ in sorted(enumerate(stride), key=lambda x: (x[1], shape[x[0]]))) + + @staticmethod + def _numel(shape: Tuple[int, ...]) -> int: + numel = 1 + for size in shape: + numel *= size + return numel + + @staticmethod + def _compute_contiguous_stride(shape: Tuple[int, ...]) -> Tuple[int, ...]: + if not shape: + return () + strides = [0] * len(shape) + running = 1 + for i in range(len(shape) - 1, -1, -1): + strides[i] = running + running *= max(shape[i], 1) + return tuple(strides) + + @staticmethod + def _is_contiguous_with_order(shape: Tuple[int, ...], stride: Tuple[int, ...], order: Tuple[int, ...]) -> bool: + expected_stride = 1 + for dim in order: + size = shape[dim] + if size == 1: + continue + if stride[dim] != expected_stride: + return False + expected_stride *= size + return True + + @staticmethod + def _compute_view_stride( + old_shape: Tuple[int, ...], + old_stride: Tuple[int, ...], + new_shape: Tuple[int, ...], + ) -> Optional[Tuple[int, ...]]: + old_numel = TorchTensorDesc._numel(old_shape) + if old_numel == 0: + return TorchTensorDesc._compute_contiguous_stride(new_shape) + + new_stride = [0] * len(new_shape) + view_dim = len(new_shape) - 1 + tensor_numel = 1 + view_numel = 1 + + for tensor_dim in range(len(old_shape) - 1, -1, -1): + tensor_numel *= old_shape[tensor_dim] + is_contiguous_chunk_end = tensor_dim == 0 or ( + old_shape[tensor_dim - 1] != 1 and old_stride[tensor_dim - 1] != tensor_numel * old_stride[tensor_dim] + ) + if is_contiguous_chunk_end: + while view_dim >= 0 and (view_numel < tensor_numel or new_shape[view_dim] == 1): + new_stride[view_dim] = view_numel * old_stride[tensor_dim] + view_numel *= new_shape[view_dim] + view_dim -= 1 + + if view_numel != tensor_numel: + return None + + if tensor_dim > 0: + tensor_numel = 1 + view_numel = 1 + + if view_dim != -1: + return None + return tuple(new_stride) + + def _with_layout(self, shape: Tuple[int, ...], stride: Tuple[int, ...]) -> "TorchTensorDesc": + return TorchTensorDesc( + dtype=self.dtype, + shape=shape, + stride=stride, + stride_order=self._compute_stride_order(shape, stride), + device=self.device, + interpret_uint8_as_fp4x2=self.interpret_uint8_as_fp4x2, + packing=self.packing, + name=self.name, + ) + + def __len__(self) -> int: + if self.ndim == 0: + raise TypeError("len() of a 0-d tensor") + return self.shape[0] + + def size(self, dim: Optional[int] = None) -> torch.Size | int: + if dim is None: + return torch.Size(self.shape) + dim = self._normalize_dim(int(dim), self.ndim) + return self.shape[dim] + + def permute(self, *dims: int | Tuple[int, ...] | List[int]) -> "TensorDesc": + if len(dims) == 1 and isinstance(dims[0], (tuple, list)): + dims = tuple(dims[0]) + dims = tuple(int(d) for d in dims) + if len(dims) != self.ndim: + raise RuntimeError(f"permute(): expected {self.ndim} dims, got {len(dims)}") + dims = tuple(self._normalize_dim(d, self.ndim) for d in dims) + if len(set(dims)) != self.ndim: + raise RuntimeError(f"permute(): dims must be unique, got {dims}") + + new_shape = tuple(self.shape[d] for d in dims) + new_stride = tuple(self.stride[d] for d in dims) + return self._with_layout(new_shape, new_stride) + + def transpose(self, dim0: int, dim1: int) -> "TensorDesc": + dim0 = self._normalize_dim(dim0, self.ndim) + dim1 = self._normalize_dim(dim1, self.ndim) + if dim0 == dim1: + return self + dims = list(range(self.ndim)) + dims[dim0], dims[dim1] = dims[dim1], dims[dim0] + return self.permute(dims) + + def squeeze(self, dim: Optional[int | Tuple[int, ...] | List[int]] = None) -> "TensorDesc": + if dim is None: + keep_dims = [i for i, size in enumerate(self.shape) if size != 1] + elif isinstance(dim, (tuple, list)): + squeeze_dims = tuple(self._normalize_dim(int(d), self.ndim) for d in dim) + if len(set(squeeze_dims)) != len(squeeze_dims): + raise RuntimeError(f"squeeze(): dims must be unique, got {squeeze_dims}") + squeeze_dims = {d for d in squeeze_dims if self.shape[d] == 1} + keep_dims = [i for i in range(self.ndim) if i not in squeeze_dims] + else: + squeeze_dim = self._normalize_dim(int(dim), self.ndim) + if self.shape[squeeze_dim] != 1: + return self + keep_dims = [i for i in range(self.ndim) if i != squeeze_dim] + + new_shape = tuple(self.shape[i] for i in keep_dims) + new_stride = tuple(self.stride[i] for i in keep_dims) + if new_shape == self.shape and new_stride == self.stride: + return self + return self._with_layout(new_shape, new_stride) + + def unsqueeze(self, dim: int) -> "TensorDesc": + dim = self._normalize_dim(dim, self.ndim, allow_new_dim=True) + + if dim >= self.ndim: + inserted_stride = 1 + else: + inserted_stride = self.stride[dim] * self.shape[dim] + + new_shape = self.shape[:dim] + (1,) + self.shape[dim:] + new_stride = self.stride[:dim] + (inserted_stride,) + self.stride[dim:] + return self._with_layout(new_shape, new_stride) + + def is_contiguous(self, memory_format: torch.memory_format = torch.contiguous_format) -> bool: + if memory_format in {torch.contiguous_format, torch.preserve_format}: + if self._numel(self.shape) == 0: + return True + return self._is_contiguous_with_order(self.shape, self.stride, tuple(range(self.ndim - 1, -1, -1))) + if memory_format == torch.channels_last: + if self.ndim != 4: + return False + return self._is_contiguous_with_order(self.shape, self.stride, (1, 3, 2, 0)) + if memory_format == torch.channels_last_3d: + if self.ndim != 5: + return False + return self._is_contiguous_with_order(self.shape, self.stride, (1, 4, 3, 2, 0)) + + raise ValueError(f"Unsupported memory format: {memory_format}") + + def contiguous(self) -> "TensorDesc": + if self.is_contiguous(): + return self + contiguous_stride = self._compute_contiguous_stride(self.shape) + return self._with_layout(self.shape, contiguous_stride) + + def view(self, *shape: int | Tuple[int, ...] | List[int]) -> "TensorDesc": + if len(shape) == 1 and isinstance(shape[0], (tuple, list)): + shape = tuple(shape[0]) + new_shape = tuple(int(s) for s in shape) + + old_numel = self._numel(self.shape) + infer_dim = None + known_numel = 1 + for i, size in enumerate(new_shape): + if size == -1: + if infer_dim is not None: + raise RuntimeError("only one dimension can be inferred") + infer_dim = i + continue + if size < 0: + raise RuntimeError(f"invalid shape dimension {size}") + known_numel *= size + + if infer_dim is not None: + if known_numel == 0 or old_numel % known_numel != 0: + raise RuntimeError(f"shape '{new_shape}' is invalid for input of size {old_numel}") + inferred_size = old_numel // known_numel + new_shape = new_shape[:infer_dim] + (inferred_size,) + new_shape[infer_dim + 1 :] + known_numel *= inferred_size + + if known_numel != old_numel: + raise RuntimeError(f"shape '{new_shape}' is invalid for input of size {old_numel}") + + new_stride = self._compute_view_stride(self.shape, self.stride, new_shape) + if new_stride is None: + raise RuntimeError( + "view size is not compatible with input tensor's size and stride " "(at least one dimension spans across two contiguous subspaces)" + ) + + return self._with_layout(new_shape, new_stride) + + def as_strided( + self, + size: Tuple[int, ...] | List[int], + stride: Tuple[int, ...] | List[int], + storage_offset: int = 0, + ) -> "TensorDesc": + if storage_offset != 0: + raise RuntimeError("TensorDesc.as_strided(): non-zero storage_offset is unsupported") + if not isinstance(size, (tuple, list)) or not isinstance(stride, (tuple, list)): + raise TypeError("TensorDesc.as_strided(): size and stride must be tuple/list") + + size = tuple(int(s) for s in size) + stride = tuple(int(s) for s in stride) + if len(size) != len(stride): + raise RuntimeError(f"TensorDesc.as_strided(): mismatch in length of size ({len(size)}) and stride ({len(stride)})") + if any(s < 0 for s in size): + raise RuntimeError(f"TensorDesc.as_strided(): invalid size, got {size}") + if any(s < 0 for s in stride): + raise RuntimeError(f"TensorDesc.as_strided(): invalid stride, got {stride}") + + return self._with_layout(size, stride) + + +class ApiBaseTorch(ApiBase, ABC): + """Abstract base class for cuDNN API wrappers. + + This class defines the common interface that all API wrapper implementations + should follow, including configuration validation, compilation, and execution. + + Provides common functionality: + - Logging via self._logger + - Support validation tracking via self._is_supported + - Compiled kernel caching via self._compiled_kernel + - Stream management helpers + + Subclasses should implement the abstract methods to provide + API-specific validation logic and execution behavior. + + Example: + >>> class MyKernelAPI(ApiBaseTorch): + ... def __init__(self, sample_input, sample_output, config): + ... super().__init__() + ... self.sample_input = sample_input + ... self.sample_output = sample_output + ... self.config = config + ... self._kernel = MyKernel + ... + ... def check_support(self) -> bool: + ... # Validate inputs and configuration + ... assert self.sample_input.dtype == torch.float32 + ... self._is_supported = True + ... return True + ... + ... def compile(self): + ... self._ensure_support_checked() + ... # Create and compile kernel + ... kernel = self._kernel(self.config) + ... self._compiled_kernel = cute.compile(kernel, ...) + ... + ... def execute(self, input_tensor, output_tensor, current_stream=None): + ... current_stream = self._get_default_stream(current_stream) + ... self._compiled_kernel(input_tensor, output_tensor, current_stream) + """ + + def __init__(self): + """Initialize the API base. + + Sets up: + - self._is_supported: Flag indicating if configuration is validated + - self._kernel: Kernel instance + - self._compiled_kernel: Cache for compiled kernel + - self._logger: Logger instance for this class + """ + super().__init__() + self._kernel = None + self._compiled_kernel = None + self._interpret_uint8_as_fp4x2 = False + + @abstractmethod + def check_support(self) -> bool: + """Check if the current configuration is supported by the kernel. + + This method should validate: + - Input/output tensor shapes and strides + - Data types compatibility + - Hardware capabilities (compute capability, memory, etc.) + - Configuration parameters (tile sizes, cluster shapes, etc.) + + Implementations should set self._is_supported = True if valid. + + :return: True if the configuration is supported + :rtype: bool + :raises AssertionError: If a configuration requirement is not met + + Example: + >>> def check_support(self) -> bool: + ... self._logger.debug("Checking support") + ... assert self.input.dtype in {torch.float16, torch.float32} + ... assert self.input.shape[0] % 16 == 0, "Shape must be 16-aligned" + ... self._is_supported = True + ... return True + """ + pass + + @abstractmethod + def compile(self) -> None: + """Compile the kernel with the current configuration. + + This method should: + 1. Ensure support has been checked (use self._ensure_support_checked()) + 2. Create the underlying kernel implementation and fake cute tensors from the sample tensor descriptors + 3. Compile the kernel using cute.compile() + 4. Cache the compiled kernel in self._compiled_kernel + + :raises AssertionError: If the configuration is not supported + + Example: + >>> def compile(self): + ... self._ensure_support_checked() + ... + ... kernel = self._kernel(self.config) + ... sample_input_cute = self._make_fake_cute_tensor_from_desc(self.sample_input) + ... sample_output_cute = self._make_fake_cute_tensor_from_desc(self.sample_output) + ... + ... self._compiled_kernel = cute.compile( + ... kernel, + ... sample_input_cute, + ... sample_output_cute + ... ) + """ + pass + + @abstractmethod + def execute( + self, + *args, + current_stream: Optional[cuda.CUstream] = None, + **kwargs, + ) -> Any: + """Execute the kernel with the provided inputs. + + This method should execute using the cached compiled kernel. + + :param args: Positional arguments (typically input/output tensors) + :param current_stream: CUDA stream for execution (optional) + :type current_stream: cuda.CUstream or None + :param kwargs: Additional keyword arguments for execution + :return: Execution result (if any) + :raises AssertionError: If compiled kernel is not available + + Example: + >>> def execute(self, input_tensor, output_tensor, current_stream=None): + ... current_stream = self._get_default_stream(current_stream) + ... assert self._compiled_kernel is not None, "Kernel not compiled" + ... self._logger.debug("Executing with compiled kernel") + ... self._compiled_kernel(input_tensor, output_tensor, current_stream) + """ + pass + + def __call__(self, *args, **kwargs) -> Any: + """Convenience method to execute the kernel. + + This is a shorthand for compiling (if needed) and then executing. + + :param args: Positional arguments passed to execute() + :param kwargs: Keyword arguments passed to execute() + :return: Result from execute() + + Example: + >>> api = MyKernelAPI(...) + >>> api.check_support() + >>> api.compile() + >>> api(input_tensor, output_tensor) + """ + if self._compiled_kernel is None: + self.compile() + return self.execute(*args, **kwargs) + + def _get_default_stream(self, stream: Optional[cuda.CUstream]) -> cuda.CUstream: + """Get default CUDA stream if none provided. + + This is a convenience helper to handle optional stream parameters. + If a stream is provided, it is returned as-is. If None, the default + CUDA stream is returned. + + :param stream: CUDA stream or None + :type stream: cuda.CUstream or None + :return: CUDA stream (either provided or default) + :rtype: cuda.CUstream + + Example: + >>> def execute(self, input_tensor, output_tensor, current_stream=None): + ... current_stream = self._get_default_stream(current_stream) + ... # Now current_stream is guaranteed to be a valid stream + """ + if stream is None: + self._logger.debug(f"{self.__class__.__name__}: No CUDA stream provided, using default stream") + return cutlass.cuda.default_stream() + return stream + + def _pad_tensor_to_ndim( + self, + tensor: Optional[torch.Tensor | TensorDesc], + ndim: int, + name: str, + ) -> Optional[torch.Tensor | TensorDesc]: + """Pad a tensor/descriptor by unsqueezing at dim -1 until it reaches ndim rank. + + - If tensor is None, returns None. + - Unsqueezes at dim -1 until tensor/descriptor rank == ndim. + - Logs final reshape for traceability. + + :param tensor: The tensor/descriptor to pad (or None) + :param ndim: Target rank (pad trailing dims until reached) + :param name: Logical tensor name for logging + :return: The padded tensor/descriptor (or None) + """ + if tensor is None: + return None + + if tensor.ndim < ndim: + self._logger.info(f"Padding {name} to {ndim}D from {tensor.shape}") + for _ in range(ndim - tensor.ndim): + tensor = tensor.unsqueeze(-1) + return tensor + + def _unpad_tensor_to_ndim( + self, + tensor: Optional[torch.Tensor | TensorDesc], + ndim: int, + name: str, + ) -> Optional[torch.Tensor | TensorDesc]: + """Unpad a tensor/descriptor by squeezing at dim -1 until it reaches ndim rank. + + - If tensor is None, returns None. + - Squeezes at dim -1 until tensor/descriptor rank == ndim. + - Logs final reshape for traceability. + + :param tensor: The tensor/descriptor to unpad (or None) + :param ndim: Target rank (squeeze trailing dims until reached) + :param name: Logical tensor name for logging + :return: The unpadded tensor/descriptor (or None) + """ + if tensor is None: + return None + + if tensor.ndim > ndim: + self._logger.info(f"Unpadding {name} from {tensor.shape} to {ndim}D") + for _ in range(tensor.ndim - ndim): + if tensor.shape and tensor.shape[-1] == 1: + tensor = tensor.squeeze(-1) + else: + break + + if tensor.ndim != ndim: + self._logger.critical(f"Unpadding {name} resulted in shape {tensor.shape}, expected {ndim}D") + return tensor + + def _is_fp4x2(self, tensor_or_dtype: torch.Tensor | torch.dtype | TensorDesc) -> bool: + """Check if tensor or dtype is an FP4x2 packed datatype. + + :param tensor_or_dtype: The torch tensor or dtype to check + :type tensor_or_dtype: torch.Tensor | torch.dtype + :return: True if tensor/dtype is an FP4x2 packed type + :rtype: bool + """ + if tensor_or_dtype is None: + return False + if isinstance(tensor_or_dtype, TensorDesc): + return tensor_or_dtype.is_fp4 + elif isinstance(tensor_or_dtype, torch.Tensor): + dtype = tensor_or_dtype.dtype + interpret_uint8_as_fp4x2 = self._interpret_uint8_as_fp4x2 + else: + dtype = tensor_or_dtype + interpret_uint8_as_fp4x2 = self._interpret_uint8_as_fp4x2 + return (dtype == getattr(torch, "float4_e2m1fn_x2", None)) or (interpret_uint8_as_fp4x2 and dtype == torch.uint8) + + def _is_fp8(self, tensor_or_dtype: torch.Tensor | torch.dtype | TensorDesc) -> bool: + """Check if tensor or dtype is an FP8 datatype. + + :param tensor_or_dtype: The torch tensor or dtype to check + :type tensor_or_dtype: torch.Tensor | torch.dtype + :return: True if tensor/dtype is an FP8 type + :rtype: bool + """ + if tensor_or_dtype is None: + return False + if isinstance(tensor_or_dtype, TensorDesc): + return tensor_or_dtype.is_fp8 + dtype = tensor_or_dtype.dtype if isinstance(tensor_or_dtype, (torch.Tensor, TensorDesc)) else tensor_or_dtype + return dtype in {torch.float8_e5m2, torch.float8_e4m3fn} + + def _is_f16(self, tensor_or_dtype: torch.Tensor | torch.dtype | TensorDesc) -> bool: + """Check if tensor or dtype is an fp16 or bf16 datatype. + + :param tensor_or_dtype: The torch tensor or dtype to check + :type tensor_or_dtype: torch.Tensor | torch.dtype + :return: True if tensor/dtype is an fp16 or bf16 type + :rtype: bool + """ + if tensor_or_dtype is None: + return False + if isinstance(tensor_or_dtype, TensorDesc): + return tensor_or_dtype.is_f16 + dtype = tensor_or_dtype.dtype if isinstance(tensor_or_dtype, (torch.Tensor, TensorDesc)) else tensor_or_dtype + return dtype in {torch.float16, torch.bfloat16} + + def _get_innermost_stride_dim(self, tensor: torch.Tensor, name: str = "") -> int: + """Return index of innermost contiguous dimension (stride == 1). + + :raises RuntimeError: If no dimension with stride 1 is found. + """ + idx = next((i for i, s in enumerate(tensor.stride()) if s == 1), None) + if idx is None: + self._logger.critical( + f"tensor {name} has shape: {tensor.shape} stride {tensor.stride()} – innermost contiguous (stride == 1) dimension not found. " + ) + raise RuntimeError(f"tensor {name} has shape: {tensor.shape} stride {tensor.stride()} – innermost contiguous (stride == 1) dimension not found. ") + return idx + + def _tensor_shape( + self, + tensor: Optional[torch.Tensor | TensorDesc], + name: str = "", + ) -> Optional[Tuple[int, ...]]: + """Get the logical shape of a tensor, handling FP4x2 packed datatypes. + + For FP4x2 datatypes, two values are packed per byte. The innermost + contiguous dimension (with stride 1) contains packed values, so the + logical shape for that dimension is 2x the physical shape. + + :param tensor: The tensor to get shape from (or None) + :type tensor: torch.Tensor or None + :param name: Logical tensor name for logging + :type name: str + :return: The logical shape tuple (or None if tensor is None) + :rtype: Tuple[int, ...] or None + """ + if tensor is None: + return None + if isinstance(tensor, TensorDesc): + return tensor.shape + + if self._is_fp4x2(tensor): + innermost_dim_index = self._get_innermost_stride_dim(tensor, name=name) + shape = tuple(dim * 2 if i == innermost_dim_index else dim for i, dim in enumerate(tensor.shape)) + self._logger.debug(f"FP4x2 tensor {name}: physical shape {tensor.shape} -> logical shape {shape}") + return shape + else: + return tensor.shape + + def _tensor_stride( + self, + tensor: Optional[torch.Tensor | TensorDesc], + name: str = "", + ) -> Optional[Tuple[int, ...]]: + """Get the logical stride of a tensor, handling FP4x2 packed datatypes. + + For FP4x2 datatypes, two values are packed per byte. The strides must + be adjusted to reflect logical element spacing. All strides are + multiplied by 2 since each physical element contains 2 logical elements. + + :param tensor: The tensor to get stride from (or None) + :type tensor: torch.Tensor or None + :param name: Logical tensor name for logging + :type name: str + :return: The logical stride tuple (or None if tensor is None) + :rtype: Tuple[int, ...] or None + """ + if tensor is None: + return None + if isinstance(tensor, TensorDesc): + return tensor.stride + + if self._is_fp4x2(tensor): + innermost_dim_index = self._get_innermost_stride_dim(tensor, name=name) + strides = tuple(s * 2 if i != innermost_dim_index else s for i, s in enumerate(tensor.stride())) + self._logger.debug(f"FP4x2 tensor {name}: physical stride {tensor.stride()} -> logical stride {strides}") + return strides + else: + return tensor.stride() + + def _check_tensor_shape( + self, + tensor_or_shape: torch.Tensor | TensorDesc | Tuple[int, ...], + shape: Tuple[int, ...] | List[Tuple[int, ...]], + name: str = "", + ) -> Optional[Tuple[int, ...]]: + """Check if the shape of a tensor matches the expected shape(s). + + :param tensor_or_shape: The tensor to get shape from or the shape to check + :type tensor_or_shape: torch.Tensor | TensorDesc | Tuple[int, ...] + :param shape: expected shape or list of expected shapes + :type shape: Tuple[int, ...] | List[Tuple[int, ...]] + :param name: Logical tensor name for logging + :type name: str + :raises ValueError: If the shape of the tensor does not match the expected shape(s) + :return: The logical shape of the tensor + :rtype: Optional[Tuple[int, ...]] + """ + if tensor_or_shape is None: + return None + tensor_shape = self._tensor_shape(tensor_or_shape, name=name) if isinstance(tensor_or_shape, (torch.Tensor, TensorDesc)) else tensor_or_shape + if not isinstance(shape, (tuple, list)): + raise ValueError(f"Expected shape to be a tuple or list, got {type(shape)}") + return super().check_tensor_shape(tensor_shape, shape, name) + + def _check_tensor_stride( + self, + tensor_or_stride: torch.Tensor | TensorDesc | Tuple[int, ...], + stride: Optional[Tuple[int, ...] | List[Tuple[int, ...]]] = None, + stride_order: Optional[Tuple[int, ...] | List[Tuple[int, ...]]] = None, + name: str = "", + extra_error_msg: str = "", + ) -> Optional[Tuple[Tuple[int, ...], Tuple[int, ...]]]: + """Check if the stride of a tensor matches the expected stride(s) or stride order(s). + + :param tensor_or_stride: The tensor to get stride from or the stride to check + :type tensor_or_stride: torch.Tensor | TensorDesc | Tuple[int, ...] + :param stride: The expected stride(s) + :type stride: Tuple[int, ...] | List[Tuple[int, ...]] + :param stride_order: The expected stride order(s) + :type stride_order: Tuple[int, ...] | List[Tuple[int, ...]] + :param name: Logical tensor name for logging + :type name: str + :param extra_error_msg: Extra error message to add to the error + :type extra_error_msg: str + :raises ValueError: If the stride of the tensor does not match the expected stride order + :return: The stride and stride order of the tensor + :rtype: Optional[Tuple[Tuple[int, ...], Tuple[int, ...]]] + """ + value = self._tensor_stride(tensor_or_stride, name=name) if isinstance(tensor_or_stride, torch.Tensor) else tensor_or_stride + return super().check_tensor_stride( + value, + stride=stride, + stride_order=stride_order, + name=name, + extra_error_msg=extra_error_msg, + ) + + def _check_dtype( + self, + tensor_or_dtype: torch.Tensor | TensorDesc | torch.dtype, + dtype: torch.dtype | List[torch.dtype], + name: str = "", + extra_error_msg: str = "", + ) -> Optional[torch.dtype]: + """Check if the dtype of a tensor or dtype matches the expected dtype(s). + + :param tensor_or_dtype: The tensor to get dtype from or the dtype to check + :type tensor_or_dtype: torch.Tensor | TensorDesc | torch.dtype + :param dtype: The expected dtype(s) + :type dtype: torch.dtype | List[torch.dtype] + :param name: Logical tensor name for logging + :type name: str + :raises ValueError: If the dtype of the tensor does not match the expected dtype(s) + :return: The dtype of the tensor + :rtype: Optional[torch.dtype] + """ + if tensor_or_dtype is None: + return None + tensor_dtype = tensor_or_dtype.dtype if isinstance(tensor_or_dtype, (torch.Tensor, TensorDesc)) else tensor_or_dtype + if not isinstance(dtype, (torch.dtype, list)): + raise ValueError(f"Expected dtype to be a torch.dtype or list, got {type(dtype)}") + super().check_dtype(tensor_dtype, dtype, name, extra_error_msg) + return tensor_dtype + + def _make_fake_cute_tensor_like( + self, + tensor: torch.Tensor, + assumed_align: int = 16, + name: str = "", + ) -> cute.Pointer: + """Make a fake tensor like the provided tensor. + :param tensor: The tensor to make a fake tensor like + :type tensor: torch.Tensor + :param assumed_align: The assumed alignment of the tensor + :type assumed_align: int + :param name: Logical tensor name for logging + :type name: str + :return: A fake tensor like the provided tensor + :rtype: cute.Pointer + """ + return self._make_fake_cute_tensor_from_desc( + self._make_tensor_desc(tensor, name=name), + assumed_align=assumed_align, + ) + + def _make_tensor_desc( + self, + tensor: Optional[torch.Tensor], + name: str = "", + interpret_uint8_as_fp4x2: Optional[bool] = None, + ) -> Optional[TensorDesc]: + """Capture logical tensor metadata that is sufficient for validation/compile.""" + if tensor is None: + return None + if interpret_uint8_as_fp4x2 is None: + interpret_uint8_as_fp4x2 = self._interpret_uint8_as_fp4x2 + return TorchTensorDesc.from_tensor( + tensor, + interpret_uint8_as_fp4x2=interpret_uint8_as_fp4x2, + name=name, + ) + + def make_tensor_desc( + self, + tensor: Optional[torch.Tensor], + name: str = "", + interpret_uint8_as_fp4x2: Optional[bool] = None, + ) -> Optional[TorchTensorDesc]: + """Return the kernel-visible metadata for a Torch tensor.""" + + return self._make_tensor_desc( + tensor, + name=name, + interpret_uint8_as_fp4x2=interpret_uint8_as_fp4x2, + ) + + def _make_fake_cute_tensor_from_desc( + self, + tensor_desc: Optional[TensorDesc], + assumed_align: int = 16, + ) -> Optional[cute.Pointer]: + """Build a fake cute tensor from a descriptor.""" + if tensor_desc is None: + return None + return self._make_fake_cute_tensor( + dtype=tensor_desc.dtype, + shape=tensor_desc.shape, + stride=tensor_desc.stride, + assumed_align=assumed_align, + interpret_uint8_as_fp4x2=tensor_desc.packing == "fp4x2", + ) + + def _make_fake_cute_tensor( + self, + dtype: torch.dtype, + shape: Tuple[int, ...], + stride: Tuple[int, ...], + assumed_align: int = 16, + interpret_uint8_as_fp4x2: Optional[bool] = None, + ) -> cute.Pointer: + """Make a fake tensor. + + :param dtype: The dtype of the tensor + :type dtype: torch.dtype + :param shape: The shape of the tensor + :type shape: Tuple[int, ...] + :param stride: The stride of the tensor + :type stride: Tuple[int, ...] + :param assumed_align: The assumed alignment of the tensor + :type assumed_align: int + :return: A fake tensor + :rtype: cute.Pointer + """ + if interpret_uint8_as_fp4x2 is None: + interpret_uint8_as_fp4x2 = self._interpret_uint8_as_fp4x2 + return cute.runtime.make_fake_tensor( + dtype=_convert_to_cutlass_data_type(dtype, interpret_uint8_as_fp4x2=interpret_uint8_as_fp4x2), + shape=shape, + stride=stride, + assumed_align=assumed_align, + ) + + def _make_fake_cute_compact_tensor( + self, + dtype: torch.dtype, + shape: Tuple[int, ...], + stride_order: Tuple[int, ...], + assumed_align: int = 16, + dynamic_mode: Optional[int] = None, + divisibility: int = 16, + interpret_uint8_as_fp4x2: Optional[bool] = None, + ) -> cute.Pointer: + """Make a fake compact tensor. + :param dtype: The dtype of the tensor + :type dtype: torch.dtype + :param shape: The shape of the tensor + :type shape: Tuple[int, ...] + :param stride_order: The stride order of the tensor + :type stride_order: Tuple[int, ...] + :param assumed_align: The assumed alignment of the tensor + :type assumed_align: int + :return: A fake compact tensor + :rtype: cute.Pointer + """ + if dynamic_mode is not None: + dynamic_dim = cute.sym_int(divisibility=divisibility) + shape = shape[:dynamic_mode] + (dynamic_dim,) + shape[dynamic_mode + 1 :] + + if interpret_uint8_as_fp4x2 is None: + interpret_uint8_as_fp4x2 = self._interpret_uint8_as_fp4x2 + return cute.runtime.make_fake_compact_tensor( + dtype=_convert_to_cutlass_data_type(dtype, interpret_uint8_as_fp4x2=interpret_uint8_as_fp4x2), + shape=shape, + stride_order=stride_order, + assumed_align=assumed_align, + ) + + +# Source-compatible spelling retained for downstream imports and subclasses. +# Internal Torch wrappers use ``ApiBaseTorch`` so their framework is explicit. +APIBase = ApiBaseTorch + + +__all__ = [ + "APIBase", + "ApiBaseTorch", + "TorchTensorDesc", + "TupleDict", + "ceil_div", + "is_power_of_2", +] diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py index fbae12785..7fead752f 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py @@ -1,13 +1,15 @@ -from .api import ( - DenseIndexerBackward, - IndexerBackward, - dense_indexer_backward_wrapper, - indexer_backward_wrapper, -) +"""Lazy API exports for DSA indexer backward.""" + +from ..._operation_api import make_operation_api -__all__ = [ +_API_EXPORTS = ( "DenseIndexerBackward", "IndexerBackward", "dense_indexer_backward_wrapper", "indexer_backward_wrapper", -] +) + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py index 023c0e20a..65721398c 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py @@ -1,4 +1,4 @@ -"""APIBase wrappers for sparse and dense indexer backward. +"""ApiBaseTorch wrappers for sparse and dense indexer backward. Combines backend-specific CuTe-DSL kernels: @@ -17,7 +17,7 @@ import torch import cuda.bindings.driver as cuda -from cudnn.api_base import APIBase, TupleDict +from cudnn.api_base import ApiBaseTorch, TupleDict from cudnn.deepseek_sparse_attention.utils.runtime import ( torch_stream_context as _torch_stream_context, validate_q_causal_offsets, @@ -121,7 +121,7 @@ def _dense_shapes( return False, batch, batch * seqlen_q, batch * seqlen_k, heads, head_dim, seqlen_q, seqlen_k -class IndexerBackward(APIBase): +class IndexerBackward(ApiBaseTorch): """End-to-end indexer backward (3 fused stages). Given the forward-computed ``AttnScore`` (target distribution) and @@ -249,7 +249,7 @@ def execute( ) -class DenseIndexerBackward(APIBase): +class DenseIndexerBackward(ApiBaseTorch): """Dense full-KV indexer backward. Consumes raw dense attention/indexer score tensors and their denominators, @@ -384,12 +384,21 @@ def execute( grad_scale = float(loss_coeff) * grad_loss_value / max(int(self.normalization_tokens), 1) # Dense backward's dK path uses atomic/bulk reductions into fp32. + # SM100 clears the reduction target in its CuTe launch sequence; + # SM90 still requires a pre-zeroed buffer. + kernel_clears_d_index_k = not self._uses_current_stream_pipeline d_index_k_target = d_index_k if d_index_k.dtype == torch.float32: d_index_k_f32 = d_index_k - d_index_k_f32.zero_() + if not kernel_clears_d_index_k: + d_index_k_f32.zero_() else: - d_index_k_f32 = torch.zeros_like(d_index_k, dtype=torch.float32) + allocation = ( + torch.empty_like + if kernel_clears_d_index_k + else torch.zeros_like + ) + d_index_k_f32 = allocation(d_index_k, dtype=torch.float32) self._compiled_kernel( index_q, 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..0698d8bba 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 @@ -1,13 +1,15 @@ """ -Indexer Backward — Dense Mode SM100 CuTe-DSL, 3-kernel design. +Indexer Backward — Dense Mode SM100 CuTe-DSL. -Three kernels launched sequentially on the same stream: +GPU work is launched sequentially on the same stream: Kernel 1 (CuTe DSL): ScoreGradDense — scalar elementwise + reduction kernel. From raw score + denom, normalize → softmax backprop reduction → compute - grad_signal, in-place overwrite PredictRaw buffer. + grad_signal into a caller-provided output buffer. Grid = (max_seqlen_q, B, 1), Block = (128, 1, 1), 4 warps. + dK clear (CuTe DSL): initializes the FP32 bulk-reduction target. + Kernel 2 (CuTe DSL): DenseIndexerBackward2QGemmSm100 — warp-specialized GEMM kernel. Three GEMMs (S, dK, dQ) with elementwise dS/dW computation. Two Q tokens per CTA (K-reuse): 6 GEMMs / K block. @@ -15,7 +17,7 @@ dK accumulated in float32 via cp.reduce.async.bulk. Grid = (ceil(max_seqlen_q/2), B, 1), Block = (384, 1, 1), 12 warps. - Kernel 3 (PyTorch): dk_convert — cast dK from float32 to output dtype. + Framework cast: convert dK from FP32 to the public output dtype. Kernel 2 — Warp specialization (12 warps, 384 threads): Warp 0: Load (TMA Q×2 + TMA K 2-stage + per-block grad_signal + W) @@ -53,7 +55,6 @@ import math from functools import partial -import torch import cuda.bindings.driver as cuda import cutlass @@ -73,12 +74,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") @@ -113,7 +109,7 @@ class ScoreGradDense: (dL/dlog_predict under clipped log), accumulate local_sum. Cross-warp reduce → sum_grad. Phase 2: Re-traverse S_k, recompute predict/g (L1 hit), write - grad_signal = g - predict * sum_grad in-place to PredictRaw. + grad_signal = g - predict * sum_grad to the output tensor. SMEM: only 16 bytes (4 × fp32 cross-warp reduction scratch). """ @@ -141,12 +137,14 @@ def __init__( def __call__( self, mPredictRaw, + mGradSignal, mTargetRaw, mLSE, mDenomTarget, mCuSeqlensQ, mCuSeqlensK, mQCausalOffsets, + mGradLoss, grad_scale: Float32, max_seqlen_q: Int32, max_seqlen_k: Int32, @@ -167,6 +165,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 +188,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 +210,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 +258,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 +272,11 @@ 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 +297,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) @@ -296,7 +309,7 @@ class SharedStorage: cute.arch.sync_threads() sum_grad = sReduceScratch[0] + sReduceScratch[1] + sReduceScratch[2] + sReduceScratch[3] - # --- Phase 2: Write grad_signal in-place to PredictRaw --- + # --- Phase 2: Write grad_signal to the caller-provided output --- # Padding columns (pos >= seqlen_k_b in THD, or pos >= col_limit) get 0. # Kernel 2 processes q tokens in reversed 2Q pairs and uses the # larger q token's K-block bound. For the smaller q token, clear up @@ -324,14 +337,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) @@ -456,6 +469,20 @@ def __call__( self.q_dtype = mQ.element_type self.k_dtype = mK.element_type + # dK is reduced from every query CTA with cp.reduce.async.bulk. Clear + # it in this launch sequence so callers do not need to provide an + # initialized buffer or rely on input/output alias initialization. + num_dk_elements = cute.size(mdK_f32) + mdK_flat = cute.make_tensor( + mdK_f32.iterator, + cute.make_layout(num_dk_elements), + ) + self.clear_dk(mdK_flat).launch( + grid=(cute.ceil_div(num_dk_elements, 256), 1, 1), + block=[256, 1, 1], + stream=stream, + ) + # Layout transposes — sequence dim leading for TMA / per-batch views: # BSHD Q (B,S_q,H,D) -> (H, D, S_q, B) mode [2,3,1,0] # THD Q (T_q,H,D) -> (H, D, T_q) mode [1,2,0] @@ -639,6 +666,16 @@ def __call__( min_blocks_per_mp=1, ) + @cute.kernel + def clear_dk(self, mdK_flat: cute.Tensor): + """Clear the FP32 dK reduction target before the GEMM launch.""" + + tidx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + element_idx = block_idx * Int32(256) + tidx + if element_idx < cute.size(mdK_flat): + mdK_flat[element_idx] = Float32(0.0) + @cute.kernel def kernel_gemm_dense_2q( self, @@ -1981,6 +2018,12 @@ 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 as _resolve_stream, + ) from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor if torch.cuda.get_device_capability()[0] < 10: @@ -2031,12 +2074,14 @@ def _ensure_compiled_score_grad( _score_grad_compile_cache[compile_key] = cute.compile( score_grad_obj, to_cute_tensor(IdxScoreRaw), + to_cute_tensor(IdxScoreRaw), to_cute_tensor(AttnScoreRaw), to_cute_tensor(IdxLSE), to_cute_tensor(AttnL1Norm), cuq_arg, cuk_arg, q_offsets_arg, + None, cutlass.Float32(float(grad_scale)), cutlass.Int32(max_seqlen_q), cutlass.Int32(max_seqlen_k), @@ -2099,7 +2144,7 @@ def _run_gemm_only( QCausalOffsets=None, current_stream=None, ): - """Run only kernel 2 (2Q GEMM). Caller must have run kernel 1 and zeroed dIndexK_f32.""" + """Run the dK clear and 2Q GEMM after the caller has run kernel 1.""" # Match the compile-time is_varlen to avoid dispatching into a kernel # compiled for the other mode. if is_varlen: @@ -2194,6 +2239,7 @@ def _run( ) with torch.cuda.nvtx.range("indexer_backward_dsl_dense_score_grad"): _score_grad_compile_cache[compile_key]( + IdxScoreRaw, IdxScoreRaw, AttnScoreRaw, IdxLSE, @@ -2201,6 +2247,7 @@ def _run( CuSeqlensQ, CuSeqlensK, QCausalOffsets, + None, cutlass.Float32(float(grad_scale)), cutlass.Int32(max_seqlen_q), cutlass.Int32(max_seqlen_k), 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..d2708b809 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,12 @@ def indexer_backward_sm100( class ScoreGradSm100: - """CuTe DSL kernel for in-place score_grad precompute.""" + """CuTe DSL kernel for sparse score-gradient precompute. + + By default the results alias the two score inputs for compatibility with + the Torch pipeline. Functional callers may provide separate gradient and + sum-gradient output tensors. + """ THREADS_PER_CTA = 128 WARP_SIZE = 32 @@ -1470,14 +1470,36 @@ def __call__( mGradLoss: cute.Tensor, grad_scale: Float32 | float, stream: cuda.CUstream, + mGradSignal: Optional[cute.Tensor] = None, + mSumGrad: Optional[cute.Tensor] = None, ): + # Supplying a distinct grad-signal output selects the functional path. + # The sum is not consumed by indexer backward, so avoid materializing it + # unless the caller explicitly requests it. With neither output passed, + # preserve the Torch pipeline's historical in-place behavior. + 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 +1508,16 @@ 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 +1559,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 +1578,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 +1616,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 +1634,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 +1657,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 +1669,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 +1700,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/jax.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py new file mode 100644 index 000000000..602779f39 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py @@ -0,0 +1,758 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX APIs for fixed-shape SM100 sparse and dense indexer backward.""" + +from __future__ import annotations + +from typing import Any + +import jax.numpy as jnp +from cutlass.jax import TensorSpec + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + call_cutedsl, + require_array, +) + + +def as_grad_loss_operand(grad_loss: Any) -> Any: + """Normalize a scalar or one-element FP32 loss gradient to shape ``(1,)``.""" + + if hasattr(grad_loss, "shape") or hasattr(grad_loss, "dtype"): + shape = require_array(grad_loss, name="grad_loss", dtype=jnp.float32) + if shape == (1,): + return grad_loss + if shape == (): + return jnp.reshape(grad_loss, (1,)) + raise ValueError(f"grad_loss must be scalar or shape (1,), got {shape}") + return jnp.asarray((grad_loss,), dtype=jnp.float32) + + +def _launch( + 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, + *, + heads: int, + head_dim: int, + topk: int, + block_i: int, + sm_scale: float, + grad_scale: float, + topk_indices_global: bool, +): + from cutlass import Float32 + + from .indexer_backward_sm100 import IndexerBackwardSm100, ScoreGradSm100 + + score_grad = ScoreGradSm100(topk=topk) + backward = IndexerBackwardSm100( + head_dim=head_dim, + heads=heads, + block_I=block_i, + topk=topk, + topk_indices_global=topk_indices_global, + ) + + score_grad( + attn_score, + index_score, + grad_loss, + Float32(grad_scale), + stream, + grad_signal, + None, + ) + backward( + index_q, + weights, + index_k, + d_index_q, + d_weights, + d_index_k_accum, + grad_signal, + topk_indices, + Float32(sm_scale), + stream, + ) + + +def _launch_dense( + stream, + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + grad_loss, + d_index_q, + d_weights, + d_index_k_accum, + grad_signal, + *, + heads: int, + head_dim: int, + block_i: int, + ratio: int, + sm_scale: float, + grad_scale: float, + max_seqlen_q: int, + max_seqlen_k: int, +): + from cutlass import Float32, Int32 + + from .dense_indexer_backward_sm100 import ( + DenseIndexerBackward2QGemmSm100, + ScoreGradDense, + ) + + score_grad = ScoreGradDense(ratio=ratio, block_I=block_i) + backward = DenseIndexerBackward2QGemmSm100( + head_dim=head_dim, + heads=heads, + block_I=block_i, + ratio=ratio, + ) + + score_grad( + index_score, + grad_signal, + attn_score, + index_lse, + attn_l1norm, + None, + None, + None, + grad_loss, + Float32(grad_scale), + Int32(max_seqlen_q), + Int32(max_seqlen_k), + stream, + ) + backward( + index_q, + weights, + index_k, + d_index_q, + d_weights, + d_index_k_accum, + grad_signal, + None, + None, + None, + Float32(sm_scale), + Int32(max_seqlen_q), + Int32(max_seqlen_k), + stream, + ) + + +def _indexer_backward_impl( + 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 = 128, + topk_indices_global: bool = False, + _validate_only: bool = False, +) -> TupleDict: + """Compute sparse indexer gradients with fixed-shape SM100 kernels. + + Inputs use BSHD shapes ``index_q=(B,S_q,64,128)``, + ``weights=(B,S_q,64)``, and ``index_k=(B,S_k,128)``. Target scores, + predicted scores, and top-K indices use shape ``(B,S_q,topk)``. Floating + model inputs use ``bfloat16``, scores and ``grad_loss`` use ``float32``, + and indices use ``int32``. + + The score-gradient and GEMM stages execute in one custom call. Caller score + tensors remain immutable; a hidden XLA-owned buffer carries the gradient + signal between stages. ``grad_loss`` remains a runtime array operand while + configuration and ``loss_coeff`` are compile-time values. + + This is a standalone backward operation, not a custom VJP. Runtime top-K + indices are trusted to follow the selected local/global index convention. + """ + + q_shape = require_array(index_q, name="index_q", rank=4, dtype=jnp.bfloat16) + k_shape = require_array(index_k, name="index_k", rank=3, dtype=jnp.bfloat16) + topk_shape = require_array( + topk_indices, + name="topk_indices", + rank=3, + dtype=jnp.int32, + ) + + batch, seqlen_q, heads, head_dim = q_shape + k_batch, seqlen_k, k_head_dim = k_shape + dimensions = { + "batch": batch, + "S_q": seqlen_q, + "S_k": seqlen_k, + "heads": heads, + "head_dim": head_dim, + "topk": topk_shape[-1], + } + nonpositive = [f"{name}={value}" for name, value in dimensions.items() if value <= 0] + if nonpositive: + raise ValueError("Indexer-backward dimensions must be positive, got " + ", ".join(nonpositive)) + if heads != 64 or head_dim != 128: + raise ValueError("The JAX SM100 sparse indexer-backward path requires " f"heads=64 and head_dim=128, got {heads} and {head_dim}") + if (k_batch, k_head_dim) != (batch, head_dim): + raise ValueError("index_k batch and head dimensions must match index_q, got " f"{(k_batch, k_head_dim)} and {(batch, head_dim)}") + score_shape = (batch, seqlen_q, topk_shape[-1]) + weights_shape = require_array( + weights, + name="weights", + shape=(batch, seqlen_q, heads), + dtype=jnp.bfloat16, + ) + require_array( + attn_score, + name="attn_score", + shape=score_shape, + dtype=jnp.float32, + ) + require_array( + index_score, + name="index_score", + shape=score_shape, + dtype=jnp.float32, + ) + if topk_shape[:2] != score_shape[:2]: + raise ValueError("topk_indices leading dimensions must match index_q's batch and " f"sequence dimensions {score_shape[:2]}, got {topk_shape[:2]}") + + topk = topk_shape[-1] + if block_I != 128: + raise ValueError(f"block_I must be 128, got {block_I}") + if topk % block_I: + raise ValueError(f"topk ({topk}) must be divisible by block_I ({block_I})") + + grad_loss_operand = as_grad_loss_operand(grad_loss) + sm_scale = float(sm_scale) + grad_scale = float(loss_coeff) / (batch * seqlen_q) + if _validate_only: + return None + + tensor_spec = TensorSpec(divisibility=head_dim) + + d_index_q, d_weights, d_index_k_accum = call_cutedsl( + _launch, + ( + index_q, + weights, + index_k, + attn_score, + index_score, + topk_indices, + grad_loss_operand, + ), + outputs=( + BufferSpec("d_index_q", q_shape, jnp.bfloat16, tensor_spec=tensor_spec), + BufferSpec("d_weights", weights_shape, jnp.bfloat16), + BufferSpec( + "d_index_k_accum", + k_shape, + jnp.float32, + tensor_spec=tensor_spec, + fill_value=0.0, + ), + ), + workspaces=(BufferSpec("grad_signal", score_shape, jnp.float32),), + input_specs=(tensor_spec, None, tensor_spec, None, None, None, None), + static_args={ + "heads": int(heads), + "head_dim": int(head_dim), + "topk": int(topk), + "block_i": int(block_I), + "sm_scale": float(sm_scale), + "grad_scale": float(grad_scale), + "topk_indices_global": bool(topk_indices_global), + }, + ) + return TupleDict( + d_index_q=d_index_q, + d_weights=d_weights, + d_index_k=d_index_k_accum.astype(jnp.bfloat16), + ) + + +def _dense_indexer_backward_impl( + 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 = 128, + ratio: int = 1, + _validate_only: bool = False, +) -> TupleDict: + """Compute fixed-shape dense indexer gradients on SM100. + + This binding covers compact BSHD inputs. Score tensors remain immutable; + an XLA-owned workspace carries the dense score gradient between the two + CuTe launches. ``grad_loss`` is a runtime FP32 operand while the remaining + configuration is static when traced. + """ + + q_shape = require_array(index_q, name="index_q", rank=4, dtype=jnp.bfloat16) + k_shape = require_array(index_k, name="index_k", rank=3, dtype=jnp.bfloat16) + + batch, seqlen_q, heads, head_dim = q_shape + k_batch, seqlen_k, k_head_dim = k_shape + dimensions = { + "batch": batch, + "S_q": seqlen_q, + "S_k": seqlen_k, + "heads": heads, + "head_dim": head_dim, + } + nonpositive = [f"{name}={value}" for name, value in dimensions.items() if value <= 0] + if nonpositive: + raise ValueError("Dense indexer-backward dimensions must be positive, got " + ", ".join(nonpositive)) + if heads != 64 or head_dim != 128: + raise ValueError("The JAX SM100 dense indexer-backward path requires " f"heads=64 and head_dim=128, got {heads} and {head_dim}") + if (k_batch, k_head_dim) != (batch, head_dim): + raise ValueError("index_k batch and head dimensions must match index_q, got " f"{(k_batch, k_head_dim)} and {(batch, head_dim)}") + + score_shape = (batch, seqlen_q, seqlen_k) + denom_shape = (batch, seqlen_q) + weights_shape = require_array( + weights, + name="weights", + shape=(batch, seqlen_q, heads), + dtype=jnp.bfloat16, + ) + require_array( + attn_score, + name="attn_score", + shape=score_shape, + dtype=jnp.float32, + ) + require_array( + attn_l1norm, + name="attn_l1norm", + shape=denom_shape, + dtype=jnp.float32, + ) + require_array( + index_score, + name="index_score", + shape=score_shape, + dtype=jnp.float32, + ) + require_array( + index_lse, + name="index_lse", + shape=denom_shape, + dtype=jnp.float32, + ) + + if block_I != 128: + raise ValueError(f"block_I must be 128, got {block_I}") + if ratio < 1: + raise ValueError(f"ratio must be at least 1, got {ratio}") + + grad_loss_operand = as_grad_loss_operand(grad_loss) + sm_scale = float(sm_scale) + grad_scale = float(loss_coeff) / (batch * seqlen_q) + if _validate_only: + return None + + q_tensor_spec = TensorSpec( + layout=(3, 2, 1, 0), + divisibility=head_dim, + ) + k_tensor_spec = TensorSpec( + layout=(2, 1, 0), + divisibility=head_dim, + ) + d_index_q, d_weights, d_index_k_accum = call_cutedsl( + _launch_dense, + ( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + grad_loss_operand, + ), + outputs=( + BufferSpec( + "d_index_q", + q_shape, + jnp.bfloat16, + tensor_spec=q_tensor_spec, + ), + BufferSpec("d_weights", weights_shape, jnp.bfloat16), + BufferSpec( + "d_index_k_accum", + k_shape, + jnp.float32, + tensor_spec=k_tensor_spec, + ), + ), + workspaces=(BufferSpec("grad_signal", score_shape, jnp.float32),), + input_specs=( + q_tensor_spec, + None, + k_tensor_spec, + None, + None, + None, + None, + None, + ), + static_args={ + "heads": int(heads), + "head_dim": int(head_dim), + "block_i": int(block_I), + "ratio": int(ratio), + "sm_scale": sm_scale, + "grad_scale": grad_scale, + "max_seqlen_q": int(seqlen_q), + "max_seqlen_k": int(seqlen_k), + }, + ) + return TupleDict( + d_index_q=d_index_q, + d_weights=d_weights, + d_index_k=d_index_k_accum.astype(jnp.bfloat16), + ) + + +class IndexerBackward(ApiBaseJax): + """Sample-signature-bound JAX callable for SM100 sparse indexer backward.""" + + 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, + sm_scale: float = 1.0, + loss_coeff: float = 1.0, + sample_grad_loss: Any = 1.0, + block_I: int = 128, + topk_indices_global: bool = False, + ) -> None: + super().__init__() + self.index_q_desc = self.make_tensor_desc(sample_index_q, name="sample_index_q") + self.weights_desc = self.make_tensor_desc(sample_weights, name="sample_weights") + self.index_k_desc = self.make_tensor_desc(sample_index_k, name="sample_index_k") + self.attn_score_desc = self.make_tensor_desc(sample_attn_score, name="sample_attn_score") + self.index_score_desc = self.make_tensor_desc(sample_index_score, name="sample_index_score") + self.topk_indices_desc = self.make_tensor_desc(sample_topk_indices, name="sample_topk_indices") + self.grad_loss_desc = self.make_tensor_desc(as_grad_loss_operand(sample_grad_loss), name="sample_grad_loss") + self.sm_scale = sm_scale + self.loss_coeff = loss_coeff + self.block_I = block_I + self.topk_indices_global = topk_indices_global + + def _check_support(self) -> None: + _indexer_backward_impl( + self.index_q_desc, + self.weights_desc, + self.index_k_desc, + self.attn_score_desc, + self.index_score_desc, + self.topk_indices_desc, + self.sm_scale, + self.loss_coeff, + self.grad_loss_desc, + self.block_I, + self.topk_indices_global, + _validate_only=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: + return super().__call__(index_q, weights, index_k, attn_score, index_score, topk_indices, grad_loss) + + def _call_impl( + self, + index_q: Any, + weights: Any, + index_k: Any, + attn_score: Any, + index_score: Any, + topk_indices: Any, + grad_loss: Any = 1.0, + ) -> TupleDict: + for value, expected, name in ( + (index_q, self.index_q_desc, "index_q"), + (weights, self.weights_desc, "weights"), + (index_k, self.index_k_desc, "index_k"), + (attn_score, self.attn_score_desc, "attn_score"), + (index_score, self.index_score_desc, "index_score"), + (topk_indices, self.topk_indices_desc, "topk_indices"), + ): + self.check_tensor_signature(value, expected, name=name) + grad_loss_operand = as_grad_loss_operand(grad_loss) + self.check_tensor_signature(grad_loss_operand, self.grad_loss_desc, name="grad_loss") + return _indexer_backward_impl( + index_q, + weights, + index_k, + attn_score, + index_score, + topk_indices, + self.sm_scale, + self.loss_coeff, + grad_loss_operand, + self.block_I, + self.topk_indices_global, + ) + + +class DenseIndexerBackward(ApiBaseJax): + """Sample-signature-bound JAX callable for dense SM100 indexer backward.""" + + 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, + sm_scale: float = 1.0, + loss_coeff: float = 1.0, + sample_grad_loss: Any = 1.0, + block_I: int = 128, + ratio: int = 1, + ) -> None: + super().__init__() + self.index_q_desc = self.make_tensor_desc( + sample_index_q, + name="sample_index_q", + ) + self.weights_desc = self.make_tensor_desc( + sample_weights, + name="sample_weights", + ) + self.index_k_desc = self.make_tensor_desc( + sample_index_k, + name="sample_index_k", + ) + self.attn_score_desc = self.make_tensor_desc( + sample_attn_score, + name="sample_attn_score", + ) + self.attn_l1norm_desc = self.make_tensor_desc( + sample_attn_l1norm, + name="sample_attn_l1norm", + ) + self.index_score_desc = self.make_tensor_desc( + sample_index_score, + name="sample_index_score", + ) + self.index_lse_desc = self.make_tensor_desc( + sample_index_lse, + name="sample_index_lse", + ) + self.grad_loss_desc = self.make_tensor_desc( + as_grad_loss_operand(sample_grad_loss), + name="sample_grad_loss", + ) + self.sm_scale = sm_scale + self.loss_coeff = loss_coeff + self.block_I = block_I + self.ratio = ratio + + def _check_support(self) -> None: + _dense_indexer_backward_impl( + self.index_q_desc, + self.weights_desc, + self.index_k_desc, + self.attn_score_desc, + self.attn_l1norm_desc, + self.index_score_desc, + self.index_lse_desc, + self.sm_scale, + self.loss_coeff, + self.grad_loss_desc, + self.block_I, + self.ratio, + _validate_only=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, + ) -> TupleDict: + return super().__call__( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + grad_loss, + ) + + def _call_impl( + 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, + ) -> TupleDict: + for value, expected, name in ( + (index_q, self.index_q_desc, "index_q"), + (weights, self.weights_desc, "weights"), + (index_k, self.index_k_desc, "index_k"), + (attn_score, self.attn_score_desc, "attn_score"), + (attn_l1norm, self.attn_l1norm_desc, "attn_l1norm"), + (index_score, self.index_score_desc, "index_score"), + (index_lse, self.index_lse_desc, "index_lse"), + ): + self.check_tensor_signature(value, expected, name=name) + grad_loss_operand = as_grad_loss_operand(grad_loss) + self.check_tensor_signature( + grad_loss_operand, + self.grad_loss_desc, + name="grad_loss", + ) + return _dense_indexer_backward_impl( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + self.sm_scale, + self.loss_coeff, + grad_loss_operand, + self.block_I, + self.ratio, + ) + + +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 = 128, + topk_indices_global: bool = False, +) -> TupleDict: + """Compute sparse indexer gradients with fixed-shape SM100 kernels.""" + + return IndexerBackward( + index_q, + weights, + index_k, + attn_score, + index_score, + topk_indices, + sm_scale=sm_scale, + loss_coeff=loss_coeff, + sample_grad_loss=grad_loss, + block_I=block_I, + topk_indices_global=topk_indices_global, + )(index_q, weights, index_k, attn_score, index_score, topk_indices, grad_loss) + + +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 = 128, + ratio: int = 1, +) -> TupleDict: + """Compute dense indexer gradients for compact BSHD inputs on SM100.""" + + return DenseIndexerBackward( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + sm_scale=sm_scale, + loss_coeff=loss_coeff, + sample_grad_loss=grad_loss, + block_I=block_I, + ratio=ratio, + )( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + grad_loss, + ) + + +__all__ = [ + "DenseIndexerBackward", + "IndexerBackward", + "dense_indexer_backward_wrapper", + "indexer_backward_wrapper", +] diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py index 4b8df4c14..c030d750a 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py @@ -1,3 +1,15 @@ -from .api import IndexerForward, indexer_forward_wrapper +"""Lazy API exports for the DeepSeek indexer-forward operation. -__all__ = ["IndexerForward", "indexer_forward_wrapper"] +Unqualified symbols always come from the Torch ``api.py``. The JAX API is +available explicitly through the sibling ``jax`` namespace. +""" + +from ..._operation_api import make_operation_api + + +_API_EXPORTS = ("IndexerForward", "indexer_forward_wrapper") + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py index d35e0da8b..be3819b22 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py @@ -1,4 +1,4 @@ -"""APIBase wrapper and dispatcher for indexer forward CuTe DSL score kernels. +"""PyTorch ApiBaseTorch wrapper and dispatcher for indexer-forward score kernels. Produces dense indexer scores Q @ K^T with per-head ReLU, weighted head reduction, and a ratio causal mask. Does NOT fuse top-K — pair with @@ -18,7 +18,7 @@ import cutlass.cute as cute from cutlass.cute.runtime import make_fake_stream -from cudnn.api_base import APIBase, TupleDict +from cudnn.api_base import ApiBaseTorch, TupleDict from cudnn.deepseek_sparse_attention.utils.compiler import compile_options from cudnn.deepseek_sparse_attention.utils.runtime import device_major, resolve_stream @@ -30,8 +30,8 @@ TMA_ALIGN_ELEMS = 4 # FP32 output => seqlen_k padded to multiples of 4 (16 B) -class IndexerForward(APIBase): - """SM100+ APIBase implementation used by ``indexer_forward_wrapper``. +class IndexerForward(ApiBaseTorch): + """SM100+ ApiBaseTorch implementation used by ``indexer_forward_wrapper``. Hopper dispatch uses the direct SM90 wrapper in ``_interface_sm90.py``. """ 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..ba9799264 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py @@ -0,0 +1,311 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX API for the DeepSeek indexer-forward score kernel.""" + +from __future__ import annotations + +from typing import Any, Optional + +import jax.numpy as jnp +from cutlass.jax import TensorSpec + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + call_cutedsl, + require_array, +) + +_TMA_ALIGN_ELEMENTS = 4 + + +def _launch( + stream, + q, + k, + w, + scores, + *, + head_dim: int, + qhead_per_kv_head: int, + ratio: int, + m_block_size: int, + n_block_size: int, + q_stage: int, + kv_stage: int, + num_kv_heads: int, + max_seqlen_q: int, + max_seqlen_k: int, + sm_scale: float, +): + # Load the configuration-specific kernel only when tracing the operation. + from cutlass import Float32, Int32 + + from .indexer_fwd_sm100 import IndexerForwardSm100 + + kernel = IndexerForwardSm100( + head_dim=head_dim, + qhead_per_kvhead=qhead_per_kv_head, + ratio=ratio, + is_varlen=False, + m_block_size=m_block_size, + n_block_size=n_block_size, + q_stage=q_stage, + kv_stage=kv_stage, + ) + + kernel( + q, + k, + w, + scores, + Int32(num_kv_heads), + Int32(max_seqlen_q), + Int32(max_seqlen_k), + Float32(sm_scale), + None, + None, + None, + stream, + ) + + +def _require_supported_config( + *, + m_block_size: int, + n_block_size: int, + q_stage: int, + kv_stage: int, +) -> None: + supported = { + "m_block_size": (m_block_size, 128), + "n_block_size": (n_block_size, 128), + "q_stage": (q_stage, 2), + "kv_stage": (kv_stage, 4), + } + unsupported = [f"{name}={actual} (expected {expected})" for name, (actual, expected) in supported.items() if actual != expected] + if unsupported: + raise ValueError("The JAX indexer-forward API supports only the validated SM100 " "kernel configuration: " + ", ".join(unsupported)) + + +def _indexer_forward_impl( + q: Any, + k: Any, + w: Any, + *, + ratio: int = 4, + qhead_per_kv_head: Optional[int] = None, + m_block_size: int = 128, + n_block_size: int = 128, + q_stage: int = 2, + kv_stage: int = 4, + sm_scale: float = 1.0, + _validate_only: bool = False, +) -> TupleDict: + """Compute fixed-shape BSHD indexer scores with the SM100 CuTe kernel. + + ``q`` and ``k`` must have shapes ``(B, S_q, H_q, 128)`` and + ``(B, S_k, H_kv, 128)``. ``w`` must have shape ``(B, S_q, H_q)``. + All three inputs use ``bfloat16`` and the returned scores use ``float32``. + + This API supports fixed-shape BSHD inputs on SM100 only. + Variable-length THD inputs and the SM90 implementation remain available + only through the existing PyTorch API. All configuration arguments are + compile-time values; close them over a jitted function or mark them static + with :func:`jax.jit`. + + The kernel intentionally skips causal and padded score positions. The + wrapper therefore initializes its XLA-owned output storage to ``-inf`` + before launching the kernel. + """ + + q_shape = require_array(q, name="q", rank=4, dtype=jnp.bfloat16) + k_shape = require_array(k, name="k", rank=4, dtype=jnp.bfloat16) + + batch, seqlen_q, num_query_heads, head_dim = q_shape + k_batch, seqlen_k, num_kv_heads, k_head_dim = k_shape + + dimensions = { + "batch": batch, + "S_q": seqlen_q, + "S_k": seqlen_k, + "H_q": num_query_heads, + "H_kv": num_kv_heads, + "head dimension": head_dim, + } + nonpositive = [f"{name}={value}" for name, value in dimensions.items() if value <= 0] + if nonpositive: + raise ValueError("Indexer-forward dimensions must be positive, got " + ", ".join(nonpositive)) + + 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 != 128: + raise ValueError(f"head dimension must be 128, got {head_dim}") + require_array( + w, + name="w", + shape=(batch, seqlen_q, num_query_heads), + dtype=jnp.bfloat16, + ) + + if ratio < 1: + raise ValueError(f"ratio must be at least 1, got {ratio}") + if seqlen_q > seqlen_k * ratio: + raise ValueError(f"S_q ({seqlen_q}) must be no greater than " f"S_k * ratio ({seqlen_k * ratio})") + + if qhead_per_kv_head is None: + if num_query_heads % num_kv_heads != 0: + raise ValueError(f"H_q ({num_query_heads}) must be divisible by H_kv " f"({num_kv_heads})") + qhead_per_kv_head = num_query_heads // num_kv_heads + if qhead_per_kv_head * num_kv_heads != num_query_heads: + raise ValueError("qhead_per_kv_head * H_kv must equal H_q, got " f"{qhead_per_kv_head} * {num_kv_heads} != {num_query_heads}") + if qhead_per_kv_head not in (32, 64): + raise ValueError("qhead_per_kv_head must be 32 or 64, " f"got {qhead_per_kv_head}") + + _require_supported_config( + m_block_size=m_block_size, + n_block_size=n_block_size, + q_stage=q_stage, + kv_stage=kv_stage, + ) + + seqlen_k_padded = ((seqlen_k + _TMA_ALIGN_ELEMENTS - 1) // _TMA_ALIGN_ELEMENTS) * _TMA_ALIGN_ELEMENTS + if _validate_only: + return None + + # Constrain the initialized physical result to the compact row-major ABI + # expected by the TMA store. CUTLASS infers equivalent defaults for inputs. + scores_spec = TensorSpec( + layout=(2, 1, 0), + mode=(0, 1, 2), + divisibility=(None, None, _TMA_ALIGN_ELEMENTS), + ) + + (scores_padded,) = call_cutedsl( + _launch, + (q, k, w), + outputs=( + BufferSpec( + "scores", + (batch, seqlen_q, seqlen_k_padded), + jnp.float32, + tensor_spec=scores_spec, + fill_value=float("-inf"), + ), + ), + static_args={ + "head_dim": int(head_dim), + "qhead_per_kv_head": int(qhead_per_kv_head), + "ratio": int(ratio), + "m_block_size": int(m_block_size), + "n_block_size": int(n_block_size), + "q_stage": int(q_stage), + "kv_stage": int(kv_stage), + "num_kv_heads": int(num_kv_heads), + "max_seqlen_q": int(seqlen_q), + "max_seqlen_k": int(seqlen_k), + "sm_scale": float(sm_scale), + }, + ) + return TupleDict(scores=scores_padded[..., :seqlen_k]) + + +class IndexerForward(ApiBaseJax): + """Sample-signature-bound JAX callable for SM100 indexer forward.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_w: Any, + *, + ratio: int = 4, + qhead_per_kv_head: Optional[int] = None, + m_block_size: int = 128, + n_block_size: int = 128, + q_stage: int = 2, + kv_stage: int = 4, + sm_scale: float = 1.0, + ) -> None: + super().__init__() + self.q_desc = self.make_tensor_desc(sample_q, name="sample_q") + self.k_desc = self.make_tensor_desc(sample_k, name="sample_k") + self.w_desc = self.make_tensor_desc(sample_w, name="sample_w") + self.ratio = ratio + self.qhead_per_kv_head = qhead_per_kv_head + self.m_block_size = m_block_size + self.n_block_size = n_block_size + self.q_stage = q_stage + self.kv_stage = kv_stage + self.sm_scale = sm_scale + + def _check_support(self) -> None: + _indexer_forward_impl( + self.q_desc, + self.k_desc, + self.w_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, + _validate_only=True, + ) + + def __call__(self, q: Any, k: Any, w: Any) -> TupleDict: + return super().__call__(q, k, w) + + def _call_impl(self, q: Any, k: Any, w: Any) -> TupleDict: + self.check_tensor_signature(q, self.q_desc, name="Q") + self.check_tensor_signature(k, self.k_desc, name="K") + self.check_tensor_signature(w, self.w_desc, name="W") + return _indexer_forward_impl( + q, + k, + w, + 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, + ) + + +def indexer_forward_wrapper( + q: Any, + k: Any, + w: Any, + *, + ratio: int = 4, + qhead_per_kv_head: Optional[int] = None, + m_block_size: int = 128, + n_block_size: int = 128, + q_stage: int = 2, + kv_stage: int = 4, + sm_scale: float = 1.0, +) -> TupleDict: + """Compute fixed-shape BSHD indexer scores with the SM100 CuTe kernel.""" + + return IndexerForward( + q, + k, + w, + ratio=ratio, + qhead_per_kv_head=qhead_per_kv_head, + 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, k, w) + + +__all__ = ["IndexerForward", "indexer_forward_wrapper"] 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..49669bb53 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 API exports for the DSA indexer top-K operation. + +Unqualified symbols always come from the Torch ``api.py``. The JAX API is +available explicitly through the sibling ``jax`` namespace. +""" + +from ..._operation_api import make_operation_api + -__all__ = [ +_API_EXPORTS = ( "IndexerTopK", "indexer_top_k_wrapper", "local_to_global_wrapper", "compactify_wrapper", -] +) + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) 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..0bddd1f6d 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/api.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/api.py @@ -1,4 +1,4 @@ -"""APIBase wrapper for the DSA indexer top-K CuTe DSL kernel.""" +"""PyTorch ApiBaseTorch wrapper for the DSA indexer top-K CuTe DSL kernel.""" from __future__ import annotations @@ -7,7 +7,7 @@ import torch import cuda.bindings.driver as cuda -from cudnn.api_base import APIBase, TupleDict +from cudnn.api_base import ApiBaseTorch, TupleDict from .local_to_global_dsl import local_to_global as _local_to_global from .compactify import compactify as _compactify @@ -21,7 +21,7 @@ def _get_cute_dsl_topk_wrapper(): return cute_dsl_topk_wrapper -class IndexerTopK(APIBase): +class IndexerTopK(ApiBaseTorch): """Top-K filter using the SM90+ CuTe-DSL radix kernel. Selects the ``top_k`` largest entries from each row of ``input_values``. @@ -121,7 +121,7 @@ def compile(self) -> None: self._ensure_support_checked() # cute_dsl_topk_wrapper compiles lazily on first call and caches # internally; nothing to do here. Keep the _compiled_kernel sentinel - # non-None so APIBase.__call__() skips a second compile() attempt. + # non-None so ApiBaseTorch.__call__() skips a second compile() attempt. self._compiled_kernel = _get_cute_dsl_topk_wrapper() def execute( 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..912ae8d07 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 + +if TYPE_CHECKING: + import torch -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 - 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,15 @@ 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..1b79b0e61 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 @@ -20,11 +20,8 @@ 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 +570,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 +593,18 @@ def cute_dsl_topk_wrapper( load_balance=False, num_copy_bits=256, ): + # Keep Torch out of the framework-neutral kernel import path. The JAX + # frontend imports ``IndexerTopKKernelVarlenDecode`` from this module. + 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..08ce11545 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py @@ -0,0 +1,443 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX API for the DSA indexer top-K CuTe DSL kernel.""" + +from __future__ import annotations + +from typing import Any + +import jax.numpy as jnp +from cutlass.jax import jax_to_cutlass_dtype + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + as_dtype, + call_cutedsl, + require_array, +) + +_INT32_MAX = (1 << 31) - 1 + + +def _launch( + stream, + input_values, + seq_lens, + output_indices, + output_values, + extra_buffer, + *, + cutlass_dtype: Any, + dtype_bits: int, + num_cols: int, + top_k: int, + next_n: int, + num_copy_bits: int, + large_occupancy: bool, +): + # Load the configuration-specific kernel only when tracing the operation. + from .indexer_top_k_decode_varlen import ( + IndexerTopKKernelVarlenDecode, + _bucket_num_cols, + ) + + kernel = IndexerTopKKernelVarlenDecode( + cutlass_dtype, + _bucket_num_cols(num_cols), + top_k, + next_n, + num_copy_bits=num_copy_bits, + return_val=True, + large_occupancy=large_occupancy, + ) + require_supported_top_k_output( + top_k=top_k, + num_threads_per_cta=kernel.num_threads_per_cta, + num_copy_bits=num_copy_bits, + dtype_bits=dtype_bits, + ) + + kernel( + input_values, + None, # Only used by the unsupported multi-CTA merge path. + extra_buffer, + None, # Only used by the unsupported persistent scheduler. + seq_lens, + output_indices, + output_values, + stream, + enable_persistent_dynamic_scheduling=False, + min_blocks_per_mp=1, + ) + + +def _launch_local_to_global_fixed( + stream, + local_indices, + global_indices, + *, + seqlen_k: int, +): + from cutlass import Int32 + + from .local_to_global_dsl import LocalToGlobalTopK + + kernel = LocalToGlobalTopK(is_varlen=False) + kernel( + local_indices, + global_indices, + Int32(seqlen_k), + None, + None, + stream, + ) + + +def _launch_local_to_global_varlen( + stream, + local_indices, + cu_seqlens_q, + cu_seqlens_k, + global_indices, + *, + seqlen_k: int, +): + from cutlass import Int32 + + from .local_to_global_dsl import LocalToGlobalTopK + + kernel = LocalToGlobalTopK(is_varlen=True) + kernel( + local_indices, + global_indices, + Int32(seqlen_k), + cu_seqlens_q, + cu_seqlens_k, + stream, + ) + + +def _launch_compactify( + stream, + indices, + compact_indices, + topk_length, + *, + rows: int, + cols: int, +): + from cutlass import Int32 + + from .compactify import CompactifyKernel + + kernel = CompactifyKernel(cols=cols) + kernel(indices, compact_indices, topk_length, Int32(rows), stream) + + +def require_supported_top_k_output( + *, + top_k: int, + num_threads_per_cta: int, + num_copy_bits: int, + dtype_bits: int, +) -> None: + """Validate the output vector width selected by the native kernel.""" + + vector_size = min( + top_k, + (top_k + num_threads_per_cta - 1) // num_threads_per_cta, + num_copy_bits // dtype_bits, + 2, + ) + if top_k % vector_size: + raise ValueError(f"top_k ({top_k}) must be divisible by the selected output vector " f"width ({vector_size}); adjust top_k or num_copy_bits") + + +def _indexer_top_k_impl( + input_values: Any, + seq_lens: Any, + top_k: int, + next_n: int = 1, + return_val: bool = True, + num_copy_bits: int = 256, + _validate_only: bool = False, +) -> TupleDict: + """Select the largest values from each row with variable valid lengths. + + This is the JAX counterpart of the Torch ``indexer_top_k_wrapper`` and is + intended for use inside :func:`jax.jit`. ``input_values`` has shape + ``(num_rows, num_cols)`` and ``seq_lens`` has shape ``(batch_size,)``, with + ``num_rows == batch_size * next_n``. The effective valid length for row + ``r`` is ``seq_lens[r // next_n] - next_n + (r % next_n) + 1``. + Runtime lengths are trusted kernel inputs. Every ``seq_lens[b]`` must + satisfy ``top_k + next_n - 1 <= seq_lens[b] <= num_cols`` so every + staggered row has at least ``top_k`` valid entries and no row reads beyond + ``input_values``. These values are not copied to the host for validation + during tracing. + + Shapes and configuration arguments must be concrete while tracing. The + API supports the kernel's single-launch path and always returns both + indices and values. + """ + + if not return_val: + raise NotImplementedError("The JAX indexer_top_k_wrapper requires return_val=True") + + input_shape = require_array( + input_values, + name="input_values", + rank=2, + dtype=(jnp.float16, jnp.bfloat16, jnp.float32), + ) + seq_lens_shape = require_array( + seq_lens, + name="seq_lens", + rank=1, + dtype=jnp.int32, + ) + + num_rows, num_cols = input_shape + (batch_size,) = 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(f"seq_lens must not be empty, got shape {seq_lens.shape}") + if next_n <= 0: + raise ValueError(f"next_n must be positive, got {next_n}") + if num_rows != batch_size * next_n: + raise ValueError(f"num_rows ({num_rows}) must equal seq_lens.size * next_n " f"({batch_size} * {next_n} = {batch_size * next_n})") + if top_k <= 0 or top_k > min(2048, num_cols): + raise ValueError(f"top_k must be in (0, min(2048, num_cols={num_cols})], got {top_k}") + if num_copy_bits <= 0 or num_copy_bits % 8: + raise ValueError(f"num_copy_bits must be a positive whole-byte width, got {num_copy_bits}") + copy_bytes = num_copy_bits // 8 + if copy_bytes & (copy_bytes - 1): + raise ValueError("num_copy_bits must describe a power-of-two byte alignment, " f"got {num_copy_bits} bits ({copy_bytes} bytes)") + + input_dtype = as_dtype(input_values) + + dtype_bits = input_dtype.itemsize * 8 + if num_copy_bits % dtype_bits != 0: + raise ValueError(f"num_copy_bits ({num_copy_bits}) must be divisible by the " f"input dtype width ({dtype_bits})") + + workspace_buffers = 2 if input_dtype == jnp.dtype(jnp.float32) else 1 + workspace_elements = num_rows * workspace_buffers * num_cols + if workspace_elements > _INT32_MAX: + raise NotImplementedError( + "The JAX indexer_top_k_wrapper does not support the Torch " + "row-chunking fallback used when the int32 workspace contains " + f"more than {_INT32_MAX} elements (requested {workspace_elements})" + ) + if _validate_only: + return None + + output_shape = (num_rows, top_k) + output_indices, output_values = call_cutedsl( + _launch, + (input_values, seq_lens), + outputs=( + BufferSpec("indices", output_shape, jnp.int32), + BufferSpec("values", output_shape, input_dtype), + ), + workspaces=( + BufferSpec( + "extra_buffer", + (num_rows, workspace_buffers, num_cols), + jnp.int32, + ), + ), + static_args={ + "cutlass_dtype": jax_to_cutlass_dtype(input_dtype), + "dtype_bits": int(dtype_bits), + "num_cols": int(num_cols), + "top_k": int(top_k), + "next_n": int(next_n), + "num_copy_bits": int(num_copy_bits), + "large_occupancy": bool(num_rows > 148), + }, + ) + return TupleDict(indices=output_indices, values=output_values) + + +class IndexerTopK(ApiBaseJax): + """Sample-signature-bound JAX callable for the DSA indexer top-K kernel.""" + + 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, + ) -> None: + super().__init__() + self.input_desc = self.make_tensor_desc(sample_input_values, name="sample_input_values") + self.seq_lens_desc = self.make_tensor_desc(sample_seq_lens, name="sample_seq_lens") + self.top_k = top_k + self.next_n = next_n + self.return_val = return_val + self.num_copy_bits = num_copy_bits + + def _check_support(self) -> None: + _indexer_top_k_impl( + self.input_desc, + self.seq_lens_desc, + self.top_k, + self.next_n, + self.return_val, + self.num_copy_bits, + _validate_only=True, + ) + + def __call__(self, input_values: Any, seq_lens: Any) -> TupleDict: + return super().__call__(input_values, seq_lens) + + def _call_impl(self, input_values: Any, seq_lens: Any) -> TupleDict: + self.check_tensor_signature(input_values, self.input_desc, name="input_values") + self.check_tensor_signature(seq_lens, self.seq_lens_desc, name="seq_lens") + return _indexer_top_k_impl( + input_values, + seq_lens, + self.top_k, + self.next_n, + self.return_val, + self.num_copy_bits, + ) + + +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, +) -> TupleDict: + """Select the largest values from each row with variable valid lengths.""" + + return IndexerTopK( + input_values, + seq_lens, + top_k, + next_n=next_n, + return_val=return_val, + num_copy_bits=num_copy_bits, + )(input_values, seq_lens) + + +def local_to_global_wrapper( + local_indices: Any, + seqlen_k: int, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, +) -> TupleDict: + """Convert local top-K indices to the global flattened KV index space. + + Fixed-shape input has shape ``(B, S_q, topk)`` and uses + ``batch_index * seqlen_k`` as its offset. Packed input has shape + ``(total_q, topk)`` and requires matching ``cu_seqlens_q`` and + ``cu_seqlens_k`` int32 arrays. Runtime cumulative lengths are consumed by + the kernel and are not copied to the host while tracing. + """ + + if seqlen_k <= 0: + raise ValueError(f"seqlen_k must be positive, got {seqlen_k}") + + is_varlen = cu_seqlens_q is not None or cu_seqlens_k is not None + if is_varlen: + if cu_seqlens_q is None or cu_seqlens_k is None: + raise ValueError("Packed local-to-global conversion requires both cu_seqlens_q " "and cu_seqlens_k") + local_shape = require_array( + local_indices, + name="local_indices", + rank=2, + dtype=(jnp.int32, jnp.int64), + ) + cu_seqlens_shape = require_array( + cu_seqlens_q, + name="cu_seqlens_q", + rank=1, + dtype=jnp.int32, + ) + require_array( + cu_seqlens_k, + name="cu_seqlens_k", + shape=cu_seqlens_shape, + dtype=jnp.int32, + ) + inputs = (local_indices, cu_seqlens_q, cu_seqlens_k) + else: + local_shape = require_array( + local_indices, + name="local_indices", + rank=3, + dtype=(jnp.int32, jnp.int64), + ) + inputs = (local_indices,) + + launcher = _launch_local_to_global_varlen if is_varlen else _launch_local_to_global_fixed + (global_indices,) = call_cutedsl( + launcher, + inputs, + outputs=( + BufferSpec( + "indices", + local_shape, + jnp.int32, + ), + ), + static_args={"seqlen_k": int(seqlen_k)}, + ) + return TupleDict(indices=global_indices) + + +def compactify_wrapper(indices: Any) -> TupleDict: + """Pack nonnegative indices to the front of each row. + + ``indices`` may have shape ``(rows, topk)`` or ``(B, S_q, topk)``. As in + the Torch API, a rank-3 input is flattened batch-major and the returned + index array has shape ``(B * S_q, topk)``. ``topk_length`` contains the + number of nonnegative entries in every returned row. + """ + + indices_shape = require_array( + indices, + name="indices", + rank=(2, 3), + dtype=jnp.int32, + ) + + cols = indices_shape[-1] + rows = 1 + for extent in indices_shape[:-1]: + rows *= extent + if rows <= 0 or cols <= 0: + raise ValueError(f"indices dimensions must be positive, got {indices_shape}") + + flat_indices = jnp.reshape(indices, (rows, cols)) + compact_indices, topk_length = call_cutedsl( + _launch_compactify, + (flat_indices,), + outputs=( + BufferSpec("indices", (rows, cols), jnp.int32), + BufferSpec("topk_length", (rows,), jnp.int32), + ), + static_args={"rows": int(rows), "cols": int(cols)}, + ) + return TupleDict( + indices=compact_indices, + topk_length=topk_length, + ) + + +__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..56b5aa133 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: @@ -120,6 +117,12 @@ def kernel( def is_available() -> bool: + import torch + + from cudnn.deepseek_sparse_attention.utils.runtime import ( + device_major as _get_device_capability, + ) + # 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. @@ -133,6 +136,14 @@ 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/score_recompute/__init__.py b/python/cudnn/deepseek_sparse_attention/score_recompute/__init__.py index d5434b6ff..908216a17 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/__init__.py @@ -1,15 +1,12 @@ -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 API exports for the DSA score-recompute operations. + +Unqualified symbols always come from the Torch ``api.py``. The JAX API is +available explicitly through the sibling ``jax`` namespace. +""" + +from ..._operation_api import make_operation_api -__all__ = [ +_API_EXPORTS = ( "SparseIndexerScoreRecompute", "sparse_indexer_score_recompute_wrapper", "SparseAttnScoreRecompute", @@ -18,4 +15,9 @@ "dense_indexer_score_recompute_wrapper", "DenseAttnScoreRecompute", "dense_attn_score_recompute_wrapper", -] +) + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) 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..86123ef37 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py @@ -21,6 +21,10 @@ from .sparse_score_recompute_sm100 import SparseScoreRecomputeSm100 from .dense_score_recompute_sm100 import DenseScoreRecomputeSm100 +from .config import ( + dispatch_sparse_attn_tile_params, + 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 +101,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 +308,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() @@ -416,57 +389,6 @@ def _sparse_attn_score_recompute( _sparse_attn_score_recompute.compile_cache = {} -def _dispatch_sparse_attn_tile_params( - head_dim: int, - qhead_per_kv_head: int, - topk: int, - compact: bool = False, -): - """Select (m_block_size, n_block_size, k_block_size) for sparse attention backward. - - Returns (m_block_size, n_block_size, k_block_size) where k_block_size=None - means no head_dim splitting. - - Compact (have_topk_length) and non-compact have different optimal configs: - compact benefits from finer-grained early termination (smaller n, less - k-split overhead), non-compact benefits from n=128 + k-split for better - pipeline utilization. - - 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 - - def sparse_attn_score_recompute( q_attn: torch.Tensor, k_attn: torch.Tensor, @@ -506,7 +428,7 @@ def sparse_attn_score_recompute( topk = topk_indices.shape[2] compact = topk_length is not None - m_block_size, n_block_size, k_block_size = _dispatch_sparse_attn_tile_params( + m_block_size, n_block_size, k_block_size = dispatch_sparse_attn_tile_params( head_dim, qhead_per_kv_head, topk, @@ -520,7 +442,7 @@ def sparse_attn_score_recompute( # Invalid positions (topk_idx=-1) are skipped in K load and masked # in epilogue, preserving correctness. if compact: - nc_params = _dispatch_sparse_attn_tile_params( + nc_params = dispatch_sparse_attn_tile_params( head_dim, qhead_per_kv_head, topk, diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/api.py b/python/cudnn/deepseek_sparse_attention/score_recompute/api.py index b8d044063..131b5e156 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/api.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/api.py @@ -1,10 +1,10 @@ -"""APIBase wrappers for the four DSA score-recompute operations. +"""ApiBaseTorch wrappers for the four DSA score-recompute operations. Wraps the SM100 and SM90 CuTe-DSL score kernels. Each backend provides indexer and attention score variants, giving four public classes. Tile / SMEM dispatch logic and compile caching live in the backend interface -modules. This module adapts those entry points to the APIBase contract. +modules. This module adapts those entry points to the ApiBaseTorch contract. """ from __future__ import annotations @@ -14,7 +14,7 @@ import torch import cuda.bindings.driver as cuda -from cudnn.api_base import APIBase, TupleDict +from cudnn.api_base import ApiBaseTorch, TupleDict from . import _interface_sm100 as _iface_sm100 @@ -23,7 +23,7 @@ # --------------------------------------------------------------------------- -def _check_score_arch(api: APIBase) -> None: +def _check_score_arch(api: ApiBaseTorch) -> None: major, _ = torch.cuda.get_device_capability() api._runtime_error_if( major != 9 and major < 10, @@ -31,8 +31,8 @@ def _check_score_arch(api: APIBase) -> None: ) -class _ScoreRecomputeBase(APIBase): - """Common APIBase shell for score-recompute ops. +class _ScoreRecomputeBase(ApiBaseTorch): + """Common ApiBaseTorch shell for score-recompute ops. ``_interface.py`` owns compile caches keyed per-kernel; ``check_support`` and ``compile`` are thin markers that simply gate invocation. ``execute`` @@ -53,7 +53,7 @@ def compile(self) -> None: def _check_sparse_score_shapes( - api: APIBase, + api: ApiBaseTorch, q_desc, k_desc, aux_desc, @@ -83,7 +83,7 @@ def _check_sparse_score_shapes( def _check_dense_score_shapes( - api: APIBase, + api: ApiBaseTorch, q_desc, k_desc, aux_desc, 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..7fa63d83b --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/config.py @@ -0,0 +1,297 @@ +# 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 :class:`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 :class:`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 attention tile without framework imports. + + ``None`` for ``k_block_size`` means the full padded head dimension. Compact + inputs favor finer-grained early termination; non-compact inputs favor + larger N tiles and K splitting. These choices were tuned on B200. + """ + + m = qhead_per_kv_head + if compact: + # Compact configurations from tests/sweep_tile_params.py --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 + + # Non-compact configurations from tests/sweep_tile_params.py. + 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("m_block_size and n_block_size must be positive, got " f"{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 " f"n_block_size ({n_block_size})") + + 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 + 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 " f"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 + s_k_per_stage = n_block_size * k_block_size_eff * 2 + s_q_size = m_block_size * head_dim_padded * 2 + s_topk_indices_size = topk * 2 * 4 + s_per_head_size = m_block_size * 2 * per_head_element_bytes + fixed_overhead = s_per_head_size + 2048 + + topk_in_smem = True + smem_overhead = s_topk_indices_size + fixed_overhead + kv_stage = min( + 4, + max(1, (_SM100_SMEM_BYTES - s_q_size - smem_overhead) // s_k_per_stage), + ) + total_smem = s_q_size + s_k_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 - s_q_size - smem_overhead) // s_k_per_stage), + ) + total_smem = s_q_size + s_k_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 tile and SMEM choices from concrete input metadata.""" + + 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("qhead_per_kv_head must be 32, 64, or 128 for the SM100 " f"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 = qhead_per_kv_head + n_block_size = 128 + k_block_size = 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 " f"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: + """Select a K tile that divides the padded head dimension and fits SMEM.""" + + 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}, " + f"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: + """Use two query tokens per tile when the minimum K tile fits in SMEM.""" + + head_dim_padded = int(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_kernel_config( + *, + score_type: str, + head_dim: int, + qhead_per_kv_head: int, +) -> DenseScoreKernelConfig: + """Resolve the tuned SM100 dense-score tile from concrete metadata.""" + + 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("qhead_per_kv_head must be 32, 64, or 128 for the SM100 " f"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 = int(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_bytes_per_stage = n_block_size * 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 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_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..5dc22a8a7 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py @@ -0,0 +1,849 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX API for the DSA sparse score-recompute kernels.""" + +from __future__ import annotations + +from typing import Any, Optional + +import jax.numpy as jnp + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + call_cutedsl, + require_array, +) +from .config import ( + resolve_dense_score_kernel_config, + resolve_sparse_score_kernel_config, +) + + +def _launch_sparse_with_topk_length( + stream, + q, + k, + per_head, + topk_indices, + topk_length, + out, + *, + score_type: str, + head_dim: int, + qhead_per_kv_head: int, + topk: int, + m_block_size: int, + n_block_size: int, + k_block_size: int | None, + kv_stage: int, + topk_in_smem: bool, + topk_indices_global: bool, + softmax_scale: float, +): + # Load the architecture-specific kernel only when tracing the operation. + from cutlass import Float32 + + from .sparse_score_recompute_sm100 import SparseScoreRecomputeSm100 + + kernel = SparseScoreRecomputeSm100( + head_dim=head_dim, + qhead_per_kvhead=qhead_per_kv_head, + m_block_size=m_block_size, + n_block_size=n_block_size, + topk=topk, + kv_stage=kv_stage, + score_type=score_type, + have_topk_length=True, + topk_in_smem=topk_in_smem, + k_block_size=k_block_size, + topk_indices_global=topk_indices_global, + ) + + kernel( + q, + k, + per_head, + topk_indices, + out, + topk_length, + Float32(softmax_scale), + stream, + ) + + +def _launch_sparse_without_topk_length( + stream, + q, + k, + per_head, + topk_indices, + out, + topk_length_workspace, + *, + score_type: str, + head_dim: int, + qhead_per_kv_head: int, + topk: int, + m_block_size: int, + n_block_size: int, + k_block_size: int | None, + kv_stage: int, + topk_in_smem: bool, + topk_indices_global: bool, + softmax_scale: float, +): + # Load the architecture-specific kernel only when tracing the operation. + from cutlass import Float32 + + from .sparse_score_recompute_sm100 import SparseScoreRecomputeSm100 + + kernel = SparseScoreRecomputeSm100( + head_dim=head_dim, + qhead_per_kvhead=qhead_per_kv_head, + m_block_size=m_block_size, + n_block_size=n_block_size, + topk=topk, + kv_stage=kv_stage, + score_type=score_type, + have_topk_length=False, + topk_in_smem=topk_in_smem, + k_block_size=k_block_size, + topk_indices_global=topk_indices_global, + ) + kernel( + q, + k, + per_head, + topk_indices, + out, + topk_length_workspace, + Float32(softmax_scale), + stream, + ) + + +def _launch_dense_with_q_causal_offsets( + stream, + q, + k, + per_head, + q_causal_offsets, + out, + denom, + *, + score_type: str, + head_dim: int, + qhead_per_kv_head: int, + ratio: int, + max_seqlen_q: int, + max_seqlen_k: int, + scale: float, +): + from cutlass import Float32, Int32 + + from .dense_score_recompute_sm100 import DenseScoreRecomputeSm100 + + config = resolve_dense_score_kernel_config( + score_type=score_type, + head_dim=head_dim, + qhead_per_kv_head=qhead_per_kv_head, + ) + kernel = DenseScoreRecomputeSm100( + head_dim=head_dim, + qhead_per_kvhead=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=score_type, + ratio=ratio, + is_varlen=False, + ) + kernel( + q, + k, + per_head, + out, + denom, + Float32(scale), + Int32(max_seqlen_q), + Int32(max_seqlen_k), + None, + None, + q_causal_offsets, + stream, + ) + + +def _launch_dense_without_q_causal_offsets( + stream, + q, + k, + per_head, + out, + denom, + *, + score_type: str, + head_dim: int, + qhead_per_kv_head: int, + ratio: int, + max_seqlen_q: int, + max_seqlen_k: int, + scale: float, +): + from cutlass import Float32, Int32 + + from .dense_score_recompute_sm100 import DenseScoreRecomputeSm100 + + config = resolve_dense_score_kernel_config( + score_type=score_type, + head_dim=head_dim, + qhead_per_kv_head=qhead_per_kv_head, + ) + kernel = DenseScoreRecomputeSm100( + head_dim=head_dim, + qhead_per_kvhead=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=score_type, + ratio=ratio, + is_varlen=False, + ) + kernel( + q, + k, + per_head, + out, + denom, + Float32(scale), + Int32(max_seqlen_q), + Int32(max_seqlen_k), + None, + None, + None, + stream, + ) + + +def _dense_score_recompute( + q: Any, + k: Any, + per_head: Any, + *, + score_type: str, + per_head_name: str, + per_head_dtype: Any, + scale: float, + qhead_per_kv_head: Optional[int], + ratio: int, + q_causal_offsets: Any | None, + _validate_only: bool = False, +) -> TupleDict: + q_shape = require_array(q, name="q", rank=4, dtype=jnp.bfloat16) + k_shape = require_array(k, name="k", rank=4, dtype=jnp.bfloat16) + + batch, seqlen_q, num_query_heads, head_dim = q_shape + k_batch, seqlen_k, num_kv_heads, k_head_dim = k_shape + dimensions = { + "batch": batch, + "S_q": seqlen_q, + "S_k": seqlen_k, + "H_q": num_query_heads, + "H_kv": num_kv_heads, + "head dimension": head_dim, + } + nonpositive = [f"{name}={value}" for name, value in dimensions.items() if value <= 0] + if nonpositive: + raise ValueError("Dense score-recompute dimensions must be positive, got " + ", ".join(nonpositive)) + 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("q and k head dimensions must match, got " f"{head_dim} and {k_head_dim}") + require_array( + per_head, + name=per_head_name, + shape=(batch, seqlen_q, num_query_heads), + dtype=per_head_dtype, + ) + if num_query_heads % num_kv_heads: + raise ValueError(f"H_q ({num_query_heads}) must be divisible by H_kv ({num_kv_heads})") + + inferred_qhead_per_kv_head = num_query_heads // num_kv_heads + if qhead_per_kv_head is None: + qhead_per_kv_head = inferred_qhead_per_kv_head + if qhead_per_kv_head != inferred_qhead_per_kv_head: + raise ValueError("qhead_per_kv_head must equal H_q / H_kv, got " f"{qhead_per_kv_head} and {num_query_heads} / {num_kv_heads}") + if ratio < 1: + raise ValueError(f"ratio must be at least 1, got {ratio}") + + inputs = (q, k, per_head) + if q_causal_offsets is not None: + require_array( + q_causal_offsets, + name="q_causal_offsets", + shape=(batch,), + dtype=jnp.int32, + ) + inputs += (q_causal_offsets,) + + resolved_scale = float(scale) + if _validate_only: + return None + + launcher = _launch_dense_with_q_causal_offsets if q_causal_offsets is not None else _launch_dense_without_q_causal_offsets + out, denom = call_cutedsl( + launcher, + inputs, + outputs=( + BufferSpec( + "out", + (batch, seqlen_q, seqlen_k), + jnp.float32, + fill_value=float("-inf"), + ), + BufferSpec("denom", (batch, seqlen_q), jnp.float32), + ), + static_args={ + "score_type": str(score_type), + "head_dim": int(head_dim), + "qhead_per_kv_head": int(qhead_per_kv_head), + "ratio": int(ratio), + "max_seqlen_q": int(seqlen_q), + "max_seqlen_k": int(seqlen_k), + "scale": resolved_scale, + }, + ) + return TupleDict(out=out, denom=denom) + + +def _sparse_score_recompute( + q: Any, + k: Any, + per_head: Any, + topk_indices: Any, + *, + score_type: str, + output_name: str, + per_head_name: str, + per_head_dtype: Any, + softmax_scale: float, + qhead_per_kv_head: Optional[int], + topk_length: Optional[Any], + topk_indices_global: bool, + _validate_only: bool = False, +) -> Any: + q_shape = require_array(q, name="q", rank=4, dtype=jnp.bfloat16) + k_shape = require_array(k, name="k", rank=3, dtype=jnp.bfloat16) + topk_shape = require_array( + topk_indices, + name="topk_indices", + rank=3, + dtype=jnp.int32, + ) + + batch, seqlen_q, num_query_heads, head_dim = q_shape + k_batch, seqlen_k, k_head_dim = k_shape + topk_batch, topk_seqlen_q, topk = topk_shape + dimensions = { + "batch": batch, + "S_q": seqlen_q, + "S_k": seqlen_k, + "H_q": num_query_heads, + "head dimension": head_dim, + "topk": topk, + } + nonpositive = [f"{name}={value}" for name, value in dimensions.items() if value <= 0] + if nonpositive: + raise ValueError("Sparse score-recompute dimensions must be positive, got " + ", ".join(nonpositive)) + 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}") + require_array( + per_head, + name=per_head_name, + shape=(batch, seqlen_q, num_query_heads), + dtype=per_head_dtype, + ) + if (topk_batch, topk_seqlen_q) != (batch, seqlen_q): + raise ValueError( + "topk_indices leading dimensions must match q's batch and sequence " f"dimensions {(batch, seqlen_q)}, got {(topk_batch, topk_seqlen_q)}" + ) + + if qhead_per_kv_head is None: + qhead_per_kv_head = num_query_heads + if qhead_per_kv_head != num_query_heads: + raise ValueError("qhead_per_kv_head must equal H_q for the MQA sparse score kernel, " f"got {qhead_per_kv_head} and H_q={num_query_heads}") + + if topk_length is not None: + require_array( + topk_length, + name="topk_length", + shape=(batch, seqlen_q), + dtype=jnp.int32, + ) + + # Keep an explicit length on the compact path. The Torch-only optimization + # that sometimes drops it assumes every tail index was already set to -1; + # the JAX contract treats topk_length itself as authoritative. + config = resolve_sparse_score_kernel_config( + score_type=score_type, + head_dim=head_dim, + qhead_per_kv_head=qhead_per_kv_head, + topk=topk, + have_topk_length=topk_length is not None, + ) + if _validate_only: + return None + + inputs = (q, k, per_head, topk_indices) + workspaces = () + if topk_length is not None: + inputs += (topk_length,) + else: + # The SM100 kernel has a mandatory tensor argument even when its static + # ``have_topk_length`` mode is false. It never reads this dummy buffer. + workspaces = (BufferSpec("topk_length_workspace", (1, 1), jnp.int32),) + + launcher = _launch_sparse_with_topk_length if config.have_topk_length else _launch_sparse_without_topk_length + (out,) = call_cutedsl( + launcher, + inputs, + outputs=(BufferSpec(output_name, (batch, seqlen_q, topk), jnp.float32),), + workspaces=workspaces, + static_args={ + "score_type": str(score_type), + "head_dim": int(head_dim), + "qhead_per_kv_head": int(qhead_per_kv_head), + "topk": int(topk), + "m_block_size": int(config.m_block_size), + "n_block_size": int(config.n_block_size), + "k_block_size": (None if config.k_block_size is None else int(config.k_block_size)), + "kv_stage": int(config.kv_stage), + "topk_in_smem": bool(config.topk_in_smem), + "topk_indices_global": bool(topk_indices_global), + "softmax_scale": float(softmax_scale), + }, + ) + return out + + +class SparseIndexerScoreRecompute(ApiBaseJax): + """Sample-signature-bound JAX callable for sparse indexer score recompute.""" + + def __init__( + self, + sample_q_indexer: Any, + sample_k_indexer: Any, + sample_weights: Any, + sample_topk_indices: Any, + qhead_per_kv_head: Optional[int] = None, + sample_topk_length: Optional[Any] = None, + topk_indices_global: bool = False, + ) -> None: + super().__init__() + self.q_desc = self.make_tensor_desc(sample_q_indexer, name="sample_q_indexer") + self.k_desc = self.make_tensor_desc(sample_k_indexer, name="sample_k_indexer") + self.weights_desc = self.make_tensor_desc(sample_weights, name="sample_weights") + self.topk_indices_desc = self.make_tensor_desc(sample_topk_indices, name="sample_topk_indices") + self.topk_length_desc = self.make_optional_tensor_desc(sample_topk_length, name="sample_topk_length") + self.qhead_per_kv_head = qhead_per_kv_head + self.topk_indices_global = topk_indices_global + + def _check_support(self) -> None: + _sparse_score_recompute( + self.q_desc, + self.k_desc, + self.weights_desc, + self.topk_indices_desc, + score_type="indexer", + output_name="predict", + per_head_name="weights", + per_head_dtype=jnp.bfloat16, + softmax_scale=1.0, + qhead_per_kv_head=self.qhead_per_kv_head, + topk_length=self.topk_length_desc, + topk_indices_global=self.topk_indices_global, + _validate_only=True, + ) + + def __call__( + self, + q_indexer: Any, + k_indexer: Any, + weights: Any, + topk_indices: Any, + topk_length: Optional[Any] = None, + ) -> TupleDict: + return super().__call__(q_indexer, k_indexer, weights, topk_indices, topk_length) + + def _call_impl( + self, + q_indexer: Any, + k_indexer: Any, + weights: Any, + topk_indices: Any, + topk_length: Optional[Any] = None, + ) -> TupleDict: + for value, expected, name in ( + (q_indexer, self.q_desc, "Q"), + (k_indexer, self.k_desc, "K"), + (weights, self.weights_desc, "weights"), + (topk_indices, self.topk_indices_desc, "topk_indices"), + ): + self.check_tensor_signature(value, expected, name=name) + self.check_optional_tensor_signature(topk_length, self.topk_length_desc, name="topk_length") + predict = _sparse_score_recompute( + q_indexer, + k_indexer, + weights, + topk_indices, + score_type="indexer", + output_name="predict", + per_head_name="weights", + per_head_dtype=jnp.bfloat16, + softmax_scale=1.0, + qhead_per_kv_head=self.qhead_per_kv_head, + topk_length=topk_length, + topk_indices_global=self.topk_indices_global, + ) + return TupleDict(predict=predict) + + +class SparseAttnScoreRecompute(ApiBaseJax): + """Sample-signature-bound JAX callable for sparse attention score recompute.""" + + def __init__( + self, + sample_q_attn: Any, + sample_k_attn: Any, + sample_lse: Any, + sample_topk_indices: Any, + softmax_scale: float, + qhead_per_kv_head: Optional[int] = None, + sample_topk_length: Optional[Any] = None, + topk_indices_global: bool = False, + ) -> None: + super().__init__() + self.q_desc = self.make_tensor_desc(sample_q_attn, name="sample_q_attn") + self.k_desc = self.make_tensor_desc(sample_k_attn, name="sample_k_attn") + self.lse_desc = self.make_tensor_desc(sample_lse, name="sample_lse") + self.topk_indices_desc = self.make_tensor_desc(sample_topk_indices, name="sample_topk_indices") + self.topk_length_desc = self.make_optional_tensor_desc(sample_topk_length, name="sample_topk_length") + self.softmax_scale = softmax_scale + self.qhead_per_kv_head = qhead_per_kv_head + self.topk_indices_global = topk_indices_global + + def _check_support(self) -> None: + _sparse_score_recompute( + self.q_desc, + self.k_desc, + self.lse_desc, + self.topk_indices_desc, + score_type="attention", + output_name="target", + per_head_name="lse", + per_head_dtype=jnp.float32, + softmax_scale=self.softmax_scale, + qhead_per_kv_head=self.qhead_per_kv_head, + topk_length=self.topk_length_desc, + topk_indices_global=self.topk_indices_global, + _validate_only=True, + ) + + def __call__( + self, + q_attn: Any, + k_attn: Any, + lse: Any, + topk_indices: Any, + topk_length: Optional[Any] = None, + ) -> TupleDict: + return super().__call__(q_attn, k_attn, lse, topk_indices, topk_length) + + def _call_impl( + self, + q_attn: Any, + k_attn: Any, + lse: Any, + topk_indices: Any, + topk_length: Optional[Any] = None, + ) -> TupleDict: + for value, expected, name in ( + (q_attn, self.q_desc, "Q"), + (k_attn, self.k_desc, "K"), + (lse, self.lse_desc, "LSE"), + (topk_indices, self.topk_indices_desc, "topk_indices"), + ): + self.check_tensor_signature(value, expected, name=name) + self.check_optional_tensor_signature(topk_length, self.topk_length_desc, name="topk_length") + target = _sparse_score_recompute( + q_attn, + k_attn, + lse, + topk_indices, + score_type="attention", + output_name="target", + per_head_name="lse", + per_head_dtype=jnp.float32, + softmax_scale=self.softmax_scale, + qhead_per_kv_head=self.qhead_per_kv_head, + topk_length=topk_length, + topk_indices_global=self.topk_indices_global, + ) + return TupleDict(target=target) + + +class DenseIndexerScoreRecompute(ApiBaseJax): + """Sample-signature-bound JAX callable for dense indexer score recompute.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_weights: Any, + qhead_per_kv_head: Optional[int] = None, + sm_scale: float = 1.0, + ratio: int = 1, + sample_q_causal_offsets: Any | None = None, + ) -> None: + super().__init__() + self.q_desc = self.make_tensor_desc(sample_q, name="sample_q") + self.k_desc = self.make_tensor_desc(sample_k, name="sample_k") + self.weights_desc = self.make_tensor_desc(sample_weights, name="sample_weights") + self.q_causal_offsets_desc = self.make_optional_tensor_desc(sample_q_causal_offsets, name="sample_q_causal_offsets") + self.qhead_per_kv_head = qhead_per_kv_head + self.sm_scale = sm_scale + self.ratio = ratio + + def _check_support(self) -> None: + _dense_score_recompute( + self.q_desc, + self.k_desc, + self.weights_desc, + score_type="indexer", + per_head_name="weights", + per_head_dtype=jnp.bfloat16, + scale=self.sm_scale, + qhead_per_kv_head=self.qhead_per_kv_head, + ratio=self.ratio, + q_causal_offsets=self.q_causal_offsets_desc, + _validate_only=True, + ) + + def __call__(self, q: Any, k: Any, weights: Any, q_causal_offsets: Any | None = None) -> TupleDict: + return super().__call__(q, k, weights, q_causal_offsets) + + def _call_impl(self, q: Any, k: Any, weights: Any, q_causal_offsets: Any | None = None) -> TupleDict: + self.check_tensor_signature(q, self.q_desc, name="Q") + self.check_tensor_signature(k, self.k_desc, name="K") + self.check_tensor_signature(weights, self.weights_desc, name="weights") + self.check_optional_tensor_signature(q_causal_offsets, self.q_causal_offsets_desc, name="q_causal_offsets") + return _dense_score_recompute( + q, + k, + weights, + score_type="indexer", + per_head_name="weights", + per_head_dtype=jnp.bfloat16, + scale=self.sm_scale, + qhead_per_kv_head=self.qhead_per_kv_head, + ratio=self.ratio, + q_causal_offsets=q_causal_offsets, + ) + + +class DenseAttnScoreRecompute(ApiBaseJax): + """Sample-signature-bound JAX callable for dense attention score recompute.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_lse: Any, + softmax_scale: float, + qhead_per_kv_head: Optional[int] = None, + ratio: int = 1, + sample_q_causal_offsets: Any | None = None, + ) -> None: + super().__init__() + self.q_desc = self.make_tensor_desc(sample_q, name="sample_q") + self.k_desc = self.make_tensor_desc(sample_k, name="sample_k") + self.lse_desc = self.make_tensor_desc(sample_lse, name="sample_lse") + self.q_causal_offsets_desc = self.make_optional_tensor_desc(sample_q_causal_offsets, name="sample_q_causal_offsets") + self.softmax_scale = softmax_scale + self.qhead_per_kv_head = qhead_per_kv_head + self.ratio = ratio + + def _check_support(self) -> None: + _dense_score_recompute( + self.q_desc, + self.k_desc, + self.lse_desc, + score_type="attention", + per_head_name="lse", + per_head_dtype=jnp.float32, + scale=self.softmax_scale, + qhead_per_kv_head=self.qhead_per_kv_head, + ratio=self.ratio, + q_causal_offsets=self.q_causal_offsets_desc, + _validate_only=True, + ) + + def __call__(self, q: Any, k: Any, lse: Any, q_causal_offsets: Any | None = None) -> TupleDict: + return super().__call__(q, k, lse, q_causal_offsets) + + def _call_impl(self, q: Any, k: Any, lse: Any, q_causal_offsets: Any | None = None) -> TupleDict: + self.check_tensor_signature(q, self.q_desc, name="Q") + self.check_tensor_signature(k, self.k_desc, name="K") + self.check_tensor_signature(lse, self.lse_desc, name="LSE") + self.check_optional_tensor_signature(q_causal_offsets, self.q_causal_offsets_desc, name="q_causal_offsets") + return _dense_score_recompute( + q, + k, + lse, + score_type="attention", + per_head_name="lse", + per_head_dtype=jnp.float32, + scale=self.softmax_scale, + qhead_per_kv_head=self.qhead_per_kv_head, + ratio=self.ratio, + q_causal_offsets=q_causal_offsets, + ) + + +def sparse_indexer_score_recompute_wrapper( + q_indexer: Any, + k_indexer: Any, + weights: Any, + topk_indices: Any, + qhead_per_kv_head: Optional[int] = None, + topk_length: Optional[Any] = None, + topk_indices_global: bool = False, +) -> TupleDict: + """Recompute normalized sparse indexer scores with the SM100 kernel. + + Inputs have shapes ``(B, S_q, H_q, D)``, ``(B, S_k, D)``, + ``(B, S_q, H_q)``, and ``(B, S_q, topk)``. The first three tensors use + ``bfloat16`` and ``topk_indices`` uses ``int32``. ``topk_length``, when + present, is an ``int32`` tensor of shape ``(B, S_q)``. The returned + ``predict`` tensor has shape ``(B, S_q, topk)`` and dtype ``float32``. + + This API supports fixed-shape SM100 execution. Configuration and + ``topk_indices_global`` are compile-time values. + """ + + return SparseIndexerScoreRecompute( + q_indexer, + k_indexer, + weights, + topk_indices, + qhead_per_kv_head=qhead_per_kv_head, + sample_topk_length=topk_length, + topk_indices_global=topk_indices_global, + )(q_indexer, k_indexer, weights, topk_indices, topk_length) + + +def sparse_attn_score_recompute_wrapper( + q_attn: Any, + k_attn: Any, + lse: Any, + topk_indices: Any, + softmax_scale: float, + qhead_per_kv_head: Optional[int] = None, + topk_length: Optional[Any] = None, + topk_indices_global: bool = False, +) -> TupleDict: + """Recompute normalized sparse attention scores with the SM100 kernel. + + ``q_attn`` and ``k_attn`` use ``bfloat16``; ``lse`` uses ``float32``; + indices and optional lengths use ``int32``. Shapes match + :func:`sparse_indexer_score_recompute_wrapper`, and the returned ``target`` + is ``(B, S_q, topk)`` ``float32``. ``softmax_scale`` and all configuration + values are compile-time values. + """ + + return SparseAttnScoreRecompute( + q_attn, + k_attn, + lse, + topk_indices, + softmax_scale=softmax_scale, + qhead_per_kv_head=qhead_per_kv_head, + sample_topk_length=topk_length, + topk_indices_global=topk_indices_global, + )(q_attn, k_attn, lse, topk_indices, topk_length) + + +def dense_indexer_score_recompute_wrapper( + q: Any, + k: Any, + weights: Any, + qhead_per_kv_head: Optional[int] = None, + sm_scale: float = 1.0, + ratio: int = 1, + q_causal_offsets: Any | None = None, +) -> TupleDict: + """Compute dense indexer scores and their log-sum-exp denominator. + + This binding currently supports fixed-shape SM100 BSHD tensors. ``q`` and + ``k`` have shapes ``(B, S_q, H_q, D)`` and ``(B, S_k, H_kv, D)``; + ``weights`` has shape ``(B, S_q, H_q)``. All inputs use ``bfloat16``. + The returned ``out`` and ``denom`` arrays use ``float32`` and have shapes + ``(B, S_q, S_k)`` and ``(B, S_q)``. + """ + + return DenseIndexerScoreRecompute( + q, + k, + weights, + qhead_per_kv_head=qhead_per_kv_head, + sm_scale=sm_scale, + ratio=ratio, + sample_q_causal_offsets=q_causal_offsets, + )(q, k, weights, q_causal_offsets) + + +def dense_attn_score_recompute_wrapper( + q: Any, + k: Any, + lse: Any, + softmax_scale: float, + qhead_per_kv_head: Optional[int] = None, + ratio: int = 1, + q_causal_offsets: Any | None = None, +) -> TupleDict: + """Compute dense attention scores and their L1 denominator. + + This binding currently supports fixed-shape SM100 BSHD tensors. ``q`` and + ``k`` use ``bfloat16`` and ``lse`` uses ``float32``. The returned ``out`` + and ``denom`` arrays use ``float32`` and have shapes ``(B, S_q, S_k)`` and + ``(B, S_q)``. + """ + + return DenseAttnScoreRecompute( + q, + k, + lse, + softmax_scale=softmax_scale, + qhead_per_kv_head=qhead_per_kv_head, + ratio=ratio, + sample_q_causal_offsets=q_causal_offsets, + )(q, k, lse, 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/sparse_attention_backward/__init__.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/__init__.py index 039c5be6a..7592cb8cb 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,13 @@ -from .api import SparseAttentionBackward, sparse_attention_backward_wrapper +"""Lazy API exports for DSA sparse-attention backward.""" -__all__ = ["SparseAttentionBackward", "sparse_attention_backward_wrapper"] +from ..._operation_api import make_operation_api + +_API_EXPORTS = ( + "SparseAttentionBackward", + "sparse_attention_backward_wrapper", +) + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) 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..29f518be2 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 @@ -99,7 +99,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 +112,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..4e19a5380 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py @@ -1,4 +1,4 @@ -"""APIBase wrapper for DeepSeek Sparse Attention backward. +"""ApiBaseTorch wrapper for DeepSeek Sparse Attention backward. The wrapper dispatches to the Hopper (SM90) or Blackwell (SM100) CuTe DSL implementation based on the active CUDA device. It consumes the ``out`` and @@ -12,12 +12,12 @@ import torch import cuda.bindings.driver as cuda -from cudnn.api_base import APIBase, TupleDict +from cudnn.api_base import ApiBaseTorch, TupleDict from . import _interface_sm100 as _iface_sm100 -class SparseAttentionBackward(APIBase): +class SparseAttentionBackward(ApiBaseTorch): def __init__( self, sample_q: torch.Tensor, # (total_S_q, H, D) BF16 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..96f4dcbf5 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 callers using the original internal names. + _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]], 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..c00a83264 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py @@ -0,0 +1,425 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for fixed-shape SM100 DSA sparse-attention backward.""" + +from __future__ import annotations + +import math +from typing import Any, Optional + +import jax.numpy as jnp +from cutlass.jax import TensorSpec + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + call_cutedsl, + require_array, +) + + +def _launch_with_topk_length( + stream, + q, + kv, + output, + doutput, + lse, + attn_sink, + topk_idxs, + topk_length, + dq, + dkv, + d_sink, + workspace_lse_odo, + workspace_dkv, + *, + total_seqlen_q: int, + total_seqlen_kv: int, + num_heads: int, + head_dim: int, + block_tile: int, + softmax_scale: float, +): + from cutlass import Float32, Int32 + + from .dsa_bwd_sm100 import FlashAttentionDSABackwardSm100 + + kernel = FlashAttentionDSABackwardSm100( + head_dim=head_dim, + head_dim_v=head_dim, + block_tile=block_tile, + ) + + problem_shape = ( + Int32(total_seqlen_q), + Int32(total_seqlen_kv), + Int32(head_dim), + (Int32(num_heads), Int32(1)), + ) + kernel( + problem_shape, + q, + kv, + output, + doutput, + lse, + attn_sink, + topk_idxs, + topk_length, + dq, + dkv, + d_sink, + workspace_lse_odo, + workspace_dkv, + Float32(softmax_scale), + stream, + ) + + +def _launch_without_topk_length( + stream, + q, + kv, + output, + doutput, + lse, + attn_sink, + topk_idxs, + dq, + dkv, + d_sink, + workspace_lse_odo, + workspace_dkv, + *, + total_seqlen_q: int, + total_seqlen_kv: int, + num_heads: int, + head_dim: int, + block_tile: int, + softmax_scale: float, +): + from cutlass import Float32, Int32 + + from .dsa_bwd_sm100 import FlashAttentionDSABackwardSm100 + + kernel = FlashAttentionDSABackwardSm100( + head_dim=head_dim, + head_dim_v=head_dim, + block_tile=block_tile, + ) + problem_shape = ( + Int32(total_seqlen_q), + Int32(total_seqlen_kv), + Int32(head_dim), + (Int32(num_heads), Int32(1)), + ) + kernel( + problem_shape, + q, + kv, + output, + doutput, + lse, + attn_sink, + topk_idxs, + None, + dq, + dkv, + d_sink, + workspace_lse_odo, + workspace_dkv, + Float32(softmax_scale), + stream, + ) + + +def _sparse_attention_backward_impl( + q: Any, + kv: Any, + out: Any, + dout: Any, + lse: Any, + attn_sink: Any, + topk_idxs: Any, + softmax_scale: Optional[float] = None, + topk_length: Optional[Any] = None, + block_tile: int = 64, + _validate_only: bool = False, +) -> TupleDict: + """Compute fixed-shape DSA sparse-attention gradients on SM100. + + Inputs follow the Torch wrapper's flat MQA contract: Q, O, and dO have + shape ``(S_q, H, 512)``, KV has shape ``(S_kv, 512)``, LSE has shape + ``(S_q, H)``, the attention sink has shape ``(H,)``, and global top-K + indices have shape ``(S_q, topk)``. All floating inputs except LSE and the + sink use ``bfloat16``; LSE and the sink use ``float32``. + + ``topk_length``, when present, is an ``int32`` vector of shape ``(S_q,)``. + Runtime indices must be ``-1`` or in ``[0, S_kv)``. Runtime lengths must be + in ``[1, topk]``; those value constraints are trusted while tracing. + + This is a standalone functional backward operation, not a custom VJP. + SM90 and packed variable-length/batched layouts are not supported. + Configuration values must be static under :func:`jax.jit`. + """ + + q_shape = require_array(q, name="q", rank=3, dtype=jnp.bfloat16) + kv_shape = require_array(kv, name="kv", rank=2, dtype=jnp.bfloat16) + + total_seqlen_q, num_heads, head_dim = q_shape + total_seqlen_kv, kv_head_dim = kv_shape + dimensions = { + "S_q": total_seqlen_q, + "S_kv": total_seqlen_kv, + "H": num_heads, + "D": head_dim, + } + nonpositive = [f"{name}={value}" for name, value in dimensions.items() if value <= 0] + if nonpositive: + raise ValueError("Sparse-attention dimensions must be positive, got " + ", ".join(nonpositive)) + if head_dim != 512 or kv_head_dim != 512: + raise ValueError("The JAX SM100 sparse-attention backward API requires Q and KV head " f"dimensions of 512, got {head_dim} and {kv_head_dim}") + if num_heads % 64: + raise ValueError(f"H must be divisible by 64, got {num_heads}") + if block_tile != 64: + raise ValueError(f"block_tile must be 64, got {block_tile}") + + require_array(out, name="out", shape=q_shape, dtype=jnp.bfloat16) + require_array(dout, name="dout", shape=q_shape, dtype=jnp.bfloat16) + require_array( + lse, + name="lse", + shape=(total_seqlen_q, num_heads), + dtype=jnp.float32, + ) + require_array( + attn_sink, + name="attn_sink", + shape=(num_heads,), + dtype=jnp.float32, + ) + topk_shape = require_array( + topk_idxs, + name="topk_idxs", + rank=2, + dtype=jnp.int32, + ) + if topk_shape[0] != total_seqlen_q: + raise ValueError(f"topk_idxs leading dimension must be S_q ({total_seqlen_q}), got {topk_shape[0]}") + if topk_shape[1] <= 0: + raise ValueError(f"topk_idxs must contain at least one index per row, got {topk_shape}") + + if topk_length is not None: + require_array( + topk_length, + name="topk_length", + shape=(total_seqlen_q,), + dtype=jnp.int32, + ) + + resolved_scale = 1.0 / math.sqrt(head_dim) if softmax_scale is None or softmax_scale == 0.0 else float(softmax_scale) + if _validate_only: + return None + + from cutlass import Float32 + + from .dsa_bwd_sm100 import FlashAttentionDSABackwardSm100 + + workspace_lse_odo_shape = FlashAttentionDSABackwardSm100.get_workspace_size_lse_odo( + total_seqlen_q, + head_dim, + num_heads, + 1, + Float32, + ) + workspace_dkv_shape = FlashAttentionDSABackwardSm100.get_workspace_size_dkv( + total_seqlen_kv, + head_dim, + 1, + Float32, + ) + tensor_spec = TensorSpec(divisibility=head_dim) + + inputs = (q, kv, out, dout, lse, attn_sink, topk_idxs) + input_specs = ( + tensor_spec, + tensor_spec, + tensor_spec, + tensor_spec, + None, + None, + None, + ) + if topk_length is not None: + inputs += (topk_length,) + input_specs += (None,) + + launcher = _launch_with_topk_length if topk_length is not None else _launch_without_topk_length + dq, dkv, d_sink = call_cutedsl( + launcher, + inputs, + outputs=( + BufferSpec("dq", q_shape, jnp.bfloat16, tensor_spec=tensor_spec), + BufferSpec( + "dkv", + (total_seqlen_kv, head_dim), + jnp.bfloat16, + tensor_spec=tensor_spec, + fill_value=0, + ), + BufferSpec("d_sink", (num_heads,), jnp.float32, fill_value=0.0), + ), + workspaces=( + BufferSpec( + "workspace_lse_odo", + workspace_lse_odo_shape, + jnp.uint8, + fill_value=0, + ), + BufferSpec( + "workspace_dkv", + workspace_dkv_shape, + jnp.uint8, + fill_value=0, + ), + ), + input_specs=input_specs, + static_args={ + "total_seqlen_q": int(total_seqlen_q), + "total_seqlen_kv": int(total_seqlen_kv), + "num_heads": int(num_heads), + "head_dim": int(head_dim), + "block_tile": int(block_tile), + "softmax_scale": float(resolved_scale), + }, + ) + return TupleDict(dq=dq, dkv=dkv, d_sink=d_sink) + + +class SparseAttentionBackward(ApiBaseJax): + """Sample-signature-bound JAX callable for SM100 sparse-attention backward.""" + + 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, + softmax_scale: Optional[float] = None, + sample_topk_length: Optional[Any] = None, + block_tile: int = 64, + ) -> None: + super().__init__() + self.q_desc = self.make_tensor_desc(sample_q, name="sample_q") + self.kv_desc = self.make_tensor_desc(sample_kv, name="sample_kv") + self.out_desc = self.make_tensor_desc(sample_out, name="sample_out") + self.dout_desc = self.make_tensor_desc(sample_dout, name="sample_dout") + self.lse_desc = self.make_tensor_desc(sample_lse, name="sample_lse") + 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_optional_tensor_desc(sample_topk_length, name="sample_topk_length") + self.softmax_scale = softmax_scale + self.block_tile = block_tile + + def _check_support(self) -> None: + _sparse_attention_backward_impl( + self.q_desc, + self.kv_desc, + self.out_desc, + self.dout_desc, + self.lse_desc, + self.attn_sink_desc, + self.topk_idxs_desc, + self.softmax_scale, + self.topk_length_desc, + self.block_tile, + _validate_only=True, + ) + + def __call__( + self, + q: Any, + kv: Any, + out: Any, + dout: Any, + lse: Any, + attn_sink: Any, + topk_idxs: Any, + topk_length: Optional[Any] = None, + ) -> TupleDict: + return super().__call__(q, kv, out, dout, lse, attn_sink, topk_idxs, topk_length) + + def _call_impl( + self, + q: Any, + kv: Any, + out: Any, + dout: Any, + lse: Any, + attn_sink: Any, + topk_idxs: Any, + topk_length: Optional[Any] = None, + ) -> TupleDict: + for value, expected, name in ( + (q, self.q_desc, "Q"), + (kv, self.kv_desc, "KV"), + (out, self.out_desc, "O"), + (dout, self.dout_desc, "dO"), + (lse, self.lse_desc, "LSE"), + (attn_sink, self.attn_sink_desc, "attn_sink"), + (topk_idxs, self.topk_idxs_desc, "topk_idxs"), + ): + self.check_tensor_signature(value, expected, name=name) + self.check_optional_tensor_signature(topk_length, self.topk_length_desc, name="topk_length") + return _sparse_attention_backward_impl( + q, + kv, + out, + dout, + lse, + attn_sink, + topk_idxs, + self.softmax_scale, + topk_length, + self.block_tile, + ) + + +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: Optional[Any] = None, + block_tile: int = 64, +) -> TupleDict: + """Compute fixed-shape DSA sparse-attention gradients on SM100.""" + + return SparseAttentionBackward( + q, + kv, + out, + dout, + lse, + attn_sink, + topk_idxs, + softmax_scale=softmax_scale, + sample_topk_length=topk_length, + block_tile=block_tile, + )(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/utils/compiler.py b/python/cudnn/deepseek_sparse_attention/utils/compiler.py index 65b68e81f..20a313621 100644 --- a/python/cudnn/deepseek_sparse_attention/utils/compiler.py +++ b/python/cudnn/deepseek_sparse_attention/utils/compiler.py @@ -1,16 +1,10 @@ """Shared cute.compile option helpers. -The cute DSL compiler accepts ``--gpu-arch `` to lock SASS to a -specific architecture. Without it, the compiler falls back to the device -arch reported by ``torch.cuda.get_device_capability()`` via the cute DSL's -internal map (see ``cutlass/base_dsl/runtime/cuda.py``). That map currently -hardcodes ``(10, 0) → "sm_100a"`` (B200) but treats unknown caps as -``"sm_"`` *without* the architecture-specific ``a`` suffix — -which silently drops sm_X-a-only features (TMA bulk, tcgen05, etc.) on -B300 and beyond. - -So we always pass an explicit ``--gpu-arch`` chosen at runtime from the -device capability. ``compile_options(extra)`` is the single entry point; +The CuTe DSL compiler accepts ``--gpu-arch `` to lock SASS to a +specific architecture. CUTLASS 4.5 auto-detects the known SM90a, SM100a, and +SM103a targets, but the Torch DSA path still passes the target explicitly so +it is visible in compile options/cache keys and unsupported devices fail with +a DSA-specific error. ``compile_options(extra)`` is the single entry point; DSA ``cute.compile`` call sites should route through it. """ diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/api.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/api.py index d04fd85dc..7f515cb24 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/api.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/api.py @@ -48,11 +48,11 @@ from cutlass.cute.runtime import make_fake_stream from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.api_base import ApiBaseTorch, TupleDict, ceil_div, is_power_of_2 from cudnn.discrete_grouped_gemm.discrete_kernel_utils import _require_pointer_tensor -class DiscreteGroupedGemmDswigluSm100(APIBase): +class DiscreteGroupedGemmDswigluSm100(ApiBaseTorch): """API class for discrete-weight grouped GEMM dGLU backward operation on SM100+ GPUs. This kernel performs the backward pass of the discrete-weight grouped GEMM diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/api.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/api.py index 86f39937b..89cee7f6a 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/api.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/api.py @@ -53,11 +53,11 @@ from cutlass.cute.runtime import make_fake_stream from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.api_base import ApiBaseTorch, TupleDict, ceil_div, is_power_of_2 from cudnn.discrete_grouped_gemm.discrete_kernel_utils import _require_pointer_tensor -class DiscreteGroupedGemmSwigluSm100(APIBase): +class DiscreteGroupedGemmSwigluSm100(ApiBaseTorch): """API class for discrete-weight grouped GEMM GLU forward operation on SM100+ GPUs. This kernel performs discrete-weight block-scaled grouped GEMM with GLU diff --git a/python/cudnn/gemm_amax/__init__.py b/python/cudnn/gemm_amax/__init__.py index b65cf491e..7632ffee9 100644 --- a/python/cudnn/gemm_amax/__init__.py +++ b/python/cudnn/gemm_amax/__init__.py @@ -1,9 +1,16 @@ -from .api import ( - GemmAmaxSm100, - gemm_amax_wrapper_sm100, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy API exports for block-scaled dense GEMM + amax.""" + +from .._operation_api import make_operation_api -__all__ = [ +_API_EXPORTS = ( "GemmAmaxSm100", "gemm_amax_wrapper_sm100", -] +) + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/gemm_amax/api.py b/python/cudnn/gemm_amax/api.py index 62fc475bb..6255a480b 100644 --- a/python/cudnn/gemm_amax/api.py +++ b/python/cudnn/gemm_amax/api.py @@ -3,6 +3,7 @@ ) from cuda.bindings import driver as cuda +import logging import os import torch from typing import Tuple, Optional @@ -11,11 +12,13 @@ import cutlass.cute as cute from cutlass.cute.runtime import make_fake_stream -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.api_base import ApiBaseTorch, TupleDict +from cudnn.gemm_validation import resolve_max_active_clusters + +from .validation import validate_gemm_amax -class GemmAmaxSm100(APIBase): +class GemmAmaxSm100(ApiBaseTorch): def __init__( self, sample_a: torch.Tensor, @@ -30,9 +33,11 @@ def __init__( sf_vec_size: int = 32, ): super().__init__() + self._interpret_uint8_as_fp4x2 = True self._warn_experimental_api() self._logger.debug("Entering __init__") + self._kernel = Sm100BlockScaledPersistentDenseGemmKernel self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") @@ -49,11 +54,6 @@ def __init__( 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}" ) @@ -61,192 +61,35 @@ def __init__( 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", - ) - if ab_dtype == torch.uint8: - self._logger.warning("Uint8 ab_dtype will be interpreted as packed fp4, not as native uint8") - - 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}}", - ) - - 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: - 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._logger.debug("Resolving kernel configuration") 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}", + self.mma_tiler_mn = self._kernel.require_mma_tiler(self.mma_tiler_mn) + self.cluster_shape_mn = self._kernel.require_cluster_shape( + self.cluster_shape_mn, + mma_tiler_mn=self.mma_tiler_mn, ) - # Check tensor strides - _ = self._check_tensor_stride( + self._logger.debug("Checking shared tensor and configuration contract") + plan = validate_gemm_amax( 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.sfa_desc, + self.sfb_desc, 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.amax_desc, + acc_dtype=self.acc_dtype, + sf_vec_size=self.sf_vec_size, + supported_sf_vec_sizes=self._kernel.SF_VEC_SIZES, + mma_tiler_mn=self.mma_tiler_mn, ) + self.ab_dtype = self.a_desc.dtype + self.c_dtype = self.c_desc.dtype + self.a_major = plan.a_major + self.b_major = plan.b_major + self.c_major = plan.c_major - 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", - ) + if self.a_desc.dtype == torch.uint8: + self._logger.warning("Uint8 ab_dtype will be interpreted as packed fp4, not as native uint8") self._logger.debug("Checking environment") self._runtime_error_if(not torch.cuda.is_available(), "CUDA is not available") @@ -258,8 +101,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 @@ -277,11 +118,9 @@ def compile(self) -> None: cluster_shape_mn=self.cluster_shape_mn, ) hardware_info = cutlass.utils.HardwareInfo() - max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) - 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 = resolve_max_active_clusters( + hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]), + self.num_cluster_overlap_margin, ) self._logger.debug("Compiling gemm_amax") @@ -366,8 +205,6 @@ def execute( self._logger.debug("Executed with compiled kernel successfully") -import logging - _logger = logging.getLogger(__name__) _cache_of_GemmAmaxSm100Objects = {} @@ -388,13 +225,13 @@ def gemm_amax_wrapper_sm100( _logger.debug("gemm_amax_wrapper_sm100: Creating empty output tensors c and amax") - m, _, l = a_tensor.shape - n, _, l = b_tensor.shape + m, _, batch = a_tensor.shape + n, _, _b_batch = 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, batch), (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, batch), (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) diff --git a/python/cudnn/gemm_amax/dense_blockscaled_gemm_persistent_amax.py b/python/cudnn/gemm_amax/dense_blockscaled_gemm_persistent_amax.py index bf2985c16..9ac676a0e 100644 --- a/python/cudnn/gemm_amax/dense_blockscaled_gemm_persistent_amax.py +++ b/python/cudnn/gemm_amax/dense_blockscaled_gemm_persistent_amax.py @@ -40,6 +40,9 @@ import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.utils.blockscaled_layout as blockscaled_utils +from ..gemm_validation import require_cluster_shape as _require_cluster_shape +from ..gemm_validation import require_mma_tiler as _require_mma_tiler + """ This example provides an experimental implementation of the SM100 batched dense blockscaled GEMM kernel, please note that the APIs and implementation details related to this kernel may change in future releases. @@ -157,6 +160,45 @@ class Sm100BlockScaledPersistentDenseGemmKernel: >>> gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, amax_tensor, max_active_clusters, stream) """ + # Configuration values supported by the FE Torch and JAX wrappers. + MMA_TILER_M = (128, 256) + MMA_TILER_N = (128, 256) + TWO_CTA_MMA_TILER_M = 256 + MAX_CLUSTER_CTAS = 16 + MAX_CLUSTER_DIMENSION = 4 + SF_VEC_SIZES = (16, 32) + KNOWN_HANG_MMA_TILER_M = (256,) + + @classmethod + def require_mma_tiler(cls, mma_tiler_mn: Tuple[int, int]) -> Tuple[int, int]: + """Validate an FE-supported MMA tile.""" + + mma_tiler_mn = _require_mma_tiler( + mma_tiler_mn, + allowed_m=cls.MMA_TILER_M, + allowed_n=cls.MMA_TILER_N, + ) + if mma_tiler_mn[0] in cls.KNOWN_HANG_MMA_TILER_M: + raise NotImplementedError(f"mma_tiler_mn[0] in {cls.KNOWN_HANG_MMA_TILER_M} currently hangs") + return mma_tiler_mn + + @classmethod + def require_cluster_shape( + cls, + cluster_shape_mn: Tuple[int, int], + *, + mma_tiler_mn: Tuple[int, int], + ) -> Tuple[int, int]: + """Validate an FE-supported cluster shape for an MMA tile.""" + + return _require_cluster_shape( + cluster_shape_mn, + mma_m=mma_tiler_mn[0], + two_cta_mma_m=cls.TWO_CTA_MMA_TILER_M, + max_ctas=cls.MAX_CLUSTER_CTAS, + max_dimension=cls.MAX_CLUSTER_DIMENSION, + ) + def __init__( self, sf_vec_size: int, @@ -185,7 +227,7 @@ def __init__( self.acc_dtype = cutlass.Float32 self.sf_vec_size = sf_vec_size - self.use_2cta_instrs = mma_tiler_mn[0] == 256 + self.use_2cta_instrs = mma_tiler_mn[0] == self.TWO_CTA_MMA_TILER_M self.cluster_shape_mn = cluster_shape_mn # K dimension is deferred in _setup_attributes self.mma_tiler = (*mma_tiler_mn, 1) diff --git a/python/cudnn/gemm_amax/jax.py b/python/cudnn/gemm_amax/jax.py new file mode 100644 index 000000000..f2419e6b4 --- /dev/null +++ b/python/cudnn/gemm_amax/jax.py @@ -0,0 +1,252 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for block-scaled dense GEMM + amax on SM100.""" + +from __future__ import annotations + +import os +from typing import Any + +import jax.numpy as jnp + +from .._jax.api_base import ( + ApiBaseJax, + BufferSpec, + JaxTensorDesc, + TupleDict, + call_cutedsl, +) +from .._jax.gemm import ( + block_scale_tensor_spec, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, +) +from ..gemm_validation import require_gemm_shapes, resolve_max_active_clusters +from .validation import validate_gemm_amax + + +def _launch( + stream, + a, + b, + sfa, + sfb, + c, + amax, + *, + sf_vec_size: int, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + cluster_overlap_margin: int, +): + # These operations happen during CUDA lowering, not abstract evaluation. + import cutlass + + from .dense_blockscaled_gemm_persistent_amax import ( + Sm100BlockScaledPersistentDenseGemmKernel, + ) + + kernel = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size=sf_vec_size, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + ) + max_active_clusters = resolve_max_active_clusters( + cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]), + cluster_overlap_margin, + ) + kernel( + a, + b, + sfa, + sfb, + c, + amax, + max_active_clusters, + stream, + ) + + +class GemmAmaxSm100(ApiBaseJax): + """JAX GEMM + amax callable specialized from sample metadata.""" + + def __init__( + self, + sample_a: Any, + sample_b: Any, + sample_sfa: Any, + sample_sfb: Any, + c_layout: str = "LMN", + c_dtype: Any = None, + acc_dtype: Any = 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: + super().__init__() + self.a_desc = self.make_tensor_desc(sample_a, tensor_spec=gemm_a_tensor_spec(a_layout), name="sample_a") + self.b_desc = self.make_tensor_desc(sample_b, tensor_spec=gemm_b_tensor_spec(b_layout), name="sample_b") + self.sfa_desc = self.make_tensor_desc(sample_sfa, tensor_spec=block_scale_tensor_spec(), name="sample_sfa") + self.sfb_desc = self.make_tensor_desc(sample_sfb, tensor_spec=block_scale_tensor_spec(), name="sample_sfb") + self._c_layout = c_layout + + self._c_dtype = self.as_optional_dtype(c_dtype) + self._acc_dtype = self.as_optional_dtype(acc_dtype) + self.mma_tiler_mn = tuple(mma_tiler_mn) + self.cluster_shape_mn = tuple(cluster_shape_mn) + self.sf_vec_size = sf_vec_size + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + self._plan = None + + def _check_support(self) -> None: + supported_inputs = (jnp.float8_e4m3fn, jnp.float8_e5m2) + self.require_dtype(self.a_desc, supported_inputs) + self.require_dtype(self.b_desc, supported_inputs) + if self.sf_vec_size != 32: + raise NotImplementedError(f"The JAX MXFP8 path requires sf_vec_size=32, got {self.sf_vec_size}") + self.require_dtype(self.sfa_desc, (jnp.float8_e8m0fnu,)) + self.require_dtype(self.sfb_desc, (jnp.float8_e8m0fnu,)) + self.c_dtype = self.require_dtype( + self._c_dtype, + (jnp.float32, jnp.float16, jnp.bfloat16), + name="c_dtype", + default=jnp.float32, + ) + self.acc_dtype = self.require_dtype( + self._acc_dtype, + (jnp.float32,), + name="acc_dtype", + default=jnp.float32, + ) + + from .dense_blockscaled_gemm_persistent_amax import ( + Sm100BlockScaledPersistentDenseGemmKernel, + ) + + kernel = Sm100BlockScaledPersistentDenseGemmKernel + self.mma_tiler_mn = kernel.require_mma_tiler(self.mma_tiler_mn) + self.cluster_shape_mn = kernel.require_cluster_shape( + self.cluster_shape_mn, + mma_tiler_mn=self.mma_tiler_mn, + ) + + m, n, _, batch = require_gemm_shapes(self.a_desc.shape, self.b_desc.shape) + self.c_desc = JaxTensorDesc( + dtype=self.c_dtype, + shape=(m, n, batch), + tensor_spec=gemm_c_tensor_spec(self._c_layout), + name="c_tensor", + ) + self.amax_desc = JaxTensorDesc( + dtype=jnp.float32, + shape=(1, 1, 1), + name="amax_tensor", + ) + self._plan = validate_gemm_amax( + self.a_desc, + self.b_desc, + self.sfa_desc, + self.sfb_desc, + self.c_desc, + self.amax_desc, + acc_dtype=self.acc_dtype, + sf_vec_size=self.sf_vec_size, + supported_sf_vec_sizes=kernel.SF_VEC_SIZES, + mma_tiler_mn=self.mma_tiler_mn, + ) + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + ) -> TupleDict: + return super().__call__(a_tensor, b_tensor, sfa_tensor, sfb_tensor) + + def _call_impl( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + ) -> TupleDict: + if self._plan is None: + raise RuntimeError("check_support() did not produce a launch plan") + self.check_tensor_signature(a_tensor, self.a_desc, name="A") + self.check_tensor_signature(b_tensor, self.b_desc, name="B") + self.check_tensor_signature(sfa_tensor, self.sfa_desc, name="SFA") + self.check_tensor_signature(sfb_tensor, self.sfb_desc, name="SFB") + + c_tensor, amax_tensor = call_cutedsl( + _launch, + (a_tensor, b_tensor, sfa_tensor, sfb_tensor), + outputs=( + BufferSpec( + "c_tensor", + self.c_desc.array_shape, + self.c_desc.dtype, + tensor_spec=self.c_desc.tensor_spec, + ), + BufferSpec( + "amax_tensor", + self.amax_desc.shape, + self.amax_desc.dtype, + fill_value=-float("inf"), + ), + ), + input_specs=( + self.a_desc.tensor_spec, + self.b_desc.tensor_spec, + self.sfa_desc.tensor_spec, + self.sfb_desc.tensor_spec, + ), + static_args={ + "sf_vec_size": self.sf_vec_size, + "mma_tiler_mn": self.mma_tiler_mn, + "cluster_shape_mn": self.cluster_shape_mn, + "cluster_overlap_margin": self.num_cluster_overlap_margin, + }, + ) + return TupleDict(c_tensor=c_tensor, amax_tensor=amax_tensor) + + +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, + acc_dtype: Any = 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 FP8 block-scaled GEMM and a global max-absolute reduction.""" + + return GemmAmaxSm100( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + 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", "gemm_amax_wrapper_sm100"] diff --git a/python/cudnn/gemm_amax/validation.py b/python/cudnn/gemm_amax/validation.py new file mode 100644 index 000000000..e82f976e7 --- /dev/null +++ b/python/cudnn/gemm_amax/validation.py @@ -0,0 +1,217 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral validation for dense block-scaled GEMM + amax.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Sequence + +from ..api_base import TensorDesc, canonical_dtype_name +from ..gemm_validation import ( + require_block_scale_shapes, + require_contiguous_alignment, + require_gemm_shapes, + require_shape, +) + +_AB_DTYPES = frozenset({"float4_e2m1fn", "float8_e4m3fn", "float8_e5m2"}) +_SCALE_DTYPES = frozenset({"float8_e4m3fn", "float8_e8m0fnu"}) +_C_DTYPES = frozenset( + { + "bfloat16", + "float16", + "float32", + "float4_e2m1fn", + "float8_e4m3fn", + "float8_e5m2", + } +) +_WIDE_C_DTYPES = frozenset({"bfloat16", "float16", "float32"}) + + +@dataclass(frozen=True) +class GemmAmaxPlan: + """Validated logical signature and layout choices for GEMM + amax.""" + + m: int + n: int + k: int + batch: int + ab_dtype_name: str + scale_dtype_name: str + c_dtype_name: str + a_major: str + b_major: str + c_major: str + + @property + def c_shape(self) -> tuple[int, int, int]: + return (self.m, self.n, self.batch) + + @property + def amax_shape(self) -> tuple[int, int, int]: + return (1, 1, 1) + + +def _logical_dtype_name(desc: TensorDesc) -> str: + if getattr(desc, "packing", "native") == "fp4x2": + return "float4_e2m1fn" + return desc.dtype_name + + +def _logical_element_bits(desc: TensorDesc) -> int: + if _logical_dtype_name(desc) == "float4_e2m1fn": + return 4 + if desc.element_bits is None: + raise ValueError(f"Cannot determine the element width of {desc.name or desc.dtype_name}") + return desc.element_bits + + +def _require_supported_dtype(name: str, actual: str, supported: frozenset[str]) -> None: + if actual not in supported: + choices = ", ".join(sorted(supported)) + raise ValueError(f"{name} dtype must be one of {{{choices}}}, got {actual}") + + +def _require_same_storage_dtype(name: str, lhs: TensorDesc, rhs: TensorDesc) -> None: + lhs_storage = getattr(lhs, "storage_dtype_name", lhs.dtype_name) + rhs_storage = getattr(rhs, "storage_dtype_name", rhs.dtype_name) + lhs_packing = getattr(lhs, "packing", "native") + rhs_packing = getattr(rhs, "packing", "native") + if (lhs_storage, lhs_packing) != (rhs_storage, rhs_packing): + lhs_label = lhs_storage if lhs_packing == "native" else f"{lhs_storage}/{lhs_packing}" + rhs_label = rhs_storage if rhs_packing == "native" else f"{rhs_storage}/{rhs_packing}" + raise ValueError(f"{name} dtypes must match, got {lhs_label} and {rhs_label}") + + +def _require_compact_major( + name: str, + desc: TensorDesc, + layouts: dict[str, tuple[tuple[int, ...], tuple[int, ...]]], +) -> str: + for major, (stride, stride_order) in layouts.items(): + if desc.stride == stride and desc.stride_order == stride_order: + return major + expected = [stride for stride, _ in layouts.values()] + raise ValueError(f"{name} tensor must use one of the compact strides {expected}, got " f"stride {desc.stride} with order {desc.stride_order}") + + +def validate_gemm_amax( + a: TensorDesc, + b: TensorDesc, + sfa: TensorDesc, + sfb: TensorDesc, + c: TensorDesc, + amax: TensorDesc, + *, + acc_dtype: Any, + sf_vec_size: int, + supported_sf_vec_sizes: Sequence[int], + mma_tiler_mn: Sequence[int], +) -> GemmAmaxPlan: + """Validate the kernel-domain contract shared by Torch and JAX. + + Framework adapters remain responsible for narrowing their public dtype + surface, selecting the target device, and converting these descriptors to + their execution-specific tensor metadata. + """ + + m, n, k, batch = require_gemm_shapes(a.shape, b.shape) + _require_same_storage_dtype("a_tensor and b_tensor", a, b) + ab_dtype_name = _logical_dtype_name(a) + _require_supported_dtype("a_tensor and b_tensor", ab_dtype_name, _AB_DTYPES) + + supported_sf_vec_sizes = tuple(supported_sf_vec_sizes) + if sf_vec_size not in supported_sf_vec_sizes: + raise ValueError(f"sf_vec_size must be one of {supported_sf_vec_sizes}, got {sf_vec_size}") + + _require_same_storage_dtype("sfa_tensor and sfb_tensor", sfa, sfb) + scale_dtype_name = _logical_dtype_name(sfa) + _require_supported_dtype("sfa_tensor and sfb_tensor", scale_dtype_name, _SCALE_DTYPES) + require_block_scale_shapes( + sfa.shape, + sfb.shape, + m=m, + n=n, + k=k, + batch=batch, + sf_vec_size=sf_vec_size, + sfa_name="sfa_tensor", + sfb_name="sfb_tensor", + ) + + c_dtype_name = _logical_dtype_name(c) + _require_supported_dtype("C", c_dtype_name, _C_DTYPES) + require_shape("C", c.shape, (m, n, batch)) + require_shape("amax", amax.shape, (1, 1, 1)) + if _logical_dtype_name(amax) != "float32": + raise ValueError(f"amax dtype must be float32, got {_logical_dtype_name(amax)}") + if canonical_dtype_name(acc_dtype) != "float32": + raise ValueError(f"Accumulator dtype must be float32, got {canonical_dtype_name(acc_dtype)}") + + a_major = _require_compact_major( + "A", + a, + { + "m": ((1, m, m * k), (0, 1, 2)), + "k": ((k, 1, m * k), (1, 0, 2)), + }, + ) + b_major = _require_compact_major( + "B", + b, + { + "n": ((1, n, n * k), (0, 1, 2)), + "k": ((k, 1, n * k), (1, 0, 2)), + }, + ) + c_major = _require_compact_major( + "C", + c, + { + "m": ((1, m, m * n), (0, 1, 2)), + "n": ((n, 1, m * n), (1, 0, 2)), + }, + ) + + if ab_dtype_name == "float4_e2m1fn" and (a_major, b_major) != ("k", "k"): + raise ValueError("Float4 A and B tensors require k-major layouts, got " f"{a_major}-major and {b_major}-major") + if c_dtype_name == "float4_e2m1fn" and c_major != "n": + raise ValueError(f"Float4 C tensors require n-major layout, got {c_major}-major") + if c_dtype_name == "float4_e2m1fn" and ab_dtype_name != "float4_e2m1fn": + raise ValueError("Float4 C requires Float4 A and B, got " f"C={c_dtype_name}, A/B={ab_dtype_name}") + if c_dtype_name.startswith("float8_") and ab_dtype_name.startswith("float8_"): + raise NotImplementedError("FP8 A/B with FP8 C is unsupported because it fails to launch") + + if scale_dtype_name == "float8_e4m3fn" and sf_vec_size == 32: + raise ValueError("float8_e4m3fn scale factors do not support sf_vec_size=32") + if ab_dtype_name.startswith("float8_") and sf_vec_size == 16: + raise ValueError("FP8 A and B tensors do not support sf_vec_size=16") + + mma_tiler_mn = tuple(mma_tiler_mn) + if ab_dtype_name == "float4_e2m1fn" and mma_tiler_mn[1] == 256 and k <= 128: + raise ValueError(f"mma_tiler_mn (X, 256) requires K > 128 for Float4, got {k}") + if mma_tiler_mn == (128, 256) and sf_vec_size == 16 and c_dtype_name 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_contiguous_alignment("A", m if a_major == "m" else k, _logical_element_bits(a)) + require_contiguous_alignment("B", n if b_major == "n" else k, _logical_element_bits(b)) + require_contiguous_alignment("C", m if c_major == "m" else n, _logical_element_bits(c)) + + return GemmAmaxPlan( + m=m, + n=n, + k=k, + batch=batch, + ab_dtype_name=ab_dtype_name, + scale_dtype_name=scale_dtype_name, + c_dtype_name=c_dtype_name, + a_major=a_major, + b_major=b_major, + c_major=c_major, + ) + + +__all__ = ["GemmAmaxPlan", "validate_gemm_amax"] diff --git a/python/cudnn/gemm_dsrelu/__init__.py b/python/cudnn/gemm_dsrelu/__init__.py index ade25458b..22ad54bf1 100644 --- a/python/cudnn/gemm_dsrelu/__init__.py +++ b/python/cudnn/gemm_dsrelu/__init__.py @@ -1,9 +1,16 @@ -from .api import ( - GemmDsreluSm100, - gemm_dsrelu_wrapper_sm100, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy API exports for dense GEMM + squared-ReLU backward.""" + +from .._operation_api import make_operation_api -__all__ = [ +_API_EXPORTS = ( "GemmDsreluSm100", "gemm_dsrelu_wrapper_sm100", -] +) + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/gemm_dsrelu/api.py b/python/cudnn/gemm_dsrelu/api.py index fb02fa027..82241938e 100644 --- a/python/cudnn/gemm_dsrelu/api.py +++ b/python/cudnn/gemm_dsrelu/api.py @@ -13,8 +13,15 @@ from cuda.bindings import driver as cuda from cutlass.cute.runtime import make_fake_stream -from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.api_base import ApiBaseTorch, TupleDict, ceil_div from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.gemm_validation import ( + block_scale_shape, + require_contiguous_alignment, + require_full_mma_rows, + require_gemm_shapes, + resolve_max_active_clusters, +) from .dense_blockscaled_gemm_persistent_dsrelu_quant import ( Sm100BlockScaledPersistentDenseGemmKernel, @@ -29,7 +36,7 @@ def _major_from_stride_order(stride_order: Tuple[int, ...], mode0_label: str, mo raise ValueError(f"Unsupported stride order {stride_order}") -class GemmDsreluSm100(APIBase): +class GemmDsreluSm100(ApiBaseTorch): def __init__( self, sample_a: torch.Tensor, @@ -51,8 +58,10 @@ def __init__( vector_f32: bool = False, ): super().__init__() + self._interpret_uint8_as_fp4x2 = True self._warn_experimental_api() + self._kernel = Sm100BlockScaledPersistentDenseGemmKernel self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") @@ -73,31 +82,28 @@ def __init__( self.alpha = alpha self.acc_dtype = acc_dtype self.mma_tiler_mn = mma_tiler_mn - self.cluster_shape_mn = cluster_shape_mn if cluster_shape_mn is not None else ((2, 1) if mma_tiler_mn[0] == 256 else (1, 1)) + self.cluster_shape_mn = ( + cluster_shape_mn if cluster_shape_mn is not None else ((2, 1) if mma_tiler_mn[0] == self._kernel.TWO_CTA_MMA_TILER_M else (1, 1)) + ) self.sf_vec_size = sf_vec_size self.vector_f32 = vector_f32 self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) - self._interpret_uint8_as_fp4x2 = True - self._kernel = Sm100BlockScaledPersistentDenseGemmKernel - def check_support(self) -> bool: - m, k, l = self._tensor_shape(self.a_desc, name="sample_a") - n, b_k, b_l = self._tensor_shape(self.b_desc, name="sample_b") - - self._value_error_if((b_k, b_l) != (k, l), f"B shape mismatch: expected (*, {k}, {l}), got {(n, b_k, b_l)}") + m, n, k, l = require_gemm_shapes( + self._tensor_shape(self.a_desc, name="sample_a"), + self._tensor_shape(self.b_desc, name="sample_b"), + ) self._check_tensor_shape(self.c_desc, (m, n, l), "C") self._check_tensor_shape(self.d_desc, (m, n, l), "D") self._check_tensor_shape(self.prob_desc, (m, 1, l), "prob") self._check_tensor_shape(self.dprob_desc, (m, 1, l), "dprob") - rest_k = ceil_div(ceil_div(k, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfa_desc, (32, 4, ceil_div(m, 128), 4, rest_k, l), "SFA") - self._check_tensor_shape(self.sfb_desc, (32, 4, ceil_div(n, 128), 4, rest_k, l), "SFB") + self._check_tensor_shape(self.sfa_desc, block_scale_shape(m, k, l, self.sf_vec_size), "SFA") + self._check_tensor_shape(self.sfb_desc, block_scale_shape(n, k, l, self.sf_vec_size), "SFB") if self.sfd_desc is not None: - rest_n = ceil_div(ceil_div(n, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfd_desc, (32, 4, ceil_div(m, 128), 4, rest_n, l), "SFD") + self._check_tensor_shape(self.sfd_desc, block_scale_shape(m, n, l, self.sf_vec_size), "SFD") self._check_tensor_shape(self.amax_desc, (1,), "amax") self._check_tensor_shape(self.norm_const_desc, (1,), "norm_const") @@ -127,7 +133,10 @@ 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.sf_vec_size not in self._kernel.SF_VEC_SIZES, + f"sf_vec_size must be one of {self._kernel.SF_VEC_SIZES}, 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" ) @@ -141,21 +150,29 @@ def check_support(self) -> bool: d_major = _major_from_stride_order(self.d_desc.stride_order, "m", "n") self._value_error_if(c_major != d_major, f"C and D must share the same layout, got {c_major} and {d_major}") - self._value_error_if( - self.mma_tiler_mn[0] not in {128, 256} or self.mma_tiler_mn[1] not in {64, 128, 192, 256}, - f"Unsupported mma_tiler_mn {self.mma_tiler_mn}", + self.mma_tiler_mn = self._kernel.require_mma_tiler(self.mma_tiler_mn) + require_full_mma_rows( + m, + self.mma_tiler_mn[0], + cta_group_size=2 if self.mma_tiler_mn[0] == self._kernel.TWO_CTA_MMA_TILER_M else 1, + reason="the probability load is not predicated", ) - self._value_error_if( - not ( - self.cluster_shape_mn[0] > 0 - and self.cluster_shape_mn[1] > 0 - and self.cluster_shape_mn[0] * self.cluster_shape_mn[1] <= 16 - and is_power_of_2(self.cluster_shape_mn[0]) - and is_power_of_2(self.cluster_shape_mn[1]) - ), - f"Invalid cluster shape {self.cluster_shape_mn}", + self.cluster_shape_mn = self._kernel.require_cluster_shape( + self.cluster_shape_mn, + mma_tiler_mn=self.mma_tiler_mn, ) + ab_bits = _convert_to_cutlass_data_type( + self.ab_dtype, + interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2, + ).width + c_bits = _convert_to_cutlass_data_type(self.c_dtype).width + d_bits = _convert_to_cutlass_data_type(self.d_dtype).width + require_contiguous_alignment("A", m if a_major == "m" else k, ab_bits) + require_contiguous_alignment("B", n if b_major == "n" else k, ab_bits) + require_contiguous_alignment("C", m if c_major == "m" else n, c_bits) + require_contiguous_alignment("D", m if d_major == "m" else n, d_bits) + self._runtime_error_if(not torch.cuda.is_available(), "CUDA is not available") major, minor = torch.cuda.get_device_capability(torch.cuda.current_device()) self._runtime_error_if(major * 10 + minor < 100, f"GemmDsreluSm100 requires SM100+, found SM{major}{minor}") @@ -195,9 +212,10 @@ def compile(self) -> None: ) hardware_info = cutlass.utils.HardwareInfo() - max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) - max_active_clusters -= self.num_cluster_overlap_margin - self._value_error_if(max_active_clusters <= 0, "max_active_clusters must be > 0 after overlap margin") + max_active_clusters = resolve_max_active_clusters( + hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]), + self.num_cluster_overlap_margin, + ) fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) epilogue_op = lambda x, y: cute.where(x > 0, x, cute.full_like(x, 0)) * 2 * y 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..6efed6f99 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,7 +29,6 @@ from typing import Type, Tuple, Union, Optional import cuda.bindings.driver as cuda -import torch import cutlass import cutlass.cute as cute @@ -46,6 +45,9 @@ from cutlass._mlir.dialects.nvvm import AtomicOpKind from cutlass.cutlass_dsl import T +from ..gemm_validation import require_cluster_shape as _require_cluster_shape +from ..gemm_validation import require_mma_tiler as _require_mma_tiler + def atomic_add_float32( ptr, @@ -107,6 +109,41 @@ class Sm100BlockScaledPersistentDenseGemmKernel: >>> gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, d_tensor, prob_tensor, amax_tensor, sfd_tensor, norm_const_tensor, alpha, max_active_clusters, stream) """ + # Configuration values supported by the FE Torch and JAX wrappers. + MMA_TILER_M = (128, 256) + MMA_TILER_N = (64, 128, 192, 256) + TWO_CTA_MMA_TILER_M = 256 + MAX_CLUSTER_CTAS = 16 + MAX_CLUSTER_DIMENSION = 4 + SF_VEC_SIZES = (16, 32) + + @classmethod + def require_mma_tiler(cls, mma_tiler_mn: Tuple[int, int]) -> Tuple[int, int]: + """Validate an FE-supported MMA tile.""" + + return _require_mma_tiler( + mma_tiler_mn, + allowed_m=cls.MMA_TILER_M, + allowed_n=cls.MMA_TILER_N, + ) + + @classmethod + def require_cluster_shape( + cls, + cluster_shape_mn: Tuple[int, int], + *, + mma_tiler_mn: Tuple[int, int], + ) -> Tuple[int, int]: + """Validate an FE-supported cluster shape for an MMA tile.""" + + return _require_cluster_shape( + cluster_shape_mn, + mma_m=mma_tiler_mn[0], + two_cta_mma_m=cls.TWO_CTA_MMA_TILER_M, + max_ctas=cls.MAX_CLUSTER_CTAS, + max_dimension=cls.MAX_CLUSTER_DIMENSION, + ) + def __init__( self, sf_vec_size: int, @@ -136,7 +173,7 @@ def __init__( self.acc_dtype = cutlass.Float32 self.sf_vec_size = sf_vec_size - self.use_2cta_instrs = mma_tiler_mn[0] == 256 + self.use_2cta_instrs = mma_tiler_mn[0] == self.TWO_CTA_MMA_TILER_M self.cluster_shape_mn = cluster_shape_mn # K dimension is deferred in _setup_attributes self.mma_tiler = (*mma_tiler_mn, 1) @@ -1909,8 +1946,9 @@ def get_dtype_rcp_limits(dtype: Type[cutlass.Numeric]) -> float: return 1 / 128.0 return 1.0 - @staticmethod + @classmethod def is_valid_dtypes_and_scale_factor_vec_size( + cls, ab_dtype: Type[cutlass.Numeric], sf_dtype: Type[cutlass.Numeric], sf_vec_size: int, @@ -1943,7 +1981,7 @@ def is_valid_dtypes_and_scale_factor_vec_size( if ab_dtype in {cutlass.Float8E5M2, cutlass.Float8E4M3FN}: # Check valid sf_vec_size - if sf_vec_size not in {16, 32}: + if sf_vec_size not in cls.SF_VEC_SIZES: is_valid = False # Check valid sf_dtype if sf_dtype not in {cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN}: @@ -2002,8 +2040,9 @@ def is_valid_layouts( return is_valid - @staticmethod + @classmethod def is_valid_mma_tiler_and_cluster_shape( + cls, mma_tiler_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], ) -> bool: @@ -2018,30 +2057,12 @@ def is_valid_mma_tiler_and_cluster_shape( :return: True if the mma tiler and cluster shape are valid, False otherwise :rtype: bool """ - is_valid = True - # Skip invalid mma tile shape - if mma_tiler_mn[0] not in [128, 256]: - is_valid = False - if mma_tiler_mn[1] not in [64, 128, 192, 256]: - is_valid = False - # Skip illegal cluster shape - if cluster_shape_mn[0] % (2 if mma_tiler_mn[0] == 256 else 1) != 0: - is_valid = False - # Skip invalid cluster shape - is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 - if ( - cluster_shape_mn[0] * cluster_shape_mn[1] > 16 - or cluster_shape_mn[0] <= 0 - or cluster_shape_mn[1] <= 0 - # Special cluster shape check for scale factor multicasts. - # Due to limited size of scale factors, we can't multicast among more than 4 CTAs. - or cluster_shape_mn[0] > 4 - or cluster_shape_mn[1] > 4 - or not is_power_of_2(cluster_shape_mn[0]) - or not is_power_of_2(cluster_shape_mn[1]) - ): - is_valid = False - return is_valid + try: + mma_tiler_mn = cls.require_mma_tiler(mma_tiler_mn) + cls.require_cluster_shape(cluster_shape_mn, mma_tiler_mn=mma_tiler_mn) + except (ValueError, NotImplementedError): + return False + return True @staticmethod def is_valid_tensor_alignment( @@ -2096,8 +2117,9 @@ def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): is_valid = False return is_valid - @staticmethod + @classmethod def can_implement( + cls, ab_dtype: Type[cutlass.Numeric], sf_dtype: Type[cutlass.Numeric], sf_vec_size: int, @@ -2147,16 +2169,16 @@ def can_implement( """ can_implement = True # Skip unsupported types - if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_dtypes_and_scale_factor_vec_size(ab_dtype, sf_dtype, sf_vec_size, d_dtype): + if not cls.is_valid_dtypes_and_scale_factor_vec_size(ab_dtype, sf_dtype, sf_vec_size, d_dtype): can_implement = False # Skip unsupported layouts - if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_layouts(ab_dtype, d_dtype, a_major, b_major, d_major): + if not cls.is_valid_layouts(ab_dtype, d_dtype, a_major, b_major, d_major): can_implement = False # Skip invalid mma tile shape and cluster shape - if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_mma_tiler_and_cluster_shape(mma_tiler_mn, cluster_shape_mn): + if not cls.is_valid_mma_tiler_and_cluster_shape(mma_tiler_mn, cluster_shape_mn): can_implement = False # Skip illegal problem shape for load/store alignment - if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_tensor_alignment(m, n, k, l, ab_dtype, d_dtype, a_major, b_major, d_major): + if not cls.is_valid_tensor_alignment(m, n, k, l, ab_dtype, d_dtype, a_major, b_major, d_major): can_implement = False return can_implement diff --git a/python/cudnn/gemm_dsrelu/jax.py b/python/cudnn/gemm_dsrelu/jax.py new file mode 100644 index 000000000..c4a88e0ed --- /dev/null +++ b/python/cudnn/gemm_dsrelu/jax.py @@ -0,0 +1,348 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for block-scaled dense GEMM + squared-ReLU backward on SM100.""" + +from __future__ import annotations + +import os +from typing import Any, Optional + +import jax.numpy as jnp + +from .._jax.api_base import ( + ApiBaseJax, + BufferSpec, + JaxTensorDesc, + TupleDict, + call_cutedsl, + require_array, +) +from .._jax.gemm import ( + block_scale_tensor_spec, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + probability_tensor_spec, + require_16_byte_extent, + require_fp8_block_scales, + require_gemm_inputs, +) +from ..gemm_validation import ( + require_full_mma_rows, + resolve_max_active_clusters, +) + + +def _launch( + stream, + a, + b, + c, + sfa, + sfb, + prob, + d, + dprob, + *, + alpha: float, + sf_vec_size: int, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + vector_f32: bool, + cluster_overlap_margin: int, +): + # These operations happen during CUDA lowering, not abstract evaluation. + import cutlass + import cutlass.cute as cute + + from .dense_blockscaled_gemm_persistent_dsrelu_quant import ( + Sm100BlockScaledPersistentDenseGemmKernel, + ) + + kernel = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size=sf_vec_size, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vector_f32=vector_f32, + ) + max_active_clusters = resolve_max_active_clusters( + cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]), + cluster_overlap_margin, + ) + + 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, + None, + None, + None, + cutlass.Float32(alpha), + max_active_clusters, + stream, + epilogue_op=squared_relu_backward, + ) + + +class GemmDsreluSm100(ApiBaseJax): + """JAX squared-ReLU backward callable specialized from sample metadata.""" + + def __init__( + self, + sample_a: Any, + sample_b: Any, + sample_c: Any, + sample_sfa: Any, + sample_sfb: Any, + sample_prob: Any, + alpha: float = 1.0, + d_layout: str = "LMN", + d_dtype: Any = None, + acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sample_norm_const: Optional[Any] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + *, + a_layout: str = "LMK", + b_layout: str = "LNK", + ) -> None: + super().__init__() + self.a_desc = self.make_tensor_desc(sample_a, tensor_spec=gemm_a_tensor_spec(a_layout), name="sample_a") + self.b_desc = self.make_tensor_desc(sample_b, tensor_spec=gemm_b_tensor_spec(b_layout), name="sample_b") + self.c_desc = self.make_tensor_desc( + sample_c, + tensor_spec=gemm_c_tensor_spec(d_layout, name="d_layout"), + name="sample_c", + ) + self.sfa_desc = self.make_tensor_desc(sample_sfa, tensor_spec=block_scale_tensor_spec(), name="sample_sfa") + self.sfb_desc = self.make_tensor_desc(sample_sfb, tensor_spec=block_scale_tensor_spec(), name="sample_sfb") + self.prob_desc = self.make_tensor_desc(sample_prob, tensor_spec=probability_tensor_spec(), name="sample_prob") + self.norm_const_desc = self.make_optional_tensor_desc(sample_norm_const, name="sample_norm_const") + + self.alpha = alpha + self._d_dtype = self.as_optional_dtype(d_dtype) + self._acc_dtype = self.as_optional_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.vector_f32 = vector_f32 + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + + def _check_support(self) -> None: + if self.norm_const_desc is not None: + raise NotImplementedError("sample_norm_const is used by the FP8 output path, which is not available in the JAX squared-ReLU backward API") + + from .dense_blockscaled_gemm_persistent_dsrelu_quant import ( + Sm100BlockScaledPersistentDenseGemmKernel, + ) + + kernel = Sm100BlockScaledPersistentDenseGemmKernel + self.m, self.n, self.k, self.batch, self.ab_dtype = require_gemm_inputs(self.a_desc, self.b_desc) + supported_inputs = {jnp.dtype(jnp.float8_e4m3fn), jnp.dtype(jnp.float8_e5m2)} + if self.ab_dtype not in supported_inputs: + raise NotImplementedError("The JAX squared-ReLU backward API supports float8_e4m3fn and " f"float8_e5m2 inputs, got {self.ab_dtype}") + require_fp8_block_scales( + self.sfa_desc, + self.sfb_desc, + m=self.m, + n=self.n, + k=self.k, + batch=self.batch, + sf_vec_size=self.sf_vec_size, + ) + + require_array( + self.c_desc, + shape=(self.m, self.n, self.batch), + ) + require_array( + self.prob_desc, + shape=(self.m, 1, self.batch), + dtype=jnp.float32, + ) + + supported_outputs = (jnp.float16, jnp.bfloat16, jnp.float32) + self.c_dtype = self.require_dtype(self.c_desc, supported_outputs) + self.d_dtype = self.require_dtype( + self._d_dtype, + supported_outputs, + name="d_dtype", + default=jnp.bfloat16, + ) + self.require_dtype( + self._acc_dtype, + (jnp.float32,), + name="acc_dtype", + default=jnp.float32, + ) + + self.mma_tiler_mn = kernel.require_mma_tiler(self.mma_tiler_mn) + require_full_mma_rows( + self.m, + self.mma_tiler_mn[0], + cta_group_size=(2 if self.mma_tiler_mn[0] == kernel.TWO_CTA_MMA_TILER_M else 1), + reason="the probability load is not predicated", + ) + if self.cluster_shape_mn is None: + self.cluster_shape_mn = (2, 1) if self.mma_tiler_mn[0] == kernel.TWO_CTA_MMA_TILER_M else (1, 1) + self.cluster_shape_mn = kernel.require_cluster_shape( + self.cluster_shape_mn, + mma_tiler_mn=self.mma_tiler_mn, + ) + + require_16_byte_extent("sample_a", self.a_desc.shape[self.a_desc.stride_order[0]], self.ab_dtype) + require_16_byte_extent("sample_b", self.b_desc.shape[self.b_desc.stride_order[0]], self.ab_dtype) + require_16_byte_extent("sample_c", self.c_desc.shape[self.c_desc.stride_order[0]], self.c_dtype) + + self.d_desc = JaxTensorDesc( + dtype=self.d_dtype, + shape=(self.m, self.n, self.batch), + tensor_spec=self.c_desc.tensor_spec, + name="d_tensor", + ) + require_16_byte_extent("d_tensor", self.d_desc.shape[self.d_desc.stride_order[0]], self.d_dtype) + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_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, + prob_tensor, + norm_const_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + prob_tensor: Any, + norm_const_tensor: Optional[Any], + ) -> TupleDict: + self.check_tensor_signature(a_tensor, self.a_desc, name="A") + self.check_tensor_signature(b_tensor, self.b_desc, name="B") + self.check_tensor_signature(c_tensor, self.c_desc, name="C") + self.check_tensor_signature(sfa_tensor, self.sfa_desc, name="SFA") + self.check_tensor_signature(sfb_tensor, self.sfb_desc, name="SFB") + self.check_tensor_signature(prob_tensor, self.prob_desc, name="prob") + self.check_optional_tensor_signature(norm_const_tensor, self.norm_const_desc, name="norm_const") + + d_tensor, dprob_tensor = call_cutedsl( + _launch, + (a_tensor, b_tensor, c_tensor, sfa_tensor, sfb_tensor, prob_tensor), + outputs=( + BufferSpec( + "d_tensor", + self.d_desc.array_shape, + self.d_desc.dtype, + tensor_spec=self.d_desc.tensor_spec, + ), + BufferSpec( + "dprob_tensor", + (self.m, 1, self.batch), + jnp.float32, + tensor_spec=probability_tensor_spec(), + fill_value=0.0, + ), + ), + input_specs=( + self.a_desc.tensor_spec, + self.b_desc.tensor_spec, + self.c_desc.tensor_spec, + self.sfa_desc.tensor_spec, + self.sfb_desc.tensor_spec, + self.prob_desc.tensor_spec, + ), + static_args={ + "alpha": float(self.alpha), + "sf_vec_size": self.sf_vec_size, + "mma_tiler_mn": self.mma_tiler_mn, + "cluster_shape_mn": self.cluster_shape_mn, + "vector_f32": bool(self.vector_f32), + "cluster_overlap_margin": self.num_cluster_overlap_margin, + }, + ) + return TupleDict( + d_tensor=d_tensor, + dprob_tensor=dprob_tensor, + amax_tensor=None, + sfd_tensor=None, + ) + + +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, + acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + norm_const_tensor: Optional[Any] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + *, + a_layout: str = "LMK", + b_layout: str = "LNK", +) -> TupleDict: + """Compute the MXFP8 squared-ReLU backward fusion.""" + + return GemmDsreluSm100( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + prob_tensor, + 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, + sample_norm_const=norm_const_tensor, + 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", "gemm_dsrelu_wrapper_sm100"] diff --git a/python/cudnn/gemm_srelu/__init__.py b/python/cudnn/gemm_srelu/__init__.py index 29324f369..dbf73de4d 100644 --- a/python/cudnn/gemm_srelu/__init__.py +++ b/python/cudnn/gemm_srelu/__init__.py @@ -1,9 +1,16 @@ -from .api import ( - GemmSreluSm100, - gemm_srelu_wrapper_sm100, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy API exports for dense GEMM + squared ReLU.""" + +from .._operation_api import make_operation_api -__all__ = [ +_API_EXPORTS = ( "GemmSreluSm100", "gemm_srelu_wrapper_sm100", -] +) + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/gemm_srelu/api.py b/python/cudnn/gemm_srelu/api.py index 3bf21f95f..313b7ce0e 100644 --- a/python/cudnn/gemm_srelu/api.py +++ b/python/cudnn/gemm_srelu/api.py @@ -13,8 +13,15 @@ from cuda.bindings import driver as cuda from cutlass.cute.runtime import make_fake_stream -from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.api_base import ApiBaseTorch, TupleDict, ceil_div from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.gemm_validation import ( + block_scale_shape, + require_contiguous_alignment, + require_full_mma_rows, + require_gemm_shapes, + resolve_max_active_clusters, +) from .dense_blockscaled_gemm_persistent_srelu_quant import ( Sm100BlockScaledPersistentDenseGemmKernel, @@ -29,7 +36,7 @@ def _major_from_stride_order(stride_order: Tuple[int, ...], mode0_label: str, mo raise ValueError(f"Unsupported stride order {stride_order}") -class GemmSreluSm100(APIBase): +class GemmSreluSm100(ApiBaseTorch): def __init__( self, sample_a: torch.Tensor, @@ -50,8 +57,10 @@ def __init__( vector_f32: bool = False, ): super().__init__() + self._interpret_uint8_as_fp4x2 = True self._warn_experimental_api() + self._kernel = Sm100BlockScaledPersistentDenseGemmKernel self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") @@ -71,32 +80,27 @@ def __init__( self.alpha = alpha self.acc_dtype = acc_dtype self.mma_tiler_mn = mma_tiler_mn - self.cluster_shape_mn = cluster_shape_mn if cluster_shape_mn is not None else ((2, 1) if mma_tiler_mn[0] == 256 else (1, 1)) + self.cluster_shape_mn = ( + cluster_shape_mn if cluster_shape_mn is not None else ((2, 1) if mma_tiler_mn[0] == self._kernel.TWO_CTA_MMA_TILER_M else (1, 1)) + ) self.sf_vec_size = sf_vec_size self.vector_f32 = vector_f32 self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) - self._interpret_uint8_as_fp4x2 = True - self._kernel = Sm100BlockScaledPersistentDenseGemmKernel - def check_support(self) -> bool: - m, k, l = self._tensor_shape(self.a_desc, name="sample_a") - n, b_k, b_l = self._tensor_shape(self.b_desc, name="sample_b") - c_m, c_n, c_l = self._tensor_shape(self.c_desc, name="sample_c") - d_m, d_n, d_l = self._tensor_shape(self.d_desc, name="sample_d") - - self._value_error_if((b_k, b_l) != (k, l), f"B shape mismatch: expected (*, {k}, {l}), got {(n, b_k, b_l)}") + m, n, k, l = require_gemm_shapes( + self._tensor_shape(self.a_desc, name="sample_a"), + self._tensor_shape(self.b_desc, name="sample_b"), + ) self._check_tensor_shape(self.c_desc, (m, n, l), "C") self._check_tensor_shape(self.d_desc, (m, n, l), "D") self._check_tensor_shape(self.prob_desc, (m, 1, l), "prob") - rest_k = ceil_div(ceil_div(k, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfa_desc, (32, 4, ceil_div(m, 128), 4, rest_k, l), "SFA") - self._check_tensor_shape(self.sfb_desc, (32, 4, ceil_div(n, 128), 4, rest_k, l), "SFB") + self._check_tensor_shape(self.sfa_desc, block_scale_shape(m, k, l, self.sf_vec_size), "SFA") + self._check_tensor_shape(self.sfb_desc, block_scale_shape(n, k, l, self.sf_vec_size), "SFB") if self.sfd_desc is not None: - rest_n = ceil_div(ceil_div(n, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfd_desc, (32, 4, ceil_div(m, 128), 4, rest_n, l), "SFD") + self._check_tensor_shape(self.sfd_desc, block_scale_shape(m, n, l, self.sf_vec_size), "SFD") self._check_tensor_shape(self.amax_desc, (1,), "amax") self._check_tensor_shape(self.norm_const_desc, (1,), "norm_const") @@ -125,7 +129,10 @@ 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.sf_vec_size not in self._kernel.SF_VEC_SIZES, + f"sf_vec_size must be one of {self._kernel.SF_VEC_SIZES}, 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" ) @@ -139,21 +146,29 @@ def check_support(self) -> bool: d_major = _major_from_stride_order(self.d_desc.stride_order, "m", "n") self._value_error_if(c_major != d_major, f"C and D must share the same layout, got {c_major} and {d_major}") - self._value_error_if( - self.mma_tiler_mn[0] not in {128, 256} or self.mma_tiler_mn[1] not in {64, 128, 192, 256}, - f"Unsupported mma_tiler_mn {self.mma_tiler_mn}", + self.mma_tiler_mn = self._kernel.require_mma_tiler(self.mma_tiler_mn) + require_full_mma_rows( + m, + self.mma_tiler_mn[0], + cta_group_size=2 if self.mma_tiler_mn[0] == self._kernel.TWO_CTA_MMA_TILER_M else 1, + reason="the probability load is not predicated", ) - self._value_error_if( - not ( - self.cluster_shape_mn[0] > 0 - and self.cluster_shape_mn[1] > 0 - and self.cluster_shape_mn[0] * self.cluster_shape_mn[1] <= 16 - and is_power_of_2(self.cluster_shape_mn[0]) - and is_power_of_2(self.cluster_shape_mn[1]) - ), - f"Invalid cluster shape {self.cluster_shape_mn}", + self.cluster_shape_mn = self._kernel.require_cluster_shape( + self.cluster_shape_mn, + mma_tiler_mn=self.mma_tiler_mn, ) + ab_bits = _convert_to_cutlass_data_type( + self.ab_dtype, + interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2, + ).width + c_bits = _convert_to_cutlass_data_type(self.c_dtype).width + d_bits = _convert_to_cutlass_data_type(self.d_dtype).width + require_contiguous_alignment("A", m if a_major == "m" else k, ab_bits) + require_contiguous_alignment("B", n if b_major == "n" else k, ab_bits) + require_contiguous_alignment("C", m if c_major == "m" else n, c_bits) + require_contiguous_alignment("D", m if d_major == "m" else n, d_bits) + self._runtime_error_if(not torch.cuda.is_available(), "CUDA is not available") major, minor = torch.cuda.get_device_capability(torch.cuda.current_device()) self._runtime_error_if(major * 10 + minor < 100, f"GemmSreluSm100 requires SM100+, found SM{major}{minor}") @@ -193,9 +208,10 @@ def compile(self) -> None: ) hardware_info = cutlass.utils.HardwareInfo() - max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) - max_active_clusters -= self.num_cluster_overlap_margin - self._value_error_if(max_active_clusters <= 0, "max_active_clusters must be > 0 after overlap margin") + max_active_clusters = resolve_max_active_clusters( + hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]), + self.num_cluster_overlap_margin, + ) fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) ** 2 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..34437509a 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,7 +29,6 @@ from typing import Type, Tuple, Union, Optional import cuda.bindings.driver as cuda -import torch import cutlass import cutlass.cute as cute @@ -43,6 +42,9 @@ import cutlass.cute.math as math from cutlass.cute.typing import Float32 +from ..gemm_validation import require_cluster_shape as _require_cluster_shape +from ..gemm_validation import require_mma_tiler as _require_mma_tiler + def get_divisibility(dtype) -> int: """ @@ -103,6 +105,41 @@ class Sm100BlockScaledPersistentDenseGemmKernel: >>> gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, d_tensor, prob_tensor, amax_tensor, sfd_tensor, norm_const_tensor, alpha, max_active_clusters, stream) """ + # Configuration values supported by the FE Torch and JAX wrappers. + MMA_TILER_M = (128, 256) + MMA_TILER_N = (64, 128, 192, 256) + TWO_CTA_MMA_TILER_M = 256 + MAX_CLUSTER_CTAS = 16 + MAX_CLUSTER_DIMENSION = 4 + SF_VEC_SIZES = (16, 32) + + @classmethod + def require_mma_tiler(cls, mma_tiler_mn: Tuple[int, int]) -> Tuple[int, int]: + """Validate an FE-supported MMA tile.""" + + return _require_mma_tiler( + mma_tiler_mn, + allowed_m=cls.MMA_TILER_M, + allowed_n=cls.MMA_TILER_N, + ) + + @classmethod + def require_cluster_shape( + cls, + cluster_shape_mn: Tuple[int, int], + *, + mma_tiler_mn: Tuple[int, int], + ) -> Tuple[int, int]: + """Validate an FE-supported cluster shape for an MMA tile.""" + + return _require_cluster_shape( + cluster_shape_mn, + mma_m=mma_tiler_mn[0], + two_cta_mma_m=cls.TWO_CTA_MMA_TILER_M, + max_ctas=cls.MAX_CLUSTER_CTAS, + max_dimension=cls.MAX_CLUSTER_DIMENSION, + ) + def __init__( self, sf_vec_size: int, @@ -132,7 +169,7 @@ def __init__( self.acc_dtype = cutlass.Float32 self.sf_vec_size = sf_vec_size - self.use_2cta_instrs = mma_tiler_mn[0] == 256 + self.use_2cta_instrs = mma_tiler_mn[0] == self.TWO_CTA_MMA_TILER_M self.cluster_shape_mn = cluster_shape_mn # K dimension is deferred in _setup_attributes self.mma_tiler = (*mma_tiler_mn, 1) @@ -1375,15 +1412,18 @@ def kernel( self.epilog_sync_barrier.arrive_and_wait() + # Apply the squared-ReLU epilogue for every output mode. + # Amax generation is optional, but D always contains the + # probability-gated activation. + 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) @@ -1828,8 +1868,9 @@ def get_dtype_rcp_limits(dtype: Type[cutlass.Numeric]) -> float: return 1 / 128.0 return 1.0 - @staticmethod + @classmethod def is_valid_dtypes_and_scale_factor_vec_size( + cls, ab_dtype: Type[cutlass.Numeric], sf_dtype: Type[cutlass.Numeric], sf_vec_size: int, @@ -1862,7 +1903,7 @@ def is_valid_dtypes_and_scale_factor_vec_size( if ab_dtype in {cutlass.Float8E5M2, cutlass.Float8E4M3FN}: # Check valid sf_vec_size - if sf_vec_size not in {16, 32}: + if sf_vec_size not in cls.SF_VEC_SIZES: is_valid = False # Check valid sf_dtype if sf_dtype not in {cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN}: @@ -1921,8 +1962,9 @@ def is_valid_layouts( return is_valid - @staticmethod + @classmethod def is_valid_mma_tiler_and_cluster_shape( + cls, mma_tiler_mn: Tuple[int, int], cluster_shape_mn: Tuple[int, int], ) -> bool: @@ -1937,30 +1979,12 @@ def is_valid_mma_tiler_and_cluster_shape( :return: True if the mma tiler and cluster shape are valid, False otherwise :rtype: bool """ - is_valid = True - # Skip invalid mma tile shape - if mma_tiler_mn[0] not in [128, 256]: - is_valid = False - if mma_tiler_mn[1] not in [64, 128, 192, 256]: - is_valid = False - # Skip illegal cluster shape - if cluster_shape_mn[0] % (2 if mma_tiler_mn[0] == 256 else 1) != 0: - is_valid = False - # Skip invalid cluster shape - is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 - if ( - cluster_shape_mn[0] * cluster_shape_mn[1] > 16 - or cluster_shape_mn[0] <= 0 - or cluster_shape_mn[1] <= 0 - # Special cluster shape check for scale factor multicasts. - # Due to limited size of scale factors, we can't multicast among more than 4 CTAs. - or cluster_shape_mn[0] > 4 - or cluster_shape_mn[1] > 4 - or not is_power_of_2(cluster_shape_mn[0]) - or not is_power_of_2(cluster_shape_mn[1]) - ): - is_valid = False - return is_valid + try: + mma_tiler_mn = cls.require_mma_tiler(mma_tiler_mn) + cls.require_cluster_shape(cluster_shape_mn, mma_tiler_mn=mma_tiler_mn) + except (ValueError, NotImplementedError): + return False + return True @staticmethod def is_valid_tensor_alignment( @@ -2015,8 +2039,9 @@ def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): is_valid = False return is_valid - @staticmethod + @classmethod def can_implement( + cls, ab_dtype: Type[cutlass.Numeric], sf_dtype: Type[cutlass.Numeric], sf_vec_size: int, @@ -2066,16 +2091,16 @@ def can_implement( """ can_implement = True # Skip unsupported types - if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_dtypes_and_scale_factor_vec_size(ab_dtype, sf_dtype, sf_vec_size, d_dtype): + if not cls.is_valid_dtypes_and_scale_factor_vec_size(ab_dtype, sf_dtype, sf_vec_size, d_dtype): can_implement = False # Skip unsupported layouts - if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_layouts(ab_dtype, d_dtype, a_major, b_major, d_major): + if not cls.is_valid_layouts(ab_dtype, d_dtype, a_major, b_major, d_major): can_implement = False # Skip invalid mma tile shape and cluster shape - if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_mma_tiler_and_cluster_shape(mma_tiler_mn, cluster_shape_mn): + if not cls.is_valid_mma_tiler_and_cluster_shape(mma_tiler_mn, cluster_shape_mn): can_implement = False # Skip illegal problem shape for load/store alignment - if not Sm100BlockScaledPersistentDenseGemmKernel.is_valid_tensor_alignment(m, n, k, l, ab_dtype, d_dtype, a_major, b_major, d_major): + if not cls.is_valid_tensor_alignment(m, n, k, l, ab_dtype, d_dtype, a_major, b_major, d_major): can_implement = False return can_implement diff --git a/python/cudnn/gemm_srelu/jax.py b/python/cudnn/gemm_srelu/jax.py new file mode 100644 index 000000000..1b84b2122 --- /dev/null +++ b/python/cudnn/gemm_srelu/jax.py @@ -0,0 +1,330 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for block-scaled dense GEMM + squared ReLU on SM100.""" + +from __future__ import annotations + +import os +from typing import Any, Optional + +import jax.numpy as jnp + +from .._jax.api_base import ( + ApiBaseJax, + BufferSpec, + JaxTensorDesc, + TupleDict, + call_cutedsl, + require_array, +) +from .._jax.gemm import ( + block_scale_tensor_spec, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + probability_tensor_spec, + require_16_byte_extent, + require_fp8_block_scales, + require_gemm_inputs, +) +from ..gemm_validation import ( + require_full_mma_rows, + resolve_max_active_clusters, +) + + +def _launch( + stream, + a, + b, + sfa, + sfb, + prob, + c, + d, + *, + alpha: float, + sf_vec_size: int, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + vector_f32: bool, + cluster_overlap_margin: int, +): + # These operations happen during CUDA lowering, not abstract evaluation. + import cutlass + import cutlass.cute as cute + + from .dense_blockscaled_gemm_persistent_srelu_quant import ( + Sm100BlockScaledPersistentDenseGemmKernel, + ) + + kernel = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size=sf_vec_size, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vector_f32=vector_f32, + ) + max_active_clusters = resolve_max_active_clusters( + cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]), + cluster_overlap_margin, + ) + + def squared_relu(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) ** 2 + + kernel( + a, + b, + sfa, + sfb, + c, + d, + prob, + None, + None, + None, + cutlass.Float32(alpha), + max_active_clusters, + stream, + epilogue_op=squared_relu, + ) + + +class GemmSreluSm100(ApiBaseJax): + """JAX GEMM + squared-ReLU callable specialized from sample metadata.""" + + def __init__( + self, + sample_a: Any, + sample_b: Any, + sample_sfa: Any, + sample_sfb: Any, + sample_prob: Any, + alpha: float = 1.0, + c_layout: str = "LMN", + c_dtype: Any = None, + d_dtype: Any = None, + acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sample_norm_const: Optional[Any] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + *, + a_layout: str = "LMK", + b_layout: str = "LNK", + ) -> None: + super().__init__() + self.a_desc = self.make_tensor_desc(sample_a, tensor_spec=gemm_a_tensor_spec(a_layout), name="sample_a") + self.b_desc = self.make_tensor_desc(sample_b, tensor_spec=gemm_b_tensor_spec(b_layout), name="sample_b") + self.sfa_desc = self.make_tensor_desc(sample_sfa, tensor_spec=block_scale_tensor_spec(), name="sample_sfa") + self.sfb_desc = self.make_tensor_desc(sample_sfb, tensor_spec=block_scale_tensor_spec(), name="sample_sfb") + self.prob_desc = self.make_tensor_desc(sample_prob, tensor_spec=probability_tensor_spec(), name="sample_prob") + self.norm_const_desc = self.make_optional_tensor_desc(sample_norm_const, name="sample_norm_const") + self._c_layout = c_layout + + self.alpha = alpha + self._c_dtype = self.as_optional_dtype(c_dtype) + self._d_dtype = self.as_optional_dtype(d_dtype) + self._acc_dtype = self.as_optional_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.vector_f32 = vector_f32 + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + + def _check_support(self) -> None: + if self.norm_const_desc is not None: + raise NotImplementedError("sample_norm_const is used by the FP8 output path, which is not available in the JAX GEMM + squared-ReLU API") + + from .dense_blockscaled_gemm_persistent_srelu_quant import ( + Sm100BlockScaledPersistentDenseGemmKernel, + ) + + kernel = Sm100BlockScaledPersistentDenseGemmKernel + self.m, self.n, self.k, self.batch, self.ab_dtype = require_gemm_inputs(self.a_desc, self.b_desc) + supported_inputs = {jnp.dtype(jnp.float8_e4m3fn), jnp.dtype(jnp.float8_e5m2)} + if self.ab_dtype not in supported_inputs: + raise NotImplementedError("The JAX GEMM + squared-ReLU API supports float8_e4m3fn and " f"float8_e5m2 inputs, got {self.ab_dtype}") + require_fp8_block_scales( + self.sfa_desc, + self.sfb_desc, + m=self.m, + n=self.n, + k=self.k, + batch=self.batch, + sf_vec_size=self.sf_vec_size, + ) + require_array( + self.prob_desc, + shape=(self.m, 1, self.batch), + dtype=jnp.float32, + ) + + supported_outputs = (jnp.float16, jnp.bfloat16, jnp.float32) + self.c_dtype = self.require_dtype( + self._c_dtype, + supported_outputs, + name="c_dtype", + default=jnp.bfloat16, + ) + self.d_dtype = self.require_dtype( + self._d_dtype, + supported_outputs, + name="d_dtype", + default=jnp.bfloat16, + ) + self.require_dtype( + self._acc_dtype, + (jnp.float32,), + name="acc_dtype", + default=jnp.float32, + ) + + self.mma_tiler_mn = kernel.require_mma_tiler(self.mma_tiler_mn) + require_full_mma_rows( + self.m, + self.mma_tiler_mn[0], + cta_group_size=(2 if self.mma_tiler_mn[0] == kernel.TWO_CTA_MMA_TILER_M else 1), + reason="the probability load is not predicated", + ) + if self.cluster_shape_mn is None: + self.cluster_shape_mn = (2, 1) if self.mma_tiler_mn[0] == kernel.TWO_CTA_MMA_TILER_M else (1, 1) + self.cluster_shape_mn = kernel.require_cluster_shape( + self.cluster_shape_mn, + mma_tiler_mn=self.mma_tiler_mn, + ) + + require_16_byte_extent("sample_a", self.a_desc.shape[self.a_desc.stride_order[0]], self.ab_dtype) + require_16_byte_extent("sample_b", self.b_desc.shape[self.b_desc.stride_order[0]], self.ab_dtype) + + output_spec = gemm_c_tensor_spec(self._c_layout) + output_shape = (self.m, self.n, self.batch) + self.c_desc = JaxTensorDesc( + dtype=self.c_dtype, + shape=output_shape, + tensor_spec=output_spec, + name="c_tensor", + ) + self.d_desc = JaxTensorDesc( + dtype=self.d_dtype, + shape=output_shape, + tensor_spec=output_spec, + name="d_tensor", + ) + require_16_byte_extent("c_tensor", self.c_desc.shape[self.c_desc.stride_order[0]], self.c_dtype) + require_16_byte_extent("d_tensor", self.d_desc.shape[self.d_desc.stride_order[0]], self.d_dtype) + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + prob_tensor: Any, + norm_const_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__(a_tensor, b_tensor, sfa_tensor, sfb_tensor, prob_tensor, norm_const_tensor) + + def _call_impl( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + prob_tensor: Any, + norm_const_tensor: Optional[Any], + ) -> TupleDict: + self.check_tensor_signature(a_tensor, self.a_desc, name="A") + self.check_tensor_signature(b_tensor, self.b_desc, name="B") + self.check_tensor_signature(sfa_tensor, self.sfa_desc, name="SFA") + self.check_tensor_signature(sfb_tensor, self.sfb_desc, name="SFB") + self.check_tensor_signature(prob_tensor, self.prob_desc, name="prob") + self.check_optional_tensor_signature(norm_const_tensor, self.norm_const_desc, name="norm_const") + + c_tensor, d_tensor = call_cutedsl( + _launch, + (a_tensor, b_tensor, sfa_tensor, sfb_tensor, prob_tensor), + outputs=( + BufferSpec( + "c_tensor", + self.c_desc.array_shape, + self.c_desc.dtype, + tensor_spec=self.c_desc.tensor_spec, + ), + BufferSpec( + "d_tensor", + self.d_desc.array_shape, + self.d_desc.dtype, + tensor_spec=self.d_desc.tensor_spec, + ), + ), + input_specs=( + self.a_desc.tensor_spec, + self.b_desc.tensor_spec, + self.sfa_desc.tensor_spec, + self.sfb_desc.tensor_spec, + self.prob_desc.tensor_spec, + ), + static_args={ + "alpha": float(self.alpha), + "sf_vec_size": self.sf_vec_size, + "mma_tiler_mn": self.mma_tiler_mn, + "cluster_shape_mn": self.cluster_shape_mn, + "vector_f32": bool(self.vector_f32), + "cluster_overlap_margin": self.num_cluster_overlap_margin, + }, + ) + return TupleDict( + c_tensor=c_tensor, + d_tensor=d_tensor, + amax_tensor=None, + sfd_tensor=None, + ) + + +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, + d_dtype: Any = None, + acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + norm_const_tensor: Optional[Any] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + *, + a_layout: str = "LMK", + b_layout: str = "LNK", +) -> TupleDict: + """Compute MXFP8 GEMM and ``D = relu(C)**2 * prob``.""" + + return GemmSreluSm100( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + prob_tensor, + 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, + sample_norm_const=norm_const_tensor, + 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", "gemm_srelu_wrapper_sm100"] diff --git a/python/cudnn/gemm_swiglu/__init__.py b/python/cudnn/gemm_swiglu/__init__.py index cf2170c32..9855b0e8a 100644 --- a/python/cudnn/gemm_swiglu/__init__.py +++ b/python/cudnn/gemm_swiglu/__init__.py @@ -1,9 +1,16 @@ -from .api import ( - GemmSwigluSm100, - gemm_swiglu_wrapper_sm100, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy API exports for dense GEMM + SwiGLU.""" + +from .._operation_api import make_operation_api -__all__ = [ +_API_EXPORTS = ( "GemmSwigluSm100", "gemm_swiglu_wrapper_sm100", -] +) + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/gemm_swiglu/api.py b/python/cudnn/gemm_swiglu/api.py index 2daeaa740..c001af220 100644 --- a/python/cudnn/gemm_swiglu/api.py +++ b/python/cudnn/gemm_swiglu/api.py @@ -41,12 +41,18 @@ import cutlass.cute as cute from cutlass.cute.runtime import make_fake_stream +from cudnn.api_base import ApiBaseTorch, TupleDict, ceil_div from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.gemm_validation import ( + require_contiguous_alignment, + require_gemm_shapes, + resolve_max_active_clusters, +) +from .validation import validate_quantized_gemm_swiglu import os -class GemmSwigluSm100(APIBase): +class GemmSwigluSm100(ApiBaseTorch): def __init__( self, sample_a: torch.Tensor, @@ -68,6 +74,7 @@ def __init__( ab12_stages: int = 4, ): super().__init__() + self._interpret_uint8_as_fp4x2 = True self._warn_experimental_api() self._logger.debug("Entering __init__") @@ -79,10 +86,6 @@ def __init__( self.alpha = alpha self.acc_dtype = acc_dtype self.mma_tiler_mn = mma_tiler_mn - if cluster_shape_mn is None: - self.cluster_shape_mn = (1, 1) if not self.mma_tiler_mn[0] == 256 else (2, 2) - else: - self.cluster_shape_mn = cluster_shape_mn ### Quantize only arguments self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa") @@ -104,43 +107,35 @@ def __init__( self._logger.debug("Quantization arguments provided, using quantized GEMM swiglu kernel") self._kernel = Sm100BlockScaledPersistentDenseGemmKernel + if cluster_shape_mn is None: + self.cluster_shape_mn = (2, 2) if self.mma_tiler_mn[0] == self._kernel.TWO_CTA_MMA_TILER_M else (1, 1) + else: + self.cluster_shape_mn = 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}" ) - self._interpret_uint8_as_fp4x2 = True - def check_support(self) -> bool: self._logger.debug("Entering check_support") self._logger.debug("Checking tensor shapes, strides, and dtypes") - m, k, l = self.a_desc.shape - n, k, l = self.b_desc.shape - m, n, l = self.ab12_desc.shape - m, n_2, l = self.c_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.ab12_desc, (m, n, l), "AB12") - self._check_tensor_shape(self.c_desc, (m, n // 2, l), "C") - - if self._kernel is Sm100BlockScaledPersistentDenseGemmKernel: - rest_k = ceil_div(ceil_div(k, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfa_desc, (32, 4, ceil_div(m, 128), 4, rest_k, l), "SFA") - self._check_tensor_shape(self.sfb_desc, (32, 4, ceil_div(n, 128), 4, rest_k, l), "SFB") - self._check_tensor_shape(self.amax_desc, (1,), "amax") - rest_n2 = ceil_div(ceil_div(n // 2, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfc_desc, (32, 4, ceil_div(m, 128), 4, rest_n2, l), "SFC") - self._check_tensor_shape(self.norm_const_desc, (1,), "norm_const") - - _ = self._check_tensor_stride(self.a_desc, stride=[(1, m, m * k), (k, 1, m * k)]) - _ = self._check_tensor_stride(self.b_desc, stride=[(1, n, n * k), (k, 1, n * k)]) - _ = self._check_tensor_stride(self.ab12_desc, stride=[(1, m, m * n), (n, 1, m * n)]) - _ = self._check_tensor_stride(self.c_desc, stride=[(1, m, m * n_2), (n_2, 1, m * n_2)]) - self._value_error_if( - self.ab12_desc.stride_order != self.c_desc.stride_order, - f"AB12 and C tensor stride orders must match, got {self.ab12_desc.stride_order} and {self.c_desc.stride_order}", + m, n, k, l = require_gemm_shapes( + self._tensor_shape(self.a_desc, name="sample_a"), + self._tensor_shape(self.b_desc, name="sample_b"), ) + n_2 = self._kernel.get_output_n(n) + if self._kernel is PersistentDenseGemmKernel: + self._check_tensor_shape(self.ab12_desc, (m, n, l), "AB12") + self._check_tensor_shape(self.c_desc, (m, n_2, l), "C") + self._check_tensor_stride(self.a_desc, stride=[(1, m, m * k), (k, 1, m * k)]) + self._check_tensor_stride(self.b_desc, stride=[(1, n, n * k), (k, 1, n * k)]) + self._check_tensor_stride(self.ab12_desc, stride=[(1, m, m * n), (n, 1, m * n)]) + self._check_tensor_stride(self.c_desc, stride=[(1, m, m * n_2), (n_2, 1, m * n_2)]) + self._value_error_if( + self.ab12_desc.stride_order != self.c_desc.stride_order, + f"AB12 and C tensor stride orders must match, got {self.ab12_desc.stride_order} and {self.c_desc.stride_order}", + ) self._logger.debug("Checking data types") if self._kernel is PersistentDenseGemmKernel: @@ -186,192 +181,65 @@ def check_support(self) -> bool: case _: raise ValueError(f"Unsupported acc_dtype: expected one of {{torch.float32, torch.float16}}, got {self.acc_dtype}") self.c_dtype = self._check_dtype(self.c_desc, dtype=[torch.float16, torch.bfloat16], name="C") - elif self._kernel is Sm100BlockScaledPersistentDenseGemmKernel: - self._value_error_if( - self.sfa_desc is None or self.sfb_desc is None, - "sfa and sfb must be provided for quantized GEMM swiglu kernel", - ) - - self.ab_dtype = self._check_dtype( - self.a_desc, - dtype=[ - torch.float4_e2m1fn_x2, - torch.uint8, - torch.float8_e5m2, - torch.float8_e4m3fn, - ], - name="A (for quantized GEMM swiglu kernel)", - ) - self.acc_dtype = self._check_dtype( - self.acc_dtype, - dtype=torch.float32, - name="Accumulator (for quantized GEMM swiglu kernel)", - ) - self.ab12_dtype = self._check_dtype( - self.ab12_desc, - dtype=[ - torch.float32, - torch.float16, - torch.bfloat16, - torch.float8_e4m3fn, - torch.float8_e5m2, - ], - name="AB12 (for quantized GEMM swiglu kernel)", - ) - self.c_dtype = self._check_dtype( - self.c_desc, - dtype=[ - torch.float32, - torch.float16, - torch.bfloat16, - torch.float8_e4m3fn, - torch.float8_e5m2, - ], - name="C (for quantized GEMM swiglu kernel)", - ) - - self._value_error_if( - self._is_fp4x2(self.ab_dtype) and self._is_fp8(self.c_dtype), - "Invalid dtype combination: fp4 ab_dtype is not compatible with fp8 c_dtype (recommended bf16)", - ) - - self._value_error_if( - self._is_fp8(self.c_dtype) and (self.sfc_desc is None or self.norm_const_desc is None), - "sfc and norm_const must be provided when c_dtype is fp8", - ) - self._value_error_if( - (self._is_fp4x2(self.ab_dtype) and self.c_dtype == torch.bfloat16) and (self.amax_desc is None), - "amax must be provided when ab_dtype is fp4 and c_dtype is bf16", - ) - - self._not_implemented_error_if( - self.c_dtype == torch.float32 and self.ab12_dtype == torch.float32, - "float32 c_dtype and float32 ab12_dtype currently disabled due to kernel bug", - ) + elif self._kernel is not Sm100BlockScaledPersistentDenseGemmKernel: + raise NotImplementedError(f"Unreachable: invalid kernel type {self._kernel}") - self._value_error_if( - self.sf_vec_size not in {16, 32}, - f"sf_vec_size must be 16 or 32 when ab_dtype is {{torch.float8_e5m2, torch.float8_e4m3fn}}, got {self.sf_vec_size}", - ) - self.sf_dtype = self._check_dtype( - self.sfa_desc, - dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn], - name="SFA", - ) - self._check_dtype( - self.sfb_desc, - dtype=self.sf_dtype, - name="SFB", - extra_error_msg="SFB must have the same dtype as SFA", - ) + if self._kernel is PersistentDenseGemmKernel: self._check_dtype( - self.sfc_desc, - dtype=self.sf_dtype, - name="SFC", - extra_error_msg="SFC must have the same dtype as SFA", + self.b_desc, + dtype=self.ab_dtype, + name="B", + extra_error_msg="A and B must have the same dtype", ) - if self._is_fp8(self.ab_dtype): - self._value_error_if( - not (self.sf_dtype == torch.float8_e8m0fnu and self.sf_vec_size == 32), - "Invalid ab_dtype and sf_dtype/sf_vec_size combination: fp8 ab_dtype requires float8_e8m0fnu sf_dtype and 32 sf_vec_size", - ) - elif self._is_fp4x2(self.ab_dtype): - self._value_error_if( - self.sf_dtype == torch.float8_e4m3fn and self.sf_vec_size == 32, - "Invalid ab_dtype and sf_dtype/sf_vec_size combination: fp4 ab_dtype not supported with float8_e4m3fn sf_dtype and 32 sf_vec_size", - ) - - if self._is_fp4x2(self.ab_dtype): - self._value_error_if( - self.a_desc.stride_order != (1, 0, 2) or self.b_desc.stride_order != (1, 0, 2), - "Invalid A or B tensor stride: fp4 dtype requires k-major layout", - ) - self._value_error_if( - self.ab12_desc.stride_order != (1, 0, 2), - "Invalid AB12 tensor stride: fp4 dtype requires n-major layout", - ) - self._check_dtype( - self.b_desc, - dtype=self.ab_dtype, - name="B", - extra_error_msg="A and B must have the same dtype", - ) self._logger.debug("Checking MMA tile shape and cluster shape") - self._value_error_if( - self.mma_tiler_mn[0] not in [128, 256], - f"Invalid MMA tile shape: expected mma_tiler_mn[0] in {{128, 256}}, got {self.mma_tiler_mn[0]}", - ) - if self._kernel is PersistentDenseGemmKernel: - self._value_error_if( - self.mma_tiler_mn[1] not in range(32, 257, 32), - f"Invalid MMA tile shape: expected mma_tiler_mn[1] in {{32, 64, ..., 224, 256}}, got {self.mma_tiler_mn[1]}", - ) - - elif self._kernel is Sm100BlockScaledPersistentDenseGemmKernel: - if self._is_fp4x2(self.ab_dtype): - self._value_error_if( - self.mma_tiler_mn[1] not in range(64, 257, 64), - f"Invalid MMA tile shape: expected mma_tiler_mn[1] in {{64, 128, 192, 256}}, got {self.mma_tiler_mn[1]}", - ) - else: - if self._is_fp8(self.ab_dtype): - self._value_error_if( - self._is_fp8(self.c_dtype) or self._is_fp8(self.ab12_dtype) or self.ab12_dtype == torch.float32, - "For MXFP8 inputs for blockscaled quantized GEMM swiglu kernel, ab12_dtype and c_dtype cannot be FP8. ab12_dtype also cannot be float32", - ) - - self._value_error_if( - self.cluster_shape_mn[0] % (2 if self.mma_tiler_mn[0] == 256 else 1) != 0, - "Invalid cluster shape: cluster_shape_mn[0] must be divisible by 2 if mma_tiler_mn[0] == 256", - ) - self._value_error_if( - not ( - self.cluster_shape_mn[0] * self.cluster_shape_mn[1] <= 16 - 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 values to be powers of 2 and cluster_shape_mn[0] * cluster_shape_mn[1] <= 16, got {self.cluster_shape_mn[0]},{self.cluster_shape_mn[1]}", - ) - - if self._kernel is PersistentDenseGemmKernel: - use_2cta_instrs = self.mma_tiler_mn[0] == 256 - self._value_error_if( - not use_2cta_instrs and self.cluster_shape_mn != (1, 1), - "Invalid cluster shape: cluster_shape must be (1, 1) when use_2cta_instrs=False", - ) - if self.cluster_shape_mn != (1, 1) and self.mma_tiler_mn[0] == 128: - self._value_error_if( - self.mma_tiler_mn != (128, 128), - "Invalid MMA tile shape: for non-1x1 cluster shape and 128xmma tile shape, mma_tiler_mn must be (128, 128)", - ) - - self._logger.debug("Checking tensor alignment") - - def check_contigous_16B_alignment(dtype, stride_order, tensor_shape): - is_mode0_major = stride_order == (0, 1, 2) - 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, interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2).width) - return num_major_elements % num_contiguous_elements == 0 - - self._value_error_if( - not ( - check_contigous_16B_alignment(self.ab_dtype, self.a_desc.stride_order, (m, k, l)) - and check_contigous_16B_alignment(self.ab_dtype, self.b_desc.stride_order, (n, k, l)) - and check_contigous_16B_alignment(self.ab12_dtype, self.ab12_desc.stride_order, (m, n, l)) - ), - "Invalid tensor alignment: tensors must be 16B aligned", + self.mma_tiler_mn = self._kernel.require_mma_tiler(self.mma_tiler_mn) + self.cluster_shape_mn = self._kernel.require_cluster_shape( + self.cluster_shape_mn, + mma_tiler_mn=self.mma_tiler_mn, ) if self._kernel is Sm100BlockScaledPersistentDenseGemmKernel: - self._value_error_if( - m % self.mma_tiler_mn[0] != 0 or n % self.mma_tiler_mn[1] != 0, - "Invalid tensor alignment: m and n must be aligned to mma_tiler_mn", + plan = validate_quantized_gemm_swiglu( + self.a_desc, + self.b_desc, + self.ab12_desc, + self.c_desc, + sfa=self.sfa_desc, + sfb=self.sfb_desc, + amax=self.amax_desc, + sfc=self.sfc_desc, + norm_const=self.norm_const_desc, + acc_dtype=self.acc_dtype, + output_n=n_2, + sf_vec_size=self.sf_vec_size, + supported_sf_vec_sizes=self._kernel.SF_VEC_SIZES, + mma_tiler_mn=self.mma_tiler_mn, ) + self.ab_dtype = self.a_desc.dtype + self.ab12_dtype = self.ab12_desc.dtype + self.c_dtype = self.c_desc.dtype + self.sf_dtype = self.sfa_desc.dtype + self._logger.debug("Resolved quantized GEMM + SwiGLU plan: %s", plan) + else: + self._logger.debug("Checking tensor alignment") + ab_bits = _convert_to_cutlass_data_type( + self.ab_dtype, + interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2, + ).width + ab12_bits = _convert_to_cutlass_data_type( + self.ab12_dtype, + interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2, + ).width + c_bits = _convert_to_cutlass_data_type( + self.c_dtype, + interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2, + ).width + require_contiguous_alignment("A", m if self.a_desc.stride_order == (0, 1, 2) else k, ab_bits) + require_contiguous_alignment("B", n if self.b_desc.stride_order == (0, 1, 2) else k, ab_bits) + require_contiguous_alignment("AB12", m if self.ab12_desc.stride_order == (0, 1, 2) else n, ab12_bits) + require_contiguous_alignment("C", m if self.c_desc.stride_order == (0, 1, 2) else n_2, c_bits) self._logger.debug("Checking environment") if not torch.cuda.is_available(): @@ -396,7 +264,7 @@ def compile(self) -> None: if self._kernel is PersistentDenseGemmKernel: gemm_swiglu = self._kernel( acc_dtype=_convert_to_cutlass_data_type(self.acc_dtype), - use_2cta_instrs=(self.mma_tiler_mn[0] == 256), + use_2cta_instrs=(self.mma_tiler_mn[0] == self._kernel.TWO_CTA_MMA_TILER_M), mma_tiler_mn=self.mma_tiler_mn, cluster_shape_mn=self.cluster_shape_mn, ) @@ -412,11 +280,9 @@ def compile(self) -> None: raise NotImplementedError(f"Unreachable: invalid kernel type {self._kernel}") hardware_info = cutlass.utils.HardwareInfo() - max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) - 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 = resolve_max_active_clusters( + hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]), + self.num_cluster_overlap_margin, ) fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) diff --git a/python/cudnn/gemm_swiglu/dense_blockscaled_gemm_persistent_swiglu_interleaved_quant.py b/python/cudnn/gemm_swiglu/dense_blockscaled_gemm_persistent_swiglu_interleaved_quant.py index 71d225b4f..54276c22a 100644 --- a/python/cudnn/gemm_swiglu/dense_blockscaled_gemm_persistent_swiglu_interleaved_quant.py +++ b/python/cudnn/gemm_swiglu/dense_blockscaled_gemm_persistent_swiglu_interleaved_quant.py @@ -42,6 +42,9 @@ import cutlass.cute.math as math from cutlass.cute.typing import Float32 +from ..gemm_validation import require_cluster_shape as _require_cluster_shape +from ..gemm_validation import require_mma_tiler as _require_mma_tiler + def sigmoid_f32(a: Union[float, Float32], fastmath: bool = False) -> Union[float, Float32]: """ @@ -153,6 +156,52 @@ class Sm100BlockScaledPersistentDenseGemmKernel: >>> gemm(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, max_active_clusters, stream) """ + # Configuration values supported by the FE Torch wrapper. + MMA_TILER_M = (128, 256) + MMA_TILER_N = (64, 128, 192, 256) + TWO_CTA_MMA_TILER_M = 256 + MAX_CLUSTER_CTAS = 16 + MAX_CLUSTER_DIMENSION = 4 + SF_VEC_SIZES = (16, 32) + SWIGLU_BLOCK_COLUMNS = 32 + SWIGLU_BLOCKS_PER_PAIR = 2 + + @classmethod + def require_mma_tiler(cls, mma_tiler_mn: Tuple[int, int]) -> Tuple[int, int]: + """Validate an FE-supported MMA tile.""" + + return _require_mma_tiler( + mma_tiler_mn, + allowed_m=cls.MMA_TILER_M, + allowed_n=cls.MMA_TILER_N, + ) + + @classmethod + def require_cluster_shape( + cls, + cluster_shape_mn: Tuple[int, int], + *, + mma_tiler_mn: Tuple[int, int], + ) -> Tuple[int, int]: + """Validate an FE-supported cluster shape for an MMA tile.""" + + return _require_cluster_shape( + cluster_shape_mn, + mma_m=mma_tiler_mn[0], + two_cta_mma_m=cls.TWO_CTA_MMA_TILER_M, + max_ctas=cls.MAX_CLUSTER_CTAS, + max_dimension=cls.MAX_CLUSTER_DIMENSION, + ) + + @classmethod + def get_output_n(cls, n: int) -> int: + """Validate the input N dimension and return the SwiGLU output N.""" + + required_columns = cls.SWIGLU_BLOCK_COLUMNS * cls.SWIGLU_BLOCKS_PER_PAIR + if n % required_columns: + raise ValueError(f"N must be divisible by {required_columns} for {cls.SWIGLU_BLOCK_COLUMNS}-column SwiGLU block pairs, got {n}") + return n // cls.SWIGLU_BLOCKS_PER_PAIR + def __init__( self, sf_vec_size: int, @@ -185,7 +234,7 @@ def __init__( self.acc_dtype = cutlass.Float32 self.sf_vec_size = sf_vec_size - self.use_2cta_instrs = mma_tiler_mn[0] == 256 + self.use_2cta_instrs = mma_tiler_mn[0] == self.TWO_CTA_MMA_TILER_M self.cluster_shape_mn = cluster_shape_mn self.ab12_stages = ab12_stages # K dimension is deferred in _setup_attributes diff --git a/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py b/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py index 28f6aff86..3727bcf0b 100644 --- a/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py +++ b/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py @@ -38,6 +38,9 @@ import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.cute.math as math +from ..gemm_validation import require_cluster_shape as _require_cluster_shape +from ..gemm_validation import require_mma_tiler as _require_mma_tiler + # Mathematical constant: log2(e) for converting exp(x) to exp2(x * log2(e)) LOG2_E = 1.4426950408889634 @@ -78,7 +81,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 * 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 +128,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 - 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 @@ -139,6 +142,55 @@ class PersistentDenseGemmKernel: >>> gemm(a_tensor, b_tensor, c_tensor, max_active_clusters, stream) """ + # Configuration values supported by the FE Torch and JAX wrappers. + MMA_TILER_M = (128, 256) + MMA_TILER_N = (64, 128, 192, 256) + TWO_CTA_MMA_TILER_M = 256 + MAX_CLUSTER_CTAS = 16 + MAX_CLUSTER_DIMENSION = None + SINGLE_CTA_CLUSTER_SHAPE = (1, 1) + SWIGLU_BLOCK_COLUMNS = 32 + SWIGLU_BLOCKS_PER_PAIR = 2 + + @classmethod + def require_mma_tiler(cls, mma_tiler_mn: Tuple[int, int]) -> Tuple[int, int]: + """Validate an FE-supported MMA tile.""" + + return _require_mma_tiler( + mma_tiler_mn, + allowed_m=cls.MMA_TILER_M, + allowed_n=cls.MMA_TILER_N, + ) + + @classmethod + def require_cluster_shape( + cls, + cluster_shape_mn: Tuple[int, int], + *, + mma_tiler_mn: Tuple[int, int], + ) -> Tuple[int, int]: + """Validate an FE-supported cluster shape for an MMA tile.""" + + cluster_shape_mn = _require_cluster_shape( + cluster_shape_mn, + mma_m=mma_tiler_mn[0], + two_cta_mma_m=cls.TWO_CTA_MMA_TILER_M, + max_ctas=cls.MAX_CLUSTER_CTAS, + max_dimension=cls.MAX_CLUSTER_DIMENSION, + ) + if mma_tiler_mn[0] != cls.TWO_CTA_MMA_TILER_M and cluster_shape_mn != cls.SINGLE_CTA_CLUSTER_SHAPE: + raise ValueError(f"cluster_shape_mn must be {cls.SINGLE_CTA_CLUSTER_SHAPE} for a single-CTA MMA tile") + return cluster_shape_mn + + @classmethod + def get_output_n(cls, n: int) -> int: + """Validate the input N dimension and return the SwiGLU output N.""" + + required_columns = cls.SWIGLU_BLOCK_COLUMNS * cls.SWIGLU_BLOCKS_PER_PAIR + if n % required_columns: + raise ValueError(f"N must be divisible by {required_columns} for {cls.SWIGLU_BLOCK_COLUMNS}-column SwiGLU block pairs, got {n}") + return n // cls.SWIGLU_BLOCKS_PER_PAIR + def __init__( self, acc_dtype: Type[cutlass.Numeric], @@ -173,7 +225,7 @@ def __init__( """ self.acc_dtype: Type[cutlass.Numeric] = acc_dtype - self.use_2cta_instrs = mma_tiler_mn[0] == 256 + self.use_2cta_instrs = mma_tiler_mn[0] == self.TWO_CTA_MMA_TILER_M self.cluster_shape_mn = cluster_shape_mn # K dimension is deferred in _setup_attributes self.mma_tiler = (*mma_tiler_mn, 1) diff --git a/python/cudnn/gemm_swiglu/jax.py b/python/cudnn/gemm_swiglu/jax.py new file mode 100644 index 000000000..8e439e4b0 --- /dev/null +++ b/python/cudnn/gemm_swiglu/jax.py @@ -0,0 +1,524 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for dense GEMM + SwiGLU on SM100.""" + +from __future__ import annotations + +import os +from typing import Any, Optional + +import jax.numpy as jnp + +from .._jax.api_base import ( + ApiBaseJax, + BufferSpec, + JaxTensorDesc, + TupleDict, + call_cutedsl, +) +from .._jax.gemm import ( + block_scale_tensor_spec, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + require_16_byte_extent, + require_gemm_inputs, +) +from ..gemm_validation import ( + require_full_mma_rows, + resolve_max_active_clusters, +) +from .validation import validate_quantized_gemm_swiglu + + +def _launch( + stream, + a, + b, + ab12, + c, + *, + alpha: float, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + cluster_overlap_margin: int, +): + # These operations happen during CUDA lowering, not abstract evaluation. + 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(acc_dtype), + use_2cta_instrs=mma_tiler_mn[0] == PersistentDenseGemmKernel.TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + ) + max_active_clusters = resolve_max_active_clusters( + cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]), + cluster_overlap_margin, + ) + kernel( + a, + b, + ab12, + c, + cutlass.Float32(alpha), + max_active_clusters, + stream, + ) + + +def _launch_quantized( + stream, + a, + b, + sfa, + sfb, + ab12, + c, + *optional_outputs, + alpha: float, + sf_vec_size: int, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + vector_f32: bool, + ab12_stages: Optional[int], + has_amax: bool, + cluster_overlap_margin: int, +): + # These operations happen during CUDA lowering, not abstract evaluation. + import cutlass + + from .dense_blockscaled_gemm_persistent_swiglu_interleaved_quant import ( + Sm100BlockScaledPersistentDenseGemmKernel, + ) + + expected_optional_outputs = 1 if has_amax else 0 + if len(optional_outputs) != expected_optional_outputs: + raise RuntimeError("Unexpected quantized GEMM + SwiGLU output count: " f"expected {expected_optional_outputs}, got {len(optional_outputs)}") + amax = optional_outputs[0] if has_amax else None + + kernel = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size=sf_vec_size, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vector_f32=vector_f32, + ab12_stages=ab12_stages, + ) + max_active_clusters = resolve_max_active_clusters( + cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]), + cluster_overlap_margin, + ) + kernel( + a, + b, + sfa, + sfb, + c, + ab12, + amax, + None, + None, + cutlass.Float32(alpha), + max_active_clusters, + stream, + ) + + +class GemmSwigluSm100(ApiBaseJax): + """JAX GEMM + SwiGLU callable specialized from sample metadata.""" + + def __init__( + self, + sample_a: Any, + sample_b: Any, + alpha: float = 1.0, + c_layout: str = "LMN", + ab12_dtype: Any = None, + c_dtype: Any = None, + acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sample_sfa: Optional[Any] = None, + sample_sfb: Optional[Any] = None, + sample_norm_const: Optional[Any] = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + ab12_stages: Optional[int] = 4, + *, + a_layout: str = "LMK", + b_layout: str = "LNK", + ) -> None: + super().__init__() + scale_spec = block_scale_tensor_spec() + self.a_desc = self.make_tensor_desc(sample_a, tensor_spec=gemm_a_tensor_spec(a_layout), name="sample_a") + self.b_desc = self.make_tensor_desc(sample_b, tensor_spec=gemm_b_tensor_spec(b_layout), name="sample_b") + self.sfa_desc = self.make_optional_tensor_desc(sample_sfa, tensor_spec=scale_spec, name="sample_sfa") + self.sfb_desc = self.make_optional_tensor_desc(sample_sfb, tensor_spec=scale_spec, name="sample_sfb") + self.norm_const_desc = self.make_optional_tensor_desc(sample_norm_const, name="sample_norm_const") + self._c_layout = c_layout + + self.alpha = alpha + self._ab12_dtype = self.as_optional_dtype(ab12_dtype) + self._c_dtype = self.as_optional_dtype(c_dtype) + self._acc_dtype = self.as_optional_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.vector_f32 = bool(vector_f32) + self.ab12_stages = ab12_stages + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + self._quantized = False + self._quantized_plan = None + self.amax_desc = None + + def _check_support(self) -> None: + self._quantized = False + self._quantized_plan = None + self.amax_desc = None + + if self.norm_const_desc is not None: + raise NotImplementedError("The JAX GEMM + SwiGLU API does not yet support FP8 C output or " "sample_norm_const") + if (self.sfa_desc is None) != (self.sfb_desc is None): + raise ValueError("sample_sfa and sample_sfb must be provided together") + if self.sfa_desc is not None: + self._quantized = True + self._check_quantized_support() + return + + if self.sf_vec_size != 16: + raise NotImplementedError("sf_vec_size applies only when sample_sfa and sample_sfb select the block-scaled path") + if self.vector_f32: + raise NotImplementedError("vector_f32 applies only when sample_sfa and sample_sfb select the block-scaled path") + if self.ab12_stages != 4: + raise NotImplementedError("ab12_stages applies only when sample_sfa and sample_sfb select the block-scaled path") + + from .dense_gemm_persistent_swiglu import PersistentDenseGemmKernel + + kernel = PersistentDenseGemmKernel + self.m, self.n, self.k, self.batch, a_dtype = require_gemm_inputs(self.a_desc, self.b_desc) + self.output_n = kernel.get_output_n(self.n) + self.a_dtype = self.require_dtype( + a_dtype, + ( + jnp.float16, + jnp.bfloat16, + jnp.float32, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + ), + name="sample_a.dtype", + ) + self.acc_dtype = self.require_dtype( + self._acc_dtype, + (jnp.float32, jnp.float16), + name="acc_dtype", + default=jnp.float32, + ) + if self.acc_dtype == jnp.dtype(jnp.float32): + supported_ab12 = (jnp.float32, jnp.float16, jnp.bfloat16) + else: + supported_ab12 = (jnp.float16, jnp.bfloat16) + if self.a_dtype not in { + jnp.dtype(jnp.float16), + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + }: + raise ValueError(f"float16 accumulation does not support input dtype {self.a_dtype}") + self.ab12_dtype = self.require_dtype( + self._ab12_dtype, + supported_ab12, + name="ab12_dtype", + default=jnp.float32, + ) + self.c_dtype = self.require_dtype( + self._c_dtype, + (jnp.float16, jnp.bfloat16), + name="c_dtype", + default=jnp.float16, + ) + + self.mma_tiler_mn = kernel.require_mma_tiler(self.mma_tiler_mn) + if self.mma_tiler_mn[0] == kernel.TWO_CTA_MMA_TILER_M: + require_full_mma_rows( + self.m, + self.mma_tiler_mn[0], + reason="2-CTA MMA requires a complete CTA pair", + ) + if self.cluster_shape_mn is None: + self.cluster_shape_mn = (2, 2) if self.mma_tiler_mn[0] == kernel.TWO_CTA_MMA_TILER_M else (1, 1) + self.cluster_shape_mn = kernel.require_cluster_shape( + self.cluster_shape_mn, + mma_tiler_mn=self.mma_tiler_mn, + ) + + require_16_byte_extent("sample_a", self.a_desc.shape[self.a_desc.stride_order[0]], self.a_dtype) + require_16_byte_extent("sample_b", self.b_desc.shape[self.b_desc.stride_order[0]], self.a_dtype) + + output_spec = gemm_c_tensor_spec(self._c_layout) + self.ab12_desc = JaxTensorDesc( + dtype=self.ab12_dtype, + shape=(self.m, self.n, self.batch), + tensor_spec=output_spec, + name="ab12_tensor", + ) + self.c_desc = JaxTensorDesc( + dtype=self.c_dtype, + shape=(self.m, self.output_n, self.batch), + tensor_spec=output_spec, + name="c_tensor", + ) + require_16_byte_extent( + "ab12_tensor", + self.ab12_desc.shape[self.ab12_desc.stride_order[0]], + self.ab12_dtype, + ) + require_16_byte_extent("c_tensor", self.c_desc.shape[self.c_desc.stride_order[0]], self.c_dtype) + + def _check_quantized_support(self) -> None: + from .dense_blockscaled_gemm_persistent_swiglu_interleaved_quant import ( + Sm100BlockScaledPersistentDenseGemmKernel, + ) + + kernel = Sm100BlockScaledPersistentDenseGemmKernel + self.m, self.n, self.k, self.batch, a_dtype = require_gemm_inputs(self.a_desc, self.b_desc) + self.output_n = kernel.get_output_n(self.n) + self.a_dtype = self.require_dtype( + a_dtype, + (jnp.float4_e2m1fn, jnp.float8_e4m3fn, jnp.float8_e5m2), + name="sample_a.dtype", + ) + self.acc_dtype = self.require_dtype( + self._acc_dtype, + (jnp.float32,), + name="acc_dtype", + default=jnp.float32, + ) + self.ab12_dtype = self.require_dtype( + self._ab12_dtype, + ( + jnp.float32, + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + ), + name="ab12_dtype", + default=jnp.float32, + ) + self.c_dtype = self.require_dtype( + self._c_dtype, + (jnp.float32, jnp.float16, jnp.bfloat16), + name="c_dtype", + default=jnp.float16, + ) + + self.mma_tiler_mn = kernel.require_mma_tiler(self.mma_tiler_mn) + if self.cluster_shape_mn is None: + self.cluster_shape_mn = (2, 2) if self.mma_tiler_mn[0] == kernel.TWO_CTA_MMA_TILER_M else (1, 1) + self.cluster_shape_mn = kernel.require_cluster_shape( + self.cluster_shape_mn, + mma_tiler_mn=self.mma_tiler_mn, + ) + + output_spec = gemm_c_tensor_spec(self._c_layout) + self.ab12_desc = JaxTensorDesc( + dtype=self.ab12_dtype, + shape=(self.m, self.n, self.batch), + tensor_spec=output_spec, + name="ab12_tensor", + ) + self.c_desc = JaxTensorDesc( + dtype=self.c_dtype, + shape=(self.m, self.output_n, self.batch), + tensor_spec=output_spec, + name="c_tensor", + ) + if self.a_desc.dtype_name == "float4_e2m1fn" and self.c_desc.dtype_name == "bfloat16": + self.amax_desc = JaxTensorDesc( + dtype=jnp.float32, + shape=(1, 1, 1), + name="amax_tensor", + ) + + self._quantized_plan = validate_quantized_gemm_swiglu( + self.a_desc, + self.b_desc, + self.ab12_desc, + self.c_desc, + sfa=self.sfa_desc, + sfb=self.sfb_desc, + amax=self.amax_desc, + sfc=None, + norm_const=None, + acc_dtype=self.acc_dtype, + output_n=self.output_n, + sf_vec_size=self.sf_vec_size, + supported_sf_vec_sizes=kernel.SF_VEC_SIZES, + mma_tiler_mn=self.mma_tiler_mn, + ) + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Optional[Any] = None, + sfb_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__(a_tensor, b_tensor, sfa_tensor, sfb_tensor, norm_const_tensor) + + def _call_impl( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Optional[Any], + sfb_tensor: Optional[Any], + norm_const_tensor: Optional[Any], + ) -> TupleDict: + self.check_tensor_signature(a_tensor, self.a_desc, name="A") + self.check_tensor_signature(b_tensor, self.b_desc, name="B") + self.check_optional_tensor_signature(sfa_tensor, self.sfa_desc, name="SFA") + self.check_optional_tensor_signature(sfb_tensor, self.sfb_desc, name="SFB") + self.check_optional_tensor_signature(norm_const_tensor, self.norm_const_desc, name="norm_const") + + if self._quantized: + if self._quantized_plan is None or self.sfa_desc is None or self.sfb_desc is None: + raise RuntimeError("check_support() did not produce a quantized launch plan") + + outputs = [ + BufferSpec( + "ab12_tensor", + self.ab12_desc.array_shape, + self.ab12_desc.dtype, + tensor_spec=self.ab12_desc.tensor_spec, + ), + BufferSpec( + "c_tensor", + self.c_desc.array_shape, + self.c_desc.dtype, + tensor_spec=self.c_desc.tensor_spec, + ), + ] + if self.amax_desc is not None: + outputs.append( + BufferSpec( + "amax_tensor", + self.amax_desc.array_shape, + self.amax_desc.dtype, + fill_value=-float("inf"), + ) + ) + + results = call_cutedsl( + _launch_quantized, + (a_tensor, b_tensor, sfa_tensor, sfb_tensor), + outputs=outputs, + input_specs=( + self.a_desc.tensor_spec, + self.b_desc.tensor_spec, + self.sfa_desc.tensor_spec, + self.sfb_desc.tensor_spec, + ), + static_args={ + "alpha": float(self.alpha), + "sf_vec_size": self.sf_vec_size, + "mma_tiler_mn": self.mma_tiler_mn, + "cluster_shape_mn": self.cluster_shape_mn, + "vector_f32": self.vector_f32, + "ab12_stages": self.ab12_stages, + "has_amax": self.amax_desc is not None, + "cluster_overlap_margin": self.num_cluster_overlap_margin, + }, + ) + ab12_tensor, c_tensor = results[:2] + amax_tensor = results[2] 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, + ) + + ab12_tensor, c_tensor = call_cutedsl( + _launch, + (a_tensor, b_tensor), + outputs=( + BufferSpec( + "ab12_tensor", + self.ab12_desc.array_shape, + self.ab12_desc.dtype, + tensor_spec=self.ab12_desc.tensor_spec, + ), + BufferSpec( + "c_tensor", + self.c_desc.array_shape, + self.c_desc.dtype, + tensor_spec=self.c_desc.tensor_spec, + ), + ), + input_specs=(self.a_desc.tensor_spec, self.b_desc.tensor_spec), + static_args={ + "alpha": float(self.alpha), + "acc_dtype": self.acc_dtype, + "mma_tiler_mn": self.mma_tiler_mn, + "cluster_shape_mn": self.cluster_shape_mn, + "cluster_overlap_margin": self.num_cluster_overlap_margin, + }, + ) + return TupleDict( + ab12_tensor=ab12_tensor, + c_tensor=c_tensor, + sfc_tensor=None, + amax_tensor=None, + ) + + +def gemm_swiglu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + alpha: float = 1.0, + c_layout: str = "LMN", + ab12_dtype: Any = None, + c_dtype: Any = None, + acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sfa_tensor: Optional[Any] = None, + sfb_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + ab12_stages: Optional[int] = 4, + *, + a_layout: str = "LMK", + b_layout: str = "LNK", +) -> TupleDict: + """Compute a dense batched GEMM and its fused SwiGLU projection.""" + + return GemmSwigluSm100( + a_tensor, + b_tensor, + 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, + sample_sfa=sfa_tensor, + sample_sfb=sfb_tensor, + sample_norm_const=norm_const_tensor, + 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, sfa_tensor, sfb_tensor, norm_const_tensor) + + +__all__ = ["GemmSwigluSm100", "gemm_swiglu_wrapper_sm100"] diff --git a/python/cudnn/gemm_swiglu/validation.py b/python/cudnn/gemm_swiglu/validation.py new file mode 100644 index 000000000..3e02dbde4 --- /dev/null +++ b/python/cudnn/gemm_swiglu/validation.py @@ -0,0 +1,266 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral validation for block-scaled GEMM + SwiGLU.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Sequence + +from ..api_base import TensorDesc, canonical_dtype_name +from ..gemm_validation import ( + block_scale_shape, + require_block_scale_shapes, + require_contiguous_alignment, + require_gemm_shapes, + require_shape, +) + +_AB_DTYPES = frozenset({"float4_e2m1fn", "float8_e4m3fn", "float8_e5m2"}) +_SCALE_DTYPES = frozenset({"float8_e4m3fn", "float8_e8m0fnu"}) +_OUTPUT_DTYPES = frozenset( + { + "bfloat16", + "float16", + "float32", + "float8_e4m3fn", + "float8_e5m2", + } +) + + +@dataclass(frozen=True) +class QuantizedGemmSwigluPlan: + """Validated logical signature for block-scaled GEMM + SwiGLU.""" + + m: int + n: int + k: int + batch: int + output_n: int + ab_dtype_name: str + scale_dtype_name: str + ab12_dtype_name: str + c_dtype_name: str + a_major: str + b_major: str + output_major: str + generate_amax: bool + generate_sfc: bool + + @property + def ab12_shape(self) -> tuple[int, int, int]: + return (self.m, self.n, self.batch) + + @property + def c_shape(self) -> tuple[int, int, int]: + return (self.m, self.output_n, self.batch) + + @property + def requires_amax(self) -> bool: + return self.ab_dtype_name == "float4_e2m1fn" and self.c_dtype_name == "bfloat16" + + +def _logical_dtype_name(desc: TensorDesc) -> str: + if getattr(desc, "packing", "native") == "fp4x2": + return "float4_e2m1fn" + return desc.dtype_name + + +def _logical_element_bits(desc: TensorDesc) -> int: + if _logical_dtype_name(desc) == "float4_e2m1fn": + return 4 + if desc.element_bits is None: + raise ValueError(f"Cannot determine the element width of {desc.name or desc.dtype_name}") + return desc.element_bits + + +def _require_supported_dtype(name: str, actual: str, supported: frozenset[str]) -> None: + if actual not in supported: + choices = ", ".join(sorted(supported)) + raise ValueError(f"{name} dtype must be one of {{{choices}}}, got {actual}") + + +def _require_same_storage_dtype(name: str, lhs: TensorDesc, rhs: TensorDesc) -> None: + lhs_storage = getattr(lhs, "storage_dtype_name", lhs.dtype_name) + rhs_storage = getattr(rhs, "storage_dtype_name", rhs.dtype_name) + lhs_packing = getattr(lhs, "packing", "native") + rhs_packing = getattr(rhs, "packing", "native") + if (lhs_storage, lhs_packing) != (rhs_storage, rhs_packing): + lhs_label = lhs_storage if lhs_packing == "native" else f"{lhs_storage}/{lhs_packing}" + rhs_label = rhs_storage if rhs_packing == "native" else f"{rhs_storage}/{rhs_packing}" + raise ValueError(f"{name} dtypes must match, got {lhs_label} and {rhs_label}") + + +def _require_compact_major( + name: str, + desc: TensorDesc, + layouts: dict[str, tuple[tuple[int, ...], tuple[int, ...]]], +) -> str: + for major, (stride, stride_order) in layouts.items(): + if desc.stride == stride and desc.stride_order == stride_order: + return major + expected = [stride for stride, _ in layouts.values()] + raise ValueError(f"{name} tensor must use one of the compact strides {expected}, got " f"stride {desc.stride} with order {desc.stride_order}") + + +def _require_singleton(name: str, desc: TensorDesc) -> None: + if tuple(desc.shape) not in {(1,), (1, 1, 1)}: + raise ValueError(f"{name} must have shape (1,) or (1, 1, 1), got {tuple(desc.shape)}") + + +def validate_quantized_gemm_swiglu( + a: TensorDesc, + b: TensorDesc, + ab12: TensorDesc, + c: TensorDesc, + *, + sfa: TensorDesc | None, + sfb: TensorDesc | None, + amax: TensorDesc | None, + sfc: TensorDesc | None, + norm_const: TensorDesc | None, + acc_dtype: Any, + output_n: int, + sf_vec_size: int, + supported_sf_vec_sizes: Sequence[int], + mma_tiler_mn: Sequence[int], +) -> QuantizedGemmSwigluPlan: + """Validate the block-scaled contract shared by Torch and JAX adapters.""" + + if sfa is None or sfb is None: + raise ValueError("sfa_tensor and sfb_tensor are required for quantized GEMM + SwiGLU") + + m, n, k, batch = require_gemm_shapes(a.shape, b.shape) + require_shape("AB12", ab12.shape, (m, n, batch)) + require_shape("C", c.shape, (m, output_n, batch)) + + _require_same_storage_dtype("a_tensor and b_tensor", a, b) + ab_dtype_name = _logical_dtype_name(a) + _require_supported_dtype("a_tensor and b_tensor", ab_dtype_name, _AB_DTYPES) + if canonical_dtype_name(acc_dtype) != "float32": + raise ValueError(f"Accumulator dtype must be float32, got {canonical_dtype_name(acc_dtype)}") + + supported_sf_vec_sizes = tuple(supported_sf_vec_sizes) + if sf_vec_size not in supported_sf_vec_sizes: + raise ValueError(f"sf_vec_size must be one of {supported_sf_vec_sizes}, got {sf_vec_size}") + + _require_same_storage_dtype("sfa_tensor and sfb_tensor", sfa, sfb) + scale_dtype_name = _logical_dtype_name(sfa) + _require_supported_dtype("sfa_tensor and sfb_tensor", scale_dtype_name, _SCALE_DTYPES) + require_block_scale_shapes( + sfa.shape, + sfb.shape, + m=m, + n=n, + k=k, + batch=batch, + sf_vec_size=sf_vec_size, + ) + + ab12_dtype_name = _logical_dtype_name(ab12) + c_dtype_name = _logical_dtype_name(c) + _require_supported_dtype("AB12", ab12_dtype_name, _OUTPUT_DTYPES) + _require_supported_dtype("C", c_dtype_name, _OUTPUT_DTYPES) + + if sfc is not None: + _require_same_storage_dtype("sfa_tensor and sfc_tensor", sfa, sfc) + require_shape( + "sfc_tensor", + sfc.shape, + block_scale_shape(m, output_n, batch, sf_vec_size), + ) + if norm_const is not None: + require_shape("norm_const_tensor", norm_const.shape, (1,)) + if amax is not None: + _require_singleton("amax_tensor", amax) + if _logical_dtype_name(amax) != "float32": + raise ValueError(f"amax_tensor dtype must be float32, got {_logical_dtype_name(amax)}") + + is_fp4 = ab_dtype_name == "float4_e2m1fn" + is_fp8 = ab_dtype_name.startswith("float8_") + c_is_fp8 = c_dtype_name.startswith("float8_") + ab12_is_fp8 = ab12_dtype_name.startswith("float8_") + + if is_fp4 and c_is_fp8: + raise ValueError("FP4 A and B are not compatible with an FP8 C dtype") + if c_is_fp8 and (sfc is None or norm_const is None): + raise ValueError("sfc_tensor and norm_const_tensor are required when C is FP8") + if is_fp4 and c_dtype_name == "bfloat16" and amax is None: + raise ValueError("amax_tensor is required when A and B are FP4 and C is bfloat16") + if c_dtype_name == "float32" and ab12_dtype_name == "float32": + raise NotImplementedError("float32 C with float32 AB12 is disabled because of a kernel bug") + + if is_fp8 and not (scale_dtype_name == "float8_e8m0fnu" and sf_vec_size == 32): + raise ValueError("FP8 A and B require float8_e8m0fnu scale factors and sf_vec_size=32") + if is_fp4 and scale_dtype_name == "float8_e4m3fn" and sf_vec_size == 32: + raise ValueError("FP4 A and B do not support float8_e4m3fn scale factors with sf_vec_size=32") + if is_fp8 and (c_is_fp8 or ab12_is_fp8 or ab12_dtype_name == "float32"): + raise ValueError("MXFP8 requires float16 or bfloat16 AB12 and a non-FP8 C dtype") + + a_major = _require_compact_major( + "A", + a, + { + "m": ((1, m, m * k), (0, 1, 2)), + "k": ((k, 1, m * k), (1, 0, 2)), + }, + ) + b_major = _require_compact_major( + "B", + b, + { + "n": ((1, n, n * k), (0, 1, 2)), + "k": ((k, 1, n * k), (1, 0, 2)), + }, + ) + ab12_major = _require_compact_major( + "AB12", + ab12, + { + "m": ((1, m, m * n), (0, 1, 2)), + "n": ((n, 1, m * n), (1, 0, 2)), + }, + ) + c_major = _require_compact_major( + "C", + c, + { + "m": ((1, m, m * output_n), (0, 1, 2)), + "n": ((output_n, 1, m * output_n), (1, 0, 2)), + }, + ) + if ab12_major != c_major: + raise ValueError(f"AB12 and C must use the same major mode, got {ab12_major} and {c_major}") + if is_fp4 and (a_major, b_major, ab12_major) != ("k", "k", "n"): + raise ValueError("FP4 requires k-major A and B and n-major outputs, got " f"{a_major}-major, {b_major}-major, and {ab12_major}-major") + + mma_m, mma_n = tuple(mma_tiler_mn) + if m % mma_m or n % mma_n: + raise ValueError(f"M and N must be divisible by mma_tiler_mn {tuple(mma_tiler_mn)}, got M={m}, N={n}") + + require_contiguous_alignment("A", m if a_major == "m" else k, _logical_element_bits(a)) + require_contiguous_alignment("B", n if b_major == "n" else k, _logical_element_bits(b)) + require_contiguous_alignment("AB12", m if ab12_major == "m" else n, _logical_element_bits(ab12)) + require_contiguous_alignment("C", m if c_major == "m" else output_n, _logical_element_bits(c)) + + return QuantizedGemmSwigluPlan( + m=m, + n=n, + k=k, + batch=batch, + output_n=output_n, + ab_dtype_name=ab_dtype_name, + scale_dtype_name=scale_dtype_name, + ab12_dtype_name=ab12_dtype_name, + c_dtype_name=c_dtype_name, + a_major=a_major, + b_major=b_major, + output_major=ab12_major, + generate_amax=amax is not None, + generate_sfc=sfc is not None and norm_const is not None, + ) + + +__all__ = ["QuantizedGemmSwigluPlan", "validate_quantized_gemm_swiglu"] diff --git a/python/cudnn/gemm_validation.py b/python/cudnn/gemm_validation.py new file mode 100644 index 000000000..9d2fbda6a --- /dev/null +++ b/python/cudnn/gemm_validation.py @@ -0,0 +1,181 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral validation rules shared by dense GEMM frontends.""" + +from __future__ import annotations + +from collections.abc import Collection, Sequence +from typing import Optional + + +def ceil_div(value: int, divisor: int) -> int: + """Return ``ceil(value / divisor)`` for positive integers.""" + + if divisor <= 0: + raise ValueError(f"divisor must be positive, got {divisor}") + return (value + divisor - 1) // divisor + + +def require_shape(name: str, actual: Sequence[int], expected: Sequence[int]) -> tuple[int, ...]: + """Require an exact logical tensor shape and return it as a tuple.""" + + actual = tuple(actual) + expected = tuple(expected) + if actual != expected: + raise ValueError(f"{name} must have shape {expected}, got {actual}") + return actual + + +def require_gemm_shapes( + a_shape: Sequence[int], + b_shape: Sequence[int], +) -> tuple[int, int, int, int]: + """Validate dense GEMM input shapes and return ``M, N, K, L``.""" + + a_shape = tuple(a_shape) + b_shape = tuple(b_shape) + if len(a_shape) != 3: + raise ValueError(f"a_tensor must have rank 3, got shape {a_shape}") + if len(b_shape) != 3: + raise ValueError(f"b_tensor must have rank 3, got shape {b_shape}") + + m, k, batch = a_shape + n, b_k, b_batch = b_shape + if (b_k, b_batch) != (k, batch): + raise ValueError("a_tensor and b_tensor must have matching K and L dimensions, " f"got {a_shape} and {b_shape}") + + dimensions = {"M": m, "N": n, "K": k, "L": batch} + nonpositive = [f"{name}={value}" for name, value in dimensions.items() if value <= 0] + if nonpositive: + raise ValueError("GEMM dimensions must be positive, got " + ", ".join(nonpositive)) + return m, n, k, batch + + +def block_scale_shape( + rows: int, + k: int, + batch: int, + sf_vec_size: int, +) -> tuple[int, int, int, int, int, int]: + """Return the six-dimensional native block-scale shape.""" + + if sf_vec_size <= 0: + raise ValueError(f"sf_vec_size must be positive, got {sf_vec_size}") + scale_k_tiles = ceil_div(ceil_div(k, sf_vec_size), 4) + return (32, 4, ceil_div(rows, 128), 4, scale_k_tiles, batch) + + +def require_block_scale_shapes( + sfa_shape: Sequence[int], + sfb_shape: Sequence[int], + *, + m: int, + n: int, + k: int, + batch: int, + sf_vec_size: int, + sfa_name: str = "sfa_tensor", + sfb_name: str = "sfb_tensor", +) -> None: + """Validate the A and B block-scale tensor shapes.""" + + require_shape(sfa_name, sfa_shape, block_scale_shape(m, k, batch, sf_vec_size)) + require_shape(sfb_name, sfb_shape, block_scale_shape(n, k, batch, sf_vec_size)) + + +def _format_values(values: Sequence[int]) -> str: + return "{" + ", ".join(str(value) for value in values) + "}" + + +def require_mma_tiler( + mma_tiler_mn: Sequence[int], + *, + allowed_m: Collection[int], + allowed_n: Collection[int], +) -> tuple[int, int]: + """Validate an MMA tile against explicit M and N domains.""" + + mma_tiler_mn = tuple(mma_tiler_mn) + allowed_m = tuple(sorted(allowed_m)) + allowed_n = tuple(sorted(allowed_n)) + if len(mma_tiler_mn) != 2 or mma_tiler_mn[0] not in allowed_m or mma_tiler_mn[1] not in allowed_n: + raise ValueError("mma_tiler_mn must have M in " f"{_format_values(allowed_m)} and N in {_format_values(allowed_n)}, " f"got {mma_tiler_mn}") + return mma_tiler_mn + + +def is_power_of_two(value: int) -> bool: + """Return whether ``value`` is a positive power of two.""" + + return value > 0 and value & (value - 1) == 0 + + +def require_cluster_shape( + cluster_shape_mn: Sequence[int], + *, + mma_m: int, + two_cta_mma_m: int, + max_ctas: int, + max_dimension: Optional[int] = None, +) -> tuple[int, int]: + """Validate a cluster shape and its 2-CTA MMA requirement.""" + + cluster_shape_mn = tuple(cluster_shape_mn) + if len(cluster_shape_mn) != 2: + raise ValueError(f"cluster_shape_mn must have two dimensions, got {cluster_shape_mn}") + + cluster_m, cluster_n = cluster_shape_mn + dimensions_valid = is_power_of_two(cluster_m) and is_power_of_two(cluster_n) + if max_dimension is not None: + dimensions_valid = dimensions_valid and cluster_m <= max_dimension and cluster_n <= max_dimension + if not dimensions_valid or cluster_m * cluster_n > max_ctas: + dimension_limit = "" if max_dimension is None else f", each at most {max_dimension}" + raise ValueError( + "cluster_shape_mn dimensions must be positive powers of two" f"{dimension_limit} with product at most {max_ctas}, got {cluster_shape_mn}" + ) + if mma_m == two_cta_mma_m and cluster_m % 2: + raise ValueError("cluster_shape_mn[0] must be divisible by 2 with a " f"{two_cta_mma_m}-wide M tile") + return cluster_shape_mn + + +def require_contiguous_alignment( + name: str, + elements: int, + element_bits: int, + *, + alignment_bytes: int = 16, +) -> None: + """Require a contiguous tensor extent to meet a byte alignment.""" + + if element_bits <= 0: + raise ValueError(f"element_bits must be positive, got {element_bits}") + if alignment_bytes <= 0: + raise ValueError(f"alignment_bytes must be positive, got {alignment_bytes}") + if elements * element_bits % (alignment_bytes * 8): + raise ValueError(f"{name}'s contiguous extent must be {alignment_bytes}-byte aligned, " f"got {elements} elements at {element_bits} bits each") + + +def require_full_mma_rows( + m: int, + mma_m: int, + *, + cta_group_size: int = 1, + reason: str = "", +) -> None: + """Require M to contain complete per-CTA rows for an MMA tile.""" + + if cta_group_size <= 0 or mma_m % cta_group_size: + raise ValueError(f"cta_group_size must divide mma_m, got {cta_group_size} and {mma_m}") + rows_per_cta = mma_m // cta_group_size + if m % rows_per_cta: + suffix = "" if not reason else f" because {reason}" + raise ValueError(f"M must be divisible by {rows_per_cta} (CTA_M for TILE_M={mma_m})" f"{suffix}, got {m}") + + +def resolve_max_active_clusters(max_active_clusters: int, overlap_margin: int) -> int: + """Apply the configured cluster overlap margin and validate the result.""" + + resolved = max_active_clusters - overlap_margin + if resolved <= 0: + raise ValueError("max_active_clusters must be positive after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") + return resolved diff --git a/python/cudnn/grouped_gemm/__init__.py b/python/cudnn/grouped_gemm/__init__.py index 818a1ab99..c52e70425 100644 --- a/python/cudnn/grouped_gemm/__init__.py +++ b/python/cudnn/grouped_gemm/__init__.py @@ -1,68 +1,41 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .grouped_gemm_swiglu.api import ( - GroupedGemmSwigluSm100, - grouped_gemm_swiglu_wrapper_sm100, -) - -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", -] +"""Lazy exports for grouped GEMM operation packages.""" + +from importlib import import_module + +_SYMBOLS = { + "GroupedGemmSwigluSm100": (".grouped_gemm_swiglu", "GroupedGemmSwigluSm100"), + "grouped_gemm_swiglu_wrapper_sm100": (".grouped_gemm_swiglu", "grouped_gemm_swiglu_wrapper_sm100"), + "GroupedGemmDswigluSm100": (".grouped_gemm_dswiglu", "GroupedGemmDswigluSm100"), + "grouped_gemm_dswiglu_wrapper_sm100": (".grouped_gemm_dswiglu", "grouped_gemm_dswiglu_wrapper_sm100"), + "GroupedGemmQuantSm100": (".grouped_gemm_quant", "GroupedGemmQuantSm100"), + "grouped_gemm_quant_wrapper_sm100": (".grouped_gemm_quant", "grouped_gemm_quant_wrapper_sm100"), + "GroupedGemmSreluSm100": (".grouped_gemm_srelu", "GroupedGemmSreluSm100"), + "grouped_gemm_srelu_wrapper_sm100": (".grouped_gemm_srelu", "grouped_gemm_srelu_wrapper_sm100"), + "GroupedGemmDsreluSm100": (".grouped_gemm_dsrelu", "GroupedGemmDsreluSm100"), + "grouped_gemm_dsrelu_wrapper_sm100": (".grouped_gemm_dsrelu", "grouped_gemm_dsrelu_wrapper_sm100"), + "GroupedGemmGluSm100": (".grouped_gemm_glu", "GroupedGemmGluSm100"), + "grouped_gemm_glu_wrapper_sm100": (".grouped_gemm_glu", "grouped_gemm_glu_wrapper_sm100"), + "GroupedGemmGluHadamardSm100": (".grouped_gemm_glu_hadamard", "GroupedGemmGluHadamardSm100"), + "grouped_gemm_glu_hadamard_wrapper_sm100": (".grouped_gemm_glu_hadamard", "grouped_gemm_glu_hadamard_wrapper_sm100"), + "GroupedGemmDgluSm100": (".grouped_gemm_dglu", "GroupedGemmDgluSm100"), + "grouped_gemm_dglu_wrapper_sm100": (".grouped_gemm_dglu", "grouped_gemm_dglu_wrapper_sm100"), + "GroupedGemmWgradSm100": (".grouped_gemm_wgrad", "GroupedGemmWgradSm100"), + "grouped_gemm_wgrad_wrapper_sm100": (".grouped_gemm_wgrad", "grouped_gemm_wgrad_wrapper_sm100"), +} + + +def __getattr__(name): + if name in _SYMBOLS: + module_name, symbol_name = _SYMBOLS[name] + value = getattr(import_module(module_name, __name__), symbol_name) + else: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + globals()[name] = value + return value + + +__all__ = list(_SYMBOLS) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py index 748246694..e6855c90c 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py @@ -1,12 +1,13 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmDgluSm100, - grouped_gemm_dglu_wrapper_sm100, -) +"""Grouped GEMM dGLU APIs for SM100.""" + +from ..._operation_api import make_operation_api -__all__ = [ +_API_EXPORTS = ( "GroupedGemmDgluSm100", "grouped_gemm_dglu_wrapper_sm100", -] +) + +__all__, __getattr__ = make_operation_api(globals(), exports=_API_EXPORTS) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dglu/api.py b/python/cudnn/grouped_gemm/grouped_gemm_dglu/api.py index 1d026812d..288ba1be5 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dglu/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dglu/api.py @@ -57,10 +57,10 @@ from cutlass.cute.runtime import from_dlpack, make_fake_stream from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.api_base import ApiBaseTorch, TupleDict, ceil_div, is_power_of_2 -class GroupedGemmDgluSm100(APIBase): +class GroupedGemmDgluSm100(ApiBaseTorch): """Unified API for grouped GEMM dGLU backward operation on SM100+ GPUs. This kernel performs block-scaled grouped GEMM with dGLU activation 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..310e16406 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_dglu/jax.py @@ -0,0 +1,655 @@ +# 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 + +import os +from typing import Any, Optional + +import jax.numpy as jnp + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + as_dtype, + call_cutedsl, + require_array, + require_dtype, +) +from ..._jax.gemm import ( + as_gemm_tensor_desc, + block_scale_tensor_spec, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + probability_tensor_spec, + require_16_byte_extent, + require_layout, +) +from ..._jax.grouped_gemm import ( + grouped_workspace_tensor_spec, + require_grouped_fp8_scales, + require_grouped_gemm_inputs, + require_grouped_probability, + require_grouped_vector, +) +from ...gemm_validation import ( + block_scale_shape, + resolve_max_active_clusters, +) + + +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, + use_dynamic_sched: bool, + act_func: str, + epilogue_op: str, + linear_offset: float, + geglu_alpha: float, + glu_clamp_max: float, + glu_clamp_min: float, + cluster_overlap_margin: 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() + + d_row = take() + d_col = take() + dprob = take() + dbias = take() if generate_dbias else None + sfd_row = take() + sfd_col = take() + 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] == BlockScaledMoEGroupedGemmDgluDbiasKernel.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, + ) + max_active_clusters = resolve_max_active_clusters( + cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]), + cluster_overlap_margin, + ) + 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, + None, + 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 an MXFP8 dense-weight grouped GEMM with fused dGLU. + + ``dprob_tensor`` and optional ``dbias_tensor`` are fresh, zero-initialized + JAX results. The FP8 output and its E8M0 scale factors are public results; + temporary scheduler storage is owned by XLA. + """ + + from .moe_blockscaled_grouped_gemm_dglu_dbias import ( + BlockScaledMoEGroupedGemmDgluDbiasKernel, + ) + + kernel = BlockScaledMoEGroupedGemmDgluDbiasKernel + 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=kernel.MAX_EXPERTS, + ) + if n % 32: + raise ValueError(f"b_tensor N must be divisible by 32 for dGLU, got {n}") + if sf_vec_size != kernel.FP8_SF_VEC_SIZE: + raise ValueError(f"FP8 grouped GEMM requires sf_vec_size={kernel.FP8_SF_VEC_SIZE}, got {sf_vec_size}") + if m_aligned != kernel.FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {kernel.FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_fp8_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ) + + output_n = 2 * n + require_array( + c_desc, + shape=(m, output_n, 1), + dtype=( + jnp.float32, + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + ), + ) + 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) + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required for an FP8 output") + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + + acc_dtype = require_dtype(acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32) + d_dtype = require_dtype( + d_dtype, + (jnp.float8_e4m3fn, jnp.float8_e5m2), + name="d_dtype", + default=ab_dtype, + ) + 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 = kernel.require_mma_tiler(mma_tiler_mn) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == kernel.TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = kernel.require_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 = [ + BufferSpec("d_row_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec), + BufferSpec("d_col_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec), + BufferSpec( + "dprob_tensor", + (m, 1, 1), + jnp.float32, + tensor_spec=probability_tensor_spec(), + fill_value=0.0, + ), + ] + if generate_dbias: + outputs.append( + BufferSpec( + "dbias_tensor", + (experts, output_n, 1), + jnp.bfloat16, + fill_value=0.0, + ) + ) + outputs.extend( + ( + BufferSpec( + "sfd_row_tensor", + block_scale_shape(m, output_n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + BufferSpec( + "sfd_col_tensor", + block_scale_shape(output_n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + ) + ) + + workspace_bytes = max(kernel.get_dense_workspace_bytes(bool(use_dynamic_sched)), 1) + results = call_cutedsl( + _launch, + ( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + norm_const_tensor, + ), + 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), + "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=( + BufferSpec( + "workspace", + (workspace_bytes,), + jnp.uint8, + tensor_spec=grouped_workspace_tensor_spec(), + ), + ), + input_specs=( + a_spec, + b_spec, + output_spec, + scale_spec, + scale_spec, + None, + None, + None, + probability_tensor_spec(), + None, + ), + ) + 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)) + sfd_row_tensor = results[result_idx] + sfd_col_tensor = results[result_idx + 1] + return TupleDict( + d_row_tensor=d_row_tensor, + d_col_tensor=d_col_tensor, + dprob_tensor=dprob_tensor, + dbias_tensor=dbias_tensor, + amax_tensor=None, + 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) + + +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_dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py b/python/cudnn/grouped_gemm/grouped_gemm_dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py index 27a42060f..9917aaf0c 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py @@ -60,6 +60,10 @@ from cutlass._mlir.dialects import math, llvm from cutlass._mlir.dialects import vector, arith +from ...gemm_validation import ( + require_cluster_shape as _require_cluster_shape, + require_mma_tiler as _require_mma_tiler, +) from ..moe_persistent_scheduler import ( MoEPersistentTileScheduler, MoESchedulerParams, @@ -232,9 +236,59 @@ class BlockScaledMoEGroupedGemmDgluDbiasKernel: """ - # Fixed pad size for user-side padding (decoupled from kernel tile size) + MMA_TILER_M = (128, 256) + MMA_TILER_N = (256,) + TWO_CTA_MMA_TILER_M = 256 + MAX_CLUSTER_CTAS = 16 + MAX_CLUSTER_DIMENSION = 4 + CLUSTER_TILER_M = (128, 256) + SF_VEC_SIZES = (16, 32) + FP8_SF_VEC_SIZE = 32 + MAX_EXPERTS = 1024 + DYNAMIC_SCHED_WORKSPACE_BYTES = 4 + FIX_PAD_SIZE = 256 + @classmethod + def require_mma_tiler(cls, mma_tiler_mn: Tuple[int, int]) -> Tuple[int, int]: + """Validate an FE-supported grouped GEMM MMA tile.""" + + return _require_mma_tiler( + mma_tiler_mn, + allowed_m=cls.MMA_TILER_M, + allowed_n=cls.MMA_TILER_N, + ) + + @classmethod + def require_cluster_shape( + cls, + cluster_shape_mn: Tuple[int, int], + *, + mma_tiler_mn: Tuple[int, int], + ) -> Tuple[int, int]: + """Validate the cluster and grouped scheduler tile shape.""" + + cluster_shape_mn = _require_cluster_shape( + cluster_shape_mn, + mma_m=mma_tiler_mn[0], + two_cta_mma_m=cls.TWO_CTA_MMA_TILER_M, + max_ctas=cls.MAX_CLUSTER_CTAS, + max_dimension=cls.MAX_CLUSTER_DIMENSION, + ) + cta_group_size = 2 if mma_tiler_mn[0] == cls.TWO_CTA_MMA_TILER_M else 1 + cluster_tiler_m = cluster_shape_mn[0] // cta_group_size * mma_tiler_mn[0] + if cluster_tiler_m not in cls.CLUSTER_TILER_M: + raise ValueError( + f"cluster M tile must be one of {cls.CLUSTER_TILER_M}, " f"got {cluster_tiler_m} from MMA tile {mma_tiler_mn} and cluster {cluster_shape_mn}" + ) + return cluster_shape_mn + + @classmethod + def get_dense_workspace_bytes(cls, use_dynamic_sched: bool) -> int: + """Return the temporary workspace size for dense-weight mode.""" + + return cls.DYNAMIC_SCHED_WORKSPACE_BYTES if use_dynamic_sched else 0 + @staticmethod def can_implement( ab_dtype: Type[cutlass.Numeric], diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py index e15137794..04fbfe719 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py @@ -1,12 +1,13 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmDsreluSm100, - grouped_gemm_dsrelu_wrapper_sm100, -) +"""Grouped GEMM dSReLU APIs for SM100.""" + +from ..._operation_api import make_operation_api -__all__ = [ +_API_EXPORTS = ( "GroupedGemmDsreluSm100", "grouped_gemm_dsrelu_wrapper_sm100", -] +) + +__all__, __getattr__ = make_operation_api(globals(), exports=_API_EXPORTS) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/api.py b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/api.py index 2d65a12ef..db7a5d402 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/api.py @@ -60,7 +60,7 @@ from cutlass.cute.runtime import from_dlpack, make_fake_stream from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.api_base import ApiBaseTorch, TupleDict, ceil_div, is_power_of_2 def _reinterpret_raw_grouped_fp4_tensor(tensor: torch.Tensor) -> torch.Tensor: @@ -71,7 +71,7 @@ def _reinterpret_raw_grouped_fp4_tensor(tensor: torch.Tensor) -> torch.Tensor: return tensor -class GroupedGemmDsreluSm100(APIBase): +class GroupedGemmDsreluSm100(ApiBaseTorch): """Unified API for grouped GEMM dSReLU backward operation on SM100+ GPUs. This kernel performs block-scaled grouped GEMM with dSReLU activation 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..0ce784249 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/jax.py @@ -0,0 +1,591 @@ +# 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 + +import os +from typing import Any, Optional + +import jax.numpy as jnp + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + as_dtype, + call_cutedsl, + require_array, + require_dtype, +) +from ..._jax.gemm import ( + as_gemm_tensor_desc, + block_scale_tensor_spec, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + probability_tensor_spec, + require_16_byte_extent, + require_layout, +) +from ..._jax.grouped_gemm import ( + grouped_workspace_tensor_spec, + require_grouped_fp8_scales, + require_grouped_gemm_inputs, + require_grouped_probability, + require_grouped_vector, +) +from ...gemm_validation import ( + block_scale_shape, + resolve_max_active_clusters, +) + + +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, + use_dynamic_sched: bool, + use_dsrelu_reuse: bool, + cluster_overlap_margin: 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() + + d_row = take() + d_col = take() + d_srelu = take() + dprob = take() + dbias = take() if generate_dbias else None + sfd_row = take() + sfd_col = take() + sfd_col_d_srelu = take() + 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] == BlockScaledMoEGroupedGemmQuantBwdKernel.TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + generate_sfd=True, + 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, + ) + max_active_clusters = resolve_max_active_clusters( + cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]), + cluster_overlap_margin, + ) + 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, + None, + 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 an MXFP8 dense-weight 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. Temporary scheduler storage is owned by XLA. + """ + + from .moe_blockscaled_grouped_gemm_dsrelu_quant import ( + BlockScaledMoEGroupedGemmQuantBwdKernel, + ) + + kernel = BlockScaledMoEGroupedGemmQuantBwdKernel + 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=kernel.MAX_EXPERTS, + ) + if sf_vec_size != kernel.FP8_SF_VEC_SIZE: + raise ValueError(f"FP8 grouped GEMM requires sf_vec_size={kernel.FP8_SF_VEC_SIZE}, got {sf_vec_size}") + if m_aligned != kernel.FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {kernel.FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_fp8_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ) + + require_array( + c_desc, + shape=(m, n, 1), + dtype=( + jnp.float32, + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + ), + ) + 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) + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required for an FP8 output") + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + + acc_dtype = require_dtype(acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32) + d_dtype = require_dtype( + d_dtype, + (jnp.float8_e4m3fn, jnp.float8_e5m2), + name="d_dtype", + default=ab_dtype, + ) + mma_tiler_mn = kernel.require_mma_tiler(mma_tiler_mn) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == kernel.TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = kernel.require_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 = [ + BufferSpec("d_row_tensor", (1, m, n), d_dtype, tensor_spec=output_spec), + BufferSpec("d_col_tensor", (1, m, n), d_dtype, tensor_spec=output_spec), + BufferSpec("d_srelu_tensor", (1, m, n), d_dtype, tensor_spec=output_spec), + BufferSpec( + "dprob_tensor", + (m, 1, 1), + jnp.float32, + tensor_spec=probability_tensor_spec(), + fill_value=0.0, + ), + ] + if generate_dbias: + outputs.append( + BufferSpec( + "dbias_tensor", + (experts, n, 1), + jnp.bfloat16, + fill_value=0.0, + ) + ) + outputs.extend( + ( + BufferSpec( + "sfd_row_tensor", + block_scale_shape(m, n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + BufferSpec( + "sfd_col_tensor", + block_scale_shape(n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + BufferSpec( + "sfd_col_d_srelu_tensor", + block_scale_shape(n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + ) + ) + + workspace_bytes = max(kernel.get_dense_workspace_bytes(bool(use_dynamic_sched)), 1) + results = call_cutedsl( + _launch, + ( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + norm_const_tensor, + ), + 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), + "use_dynamic_sched": bool(use_dynamic_sched), + "use_dsrelu_reuse": bool(use_dsrelu_reuse), + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + workspaces=( + BufferSpec( + "workspace", + (workspace_bytes,), + jnp.uint8, + tensor_spec=grouped_workspace_tensor_spec(), + ), + ), + input_specs=( + a_spec, + b_spec, + output_spec, + scale_spec, + scale_spec, + None, + None, + probability_tensor_spec(), + None, + ), + ) + 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)) + sfd_row_tensor = results[result_idx] + sfd_col_tensor = results[result_idx + 1] + sfd_col_d_srelu_tensor = results[result_idx + 2] + 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=None, + 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) + + +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_dsrelu/moe_blockscaled_grouped_gemm_dsrelu_quant.py b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/moe_blockscaled_grouped_gemm_dsrelu_quant.py index b6d5fe91f..1d0d2350b 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/moe_blockscaled_grouped_gemm_dsrelu_quant.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/moe_blockscaled_grouped_gemm_dsrelu_quant.py @@ -39,6 +39,10 @@ from cutlass._mlir.dialects.nvvm import ReduxKind from cutlass._mlir.dialects import llvm from cutlass.cute.typing import Float32, Int32, AddressSpace +from ...gemm_validation import ( + require_cluster_shape as _require_cluster_shape, + require_mma_tiler as _require_mma_tiler, +) def atomic_add_bf16x2(ptr, val_fp32_lo, val_fp32_hi, *, loc=None, ip=None): @@ -120,8 +124,58 @@ class BlockScaledMoEGroupedGemmQuantBwdKernel: :param use_dsrelu_reuse: Reuse relu(C)^2 between d_srelu and dprob. """ + MMA_TILER_M = (128, 256) + MMA_TILER_N = (256,) + TWO_CTA_MMA_TILER_M = 256 + MAX_CLUSTER_CTAS = 16 + MAX_CLUSTER_DIMENSION = 4 + CLUSTER_TILER_M = (128, 256) + SF_VEC_SIZES = (16, 32) + FP8_SF_VEC_SIZE = 32 + MAX_EXPERTS = 1024 + DYNAMIC_SCHED_WORKSPACE_BYTES = 4 FIX_PAD_SIZE = 256 + @classmethod + def require_mma_tiler(cls, mma_tiler_mn: Tuple[int, int]) -> Tuple[int, int]: + """Validate an FE-supported grouped GEMM MMA tile.""" + + return _require_mma_tiler( + mma_tiler_mn, + allowed_m=cls.MMA_TILER_M, + allowed_n=cls.MMA_TILER_N, + ) + + @classmethod + def require_cluster_shape( + cls, + cluster_shape_mn: Tuple[int, int], + *, + mma_tiler_mn: Tuple[int, int], + ) -> Tuple[int, int]: + """Validate the cluster and grouped scheduler tile shape.""" + + cluster_shape_mn = _require_cluster_shape( + cluster_shape_mn, + mma_m=mma_tiler_mn[0], + two_cta_mma_m=cls.TWO_CTA_MMA_TILER_M, + max_ctas=cls.MAX_CLUSTER_CTAS, + max_dimension=cls.MAX_CLUSTER_DIMENSION, + ) + cta_group_size = 2 if mma_tiler_mn[0] == cls.TWO_CTA_MMA_TILER_M else 1 + cluster_tiler_m = cluster_shape_mn[0] // cta_group_size * mma_tiler_mn[0] + if cluster_tiler_m not in cls.CLUSTER_TILER_M: + raise ValueError( + f"cluster M tile must be one of {cls.CLUSTER_TILER_M}, " f"got {cluster_tiler_m} from MMA tile {mma_tiler_mn} and cluster {cluster_shape_mn}" + ) + return cluster_shape_mn + + @classmethod + def get_dense_workspace_bytes(cls, use_dynamic_sched: bool) -> int: + """Return the temporary workspace size for dense-weight mode.""" + + return cls.DYNAMIC_SCHED_WORKSPACE_BYTES if use_dynamic_sched else 0 + @staticmethod def can_implement( ab_dtype: Type[cutlass.Numeric], diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py index f0261951a..3e28f274b 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py @@ -1,12 +1,16 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmDswigluSm100, - grouped_gemm_dswiglu_wrapper_sm100, -) +"""Grouped GEMM dSwiGLU APIs for SM100.""" + +from ..._operation_api import make_operation_api -__all__ = [ +_API_EXPORTS = ( "GroupedGemmDswigluSm100", "grouped_gemm_dswiglu_wrapper_sm100", -] +) + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/api.py b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/api.py index f6b518223..16188b158 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/api.py @@ -45,10 +45,10 @@ from cutlass.cute.runtime import make_fake_stream from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.api_base import ApiBaseTorch, TupleDict, ceil_div, is_power_of_2 -class GroupedGemmDswigluSm100(APIBase): +class GroupedGemmDswigluSm100(ApiBaseTorch): """API class for Grouped GEMM dSwiGLU backward operation on SM100+ GPUs. This kernel performs contiguous grouped block-scaled GEMM backward pass diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/grouped_gemm_dswiglu_quant.py b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/grouped_gemm_dswiglu_quant.py index c70db0023..c9a2c8fd6 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/grouped_gemm_dswiglu_quant.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/grouped_gemm_dswiglu_quant.py @@ -39,6 +39,11 @@ import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.utils.blockscaled_layout as blockscaled_utils +from ...gemm_validation import ( + require_cluster_shape as _require_cluster_shape, + require_mma_tiler as _require_mma_tiler, +) + from cutlass.cute.typing import Float32 from cutlass.cutlass_dsl import T from cutlass._mlir import ir @@ -190,9 +195,53 @@ class BlockScaledContiguousGroupedGemmKernel: """ + MMA_TILER_M = (64, 128, 256) + MMA_TILER_N = (128, 256) + TWO_CTA_MMA_TILER_M = 256 + MAX_CLUSTER_CTAS = 16 + MAX_CLUSTER_DIMENSION = 4 + CLUSTER_TILER_M = (128, 256) + SF_VEC_SIZES = (16, 32) + FP8_SF_VEC_SIZE = 32 + MAX_EXPERTS = 1024 + # Fixed pad size for user-side padding (decoupled from kernel tile size) FIX_PAD_SIZE = 256 + @classmethod + def require_mma_tiler(cls, mma_tiler_mn: Tuple[int, int]) -> Tuple[int, int]: + """Validate an FE-supported grouped GEMM MMA tile.""" + + return _require_mma_tiler( + mma_tiler_mn, + allowed_m=cls.MMA_TILER_M, + allowed_n=cls.MMA_TILER_N, + ) + + @classmethod + def require_cluster_shape( + cls, + cluster_shape_mn: Tuple[int, int], + *, + mma_tiler_mn: Tuple[int, int], + ) -> Tuple[int, int]: + """Validate the cluster and grouped scheduler tile shape.""" + + cluster_shape_mn = _require_cluster_shape( + cluster_shape_mn, + mma_m=mma_tiler_mn[0], + two_cta_mma_m=cls.TWO_CTA_MMA_TILER_M, + max_ctas=cls.MAX_CLUSTER_CTAS, + max_dimension=cls.MAX_CLUSTER_DIMENSION, + ) + cta_group_size = 2 if mma_tiler_mn[0] == cls.TWO_CTA_MMA_TILER_M else 1 + cluster_tiler_m = cluster_shape_mn[0] // cta_group_size * mma_tiler_mn[0] + if cluster_tiler_m not in cls.CLUSTER_TILER_M: + raise ValueError( + f"cluster M tile must be one of {cls.CLUSTER_TILER_M}, " f"got {cluster_tiler_m} from MMA tile {mma_tiler_mn} and cluster {cluster_shape_mn}" + ) + return cluster_shape_mn + def __init__( self, sf_vec_size: int, 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..78a786160 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/jax.py @@ -0,0 +1,523 @@ +# 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 + +import os +from typing import Any, Optional + +import jax.numpy as jnp + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + as_dtype, + call_cutedsl, + require_array, + require_dtype, +) +from ..._jax.gemm import ( + as_gemm_tensor_desc, + block_scale_tensor_spec, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + probability_tensor_spec, + require_16_byte_extent, + require_layout, +) +from ..._jax.grouped_gemm import ( + require_grouped_fp8_scales, + require_grouped_gemm_inputs, + require_grouped_probability, + require_grouped_vector, +) +from ...gemm_validation import ( + block_scale_shape, + resolve_max_active_clusters, +) + + +def _launch( + stream, + a, + b, + c, + sfa, + sfb, + padded_offsets, + alpha, + beta, + prob, + norm_const, + d_row, + d_col, + dprob, + sfd_row, + sfd_col, + *, + 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, + cluster_overlap_margin: int, +): + import cutlass + import cutlass.cute as cute + from cutlass.jax import jax_to_cutlass_dtype + + from .grouped_gemm_dswiglu_quant import BlockScaledContiguousGroupedGemmKernel + + 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] == BlockScaledContiguousGroupedGemmKernel.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, + ) + max_active_clusters = resolve_max_active_clusters( + cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]), + cluster_overlap_margin, + ) + kernel( + a, + b, + c, + d_row, + d_col, + sfa, + sfb, + sfd_row, + sfd_col, + None, + 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: Any, + 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, + cluster_overlap_margin: int = 0, + *, + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute the MXFP8 contiguous 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. + Configuration values are static when traced by :func:`jax.jit`. + """ + + from .grouped_gemm_dswiglu_quant import BlockScaledContiguousGroupedGemmKernel + + kernel = BlockScaledContiguousGroupedGemmKernel + 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) + 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=kernel.MAX_EXPERTS, + ) + if sf_vec_size != kernel.FP8_SF_VEC_SIZE: + raise ValueError(f"FP8 grouped GEMM requires sf_vec_size={kernel.FP8_SF_VEC_SIZE}, got {sf_vec_size}") + if m_aligned != kernel.FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {kernel.FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_fp8_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ) + + output_n = 2 * n + require_array( + c_desc, + shape=(m, output_n, 1), + dtype=( + jnp.float32, + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + ), + ) + 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) + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + 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) + d_dtype = require_dtype( + d_dtype, + (jnp.float8_e4m3fn, jnp.float8_e5m2), + name="d_dtype", + default=ab_dtype, + ) + 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 = kernel.require_mma_tiler(mma_tiler_mn) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == kernel.TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = kernel.require_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, + } + + d_row_tensor, d_col_tensor, dprob_tensor, sfd_row_tensor, sfd_col_tensor = call_cutedsl( + _launch, + ( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + norm_const_tensor, + ), + 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, + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=( + BufferSpec("d_row_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec), + BufferSpec("d_col_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec), + BufferSpec( + "dprob_tensor", + (m, 1, 1), + jnp.float32, + tensor_spec=probability_tensor_spec(), + fill_value=0.0, + ), + BufferSpec( + "sfd_row_tensor", + block_scale_shape(m, output_n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + BufferSpec( + "sfd_col_tensor", + block_scale_shape(output_n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + ), + input_specs=( + a_spec, + b_spec, + output_spec, + scale_spec, + scale_spec, + None, + None, + None, + probability_tensor_spec(), + None, + ), + ) + return TupleDict( + d_row_tensor=d_row_tensor, + d_col_tensor=d_col_tensor, + dprob_tensor=dprob_tensor, + amax_tensor=None, + 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: Any, + 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, + ) -> 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") + 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_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, + "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: Any, + ) -> 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: Any, + ) -> 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) + + +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: Any, + 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, +) -> TupleDict: + """Compute the MXFP8 contiguous 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, + ) + 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..c21c9738f 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu/__init__.py @@ -1,12 +1,13 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmGluSm100, - grouped_gemm_glu_wrapper_sm100, -) +"""Grouped GEMM GLU APIs for SM100.""" + +from ..._operation_api import make_operation_api -__all__ = [ +_API_EXPORTS = ( "GroupedGemmGluSm100", "grouped_gemm_glu_wrapper_sm100", -] +) + +__all__, __getattr__ = make_operation_api(globals(), exports=_API_EXPORTS) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu/api.py b/python/cudnn/grouped_gemm/grouped_gemm_glu/api.py index b814b046f..e92c5f69a 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu/api.py @@ -57,10 +57,10 @@ from cutlass.cute.runtime import from_dlpack, make_fake_stream from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.api_base import ApiBaseTorch, TupleDict, ceil_div, is_power_of_2 -class GroupedGemmGluSm100(APIBase): +class GroupedGemmGluSm100(ApiBaseTorch): """Unified API for grouped GEMM GLU forward operation on SM100+ GPUs. This kernel performs block-scaled grouped GEMM with GLU activation 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..49c554b9e --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu/jax.py @@ -0,0 +1,632 @@ +# 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 + +import os +from typing import Any, Optional + +import jax.numpy as jnp + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + call_cutedsl, + require_array, + require_dtype, +) +from ..._jax.gemm import ( + as_gemm_tensor_desc, + block_scale_tensor_spec, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + probability_tensor_spec, + require_16_byte_extent, + require_layout, +) +from ..._jax.grouped_gemm import ( + grouped_bias_tensor_spec, + grouped_workspace_tensor_spec, + require_grouped_fp8_scales, + require_grouped_gemm_inputs, + require_grouped_probability, + require_grouped_vector, +) +from ...gemm_validation import ( + block_scale_shape, + resolve_max_active_clusters, +) + + +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, + quantized_output: bool, + use_dynamic_sched: bool, + act_func: str, + linear_offset: float, + geglu_alpha: float, + glu_clamp_max: float, + glu_clamp_min: float, + cluster_overlap_margin: 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 quantized_output else None + + c = take() + d = take() + if quantized_output: + d_col = take() + amax = None + sfd_row = take() + sfd_col = take() + else: + d_col = None + amax = take() + sfd_row = None + sfd_col = None + + if not quantized_output: + d_col = take() + 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] == BlockScaledMoEGroupedGemmGluBiasKernel.TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + generate_sfd=quantized_output, + 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, + ) + max_active_clusters = resolve_max_active_clusters( + cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]), + cluster_overlap_margin, + ) + 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 an MXFP8 dense-weight grouped GEMM with fused GLU. + + Dense expert weights use public shape ``(L, N, K)``. ``swiglu`` and ``geglu`` + produce ``N / 2`` output columns. FP8 output returns row/column E8M0 + scale factors; FP16/BF16 output returns an initialized per-expert amax. + Temporary output and scheduler storage is owned by XLA. + """ + + from .moe_blockscaled_grouped_gemm_glu_bias import ( + BlockScaledMoEGroupedGemmGluBiasKernel, + ) + + kernel = BlockScaledMoEGroupedGemmGluBiasKernel + 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=kernel.MAX_EXPERTS, + ) + if n % 64: + raise ValueError(f"b_tensor N must be divisible by 64 for GLU, got {n}") + if sf_vec_size != kernel.FP8_SF_VEC_SIZE: + raise ValueError(f"FP8 grouped GEMM requires sf_vec_size={kernel.FP8_SF_VEC_SIZE}, got {sf_vec_size}") + if m_aligned != kernel.FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {kernel.FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_fp8_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ) + 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=(n, experts), + 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, + ) + d_dtype = require_dtype( + d_dtype, + (jnp.float16, jnp.bfloat16, jnp.float8_e4m3fn, jnp.float8_e5m2), + name="d_dtype", + default=jnp.bfloat16, + ) + fp8_dtypes = { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + } + quantized_output = d_dtype in fp8_dtypes + if quantized_output: + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required for an FP8 output") + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + 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 = kernel.require_mma_tiler(mma_tiler_mn) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == kernel.TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = kernel.require_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 quantized_output: + inputs.append(norm_const_tensor) + input_specs.append(None) + + outputs = [ + BufferSpec("c_tensor", (1, m, n), c_dtype, tensor_spec=output_spec), + BufferSpec("d_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec), + ] + workspaces = [] + if quantized_output: + outputs.extend( + ( + BufferSpec( + "d_col_tensor", + (1, m, output_n), + d_dtype, + tensor_spec=output_spec, + ), + BufferSpec( + "sfd_row_tensor", + block_scale_shape(m, output_n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + BufferSpec( + "sfd_col_tensor", + block_scale_shape(output_n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + ) + ) + else: + outputs.append(BufferSpec("amax_tensor", (experts, 1), jnp.float32, fill_value=-float("inf"))) + workspaces.append( + BufferSpec( + "d_col_scratch", + (1, m, output_n), + d_dtype, + tensor_spec=output_spec, + ) + ) + + workspace_bytes = max(kernel.get_dense_workspace_bytes(bool(use_dynamic_sched)), 1) + workspaces.append( + BufferSpec( + "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, + "quantized_output": quantized_output, + "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, + ) + if quantized_output: + c_tensor, d_tensor, d_col_tensor, sfd_row_tensor, sfd_col_tensor = results + amax_tensor = None + else: + c_tensor, d_tensor, amax_tensor = results + d_col_tensor = None + sfd_row_tensor = None + sfd_col_tensor = 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) + + +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/moe_blockscaled_grouped_gemm_glu_bias.py b/python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_bias.py index 26da93362..15b23d300 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_bias.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_bias.py @@ -55,6 +55,10 @@ import cutlass.utils.blockscaled_layout as blockscaled_utils from cutlass._mlir.dialects.nvvm import ReduxKind from cutlass.cute.typing import Float32, Int32, AddressSpace +from ...gemm_validation import ( + require_cluster_shape as _require_cluster_shape, + require_mma_tiler as _require_mma_tiler, +) from ..moe_persistent_scheduler import ( MoEPersistentTileScheduler, MoESchedulerParams, @@ -235,9 +239,59 @@ class BlockScaledMoEGroupedGemmGluBiasKernel: """ - # Fixed pad size for user-side padding (decoupled from kernel tile size) + MMA_TILER_M = (128, 256) + MMA_TILER_N = (256,) + TWO_CTA_MMA_TILER_M = 256 + MAX_CLUSTER_CTAS = 16 + MAX_CLUSTER_DIMENSION = 4 + CLUSTER_TILER_M = (128, 256) + SF_VEC_SIZES = (16, 32) + FP8_SF_VEC_SIZE = 32 + MAX_EXPERTS = 1024 + DYNAMIC_SCHED_WORKSPACE_BYTES = 4 + FIX_PAD_SIZE = 256 + @classmethod + def require_mma_tiler(cls, mma_tiler_mn: Tuple[int, int]) -> Tuple[int, int]: + """Validate an FE-supported grouped GEMM MMA tile.""" + + return _require_mma_tiler( + mma_tiler_mn, + allowed_m=cls.MMA_TILER_M, + allowed_n=cls.MMA_TILER_N, + ) + + @classmethod + def require_cluster_shape( + cls, + cluster_shape_mn: Tuple[int, int], + *, + mma_tiler_mn: Tuple[int, int], + ) -> Tuple[int, int]: + """Validate the cluster and grouped scheduler tile shape.""" + + cluster_shape_mn = _require_cluster_shape( + cluster_shape_mn, + mma_m=mma_tiler_mn[0], + two_cta_mma_m=cls.TWO_CTA_MMA_TILER_M, + max_ctas=cls.MAX_CLUSTER_CTAS, + max_dimension=cls.MAX_CLUSTER_DIMENSION, + ) + cta_group_size = 2 if mma_tiler_mn[0] == cls.TWO_CTA_MMA_TILER_M else 1 + cluster_tiler_m = cluster_shape_mn[0] // cta_group_size * mma_tiler_mn[0] + if cluster_tiler_m not in cls.CLUSTER_TILER_M: + raise ValueError( + f"cluster M tile must be one of {cls.CLUSTER_TILER_M}, " f"got {cluster_tiler_m} from MMA tile {mma_tiler_mn} and cluster {cluster_shape_mn}" + ) + return cluster_shape_mn + + @classmethod + def get_dense_workspace_bytes(cls, use_dynamic_sched: bool) -> int: + """Return the temporary workspace size for dense-weight mode.""" + + return cls.DYNAMIC_SCHED_WORKSPACE_BYTES if use_dynamic_sched else 0 + @staticmethod def can_implement( ab_dtype: Type[cutlass.Numeric], 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..4ac5fe371 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,16 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmGluHadamardSm100, - grouped_gemm_glu_hadamard_wrapper_sm100, -) +"""Lazy API exports for grouped GEMM + GLU + Hadamard.""" + +from ..._operation_api import make_operation_api -__all__ = [ +_API_EXPORTS = ( "GroupedGemmGluHadamardSm100", "grouped_gemm_glu_hadamard_wrapper_sm100", -] +) + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/api.py b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/api.py index a72b6fba5..1ed5de17b 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/api.py @@ -14,7 +14,7 @@ from cutlass.cute.nvgpu import OperandMajorMode from cutlass.cute.runtime import from_dlpack, make_fake_stream -from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.api_base import ApiBaseTorch, TupleDict, ceil_div, is_power_of_2 from cudnn.datatypes import _convert_to_cutlass_data_type from ..moe_utils import MoEWeightMode @@ -30,7 +30,7 @@ def _reinterpret_raw_grouped_fp4_tensor(tensor: torch.Tensor) -> torch.Tensor: return tensor -class GroupedGemmGluHadamardSm100(APIBase): +class GroupedGemmGluHadamardSm100(ApiBaseTorch): """Grouped GEMM GLU forward kernel with fused Hadamard transform.""" def __init__( 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..6fcd0c5ab --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/jax.py @@ -0,0 +1,586 @@ +# 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 + +import os +from typing import Any, Optional + +import jax.numpy as jnp + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + call_cutedsl, + require_array, + require_dtype, +) +from ..._jax.gemm import ( + as_gemm_tensor_desc, + block_scale_tensor_spec, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + probability_tensor_spec, + require_layout, +) +from ..._jax.grouped_gemm import ( + grouped_bias_tensor_spec, + grouped_workspace_tensor_spec, + require_grouped_block_scales, + require_grouped_gemm_inputs, + require_grouped_probability, +) +from ...gemm_validation import ( + require_contiguous_alignment, + resolve_max_active_clusters, +) + + +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, + cluster_overlap_margin: 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, + ) + max_active_clusters = resolve_max_active_clusters( + cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]), + cluster_overlap_margin, + ) + 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=kernel.MAX_EXPERTS, + valid_ab_dtypes=(jnp.float4_e2m1fn,), + ) + if m % kernel.FIX_PAD_SIZE: + raise ValueError(f"M must be divisible by {kernel.FIX_PAD_SIZE}, got {m}") + if n % 64: + raise ValueError(f"N must be divisible by 64, got {n}") + if m_aligned != kernel.FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {kernel.FIX_PAD_SIZE}, got {m_aligned}") + if sf_vec_size not in kernel.SF_VEC_SIZES: + raise ValueError(f"sf_vec_size must be one of {kernel.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=(n, experts), + 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 = kernel.require_mma_tiler(mma_tiler_mn) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) + cluster_shape_mn = kernel.require_cluster_shape(cluster_shape_mn, mma_tiler_mn=mma_tiler_mn) + + output_n = n if act_func == "srelu" else n // 2 + if output_n % kernel.HADAMARD_SIZE: + raise ValueError(f"D's N dimension must be divisible by {kernel.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=(kernel.HADAMARD_SIZE, kernel.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(kernel.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(kernel.get_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=( + BufferSpec("c_tensor", (1, m, n), c_dtype, tensor_spec=output_spec, fill_value=0), + BufferSpec( + "d_tensor", + (1, m, output_n), + d_dtype, + tensor_spec=output_spec, + fill_value=0, + ), + BufferSpec("amax_tensor", (experts, 1), jnp.float32, fill_value=0.0), + BufferSpec("post_rht_amax_tensor", (experts, 1), jnp.float32, fill_value=0.0), + ), + workspaces=( + BufferSpec( + "workspace", + (workspace_bytes,), + jnp.uint8, + tensor_spec=grouped_workspace_tensor_spec(), + fill_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) + + +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_glu_hadamard/moe_blockscaled_grouped_gemm_glu_hadamard.py b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/moe_blockscaled_grouped_gemm_glu_hadamard.py index 3e1b345b8..3b44c7f95 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/moe_blockscaled_grouped_gemm_glu_hadamard.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/moe_blockscaled_grouped_gemm_glu_hadamard.py @@ -63,6 +63,10 @@ import cutlass.utils.blockscaled_layout as blockscaled_utils from cutlass._mlir.dialects.nvvm import ReduxKind from cutlass.cute.typing import Float32, Int32, AddressSpace +from ...gemm_validation import ( + require_cluster_shape as _require_cluster_shape, + require_mma_tiler as _require_mma_tiler, +) from ..moe_persistent_scheduler import ( MoEPersistentTileScheduler, MoESchedulerParams, @@ -117,8 +121,58 @@ class BlockScaledMoEGroupedGemmGluHadamardKernel: :param enable_bias: Enable bias addition. """ + MMA_TILER_M = (256,) + MMA_TILER_N = (256,) + TWO_CTA_MMA_TILER_M = 256 + MAX_CLUSTER_CTAS = 16 + MAX_CLUSTER_DIMENSION = 4 + CLUSTER_TILER_M = (256,) + SF_VEC_SIZES = (16, 32) + MAX_EXPERTS = 1024 + DYNAMIC_SCHED_WORKSPACE_BYTES = 4 + HADAMARD_SIZE = HADAMARD_SIZE FIX_PAD_SIZE = 256 + @classmethod + def require_mma_tiler(cls, mma_tiler_mn: Tuple[int, int]) -> Tuple[int, int]: + """Validate the MMA tile supported by the Hadamard fusion.""" + + return _require_mma_tiler( + mma_tiler_mn, + allowed_m=cls.MMA_TILER_M, + allowed_n=cls.MMA_TILER_N, + ) + + @classmethod + def require_cluster_shape( + cls, + cluster_shape_mn: Tuple[int, int], + *, + mma_tiler_mn: Tuple[int, int], + ) -> Tuple[int, int]: + """Validate the cluster shape supported by the Hadamard scheduler.""" + + cluster_shape_mn = _require_cluster_shape( + cluster_shape_mn, + mma_m=mma_tiler_mn[0], + two_cta_mma_m=cls.TWO_CTA_MMA_TILER_M, + max_ctas=cls.MAX_CLUSTER_CTAS, + max_dimension=cls.MAX_CLUSTER_DIMENSION, + ) + cta_group_size = 2 if mma_tiler_mn[0] == cls.TWO_CTA_MMA_TILER_M else 1 + cluster_tiler_m = cluster_shape_mn[0] // cta_group_size * mma_tiler_mn[0] + if cluster_tiler_m not in cls.CLUSTER_TILER_M: + raise ValueError( + f"cluster M tile must be one of {cls.CLUSTER_TILER_M}, " f"got {cluster_tiler_m} from MMA tile {mma_tiler_mn} and cluster {cluster_shape_mn}" + ) + return cluster_shape_mn + + @classmethod + def get_dense_workspace_bytes(cls, use_dynamic_sched: bool) -> int: + """Return the temporary workspace size for dense-weight mode.""" + + return cls.DYNAMIC_SCHED_WORKSPACE_BYTES if use_dynamic_sched else 0 + @staticmethod def can_implement( ab_dtype: Type[cutlass.Numeric], @@ -457,7 +511,7 @@ def get_desc_workspace_bytes(self) -> int: def get_workspace_bytes(self) -> int: """Return total workspace size in bytes.""" desc_workspace_bytes = self.get_desc_workspace_bytes() - dynamic_sched_bytes = 4 if self.use_dynamic_sched else 0 + dynamic_sched_bytes = self.get_dense_workspace_bytes(self.use_dynamic_sched) return desc_workspace_bytes + dynamic_sched_bytes @cute.jit diff --git a/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py index 9f88a8d00..9a027f02f 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py @@ -1,20 +1,13 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -""" -Grouped GEMM Quant Kernel Module +"""Grouped GEMM quantization APIs for SM100.""" -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 ..._operation_api import make_operation_api -from .api import ( - GroupedGemmQuantSm100, - grouped_gemm_quant_wrapper_sm100, -) - -__all__ = [ +_API_EXPORTS = ( "GroupedGemmQuantSm100", "grouped_gemm_quant_wrapper_sm100", -] +) + +__all__, __getattr__ = make_operation_api(globals(), exports=_API_EXPORTS) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py b/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py index eaf7d7d98..b46dda4f3 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py @@ -43,7 +43,7 @@ from cuda.bindings import driver as cuda from cutlass.cute.runtime import make_fake_stream -from cudnn.api_base import APIBase, TensorDesc, TupleDict, ceil_div, is_power_of_2 +from cudnn.api_base import ApiBaseTorch, TorchTensorDesc, TupleDict, ceil_div, is_power_of_2 from cudnn.datatypes import _convert_to_cutlass_data_type from .grouped_gemm_quant import ( @@ -54,7 +54,7 @@ from cutlass.cute.runtime import from_dlpack -class GroupedGemmQuantSm100(APIBase): +class GroupedGemmQuantSm100(ApiBaseTorch): """Unified API for grouped GEMM quant operation on SM100+ GPUs. This kernel performs block-scaled grouped GEMM with output quantization @@ -161,12 +161,13 @@ def __init__( self._has_d_col = sample_d_col is not None self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col") if self.d_col_desc is None: - self.d_col_desc = TensorDesc( + self.d_col_desc = TorchTensorDesc( dtype=self.d_desc.dtype, shape=self.d_desc.shape, stride=self.d_desc.stride, stride_order=self.d_desc.stride_order, device=self.d_desc.device, + packing=self.d_desc.packing, name="sample_d_col", ) self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row") diff --git a/python/cudnn/grouped_gemm/grouped_gemm_quant/grouped_gemm_quant.py b/python/cudnn/grouped_gemm/grouped_gemm_quant/grouped_gemm_quant.py index cedc8ffc9..795096e17 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_quant/grouped_gemm_quant.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_quant/grouped_gemm_quant.py @@ -29,6 +29,10 @@ import cutlass.utils.blockscaled_layout as blockscaled_utils from cutlass._mlir.dialects.nvvm import ReduxKind from cutlass.cute.typing import Float32, Int32, AddressSpace +from ...gemm_validation import ( + require_cluster_shape as _require_cluster_shape, + require_mma_tiler as _require_mma_tiler, +) from ..moe_persistent_scheduler import ( MoEPersistentTileScheduler, MoESchedulerParams, @@ -78,8 +82,58 @@ class BlockScaledMoEGroupedGemmQuantKernel: :param use_dynamic_sched: Enable dynamic tile scheduling. """ + MMA_TILER_M = (128, 256) + MMA_TILER_N = (256,) + TWO_CTA_MMA_TILER_M = 256 + MAX_CLUSTER_CTAS = 16 + MAX_CLUSTER_DIMENSION = 4 + CLUSTER_TILER_M = (128, 256) + SF_VEC_SIZES = (16, 32) + FP8_SF_VEC_SIZE = 32 + MAX_EXPERTS = 1024 + DYNAMIC_SCHED_WORKSPACE_BYTES = 4 FIX_PAD_SIZE = 256 + @classmethod + def require_mma_tiler(cls, mma_tiler_mn: Tuple[int, int]) -> Tuple[int, int]: + """Validate an FE-supported grouped GEMM MMA tile.""" + + return _require_mma_tiler( + mma_tiler_mn, + allowed_m=cls.MMA_TILER_M, + allowed_n=cls.MMA_TILER_N, + ) + + @classmethod + def require_cluster_shape( + cls, + cluster_shape_mn: Tuple[int, int], + *, + mma_tiler_mn: Tuple[int, int], + ) -> Tuple[int, int]: + """Validate the cluster and grouped scheduler tile shape.""" + + cluster_shape_mn = _require_cluster_shape( + cluster_shape_mn, + mma_m=mma_tiler_mn[0], + two_cta_mma_m=cls.TWO_CTA_MMA_TILER_M, + max_ctas=cls.MAX_CLUSTER_CTAS, + max_dimension=cls.MAX_CLUSTER_DIMENSION, + ) + cta_group_size = 2 if mma_tiler_mn[0] == cls.TWO_CTA_MMA_TILER_M else 1 + cluster_tiler_m = cluster_shape_mn[0] // cta_group_size * mma_tiler_mn[0] + if cluster_tiler_m not in cls.CLUSTER_TILER_M: + raise ValueError( + f"cluster M tile must be one of {cls.CLUSTER_TILER_M}, " f"got {cluster_tiler_m} from MMA tile {mma_tiler_mn} and cluster {cluster_shape_mn}" + ) + return cluster_shape_mn + + @classmethod + def get_dense_workspace_bytes(cls, use_dynamic_sched: bool) -> int: + """Return the temporary workspace size for dense-weight mode.""" + + return cls.DYNAMIC_SCHED_WORKSPACE_BYTES if use_dynamic_sched else 0 + @staticmethod def can_implement( ab_dtype: Type[cutlass.Numeric], 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..f4871a705 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_quant/jax.py @@ -0,0 +1,555 @@ +# 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 + +import os +from typing import Any, Optional + +import jax.numpy as jnp + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + call_cutedsl, + require_array, + require_dtype, +) +from ..._jax.gemm import ( + as_gemm_tensor_desc, + block_scale_tensor_spec, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + probability_tensor_spec, + require_16_byte_extent, + require_layout, +) +from ..._jax.grouped_gemm import ( + grouped_bias_tensor_spec, + grouped_workspace_tensor_spec, + require_grouped_fp8_scales, + require_grouped_gemm_inputs, + require_grouped_probability, + require_grouped_vector, +) +from ...gemm_validation import ( + block_scale_shape, + resolve_max_active_clusters, +) + + +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, + quantized_output: bool, + use_dynamic_sched: bool, + cluster_overlap_margin: 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 quantized_output else None + d = take() + d_col = take() if quantized_output else None + amax = None if quantized_output else take() + sfd_row = take() if quantized_output else None + sfd_col = take() if quantized_output 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] == BlockScaledMoEGroupedGemmQuantKernel.TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + generate_sfd=quantized_output, + discrete_col_sfd=discrete_col_sfd, + enable_bias=has_bias, + expert_cnt=expert_cnt, + use_dynamic_sched=use_dynamic_sched, + ) + max_active_clusters = resolve_max_active_clusters( + cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]), + cluster_overlap_margin, + ) + 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 an MXFP8 dense-weight grouped GEMM with optional 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; FP16/BF16 outputs return an initialized per-expert amax reduction. + Temporary scheduler storage is owned by XLA and is not returned. + """ + + from .grouped_gemm_quant import BlockScaledMoEGroupedGemmQuantKernel + + kernel = BlockScaledMoEGroupedGemmQuantKernel + 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=kernel.MAX_EXPERTS, + ) + if sf_vec_size != kernel.FP8_SF_VEC_SIZE: + raise ValueError(f"FP8 grouped GEMM requires sf_vec_size={kernel.FP8_SF_VEC_SIZE}, got {sf_vec_size}") + if m_aligned != kernel.FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {kernel.FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_fp8_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ) + 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=(n, experts), + dtype=(jnp.float16, jnp.bfloat16, jnp.float32), + ) + + acc_dtype = require_dtype(acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32) + d_dtype = require_dtype( + d_dtype, + (jnp.float16, jnp.bfloat16, jnp.float8_e4m3fn, jnp.float8_e5m2), + name="d_dtype", + default=jnp.bfloat16, + ) + quantized_output = d_dtype in { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + } + if quantized_output: + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required for an FP8 output") + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + mma_tiler_mn = kernel.require_mma_tiler(mma_tiler_mn) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == kernel.TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = kernel.require_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 quantized_output: + inputs.append(norm_const_tensor) + input_specs.append(None) + + outputs = [BufferSpec("d_tensor", (1, m, n), d_dtype, tensor_spec=output_spec)] + if quantized_output: + outputs.extend( + ( + BufferSpec("d_col_tensor", (1, m, n), d_dtype, tensor_spec=output_spec), + BufferSpec( + "sfd_row_tensor", + block_scale_shape(m, n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + BufferSpec( + "sfd_col_tensor", + block_scale_shape(n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + ) + ) + else: + outputs.append(BufferSpec("amax_tensor", (experts, 1), jnp.float32, fill_value=-float("inf"))) + + workspace_bytes = max(kernel.get_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, + "quantized_output": quantized_output, + "use_dynamic_sched": bool(use_dynamic_sched), + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + workspaces=( + BufferSpec( + "workspace", + (workspace_bytes,), + jnp.uint8, + tensor_spec=grouped_workspace_tensor_spec(), + ), + ), + input_specs=input_specs, + ) + if quantized_output: + d_tensor, d_col_tensor, sfd_row_tensor, sfd_col_tensor = results + amax_tensor = None + else: + d_tensor, amax_tensor = results + d_col_tensor = None + sfd_row_tensor = None + sfd_col_tensor = 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) + + +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..bf63c16f0 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_srelu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_srelu/__init__.py @@ -1,19 +1,13 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -""" -Grouped GEMM SReLU Kernel Module +"""Grouped GEMM squared-ReLU APIs for SM100.""" -This module provides the forward grouped GEMM with SReLU activation -for MoE (Mixture of Experts) workloads on SM100+ GPUs. -""" +from ..._operation_api import make_operation_api -from .api import ( - GroupedGemmSreluSm100, - grouped_gemm_srelu_wrapper_sm100, -) - -__all__ = [ +_API_EXPORTS = ( "GroupedGemmSreluSm100", "grouped_gemm_srelu_wrapper_sm100", -] +) + +__all__, __getattr__ = make_operation_api(globals(), exports=_API_EXPORTS) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_srelu/api.py b/python/cudnn/grouped_gemm/grouped_gemm_srelu/api.py index 2ff670e08..c95d92199 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_srelu/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_srelu/api.py @@ -43,7 +43,7 @@ from cuda.bindings import driver as cuda from cutlass.cute.runtime import make_fake_stream -from cudnn.api_base import APIBase, TensorDesc, TupleDict, ceil_div, is_power_of_2 +from cudnn.api_base import ApiBaseTorch, TorchTensorDesc, TupleDict, ceil_div, is_power_of_2 from cudnn.datatypes import _convert_to_cutlass_data_type from .moe_blockscaled_grouped_gemm_srelu_quant import ( @@ -63,7 +63,7 @@ def _reinterpret_raw_grouped_fp4_tensor(tensor: torch.Tensor) -> torch.Tensor: return tensor -class GroupedGemmSreluSm100(APIBase): +class GroupedGemmSreluSm100(ApiBaseTorch): """Unified API for grouped GEMM SReLU operation on SM100+ GPUs. This kernel performs block-scaled grouped GEMM with output SReLU output quantization @@ -172,12 +172,13 @@ def __init__( self._has_d_col = sample_d_col is not None self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col") if self.d_col_desc is None: - self.d_col_desc = TensorDesc( + self.d_col_desc = TorchTensorDesc( dtype=self.d_desc.dtype, shape=self.d_desc.shape, stride=self.d_desc.stride, stride_order=self.d_desc.stride_order, device=self.d_desc.device, + packing=self.d_desc.packing, name="sample_d_col", ) self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row") 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..7cfcc81fc --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_srelu/jax.py @@ -0,0 +1,560 @@ +# 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 + +import os +from typing import Any, Optional + +import jax.numpy as jnp + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + call_cutedsl, + require_array, + require_dtype, +) +from ..._jax.gemm import ( + as_gemm_tensor_desc, + block_scale_tensor_spec, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + probability_tensor_spec, + require_16_byte_extent, + require_layout, +) +from ..._jax.grouped_gemm import ( + grouped_bias_tensor_spec, + grouped_workspace_tensor_spec, + require_grouped_fp8_scales, + require_grouped_gemm_inputs, + require_grouped_probability, + require_grouped_vector, +) +from ...gemm_validation import ( + block_scale_shape, + resolve_max_active_clusters, +) + + +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, + quantized_output: bool, + use_dynamic_sched: bool, + cluster_overlap_margin: 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 quantized_output else None + c = take() + d = take() + d_col = take() if quantized_output else None + amax = None if quantized_output else take() + sfd_row = take() if quantized_output else None + sfd_col = take() if quantized_output 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] == BlockScaledMoEGroupedGemmQuantKernel.TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + generate_sfd=quantized_output, + 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, + ) + max_active_clusters = resolve_max_active_clusters( + cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]), + cluster_overlap_margin, + ) + 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 an MXFP8 dense-weight grouped GEMM and squared-ReLU fusion. + + FP8 outputs return row/column E8M0 scale factors; FP16/BF16 outputs + return an initialized per-expert amax reduction. The scheduler workspace + is an internal XLA-owned result and is hidden from the public return value. + """ + + from .moe_blockscaled_grouped_gemm_srelu_quant import ( + BlockScaledMoEGroupedGemmQuantKernel, + ) + + kernel = BlockScaledMoEGroupedGemmQuantKernel + 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=kernel.MAX_EXPERTS, + ) + if sf_vec_size != kernel.FP8_SF_VEC_SIZE: + raise ValueError(f"FP8 grouped GEMM requires sf_vec_size={kernel.FP8_SF_VEC_SIZE}, got {sf_vec_size}") + if m_aligned != kernel.FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {kernel.FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_fp8_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ) + 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=(n, experts), + 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, + ) + d_dtype = require_dtype( + d_dtype, + (jnp.float16, jnp.bfloat16, jnp.float8_e4m3fn, jnp.float8_e5m2), + name="d_dtype", + default=jnp.bfloat16, + ) + quantized_output = d_dtype in { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + } + if quantized_output: + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required for an FP8 output") + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + mma_tiler_mn = kernel.require_mma_tiler(mma_tiler_mn) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == kernel.TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = kernel.require_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 quantized_output: + inputs.append(norm_const_tensor) + input_specs.append(None) + + outputs = [ + BufferSpec("c_tensor", (1, m, n), c_dtype, tensor_spec=output_spec), + BufferSpec("d_tensor", (1, m, n), d_dtype, tensor_spec=output_spec), + ] + if quantized_output: + outputs.extend( + ( + BufferSpec("d_col_tensor", (1, m, n), d_dtype, tensor_spec=output_spec), + BufferSpec( + "sfd_row_tensor", + block_scale_shape(m, n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + BufferSpec( + "sfd_col_tensor", + block_scale_shape(n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + ) + ) + else: + outputs.append(BufferSpec("amax_tensor", (experts, 1), jnp.float32, fill_value=-float("inf"))) + + workspace_bytes = max(kernel.get_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, + "quantized_output": quantized_output, + "use_dynamic_sched": bool(use_dynamic_sched), + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + workspaces=( + BufferSpec( + "workspace", + (workspace_bytes,), + jnp.uint8, + tensor_spec=grouped_workspace_tensor_spec(), + ), + ), + input_specs=input_specs, + ) + if quantized_output: + c_tensor, d_tensor, d_col_tensor, sfd_row_tensor, sfd_col_tensor = results + amax_tensor = None + else: + c_tensor, d_tensor, amax_tensor = results + d_col_tensor = None + sfd_row_tensor = None + sfd_col_tensor = 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) + + +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_srelu/moe_blockscaled_grouped_gemm_srelu_quant.py b/python/cudnn/grouped_gemm/grouped_gemm_srelu/moe_blockscaled_grouped_gemm_srelu_quant.py index 99840e976..0fed4aa50 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_srelu/moe_blockscaled_grouped_gemm_srelu_quant.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_srelu/moe_blockscaled_grouped_gemm_srelu_quant.py @@ -32,6 +32,10 @@ import cutlass.utils.blockscaled_layout as blockscaled_utils from cutlass._mlir.dialects.nvvm import ReduxKind from cutlass.cute.typing import Float32, Int32, AddressSpace +from ...gemm_validation import ( + require_cluster_shape as _require_cluster_shape, + require_mma_tiler as _require_mma_tiler, +) from ..moe_persistent_scheduler import ( MoEPersistentTileScheduler, MoESchedulerParams, @@ -90,8 +94,58 @@ class BlockScaledMoEGroupedGemmQuantKernel: :param epilogue_type: Epilogue activation type (``EpilogueType.NONE`` or ``EpilogueType.SRELU``). """ + MMA_TILER_M = (128, 256) + MMA_TILER_N = (256,) + TWO_CTA_MMA_TILER_M = 256 + MAX_CLUSTER_CTAS = 16 + MAX_CLUSTER_DIMENSION = 4 + CLUSTER_TILER_M = (128, 256) + SF_VEC_SIZES = (16, 32) + FP8_SF_VEC_SIZE = 32 + MAX_EXPERTS = 1024 + DYNAMIC_SCHED_WORKSPACE_BYTES = 4 FIX_PAD_SIZE = 256 + @classmethod + def require_mma_tiler(cls, mma_tiler_mn: Tuple[int, int]) -> Tuple[int, int]: + """Validate an FE-supported grouped GEMM MMA tile.""" + + return _require_mma_tiler( + mma_tiler_mn, + allowed_m=cls.MMA_TILER_M, + allowed_n=cls.MMA_TILER_N, + ) + + @classmethod + def require_cluster_shape( + cls, + cluster_shape_mn: Tuple[int, int], + *, + mma_tiler_mn: Tuple[int, int], + ) -> Tuple[int, int]: + """Validate the cluster and grouped scheduler tile shape.""" + + cluster_shape_mn = _require_cluster_shape( + cluster_shape_mn, + mma_m=mma_tiler_mn[0], + two_cta_mma_m=cls.TWO_CTA_MMA_TILER_M, + max_ctas=cls.MAX_CLUSTER_CTAS, + max_dimension=cls.MAX_CLUSTER_DIMENSION, + ) + cta_group_size = 2 if mma_tiler_mn[0] == cls.TWO_CTA_MMA_TILER_M else 1 + cluster_tiler_m = cluster_shape_mn[0] // cta_group_size * mma_tiler_mn[0] + if cluster_tiler_m not in cls.CLUSTER_TILER_M: + raise ValueError( + f"cluster M tile must be one of {cls.CLUSTER_TILER_M}, " f"got {cluster_tiler_m} from MMA tile {mma_tiler_mn} and cluster {cluster_shape_mn}" + ) + return cluster_shape_mn + + @classmethod + def get_dense_workspace_bytes(cls, use_dynamic_sched: bool) -> int: + """Return the temporary workspace size for dense-weight mode.""" + + return cls.DYNAMIC_SCHED_WORKSPACE_BYTES if use_dynamic_sched else 0 + @staticmethod def can_implement( ab_dtype: Type[cutlass.Numeric], diff --git a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py index 6bcd935f2..047c51cb8 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py @@ -1,19 +1,16 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -""" -Grouped GEMM SwiGLU Kernel Module +"""Grouped GEMM SwiGLU APIs for SM100.""" -This module provides the forward grouped GEMM with SwiGLU activation -for MoE (Mixture of Experts) workloads on SM100+ GPUs. -""" +from ..._operation_api import make_operation_api -from .api import ( - GroupedGemmSwigluSm100, - grouped_gemm_swiglu_wrapper_sm100, -) - -__all__ = [ +_API_EXPORTS = ( "GroupedGemmSwigluSm100", "grouped_gemm_swiglu_wrapper_sm100", -] +) + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/api.py b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/api.py index 7f876e1e2..9223e4c54 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/api.py @@ -46,10 +46,10 @@ from cutlass.cute.runtime import make_fake_stream from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from cudnn.api_base import ApiBaseTorch, TupleDict, ceil_div, is_power_of_2 -class GroupedGemmSwigluSm100(APIBase): +class GroupedGemmSwigluSm100(ApiBaseTorch): """API class for Grouped GEMM SwiGLU forward operation on SM100+ GPUs. This kernel performs contiguous grouped block-scaled GEMM with SwiGLU activation, diff --git a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/grouped_gemm_swiglu_quant.py b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/grouped_gemm_swiglu_quant.py index 3ed2a2687..44f922e70 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/grouped_gemm_swiglu_quant.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/grouped_gemm_swiglu_quant.py @@ -46,6 +46,11 @@ import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.utils.blockscaled_layout as blockscaled_utils +from ...gemm_validation import ( + require_cluster_shape as _require_cluster_shape, + require_mma_tiler as _require_mma_tiler, +) + from ..utils import ( PersistentTileSchedulerParams, StaticPersistentTileScheduler, @@ -205,9 +210,53 @@ class BlockScaledContiguousGroupedGemmKernel: """ + MMA_TILER_M = (64, 128, 256) + MMA_TILER_N = (128, 256) + TWO_CTA_MMA_TILER_M = 256 + MAX_CLUSTER_CTAS = 16 + MAX_CLUSTER_DIMENSION = 4 + CLUSTER_TILER_M = (128, 256) + SF_VEC_SIZES = (16, 32) + FP8_SF_VEC_SIZE = 32 + MAX_EXPERTS = 1024 + # Fixed pad size for user-side padding (decoupled from kernel tile size) FIX_PAD_SIZE = 256 + @classmethod + def require_mma_tiler(cls, mma_tiler_mn: Tuple[int, int]) -> Tuple[int, int]: + """Validate an FE-supported grouped GEMM MMA tile.""" + + return _require_mma_tiler( + mma_tiler_mn, + allowed_m=cls.MMA_TILER_M, + allowed_n=cls.MMA_TILER_N, + ) + + @classmethod + def require_cluster_shape( + cls, + cluster_shape_mn: Tuple[int, int], + *, + mma_tiler_mn: Tuple[int, int], + ) -> Tuple[int, int]: + """Validate the cluster and grouped scheduler tile shape.""" + + cluster_shape_mn = _require_cluster_shape( + cluster_shape_mn, + mma_m=mma_tiler_mn[0], + two_cta_mma_m=cls.TWO_CTA_MMA_TILER_M, + max_ctas=cls.MAX_CLUSTER_CTAS, + max_dimension=cls.MAX_CLUSTER_DIMENSION, + ) + cta_group_size = 2 if mma_tiler_mn[0] == cls.TWO_CTA_MMA_TILER_M else 1 + cluster_tiler_m = cluster_shape_mn[0] // cta_group_size * mma_tiler_mn[0] + if cluster_tiler_m not in cls.CLUSTER_TILER_M: + raise ValueError( + f"cluster M tile must be one of {cls.CLUSTER_TILER_M}, " f"got {cluster_tiler_m} from MMA tile {mma_tiler_mn} and cluster {cluster_shape_mn}" + ) + return cluster_shape_mn + def __init__( self, sf_vec_size: int, 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..c6dd25d39 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/jax.py @@ -0,0 +1,490 @@ +# 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 + +import os +from typing import Any, Optional + +import jax.numpy as jnp + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + call_cutedsl, + require_dtype, +) +from ..._jax.gemm import ( + as_gemm_tensor_desc, + block_scale_tensor_spec, + gemm_a_tensor_spec, + gemm_b_tensor_spec, + gemm_c_tensor_spec, + probability_tensor_spec, + require_16_byte_extent, + require_layout, +) +from ..._jax.grouped_gemm import ( + require_grouped_fp8_scales, + require_grouped_gemm_inputs, + require_grouped_probability, + require_grouped_vector, +) +from ...gemm_validation import ( + block_scale_shape, + resolve_max_active_clusters, +) + + +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_amax: bool, + cluster_overlap_margin: int, +): + import cutlass + 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() + prob = take() if has_prob else None + c = take() + d = take() + d_col = take() + amax = take() if has_amax else None + sfd_row = take() + sfd_col = take() + 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] == BlockScaledContiguousGroupedGemmKernel.TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vector_f32=vector_f32, + generate_sfd=True, + discrete_col_sfd=discrete_col_sfd, + expert_cnt=expert_cnt, + use_mono_increase_expert_idx=True, + ) + max_active_clusters = resolve_max_active_clusters( + cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]), + cluster_overlap_margin, + ) + 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: Any, + 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 MXFP8 contiguous 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. + + This binding supports FP8 inputs with E8M0 scales and ``sf_vec_size=32``. + Configuration values are static when traced by :func:`jax.jit`. + """ + + from .grouped_gemm_swiglu_quant import BlockScaledContiguousGroupedGemmKernel + + kernel = BlockScaledContiguousGroupedGemmKernel + 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=kernel.MAX_EXPERTS, + ) + if n % 64: + raise ValueError(f"b_tensor N must be divisible by 64 for SwiGLU, got {n}") + if sf_vec_size != kernel.FP8_SF_VEC_SIZE: + raise ValueError(f"FP8 grouped GEMM requires sf_vec_size={kernel.FP8_SF_VEC_SIZE}, got {sf_vec_size}") + if m_aligned != kernel.FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {kernel.FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_fp8_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ) + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + 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), + name="c_dtype", + default=jnp.bfloat16, + ) + d_dtype = require_dtype( + d_dtype, + (jnp.float16, jnp.bfloat16, jnp.float8_e4m3fn, jnp.float8_e5m2), + name="d_dtype", + default=jnp.bfloat16, + ) + 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 = kernel.require_mma_tiler(mma_tiler_mn) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == kernel.TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = kernel.require_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, + norm_const_tensor, + ] + input_specs = [a_spec, b_spec, scale_spec, scale_spec, None, None, 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 = [ + BufferSpec("c_tensor", (1, m, n), c_dtype, tensor_spec=output_spec), + BufferSpec("d_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec), + BufferSpec("d_col_tensor", (1, m, output_n), d_dtype, tensor_spec=output_spec), + ] + if has_amax: + output_specs.append(BufferSpec("amax_tensor", (experts, 1), jnp.float32, fill_value=-float("inf"))) + output_specs.extend( + ( + BufferSpec( + "sfd_row_tensor", + block_scale_shape(m, output_n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + tensor_spec=scale_spec, + ), + BufferSpec( + "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, + "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] + sfd_col_tensor = results[result_idx + 1] + 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: Any, + 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_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: Any, + 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: Any, + 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) + + +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: Any, + 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 MXFP8 contiguous 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..f74fad3b6 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/__init__.py @@ -1,12 +1,13 @@ -# 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, -) +"""Grouped GEMM weight-gradient APIs for SM100.""" + +from ..._operation_api import make_operation_api -__all__ = [ +_API_EXPORTS = ( "GroupedGemmWgradSm100", "grouped_gemm_wgrad_wrapper_sm100", -] +) + +__all__, __getattr__ = make_operation_api(globals(), exports=_API_EXPORTS) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py index 29b3d7328..516e59848 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py @@ -14,7 +14,7 @@ from cuda.bindings import driver as cuda from cutlass.cute.runtime import from_dlpack, make_fake_stream -from cudnn.api_base import APIBase, TensorDesc, TupleDict, ceil_div, is_power_of_2 +from cudnn.api_base import ApiBaseTorch, TorchTensorDesc, TupleDict, ceil_div, is_power_of_2 from cudnn.datatypes import _convert_to_cutlass_data_type from cudnn.discrete_grouped_gemm.discrete_kernel_utils import _require_pointer_tensor @@ -32,7 +32,7 @@ def _normalize_input_order(input_order: WGradInputOrder | str) -> WGradInputOrde return WGradInputOrder(input_order) -class GroupedGemmWgradSm100(APIBase): +class GroupedGemmWgradSm100(ApiBaseTorch): """Unified grouped GEMM wgrad FE API for SM100+ GPUs.""" def __init__( @@ -94,12 +94,13 @@ def __init__( self.expert_cnt = self.wgrad_desc.shape[0] self.wgrad_shape = self.wgrad_desc.shape[1:] self.wgrad_dtype = self.wgrad_desc.dtype - self.single_expert_wgrad_desc = TensorDesc( + self.single_expert_wgrad_desc = TorchTensorDesc( dtype=self.wgrad_desc.dtype, shape=self.wgrad_desc.shape[1:], stride=self.wgrad_desc.stride[1:], 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, + packing=self.wgrad_desc.packing, name="single_expert_wgrad", ) else: # MoEWeightMode.DISCRETE @@ -113,7 +114,7 @@ def __init__( name="sample_wgrad_expert", ) else: - self.single_expert_wgrad_desc = TensorDesc( + self.single_expert_wgrad_desc = TorchTensorDesc( dtype=wgrad_dtype, shape=self.wgrad_shape, stride=(self.wgrad_shape[1], 1), 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..8594835d7 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/jax.py @@ -0,0 +1,483 @@ +# 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 + +import os +from typing import Any, Optional + +import jax.numpy as jnp +from cutlass.jax import TensorSpec + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + as_dtype, + call_cutedsl, + require_array, + require_dtype, +) +from ..._jax.gemm import require_16_byte_extent +from ..._jax.grouped_gemm import grouped_workspace_tensor_spec +from ...gemm_validation import ceil_div, resolve_max_active_clusters + + +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, + cluster_overlap_margin: int, +): + import cutlass + 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] == BlockScaledMoEGroupedGemmWgradKernel.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), + ) + max_active_clusters = resolve_max_active_clusters( + cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_shape_mn[0] * cluster_shape_mn[1]), + cluster_overlap_margin, + ) + 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, + 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. The weight gradient and + TMA-descriptor workspace are XLA-owned buffers; raw per-expert pointer + outputs are not accepted. With ``accumulate_on_output=True``, XLA supplies + a zero-initialized output to the reducing epilogue. + + ``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. + """ + + from .moe_blockscaled_grouped_gemm_wgrad import ( + BlockScaledMoEGroupedGemmWgradKernel, + ) + + kernel = BlockScaledMoEGroupedGemmWgradKernel + 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 != kernel.FP8_SF_VEC_SIZE: + raise ValueError("The JAX FP8 grouped-wgrad path requires " f"sf_vec_size={kernel.FP8_SF_VEC_SIZE}, 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}") + 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, + ) + wgrad_dtype = require_dtype( + wgrad_dtype, + (jnp.bfloat16, jnp.float16, jnp.float32), + name="wgrad_dtype", + default=jnp.bfloat16, + ) + mma_tiler_mn = kernel.require_mma_tiler(mma_tiler_mn) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == kernel.TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = kernel.require_cluster_shape( + cluster_shape_mn, + mma_tiler_mn=mma_tiler_mn, + ) + input_order_value = kernel.require_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)) + + output_shape = (expert_cnt, hidden, intermediate) + (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=( + BufferSpec( + "wgrad_tensor", + output_shape, + wgrad_dtype, + tensor_spec=wgrad_output_tensor_spec(), + fill_value=0.0 if accumulate_on_output else None, + ), + ), + workspaces=( + BufferSpec( + "workspace", + ( + kernel.get_dense_workspace_bytes( + expert_cnt, + input_order_value, + ), + ), + jnp.uint8, + tensor_spec=grouped_workspace_tensor_spec(), + ), + ), + input_specs=input_specs, + ) + return TupleDict(wgrad_tensor=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", + ) -> 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"), + } + 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._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, + ) -> TupleDict: + return super().__call__( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + offsets_tensor, + global_scale_a, + global_scale_b, + ) + + 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, + ) -> 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, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_wgrad_impl(**values, **self._config) + + +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", +) -> TupleDict: + """Compute dense expert weight gradients using the SM100 grouped kernel.""" + + 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, + ) + return op( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + offsets_tensor, + global_scale_a, + global_scale_b, + ) + + +__all__ = [ + "GroupedGemmWgradSm100", + "grouped_gemm_wgrad_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_blockscaled_grouped_gemm_wgrad.py b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_blockscaled_grouped_gemm_wgrad.py index 0f9c2e195..d4980ae0a 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_blockscaled_grouped_gemm_wgrad.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_blockscaled_grouped_gemm_wgrad.py @@ -39,6 +39,10 @@ epilogue_tmem_copy_and_partition, epilogue_smem_copy_and_partition, ) +from ...gemm_validation import ( + require_cluster_shape as _require_cluster_shape, + require_mma_tiler as _require_mma_tiler, +) from ..moe_persistent_scheduler import ( MoEPersistentTileScheduler, MoESchedulerParams, @@ -70,6 +74,82 @@ class BlockScaledMoEGroupedGemmWgradKernel: :param weight_mode: ``MoEWeightMode.DENSE`` or ``MoEWeightMode.DISCRETE`` for output. """ + MMA_TILER_M = (128, 256) + MMA_TILER_N = (128, 256) + TWO_CTA_MMA_TILER_M = 256 + MAX_CLUSTER_CTAS = 16 + MAX_CLUSTER_DIMENSION = 4 + CLUSTER_TILER_M = (128, 256) + SF_VEC_SIZES = (16, 32) + FP8_SF_VEC_SIZE = 32 + + @classmethod + def require_mma_tiler(cls, mma_tiler_mn: Tuple[int, int]) -> Tuple[int, int]: + """Validate an FE-supported grouped-wgrad MMA tile.""" + + return _require_mma_tiler( + mma_tiler_mn, + allowed_m=cls.MMA_TILER_M, + allowed_n=cls.MMA_TILER_N, + ) + + @classmethod + def require_cluster_shape( + cls, + cluster_shape_mn: Tuple[int, int], + *, + mma_tiler_mn: Tuple[int, int], + ) -> Tuple[int, int]: + """Validate a cluster shape for the selected grouped-wgrad tile.""" + + cluster_shape_mn = _require_cluster_shape( + cluster_shape_mn, + mma_m=mma_tiler_mn[0], + two_cta_mma_m=cls.TWO_CTA_MMA_TILER_M, + max_ctas=cls.MAX_CLUSTER_CTAS, + max_dimension=cls.MAX_CLUSTER_DIMENSION, + ) + cta_group_size = 2 if mma_tiler_mn[0] == cls.TWO_CTA_MMA_TILER_M else 1 + cluster_tiler_m = cluster_shape_mn[0] // cta_group_size * mma_tiler_mn[0] + if cluster_tiler_m not in cls.CLUSTER_TILER_M: + raise ValueError( + f"cluster M tile must be one of {cls.CLUSTER_TILER_M}, " + f"got {cluster_tiler_m} from MMA tile {mma_tiler_mn} " + f"and cluster {cluster_shape_mn}" + ) + return cluster_shape_mn + + @staticmethod + def require_input_order( + input_order: WGradInputOrder | str, + ) -> WGradInputOrder: + """Normalize a supported grouped-wgrad input packing order.""" + + 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 + + @classmethod + def get_dense_workspace_bytes( + cls, + expert_cnt: int, + input_order: WGradInputOrder | str = WGradInputOrder.Tensor2D, + ) -> int: + """Return TMA-descriptor workspace bytes for dense output mode.""" + + if expert_cnt < 0: + raise ValueError(f"expert_cnt must be non-negative, got {expert_cnt}") + input_order = cls.require_input_order(input_order) + return WgradSfTensormapConstructor.get_workspace_size( + input_order, + MoEWeightMode.DENSE, + expert_cnt, + ) + def __init__( self, sf_vec_size: int, @@ -120,7 +200,16 @@ def __init__( # ------------------------------------------------------------------ def get_workspace_bytes(self) -> int: - return WgradSfTensormapConstructor.get_workspace_size(self.input_order, self.weight_mode, self.expert_cnt) + if self.weight_mode == MoEWeightMode.DENSE: + return self.get_dense_workspace_bytes( + self.expert_cnt, + self.input_order, + ) + return WgradSfTensormapConstructor.get_workspace_size( + self.input_order, + self.weight_mode, + self.expert_cnt, + ) # ------------------------------------------------------------------ # _setup_attributes diff --git a/python/cudnn/grouped_gemm/moe_kernel_helpers.py b/python/cudnn/grouped_gemm/moe_kernel_helpers.py index ab41414da..95d73b624 100644 --- a/python/cudnn/grouped_gemm/moe_kernel_helpers.py +++ b/python/cudnn/grouped_gemm/moe_kernel_helpers.py @@ -38,13 +38,15 @@ - 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 -import cutlass.cute.testing as testing from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cutlass_dsl import T, dsl_user_op import cutlass.utils as utils @@ -455,6 +457,9 @@ 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..2d2c77090 100644 --- a/python/cudnn/grouped_gemm/utils.py +++ b/python/cudnn/grouped_gemm/utils.py @@ -33,7 +33,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 +57,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 +232,9 @@ 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 new file mode 100644 index 000000000..f2162ffdd --- /dev/null +++ b/python/cudnn/jax/__init__.py @@ -0,0 +1,242 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX facade for co-located frontend-only operation APIs. + +Importing :mod:`cudnn.jax` is the explicit dependency boundary for JAX and +CuTe DSL. Dependencies are validated once here, public wrapper modules are then +imported, and architecture-specific kernel modules remain deferred until an +operation is traced. +""" + +from importlib import import_module +from types import SimpleNamespace + +_INSTALL_HINT = "pip install 'nvidia-cudnn-frontend[jax]'" + + +def _require_dependencies(module_names: list[str]) -> None: + for module_name in module_names: + try: + import_module(module_name) + except ModuleNotFoundError as error: + missing_name = error.name + if missing_name is None or not (module_name == missing_name or module_name.startswith(f"{missing_name}.")): + raise + raise ImportError(f"cudnn.jax requires the {module_name!r} module. Install the " f"JAX integration with `{_INSTALL_HINT}`.") from error + + +_require_dependencies(["jax", "cutlass"]) + +import cutlass.jax +import jax + +if not cutlass.jax.is_available(): + minimum_version_info = tuple(cutlass.jax.CUTE_DSL_MIN_SUPPORTED_JAX_VERSION) + minimum_version = ".".join(str(part) for part in minimum_version_info) + installed_version = getattr(jax, "__version__", "unknown") + reason = f"CUTLASS JAX support is unavailable with JAX {installed_version}; " f"the minimum supported JAX version is {minimum_version}." + raise ImportError(f"cudnn.jax cannot be imported because {reason} Install the JAX " f"integration with `{_INSTALL_HINT}`.") + +from .._jax.api_base import ApiBaseJax, JaxTensorDesc, TupleDict + +from ..deepseek_sparse_attention.indexer_forward.jax import ( + IndexerForward, + indexer_forward_wrapper, +) +from ..deepseek_sparse_attention.indexer_backward.jax import ( + DenseIndexerBackward, + IndexerBackward, + dense_indexer_backward_wrapper, + indexer_backward_wrapper, +) +from ..deepseek_sparse_attention.indexer_top_k.jax import ( + IndexerTopK, + compactify_wrapper, + indexer_top_k_wrapper, + local_to_global_wrapper, +) +from ..deepseek_sparse_attention.score_recompute.jax import ( + DenseAttnScoreRecompute, + DenseIndexerScoreRecompute, + SparseAttnScoreRecompute, + SparseIndexerScoreRecompute, + dense_attn_score_recompute_wrapper, + dense_indexer_score_recompute_wrapper, + sparse_attn_score_recompute_wrapper, + sparse_indexer_score_recompute_wrapper, +) +from ..deepseek_sparse_attention.sparse_attention_backward.jax import ( + SparseAttentionBackward, + sparse_attention_backward_wrapper, +) +from ..gemm_amax.jax import GemmAmaxSm100, gemm_amax_wrapper_sm100 +from ..gemm_dsrelu.jax import ( + GemmDsreluSm100, + gemm_dsrelu_wrapper_sm100, +) +from ..gemm_srelu.jax import GemmSreluSm100, gemm_srelu_wrapper_sm100 +from ..gemm_swiglu.jax import ( + GemmSwigluSm100, + gemm_swiglu_wrapper_sm100, +) +from ..grouped_gemm.grouped_gemm_dswiglu.jax import ( + GroupedGemmDswigluSm100, + grouped_gemm_dswiglu_wrapper_sm100, +) +from ..grouped_gemm.grouped_gemm_dglu.jax import ( + GroupedGemmDgluSm100, + grouped_gemm_dglu_wrapper_sm100, +) +from ..grouped_gemm.grouped_gemm_dsrelu.jax import ( + GroupedGemmDsreluSm100, + grouped_gemm_dsrelu_wrapper_sm100, +) +from ..grouped_gemm.grouped_gemm_glu.jax import ( + GroupedGemmGluSm100, + grouped_gemm_glu_wrapper_sm100, +) +from ..grouped_gemm.grouped_gemm_glu_hadamard.jax import ( + GroupedGemmGluHadamardSm100, + grouped_gemm_glu_hadamard_wrapper_sm100, +) +from ..grouped_gemm.grouped_gemm_quant.jax import ( + GroupedGemmQuantSm100, + grouped_gemm_quant_wrapper_sm100, +) +from ..grouped_gemm.grouped_gemm_srelu.jax import ( + GroupedGemmSreluSm100, + grouped_gemm_srelu_wrapper_sm100, +) +from ..grouped_gemm.grouped_gemm_swiglu.jax import ( + GroupedGemmSwigluSm100, + grouped_gemm_swiglu_wrapper_sm100, +) +from ..grouped_gemm.grouped_gemm_wgrad.jax import ( + GroupedGemmWgradSm100, + grouped_gemm_wgrad_wrapper_sm100, +) +from ..native_sparse_attention.compression.jax import ( + CompressionAttention, + compression_attention_wrapper, +) +from ..native_sparse_attention.selection.jax import ( + SelectionAttention, + selection_attention_wrapper, +) +from ..native_sparse_attention.sliding_window_attention.jax import ( + SlidingWindowAttention, + sliding_window_attention_wrapper, +) +from ..native_sparse_attention.top_k.jax import ( + TopKReduction, + topk_reduction_wrapper, +) +from ..rmsnorm_rht_amax.jax import ( + RmsNormRhtAmaxSm100, + rmsnorm_rht_amax_sm100, +) +from ..sdpa.bwd.jax import SdpabwdSm100D256, sdpa_bwd_wrapper_sm100_d256 +from ..sdpa.fwd.jax import SdpafwdSm100D256, sdpa_fwd_wrapper_sm100_d256 + +DSA = SimpleNamespace( + DenseIndexerBackward=DenseIndexerBackward, + IndexerBackward=IndexerBackward, + IndexerForward=IndexerForward, + IndexerTopK=IndexerTopK, + SparseIndexerScoreRecompute=SparseIndexerScoreRecompute, + SparseAttnScoreRecompute=SparseAttnScoreRecompute, + DenseIndexerScoreRecompute=DenseIndexerScoreRecompute, + DenseAttnScoreRecompute=DenseAttnScoreRecompute, + SparseAttentionBackward=SparseAttentionBackward, + dense_indexer_backward_wrapper=dense_indexer_backward_wrapper, + indexer_backward_wrapper=indexer_backward_wrapper, + indexer_forward_wrapper=indexer_forward_wrapper, + indexer_top_k_wrapper=indexer_top_k_wrapper, + local_to_global_wrapper=local_to_global_wrapper, + compactify_wrapper=compactify_wrapper, + sparse_indexer_score_recompute_wrapper=sparse_indexer_score_recompute_wrapper, + sparse_attn_score_recompute_wrapper=sparse_attn_score_recompute_wrapper, + dense_indexer_score_recompute_wrapper=dense_indexer_score_recompute_wrapper, + dense_attn_score_recompute_wrapper=dense_attn_score_recompute_wrapper, + sparse_attention_backward_wrapper=sparse_attention_backward_wrapper, +) + +NSA = SimpleNamespace( + CompressionAttention=CompressionAttention, + SelectionAttention=SelectionAttention, + SlidingWindowAttention=SlidingWindowAttention, + TopKReduction=TopKReduction, + compression_attention_wrapper=compression_attention_wrapper, + selection_attention_wrapper=selection_attention_wrapper, + sliding_window_attention_wrapper=sliding_window_attention_wrapper, + topk_reduction_wrapper=topk_reduction_wrapper, +) + +__all__ = [ + "ApiBaseJax", + "DSA", + "NSA", + "CompressionAttention", + "DenseAttnScoreRecompute", + "DenseIndexerBackward", + "DenseIndexerScoreRecompute", + "GemmAmaxSm100", + "GemmDsreluSm100", + "GemmSreluSm100", + "GemmSwigluSm100", + "GroupedGemmDswigluSm100", + "GroupedGemmDgluSm100", + "GroupedGemmDsreluSm100", + "GroupedGemmGluHadamardSm100", + "GroupedGemmGluSm100", + "GroupedGemmQuantSm100", + "GroupedGemmSreluSm100", + "GroupedGemmSwigluSm100", + "GroupedGemmWgradSm100", + "IndexerBackward", + "IndexerForward", + "IndexerTopK", + "JaxTensorDesc", + "RmsNormRhtAmaxSm100", + "SdpabwdSm100D256", + "SdpafwdSm100D256", + "SelectionAttention", + "SlidingWindowAttention", + "SparseAttentionBackward", + "SparseAttnScoreRecompute", + "SparseIndexerScoreRecompute", + "TopKReduction", + "TupleDict", + "compactify_wrapper", + "compression_attention_wrapper", + "dense_attn_score_recompute_wrapper", + "dense_indexer_backward_wrapper", + "dense_indexer_score_recompute_wrapper", + "indexer_forward_wrapper", + "indexer_backward_wrapper", + "indexer_top_k_wrapper", + "gemm_amax_wrapper_sm100", + "gemm_dsrelu_wrapper_sm100", + "gemm_srelu_wrapper_sm100", + "gemm_swiglu_wrapper_sm100", + "grouped_gemm_dswiglu_wrapper_sm100", + "grouped_gemm_dglu_wrapper_sm100", + "grouped_gemm_dsrelu_wrapper_sm100", + "grouped_gemm_glu_hadamard_wrapper_sm100", + "grouped_gemm_glu_wrapper_sm100", + "grouped_gemm_quant_wrapper_sm100", + "grouped_gemm_srelu_wrapper_sm100", + "grouped_gemm_swiglu_wrapper_sm100", + "grouped_gemm_wgrad_wrapper_sm100", + "local_to_global_wrapper", + "rmsnorm_rht_amax_sm100", + "sdpa_bwd_wrapper_sm100_d256", + "sdpa_fwd_wrapper_sm100_d256", + "selection_attention_wrapper", + "sliding_window_attention_wrapper", + "sparse_attention_backward_wrapper", + "sparse_attn_score_recompute_wrapper", + "sparse_indexer_score_recompute_wrapper", + "topk_reduction_wrapper", +] diff --git a/python/cudnn/native_sparse_attention/__init__.py b/python/cudnn/native_sparse_attention/__init__.py index 6d4f8220b..d6d7c0b12 100644 --- a/python/cudnn/native_sparse_attention/__init__.py +++ b/python/cudnn/native_sparse_attention/__init__.py @@ -1,28 +1,51 @@ -from .selection import SelectionAttention, selection_attention_wrapper -from .compression import CompressionAttention, compression_attention_wrapper -from .sliding_window_attention import ( - SlidingWindowAttention, - sliding_window_attention_wrapper, -) -from .top_k import TopKReduction, topk_reduction_wrapper +"""Lazy public namespace for Native Sparse Attention operations.""" + +from importlib import import_module + +_SYMBOLS = { + "SelectionAttention": (".selection", "SelectionAttention"), + "selection_attention_wrapper": (".selection", "selection_attention_wrapper"), + "CompressionAttention": (".compression", "CompressionAttention"), + "compression_attention_wrapper": ( + ".compression", + "compression_attention_wrapper", + ), + "SlidingWindowAttention": ( + ".sliding_window_attention", + "SlidingWindowAttention", + ), + "sliding_window_attention_wrapper": ( + ".sliding_window_attention", + "sliding_window_attention_wrapper", + ), + "TopKReduction": (".top_k", "TopKReduction"), + "topk_reduction_wrapper": (".top_k", "topk_reduction_wrapper"), +} + + +def _load_symbol(name): + module_name, symbol_name = _SYMBOLS[name] + module = import_module(module_name, package=__name__) + symbol = getattr(module, symbol_name) + globals()[name] = symbol + return symbol + + +def __getattr__(name): + if name == "NSA": + return NSA + if name in _SYMBOLS: + return _load_symbol(name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") class NSANamespace: - SelectionAttention = staticmethod(SelectionAttention) - selection_attention_wrapper = staticmethod(selection_attention_wrapper) - - 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 __getattr__(self, name): + if name in _SYMBOLS: + return _load_symbol(name) + raise AttributeError(f"NSA has no attribute {name!r}") NSA = NSANamespace() -__all__ = [ - "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..f6e39fbfc 100644 --- a/python/cudnn/native_sparse_attention/compression/__init__.py +++ b/python/cudnn/native_sparse_attention/compression/__init__.py @@ -1,6 +1,10 @@ -from .api import CompressionAttention, compression_attention_wrapper +"""Lazy API exports for NSA compression attention.""" -__all__ = [ - "CompressionAttention", - "compression_attention_wrapper", -] +from ..._operation_api import make_operation_api + +_API_EXPORTS = ("CompressionAttention", "compression_attention_wrapper") + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/native_sparse_attention/compression/api.py b/python/cudnn/native_sparse_attention/compression/api.py index fc89c004d..879592684 100644 --- a/python/cudnn/native_sparse_attention/compression/api.py +++ b/python/cudnn/native_sparse_attention/compression/api.py @@ -9,7 +9,7 @@ from cutlass.cute.runtime import make_fake_stream from cutlass.cute.typing import Int32 -from cudnn.api_base import APIBase, TupleDict +from cudnn.api_base import ApiBaseTorch, TupleDict from cudnn.datatypes import _convert_to_cutlass_data_type from ..utils import make_tensor_strided_like @@ -17,7 +17,7 @@ from . import fmha_helpers as fmha_utils -class CompressionAttention(APIBase): +class CompressionAttention(ApiBaseTorch): def __init__( self, sample_q: torch.Tensor, 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..66e974b0f --- /dev/null +++ b/python/cudnn/native_sparse_attention/compression/jax.py @@ -0,0 +1,321 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for fixed-shape SM100 NSA compression attention.""" + +from __future__ import annotations + +import math +from typing import Any + +import jax.numpy as jnp + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + call_cutedsl, + require_dtype, +) +from ..jax_utils import ( + bhs_lse_as_bsh_spec, + bhsd_storage_spec, + require_bhsd_qkv, +) + + +def _launch( + stream, + q, + k, + v, + output, + lse=None, + *, + batch: int, + seqlen_q: int, + seqlen_k: int, + num_query_heads: int, + num_kv_heads: int, + head_dim: int, + enable_lse: bool, + is_persistent: bool, + scale_softmax: float, + scale_output: float, +): + 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, head_dim), + is_persistent=is_persistent, + mask_type=MaskType.COMPRESSED_CAUSAL_MASK, + ) + problem_size = tuple( + Int32(value) + for value in ( + batch, + seqlen_q, + seqlen_q, + seqlen_k, + num_query_heads, + num_kv_heads, + head_dim, + ) + ) + + kernel( + q, + k, + v, + output, + problem_size, + None, + None, + lse if enable_lse else None, + Float32(scale_softmax * math.log2(math.e)), + Float32(scale_softmax), + Float32(scale_output), + None, + Int32(0), + stream, + ) + + +def _compression_attention_impl( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + 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, + *, + _validate_only: bool = False, +) -> TupleDict: + """Compute fixed-shape BHSD compression attention on SM100. + + Q, K, and V use logical shapes ``(B, H, S, D)`` and a shared ``float16`` + or ``bfloat16`` dtype. The head dimension must be 32, 64, or 128. The + returned O has Q's shape and ``o_dtype``; optional LSE has shape + ``(B, H_q, S_q)`` and dtype ``float32``. + + Variable-length THD inputs are not part of this API. Configuration values + must be static while tracing with :func:`jax.jit`. + """ + + ( + batch, + num_query_heads, + num_kv_heads, + seqlen_q, + seqlen_k, + head_dim, + input_dtype, + ) = require_bhsd_qkv(q_tensor, k_tensor, v_tensor) + + require_dtype( + qk_acc_dtype, + (jnp.float32,), + name="qk_acc_dtype", + default=jnp.float32, + ) + require_dtype( + pv_acc_dtype, + (jnp.float32,), + name="pv_acc_dtype", + default=jnp.float32, + ) + output_dtype = require_dtype( + o_dtype, + (jnp.float16, jnp.bfloat16), + name="o_dtype", + default=input_dtype, + ) + if mma_tiler_mn != (128, 128): + raise ValueError(f"mma_tiler_mn must be (128, 128), got {mma_tiler_mn}") + if seqlen_q < seqlen_k or seqlen_q % seqlen_k: + raise ValueError("Compression attention requires S_q to be an integer multiple of " f"S_k, got S_q={seqlen_q} and S_k={seqlen_k}") + + base_softmax_scale = 1.0 / math.sqrt(head_dim) if scale_softmax is None else float(scale_softmax) + resolved_softmax_scale = float(scale_q) * float(scale_k) * base_softmax_scale + resolved_output_scale = float(scale_v) * float(inv_scale_o) + if _validate_only: + return None + + input_spec = bhsd_storage_spec(present_as_bshd=True) + lse_spec = bhs_lse_as_bsh_spec() + + outputs = [ + BufferSpec( + "o_tensor", + tuple(q_tensor.shape), + output_dtype, + tensor_spec=input_spec, + ) + ] + if enable_lse: + outputs.append( + BufferSpec( + "lse_tensor", + (batch, num_query_heads, seqlen_q), + jnp.float32, + tensor_spec=lse_spec, + ) + ) + + results = call_cutedsl( + _launch, + (q_tensor, k_tensor, v_tensor), + outputs=tuple(outputs), + input_specs=(input_spec, input_spec, input_spec), + static_args={ + "batch": batch, + "seqlen_q": seqlen_q, + "seqlen_k": seqlen_k, + "num_query_heads": num_query_heads, + "num_kv_heads": num_kv_heads, + "head_dim": head_dim, + "enable_lse": bool(enable_lse), + "is_persistent": bool(is_persistent), + "scale_softmax": resolved_softmax_scale, + "scale_output": resolved_output_scale, + }, + ) + return TupleDict( + o_tensor=results[0], + lse_tensor=results[1] if enable_lse else None, + ) + + +class CompressionAttention(ApiBaseJax): + """Sample-signature-bound JAX callable for SM100 compression attention.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + 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, + ) -> None: + super().__init__() + self.q_desc = self.make_tensor_desc(sample_q, name="sample_q") + self.k_desc = self.make_tensor_desc(sample_k, name="sample_k") + self.v_desc = self.make_tensor_desc(sample_v, name="sample_v") + self.enable_lse = enable_lse + self.o_dtype = self.as_optional_dtype(o_dtype) + self.qk_acc_dtype = self.as_optional_dtype(qk_acc_dtype) + self.pv_acc_dtype = self.as_optional_dtype(pv_acc_dtype) + self.mma_tiler_mn = tuple(mma_tiler_mn) + self.is_persistent = is_persistent + self.scale_q = scale_q + self.scale_k = scale_k + self.scale_v = scale_v + self.inv_scale_o = inv_scale_o + self.scale_softmax = scale_softmax + + def _check_support(self) -> None: + _compression_attention_impl( + self.q_desc, + self.k_desc, + self.v_desc, + self.enable_lse, + self.o_dtype, + self.qk_acc_dtype, + self.pv_acc_dtype, + self.mma_tiler_mn, + self.is_persistent, + self.scale_q, + self.scale_k, + self.scale_v, + self.inv_scale_o, + self.scale_softmax, + _validate_only=True, + ) + + def __call__(self, q_tensor: Any, k_tensor: Any, v_tensor: Any) -> TupleDict: + return super().__call__(q_tensor, k_tensor, v_tensor) + + def _call_impl(self, q_tensor: Any, k_tensor: Any, v_tensor: Any) -> TupleDict: + self.check_tensor_signature(q_tensor, self.q_desc, name="Q") + self.check_tensor_signature(k_tensor, self.k_desc, name="K") + self.check_tensor_signature(v_tensor, self.v_desc, name="V") + return _compression_attention_impl( + q_tensor, + k_tensor, + v_tensor, + self.enable_lse, + self.o_dtype, + self.qk_acc_dtype, + self.pv_acc_dtype, + self.mma_tiler_mn, + self.is_persistent, + self.scale_q, + self.scale_k, + self.scale_v, + self.inv_scale_o, + self.scale_softmax, + ) + + +def compression_attention_wrapper( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + 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, +) -> TupleDict: + """Compute fixed-shape BHSD compression attention on SM100.""" + + return CompressionAttention( + q_tensor, + k_tensor, + v_tensor, + 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, + )(q_tensor, k_tensor, v_tensor) + + +__all__ = [ + "CompressionAttention", + "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..b802f47ef --- /dev/null +++ b/python/cudnn/native_sparse_attention/jax_utils.py @@ -0,0 +1,106 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared validation and layout helpers for JAX NSA operations.""" + +from __future__ import annotations + +from typing import Any + +import jax.numpy as jnp +from cutlass.jax import TensorSpec + +from .._jax.api_base import as_dtype, require_array + + +def bhsd_storage_spec(*, present_as_bshd: bool) -> TensorSpec: + """Describe logical BHSD data backed by compact BSHD storage. + + NSA Torch tensors use logical ``(B, H, S, D)`` shapes with H inside S in + physical memory. Some kernels consume that BHSD mode order directly while + compression attention consumes a ``(B, S, H, D)`` view. + """ + + return TensorSpec( + layout=(3, 1, 2, 0), + mode=(0, 2, 1, 3) if present_as_bshd else (0, 1, 2, 3), + ) + + +def bhs_lse_as_bsh_spec() -> TensorSpec: + """Present a public ``(B, H, S)`` LSE array as ``(B, S, H)``.""" + + return TensorSpec(layout=(2, 1, 0), mode=(0, 2, 1)) + + +def require_bhsd_qkv( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any | None = None, +) -> tuple[int, int, int, int, int, int, Any]: + """Validate fixed-shape NSA BHSD inputs and return their dimensions.""" + + q_shape = require_array( + q_tensor, + name="q_tensor", + rank=4, + dtype=(jnp.float16, jnp.bfloat16), + ) + dtype = as_dtype(q_tensor) + k_shape = require_array(k_tensor, name="k_tensor", rank=4, dtype=dtype) + v_shape = None + if v_tensor is not None: + v_shape = require_array(v_tensor, name="v_tensor", rank=4, dtype=dtype) + + batch, num_query_heads, seqlen_q, head_dim = q_shape + k_batch, num_kv_heads, seqlen_k, k_head_dim = k_shape + dimensions = { + "batch": batch, + "H_q": num_query_heads, + "H_kv": num_kv_heads, + "S_q": seqlen_q, + "S_kv": seqlen_k, + "D": head_dim, + } + nonpositive = [f"{name}={value}" for name, value in dimensions.items() if value <= 0] + if nonpositive: + raise ValueError("NSA dimensions must be positive, got " + ", ".join(nonpositive)) + if k_batch != batch: + raise ValueError(f"q_tensor and k_tensor batch dimensions must match, got {batch} and {k_batch}") + if k_head_dim != head_dim: + raise ValueError("q_tensor and k_tensor head dimensions must match, got " f"{head_dim} and {k_head_dim}") + if head_dim not in (32, 64, 128): + raise ValueError(f"head dimension must be 32, 64, or 128, 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_shape is not None: + v_batch, num_value_heads, v_seqlen, value_dim = v_shape + if (v_batch, num_value_heads, v_seqlen) != ( + batch, + num_kv_heads, + seqlen_k, + ): + raise ValueError( + "k_tensor and v_tensor batch, head, and sequence dimensions " + f"must match, got {(batch, num_kv_heads, seqlen_k)} and " + f"{(v_batch, num_value_heads, v_seqlen)}" + ) + if value_dim != head_dim: + raise ValueError(f"V head dimension must match Q/K ({head_dim}), got {value_dim}") + return ( + batch, + num_query_heads, + num_kv_heads, + seqlen_q, + seqlen_k, + head_dim, + dtype, + ) + + +__all__ = [ + "bhs_lse_as_bsh_spec", + "bhsd_storage_spec", + "require_bhsd_qkv", +] diff --git a/python/cudnn/native_sparse_attention/selection/NSA_select_attn_fwd_hmma.py b/python/cudnn/native_sparse_attention/selection/NSA_select_attn_fwd_hmma.py index 1360a66b5..d59ac9975 100644 --- a/python/cudnn/native_sparse_attention/selection/NSA_select_attn_fwd_hmma.py +++ b/python/cudnn/native_sparse_attention/selection/NSA_select_attn_fwd_hmma.py @@ -568,9 +568,10 @@ def kernel( seq_len = seq_offsets[offset_idx + 1] - seq_offsets[offset_idx] offset = seq_offsets[offset_idx] - for i in cutlass.range((block_indices.shape[2] + self.threads_per_block - 1) // self.threads_per_block): - if i * self.threads_per_block + tidx < block_indices.shape[2]: - sIDX[i * self.threads_per_block + tidx] = block_indices[offset + t, KV_head_idx, i * self.threads_per_block + tidx] + if t < seq_len: + for i in cutlass.range((block_indices.shape[2] + self.threads_per_block - 1) // self.threads_per_block): + if i * self.threads_per_block + tidx < block_indices.shape[2]: + sIDX[i * self.threads_per_block + tidx] = block_indices[offset + t, KV_head_idx, i * self.threads_per_block + tidx] cute.arch.sync_threads() seq_len_aligned = (seq_len + self.tile_shape_mnk_QK[1] - 1) // self.tile_shape_mnk_QK[1] * self.tile_shape_mnk_QK[1] diff --git a/python/cudnn/native_sparse_attention/selection/__init__.py b/python/cudnn/native_sparse_attention/selection/__init__.py index 576723053..0345b8901 100644 --- a/python/cudnn/native_sparse_attention/selection/__init__.py +++ b/python/cudnn/native_sparse_attention/selection/__init__.py @@ -1,6 +1,10 @@ -from .api import SelectionAttention, selection_attention_wrapper +"""Lazy API exports for NSA selection attention.""" -__all__ = [ - "SelectionAttention", - "selection_attention_wrapper", -] +from ..._operation_api import make_operation_api + +_API_EXPORTS = ("SelectionAttention", "selection_attention_wrapper") + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/native_sparse_attention/selection/api.py b/python/cudnn/native_sparse_attention/selection/api.py index 03ca47c2b..3793314d0 100644 --- a/python/cudnn/native_sparse_attention/selection/api.py +++ b/python/cudnn/native_sparse_attention/selection/api.py @@ -1,6 +1,6 @@ from .NSA_select_attn_fwd_hmma import HopperSelectAttentionFwd from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.api_base import APIBase, TupleDict +from cudnn.api_base import ApiBaseTorch, TupleDict import cutlass import cutlass.cute as cute @@ -11,7 +11,7 @@ import math -class SelectionAttention(APIBase): +class SelectionAttention(ApiBaseTorch): def __init__( self, sample_q: torch.Tensor, 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..f64a93eb4 --- /dev/null +++ b/python/cudnn/native_sparse_attention/selection/jax.py @@ -0,0 +1,415 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for SM90 NSA selection attention.""" + +from __future__ import annotations + +import math +from typing import Any + +import jax.numpy as jnp +from cutlass.jax import jax_to_cutlass_dtype + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + as_dtype, + call_cutedsl, + require_array, + require_dtype, +) + + +def _launch( + stream, + q, + k, + v, + block_indices, + block_counts, + cum_seqlen_q, + cum_seqlen_k, + output, + lse_sum, + row_max, + *, + element_dtype: Any, + head_dim: int, + value_dim: int, + gqa_group_size: int, + block_size: int, + max_s_q: int, + scale_softmax: float, +): + from cutlass import Float32 + + from .NSA_select_attn_fwd_hmma import HopperSelectAttentionFwd + + kernel = HopperSelectAttentionFwd( + head_dim=head_dim, + value_dim=value_dim, + GQA_group_size=gqa_group_size, + block_size=block_size, + dtype=element_dtype, + acc_dtype=Float32, + ) + + # Selection attention is currently self-attention. The second runtime + # offsets operand preserves API parity and is required to match Q's + # offsets, but the native kernel has a single seq_offsets argument. + del cum_seqlen_k + kernel( + q, + k, + v, + output, + lse_sum, + row_max, + block_indices, + block_counts, + max_s_q, + cum_seqlen_q, + Float32(scale_softmax), + stream, + ) + + +def _selection_attention_impl( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + block_indices_tensor: Any, + block_counts_tensor: Any, + cum_seqlen_q_tensor: Any, + cum_seqlen_k_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, + _validate_only: bool = False, +) -> TupleDict: + """Compute packed THD selection attention with the SM90 CuTe kernel. + + Q, K, and V use compact row-major ``(T, H, D)`` storage. Block indices + have shape ``(T, H_kv, K)`` and block counts have shape ``(T, H_kv)``. + Both cumulative-length arrays are runtime ``int32`` or ``int64`` operands of shape + ``(B + 1,)``. ``max_s_q`` and ``max_s_k`` are required static integers. + + The current kernel implements self-attention: Q and KV must have the same + packed token count, the two cumulative-length arrays must contain identical + runtime values, and both static maxima must be equal. Runtime offsets must + start at zero, be nondecreasing, end at T, and describe lengths no greater + than the static maximum. Block counts and indices are trusted to be valid + for their corresponding sequences while tracing with :func:`jax.jit`. + + O has shape ``(T, H_q, D_v)`` and the input dtype. L and M have shape + ``(T, H_q, 1)`` and dtype ``float32``. + """ + + q_shape = require_array( + q_tensor, + name="q_tensor", + rank=3, + dtype=(jnp.float16, jnp.bfloat16), + ) + input_dtype = as_dtype(q_tensor) + k_shape = require_array(k_tensor, name="k_tensor", rank=3, dtype=input_dtype) + v_shape = require_array(v_tensor, name="v_tensor", rank=3, dtype=input_dtype) + + total_tokens, num_query_heads, head_dim = q_shape + k_total_tokens, num_kv_heads, k_head_dim = k_shape + v_total_tokens, num_value_heads, value_dim = v_shape + dimensions = { + "T": total_tokens, + "H_q": num_query_heads, + "H_kv": num_kv_heads, + "D_qk": head_dim, + "D_v": value_dim, + } + nonpositive = [f"{name}={value}" for name, value in dimensions.items() if value <= 0] + if nonpositive: + raise ValueError("Selection-attention dimensions must be positive, got " + ", ".join(nonpositive)) + if (k_total_tokens, v_total_tokens) != (total_tokens, total_tokens): + raise ValueError("Q, K, and V must have the same packed token count for self-attention, " f"got {total_tokens}, {k_total_tokens}, and {v_total_tokens}") + if k_head_dim != head_dim: + raise ValueError(f"Q and K head dimensions must match, got {head_dim} and {k_head_dim}") + if num_value_heads != num_kv_heads: + raise ValueError(f"K and V head counts must match, got {num_kv_heads} and {num_value_heads}") + if head_dim % 16 or value_dim % 16: + raise ValueError(f"D_qk and D_v must be multiples of 16, got {head_dim} and {value_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})") + gqa_group_size = num_query_heads // num_kv_heads + if gqa_group_size not in (1, 2, 4, 8): + raise ValueError(f"H_q / H_kv must be one of {{1, 2, 4, 8}}, got {gqa_group_size}") + + block_indices_shape = require_array( + block_indices_tensor, + name="block_indices_tensor", + rank=3, + dtype=jnp.int32, + ) + if block_indices_shape[:2] != (total_tokens, num_kv_heads) or block_indices_shape[2] <= 0: + raise ValueError("block_indices_tensor must have shape " f"(T, H_kv, K) with K > 0, got {block_indices_shape}") + require_array( + block_counts_tensor, + name="block_counts_tensor", + shape=(total_tokens, num_kv_heads), + dtype=jnp.int32, + ) + + cum_q_shape = require_array( + cum_seqlen_q_tensor, + name="cum_seqlen_q_tensor", + rank=1, + dtype=(jnp.int32, jnp.int64), + ) + if cum_q_shape[0] < 2: + raise ValueError("cum_seqlen_q_tensor must have shape (B + 1,) with B > 0, " f"got {cum_q_shape}") + offsets_dtype = as_dtype(cum_seqlen_q_tensor) + require_array( + cum_seqlen_k_tensor, + name="cum_seqlen_k_tensor", + shape=cum_q_shape, + dtype=offsets_dtype, + ) + + if max_s_q <= 0 or max_s_k <= 0: + raise ValueError(f"max_s_q and max_s_k must be positive, got {max_s_q} and {max_s_k}") + if max_s_q != max_s_k: + raise ValueError(f"max_s_q and max_s_k must be identical, got {max_s_q} and {max_s_k}") + if block_size not in (16, 32, 64): + raise ValueError(f"block_size must be 16, 32, or 64, got {block_size}") + + output_dtype = require_dtype( + o_dtype, + (input_dtype,), + name="o_dtype", + default=input_dtype, + ) + require_dtype( + acc_dtype, + (jnp.float32,), + name="acc_dtype", + default=jnp.float32, + ) + resolved_scale = 1.0 / math.sqrt(head_dim) if scale_softmax is None else float(scale_softmax) + if _validate_only: + return None + + # The native kernel consumes the same singleton-batch views produced by + # Torch's unsqueeze(0). These reshapes are storage-preserving XLA bitcasts. + q_storage = jnp.reshape(q_tensor, (1, total_tokens, num_query_heads, head_dim)) + k_storage = jnp.reshape(k_tensor, (1, total_tokens, num_kv_heads, head_dim)) + v_storage = jnp.reshape(v_tensor, (1, total_tokens, num_kv_heads, value_dim)) + + o_storage, l_storage, m_storage = call_cutedsl( + _launch, + ( + q_storage, + k_storage, + v_storage, + block_indices_tensor, + block_counts_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ), + outputs=( + BufferSpec( + "o_tensor", + (1, total_tokens, num_query_heads, value_dim), + output_dtype, + fill_value=0, + ), + BufferSpec( + "l_tensor", + (1, total_tokens, num_query_heads), + jnp.float32, + fill_value=0.0, + ), + BufferSpec( + "m_tensor", + (1, total_tokens, num_query_heads), + jnp.float32, + fill_value=float("-inf"), + ), + ), + static_args={ + "element_dtype": jax_to_cutlass_dtype(input_dtype), + "head_dim": head_dim, + "value_dim": value_dim, + "gqa_group_size": gqa_group_size, + "block_size": block_size, + "max_s_q": max_s_q, + "scale_softmax": resolved_scale, + }, + ) + return TupleDict( + o_tensor=jnp.reshape(o_storage, (total_tokens, num_query_heads, value_dim)), + l_tensor=jnp.reshape(l_storage, (total_tokens, num_query_heads, 1)), + m_tensor=jnp.reshape(m_storage, (total_tokens, num_query_heads, 1)), + ) + + +class SelectionAttention(ApiBaseJax): + """Sample-signature-bound JAX callable for SM90 selection attention.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_block_indices: Any, + sample_block_counts: Any, + sample_cum_seqlen_q: Any, + sample_cum_seqlen_k: 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, + ) -> None: + super().__init__() + self.q_desc = self.make_tensor_desc(sample_q, name="sample_q") + self.k_desc = self.make_tensor_desc(sample_k, name="sample_k") + self.v_desc = self.make_tensor_desc(sample_v, name="sample_v") + self.block_indices_desc = self.make_tensor_desc(sample_block_indices, name="sample_block_indices") + self.block_counts_desc = self.make_tensor_desc(sample_block_counts, name="sample_block_counts") + self.cum_q_desc = self.make_tensor_desc(sample_cum_seqlen_q, name="sample_cum_seqlen_q") + self.cum_k_desc = self.make_tensor_desc(sample_cum_seqlen_k, name="sample_cum_seqlen_k") + self.max_s_q = max_s_q + self.max_s_k = max_s_k + self.block_size = block_size + self.scale_softmax = scale_softmax + self.o_dtype = self.as_optional_dtype(o_dtype) + self.acc_dtype = self.as_optional_dtype(acc_dtype) + + def _check_support(self) -> None: + _selection_attention_impl( + self.q_desc, + self.k_desc, + self.v_desc, + self.block_indices_desc, + self.block_counts_desc, + self.cum_q_desc, + self.cum_k_desc, + max_s_q=self.max_s_q, + max_s_k=self.max_s_k, + block_size=self.block_size, + scale_softmax=self.scale_softmax, + o_dtype=self.o_dtype, + acc_dtype=self.acc_dtype, + _validate_only=True, + ) + + def __call__( + self, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + block_indices_tensor: Any, + block_counts_tensor: Any, + cum_seqlen_q_tensor: Any, + cum_seqlen_k_tensor: Any, + ) -> TupleDict: + return super().__call__( + q_tensor, + k_tensor, + v_tensor, + block_indices_tensor, + block_counts_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + + def _call_impl( + self, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + block_indices_tensor: Any, + block_counts_tensor: Any, + cum_seqlen_q_tensor: Any, + cum_seqlen_k_tensor: Any, + ) -> TupleDict: + for value, expected, name in ( + (q_tensor, self.q_desc, "Q"), + (k_tensor, self.k_desc, "K"), + (v_tensor, self.v_desc, "V"), + (block_indices_tensor, self.block_indices_desc, "block_indices"), + (block_counts_tensor, self.block_counts_desc, "block_counts"), + (cum_seqlen_q_tensor, self.cum_q_desc, "cum_seqlen_q"), + (cum_seqlen_k_tensor, self.cum_k_desc, "cum_seqlen_k"), + ): + self.check_tensor_signature(value, expected, name=name) + return _selection_attention_impl( + q_tensor, + k_tensor, + v_tensor, + block_indices_tensor, + block_counts_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + max_s_q=self.max_s_q, + max_s_k=self.max_s_k, + block_size=self.block_size, + scale_softmax=self.scale_softmax, + o_dtype=self.o_dtype, + acc_dtype=self.acc_dtype, + ) + + +def selection_attention_wrapper( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + block_indices_tensor: Any, + block_counts_tensor: Any, + cum_seqlen_q_tensor: Any, + cum_seqlen_k_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, +) -> TupleDict: + """Compute packed THD selection attention with the SM90 CuTe kernel.""" + + return SelectionAttention( + q_tensor, + k_tensor, + v_tensor, + block_indices_tensor, + block_counts_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + 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, + )( + q_tensor, + k_tensor, + v_tensor, + block_indices_tensor, + block_counts_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + + +__all__ = [ + "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..6cc5e09e0 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,10 @@ -from .api import SlidingWindowAttention, sliding_window_attention_wrapper +"""Lazy API exports for NSA sliding-window attention.""" -__all__ = [ - "SlidingWindowAttention", - "sliding_window_attention_wrapper", -] +from ..._operation_api import make_operation_api + +_API_EXPORTS = ("SlidingWindowAttention", "sliding_window_attention_wrapper") + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/native_sparse_attention/sliding_window_attention/api.py b/python/cudnn/native_sparse_attention/sliding_window_attention/api.py index 4cdb5a37d..d610a3451 100644 --- a/python/cudnn/native_sparse_attention/sliding_window_attention/api.py +++ b/python/cudnn/native_sparse_attention/sliding_window_attention/api.py @@ -5,13 +5,13 @@ from cuda.bindings import driver as cuda from cudnn.datatypes import _torch_to_cudnn_data_type -from cudnn.api_base import APIBase, TupleDict +from cudnn.api_base import ApiBaseTorch, TupleDict from typing import Optional from ..utils import make_tensor_strided_like -class SlidingWindowAttention(APIBase): +class SlidingWindowAttention(ApiBaseTorch): def __init__( self, sample_q: torch.Tensor, 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..c3f652c64 --- /dev/null +++ b/python/cudnn/native_sparse_attention/sliding_window_attention/jax.py @@ -0,0 +1,167 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for fixed-shape NSA sliding-window attention inference.""" + +from __future__ import annotations + +from typing import Any + +import jax +import jax.numpy as jnp + +from ..._jax.api_base import ApiBaseJax, TupleDict, require_dtype +from ..jax_utils import require_bhsd_qkv + + +def _sliding_window_attention_impl( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + left_bound: int = 1, + right_bound: int = 0, + is_infer: bool = True, + attn_scale: float | None = None, + o_dtype: Any = None, + *, + _validate_only: bool = False, +) -> TupleDict: + """Compute fixed-shape BHSD sliding-window attention through cuDNN. + + The JAX binding covers the inference subset of the Torch API. Inputs use + logical ``(B, H, S, D)`` shapes and are transposed to JAX's public BTNH + attention contract before lowering to its registered cuDNN custom call. + Packed THD inputs and training statistics are not supported. + """ + + ( + _, + _, + _, + seqlen_q, + seqlen_k, + _, + input_dtype, + ) = require_bhsd_qkv(q_tensor, k_tensor, v_tensor) + output_dtype = require_dtype( + o_dtype, + (jnp.float16, jnp.bfloat16), + name="o_dtype", + default=input_dtype, + ) + + if left_bound < 1: + raise ValueError(f"left_bound must be at least 1, got {left_bound}") + if seqlen_q != seqlen_k: + raise NotImplementedError("JAX sliding-window attention currently requires S_q == S_k") + if right_bound != 0: + raise NotImplementedError("JAX cuDNN sliding-window attention currently requires right_bound=0") + if not is_infer: + raise NotImplementedError("JAX sliding-window attention currently supports inference only") + + scale = None if attn_scale is None else float(attn_scale) + if _validate_only: + return None + + 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)) + output_btnh = jax.nn.dot_product_attention( + q_btnh, + k_btnh, + v_btnh, + scale=scale, + is_causal=True, + local_window_size=(int(left_bound) - 1, 0), + implementation="cudnn", + ) + output = jnp.transpose(output_btnh, (0, 2, 1, 3)).astype(output_dtype) + return TupleDict(o_tensor=output, stats_tensor=None) + + +class SlidingWindowAttention(ApiBaseJax): + """Sample-signature-bound JAX callable for sliding-window inference.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + left_bound: int = 1, + right_bound: int = 0, + is_infer: bool = True, + attn_scale: float | None = None, + o_dtype: Any = None, + ) -> None: + super().__init__() + self.q_desc = self.make_tensor_desc(sample_q, name="sample_q") + self.k_desc = self.make_tensor_desc(sample_k, name="sample_k") + self.v_desc = self.make_tensor_desc(sample_v, name="sample_v") + self.left_bound = left_bound + self.right_bound = right_bound + self.is_infer = is_infer + self.attn_scale = attn_scale + self.o_dtype = self.as_optional_dtype(o_dtype) + + def _check_support(self) -> None: + _sliding_window_attention_impl( + self.q_desc, + self.k_desc, + self.v_desc, + self.left_bound, + self.right_bound, + self.is_infer, + self.attn_scale, + self.o_dtype, + _validate_only=True, + ) + + def __call__(self, q_tensor: Any, k_tensor: Any, v_tensor: Any) -> TupleDict: + return super().__call__(q_tensor, k_tensor, v_tensor) + + def _call_impl( + self, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + ) -> TupleDict: + self.check_tensor_signature(q_tensor, self.q_desc, name="Q") + self.check_tensor_signature(k_tensor, self.k_desc, name="K") + self.check_tensor_signature(v_tensor, self.v_desc, name="V") + return _sliding_window_attention_impl( + q_tensor, + k_tensor, + v_tensor, + self.left_bound, + self.right_bound, + self.is_infer, + self.attn_scale, + self.o_dtype, + ) + + +def sliding_window_attention_wrapper( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + left_bound: int = 1, + right_bound: int = 0, + is_infer: bool = True, + attn_scale: float | None = None, + o_dtype: Any = None, +) -> TupleDict: + """Compute fixed-shape BHSD sliding-window attention inference.""" + + return SlidingWindowAttention( + q_tensor, + k_tensor, + v_tensor, + left_bound=left_bound, + right_bound=right_bound, + is_infer=is_infer, + attn_scale=attn_scale, + o_dtype=o_dtype, + )(q_tensor, k_tensor, v_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..6d0e5f2f6 100644 --- a/python/cudnn/native_sparse_attention/top_k/__init__.py +++ b/python/cudnn/native_sparse_attention/top_k/__init__.py @@ -1,6 +1,10 @@ -from .api import TopKReduction, topk_reduction_wrapper +"""Lazy API exports for NSA top-K reduction.""" -__all__ = [ - "TopKReduction", - "topk_reduction_wrapper", -] +from ..._operation_api import make_operation_api + +_API_EXPORTS = ("TopKReduction", "topk_reduction_wrapper") + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/native_sparse_attention/top_k/api.py b/python/cudnn/native_sparse_attention/top_k/api.py index 3d0a5f639..b2b6f3614 100644 --- a/python/cudnn/native_sparse_attention/top_k/api.py +++ b/python/cudnn/native_sparse_attention/top_k/api.py @@ -10,10 +10,10 @@ from cutlass.cute.runtime import make_fake_stream from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.api_base import APIBase, TupleDict +from cudnn.api_base import ApiBaseTorch, TupleDict -class TopKReduction(APIBase): +class TopKReduction(ApiBaseTorch): """ Top-K Reduction for Native Sparse Attention. 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..249d11626 --- /dev/null +++ b/python/cudnn/native_sparse_attention/top_k/jax.py @@ -0,0 +1,299 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for fixed-shape SM100 NSA top-K reduction.""" + +from __future__ import annotations + +import math +from typing import Any + +import jax.numpy as jnp +from cutlass.jax import TensorSpec, jax_to_cutlass_dtype + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + call_cutedsl, + require_array, + require_dtype, +) +from ..jax_utils import bhsd_storage_spec, require_bhsd_qkv + + +def _launch( + stream, + q, + k, + lse, + topk_scores, + topk_indices, + *, + element_dtype: Any, + batch: int, + seqlen_q: int, + seqlen_k: int, + num_query_heads: int, + num_kv_heads: int, + head_dim: int, + k_value: int, + selection_block_size: int, + compress_stride: int, + is_causal: bool, + scale_softmax: float, +): + from cutlass import Float32, Int32 + + from .nsa_top_k_reduction_fwd import FineGrainedReductionQK + + kernel = FineGrainedReductionQK( + element_dtype=element_dtype, + acc_dtype=Float32, + k_value=k_value, + selection_block_size=selection_block_size, + compress_block_sliding_stride=compress_stride, + mma_tiler=(128, 128, head_dim), + is_causal=is_causal, + ) + problem_size = tuple( + Int32(value) + for value in ( + batch, + seqlen_q, + seqlen_k, + num_query_heads, + num_kv_heads, + head_dim, + ) + ) + + kernel( + problem_size, + q, + k, + lse, + topk_scores, + topk_indices, + Float32(scale_softmax * math.log2(math.e)), + None, + None, + stream, + ) + + +def _require_topk_config( + *, + k_value: int, + selection_block_size: int, + compress_stride: int, + mma_tiler_mn: tuple[int, int], +) -> None: + if mma_tiler_mn != (128, 128): + raise ValueError(f"mma_tiler_mn must be (128, 128), got {mma_tiler_mn}") + if selection_block_size <= 0 or compress_stride <= 0: + raise ValueError("selection_block_size and compress_stride must be positive, got " f"{selection_block_size} and {compress_stride}") + if selection_block_size % compress_stride: + raise ValueError("selection_block_size must be divisible by compress_stride, got " f"{selection_block_size} and {compress_stride}") + reduction_width = selection_block_size // compress_stride + if reduction_width > mma_tiler_mn[1] or mma_tiler_mn[1] % reduction_width: + raise ValueError("selection_block_size / compress_stride must divide the MMA N tile, " f"got {reduction_width} and {mma_tiler_mn[1]}") + candidate_count = mma_tiler_mn[1] // reduction_width + if k_value <= 0 or k_value % 4 or 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 {k_value}") + + +def _topk_reduction_impl( + q_tensor: Any, + k_tensor: Any, + lse_tensor: Any, + 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, + *, + _validate_only: bool = False, +) -> TupleDict: + """Select top compressed-KV blocks for fixed-shape BHSD inputs on SM100. + + ``q_tensor`` and ``k_tensor`` have logical shapes ``(B, H, S, D)`` and a + shared ``float16`` or ``bfloat16`` dtype. ``lse_tensor`` has shape + ``(B, H_q, S_q)`` and dtype ``float32``. Outputs have logical shape + ``(B, H_kv, S_q, k_value)`` and dtypes ``float32`` and ``int32``. + + Variable-length THD inputs are not part of this API. Configuration values + must be static while tracing with :func:`jax.jit`. + """ + + ( + batch, + num_query_heads, + num_kv_heads, + seqlen_q, + seqlen_k, + head_dim, + input_dtype, + ) = require_bhsd_qkv(q_tensor, k_tensor) + require_dtype( + acc_dtype, + (jnp.float32,), + name="acc_dtype", + default=jnp.float32, + ) + require_array( + lse_tensor, + name="lse_tensor", + shape=(batch, num_query_heads, seqlen_q), + dtype=jnp.float32, + ) + _require_topk_config( + k_value=int(k_value), + selection_block_size=int(selection_block_size), + compress_stride=int(compress_stride), + mma_tiler_mn=mma_tiler_mn, + ) + + resolved_scale = 1.0 / math.sqrt(head_dim) if scale_softmax is None else float(scale_softmax) + if _validate_only: + return None + + bhsd_spec = bhsd_storage_spec(present_as_bshd=False) + lse_spec = TensorSpec(layout=(2, 1, 0), mode=(0, 1, 2)) + output_shape = (batch, num_kv_heads, seqlen_q, int(k_value)) + topk_scores, topk_indices = call_cutedsl( + _launch, + (q_tensor, k_tensor, lse_tensor), + outputs=( + BufferSpec( + "topk_scores_tensor", + output_shape, + jnp.float32, + tensor_spec=bhsd_spec, + fill_value=float("-inf"), + ), + BufferSpec( + "topk_indices_tensor", + output_shape, + jnp.int32, + tensor_spec=bhsd_spec, + fill_value=-1, + ), + ), + input_specs=(bhsd_spec, bhsd_spec, lse_spec), + static_args={ + "element_dtype": jax_to_cutlass_dtype(input_dtype), + "batch": batch, + "seqlen_q": seqlen_q, + "seqlen_k": seqlen_k, + "num_query_heads": num_query_heads, + "num_kv_heads": num_kv_heads, + "head_dim": head_dim, + "k_value": int(k_value), + "selection_block_size": int(selection_block_size), + "compress_stride": int(compress_stride), + "is_causal": bool(is_causal), + "scale_softmax": resolved_scale, + }, + ) + return TupleDict( + topk_scores_tensor=topk_scores, + topk_indices_tensor=topk_indices, + ) + + +class TopKReduction(ApiBaseJax): + """Sample-signature-bound JAX callable for SM100 NSA top-K reduction.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_lse: Any, + 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, + ) -> None: + super().__init__() + self.q_desc = self.make_tensor_desc(sample_q, name="sample_q") + self.k_desc = self.make_tensor_desc(sample_k, name="sample_k") + self.lse_desc = self.make_tensor_desc(sample_lse, name="sample_lse") + self.acc_dtype = self.as_optional_dtype(acc_dtype) + self.k_value = k_value + self.selection_block_size = selection_block_size + self.compress_stride = compress_stride + self.is_causal = is_causal + self.mma_tiler_mn = tuple(mma_tiler_mn) + self.scale_softmax = scale_softmax + + def _check_support(self) -> None: + _topk_reduction_impl( + self.q_desc, + self.k_desc, + self.lse_desc, + self.acc_dtype, + self.k_value, + self.selection_block_size, + self.compress_stride, + self.is_causal, + self.mma_tiler_mn, + self.scale_softmax, + _validate_only=True, + ) + + def __call__(self, q_tensor: Any, k_tensor: Any, lse_tensor: Any) -> TupleDict: + return super().__call__(q_tensor, k_tensor, lse_tensor) + + def _call_impl(self, q_tensor: Any, k_tensor: Any, lse_tensor: Any) -> TupleDict: + self.check_tensor_signature(q_tensor, self.q_desc, name="Q") + self.check_tensor_signature(k_tensor, self.k_desc, name="K") + self.check_tensor_signature(lse_tensor, self.lse_desc, name="LSE") + return _topk_reduction_impl( + q_tensor, + k_tensor, + lse_tensor, + self.acc_dtype, + self.k_value, + self.selection_block_size, + self.compress_stride, + self.is_causal, + self.mma_tiler_mn, + self.scale_softmax, + ) + + +def topk_reduction_wrapper( + q_tensor: Any, + k_tensor: Any, + lse_tensor: Any, + 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, +) -> TupleDict: + """Select top compressed-KV blocks for fixed-shape BHSD inputs on SM100.""" + + return TopKReduction( + q_tensor, + k_tensor, + lse_tensor, + 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, + )(q_tensor, k_tensor, lse_tensor) + + +__all__ = ["TopKReduction", "topk_reduction_wrapper"] diff --git a/python/cudnn/rmsnorm_rht_amax/__init__.py b/python/cudnn/rmsnorm_rht_amax/__init__.py index 9da5cae15..b5734dbc5 100644 --- a/python/cudnn/rmsnorm_rht_amax/__init__.py +++ b/python/cudnn/rmsnorm_rht_amax/__init__.py @@ -1,16 +1,24 @@ # 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 API exports for RMSNorm + RHT + amax. + +Unqualified symbols always come from the Torch ``api.py``. The JAX API is +available explicitly through the sibling ``jax`` namespace. +""" + +from .._operation_api import make_operation_api + -__all__ = [ +_API_EXPORTS = ( "RmsNormRhtAmaxSm100", "best_num_threads", "pick_rows_per_cta", + "rmsnorm_rht_amax_sm100", "rmsnorm_rht_amax_wrapper_sm100", -] +) + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/rmsnorm_rht_amax/api.py b/python/cudnn/rmsnorm_rht_amax/api.py index 0294a3a22..b67fb821e 100644 --- a/python/cudnn/rmsnorm_rht_amax/api.py +++ b/python/cudnn/rmsnorm_rht_amax/api.py @@ -1,55 +1,28 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -"""FE API for fused RMSNorm + RHT + per-CTA amax.""" +"""PyTorch FE API for fused RMSNorm + RHT + per-CTA amax.""" import logging 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 cudnn.api_base import ApiBaseTorch, TensorDesc, TupleDict +from .config import ( + best_num_threads, + pick_rows_per_cta, + validate_rmsnorm_rht_amax, +) from .kernel import RMSNormRHTAmaxKernel -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 - num_ctas = m // rows_per_cta - if num_ctas >= TARGET_MIN_CTAS: - return rows_per_cta - return RPC_CANDIDATES[0] - - -class RmsNormRhtAmaxSm100(APIBase): + +class RmsNormRhtAmaxSm100(ApiBaseTorch): """Class API for the RMSNorm + RHT + amax kernel.""" def __init__( @@ -79,52 +52,28 @@ def __init__( self.n = None 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}", + plan = validate_rmsnorm_rht_amax( + self.x_desc, + self.w_desc, + output=self.o_desc, + amax=self.amax_desc, + num_threads=self.requested_num_threads, + rows_per_cta=self.requested_rows_per_cta, ) - 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) - - 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}") - - 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") - - expected_num_ctas = m // resolved_rows_per_cta - self._check_tensor_shape(self.amax_desc, (expected_num_ctas,), "Amax") + self._check_tensor_stride( + self.x_desc, + stride=(plan.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=(plan.n, 1), + name="O", + extra_error_msg="O must be row-major contiguous", + ) 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 +83,9 @@ 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.num_threads = plan.num_threads + self.rows_per_cta = plan.rows_per_cta + self.n = plan.n self._is_supported = True return True @@ -215,7 +164,10 @@ def execute( amax_tensor: torch.Tensor, current_stream: Optional[cuda.CUstream] = None, ) -> None: - self._runtime_error_if(self._compiled_kernel is None, "RmsNormRhtAmaxSm100 kernel not compiled; call compile() first") + 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") @@ -251,19 +203,18 @@ 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}") + plan = validate_rmsnorm_rht_amax( + TensorDesc(dtype=x_tensor.dtype, shape=tuple(x_tensor.shape), name="X"), + TensorDesc(dtype=w_tensor.dtype, shape=tuple(w_tensor.shape), name="W"), + num_threads=num_threads, + rows_per_cta=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) + amax_tensor = torch.full(plan.amax_shape, float("-inf"), dtype=torch.float32, device=x_tensor.device) cache_key = ( - n, + plan.n, x_tensor.dtype, w_tensor.dtype, o_tensor.dtype, @@ -271,8 +222,8 @@ def rmsnorm_rht_amax_wrapper_sm100( tuple(w_tensor.stride()), tuple(o_tensor.stride()), eps, - resolved_num_threads, - resolved_rows_per_cta, + plan.num_threads, + plan.rows_per_cta, ) if cache_key in _cache_of_RmsNormRhtAmaxSm100Objects: @@ -284,8 +235,8 @@ def rmsnorm_rht_amax_wrapper_sm100( sample_o=o_tensor, sample_amax=amax_tensor, eps=eps, - num_threads=resolved_num_threads, - rows_per_cta=resolved_rows_per_cta, + num_threads=plan.num_threads, + rows_per_cta=plan.rows_per_cta, ) assert api.check_support(), "Unsupported configuration" api.compile() @@ -300,3 +251,39 @@ def rmsnorm_rht_amax_wrapper_sm100( ) return TupleDict(o_tensor=o_tensor, amax_tensor=amax_tensor) + + +def rmsnorm_rht_amax_sm100( + x: torch.Tensor, + weight: torch.Tensor, + *, + eps: float = 1e-5, + num_threads: Optional[int] = None, + rows_per_cta: Optional[int] = None, + current_stream: Optional[cuda.CUstream] = None, +) -> TupleDict: + """Apply RMSNorm + RHT + amax through the aligned Torch API. + + This additive facade shares its operation name, semantic operands, options, + and result roles with ``cudnn.jax.rmsnorm_rht_amax_sm100``. The existing + ``rmsnorm_rht_amax_wrapper_sm100`` API remains unchanged for compatibility. + """ + + result = rmsnorm_rht_amax_wrapper_sm100( + x_tensor=x, + w_tensor=weight, + eps=eps, + num_threads=num_threads, + rows_per_cta=rows_per_cta, + current_stream=current_stream, + ) + return TupleDict(output=result["o_tensor"], amax=result["amax_tensor"]) + + +__all__ = [ + "RmsNormRhtAmaxSm100", + "best_num_threads", + "pick_rows_per_cta", + "rmsnorm_rht_amax_sm100", + "rmsnorm_rht_amax_wrapper_sm100", +] diff --git a/python/cudnn/rmsnorm_rht_amax/config.py b/python/cudnn/rmsnorm_rht_amax/config.py new file mode 100644 index 000000000..c49cbe5d7 --- /dev/null +++ b/python/cudnn/rmsnorm_rht_amax/config.py @@ -0,0 +1,176 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral validation for RMSNorm + RHT + amax.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional, Tuple + +if TYPE_CHECKING: + from cudnn.api_base 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 + + +@dataclass(frozen=True) +class RmsNormRhtAmaxPlan: + """Validated, framework-independent operation metadata.""" + + m: int + n: int + num_threads: int + rows_per_cta: int + output_shape: tuple[int, int] + amax_shape: tuple[int] + + +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 + num_ctas = m // rows_per_cta + if num_ctas >= TARGET_MIN_CTAS: + return rows_per_cta + return RPC_CANDIDATES[0] + + +def resolve_launch_config( + m: int, + n: int, + *, + num_threads: Optional[int] = None, + rows_per_cta: Optional[int] = None, +) -> Tuple[int, int]: + """Validate dimensions and resolve target-independent launch parameters.""" + + 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 n % 16 != 0: + raise ValueError(f"N must be divisible by 16 for the Hadamard block size, got {n}") + + resolved_num_threads = num_threads + if resolved_num_threads is None: + resolved_num_threads = 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}") + if resolved_num_threads <= 0: + raise ValueError(f"num_threads must be positive, got {resolved_num_threads}") + if resolved_num_threads % 32 != 0: + raise ValueError(f"num_threads must be warp-aligned, got {resolved_num_threads}") + if resolved_num_threads > 1024: + raise ValueError("num_threads must not exceed the CUDA block size limit, " f"got {resolved_num_threads}") + + resolved_rows_per_cta = rows_per_cta + if resolved_rows_per_cta is None: + resolved_rows_per_cta = 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("M must be divisible by rows_per_cta, " f"got M={m}, rows_per_cta={resolved_rows_per_cta}") + if n % resolved_num_threads != 0: + raise ValueError(f"N={n} must be divisible by num_threads={resolved_num_threads}") + + ept = n // resolved_num_threads + if ept < 8 or ept % 8 != 0: + raise ValueError(f"EPT={ept} must be >= 8 and divisible by 8") + + return resolved_num_threads, resolved_rows_per_cta + + +def _require_rank(tensor: TensorDesc, rank: int, name: str) -> None: + if tensor.ndim != rank: + raise ValueError(f"{name} must have rank {rank}, got shape {tensor.shape}") + + +def _require_shape(tensor: TensorDesc, shape: tuple[int, ...], name: str) -> None: + if tensor.shape != shape: + raise ValueError(f"{name} must have shape {shape}, got {tensor.shape}") + + +def _require_dtype(tensor: TensorDesc, dtype_name: str, name: str) -> None: + if tensor.dtype_name != dtype_name: + raise ValueError(f"{name} must have dtype {dtype_name}, got {tensor.dtype_name}") + + +def validate_rmsnorm_rht_amax( + x: TensorDesc, + weight: TensorDesc, + *, + output: Optional[TensorDesc] = None, + amax: Optional[TensorDesc] = None, + num_threads: Optional[int] = None, + rows_per_cta: Optional[int] = None, +) -> RmsNormRhtAmaxPlan: + """Validate logical tensor metadata and infer the operation outputs. + + Physical layout and device capability remain adapter responsibilities: + Torch validates observed strides and its CUDA device, while JAX declares a + compact custom-call layout with ``cutlass.jax.TensorSpec``. + """ + + _require_rank(x, 2, "X") + _require_rank(weight, 1, "W") + _require_dtype(x, "bfloat16", "X") + _require_dtype(weight, "bfloat16", "W") + + m, n = x.shape + _require_shape(weight, (n,), "W") + resolved_num_threads, resolved_rows_per_cta = resolve_launch_config( + m, + n, + num_threads=num_threads, + rows_per_cta=rows_per_cta, + ) + + output_shape = (m, n) + amax_shape = (m // resolved_rows_per_cta,) + if output is not None: + _require_shape(output, output_shape, "O") + _require_dtype(output, "bfloat16", "O") + if amax is not None: + _require_shape(amax, amax_shape, "Amax") + _require_dtype(amax, "float32", "Amax") + + return RmsNormRhtAmaxPlan( + m=m, + n=n, + num_threads=resolved_num_threads, + rows_per_cta=resolved_rows_per_cta, + output_shape=output_shape, + amax_shape=amax_shape, + ) + + +__all__ = [ + "DEFAULT_NUM_THREADS_BY_N", + "RPC_CANDIDATES", + "RmsNormRhtAmaxPlan", + "TARGET_MIN_CTAS", + "best_num_threads", + "pick_rows_per_cta", + "resolve_launch_config", + "validate_rmsnorm_rht_amax", +] diff --git a/python/cudnn/rmsnorm_rht_amax/jax.py b/python/cudnn/rmsnorm_rht_amax/jax.py new file mode 100644 index 000000000..4bb8e656a --- /dev/null +++ b/python/cudnn/rmsnorm_rht_amax/jax.py @@ -0,0 +1,163 @@ +# 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 + +import jax.numpy as jnp +from cutlass.jax import TensorSpec + +from .._jax.api_base import ApiBaseJax, BufferSpec, TupleDict, call_cutedsl +from .config import RmsNormRhtAmaxPlan, validate_rmsnorm_rht_amax + + +def _launch( + stream, + x, + weight, + output, + amax, + *, + n: int, + num_threads: int, + rows_per_cta: int, + eps: float, +): + # Load the configuration-specific kernel only when tracing the operation. + from cutlass import Float32 + from .kernel import RMSNormRHTAmaxKernel + + kernel = RMSNormRHTAmaxKernel( + n=n, + num_threads=num_threads, + eps=eps, + rows_per_cta=rows_per_cta, + ) + kernel(x, weight, output, amax, Float32(eps), stream) + + +class RmsNormRhtAmaxSm100(ApiBaseJax): + """JAX callable specialized from sample shape and dtype metadata. + + Sample values are converted to descriptors during construction and are not + retained. Actual arrays are passed to :meth:`__call__`, so this object can + be used directly with ``jax.jit`` without capturing array constants. + ``check_support()`` validates the abstract signature and static kernel + configuration; final device capability is determined during lowering. + """ + + 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: + super().__init__() + self.x_desc = self.make_tensor_desc(sample_x, name="sample_x") + self.w_desc = self.make_tensor_desc(sample_w, name="sample_w") + self.eps = eps + self.requested_num_threads = num_threads + self.requested_rows_per_cta = rows_per_cta + self._plan: Optional[RmsNormRhtAmaxPlan] = None + + @property + def num_threads(self) -> Optional[int]: + return None if self._plan is None else self._plan.num_threads + + @property + def rows_per_cta(self) -> Optional[int]: + return None if self._plan is None else self._plan.rows_per_cta + + @property + def n(self) -> Optional[int]: + return None if self._plan is None else self._plan.n + + def _check_support(self) -> None: + self._plan = validate_rmsnorm_rht_amax( + self.x_desc, + self.w_desc, + num_threads=self.requested_num_threads, + rows_per_cta=self.requested_rows_per_cta, + ) + + def __call__(self, x: Any, weight: Any) -> TupleDict: + """Run with arrays matching the validated sample signature.""" + + return super().__call__(x, weight) + + def _call_impl(self, x: Any, weight: Any) -> TupleDict: + plan = self._plan + if plan is None: + raise RuntimeError("check_support() did not produce a launch plan") + + self.check_tensor_signature(x, self.x_desc, name="X") + self.check_tensor_signature(weight, self.w_desc, name="W") + + # JAX/XLA owns physical buffers. These specs constrain ordinary + # row-major storage while providing the divisibility facts used by the + # CuTe kernel. + x_spec = TensorSpec(divisibility=(plan.rows_per_cta, 16)) + weight_spec = TensorSpec(divisibility=(16,)) + + output, amax = call_cutedsl( + _launch, + (x, weight), + outputs=( + BufferSpec( + "output", + plan.output_shape, + jnp.bfloat16, + tensor_spec=x_spec, + ), + BufferSpec( + "amax", + plan.amax_shape, + jnp.float32, + ), + ), + input_specs=(x_spec, weight_spec), + static_args={ + "n": plan.n, + "num_threads": plan.num_threads, + "rows_per_cta": plan.rows_per_cta, + "eps": self.eps, + }, + ) + return TupleDict(output=output, amax=amax) + + +def rmsnorm_rht_amax_sm100( + x: Any, + weight: Any, + *, + eps: float = 1e-5, + num_threads: Optional[int] = None, + rows_per_cta: Optional[int] = None, +) -> TupleDict: + """Apply fused RMSNorm, 16-wide RHT, and per-CTA amax from JAX. + + ``x`` must be a row-major ``bfloat16`` array of shape ``(M, N)`` and + ``weight`` a ``bfloat16`` array of shape ``(N,)``. ``M`` and ``N`` must be + concrete; shape-polymorphic export is unsupported. The underlying callable + is not wrapped in ``jax.jit``; callers own JIT and sharding policy. + + ``eps``, ``num_threads``, and ``rows_per_cta`` are compile-time + configuration. Returns ``(output, amax)`` with shapes ``(M, N)`` and + ``(M // rows_per_cta,)`` respectively. + """ + + 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/sdpa/__init__.py b/python/cudnn/sdpa/__init__.py index b200dcbda..8ee804a77 100644 --- a/python/cudnn/sdpa/__init__.py +++ b/python/cudnn/sdpa/__init__.py @@ -1,12 +1,27 @@ # 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 - -__all__ = [ - "SdpafwdSm100D256", - "sdpa_fwd_wrapper_sm100_d256", - "SdpabwdSm100D256", - "sdpa_bwd_wrapper_sm100_d256", -] +"""Lazy exports for SDPA operation packages.""" + +from importlib import import_module + +_SYMBOLS = { + "SdpafwdSm100D256": (".fwd", "SdpafwdSm100D256"), + "sdpa_fwd_wrapper_sm100_d256": (".fwd", "sdpa_fwd_wrapper_sm100_d256"), + "SdpabwdSm100D256": (".bwd", "SdpabwdSm100D256"), + "sdpa_bwd_wrapper_sm100_d256": (".bwd", "sdpa_bwd_wrapper_sm100_d256"), +} + + +def __getattr__(name): + if name in _SYMBOLS: + module_name, symbol_name = _SYMBOLS[name] + value = getattr(import_module(module_name, __name__), symbol_name) + else: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + globals()[name] = value + return value + + +__all__ = list(_SYMBOLS) diff --git a/python/cudnn/sdpa/bwd/__init__.py b/python/cudnn/sdpa/bwd/__init__.py index 0d0c368ed..877b4bc23 100644 --- a/python/cudnn/sdpa/bwd/__init__.py +++ b/python/cudnn/sdpa/bwd/__init__.py @@ -1,9 +1,16 @@ -# 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 SdpabwdSm100D256, sdpa_bwd_wrapper_sm100_d256 +"""Lazy API exports for d=256 SDPA backward.""" -__all__ = [ +from ..._operation_api import make_operation_api + +_API_EXPORTS = ( "SdpabwdSm100D256", "sdpa_bwd_wrapper_sm100_d256", -] +) + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/sdpa/bwd/api.py b/python/cudnn/sdpa/bwd/api.py index bd99fe3e9..81a850f57 100644 --- a/python/cudnn/sdpa/bwd/api.py +++ b/python/cudnn/sdpa/bwd/api.py @@ -11,14 +11,14 @@ import cutlass.cute as cute from cutlass.cute.runtime import make_fake_stream -from cudnn.api_base import APIBase, TupleDict +from cudnn.api_base import ApiBaseTorch, TupleDict from cudnn.datatypes import _convert_to_cutlass_data_type from .fmha_backward_sm100_2kernel import BlackwellFusedMultiHeadAttentionBackward from ..fmha_utils import MaskEnum -class SdpabwdSm100D256(APIBase): +class SdpabwdSm100D256(ApiBaseTorch): """API class for d=256 SDPA backward (SM100+) using 2-kernel implementation. Input/output layout follows the kernel contract: diff --git a/python/cudnn/sdpa/bwd/jax.py b/python/cudnn/sdpa/bwd/jax.py new file mode 100644 index 000000000..35a1d13a6 --- /dev/null +++ b/python/cudnn/sdpa/bwd/jax.py @@ -0,0 +1,369 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for fixed-shape SM100 d=256 SDPA backward.""" + +from __future__ import annotations + +from typing import Any + +import jax.numpy as jnp +from cutlass.jax import jax_to_cutlass_dtype + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + call_cutedsl, + require_array, + require_dtype, +) +from ..jax_utils import ( + bhsd_tensor_spec, + require_bhsd_qkv, + resolve_sdpa_config, +) + + +def _launch( + stream, + q, + k, + v, + output, + doutput, + lse, + dq, + dk, + dv, + workspace, + *, + batch: int, + seqlen_q: int, + seqlen_k: int, + num_query_heads: int, + num_kv_heads: int, + element_dtype: Any, + scale_softmax: float, + is_causal: bool, + window_size_left: int, + window_size_right: int, + mask_kind: str, +): + from cutlass import Float32, Int32 + + from ..fmha_utils import MaskEnum + from .fmha_backward_sm100_2kernel import BlackwellFusedMultiHeadAttentionBackward + + mask_type = { + "residual": MaskEnum.RESIDUAL_MASK, + "window": MaskEnum.WINDOW_MASK_INFERENCE, + }[mask_kind] + kernel = BlackwellFusedMultiHeadAttentionBackward( + element_dtype=element_dtype, + acc_dtype=Float32, + mma_tiler=(128, 128, 256), + dkdv_mma_tiler=(128, 64, 256), + varlen=False, + is_causal=is_causal, + mask_type=mask_type, + window_size_left=window_size_left, + window_size_right=window_size_right, + ) + + problem_shape = ( + Int32(seqlen_q), + Int32(seqlen_k), + Int32(256), + ( + (Int32(num_query_heads // num_kv_heads), Int32(num_kv_heads)), + Int32(batch), + ), + ) + kernel( + problem_shape, + q, + k, + v, + output, + dq, + dk, + dv, + doutput, + lse, + None, + None, + Float32(scale_softmax), + workspace, + stream, + ) + + +def _sdpa_bwd_impl( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + o_tensor: Any, + do_tensor: Any, + lse_tensor: Any, + 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, + *, + _validate_only: bool = False, +) -> TupleDict: + """Compute fixed-shape BHSD SDPA gradients with the SM100 d=256 kernel. + + Q, K, V, O, and dO use logical JAX shapes ``(B, H, S, 256)``. + ``lse_tensor`` uses shape ``(B, H_q, S_q)`` and dtype ``float32``. + The returned dQ, dK, and dV arrays match Q, K, and V respectively. + + The multi-kernel implementation uses a hidden, zero-initialized workspace + owned by XLA. Variable-length THD inputs are not part of this API. + """ + + ( + batch, + num_query_heads, + num_kv_heads, + seqlen_q, + seqlen_k, + head_dim, + dtype, + ) = require_bhsd_qkv(q_tensor, k_tensor, v_tensor) + require_array( + o_tensor, + name="o_tensor", + shape=tuple(q_tensor.shape), + dtype=dtype, + ) + require_array( + do_tensor, + name="do_tensor", + shape=tuple(q_tensor.shape), + dtype=dtype, + ) + require_array( + lse_tensor, + name="lse_tensor", + shape=(batch, num_query_heads, seqlen_q), + dtype=jnp.float32, + ) + + require_dtype( + acc_dtype, + (jnp.float32,), + name="acc_dtype", + default=jnp.float32, + ) + if mma_tiler_mn != (128, 128): + raise ValueError(f"mma_tiler_mn must be (128, 128), got {mma_tiler_mn}") + if dkdv_mma_tiler_mn != (128, 64): + raise ValueError(f"dkdv_mma_tiler_mn must be (128, 64), got {dkdv_mma_tiler_mn}") + + scale_softmax, window_size_left, window_size_right, mask_kind = resolve_sdpa_config( + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + tile_extent=seqlen_q, + is_causal=bool(is_causal), + window_size=window_size, + scale_softmax=scale_softmax, + ) + if _validate_only: + return None + + from cutlass import Float32 + + from .fmha_backward_sm100_2kernel import BlackwellFusedMultiHeadAttentionBackward + + workspace_shape = BlackwellFusedMultiHeadAttentionBackward.get_workspace_size( + seqlen_q, + head_dim, + num_query_heads, + batch, + Float32, + ) + bhsd_spec = bhsd_tensor_spec() + dq_tensor, dk_tensor, dv_tensor = call_cutedsl( + _launch, + (q_tensor, k_tensor, v_tensor, o_tensor, do_tensor, lse_tensor), + outputs=( + BufferSpec( + "dq_tensor", + tuple(q_tensor.shape), + dtype, + tensor_spec=bhsd_spec, + ), + BufferSpec( + "dk_tensor", + tuple(k_tensor.shape), + dtype, + tensor_spec=bhsd_spec, + ), + BufferSpec( + "dv_tensor", + tuple(v_tensor.shape), + dtype, + tensor_spec=bhsd_spec, + ), + ), + workspaces=( + BufferSpec( + "workspace", + workspace_shape, + jnp.uint8, + fill_value=0, + ), + ), + input_specs=(bhsd_spec, bhsd_spec, bhsd_spec, bhsd_spec, bhsd_spec, None), + static_args={ + "batch": batch, + "seqlen_q": seqlen_q, + "seqlen_k": seqlen_k, + "num_query_heads": num_query_heads, + "num_kv_heads": num_kv_heads, + "element_dtype": jax_to_cutlass_dtype(dtype), + "scale_softmax": scale_softmax, + "is_causal": bool(is_causal), + "window_size_left": window_size_left, + "window_size_right": window_size_right, + "mask_kind": mask_kind, + }, + ) + return TupleDict( + dq_tensor=dq_tensor, + dk_tensor=dk_tensor, + dv_tensor=dv_tensor, + ) + + +class SdpabwdSm100D256(ApiBaseJax): + """Sample-signature-bound JAX callable for SM100 d=256 SDPA backward.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_o: Any, + sample_do: Any, + sample_lse: Any, + 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, + ) -> None: + super().__init__() + self.q_desc = self.make_tensor_desc(sample_q, name="sample_q") + self.k_desc = self.make_tensor_desc(sample_k, name="sample_k") + self.v_desc = self.make_tensor_desc(sample_v, name="sample_v") + self.o_desc = self.make_tensor_desc(sample_o, name="sample_o") + self.do_desc = self.make_tensor_desc(sample_do, name="sample_do") + self.lse_desc = self.make_tensor_desc(sample_lse, name="sample_lse") + self.acc_dtype = self.as_optional_dtype(acc_dtype) + self.mma_tiler_mn = tuple(mma_tiler_mn) + self.dkdv_mma_tiler_mn = tuple(dkdv_mma_tiler_mn) + self.is_causal = is_causal + self.window_size = tuple(window_size) + self.scale_softmax = scale_softmax + + def _check_support(self) -> None: + _sdpa_bwd_impl( + self.q_desc, + self.k_desc, + self.v_desc, + self.o_desc, + self.do_desc, + self.lse_desc, + self.acc_dtype, + self.mma_tiler_mn, + self.dkdv_mma_tiler_mn, + self.is_causal, + self.window_size, + self.scale_softmax, + _validate_only=True, + ) + + def __call__( + self, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + o_tensor: Any, + do_tensor: Any, + lse_tensor: Any, + ) -> TupleDict: + return super().__call__(q_tensor, k_tensor, v_tensor, o_tensor, do_tensor, lse_tensor) + + def _call_impl( + self, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + o_tensor: Any, + do_tensor: Any, + lse_tensor: Any, + ) -> TupleDict: + for value, expected, name in ( + (q_tensor, self.q_desc, "Q"), + (k_tensor, self.k_desc, "K"), + (v_tensor, self.v_desc, "V"), + (o_tensor, self.o_desc, "O"), + (do_tensor, self.do_desc, "dO"), + (lse_tensor, self.lse_desc, "LSE"), + ): + self.check_tensor_signature(value, expected, name=name) + return _sdpa_bwd_impl( + q_tensor, + k_tensor, + v_tensor, + o_tensor, + do_tensor, + lse_tensor, + self.acc_dtype, + self.mma_tiler_mn, + self.dkdv_mma_tiler_mn, + self.is_causal, + self.window_size, + self.scale_softmax, + ) + + +def sdpa_bwd_wrapper_sm100_d256( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + o_tensor: Any, + do_tensor: Any, + lse_tensor: Any, + 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, +) -> TupleDict: + """Compute fixed-shape BHSD SDPA gradients with the SM100 d=256 kernel.""" + + return SdpabwdSm100D256( + q_tensor, + k_tensor, + v_tensor, + o_tensor, + do_tensor, + lse_tensor, + 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, + )(q_tensor, k_tensor, v_tensor, o_tensor, do_tensor, lse_tensor) + + +__all__ = ["SdpabwdSm100D256", "sdpa_bwd_wrapper_sm100_d256"] diff --git a/python/cudnn/sdpa/fwd/__init__.py b/python/cudnn/sdpa/fwd/__init__.py index bf95abac2..63691957f 100644 --- a/python/cudnn/sdpa/fwd/__init__.py +++ b/python/cudnn/sdpa/fwd/__init__.py @@ -1,9 +1,16 @@ -# 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 SdpafwdSm100D256, sdpa_fwd_wrapper_sm100_d256 +"""Lazy API exports for d=256 SDPA forward.""" -__all__ = [ +from ..._operation_api import make_operation_api + +_API_EXPORTS = ( "SdpafwdSm100D256", "sdpa_fwd_wrapper_sm100_d256", -] +) + +__all__, __getattr__ = make_operation_api( + globals(), + exports=_API_EXPORTS, +) diff --git a/python/cudnn/sdpa/fwd/api.py b/python/cudnn/sdpa/fwd/api.py index 90f4c76d3..d3e51adbd 100644 --- a/python/cudnn/sdpa/fwd/api.py +++ b/python/cudnn/sdpa/fwd/api.py @@ -12,14 +12,14 @@ import cutlass.cute as cute from cutlass.cute.runtime import make_fake_stream -from cudnn.api_base import APIBase, TupleDict +from cudnn.api_base import ApiBaseTorch, TupleDict from cudnn.datatypes import _convert_to_cutlass_data_type from .fmha_forward_sm100_d256 import BlackwellFusedMultiHeadAttentionForward from ..fmha_utils import MaskEnum -class SdpafwdSm100D256(APIBase): +class SdpafwdSm100D256(ApiBaseTorch): """API class for d=256 SDPA forward (SM100+) using the FE OSS CUTE DSL kernel.""" def __init__( diff --git a/python/cudnn/sdpa/fwd/jax.py b/python/cudnn/sdpa/fwd/jax.py new file mode 100644 index 000000000..4074c8bef --- /dev/null +++ b/python/cudnn/sdpa/fwd/jax.py @@ -0,0 +1,279 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for fixed-shape SM100 d=256 SDPA forward.""" + +from __future__ import annotations + +import math +from typing import Any + +import jax.numpy as jnp + +from ..._jax.api_base import ( + ApiBaseJax, + BufferSpec, + TupleDict, + call_cutedsl, + require_dtype, +) +from ..jax_utils import bhsd_tensor_spec, require_bhsd_qkv, resolve_sdpa_config + + +def _launch( + stream, + q, + k, + v, + output, + lse, + *, + batch: int, + seqlen_q: int, + seqlen_k: int, + num_query_heads: int, + num_kv_heads: int, + scale_softmax: float, + scale_output: float, + window_size_left: int, + window_size_right: int, + mask_kind: str, +): + 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, + }[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 ( + batch, + seqlen_q, + seqlen_k, + num_query_heads, + num_kv_heads, + 256, + ) + ) + left = None if window_size_left < 0 else Int32(window_size_left) + right = None if window_size_right < 0 else Int32(window_size_right) + kernel( + q, + k, + v, + output, + problem_size, + None, + None, + lse, + Float32(scale_softmax * math.log2(math.e)), + Float32(scale_softmax), + Float32(scale_output), + left, + right, + stream, + ) + + +def _sdpa_fwd_impl( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + 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, + *, + _validate_only: bool = False, +) -> TupleDict: + """Compute fixed-shape BHSD SDPA forward with the SM100 d=256 kernel. + + ``q_tensor``, ``k_tensor``, and ``v_tensor`` use logical JAX shapes + ``(B, H, S, 256)`` and must share a ``float16`` or ``bfloat16`` dtype. + The result contains ``o_tensor`` with the same shape and dtype as Q and + ``lse_tensor`` with shape ``(B, H_q, S_q)`` and dtype ``float32``. + + Variable-length THD inputs are not part of this API. Configuration values + must be static while tracing with :func:`jax.jit`. + """ + + ( + batch, + num_query_heads, + num_kv_heads, + seqlen_q, + seqlen_k, + _head_dim, + dtype, + ) = require_bhsd_qkv(q_tensor, k_tensor, v_tensor) + + require_dtype( + qk_acc_dtype, + (jnp.float32,), + name="qk_acc_dtype", + default=jnp.float32, + ) + require_dtype( + pv_acc_dtype, + (jnp.float32,), + name="pv_acc_dtype", + default=jnp.float32, + ) + if mma_tiler_mn != (128, 128): + raise ValueError(f"mma_tiler_mn must be (128, 128), got {mma_tiler_mn}") + + scale_softmax, window_size_left, window_size_right, mask_kind = resolve_sdpa_config( + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + tile_extent=seqlen_k, + is_causal=bool(is_causal), + window_size=window_size, + scale_softmax=scale_softmax, + ) + scale_output = float(scale_output) + if _validate_only: + return None + + bhsd_spec = bhsd_tensor_spec() + + o_tensor, lse_tensor = call_cutedsl( + _launch, + (q_tensor, k_tensor, v_tensor), + outputs=( + BufferSpec( + "o_tensor", + tuple(q_tensor.shape), + dtype, + tensor_spec=bhsd_spec, + ), + BufferSpec( + "lse_tensor", + (batch, num_query_heads, seqlen_q), + jnp.float32, + ), + ), + input_specs=(bhsd_spec, bhsd_spec, bhsd_spec), + static_args={ + "batch": batch, + "seqlen_q": seqlen_q, + "seqlen_k": seqlen_k, + "num_query_heads": num_query_heads, + "num_kv_heads": num_kv_heads, + "scale_softmax": scale_softmax, + "scale_output": scale_output, + "window_size_left": window_size_left, + "window_size_right": window_size_right, + "mask_kind": mask_kind, + }, + ) + return TupleDict(o_tensor=o_tensor, lse_tensor=lse_tensor) + + +class SdpafwdSm100D256(ApiBaseJax): + """Sample-signature-bound JAX callable for SM100 d=256 SDPA forward.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + 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, + ) -> None: + super().__init__() + self.q_desc = self.make_tensor_desc(sample_q, name="sample_q") + self.k_desc = self.make_tensor_desc(sample_k, name="sample_k") + self.v_desc = self.make_tensor_desc(sample_v, name="sample_v") + self.qk_acc_dtype = self.as_optional_dtype(qk_acc_dtype) + self.pv_acc_dtype = self.as_optional_dtype(pv_acc_dtype) + self.mma_tiler_mn = tuple(mma_tiler_mn) + self.is_causal = is_causal + self.window_size = tuple(window_size) + self.scale_softmax = scale_softmax + self.scale_output = scale_output + + def _check_support(self) -> None: + _sdpa_fwd_impl( + self.q_desc, + self.k_desc, + self.v_desc, + self.qk_acc_dtype, + self.pv_acc_dtype, + self.mma_tiler_mn, + self.is_causal, + self.window_size, + self.scale_softmax, + self.scale_output, + _validate_only=True, + ) + + def __call__(self, q_tensor: Any, k_tensor: Any, v_tensor: Any) -> TupleDict: + return super().__call__(q_tensor, k_tensor, v_tensor) + + def _call_impl(self, q_tensor: Any, k_tensor: Any, v_tensor: Any) -> TupleDict: + self.check_tensor_signature(q_tensor, self.q_desc, name="Q") + self.check_tensor_signature(k_tensor, self.k_desc, name="K") + self.check_tensor_signature(v_tensor, self.v_desc, name="V") + return _sdpa_fwd_impl( + q_tensor, + k_tensor, + v_tensor, + self.qk_acc_dtype, + self.pv_acc_dtype, + self.mma_tiler_mn, + self.is_causal, + self.window_size, + self.scale_softmax, + self.scale_output, + ) + + +def sdpa_fwd_wrapper_sm100_d256( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + 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, +) -> TupleDict: + """Compute fixed-shape BHSD SDPA forward with the SM100 d=256 kernel.""" + + return SdpafwdSm100D256( + q_tensor, + k_tensor, + v_tensor, + 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, + )(q_tensor, k_tensor, v_tensor) + + +__all__ = ["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..bf40c7b4b --- /dev/null +++ b/python/cudnn/sdpa/jax_utils.py @@ -0,0 +1,116 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared validation and layout helpers for the JAX SDPA APIs.""" + +from __future__ import annotations + +import math +from typing import Any + +import jax.numpy as jnp +from cutlass.jax import TensorSpec + +from .._jax.api_base import as_dtype, require_array + + +def bhsd_tensor_spec() -> TensorSpec: + """Present a logical BHSD JAX array to a CuTe kernel as compact BSHD. + + The public shape remains ``(B, H, S, D)``. XLA lays out the custom-call + buffer with H inside S, matching the existing Torch API's transposed view, + and ``mode`` presents the kernel with ``(B, S, H, D)`` modes. + """ + + return TensorSpec( + layout=(3, 1, 2, 0), + mode=(0, 2, 1, 3), + ) + + +def require_bhsd_qkv(q: Any, k: Any, v: Any) -> tuple[int, int, int, int, int, int, Any]: + """Validate fixed-shape BHSD Q/K/V tensors and return their dimensions.""" + + q_shape = require_array( + q, + name="q_tensor", + rank=4, + dtype=(jnp.float16, jnp.bfloat16), + ) + dtype = as_dtype(q) + k_shape = require_array(k, name="k_tensor", rank=4, dtype=dtype) + v_shape = require_array(v, name="v_tensor", rank=4, dtype=dtype) + + 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": batch, + "H_q": num_query_heads, + "H_kv": num_kv_heads, + "S_q": seqlen_q, + "S_kv": seqlen_k, + "D": head_dim, + } + nonpositive = [f"{name}={value}" for name, value in dimensions.items() if value <= 0] + if nonpositive: + raise ValueError("SDPA dimensions must be positive, got " + ", ".join(nonpositive)) + if (k_batch, v_batch) != (batch, batch): + raise ValueError("q_tensor, k_tensor, and v_tensor batch dimensions must match, " f"got {batch}, {k_batch}, and {v_batch}") + if (num_value_heads, v_seqlen) != (num_kv_heads, seqlen_k): + raise ValueError("k_tensor and v_tensor head and sequence dimensions must match, " f"got {(num_kv_heads, seqlen_k)} and {(num_value_heads, v_seqlen)}") + if (k_head_dim, value_dim) != (head_dim, head_dim): + raise ValueError("q_tensor, k_tensor, and v_tensor head dimensions must match, " f"got {head_dim}, {k_head_dim}, and {value_dim}") + 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 ( + batch, + num_query_heads, + num_kv_heads, + seqlen_q, + seqlen_k, + head_dim, + 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]: + """Resolve shared mask and scaling configuration for fixed-shape SDPA.""" + + 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("window_size must be (-1, -1) for non-causal mode, got " f"{(window_size_left, window_size_right)}") + + if window_size_left >= seqlen_k - 1: + raise ValueError(f"window_size_left must be less than S_kv - 1 ({seqlen_k - 1}), " f"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}), " f"got {window_size_right}") + + resolved_scale = 1.0 / math.sqrt(256) if scale_softmax is None or scale_softmax == 0.0 else float(scale_softmax) + mask_kind = "window" + if not is_causal and tile_extent % 128: + mask_kind = "residual" + return resolved_scale, window_size_left, window_size_right, mask_kind + + +__all__ = [ + "bhsd_tensor_spec", + "require_bhsd_qkv", + "resolve_sdpa_config", +] diff --git a/test/python/fe_api/conftest.py b/test/python/fe_api/conftest.py new file mode 100644 index 000000000..008d7e03e --- /dev/null +++ b/test/python/fe_api/conftest.py @@ -0,0 +1,73 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Pytest support for optional JAX tests.""" + +from importlib import import_module + +import pytest + +_JAX_TEST_PREFIX = "test_jax_" +_JAX_INSTALL_HINT = "pip install 'nvidia-cudnn-frontend[jax]'" + + +def _requested_module_is_missing(module_name, error): + missing_name = error.name + return missing_name is not None and ( + module_name == missing_name or module_name.startswith(f"{missing_name}.") + ) + + +def _jax_test_skip_reason(): + try: + jax = import_module("jax") + except ModuleNotFoundError as error: + if not _requested_module_is_missing("jax", error): + raise + return f"JAX tests require JAX. Install it with `{_JAX_INSTALL_HINT}`." + + try: + cutlass_jax = import_module("cutlass.jax") + except ModuleNotFoundError as error: + if not _requested_module_is_missing("cutlass.jax", error): + raise + return ( + "JAX tests require CUTLASS JAX support. " + f"Install it with `{_JAX_INSTALL_HINT}`." + ) + + is_available = getattr(cutlass_jax, "is_available", None) + if not callable(is_available): + raise AttributeError("cutlass.jax must define is_available()") + if is_available(): + return None + + installed_version = getattr(jax, "__version__", "unknown") + minimum_version_info = getattr( + cutlass_jax, + "CUTE_DSL_MIN_SUPPORTED_JAX_VERSION", + None, + ) + if minimum_version_info is None: + return f"CUTLASS JAX support is unavailable with JAX {installed_version}." + + minimum_version = ".".join(str(part) for part in minimum_version_info) + return ( + f"CUTLASS JAX support is unavailable with JAX {installed_version}; " + f"the minimum supported JAX version is {minimum_version}." + ) + + +def pytest_collection_modifyitems(config, items): + del config + jax_items = [item for item in items if item.path.name.startswith(_JAX_TEST_PREFIX)] + if not jax_items: + return + + reason = _jax_test_skip_reason() + if reason is None: + return + + skip = pytest.mark.skip(reason=reason) + for item in jax_items: + item.add_marker(skip) diff --git a/test/python/fe_api/dsa/test_DSA_dense_indexer_backward.py b/test/python/fe_api/dsa/test_DSA_dense_indexer_backward.py index f8325a73f..a2bddbd0e 100644 --- a/test/python/fe_api/dsa/test_DSA_dense_indexer_backward.py +++ b/test/python/fe_api/dsa/test_DSA_dense_indexer_backward.py @@ -56,6 +56,7 @@ def _allocate(cfg, sm_scale: float, ratio: int, q_causal_offsets: torch.Tensor): @pytest.mark.L0 @torch_fork_set_rng(seed=0) +@pytest.mark.parametrize("has_q_causal_offsets", [False, True]) @with_dsa_dense_indexer_backward_params def test_DSA_dense_indexer_backward_wrapper( dtype, @@ -64,6 +65,7 @@ def test_DSA_dense_indexer_backward_wrapper( qhead_per_kv_head, block_I, ratio, + has_q_causal_offsets, request, ): try: @@ -87,7 +89,11 @@ def test_DSA_dense_indexer_backward_wrapper( sm_scale = 1.0 b_cfg = cfg["b"] s_q_cfg = cfg["s_q"] - q_causal_offsets = torch.full((b_cfg,), 8, dtype=torch.int32, device="cuda") + q_causal_offsets = ( + torch.full((b_cfg,), 8, dtype=torch.int32, device="cuda") + if has_q_causal_offsets + else None + ) loss_coeff = float(b_cfg * s_q_cfg) grad_loss = 1.0 grad_scale_expected = (loss_coeff / (b_cfg * s_q_cfg)) * grad_loss @@ -125,6 +131,8 @@ def test_DSA_dense_indexer_backward_wrapper( stream=stream, ) except (ValueError, NotImplementedError, RuntimeError) as e: + if torch.cuda.get_device_capability()[0] >= 10: + raise pytest.skip(f"Unsupported testcase: {e}") torch_stream.synchronize() diff --git a/test/python/fe_api/test_api_base_framework_split.py b/test/python/fe_api/test_api_base_framework_split.py new file mode 100644 index 000000000..0f4c52bb7 --- /dev/null +++ b/test/python/fe_api/test_api_base_framework_split.py @@ -0,0 +1,578 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for the framework-specific API base split.""" + +from __future__ import annotations + +import ast +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CUDNN_ROOT = _REPO_ROOT / "python" / "cudnn" + + +class _DType: + def __init__(self, name, itemsize=None): + self.name = name + self.itemsize = itemsize + + +def _install_test_package(name): + package = types.ModuleType(name) + package.__path__ = [str(_CUDNN_ROOT)] + package.__package__ = name + package.__spec__ = ModuleSpec(name, loader=None, is_package=True) + sys.modules[name] = package + + +def _remove_test_package(name): + for module_name in tuple(sys.modules): + if module_name == name or module_name.startswith(f"{name}."): + sys.modules.pop(module_name, None) + + +class ApiBaseFrameworkSplitTest(unittest.TestCase): + def test_torch_wrappers_use_explicit_base_name(self): + for path in _CUDNN_ROOT.rglob("api.py"): + tree = ast.parse(path.read_text(), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef): + legacy_bases = [base for base in node.bases if isinstance(base, ast.Name) and base.id == "APIBase"] + self.assertFalse(legacy_bases, f"{path} must inherit ApiBaseTorch explicitly") + + def test_neutral_base_does_not_load_optional_frameworks(self): + package_name = "cudnn_frontend_neutral_api_base_test" + _install_test_package(package_name) + try: + with mock.patch.dict( + sys.modules, + { + "torch": None, + "jax": None, + "jax.numpy": None, + "cutlass": None, + "cuda": None, + }, + ): + module = importlib.import_module(f"{package_name}.api_base") + self.assertNotIn(f"{package_name}.api_base_torch", sys.modules) + + desc = module.TensorDesc( + dtype=_DType("bfloat16", 2), + shape=(4, 8), + stride_order=(1, 0), + name="x", + ) + self.assertEqual(desc.shape, (4, 8)) + self.assertEqual(desc.stride, (8, 1)) + self.assertEqual(desc.dtype_name, "bfloat16") + self.assertEqual(desc.element_bits, 16) + + base = module.ApiBase() + self.assertEqual(base.check_tensor_shape(desc, (4, 8), "x"), (4, 8)) + self.assertEqual(base.check_dtype(desc, (_DType("bfloat16"),), "x"), "bfloat16") + + fp4 = module.TensorDesc( + dtype=_DType("float4_e2m1fn_x2", 1), + shape=(8,), + ) + self.assertEqual(fp4.dtype_name, "float4_e2m1fn") + self.assertEqual(fp4.element_bits, 4) + self.assertTrue(fp4.is_fp4) + self.assertEqual(fp4.storage_dtype_name, "float4_e2m1fn_x2") + self.assertEqual(fp4.packing, "fp4x2") + + packed_fp4 = module.TensorDesc( + dtype=_DType("uint8", 1), + shape=(8,), + packing="fp4x2", + ) + self.assertEqual(packed_fp4.storage_dtype_name, "uint8") + self.assertEqual(packed_fp4.dtype_name, "float4_e2m1fn") + self.assertEqual(packed_fp4.element_bits, 4) + finally: + _remove_test_package(package_name) + + def test_jax_descriptor_and_callable_are_abstract_and_stable(self): + package_name = "cudnn_frontend_jax_api_base_test" + _install_test_package(package_name) + + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.dtype = lambda value: value + fake_jax = types.ModuleType("jax") + fake_jax.__path__ = [] + fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + fake_jax.numpy = fake_jnp + pytree_registrations = [] + + def register_pytree_with_keys(node_type, flatten_with_keys, unflatten, flatten): + pytree_registrations.append((node_type, flatten_with_keys, unflatten, flatten)) + + fake_jax.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=register_pytree_with_keys, + ) + + def identity_jit(fn=None, **_kwargs): + return (lambda decorated_fn: decorated_fn) if fn is None else fn + + fake_cutlass_jax = types.ModuleType("cutlass.jax") + + class _TensorSpec: + def __init__( + self, + *, + layout=None, + mode=None, + static=None, + ptr_assumed_align=256, + divisibility=None, + ): + self.layout = layout + self.mode = mode + self.static = static + self.ptr_assumed_align = ptr_assumed_align + self.divisibility = divisibility + + fake_cutlass_jax.TensorSpec = _TensorSpec + fake_cute = types.ModuleType("cutlass.cute") + fake_cute.jit = identity_jit + fake_cutlass = types.ModuleType("cutlass") + fake_cutlass.__path__ = [] + fake_cutlass.Constexpr = object + fake_cutlass.cute = fake_cute + fake_cutlass.jax = fake_cutlass_jax + + try: + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "cutlass": fake_cutlass, + "cutlass.cute": fake_cute, + "cutlass.jax": fake_cutlass_jax, + "torch": None, + }, + ): + module = importlib.import_module(f"{package_name}._jax.api_base") + + self.assertEqual(len(pytree_registrations), 1) + node_type, flatten_with_keys, unflatten, flatten = pytree_registrations[0] + self.assertIs(node_type, module.TupleDict) + result = module.TupleDict(output="output", values="values") + children, keys = flatten(result) + self.assertEqual(children, ("output", "values")) + self.assertEqual(keys, ("output", "values")) + keyed_children, keyed_aux = flatten_with_keys(result) + self.assertEqual(keyed_children, (("output", "output"), ("values", "values"))) + self.assertEqual(keyed_aux, keys) + restored = unflatten(keys, children) + self.assertIsInstance(restored, module.TupleDict) + self.assertEqual(tuple(restored.keys()), keys) + self.assertEqual(tuple(restored), children) + self.assertEqual(restored[0], "output") + self.assertEqual(restored["values"], "values") + self.assertIs(importlib.reload(module), module) + self.assertEqual(len(pytree_registrations), 1) + + value = types.SimpleNamespace( + shape=(2, 3, 5, 7, 11, 13), + dtype=_DType("float8_e4m3fn", 1), + ) + desc = module.JaxTensorDesc.from_value( + value, + layout=(2, 1, 4, 0, 3, 5), + name="scale", + ) + self.assertEqual(desc.layout, (2, 1, 4, 0, 3, 5)) + self.assertEqual(desc.stride_order, (3, 1, 0, 4, 2, 5)) + self.assertEqual(desc.stride, (21, 7, 462, 1, 42, 2310)) + self.assertEqual(desc.dtype_name, "float8_e4m3fn") + self.assertIsInstance(desc.tensor_spec, _TensorSpec) + self.assertEqual(desc.tensor_spec.layout, (2, 1, 4, 0, 3, 5)) + + mode_only = module.JaxTensorDesc.from_value( + types.SimpleNamespace(shape=(2, 3), dtype=_DType("bfloat16", 2)), + mode=(1, 0), + ) + self.assertEqual(mode_only.shape, (3, 2)) + self.assertIsNone(mode_only.tensor_spec.layout) + self.assertEqual(mode_only.tensor_spec.mode, (1, 0)) + + implicit = module.JaxTensorDesc.from_value( + types.SimpleNamespace(shape=(2, 3), dtype=_DType("bfloat16", 2)), + ) + self.assertIsNone(implicit.tensor_spec) + self.assertEqual(implicit.layout, (1, 0)) + + transposed_value = types.SimpleNamespace(shape=(2, 3, 5), dtype=_DType("bfloat16", 2)) + transposed_spec = _TensorSpec( + layout=(2, 1, 0), + mode=(1, 2, 0), + ) + transposed = module.JaxTensorDesc.from_value( + transposed_value, + tensor_spec=transposed_spec, + ) + self.assertEqual(transposed.shape, (3, 5, 2)) + self.assertEqual(transposed.stride, (5, 1, 15)) + self.assertEqual(transposed.stride_order, (1, 0, 2)) + self.assertEqual(transposed.mode, (1, 2, 0)) + self.assertIs(transposed.tensor_spec, transposed_spec) + + with self.subTest("shared array validation"): + bfloat16 = transposed.dtype + float32 = _DType("float32", 4) + array = types.SimpleNamespace(shape=[3, 5, 2], dtype=bfloat16) + + self.assertEqual( + module.require_array( + array, + name="array", + rank=3, + shape=(3, 5, 2), + dtype=bfloat16, + ), + (3, 5, 2), + ) + self.assertEqual( + module.require_array( + types.SimpleNamespace(shape=(3, 5, 2), dtype=float32), + name="array", + dtype=(bfloat16, float32), + ), + (3, 5, 2), + ) + + for incomplete in ( + types.SimpleNamespace(shape=(3, 5, 2)), + types.SimpleNamespace(dtype=bfloat16), + ): + with self.assertRaisesRegex(TypeError, "must have shape and dtype metadata"): + module.require_array(incomplete, name="array") + + with self.assertRaisesRegex(TypeError, "value must have shape and dtype metadata"): + module.require_array(object()) + + with self.assertRaisesRegex(ValueError, "array must have rank 2"): + module.require_array(array, name="array", rank=2) + with self.assertRaisesRegex(ValueError, "array must have shape"): + module.require_array(array, name="array", shape=(2, 5, 3)) + with self.assertRaisesRegex(ValueError, "array.dtype must be one of"): + module.require_array( + array, + name="array", + dtype=(float32,), + ) + + # A descriptor exposes the shape seen by the kernel; its + # public input shape remains available through array_shape. + named_transposed = module.JaxTensorDesc.from_value( + transposed_value, + tensor_spec=transposed_spec, + name="transposed", + ) + self.assertEqual( + module.require_array( + named_transposed, + shape=(3, 5, 2), + dtype=bfloat16, + ), + named_transposed.shape, + ) + self.assertEqual(named_transposed.array_shape, (2, 3, 5)) + with self.assertRaisesRegex(ValueError, "transposed must have shape"): + module.require_array( + named_transposed, + shape=named_transposed.array_shape, + ) + + inferred_spec = _TensorSpec() + default_spec = module.JaxTensorDesc.from_value( + types.SimpleNamespace(shape=(2, 3), dtype=_DType("bfloat16", 2)), + tensor_spec=inferred_spec, + ) + self.assertEqual(default_spec.layout, (1, 0)) + self.assertEqual(default_spec.mode, (0, 1)) + self.assertEqual(default_spec.stride, (3, 1)) + self.assertIs(default_spec.tensor_spec, inferred_spec) + + direct = module.JaxTensorDesc( + dtype=_DType("bfloat16", 2), + shape=(3, 5, 2), + tensor_spec=transposed_spec, + ) + self.assertEqual(direct.layout, (2, 1, 0)) + self.assertEqual(direct.mode, (1, 2, 0)) + self.assertEqual(direct.stride, (5, 1, 15)) + self.assertIs(direct.tensor_spec, transposed_spec) + + lowering_hints = _TensorSpec( + static=False, + ptr_assumed_align=128, + divisibility=(16, None), + ) + hinted = module.JaxTensorDesc.from_value( + types.SimpleNamespace(shape=(16, 3), dtype=_DType("bfloat16", 2)), + tensor_spec=lowering_hints, + ) + self.assertIs(hinted.tensor_spec, lowering_hints) + + with self.assertRaisesRegex(ValueError, "tensor_spec cannot be combined"): + module.JaxTensorDesc.from_value( + transposed_value, + tensor_spec=transposed_spec, + layout=(2, 1, 0), + ) + with self.assertRaisesRegex(ValueError, "tensor_spec cannot be combined"): + module.JaxTensorDesc.from_value( + transposed_value, + tensor_spec=transposed_spec, + mode=(1, 2, 0), + ) + + with self.assertRaisesRegex(ValueError, "stride_order must agree"): + module.JaxTensorDesc( + dtype=_DType("bfloat16", 2), + shape=(2, 3), + stride_order=(0, 1), + tensor_spec=_TensorSpec(layout=(1, 0)), + ) + + with self.assertRaisesRegex(ValueError, "stride must describe the compact layout"): + module.JaxTensorDesc( + dtype=_DType("bfloat16", 2), + shape=(2, 3), + stride=(1, 2), + tensor_spec=_TensorSpec(layout=(1, 0)), + ) + + with self.assertRaisesRegex(ValueError, "mode must be a permutation"): + module.JaxTensorDesc.from_value( + types.SimpleNamespace(shape=(2, 3, 5), dtype=_DType("bfloat16", 2)), + mode=(0, 1), + ) + + class _JaxApi(module.ApiBaseJax): + def __init__(self, sample): + super().__init__() + self.sample_desc = self.make_tensor_desc(sample, name="sample") + self.static_option = 1 + self.check_count = 0 + + def _check_support(self) -> None: + self.check_count += 1 + + def _call_impl(self, x): + self.check_tensor_signature(x, self.sample_desc, name="input") + return x + + api = _JaxApi(value) + self.assertTrue(all(stored is not value for stored in vars(api).values())) + self.assertEqual( + api.make_tensor_desc( + transposed_value, + tensor_spec=transposed_spec, + ), + transposed, + ) + self.assertEqual( + api.make_optional_tensor_desc( + transposed_value, + tensor_spec=transposed_spec, + ), + transposed, + ) + self.assertIsNone(api.make_optional_tensor_desc(None, tensor_spec=transposed_spec)) + checked = api.check_tensor_signature(transposed_value, transposed, name="input") + self.assertIs(checked.tensor_spec, transposed_spec) + checked_default = api.check_tensor_signature(value, api.sample_desc, name="input") + self.assertIsNone(checked_default.tensor_spec) + self.assertIs(api.get_jax_callable(), api) + self.assertIs(api.as_dtype(value), value.dtype) + self.assertIs(api.as_optional_dtype(value), value.dtype) + self.assertIsNone(api.as_optional_dtype(None)) + self.assertIs( + api.require_dtype(api.sample_desc, (value.dtype,)), + value.dtype, + ) + self.assertIs( + api.require_dtype( + None, + (value.dtype,), + name="output_dtype", + default=value.dtype, + ), + value.dtype, + ) + with self.assertRaisesRegex(ValueError, "output_dtype must not be None"): + api.require_dtype(None, (value.dtype,), name="output_dtype") + with self.assertRaisesRegex(ValueError, "sample.dtype must be one of"): + api.require_dtype( + api.sample_desc, + (_DType("bfloat16", 2),), + ) + self.assertEqual(api.check_tensor_signature(value, desc, name="scale"), desc) + with self.assertRaisesRegex(ValueError, "scale tensor shape mismatch"): + api.check_tensor_signature( + types.SimpleNamespace(shape=(2, 3), dtype=_DType("float8_e4m3fn", 1)), + desc, + name="scale", + ) + with self.assertRaisesRegex(ValueError, "scale dtype mismatch"): + api.check_tensor_signature( + types.SimpleNamespace(shape=value.shape, dtype=_DType("bfloat16", 2)), + desc, + name="scale", + ) + self.assertIsNone(api.make_optional_tensor_desc(None, name="bias")) + self.assertIsNone(api.check_optional_tensor_signature(None, None, name="bias")) + with self.assertRaisesRegex(ValueError, "bias presence mismatch"): + api.check_optional_tensor_signature(value, None, name="bias") + self.assertIs(api(value), value) + self.assertIs(api(value), value) + self.assertEqual(api.check_count, 1) + self.assertTrue(api.check_support()) + self.assertEqual(api.check_count, 1) + with self.assertRaisesRegex(AttributeError, "immutable after its first call"): + api.static_option = 2 + + mutable_before_call = _JaxApi(value) + self.assertTrue(mutable_before_call.check_support()) + mutable_before_call.static_option = 2 + self.assertFalse(mutable_before_call._is_supported) + self.assertIs(mutable_before_call(value), value) + self.assertEqual(mutable_before_call.check_count, 2) + self.assertFalse(hasattr(fake_jax, "jit")) + finally: + _remove_test_package(package_name) + + def test_legacy_api_base_lazily_resolves_to_torch_alias(self): + package_name = "cudnn_frontend_torch_api_base_test" + _install_test_package(package_name) + + fake_torch = types.ModuleType("torch") + + class _Device: + def __init__(self, value): + self.value = value + + def __eq__(self, other): + return isinstance(other, _Device) and self.value == other.value + + fake_torch.device = _Device + fake_torch.Tensor = type("Tensor", (), {}) + fake_torch.dtype = _DType + fake_torch.Size = tuple + fake_torch.memory_format = type("memory_format", (), {}) + fake_torch.contiguous_format = object() + fake_torch.preserve_format = object() + fake_torch.channels_last = object() + fake_torch.channels_last_3d = object() + fake_torch.uint8 = _DType("uint8", 1) + fake_torch.float4_e2m1fn_x2 = _DType("float4_e2m1fn_x2", 1) + + fake_cuda = types.ModuleType("cuda") + fake_cuda.__path__ = [] + fake_bindings = types.ModuleType("cuda.bindings") + fake_bindings.__path__ = [] + fake_driver = types.ModuleType("cuda.bindings.driver") + fake_cuda.bindings = fake_bindings + fake_bindings.driver = fake_driver + + fake_cutlass = types.ModuleType("cutlass") + fake_cutlass.__path__ = [] + fake_cute = types.ModuleType("cutlass.cute") + fake_cutlass.cute = fake_cute + + fake_datatypes = types.ModuleType(f"{package_name}.datatypes") + fake_datatypes._convert_to_cutlass_data_type = lambda dtype, **_kwargs: dtype + + try: + with mock.patch.dict( + sys.modules, + { + "torch": fake_torch, + "cuda": fake_cuda, + "cuda.bindings": fake_bindings, + "cuda.bindings.driver": fake_driver, + "cutlass": fake_cutlass, + "cutlass.cute": fake_cute, + f"{package_name}.datatypes": fake_datatypes, + }, + ): + neutral = importlib.import_module(f"{package_name}.api_base") + self.assertNotIn(f"{package_name}.api_base_torch", sys.modules) + + legacy = neutral.APIBase + self.assertIn(f"{package_name}.api_base_torch", sys.modules) + self.assertIs(legacy, neutral.ApiBaseTorch) + self.assertTrue(issubclass(neutral.ApiBaseTorch, neutral.ApiBase)) + self.assertEqual(legacy.__name__, "ApiBaseTorch") + + desc = neutral.TorchTensorDesc( + dtype=_DType("float32", 4), + shape=(2, 3), + stride=(3, 1), + stride_order=(1, 0), + device="cuda:0", + ) + self.assertIsInstance(desc, neutral.TensorDesc) + self.assertEqual(desc.device, _Device("cuda:0")) + self.assertEqual(desc.size(), (2, 3)) + transposed = desc.transpose(0, 1) + self.assertIsInstance(transposed, neutral.TorchTensorDesc) + self.assertEqual(transposed.shape, (3, 2)) + self.assertEqual(transposed.device, desc.device) + + class _TorchApi(neutral.ApiBaseTorch): + def check_support(self): + return True + + def compile(self): + return None + + def execute(self, *_args, **_kwargs): + return None + + api = _TorchApi() + self.assertEqual(api._check_tensor_shape(desc, (2, 3), "x"), (2, 3)) + self.assertEqual( + api._check_tensor_stride(desc, stride=(3, 1), name="x"), + ((3, 1), (1, 0)), + ) + self.assertEqual( + api._check_tensor_stride(None, stride=[(3, 1)], name="optional"), + (None, None), + ) + self.assertIs(api._check_dtype(desc, _DType("float32", 4), "x"), desc.dtype) + with self.assertRaisesRegex(ValueError, "must match another operand"): + api._check_dtype(desc, _DType("float16", 2), "x", "must match another operand") + + packed_tensor = fake_torch.Tensor() + packed_tensor.dtype = fake_torch.uint8 + packed_tensor.shape = (2, 4) + packed_tensor.device = _Device("cuda:0") + packed_tensor.stride = lambda: (4, 1) + packed_desc = neutral.TorchTensorDesc.from_tensor( + packed_tensor, + interpret_uint8_as_fp4x2=True, + ) + self.assertEqual(packed_desc.shape, (2, 8)) + self.assertEqual(packed_desc.stride, (8, 1)) + self.assertEqual(packed_desc.packing, "fp4x2") + self.assertEqual(packed_desc.element_bits, 4) + finally: + _remove_test_package(package_name) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_gemm_amax_validation.py b/test/python/fe_api/test_gemm_amax_validation.py new file mode 100644 index 000000000..bf7d8572d --- /dev/null +++ b/test/python/fe_api/test_gemm_amax_validation.py @@ -0,0 +1,197 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for shared GEMM + amax validation.""" + +from __future__ import annotations + +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 + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CUDNN_ROOT = _REPO_ROOT / "python" / "cudnn" +_PACKAGE = "cudnn_frontend_gemm_amax_validation_test" + + +def _canonical_dtype_name(dtype): + if not isinstance(dtype, type) and hasattr(dtype, "dtype_name"): + return dtype.dtype_name + name = getattr(dtype, "name", str(dtype)).rsplit(".", 1)[-1].lower() + return {"float4_e2m1fn_x2": "float4_e2m1fn"}.get(name, name) + + +def _load_validation_modules(): + 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 + + fake_api_base = types.ModuleType(f"{_PACKAGE}.api_base") + fake_api_base.TensorDesc = object + fake_api_base.canonical_dtype_name = _canonical_dtype_name + sys.modules[fake_api_base.__name__] = fake_api_base + + operation = types.ModuleType(f"{_PACKAGE}.gemm_amax") + operation.__path__ = [str(_CUDNN_ROOT / "gemm_amax")] + operation.__package__ = operation.__name__ + operation.__spec__ = ModuleSpec(operation.__name__, loader=None, is_package=True) + sys.modules[operation.__name__] = operation + + gemm_validation = importlib.import_module(f"{_PACKAGE}.gemm_validation") + operation_validation = importlib.import_module(f"{_PACKAGE}.gemm_amax.validation") + return gemm_validation, operation_validation + + +_GEMM_VALIDATION, _VALIDATION = _load_validation_modules() + + +def _compact_stride(shape, order): + stride = [None] * len(shape) + running = 1 + for dim in order: + stride[dim] = running + running *= shape[dim] + return tuple(stride) + + +def _desc(shape, dtype_name, *, order=None, packed_fp4=False, name=""): + shape = tuple(shape) + bits = { + "int8": 8, + "uint8": 8, + "float4_e2m1fn": 4, + "float8_e4m3fn": 8, + "float8_e5m2": 8, + "float8_e8m0fnu": 8, + "bfloat16": 16, + "float16": 16, + "float32": 32, + }[dtype_name] + return types.SimpleNamespace( + shape=shape, + ndim=len(shape), + dtype_name=dtype_name, + storage_dtype_name=dtype_name, + element_bits=bits, + stride=None if order is None else _compact_stride(shape, order), + stride_order=order, + packing="fp4x2" if packed_fp4 else "native", + interpret_uint8_as_fp4x2=packed_fp4, + name=name, + ) + + +def _valid_descriptors(*, ab_dtype="float8_e4m3fn", c_dtype="float32", k=128, sf_vec_size=32): + m = n = 128 + batch = 1 + scale_a = _GEMM_VALIDATION.block_scale_shape(m, k, batch, sf_vec_size) + scale_b = _GEMM_VALIDATION.block_scale_shape(n, k, batch, sf_vec_size) + return { + "a": _desc((m, k, batch), ab_dtype, order=(1, 0, 2), name="A"), + "b": _desc((n, k, batch), ab_dtype, order=(1, 0, 2), name="B"), + "sfa": _desc(scale_a, "float8_e8m0fnu", name="sfa"), + "sfb": _desc(scale_b, "float8_e8m0fnu", name="sfb"), + "c": _desc((m, n, batch), c_dtype, order=(1, 0, 2), name="C"), + "amax": _desc((1, 1, 1), "float32", name="amax"), + } + + +def _validate(descs, **kwargs): + return _VALIDATION.validate_gemm_amax( + descs["a"], + descs["b"], + descs["sfa"], + descs["sfb"], + descs["c"], + descs["amax"], + acc_dtype=kwargs.pop("acc_dtype", "float32"), + sf_vec_size=kwargs.pop("sf_vec_size", 32), + supported_sf_vec_sizes=(16, 32), + mma_tiler_mn=kwargs.pop("mma_tiler_mn", (128, 128)), + **kwargs, + ) + + +class GemmAmaxValidationTest(unittest.TestCase): + def test_returns_shared_logical_plan(self): + plan = _validate(_valid_descriptors()) + + self.assertEqual((plan.m, plan.n, plan.k, plan.batch), (128, 128, 128, 1)) + self.assertEqual(plan.c_shape, (128, 128, 1)) + self.assertEqual(plan.amax_shape, (1, 1, 1)) + self.assertEqual((plan.a_major, plan.b_major, plan.c_major), ("k", "k", "n")) + self.assertEqual(plan.ab_dtype_name, "float8_e4m3fn") + + def test_accepts_torch_packed_fp4_as_logical_fp4(self): + descs = _valid_descriptors(ab_dtype="uint8") + descs["a"] = _desc((128, 128, 1), "uint8", order=(1, 0, 2), packed_fp4=True, name="A") + descs["b"] = _desc((128, 128, 1), "uint8", order=(1, 0, 2), packed_fp4=True, name="B") + + plan = _validate(descs) + + self.assertEqual(plan.ab_dtype_name, "float4_e2m1fn") + + def test_rejects_invalid_shapes_dtypes_and_layouts(self): + cases = [] + + descs = _valid_descriptors() + descs["b"] = _desc((128, 128, 1), "float8_e5m2", order=(1, 0, 2)) + cases.append((descs, {}, "a_tensor and b_tensor dtypes must match")) + + descs = _valid_descriptors() + descs["sfa"] = _desc((32, 4, 1, 4, 2, 1), "float8_e8m0fnu") + cases.append((descs, {}, "sfa_tensor must have shape")) + + descs = _valid_descriptors() + descs["c"] = _desc((128, 64, 1), "float32", order=(1, 0, 2)) + cases.append((descs, {}, "C must have shape")) + + descs = _valid_descriptors() + descs["a"].stride = (256, 2, 32768) + cases.append((descs, {}, "A tensor must use one of the compact strides")) + + for descs, kwargs, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + _validate(descs, **kwargs) + + def test_rejects_kernel_domain_combinations(self): + descs = _valid_descriptors(c_dtype="float8_e5m2") + with self.assertRaisesRegex(NotImplementedError, "FP8 A/B with FP8 C"): + _validate(descs) + + descs = _valid_descriptors(sf_vec_size=16) + with self.assertRaisesRegex(ValueError, "do not support sf_vec_size=16"): + _validate(descs, sf_vec_size=16) + + descs = _valid_descriptors(ab_dtype="float4_e2m1fn", k=128) + with self.assertRaisesRegex(ValueError, "requires K > 128"): + _validate(descs, mma_tiler_mn=(128, 256)) + + descs = _valid_descriptors(ab_dtype="float4_e2m1fn", k=128, sf_vec_size=16) + descs["sfa"] = _desc(descs["sfa"].shape, "int8") + descs["sfb"] = _desc(descs["sfb"].shape, "int8") + with self.assertRaisesRegex(ValueError, "sfa_tensor and sfb_tensor dtype must be one of"): + _validate(descs, sf_vec_size=16) + + def test_rejects_unaligned_contiguous_extent(self): + descs = _valid_descriptors(k=120) + with self.assertRaisesRegex(ValueError, "16-byte aligned"): + _validate(descs) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_gemm_swiglu_validation.py b/test/python/fe_api/test_gemm_swiglu_validation.py new file mode 100644 index 000000000..0333b7e9f --- /dev/null +++ b/test/python/fe_api/test_gemm_swiglu_validation.py @@ -0,0 +1,282 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for shared GEMM + SwiGLU validation.""" + +from __future__ import annotations + +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 + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CUDNN_ROOT = _REPO_ROOT / "python" / "cudnn" +_PACKAGE = "cudnn_frontend_gemm_swiglu_validation_test" + + +def _canonical_dtype_name(dtype): + if not isinstance(dtype, type) and hasattr(dtype, "dtype_name"): + return dtype.dtype_name + name = getattr(dtype, "name", str(dtype)).rsplit(".", 1)[-1].lower() + return {"float4_e2m1fn_x2": "float4_e2m1fn"}.get(name, name) + + +def _load_validation_modules(): + 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 + + fake_api_base = types.ModuleType(f"{_PACKAGE}.api_base") + fake_api_base.TensorDesc = object + fake_api_base.canonical_dtype_name = _canonical_dtype_name + sys.modules[fake_api_base.__name__] = fake_api_base + + operation = types.ModuleType(f"{_PACKAGE}.gemm_swiglu") + operation.__path__ = [str(_CUDNN_ROOT / "gemm_swiglu")] + operation.__package__ = operation.__name__ + operation.__spec__ = ModuleSpec(operation.__name__, loader=None, is_package=True) + sys.modules[operation.__name__] = operation + + gemm_validation = importlib.import_module(f"{_PACKAGE}.gemm_validation") + operation_validation = importlib.import_module(f"{_PACKAGE}.gemm_swiglu.validation") + return gemm_validation, operation_validation + + +_GEMM_VALIDATION, _VALIDATION = _load_validation_modules() + + +def _compact_stride(shape, order): + stride = [None] * len(shape) + running = 1 + for dim in order: + stride[dim] = running + running *= shape[dim] + return tuple(stride) + + +def _desc(shape, dtype_name, *, order=None, packing="native", name=""): + shape = tuple(shape) + bits = { + "uint8": 8, + "float4_e2m1fn": 4, + "float8_e4m3fn": 8, + "float8_e5m2": 8, + "float8_e8m0fnu": 8, + "bfloat16": 16, + "float16": 16, + "float32": 32, + }[dtype_name] + return types.SimpleNamespace( + shape=shape, + ndim=len(shape), + dtype_name=dtype_name, + storage_dtype_name=dtype_name, + element_bits=bits, + stride=None if order is None else _compact_stride(shape, order), + stride_order=order, + packing=packing, + name=name, + ) + + +def _valid_descriptors( + *, + ab_dtype="float8_e4m3fn", + scale_dtype="float8_e8m0fnu", + ab12_dtype="bfloat16", + c_dtype="bfloat16", + k=128, + sf_vec_size=32, + packing="native", + a_order=(1, 0, 2), + b_order=(1, 0, 2), + output_order=(1, 0, 2), + include_amax=False, +): + m = n = 128 + batch = 1 + output_n = n // 2 + return { + "a": _desc( + (m, k, batch), + ab_dtype, + order=a_order, + packing=packing, + name="A", + ), + "b": _desc( + (n, k, batch), + ab_dtype, + order=b_order, + packing=packing, + name="B", + ), + "ab12": _desc( + (m, n, batch), + ab12_dtype, + order=output_order, + name="AB12", + ), + "c": _desc( + (m, output_n, batch), + c_dtype, + order=output_order, + name="C", + ), + "sfa": _desc( + _GEMM_VALIDATION.block_scale_shape(m, k, batch, sf_vec_size), + scale_dtype, + name="SFA", + ), + "sfb": _desc( + _GEMM_VALIDATION.block_scale_shape(n, k, batch, sf_vec_size), + scale_dtype, + name="SFB", + ), + "amax": (_desc((1, 1, 1), "float32", name="amax") if include_amax else None), + "sfc": None, + "norm_const": None, + } + + +def _validate(descs, **kwargs): + return _VALIDATION.validate_quantized_gemm_swiglu( + descs["a"], + descs["b"], + descs["ab12"], + descs["c"], + sfa=descs["sfa"], + sfb=descs["sfb"], + amax=descs["amax"], + sfc=descs["sfc"], + norm_const=descs["norm_const"], + acc_dtype=kwargs.pop("acc_dtype", "float32"), + output_n=64, + sf_vec_size=kwargs.pop("sf_vec_size", 32), + supported_sf_vec_sizes=(16, 32), + mma_tiler_mn=kwargs.pop("mma_tiler_mn", (128, 128)), + **kwargs, + ) + + +class GemmSwigluValidationTest(unittest.TestCase): + def test_returns_mxfp8_plan(self): + plan = _validate(_valid_descriptors()) + + self.assertEqual((plan.m, plan.n, plan.k, plan.batch), (128, 128, 128, 1)) + self.assertEqual(plan.ab12_shape, (128, 128, 1)) + self.assertEqual(plan.c_shape, (128, 64, 1)) + self.assertEqual((plan.a_major, plan.b_major, plan.output_major), ("k", "k", "n")) + self.assertEqual(plan.ab_dtype_name, "float8_e4m3fn") + self.assertFalse(plan.generate_amax) + self.assertFalse(plan.generate_sfc) + + def test_accepts_native_and_torch_packed_fp4(self): + cases = ( + ("float4_e2m1fn", "native"), + ("uint8", "fp4x2"), + ) + for dtype_name, packing in cases: + with self.subTest(dtype_name=dtype_name, packing=packing): + descs = _valid_descriptors( + ab_dtype=dtype_name, + scale_dtype="float8_e4m3fn", + sf_vec_size=16, + packing=packing, + include_amax=True, + ) + plan = _validate(descs, sf_vec_size=16) + + self.assertEqual(plan.ab_dtype_name, "float4_e2m1fn") + self.assertTrue(plan.requires_amax) + self.assertTrue(plan.generate_amax) + + def test_requires_a_complete_scale_pair(self): + for missing in ("sfa", "sfb"): + with self.subTest(missing=missing): + descs = _valid_descriptors() + descs[missing] = None + with self.assertRaisesRegex( + ValueError, + "sfa_tensor and sfb_tensor are required", + ): + _validate(descs) + + def test_rejects_invalid_scale_formats(self): + cases = ( + ( + _valid_descriptors(sf_vec_size=16), + 16, + "FP8 A and B require float8_e8m0fnu scale factors and sf_vec_size=32", + ), + ( + _valid_descriptors(scale_dtype="float8_e4m3fn"), + 32, + "FP8 A and B require float8_e8m0fnu scale factors and sf_vec_size=32", + ), + ( + _valid_descriptors( + ab_dtype="float4_e2m1fn", + scale_dtype="float8_e4m3fn", + sf_vec_size=32, + include_amax=True, + ), + 32, + "FP4 A and B do not support float8_e4m3fn scale factors with sf_vec_size=32", + ), + ) + for descs, sf_vec_size, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + _validate(descs, sf_vec_size=sf_vec_size) + + def test_rejects_raw_uint8_and_non_k_major_fp4(self): + raw_uint8 = _valid_descriptors(ab_dtype="uint8") + with self.assertRaisesRegex( + ValueError, + "a_tensor and b_tensor dtype must be one of", + ): + _validate(raw_uint8) + + wrong_layout = _valid_descriptors( + ab_dtype="float4_e2m1fn", + sf_vec_size=16, + a_order=(0, 1, 2), + include_amax=True, + ) + with self.assertRaisesRegex(ValueError, "FP4 requires k-major A and B"): + _validate(wrong_layout, sf_vec_size=16) + + def test_requires_fp4_amax_and_uses_four_bit_alignment(self): + missing_amax = _valid_descriptors( + ab_dtype="float4_e2m1fn", + sf_vec_size=16, + ) + with self.assertRaisesRegex(ValueError, "amax_tensor is required"): + _validate(missing_amax, sf_vec_size=16) + + unaligned = _valid_descriptors( + ab_dtype="float4_e2m1fn", + k=16, + sf_vec_size=16, + include_amax=True, + ) + with self.assertRaisesRegex(ValueError, "A's contiguous extent must be 16-byte aligned"): + _validate(unaligned, sf_vec_size=16) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_gemm_validation.py b/test/python/fe_api/test_gemm_validation.py new file mode 100644 index 000000000..bf89a57da --- /dev/null +++ b/test/python/fe_api/test_gemm_validation.py @@ -0,0 +1,141 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for shared dense GEMM validation.""" + +import importlib.util +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" / "gemm_validation.py" +_SPEC = importlib.util.spec_from_file_location("cudnn_gemm_validation_test", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_MODULE = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _MODULE +_SPEC.loader.exec_module(_MODULE) + + +class GemmValidationTest(unittest.TestCase): + def test_reconciles_input_shapes(self): + self.assertEqual( + _MODULE.require_gemm_shapes((129, 257, 3), (65, 257, 3)), + (129, 65, 257, 3), + ) + + def test_rejects_invalid_input_shapes(self): + cases = ( + (((1, 2), (3, 2, 1)), "a_tensor must have rank 3"), + (((1, 2, 1), (3, 4, 1)), "matching K and L dimensions"), + (((0, 2, 1), (3, 2, 1)), "M=0"), + ) + for args, message in cases: + with self.subTest(args=args): + with self.assertRaisesRegex(ValueError, message): + _MODULE.require_gemm_shapes(*args) + + def test_calculates_and_requires_block_scale_shapes(self): + expected_a = (32, 4, 2, 4, 3, 2) + expected_b = (32, 4, 1, 4, 3, 2) + self.assertEqual(_MODULE.block_scale_shape(129, 257, 2, 32), expected_a) + _MODULE.require_block_scale_shapes( + expected_a, + expected_b, + m=129, + n=65, + k=257, + batch=2, + sf_vec_size=32, + ) + with self.assertRaisesRegex(ValueError, "sfa_tensor must have shape"): + _MODULE.require_block_scale_shapes( + (32, 4, 2, 4, 2, 2), + expected_b, + m=129, + n=65, + k=257, + batch=2, + sf_vec_size=32, + ) + + def test_validates_mma_tiler(self): + self.assertEqual( + _MODULE.require_mma_tiler( + (256, 192), + allowed_m=(128, 256), + allowed_n=(64, 128, 192, 256), + ), + (256, 192), + ) + with self.assertRaisesRegex(ValueError, r"N in \{64, 128, 192, 256\}"): + _MODULE.require_mma_tiler( + (128, 32), + allowed_m=(128, 256), + allowed_n=(64, 128, 192, 256), + ) + + def test_validates_cluster_shape(self): + for cluster_shape in ((1, 1), (2, 4), (4, 4)): + with self.subTest(cluster_shape=cluster_shape): + self.assertEqual( + _MODULE.require_cluster_shape( + cluster_shape, + mma_m=128, + two_cta_mma_m=256, + max_ctas=16, + max_dimension=4, + ), + cluster_shape, + ) + + for cluster_shape, mma_m, message in ( + ((0, 1), 128, "positive powers of two"), + ((3, 1), 128, "positive powers of two"), + ((8, 1), 128, "each at most 4"), + ((4, 8), 128, "product at most 16"), + ((1, 2), 256, "divisible by 2"), + ): + with self.subTest(cluster_shape=cluster_shape, mma_m=mma_m): + with self.assertRaisesRegex(ValueError, message): + _MODULE.require_cluster_shape( + cluster_shape, + mma_m=mma_m, + two_cta_mma_m=256, + max_ctas=16, + max_dimension=4, + ) + + def test_validates_contiguous_alignment_for_scalar_and_packed_types(self): + _MODULE.require_contiguous_alignment("fp8", 16, 8) + _MODULE.require_contiguous_alignment("fp16", 8, 16) + _MODULE.require_contiguous_alignment("packed_fp4", 32, 4) + with self.assertRaisesRegex(ValueError, "16-byte aligned"): + _MODULE.require_contiguous_alignment("fp16", 7, 16) + + def test_requires_complete_mma_rows(self): + _MODULE.require_full_mma_rows(256, 128) + _MODULE.require_full_mma_rows(384, 256, cta_group_size=2) + with self.assertRaisesRegex(ValueError, "TILE_M=256"): + _MODULE.require_full_mma_rows( + 130, + 256, + cta_group_size=2, + reason="the probability load is not predicated", + ) + + def test_resolves_cluster_overlap_margin(self): + self.assertEqual(_MODULE.resolve_max_active_clusters(12, 2), 10) + with self.assertRaisesRegex(ValueError, "CUDNNFE_CLUSTER_OVERLAP_MARGIN"): + _MODULE.resolve_max_active_clusters(2, 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_api_surface.py b/test/python/fe_api/test_jax_api_surface.py new file mode 100644 index 000000000..1fd60e2e9 --- /dev/null +++ b/test/python/fe_api/test_jax_api_surface.py @@ -0,0 +1,613 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free smoke tests for the optional JAX API surface.""" + +from __future__ import annotations + +import ast +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import re +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +def _identity_jit(fn=None, **_kwargs): + return (lambda decorated_fn: decorated_fn) if fn is None else fn + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CUDNN_ROOT = _REPO_ROOT / "python" / "cudnn" +_TEST_PACKAGE = "cudnn_frontend_jax_surface_test" + + +def _colocated_jax_files(): + """Discover operation-local JAX APIs without maintaining an inventory.""" + + facade_dir = _CUDNN_ROOT / "jax" + return tuple(path for path in sorted(_CUDNN_ROOT.rglob("jax.py")) if path.parent != facade_dir) + + +def _operation_path(jax_file): + return ".".join(jax_file.parent.relative_to(_CUDNN_ROOT).parts) + + +def _literal_assignment(path, name): + tree = ast.parse(path.read_text(), filename=str(path)) + value = next( + node.value 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(value) + + +class JaxApiSurfaceTest(unittest.TestCase): + @classmethod + def tearDownClass(cls): + for module_name in tuple(sys.modules): + if module_name == _TEST_PACKAGE or module_name.startswith(f"{_TEST_PACKAGE}."): + sys.modules.pop(module_name, None) + + def _optional_modules(self): + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.dtype = lambda value: value + fake_jnp.bfloat16 = types.SimpleNamespace(name="bfloat16", itemsize=2) + fake_jnp.float32 = types.SimpleNamespace(name="float32", itemsize=4) + 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.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, + ) + + fake_cutlass_jax = types.ModuleType("cutlass.jax") + fake_cutlass_jax.is_available = lambda: True + + class TensorSpec: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + fake_cutlass_jax.TensorSpec = TensorSpec + fake_cutlass_jax.jax_to_cutlass_dtype = lambda dtype: dtype + fake_cutlass = types.ModuleType("cutlass") + fake_cutlass.__path__ = [] + fake_cutlass.Constexpr = object + fake_cutlass_cute = types.ModuleType("cutlass.cute") + fake_cutlass_cute.jit = _identity_jit + fake_cutlass.cute = fake_cutlass_cute + fake_cutlass.jax = fake_cutlass_jax + + return mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "cutlass": fake_cutlass, + "cutlass.cute": fake_cutlass_cute, + "cutlass.jax": fake_cutlass_jax, + "torch": None, + }, + ) + + def test_explicit_namespace_loads_jax_without_torch(self): + parent = types.ModuleType(_TEST_PACKAGE) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = _TEST_PACKAGE + sys.modules[_TEST_PACKAGE] = parent + + torch_before = {name for name, module in sys.modules.items() if name.split(".", 1)[0] == "torch" and module is not None} + with self._optional_modules(): + jax_namespace = importlib.import_module(f"{_TEST_PACKAGE}.jax") + operation_modules = {} + for jax_file in _colocated_jax_files(): + operation_path = _operation_path(jax_file) + module_name = f"{_TEST_PACKAGE}.{operation_path}.jax" + self.assertIn(module_name, sys.modules) + operation_modules[operation_path] = sys.modules[module_name] + self.assertIn(f"{_TEST_PACKAGE}._jax.api_base", sys.modules) + self.assertNotIn(f"{_TEST_PACKAGE}._jax.cutedsl", sys.modules) + + expected_exports = {"ApiBaseJax", "DSA", "JaxTensorDesc", "NSA", "TupleDict"} + expected_dsa_exports = set() + expected_nsa_exports = set() + torch_dsa_exports = set( + _literal_assignment( + _CUDNN_ROOT / "deepseek_sparse_attention" / "__init__.py", + "_SYMBOLS", + ) + ) + torch_nsa_exports = set( + _literal_assignment( + _CUDNN_ROOT / "native_sparse_attention" / "__init__.py", + "_SYMBOLS", + ) + ) + for operation_path, operation_module in operation_modules.items(): + expected_exports.update(operation_module.__all__) + operation_package = importlib.import_module(f"{_TEST_PACKAGE}.{operation_path}") + torch_exports = set(operation_package._API_EXPORTS) + shared_exports = torch_exports.intersection(operation_module.__all__) + for name in shared_exports: + self.assertNotIn(name, vars(operation_package)) + jax_export = getattr(jax_namespace, name) + self.assertIs(getattr(operation_package.jax, name), jax_export) + self.assertTrue(callable(jax_export)) + if isinstance(jax_export, type): + self.assertTrue(issubclass(jax_export, jax_namespace.ApiBaseJax)) + if name in torch_dsa_exports: + expected_dsa_exports.add(name) + if name in torch_nsa_exports: + expected_nsa_exports.add(name) + + loaded_descendants = {name for name in sys.modules if name.startswith(f"{_TEST_PACKAGE}.{operation_path}.")} + self.assertLessEqual( + loaded_descendants, + { + f"{_TEST_PACKAGE}.{operation_path}.config", + f"{_TEST_PACKAGE}.{operation_path}.jax", + f"{_TEST_PACKAGE}.{operation_path}.validation", + }, + ) + + self.assertEqual(set(jax_namespace.__all__), expected_exports) + self.assertEqual(set(vars(jax_namespace.DSA)), expected_dsa_exports) + self.assertEqual(set(vars(jax_namespace.NSA)), expected_nsa_exports) + for name in expected_dsa_exports: + self.assertIs(getattr(jax_namespace.DSA, name), getattr(jax_namespace, name)) + for name in expected_nsa_exports: + self.assertIs(getattr(jax_namespace.NSA, name), getattr(jax_namespace, name)) + + dtype = types.SimpleNamespace(name="bfloat16", itemsize=2) + sample_x = types.SimpleNamespace(shape=(256, 2048), dtype=dtype) + sample_w = types.SimpleNamespace(shape=(2048,), dtype=dtype) + rmsnorm_api = jax_namespace.RmsNormRhtAmaxSm100(sample_x, sample_w) + self.assertFalse(hasattr(rmsnorm_api, "sample_x")) + self.assertFalse(hasattr(rmsnorm_api, "sample_w")) + self.assertTrue(all(value is not sample_x and value is not sample_w for value in vars(rmsnorm_api).values())) + self.assertTrue(rmsnorm_api.check_support()) + self.assertEqual(rmsnorm_api.num_threads, 128) + self.assertEqual(rmsnorm_api.rows_per_cta, 2) + self.assertIs(rmsnorm_api.get_jax_callable(), rmsnorm_api) + + rmsnorm_module = operation_modules["rmsnorm_rht_amax"] + captured = {} + + def fake_call(launcher, inputs, **options): + captured.update(launcher=launcher, inputs=inputs, **options) + return "output", "amax" + + with mock.patch.object(rmsnorm_module, "call_cutedsl", side_effect=fake_call): + result = rmsnorm_api(sample_x, sample_w) + self.assertIsInstance(result, jax_namespace.TupleDict) + self.assertEqual(tuple(result.keys()), ("output", "amax")) + self.assertEqual(tuple(result), ("output", "amax")) + self.assertIs(captured["launcher"], rmsnorm_module._launch) + self.assertEqual(captured["inputs"], (sample_x, sample_w)) + self.assertEqual( + captured["static_args"], + {"n": 2048, "num_threads": 128, "rows_per_cta": 2, "eps": 1e-5}, + ) + self.assertEqual( + [(spec.name, spec.shape) for spec in captured["outputs"]], + [("output", (256, 2048)), ("amax", (128,))], + ) + + torch_after = {name for name, module in sys.modules.items() if name.split(".", 1)[0] == "torch" and module is not None} + self.assertEqual(torch_after - torch_before, set()) + + def test_colocated_jax_modules_can_load_before_facade(self): + package_name = f"{_TEST_PACKAGE}_direct_first" + parent = types.ModuleType(package_name) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = package_name + sys.modules[package_name] = parent + + try: + with self._optional_modules(): + operation_modules = [] + for jax_file in _colocated_jax_files(): + module = importlib.import_module(f"{package_name}.{_operation_path(jax_file)}.jax") + operation_modules.append(module) + self.assertNotIn(f"{package_name}.jax", sys.modules) + + facade = importlib.import_module(f"{package_name}.jax") + + for operation_module in operation_modules: + for name in operation_module.__all__: + self.assertIs(getattr(facade, name), getattr(operation_module, name)) + finally: + 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_missing_jax_reports_optional_extra(self): + package_name = f"{_TEST_PACKAGE}_missing_jax" + parent = types.ModuleType(package_name) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = package_name + sys.modules[package_name] = parent + + try: + with mock.patch.dict(sys.modules, {"jax": None, "jax.numpy": None}): + with self.assertRaisesRegex( + ImportError, + r"requires the 'jax' module.*nvidia-cudnn-frontend\[jax\]", + ): + importlib.import_module(f"{package_name}.jax") + finally: + 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_missing_cutedsl_reports_optional_extra(self): + package_name = f"{_TEST_PACKAGE}_no_cutlass" + parent = types.ModuleType(package_name) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = package_name + sys.modules[package_name] = parent + fake_jax = types.ModuleType("jax") + fake_jnp = types.ModuleType("jax.numpy") + fake_jax.__path__ = [] + fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + fake_jax.numpy = fake_jnp + fake_jax.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, + ) + + try: + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "cutlass": None, + "cutlass.jax": None, + }, + ): + with self.assertRaisesRegex( + ImportError, + r"requires the 'cutlass' module.*nvidia-cudnn-frontend\[jax\]", + ): + importlib.import_module(f"{package_name}.jax") + finally: + 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_unavailable_cutlass_jax_reports_minimum_version(self): + package_name = f"{_TEST_PACKAGE}_unavailable_cutlass_jax" + parent = types.ModuleType(package_name) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = package_name + sys.modules[package_name] = parent + + fake_jax = types.ModuleType("jax") + fake_jax.__version__ = "0.4.0" + fake_jax.version = types.SimpleNamespace(__version_info__=(0, 4, 0)) + fake_cutlass = types.ModuleType("cutlass") + fake_cutlass.__path__ = [] + fake_cutlass_jax = types.ModuleType("cutlass.jax") + fake_cutlass_jax.CUTE_DSL_MIN_SUPPORTED_JAX_VERSION = (0, 5, 0) + fake_cutlass_jax.is_available = lambda: False + fake_cutlass.jax = fake_cutlass_jax + + try: + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "cutlass": fake_cutlass, + "cutlass.jax": fake_cutlass_jax, + }, + ): + with self.assertRaisesRegex( + ImportError, + r"CUTLASS JAX support is unavailable with JAX 0\.4\.0; " r"the minimum supported JAX version is 0\.5\.0.*nvidia-cudnn-frontend\[jax\]", + ): + importlib.import_module(f"{package_name}.jax") + finally: + 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_operation_packages_use_explicit_jax_namespaces(self): + operations = [] + for jax_file in _colocated_jax_files(): + operation_path = _operation_path(jax_file) + torch_exports = set(_literal_assignment(jax_file.parent / "__init__.py", "_API_EXPORTS")) + jax_exports = set(_literal_assignment(jax_file, "__all__")) + operations.append( + ( + operation_path, + sorted(torch_exports.intersection(jax_exports)), + sorted(jax_exports.difference(torch_exports)), + ) + ) + + package_name = f"{_TEST_PACKAGE}_explicit_jax" + parent = types.ModuleType(package_name) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = package_name + sys.modules[package_name] = parent + + operation_apis = {} + for operation_path, shared_names, jax_only_names in operations: + self.assertTrue(shared_names, f"{operation_path} has no aligned Torch/JAX symbol") + operation_name = f"{package_name}.{operation_path}" + torch_api = types.ModuleType(f"{operation_name}.api") + jax_api = types.ModuleType(f"{operation_name}.jax") + torch_values = {name: object() for name in shared_names} + jax_values = {name: object() for name in (*shared_names, *jax_only_names)} + for name, value in torch_values.items(): + setattr(torch_api, name, value) + for name, value in jax_values.items(): + setattr(jax_api, name, value) + operation_apis[operation_path] = ( + shared_names, + jax_only_names, + torch_values, + jax_values, + torch_api, + jax_api, + ) + + try: + with mock.patch( + "importlib.util.find_spec", + side_effect=AssertionError("operation packages must not probe frameworks"), + ): + for operation_path, _, _ in operations: + with self.subTest(operation=operation_path): + operation_name = f"{package_name}.{operation_path}" + operation = importlib.import_module(operation_name) + ( + shared_names, + jax_only_names, + torch_values, + jax_values, + torch_api, + jax_api, + ) = operation_apis[operation_path] + self.assertNotIn(torch_api.__name__, sys.modules) + self.assertNotIn(jax_api.__name__, sys.modules) + self.assertEqual(operation.__all__, list(operation._API_EXPORTS)) + + with mock.patch.dict( + sys.modules, + { + torch_api.__name__: torch_api, + jax_api.__name__: jax_api, + }, + ): + for name in shared_names: + self.assertIs(getattr(operation, name), torch_values[name]) + self.assertIn(name, vars(operation)) + for name in jax_only_names: + with self.assertRaises(AttributeError): + getattr(operation, name) + + self.assertIs(operation.api, torch_api) + self.assertIs(operation.jax, jax_api) + for name, value in jax_values.items(): + self.assertIs(getattr(operation.jax, name), value) + finally: + 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_unqualified_api_does_not_fall_back_to_jax(self): + package_name = f"{_TEST_PACKAGE}_jax_only" + operation_name = f"{package_name}.rmsnorm_rht_amax" + parent = types.ModuleType(package_name) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = package_name + sys.modules[package_name] = parent + + jax_api = types.ModuleType(f"{operation_name}.jax") + jax_api.rmsnorm_rht_amax_sm100 = object() + + try: + with mock.patch.dict( + sys.modules, + { + f"{operation_name}.api": None, + jax_api.__name__: jax_api, + }, + ): + operation = importlib.import_module(operation_name) + self.assertIs(operation.jax, jax_api) + with self.assertRaises(ModuleNotFoundError): + operation.rmsnorm_rht_amax_sm100 + finally: + 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_jax_implementations_are_colocated_with_torch_apis(self): + jax_files = _colocated_jax_files() + self.assertTrue(jax_files) + for jax_file in jax_files: + operation_dir = jax_file.parent + with self.subTest(operation=operation_dir.name): + self.assertTrue((operation_dir / "api.py").is_file()) + self.assertTrue((operation_dir / "__init__.py").is_file()) + + self.assertEqual( + {path.name for path in (_CUDNN_ROOT / "jax").glob("*.py")}, + {"__init__.py"}, + ) + + def test_cutedsl_calls_use_stable_module_launchers(self): + for jax_file in _colocated_jax_files(): + tree = ast.parse(jax_file.read_text(), filename=str(jax_file)) + module_functions = {node.name: node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))} + parents = {child: parent for parent in ast.walk(tree) for child in ast.iter_child_nodes(parent)} + + def resolve_launchers(value, assignments, seen=frozenset()): + if isinstance(value, ast.Name): + if value.id in module_functions and value.id.startswith("_launch"): + return {value.id} + if value.id in seen: + return set() + resolved_assignments = [ + resolve_launchers(assigned_value, assignments, seen | {value.id}) for assigned_value in assignments.get(value.id, ()) + ] + if not resolved_assignments or any(not resolved for resolved in resolved_assignments): + return set() + return set().union(*resolved_assignments) + if isinstance(value, ast.IfExp): + body_launchers = resolve_launchers(value.body, assignments, seen) + else_launchers = resolve_launchers(value.orelse, assignments, seen) + if not body_launchers or not else_launchers: + return set() + return body_launchers | else_launchers + return set() + + def same_scope_nodes(scope): + pending = list(scope.body) + while pending: + node = pending.pop() + yield node + pending.extend( + child + for child in ast.iter_child_nodes(node) + if not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)) + ) + + launcher_factories = { + node.name + for node in ast.walk(tree) + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and re.fullmatch(r"_make_?.*launcher", node.name) + } + cutedsl_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" + ] + if cutedsl_calls: + self.assertFalse(launcher_factories, f"{jax_file} must use stable module-level launchers") + for call in cutedsl_calls: + with self.subTest(operation=jax_file.parent.name, line=call.lineno): + self.assertTrue(call.args, "call_cutedsl requires a launcher") + scope = parents.get(call) + while scope is not None and not isinstance(scope, (ast.FunctionDef, ast.AsyncFunctionDef)): + scope = parents.get(scope) + self.assertIsNotNone(scope) + + assignments = {} + for node in same_scope_nodes(scope): + if not isinstance(node, ast.Assign): + continue + for target in node.targets: + if isinstance(target, ast.Name): + assignments.setdefault(target.id, []).append(node.value) + + launchers = resolve_launchers(call.args[0], assignments) + self.assertTrue(launchers, "launcher must resolve to a stable module-level function") + + static_keyword = next((keyword for keyword in call.keywords if keyword.arg == "static_args"), None) + expected_by_launcher = {launcher: {argument.arg for argument in module_functions[launcher].args.kwonlyargs} for launcher in launchers} + if not any(expected_by_launcher.values()) and static_keyword is None: + continue + self.assertIsNotNone(static_keyword) + self.assertIsInstance(static_keyword.value, ast.Dict) + static_names = {key.value for key in static_keyword.value.keys if isinstance(key, ast.Constant) and isinstance(key.value, str)} + self.assertEqual(len(static_names), len(static_keyword.value.keys)) + for launcher, expected_names in expected_by_launcher.items(): + self.assertEqual(static_names, expected_names) + + def test_torch_function_remains_a_lazy_top_level_export(self): + top_level_tree = ast.parse( + (_CUDNN_ROOT / "__init__.py").read_text(), + filename=str(_CUDNN_ROOT / "__init__.py"), + ) + lazy_imports_node = next( + node.value + for node in top_level_tree.body + if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == "_LAZY_OPTIONAL_IMPORTS" for target in node.targets) + ) + lazy_imports = ast.literal_eval(lazy_imports_node) + self.assertEqual( + lazy_imports["rmsnorm_rht_amax_sm100"], + (".rmsnorm_rht_amax", "rmsnorm_rht_amax_sm100"), + ) + self.assertEqual(lazy_imports["DSA"], (".deepseek_sparse_attention", "DSA")) + + def test_root_lazily_imports_explicit_jax_namespace(self): + top_level_tree = ast.parse( + (_CUDNN_ROOT / "__init__.py").read_text(), + filename=str(_CUDNN_ROOT / "__init__.py"), + ) + getattr_node = next(node for node in top_level_tree.body if isinstance(node, ast.FunctionDef) and node.name == "__getattr__") + jax_branch = next( + node + for node in getattr_node.body + if isinstance(node, ast.If) + and isinstance(node.test, ast.Compare) + and isinstance(node.test.left, ast.Name) + and node.test.left.id == "name" + and len(node.test.comparators) == 1 + and isinstance(node.test.comparators[0], ast.Constant) + and node.test.comparators[0].value == "jax" + ) + imported_modules = { + call.args[0].value + for call in ast.walk(jax_branch) + if isinstance(call, ast.Call) + and isinstance(call.func, ast.Attribute) + and call.func.attr == "import_module" + and call.args + and isinstance(call.args[0], ast.Constant) + } + self.assertEqual(imported_modules, {".jax"}) + + def test_dsa_operator_names_match_torch_namespace(self): + dsa_tree = ast.parse( + (_CUDNN_ROOT / "deepseek_sparse_attention" / "__init__.py").read_text(), + filename=str(_CUDNN_ROOT / "deepseek_sparse_attention" / "__init__.py"), + ) + symbols_node = next( + node.value + for node in dsa_tree.body + if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == "_SYMBOLS" for target in node.targets) + ) + torch_symbols = ast.literal_eval(symbols_node) + facade_tree = ast.parse( + (_CUDNN_ROOT / "jax" / "__init__.py").read_text(), + filename=str(_CUDNN_ROOT / "jax" / "__init__.py"), + ) + dsa_assignment = next( + node + for node in facade_tree.body + if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == "DSA" for target in node.targets) + ) + jax_dsa_symbols = {keyword.arg for keyword in dsa_assignment.value.keywords} + self.assertTrue(jax_dsa_symbols) + self.assertLessEqual(jax_dsa_symbols, set(torch_symbols)) + + def test_jax_optional_extra_does_not_install_torch(self): + pyproject = (_REPO_ROOT / "pyproject.toml").read_text() + optional_dependencies = re.search( + r"(?ms)^\[project\.optional-dependencies\]\n(?P.*?)(?=^\[|\Z)", + pyproject, + ) + self.assertIsNotNone(optional_dependencies) + body = optional_dependencies.group("body") + jax_dependencies = re.search(r"(?ms)^jax = \[\n(?P.*?)^\]$", body) + self.assertIsNotNone(jax_dependencies) + self.assertNotIn('"torch"', jax_dependencies.group("body")) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_cutedsl_adapter.py b/test/python/fe_api/test_jax_cutedsl_adapter.py new file mode 100644 index 000000000..98d84484b --- /dev/null +++ b/test/python/fe_api/test_jax_cutedsl_adapter.py @@ -0,0 +1,395 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""CPU-only contract tests for the CuTe DSL JAX tracing adapter.""" + +from __future__ import annotations + +from dataclasses import dataclass +import importlib.util +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + # Keep this contract test runnable with the standard library alone. + pass +else: + pytestmark = pytest.mark.L0 + + +def _identity_jit(fn=None, **kwargs): + def decorate(decorated_fn): + decorated_fn._cute_jit_options = kwargs + return decorated_fn + + return decorate if fn is None else decorate(fn) + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_TEST_PACKAGE = "cudnn_frontend_jax_api_base_adapter_test" +_PARENT = types.ModuleType(_TEST_PACKAGE) +_PARENT.__path__ = [str(_CUDNN_ROOT)] +_JAX_PACKAGE = types.ModuleType(f"{_TEST_PACKAGE}._jax") +_JAX_PACKAGE.__path__ = [str(_CUDNN_ROOT / "_jax")] + +_MODULE_PATH = _CUDNN_ROOT / "_jax" / "api_base.py" +_SPEC = importlib.util.spec_from_file_location(f"{_TEST_PACKAGE}._jax.api_base", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_MODULE = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _MODULE + +_bootstrap_jnp = types.ModuleType("jax.numpy") +_bootstrap_jnp.dtype = lambda value: value +_bootstrap_jax = types.ModuleType("jax") +_bootstrap_jax.__path__ = [] +_bootstrap_jax.numpy = _bootstrap_jnp +_bootstrap_jax.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, +) +_bootstrap_cutlass_jax = types.ModuleType("cutlass.jax") +_bootstrap_cutlass_jax.TensorSpec = type("TensorSpec", (), {}) +_bootstrap_cute = types.ModuleType("cutlass.cute") +_bootstrap_cute.jit = _identity_jit +_bootstrap_cutlass = types.ModuleType("cutlass") +_bootstrap_cutlass.__path__ = [] +_bootstrap_cutlass.Constexpr = object +_bootstrap_cutlass.cute = _bootstrap_cute +_bootstrap_cutlass.jax = _bootstrap_cutlass_jax +with mock.patch.dict( + sys.modules, + { + "jax": _bootstrap_jax, + "jax.numpy": _bootstrap_jnp, + "cutlass": _bootstrap_cutlass, + "cutlass.cute": _bootstrap_cute, + "cutlass.jax": _bootstrap_cutlass_jax, + _TEST_PACKAGE: _PARENT, + f"{_TEST_PACKAGE}._jax": _JAX_PACKAGE, + }, +): + _SPEC.loader.exec_module(_MODULE) + +BufferSpec = _MODULE.BufferSpec +call_cutedsl = _MODULE.call_cutedsl + + +class _Array: + def __init__(self, shape, dtype, label): + self.shape = tuple(shape) + self.dtype = dtype + self.label = label + + def __repr__(self): + return f"_Array({self.label!r}, {self.shape!r}, {self.dtype!r})" + + +class _ShapeDtypeStruct: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + +class _FakeJax: + ShapeDtypeStruct = _ShapeDtypeStruct + + +class _FakeJaxNumpy: + def __init__(self): + self.allocations = [] + + def full(self, shape, value, dtype): + result = _Array(shape, dtype, f"full({value})") + self.allocations.append(("full", result, value)) + return result + + +@dataclass(frozen=True) +class _FakeCutlassTensorSpec: + """Dependency-free stand-in with CUTLASS 4.5's TensorSpec fields.""" + + layout: object = None + mode: object = None + static: object = None + ptr_assumed_align: object = 256 + divisibility: object = None + + +class _FakeCutlassJax: + TensorSpec = _FakeCutlassTensorSpec + + def __init__(self): + self.calls = [] + + def cutlass_call(self, fn, **options): + self.calls.append((fn, options)) + known_options = { + "output_shape_dtype", + "input_spec", + "output_spec", + "input_output_aliases", + "allow_cuda_graph", + "compile_options", + "use_static_tensors", + } + static_args = {k: v for k, v in options.items() if k not in known_options} + + def invoke(*inputs): + aliases = options["input_output_aliases"] + alias_by_output = {output_idx: input_idx for input_idx, output_idx in aliases.items()} + results = [] + fresh_results = [] + for result_idx, metadata in enumerate(options["output_shape_dtype"]): + if result_idx in alias_by_output: + result = inputs[alias_by_output[result_idx]] + else: + result = _Array( + metadata.shape, + metadata.dtype, + f"result-{result_idx}", + ) + fresh_results.append(result) + results.append(result) + + fn("xla-stream", *inputs, *fresh_results, **static_args) + return tuple(results) + + return invoke + + +class CallCutedslAdapterTest(unittest.TestCase): + def _call(self, kernel, inputs, **kwargs): + fake_jnp = _FakeJaxNumpy() + fake_cutlass_jax = _FakeCutlassJax() + with ( + mock.patch.object(_MODULE, "jax", _FakeJax), + mock.patch.object(_MODULE, "jnp", fake_jnp), + mock.patch.object(_MODULE, "cutlass_jax", fake_cutlass_jax), + ): + result = call_cutedsl(kernel, inputs, **kwargs) + return result, fake_jnp, fake_cutlass_jax + + def test_launch_adapter_uses_compile_time_tracing(self): + self.assertEqual(_MODULE._launch_adapter._cute_jit_options, {"preprocess": False}) + + def test_workspace_is_passed_to_launcher_but_hidden_from_results(self): + seen = [] + + def kernel(stream, x, output, workspace, *, scale): + seen.append((stream, x, output, workspace, scale)) + + x = _Array((8,), "f32", "x") + result, _, bridge = self._call( + kernel, + (x,), + outputs=(BufferSpec("output", (8,), "f32"),), + workspaces=(BufferSpec("workspace", (128,), "u8"),), + static_args={"scale": 2.0}, + ) + + self.assertEqual(len(result), 1) + self.assertEqual(result[0].label, "result-0") + self.assertEqual(seen[0][0], "xla-stream") + self.assertIs(seen[0][1], x) + self.assertEqual(seen[0][2].label, "result-0") + self.assertEqual(seen[0][3].label, "result-1") + self.assertEqual(seen[0][4], 2.0) + self.assertEqual(len(bridge.calls[0][1]["output_shape_dtype"]), 2) + + def test_initialized_buffers_are_inputs_aliased_to_results(self): + seen = [] + + def kernel(stream, x, output, workspace): + seen.append((output, workspace)) + + x = _Array((4,), "bf16", "x") + output_tensor_spec = _FakeCutlassTensorSpec(divisibility=(4,)) + workspace_tensor_spec = _FakeCutlassTensorSpec(ptr_assumed_align=128) + result, fake_jnp, bridge = self._call( + kernel, + (x,), + outputs=( + BufferSpec( + "output", + (4,), + "bf16", + tensor_spec=output_tensor_spec, + fill_value=float("-inf"), + ), + ), + workspaces=( + BufferSpec( + "workspace", + (32,), + "u8", + tensor_spec=workspace_tensor_spec, + fill_value=0, + ), + ), + ) + + aliases = bridge.calls[0][1]["input_output_aliases"] + self.assertEqual(aliases, {1: 0, 2: 1}) + input_specs = bridge.calls[0][1]["input_spec"] + output_specs = bridge.calls[0][1]["output_spec"] + self.assertIsNone(input_specs[0]) + self.assertIs(input_specs[1], output_tensor_spec) + self.assertIs(input_specs[2], workspace_tensor_spec) + self.assertIs(output_specs[0], output_tensor_spec) + self.assertIs(output_specs[1], workspace_tensor_spec) + self.assertIs(result[0], fake_jnp.allocations[0][1]) + self.assertIs(seen[0][0], fake_jnp.allocations[0][1]) + self.assertIs(seen[0][1], fake_jnp.allocations[1][1]) + self.assertEqual(fake_jnp.allocations[0][0], "full") + self.assertEqual(fake_jnp.allocations[1][0], "full") + self.assertEqual(fake_jnp.allocations[1][2], 0) + + def test_native_tensor_specs_are_forwarded_without_conversion(self): + input_spec = _FakeCutlassTensorSpec( + layout=(1, 0), + mode=(1, 0), + static=True, + ptr_assumed_align=128, + # CUTLASS accepts -1 divisibility sentinels. The adapter must not + # reinterpret or narrow the native TensorSpec contract. + divisibility=(-1, -1), + ) + output_spec = _FakeCutlassTensorSpec( + layout=(1, 0), + mode=(1, 0), + static=True, + ptr_assumed_align=128, + divisibility=(16, 8), + ) + + result, _, bridge = self._call( + lambda stream, x, output: None, + (_Array((8, 16), "f32", "x"),), + outputs=( + BufferSpec( + "output", + (8, 16), + "f32", + tensor_spec=output_spec, + ), + ), + input_specs=(input_spec,), + allow_cuda_graph=False, + compile_options="--example-option", + ) + + self.assertEqual(len(result), 1) + options = bridge.calls[0][1] + self.assertIs(options["input_spec"][0], input_spec) + self.assertIs(options["output_spec"][0], output_spec) + self.assertFalse(options["allow_cuda_graph"]) + self.assertEqual(options["compile_options"], "--example-option") + + def test_default_result_spec_uses_cutlass_inference(self): + _, _, bridge = self._call( + lambda stream, x, output: None, + (_Array((8,), "f32", "x"),), + outputs=(BufferSpec("output", (8,), "f32"),), + ) + + self.assertIsNone(bridge.calls[0][1]["output_spec"][0]) + + def test_static_tensors_default_to_true(self): + _, _, bridge = self._call( + lambda stream, x, output: None, + (_Array((8,), "f32", "x"),), + outputs=(BufferSpec("output", (8,), "f32"),), + ) + + self.assertTrue(bridge.calls[0][1]["use_static_tensors"]) + + def test_static_tensors_can_be_disabled(self): + _, _, bridge = self._call( + lambda stream, x, output: None, + (_Array((8,), "f32", "x"),), + outputs=(BufferSpec("output", (8,), "f32"),), + use_static_tensors=False, + ) + + self.assertFalse(bridge.calls[0][1]["use_static_tensors"]) + + def test_rejects_invalid_call_plans(self): + with self.assertRaisesRegex(ValueError, "at least one public output"): + self._call(lambda stream, x: None, (_Array((1,), "f32", "x"),), outputs=()) + + with self.assertRaisesRegex(ValueError, "must be unique"): + self._call( + lambda stream, x, y, z: None, + (_Array((1,), "f32", "x"),), + outputs=(BufferSpec("same", (1,), "f32"),), + workspaces=(BufferSpec("same", (1,), "u8"),), + ) + + with self.assertRaisesRegex(ValueError, "Expected 1 input tensor spec"): + self._call( + lambda stream, x, y: None, + (_Array((1,), "f32", "x"),), + outputs=(BufferSpec("output", (1,), "f32"),), + input_specs=(), + ) + + with self.assertRaisesRegex(TypeError, "flat sequence"): + self._call( + lambda stream, x, y: None, + ((_Array((1,), "f32", "nested"),),), + outputs=(BufferSpec("output", (1,), "f32"),), + ) + + def test_kernel_static_args_do_not_collide_with_cutlass_call_options(self): + seen = [] + + def kernel(stream, x, output, *, compile_options): + seen.append((stream, x, output, compile_options)) + + x = _Array((1,), "f32", "x") + result, _, bridge = self._call( + kernel, + (x,), + outputs=(BufferSpec("output", (1,), "f32"),), + static_args={"compile_options": "kernel option"}, + ) + + self.assertEqual(len(result), 1) + self.assertEqual(seen[0][3], "kernel option") + self.assertIsNone(bridge.calls[0][1]["compile_options"]) + self.assertIs(bridge.calls[0][0], _MODULE._launch_adapter) + self.assertIsInstance(hash(bridge.calls[0][1]["config"]), int) + + def test_static_launcher_config_has_a_value_stable_cache_key(self): + def kernel(stream, x, output, *, scale): + del stream, x, output, scale + + x = _Array((1,), "f32", "x") + options = { + "outputs": (BufferSpec("output", (1,), "f32"),), + "static_args": {"scale": 2.0}, + } + _, _, first_bridge = self._call(kernel, (x,), **options) + _, _, second_bridge = self._call(kernel, (x,), **options) + _, _, different_bridge = self._call( + kernel, + (x,), + outputs=options["outputs"], + static_args={"scale": 3.0}, + ) + + first_config = first_bridge.calls[0][1]["config"] + second_config = second_bridge.calls[0][1]["config"] + different_config = different_bridge.calls[0][1]["config"] + self.assertEqual(first_config, second_config) + self.assertEqual(hash(first_config), hash(second_config)) + self.assertNotEqual(first_config, different_config) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_dsa.py b/test/python/fe_api/test_jax_dsa.py new file mode 100644 index 000000000..73a4716f6 --- /dev/null +++ b/test/python/fe_api/test_jax_dsa.py @@ -0,0 +1,759 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX integration tests for the DSA kernels.""" + +import pytest + + +def _jax_dependencies(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + pytest.importorskip("cutlass.jax") + return jax, jnp + + +def _gpu_device(jax, minimum_compute_capability): + gpu_devices = [device for device in jax.local_devices() if device.platform == "gpu"] + for device in gpu_devices: + capability = getattr(device, "compute_capability", None) + if capability is None: + continue + if isinstance(capability, (tuple, list)): + major, minor = (int(value) for value in capability[:2]) + capability_number = major * 10 + minor + elif "." in str(capability): + major, minor = (int(value) for value in str(capability).split(".", 1)) + capability_number = major * 10 + minor + else: + capability_number = int(capability) + if capability_number < 10: + capability_number *= 10 + if capability_number >= minimum_compute_capability: + return device + pytest.skip(f"JAX SM{minimum_compute_capability}+ device is not available") + + +@pytest.mark.L0 +def test_jax_dsa_abstract_shapes(): + jax, jnp = _jax_dependencies() + + from cudnn.jax import ( + TupleDict, + dense_indexer_backward_wrapper, + indexer_forward_wrapper, + indexer_top_k_wrapper, + sparse_attn_score_recompute_wrapper, + sparse_indexer_score_recompute_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) + (scores,) = jax.eval_shape( + lambda q, k, w: indexer_forward_wrapper(q, k, w, ratio=1), + q, + k, + w, + ) + assert scores.shape == (1, 4, 5) + assert scores.dtype == jnp.float32 + + input_values = jax.ShapeDtypeStruct((2, 64), jnp.float32) + seq_lens = jax.ShapeDtypeStruct((2,), jnp.int32) + topk_result = jax.eval_shape( + lambda values, lengths: indexer_top_k_wrapper( + values, + lengths, + top_k=8, + ), + input_values, + seq_lens, + ) + assert isinstance(topk_result, TupleDict) + assert tuple(topk_result.keys()) == ("indices", "values") + indices, values = topk_result + assert indices.shape == (2, 8) + assert indices.dtype == jnp.int32 + assert values.shape == (2, 8) + assert values.dtype == jnp.float32 + + sparse_q = jax.ShapeDtypeStruct((1, 4, 32, 128), jnp.bfloat16) + sparse_k = jax.ShapeDtypeStruct((1, 128, 128), jnp.bfloat16) + sparse_topk = jax.ShapeDtypeStruct((1, 4, 128), jnp.int32) + sparse_weights = jax.ShapeDtypeStruct((1, 4, 32), jnp.bfloat16) + predict = jax.eval_shape( + sparse_indexer_score_recompute_wrapper, + sparse_q, + sparse_k, + sparse_weights, + sparse_topk, + )["predict"] + assert predict.shape == (1, 4, 128) + assert predict.dtype == jnp.float32 + + sparse_lse = jax.ShapeDtypeStruct((1, 4, 32), jnp.float32) + target = jax.eval_shape( + lambda q, k, lse, indices: sparse_attn_score_recompute_wrapper( + q, + k, + lse, + indices, + softmax_scale=0.125, + ), + sparse_q, + sparse_k, + sparse_lse, + sparse_topk, + )["target"] + assert target.shape == (1, 4, 128) + assert target.dtype == jnp.float32 + + dense_q = jax.ShapeDtypeStruct((1, 2, 64, 128), jnp.bfloat16) + dense_weights = jax.ShapeDtypeStruct((1, 2, 64), jnp.bfloat16) + dense_k = jax.ShapeDtypeStruct((1, 128, 128), jnp.bfloat16) + dense_score = jax.ShapeDtypeStruct((1, 2, 128), jnp.float32) + dense_denom = jax.ShapeDtypeStruct((1, 2), jnp.float32) + dense_result = jax.eval_shape( + dense_indexer_backward_wrapper, + dense_q, + dense_weights, + dense_k, + dense_score, + dense_denom, + dense_score, + dense_denom, + ) + assert isinstance(dense_result, TupleDict) + assert dense_result["d_index_q"].shape == dense_q.shape + assert dense_result["d_index_q"].dtype == jnp.bfloat16 + assert dense_result["d_weights"].shape == dense_weights.shape + assert dense_result["d_weights"].dtype == jnp.bfloat16 + assert dense_result["d_index_k"].shape == dense_k.shape + assert dense_result["d_index_k"].dtype == jnp.bfloat16 + + +@pytest.mark.L0 +def test_jax_indexer_forward_jit(): + jax, jnp = _jax_dependencies() + device = _gpu_device(jax, 100) + + from cudnn.jax import indexer_forward_wrapper + + batch, seqlen_q, seqlen_k = 1, 4, 5 + num_query_heads, num_kv_heads, head_dim = 32, 1, 128 + sm_scale = 0.5 + 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, num_kv_heads, 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, + ) + + @jax.jit + def run(q, k, w): + return indexer_forward_wrapper( + q, + k, + w, + ratio=1, + sm_scale=sm_scale, + )["scores"] + + lowered = run.lower(q, k, w) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + compiled = lowered.compile() + q_second = -q + scores_first = compiled(q, k, w) + scores_second = compiled(q_second, k, w) + scores_second.block_until_ready() + + k_f32 = jnp.repeat(k.astype(jnp.float32), num_query_heads, axis=2) + valid = jnp.arange(seqlen_k)[None, :] < (seqlen_k - seqlen_q + jnp.arange(seqlen_q)[:, None] + 1) + + def reference(q_value): + dots = jnp.einsum("bqhd,bkhd->bqkh", q_value.astype(jnp.float32), k_f32) + result = sm_scale * jnp.sum( + jnp.maximum(dots, 0.0) * w.astype(jnp.float32)[:, :, None, :], + axis=-1, + ) + return jnp.where(valid[None, :, :], result, -jnp.inf) + + for scores, expected in ( + (scores_first, reference(q)), + (scores_second, reference(q_second)), + ): + assert scores.shape == (batch, seqlen_q, seqlen_k) + 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) + + +@pytest.mark.L0 +def test_jax_indexer_top_k_jit_uses_hidden_workspace(): + jax, jnp = _jax_dependencies() + device = _gpu_device(jax, 90) + + from cudnn.jax import indexer_top_k_wrapper + + # More than 148 rows selects the large-occupancy configuration, whose + # FP32 candidate capacity is 256. These 257 unique values all round to + # the same FP16 coarse key, so every full-length row deterministically + # spills one candidate to global scratch and reads it during refinement. + num_rows, num_cols, top_k = 149, 257, 8 + row_values = 1.0 + jnp.arange(num_cols, dtype=jnp.float32) * 1.0e-6 + input_values = jax.device_put( + jnp.broadcast_to(row_values, (num_rows, num_cols)), + device, + ) + seq_lens = jax.device_put( + jnp.full((num_rows,), num_cols, dtype=jnp.int32).at[0].set(num_cols - 1), + device, + ) + + @jax.jit + def run(values, lengths): + return indexer_top_k_wrapper(values, lengths, top_k=top_k) + + lowered = run.lower(input_values, seq_lens) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + # Float32 top-K uses two int32 scratch planes. The custom call returns this + # allocation internally even though the public JAX result has only two leaves. + assert f"{num_rows}x2x{num_cols}xi32" in stablehlo + + indices, values = lowered.compile()(input_values, seq_lens) + indices.block_until_ready() + valid = jnp.arange(num_cols)[None, :] < seq_lens[:, None] + masked_values = jnp.where(valid, input_values, -jnp.inf) + reference_values, reference_indices = jax.lax.top_k(masked_values, top_k) + + assert indices.dtype == jnp.int32 + assert values.dtype == jnp.float32 + assert jnp.array_equal(jnp.sort(indices, axis=-1), jnp.sort(reference_indices, axis=-1)) + assert jnp.allclose( + jnp.sort(values, axis=-1), + jnp.sort(reference_values, axis=-1), + atol=0.0, + rtol=0.0, + ) + + +@pytest.mark.L0 +def test_jax_sparse_indexer_score_recompute_jit(): + jax, jnp = _jax_dependencies() + device = _gpu_device(jax, 100) + + from cudnn.jax import sparse_indexer_score_recompute_wrapper + + 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, + ) + topk_indices = jax.device_put( + jnp.broadcast_to( + jnp.arange(seqlen_k - 1, -1, -1, dtype=jnp.int32), + (batch, seqlen_q, topk), + ), + device, + ) + + @jax.jit + def run(q, k, weights, indices): + return sparse_indexer_score_recompute_wrapper( + q, + k, + weights, + indices, + )["predict"] + + lowered = run.lower(q, k, weights, topk_indices) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + predict = lowered.compile()(q, k, weights, topk_indices) + predict.block_until_ready() + + k_gather = jax.vmap(lambda batch_k, batch_indices: batch_k[batch_indices])( + k.astype(jnp.float32), + topk_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 predict.shape == (batch, seqlen_q, topk) + assert predict.dtype == jnp.float32 + 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_dependencies() + device = _gpu_device(jax, 100) + + from cudnn.jax import sparse_attn_score_recompute_wrapper + + 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, + ) + topk_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)) + + @jax.jit + def run(q, k, lse, indices): + return sparse_attn_score_recompute_wrapper( + q, + k, + lse, + indices, + softmax_scale=softmax_scale, + )["target"] + + lowered = run.lower(q, k, lse, topk_indices) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + target = lowered.compile()(q, k, lse, topk_indices) + target.block_until_ready() + + k_gather = jax.vmap(lambda batch_k, batch_indices: batch_k[batch_indices])( + k.astype(jnp.float32), + topk_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 target.shape == (batch, seqlen_q, topk) + assert target.dtype == jnp.float32 + 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_dependencies() + device = _gpu_device(jax, 100) + + 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, + )["predict"] + target = sparse_attn_score_recompute_wrapper( + q, + k, + lse, + indices, + softmax_scale, + topk_length=lengths, + topk_indices_global=True, + )["target"] + return predict, target + + lowered = run.lower(q, k, weights, lse, global_indices, topk_length) + stablehlo = lowered.as_text("stablehlo") + assert 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:])) + + +@pytest.mark.L0 +def test_jax_dense_indexer_backward_jit(): + jax, jnp = _jax_dependencies() + device = _gpu_device(jax, 100) + + from cudnn.jax import dense_indexer_backward_wrapper + + batch, seqlen_q, seqlen_k = 1, 256, 512 + 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( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + grad_loss, + ): + return dense_indexer_backward_wrapper( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + grad_loss=grad_loss, + sm_scale=sm_scale, + ) + + grad_loss = jax.device_put(jnp.ones((1,), dtype=jnp.float32), device) + lowered = run.lower( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + grad_loss, + ) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + compiled = lowered.compile() + result = compiled( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + grad_loss, + ) + 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 result["d_index_q"].dtype == jnp.bfloat16 + assert result["d_weights"].dtype == jnp.bfloat16 + assert result["d_index_k"].dtype == jnp.bfloat16 + assert all(jnp.all(jnp.isfinite(value)) for value in result.values()) + assert any(jnp.any(value != 0) for value in result.values()) + assert jnp.array_equal( + result["d_index_k"][:, seqlen_q:], + jnp.zeros_like(result["d_index_k"][:, seqlen_q:]), + ) + + 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 = result[name] + actual_f32 = actual.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}; " + f"actual norm: {jnp.linalg.norm(actual_f32)}; " + f"expected norm: {jnp.linalg.norm(expected_f32)}" + ) + assert relative_rms <= 0.55, f"{name} relative RMS error: {relative_rms}" + + zero_result = compiled( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + 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_contract.py b/test/python/fe_api/test_jax_dsa_contract.py new file mode 100644 index 000000000..7e356151b --- /dev/null +++ b/test/python/fe_api/test_jax_dsa_contract.py @@ -0,0 +1,898 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contract tests for the JAX DSA wrappers.""" + +from __future__ import annotations + +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 + + +def _identity_jit(fn=None, **_kwargs): + return (lambda decorated_fn: decorated_fn) if fn is None else fn + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CUDNN_ROOT = _REPO_ROOT / "python" / "cudnn" +_TEST_PACKAGE = "cudnn_frontend_jax_dsa_contract_test" + + +class _DType: + def __init__(self, name, itemsize): + 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 not (isinstance(key, tuple) and key[:-1] == (Ellipsis,) and isinstance(key[-1], slice) and key[-1].start is None and key[-1].step is None): + raise AssertionError(f"Unexpected test slice {key!r}") + return _Array((*self.shape[:-1], key[-1].stop), self.dtype) + + +class _TensorSpec: + def __init__( + self, + *, + layout=None, + mode=None, + static=None, + ptr_assumed_align=256, + divisibility=None, + ): + self.layout = layout + self.mode = mode + self.static = static + self.ptr_assumed_align = ptr_assumed_align + self.divisibility = divisibility + + +class JaxDsaContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.bfloat16 = _DType("bfloat16", 2) + cls.float16 = _DType("float16", 2) + cls.float32 = _DType("float32", 4) + cls.int32 = _DType("int32", 4) + cls.int64 = _DType("int64", 8) + + cls.fake_jnp = types.ModuleType("jax.numpy") + cls.fake_jnp.bfloat16 = cls.bfloat16 + cls.fake_jnp.float16 = cls.float16 + cls.fake_jnp.float32 = cls.float32 + cls.fake_jnp.int32 = cls.int32 + cls.fake_jnp.int64 = cls.int64 + cls.fake_jnp.dtype = lambda value: value + cls.fake_jnp.reshape = lambda value, shape: _Array(shape, value.dtype) + + cls.fake_jax = types.ModuleType("jax") + cls.fake_jax.__path__ = [] + cls.fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + cls.fake_jax.numpy = cls.fake_jnp + cls.fake_jax.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, + ) + + cls.fake_cutlass = types.ModuleType("cutlass") + cls.fake_cutlass.__path__ = [] + cls.fake_cutlass.Constexpr = object + cls.fake_cutlass_cute = types.ModuleType("cutlass.cute") + cls.fake_cutlass_cute.jit = _identity_jit + cls.fake_cutlass.cute = cls.fake_cutlass_cute + cls.fake_cutlass_jax = types.ModuleType("cutlass.jax") + cls.fake_cutlass_jax.is_available = lambda: True + cls.fake_cutlass_jax.TensorSpec = _TensorSpec + cls.fake_cutlass_jax.jax_to_cutlass_dtype = lambda dtype: f"cutlass.{dtype.name}" + cls.fake_cutlass.jax = cls.fake_cutlass_jax + + parent = types.ModuleType(_TEST_PACKAGE) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = _TEST_PACKAGE + sys.modules[_TEST_PACKAGE] = parent + with mock.patch.dict( + sys.modules, + { + "jax": cls.fake_jax, + "jax.numpy": cls.fake_jnp, + "cutlass": cls.fake_cutlass, + "cutlass.cute": cls.fake_cutlass_cute, + "cutlass.jax": cls.fake_cutlass_jax, + }, + ): + importlib.import_module(f"{_TEST_PACKAGE}.jax") + cls.forward = importlib.import_module(f"{_TEST_PACKAGE}.deepseek_sparse_attention.indexer_forward.jax") + cls.top_k = importlib.import_module(f"{_TEST_PACKAGE}.deepseek_sparse_attention.indexer_top_k.jax") + cls.score = importlib.import_module(f"{_TEST_PACKAGE}.deepseek_sparse_attention.score_recompute.jax") + cls.score_config = importlib.import_module(f"{_TEST_PACKAGE}.deepseek_sparse_attention.score_recompute.config") + + @classmethod + def tearDownClass(cls): + for module_name in tuple(sys.modules): + if module_name == _TEST_PACKAGE or module_name.startswith(f"{_TEST_PACKAGE}."): + sys.modules.pop(module_name, None) + + def _optional_modules(self): + return mock.patch.dict( + sys.modules, + { + "jax": self.fake_jax, + "jax.numpy": self.fake_jnp, + "cutlass": self.fake_cutlass, + "cutlass.cute": self.fake_cutlass_cute, + "cutlass.jax": self.fake_cutlass_jax, + }, + ) + + def test_shared_index_helper_kernels_do_not_import_torch_eagerly(self): + helper_root = _CUDNN_ROOT / "deepseek_sparse_attention" / "indexer_top_k" + for filename in ("local_to_global_dsl.py", "compactify.py"): + tree = ast.parse((helper_root / filename).read_text()) + eager_imports = [] + for node in tree.body: + if isinstance(node, ast.Import): + eager_imports.extend(alias.name for alias in node.names if alias.name == "torch") + elif isinstance(node, ast.ImportFrom) and node.module == "torch": + eager_imports.append(node.module) + self.assertEqual(eager_imports, [], filename) + + def test_indexer_forward_declares_initialized_padded_result(self): + captured = {} + + def fake_call(launcher, inputs, **options): + captured.update(launcher=launcher, inputs=inputs, **options) + result = options["outputs"][0] + return (_Array(result.shape, result.dtype),) + + q = _Array((1, 4, 32, 128), self.bfloat16) + k = _Array((1, 5, 1, 128), self.bfloat16) + w = _Array((1, 4, 32), self.bfloat16) + launcher = object() + with ( + self._optional_modules(), + mock.patch.object(self.forward, "_launch", new=launcher), + mock.patch.object(self.forward, "call_cutedsl", side_effect=fake_call), + ): + result = self.forward.indexer_forward_wrapper(q, k, w, ratio=1) + + self.assertEqual(result["scores"].shape, (1, 4, 5)) + self.assertIs(captured["launcher"], launcher) + self.assertEqual(captured["inputs"], (q, k, w)) + self.assertEqual( + captured["static_args"], + { + "head_dim": 128, + "qhead_per_kv_head": 32, + "ratio": 1, + "m_block_size": 128, + "n_block_size": 128, + "q_stage": 2, + "kv_stage": 4, + "num_kv_heads": 1, + "max_seqlen_q": 4, + "max_seqlen_k": 5, + "sm_scale": 1.0, + }, + ) + + output = captured["outputs"][0] + self.assertEqual(output.name, "scores") + self.assertEqual(output.shape, (1, 4, 8)) + self.assertEqual(output.dtype, self.float32) + self.assertEqual(output.fill_value, float("-inf")) + self.assertEqual(output.tensor_spec.layout, (2, 1, 0)) + self.assertEqual(output.tensor_spec.mode, (0, 1, 2)) + self.assertEqual(output.tensor_spec.divisibility, (None, None, 4)) + + def test_indexer_top_k_declares_hidden_workspace(self): + captured = {} + + def fake_call(launcher, inputs, **options): + captured.update(launcher=launcher, inputs=inputs, **options) + return tuple(_Array(spec.shape, spec.dtype) for spec in options["outputs"]) + + input_values = _Array((2, 64), self.float32) + seq_lens = _Array((2,), self.int32) + launcher = object() + with ( + self._optional_modules(), + mock.patch.object(self.top_k, "_launch", new=launcher), + mock.patch.object( + self.top_k, + "call_cutedsl", + side_effect=fake_call, + ), + ): + result = self.top_k.indexer_top_k_wrapper( + input_values, + seq_lens, + top_k=8, + ) + + self.assertEqual(result["indices"].shape, (2, 8)) + self.assertEqual(result["values"].shape, (2, 8)) + self.assertIs(captured["launcher"], launcher) + self.assertEqual(captured["inputs"], (input_values, seq_lens)) + self.assertEqual( + [(spec.name, spec.shape, spec.dtype) for spec in captured["outputs"]], + [ + ("indices", (2, 8), self.int32), + ("values", (2, 8), self.float32), + ], + ) + self.assertEqual(len(captured["workspaces"]), 1) + workspace = captured["workspaces"][0] + self.assertEqual(workspace.name, "extra_buffer") + self.assertEqual(workspace.shape, (2, 2, 64)) + self.assertEqual(workspace.dtype, self.int32) + self.assertIsNone(workspace.fill_value) + self.assertEqual( + captured["static_args"], + { + "cutlass_dtype": "cutlass.float32", + "dtype_bits": 32, + "num_cols": 64, + "top_k": 8, + "next_n": 1, + "num_copy_bits": 256, + "large_occupancy": False, + }, + ) + + def test_indexer_top_k_rejects_unsupported_result_mode(self): + input_values = _Array((2, 64), self.float16) + seq_lens = _Array((2,), self.int32) + with ( + self._optional_modules(), + self.assertRaisesRegex( + NotImplementedError, + "return_val=True", + ), + ): + self.top_k.indexer_top_k_wrapper( + input_values, + seq_lens, + top_k=8, + return_val=False, + ) + + def test_indexer_top_k_validates_selected_output_vector_width(self): + with self.assertRaisesRegex(ValueError, "selected output vector width"): + self.top_k.require_supported_top_k_output( + top_k=513, + num_threads_per_cta=256, + num_copy_bits=256, + dtype_bits=32, + ) + + # Small odd top-K values stay on the scalar path and remain valid. + self.top_k.require_supported_top_k_output( + top_k=3, + num_threads_per_cta=256, + num_copy_bits=256, + dtype_bits=32, + ) + + def test_index_helpers_declare_functional_results(self): + captured = [] + + def fake_call(launcher, inputs, **options): + captured.append(dict(launcher=launcher, inputs=inputs, **options)) + return tuple(_Array(spec.shape, spec.dtype) for spec in options["outputs"]) + + local_indices = _Array((2, 3, 8), self.int64) + compact_input = _Array((2, 3, 8), self.int32) + global_launcher = object() + compact_launcher = object() + with ( + self._optional_modules(), + mock.patch.object(self.top_k, "_launch_local_to_global_fixed", new=global_launcher), + mock.patch.object(self.top_k, "_launch_compactify", new=compact_launcher), + mock.patch.object(self.top_k, "call_cutedsl", side_effect=fake_call), + ): + global_result = self.top_k.local_to_global_wrapper( + local_indices, + seqlen_k=1024, + ) + compact_result = self.top_k.compactify_wrapper(compact_input) + + self.assertEqual(global_result["indices"].shape, (2, 3, 8)) + self.assertEqual(global_result["indices"].dtype, self.int32) + self.assertEqual(compact_result["indices"].shape, (6, 8)) + self.assertEqual(compact_result["topk_length"].shape, (6,)) + + global_call, compact_call = captured + self.assertIs(global_call["launcher"], global_launcher) + self.assertEqual(global_call["inputs"], (local_indices,)) + self.assertEqual( + [(spec.name, spec.shape, spec.dtype) for spec in global_call["outputs"]], + [("indices", (2, 3, 8), self.int32)], + ) + self.assertEqual(global_call["static_args"], {"seqlen_k": 1024}) + + self.assertIs(compact_call["launcher"], compact_launcher) + self.assertEqual(compact_call["inputs"][0].shape, (6, 8)) + self.assertEqual( + [(spec.name, spec.shape, spec.dtype) for spec in compact_call["outputs"]], + [ + ("indices", (6, 8), self.int32), + ("topk_length", (6,), self.int32), + ], + ) + self.assertEqual(compact_call["static_args"], {"rows": 6, "cols": 8}) + + def test_local_to_global_packed_inputs_remain_device_values(self): + captured = {} + + def fake_call(launcher, inputs, **options): + captured.update(launcher=launcher, inputs=inputs, **options) + spec = options["outputs"][0] + return (_Array(spec.shape, spec.dtype),) + + local_indices = _Array((7, 8), self.int32) + cu_seqlens_q = _Array((3,), self.int32) + cu_seqlens_k = _Array((3,), self.int32) + launcher = object() + with ( + self._optional_modules(), + mock.patch.object(self.top_k, "_launch_local_to_global_varlen", new=launcher), + mock.patch.object(self.top_k, "call_cutedsl", side_effect=fake_call), + ): + result = self.top_k.local_to_global_wrapper( + local_indices, + seqlen_k=16, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + ) + + self.assertEqual(result["indices"].shape, local_indices.shape) + self.assertEqual( + captured["inputs"], + (local_indices, cu_seqlens_q, cu_seqlens_k), + ) + self.assertEqual(captured["static_args"], {"seqlen_k": 16}) + + def test_index_helper_launchers_preserve_native_argument_order(self): + calls = [] + + class FakeLocalToGlobalKernel: + def __init__(self, **options): + self.options = options + + def __call__(self, *args): + calls.append(("global", args)) + + class FakeCompactifyKernel: + def __init__(self, **options): + self.options = options + + def __call__(self, *args): + calls.append(("compact", args)) + + package = f"{_TEST_PACKAGE}.deepseek_sparse_attention.indexer_top_k" + global_module = types.ModuleType(f"{package}.local_to_global_dsl") + global_module.LocalToGlobalTopK = FakeLocalToGlobalKernel + compact_module = types.ModuleType(f"{package}.compactify") + compact_module.CompactifyKernel = FakeCompactifyKernel + int32 = lambda value: ("Int32", value) + + with ( + self._optional_modules(), + mock.patch.dict( + sys.modules, + { + global_module.__name__: global_module, + compact_module.__name__: compact_module, + }, + ), + mock.patch.object(self.fake_cutlass, "Int32", int32, create=True), + ): + self.top_k._launch_local_to_global_fixed("stream", "local", "global", seqlen_k=32) + self.top_k._launch_local_to_global_varlen("stream", "local", "cuq", "cuk", "global", seqlen_k=32) + self.top_k._launch_compactify("stream", "indices", "out", "length", rows=6, cols=8) + + self.assertEqual( + calls, + [ + ( + "global", + ( + "local", + "global", + ("Int32", 32), + None, + None, + "stream", + ), + ), + ( + "global", + ( + "local", + "global", + ("Int32", 32), + "cuq", + "cuk", + "stream", + ), + ), + ( + "compact", + ( + "indices", + "out", + "length", + ("Int32", 6), + "stream", + ), + ), + ], + ) + + def test_dense_score_wrappers_declare_initialized_results(self): + captured = [] + + def fake_call(launcher, inputs, **options): + captured.append(dict(launcher=launcher, inputs=inputs, **options)) + return tuple(_Array(spec.shape, spec.dtype) for spec in options["outputs"]) + + q = _Array((2, 4, 32, 128), self.bfloat16) + k = _Array((2, 8, 1, 128), self.bfloat16) + weights = _Array((2, 4, 32), self.bfloat16) + lse = _Array((2, 4, 32), self.float32) + q_causal_offsets = _Array((2,), self.int32) + indexer_launcher = object() + attention_launcher = object() + + with ( + self._optional_modules(), + mock.patch.object( + self.score, + "_launch_dense_with_q_causal_offsets", + new=indexer_launcher, + ), + mock.patch.object( + self.score, + "_launch_dense_without_q_causal_offsets", + new=attention_launcher, + ), + mock.patch.object(self.score, "call_cutedsl", side_effect=fake_call), + ): + indexer = self.score.dense_indexer_score_recompute_wrapper( + q, + k, + weights, + ratio=2, + q_causal_offsets=q_causal_offsets, + ) + attention = self.score.dense_attn_score_recompute_wrapper( + q, + k, + lse, + softmax_scale=0.125, + ) + + for result in (indexer, attention): + self.assertEqual((result["out"].shape, result["out"].dtype), ((2, 4, 8), self.float32)) + self.assertEqual((result["denom"].shape, result["denom"].dtype), ((2, 4), self.float32)) + + indexer_call, attention_call = captured + self.assertIs(indexer_call["launcher"], indexer_launcher) + self.assertEqual( + indexer_call["inputs"], + (q, k, weights, q_causal_offsets), + ) + self.assertEqual(indexer_call["outputs"][0].name, "out") + self.assertEqual(indexer_call["outputs"][0].fill_value, float("-inf")) + self.assertEqual(indexer_call["outputs"][1].name, "denom") + self.assertIsNone(indexer_call["outputs"][1].fill_value) + + self.assertIs(attention_call["launcher"], attention_launcher) + self.assertEqual(attention_call["inputs"], (q, k, lse)) + self.assertEqual(attention_call["outputs"][0].fill_value, float("-inf")) + + indexer_config = indexer_call["static_args"] + self.assertEqual(indexer_config["score_type"], "indexer") + self.assertEqual(indexer_config["ratio"], 2) + attention_config = attention_call["static_args"] + self.assertEqual(attention_config["score_type"], "attention") + self.assertEqual(attention_config["scale"], 0.125) + + def test_sparse_score_wrappers_declare_functional_results(self): + captured = [] + + def fake_call(launcher, inputs, **options): + captured.append(dict(launcher=launcher, inputs=inputs, **options)) + return tuple(_Array(spec.shape, spec.dtype) for spec in options["outputs"]) + + q = _Array((2, 4, 32, 128), self.bfloat16) + k = _Array((2, 256, 128), self.bfloat16) + weights = _Array((2, 4, 32), self.bfloat16) + lse = _Array((2, 4, 32), self.float32) + indices = _Array((2, 4, 128), self.int32) + lengths = _Array((2, 4), self.int32) + indexer_launcher = object() + attention_launcher = object() + + with ( + self._optional_modules(), + mock.patch.object( + self.score, + "_launch_sparse_without_topk_length", + new=indexer_launcher, + ), + mock.patch.object( + self.score, + "_launch_sparse_with_topk_length", + new=attention_launcher, + ), + mock.patch.object(self.score, "call_cutedsl", side_effect=fake_call), + ): + predict = self.score.sparse_indexer_score_recompute_wrapper( + q, + k, + weights, + indices, + )["predict"] + target = self.score.sparse_attn_score_recompute_wrapper( + q, + k, + lse, + indices, + 0.125, + topk_length=lengths, + )["target"] + + self.assertEqual((predict.shape, predict.dtype), ((2, 4, 128), self.float32)) + self.assertEqual((target.shape, target.dtype), ((2, 4, 128), self.float32)) + + indexer_call, attention_call = captured + self.assertIs(indexer_call["launcher"], indexer_launcher) + self.assertEqual(indexer_call["inputs"], (q, k, weights, indices)) + self.assertEqual(len(indexer_call["outputs"]), 1) + self.assertEqual(indexer_call["outputs"][0].name, "predict") + self.assertEqual(len(indexer_call["workspaces"]), 1) + self.assertEqual( + ( + indexer_call["workspaces"][0].name, + indexer_call["workspaces"][0].shape, + indexer_call["workspaces"][0].dtype, + ), + ("topk_length_workspace", (1, 1), self.int32), + ) + + self.assertIs(attention_call["launcher"], attention_launcher) + self.assertEqual(attention_call["inputs"], (q, k, lse, indices, lengths)) + self.assertEqual(attention_call["outputs"][0].name, "target") + self.assertEqual(attention_call["workspaces"], ()) + + indexer_config = indexer_call["static_args"] + self.assertEqual(indexer_config["score_type"], "indexer") + self.assertEqual(indexer_config["n_block_size"], 128) + attention_config = attention_call["static_args"] + self.assertEqual(attention_config["score_type"], "attention") + self.assertEqual(attention_config["n_block_size"], 64) + self.assertEqual(attention_config["softmax_scale"], 0.125) + + def test_sparse_score_launcher_preserves_native_argument_order(self): + calls = [] + + class FakeKernel: + def __init__(self, **options): + self.options = options + + def __call__(self, *args): + calls.append(args) + + kernel_module_name = f"{_TEST_PACKAGE}.deepseek_sparse_attention.score_recompute." "sparse_score_recompute_sm100" + kernel_module = types.ModuleType(kernel_module_name) + kernel_module.SparseScoreRecomputeSm100 = FakeKernel + float32 = lambda value: ("Float32", value) + + common = dict( + score_type="indexer", + head_dim=128, + qhead_per_kv_head=32, + topk=128, + m_block_size=32, + n_block_size=128, + k_block_size=None, + kv_stage=4, + topk_in_smem=True, + topk_indices_global=False, + softmax_scale=1.0, + ) + with ( + self._optional_modules(), + mock.patch.dict(sys.modules, {kernel_module_name: kernel_module}), + mock.patch.object(self.fake_cutlass, "Float32", float32, create=True), + ): + self.score._launch_sparse_with_topk_length( + "stream", + "q", + "k", + "aux", + "indices", + "length", + "out", + **common, + ) + self.score._launch_sparse_without_topk_length( + "stream", + "q", + "k", + "aux", + "indices", + "out", + "dummy_length", + **common, + ) + + self.assertEqual( + calls, + [ + ( + "q", + "k", + "aux", + "indices", + "out", + "length", + ("Float32", 1.0), + "stream", + ), + ( + "q", + "k", + "aux", + "indices", + "out", + "dummy_length", + ("Float32", 1.0), + "stream", + ), + ], + ) + + def test_dense_score_launcher_preserves_native_argument_order(self): + calls = [] + + class FakeKernel: + def __init__(self, **options): + self.options = options + + def __call__(self, *args): + calls.append(args) + + kernel_module_name = f"{_TEST_PACKAGE}.deepseek_sparse_attention.score_recompute." "dense_score_recompute_sm100" + kernel_module = types.ModuleType(kernel_module_name) + kernel_module.DenseScoreRecomputeSm100 = FakeKernel + float32 = lambda value: ("Float32", value) + int32 = lambda value: ("Int32", value) + + with ( + self._optional_modules(), + mock.patch.dict(sys.modules, {kernel_module_name: kernel_module}), + mock.patch.object(self.fake_cutlass, "Float32", float32, create=True), + mock.patch.object(self.fake_cutlass, "Int32", int32, create=True), + ): + self.score._launch_dense_without_q_causal_offsets( + "stream", + "q", + "k", + "weights", + "out", + "denom", + score_type="indexer", + head_dim=128, + qhead_per_kv_head=32, + ratio=1, + max_seqlen_q=4, + max_seqlen_k=8, + scale=0.5, + ) + self.score._launch_dense_with_q_causal_offsets( + "stream", + "q", + "k", + "lse", + "offsets", + "out", + "denom", + score_type="attention", + head_dim=128, + qhead_per_kv_head=32, + ratio=2, + max_seqlen_q=4, + max_seqlen_k=8, + scale=0.125, + ) + + self.assertEqual( + calls, + [ + ( + "q", + "k", + "weights", + "out", + "denom", + ("Float32", 0.5), + ("Int32", 4), + ("Int32", 8), + None, + None, + None, + "stream", + ), + ( + "q", + "k", + "lse", + "out", + "denom", + ("Float32", 0.125), + ("Int32", 4), + ("Int32", 8), + None, + None, + "offsets", + "stream", + ), + ], + ) + + def test_sparse_score_config_matches_sm100_dispatch(self): + indexer = self.score_config.resolve_sparse_score_kernel_config( + score_type="indexer", + head_dim=128, + qhead_per_kv_head=32, + topk=128, + have_topk_length=False, + ) + self.assertEqual( + (indexer.m_block_size, indexer.n_block_size, indexer.k_block_size), + (32, 128, None), + ) + self.assertEqual(indexer.kv_stage, 4) + self.assertTrue(indexer.topk_in_smem) + + large_topk = self.score_config.resolve_sparse_score_kernel_config( + score_type="indexer", + head_dim=128, + qhead_per_kv_head=32, + topk=32768, + have_topk_length=False, + ) + self.assertFalse(large_topk.topk_in_smem) + + compact_attention = self.score_config.resolve_sparse_score_kernel_config( + score_type="attention", + head_dim=512, + qhead_per_kv_head=64, + topk=512, + have_topk_length=True, + ) + full_attention = self.score_config.resolve_sparse_score_kernel_config( + score_type="attention", + head_dim=512, + qhead_per_kv_head=64, + topk=512, + have_topk_length=False, + ) + self.assertEqual( + ( + compact_attention.m_block_size, + compact_attention.n_block_size, + compact_attention.k_block_size, + ), + (64, 64, None), + ) + self.assertEqual( + ( + full_attention.m_block_size, + full_attention.n_block_size, + full_attention.k_block_size, + ), + (64, 128, 256), + ) + + with self.assertRaisesRegex(ValueError, "multiple of the selected n_block_size"): + self.score_config.resolve_sparse_score_kernel_config( + score_type="indexer", + head_dim=128, + qhead_per_kv_head=32, + topk=64, + have_topk_length=False, + ) + + dense_indexer = self.score_config.resolve_dense_score_kernel_config( + score_type="indexer", + head_dim=128, + qhead_per_kv_head=32, + ) + self.assertEqual( + ( + dense_indexer.m_block_size, + dense_indexer.n_block_size, + dense_indexer.k_block_size, + dense_indexer.kv_stage, + ), + (64, 128, None, 4), + ) + + dense_attention = self.score_config.resolve_dense_score_kernel_config( + score_type="attention", + head_dim=512, + qhead_per_kv_head=64, + ) + self.assertEqual( + ( + dense_attention.m_block_size, + dense_attention.n_block_size, + dense_attention.k_block_size, + dense_attention.kv_stage, + ), + (128, 128, 64, 4), + ) + + def test_wrapper_validation_fails_before_launch(self): + bad_q = _Array((1, 4, 32, 128), self.float16) + k = _Array((1, 5, 1, 128), self.bfloat16) + w = _Array((1, 4, 32), self.bfloat16) + with ( + self._optional_modules(), + self.assertRaisesRegex( + ValueError, + "q.dtype", + ), + mock.patch.object(self.forward, "call_cutedsl") as call_forward, + ): + self.forward.indexer_forward_wrapper(bad_q, k, w, ratio=1) + call_forward.assert_not_called() + + input_values = _Array((2, 64), self.float32) + wrong_batch = _Array((1,), self.int32) + with ( + self._optional_modules(), + self.assertRaisesRegex( + ValueError, + "num_rows.*must equal", + ), + mock.patch.object(self.top_k, "call_cutedsl") as call_top_k, + ): + self.top_k.indexer_top_k_wrapper( + input_values, + wrong_batch, + top_k=8, + ) + call_top_k.assert_not_called() + + +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 new file mode 100644 index 000000000..c0d536133 --- /dev/null +++ b/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py @@ -0,0 +1,616 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for JAX DSA indexer backward.""" + +from __future__ import annotations + +import ast +from importlib import import_module +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 + + +def _identity_jit(fn=None, **_kwargs): + return (lambda decorated_fn: decorated_fn) if fn is None else fn + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CUDNN_ROOT = _REPO_ROOT / "python" / "cudnn" +_TEST_PACKAGE = "cudnn_frontend_jax_dsa_indexer_bwd_contract_test" + + +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 _TensorSpec: + def __init__( + self, + *, + layout=None, + mode=None, + static=None, + ptr_assumed_align=256, + divisibility=None, + ): + self.layout = layout + self.mode = mode + self.static = static + self.ptr_assumed_align = ptr_assumed_align + self.divisibility = divisibility + + +class _Float32: + def __init__(self, value=0.0): + self.value = value + + +class _Int32: + def __init__(self, value=0): + self.value = value + + +class _ScoreGradKernel: + instances = [] + + def __init__(self, *, topk): + self.topk = topk + self.calls = [] + self.instances.append(self) + + def __call__(self, *args): + self.calls.append(args) + + +class _BackwardKernel: + instances = [] + + def __init__( + self, + *, + head_dim, + heads, + block_I, + topk, + topk_indices_global, + ): + self.configuration = ( + head_dim, + heads, + block_I, + topk, + topk_indices_global, + ) + self.calls = [] + self.instances.append(self) + + def __call__(self, *args): + self.calls.append(args) + + +class _DenseScoreGradKernel: + instances = [] + + def __init__(self, *, ratio, block_I): + self.configuration = (ratio, block_I) + self.calls = [] + self.instances.append(self) + + def __call__(self, *args): + self.calls.append(args) + + +class _DenseBackwardKernel: + instances = [] + + def __init__(self, *, head_dim, heads, block_I, ratio): + self.configuration = (head_dim, heads, block_I, ratio) + self.calls = [] + self.instances.append(self) + + def __call__(self, *args): + self.calls.append(args) + + +class JaxDsaIndexerBackwardContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.bfloat16 = _DType("bfloat16") + cls.float32 = _DType("float32") + cls.int32 = _DType("int32") + + cls.fake_jnp = types.ModuleType("jax.numpy") + cls.fake_jnp.bfloat16 = cls.bfloat16 + cls.fake_jnp.float32 = cls.float32 + cls.fake_jnp.int32 = cls.int32 + cls.fake_jnp.dtype = lambda value: value + cls.fake_jnp.reshape = lambda value, shape: _Array(shape, value.dtype) + cls.fake_jnp.asarray = lambda value, dtype: _Array((len(value),), dtype) + + cls.fake_jax = types.ModuleType("jax") + cls.fake_jax.__path__ = [] + cls.fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + cls.fake_jax.numpy = cls.fake_jnp + cls.fake_jax.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, + ) + cls.fake_jax.ShapeDtypeStruct = lambda shape, dtype: (shape, dtype) + + cls.fake_cutlass = types.ModuleType("cutlass") + cls.fake_cutlass.__path__ = [] + cls.fake_cutlass.Constexpr = object + cls.fake_cutlass.Float32 = _Float32 + cls.fake_cutlass.Int32 = _Int32 + cls.fake_cutlass_cute = types.ModuleType("cutlass.cute") + cls.fake_cutlass_cute.jit = _identity_jit + cls.fake_cutlass.cute = cls.fake_cutlass_cute + cls.fake_cutlass_jax = types.ModuleType("cutlass.jax") + cls.fake_cutlass_jax.TensorSpec = _TensorSpec + cls.fake_cutlass_jax.cutlass_call = None + cls.fake_cutlass.jax = cls.fake_cutlass_jax + + cls.kernel_module_name = f"{_TEST_PACKAGE}.deepseek_sparse_attention.indexer_backward." "indexer_backward_sm100" + cls.kernel_module = types.ModuleType(cls.kernel_module_name) + cls.kernel_module.ScoreGradSm100 = _ScoreGradKernel + cls.kernel_module.IndexerBackwardSm100 = _BackwardKernel + cls.dense_kernel_module_name = ( + f"{_TEST_PACKAGE}.deepseek_sparse_attention.indexer_backward." + "dense_indexer_backward_sm100" + ) + cls.dense_kernel_module = types.ModuleType(cls.dense_kernel_module_name) + cls.dense_kernel_module.ScoreGradDense = _DenseScoreGradKernel + cls.dense_kernel_module.DenseIndexerBackward2QGemmSm100 = _DenseBackwardKernel + + package_paths = { + _TEST_PACKAGE: _CUDNN_ROOT, + f"{_TEST_PACKAGE}.deepseek_sparse_attention": (_CUDNN_ROOT / "deepseek_sparse_attention"), + f"{_TEST_PACKAGE}.deepseek_sparse_attention.indexer_backward": (_CUDNN_ROOT / "deepseek_sparse_attention" / "indexer_backward"), + } + for package_name, package_path in package_paths.items(): + package = types.ModuleType(package_name) + package.__path__ = [str(package_path)] + package.__package__ = package_name + sys.modules[package_name] = package + + with cls._optional_modules(): + cls.module = import_module(f"{_TEST_PACKAGE}.deepseek_sparse_attention.indexer_backward.jax") + + @classmethod + def tearDownClass(cls): + for module_name in tuple(sys.modules): + if module_name == _TEST_PACKAGE or module_name.startswith(f"{_TEST_PACKAGE}."): + sys.modules.pop(module_name, None) + + def setUp(self): + _ScoreGradKernel.instances.clear() + _BackwardKernel.instances.clear() + _DenseScoreGradKernel.instances.clear() + _DenseBackwardKernel.instances.clear() + + @classmethod + def _optional_modules(cls, *, include_kernel=False): + modules = { + "jax": cls.fake_jax, + "jax.numpy": cls.fake_jnp, + "cutlass": cls.fake_cutlass, + "cutlass.cute": cls.fake_cutlass_cute, + "cutlass.jax": cls.fake_cutlass_jax, + } + if include_kernel: + modules[cls.kernel_module_name] = cls.kernel_module + modules[cls.dense_kernel_module_name] = cls.dense_kernel_module + return mock.patch.dict(sys.modules, modules) + + @classmethod + def _inputs(cls, *, topk=128): + q = _Array((2, 64, 64, 128), cls.bfloat16) + weights = _Array((2, 64, 64), cls.bfloat16) + k = _Array((2, 256, 128), cls.bfloat16) + score_shape = (2, 64, topk) + attn_score = _Array(score_shape, cls.float32) + index_score = _Array(score_shape, cls.float32) + topk_indices = _Array(score_shape, cls.int32) + return q, weights, k, attn_score, index_score, topk_indices + + @classmethod + def _dense_inputs(cls): + q = _Array((2, 64, 64, 128), cls.bfloat16) + weights = _Array((2, 64, 64), cls.bfloat16) + k = _Array((2, 256, 128), cls.bfloat16) + score_shape = (2, 64, 256) + denom_shape = (2, 64) + attn_score = _Array(score_shape, cls.float32) + attn_l1norm = _Array(denom_shape, cls.float32) + index_score = _Array(score_shape, cls.float32) + index_lse = _Array(denom_shape, cls.float32) + return ( + q, + weights, + k, + attn_score, + attn_l1norm, + index_score, + index_lse, + ) + + @staticmethod + def _fake_call(captured): + def call(launcher, inputs, **options): + captured.update(launcher=launcher, inputs=inputs, **options) + return tuple(_Array(spec.shape, spec.dtype) for spec in options["outputs"]) + + return call + + def test_kernel_module_is_lazy(self): + self.assertNotIn(self.kernel_module_name, sys.modules) + self.assertNotIn(self.dense_kernel_module_name, sys.modules) + + def test_dense_shared_kernel_keeps_torch_support_lazy(self): + kernel_path = ( + _CUDNN_ROOT + / "deepseek_sparse_attention" + / "indexer_backward" + / "dense_indexer_backward_sm100.py" + ) + tree = ast.parse(kernel_path.read_text()) + eager_imports = [] + for node in tree.body: + if isinstance(node, ast.Import): + eager_imports.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module is not None: + eager_imports.append(node.module) + + self.assertNotIn("torch", eager_imports) + self.assertNotIn( + "cudnn.deepseek_sparse_attention.utils.compiler", + eager_imports, + ) + self.assertNotIn( + "cudnn.deepseek_sparse_attention.utils.runtime", + eager_imports, + ) + + def test_declares_functional_outputs_and_hidden_workspace(self): + captured = {} + launcher = object() + inputs = self._inputs() + + with ( + self._optional_modules(include_kernel=True), + mock.patch.object(self.module, "_launch", new=launcher), + mock.patch.object( + self.module, + "call_cutedsl", + side_effect=self._fake_call(captured), + ), + ): + result = self.module.indexer_backward_wrapper(*inputs) + + q, weights, k, attn_score, index_score, topk_indices = inputs + self.assertEqual(result["d_index_q"].shape, q.shape) + self.assertEqual(result["d_weights"].shape, weights.shape) + self.assertEqual(result["d_index_k"].shape, k.shape) + self.assertIs(result["d_index_k"].dtype, self.bfloat16) + self.assertEqual(captured["inputs"][:-1], inputs) + grad_loss = captured["inputs"][-1] + self.assertEqual((grad_loss.shape, grad_loss.dtype), ((1,), self.float32)) + self.assertIs(captured["launcher"], launcher) + self.assertEqual( + [(spec.name, spec.shape, spec.dtype, spec.fill_value) for spec in captured["outputs"]], + [ + ("d_index_q", q.shape, self.bfloat16, None), + ("d_weights", weights.shape, self.bfloat16, None), + ("d_index_k_accum", k.shape, self.float32, 0.0), + ], + ) + (workspace,) = captured["workspaces"] + self.assertEqual( + (workspace.name, workspace.shape, workspace.dtype, workspace.fill_value), + ("grad_signal", attn_score.shape, self.float32, None), + ) + self.assertEqual(len(captured["input_specs"]), 7) + self.assertIs(captured["inputs"][3], attn_score) + self.assertIs(captured["inputs"][4], index_score) + self.assertIs(captured["inputs"][5], topk_indices) + self.assertEqual( + captured["static_args"], + { + "heads": 64, + "head_dim": 128, + "topk": 128, + "block_i": 128, + "sm_scale": 1.0, + "grad_scale": 1.0 / (2 * 64), + "topk_indices_global": False, + }, + ) + + def test_grad_loss_array_remains_a_runtime_operand(self): + captured = {} + grad_loss = _Array((), self.float32) + with ( + self._optional_modules(include_kernel=True), + mock.patch.object(self.module, "_launch", new=object()), + mock.patch.object( + self.module, + "call_cutedsl", + side_effect=self._fake_call(captured), + ), + ): + self.module.indexer_backward_wrapper(*self._inputs(), grad_loss=grad_loss) + + operand = captured["inputs"][-1] + self.assertEqual((operand.shape, operand.dtype), ((1,), self.float32)) + + def test_launcher_orders_stages_without_mutating_score_inputs(self): + placeholders = [object() for _ in range(12)] + with self._optional_modules(include_kernel=True): + self.module._launch( + *placeholders, + heads=64, + head_dim=128, + topk=128, + block_i=128, + sm_scale=0.125, + grad_scale=0.25, + topk_indices_global=False, + ) + + score_grad = _ScoreGradKernel.instances[-1] + backward = _BackwardKernel.instances[-1] + self.assertEqual(score_grad.topk, 128) + self.assertEqual(backward.configuration, (128, 64, 128, 128, False)) + + score_args = score_grad.calls[-1] + self.assertIs(score_args[0], placeholders[4]) + self.assertIs(score_args[1], placeholders[5]) + self.assertIs(score_args[2], placeholders[7]) + self.assertEqual(score_args[3].value, 0.25) + self.assertIs(score_args[4], placeholders[0]) + self.assertIs(score_args[5], placeholders[11]) + self.assertIsNone(score_args[6]) + + backward_args = backward.calls[-1] + self.assertEqual( + backward_args[:8], + ( + placeholders[1], + placeholders[2], + placeholders[3], + placeholders[8], + placeholders[9], + placeholders[10], + placeholders[11], + placeholders[6], + ), + ) + self.assertEqual(backward_args[8].value, 0.125) + self.assertIs(backward_args[9], placeholders[0]) + + def test_dense_declares_functional_outputs_and_hidden_workspace(self): + captured = {} + launcher = object() + inputs = self._dense_inputs() + + with ( + self._optional_modules(include_kernel=True), + mock.patch.object(self.module, "_launch_dense", new=launcher), + mock.patch.object( + self.module, + "call_cutedsl", + side_effect=self._fake_call(captured), + ), + ): + result = self.module.dense_indexer_backward_wrapper(*inputs) + + q, weights, k, attn_score, attn_l1norm, index_score, index_lse = inputs + self.assertEqual(result["d_index_q"].shape, q.shape) + self.assertEqual(result["d_weights"].shape, weights.shape) + self.assertEqual(result["d_index_k"].shape, k.shape) + self.assertIs(result["d_index_k"].dtype, self.bfloat16) + self.assertEqual(captured["inputs"][:-1], inputs) + grad_loss = captured["inputs"][-1] + self.assertEqual((grad_loss.shape, grad_loss.dtype), ((1,), self.float32)) + self.assertIs(captured["launcher"], launcher) + self.assertEqual( + [ + (spec.name, spec.shape, spec.dtype, spec.fill_value) + for spec in captured["outputs"] + ], + [ + ("d_index_q", q.shape, self.bfloat16, None), + ("d_weights", weights.shape, self.bfloat16, None), + ("d_index_k_accum", k.shape, self.float32, None), + ], + ) + (workspace,) = captured["workspaces"] + self.assertEqual( + (workspace.name, workspace.shape, workspace.dtype, workspace.fill_value), + ("grad_signal", attn_score.shape, self.float32, None), + ) + self.assertEqual(len(captured["input_specs"]), 8) + q_spec = captured["input_specs"][0] + k_spec = captured["input_specs"][2] + self.assertEqual((q_spec.layout, q_spec.divisibility), ((3, 2, 1, 0), 128)) + self.assertEqual((k_spec.layout, k_spec.divisibility), ((2, 1, 0), 128)) + self.assertIs(captured["outputs"][0].tensor_spec, q_spec) + self.assertIs(captured["outputs"][2].tensor_spec, k_spec) + self.assertIs(captured["inputs"][3], attn_score) + self.assertIs(captured["inputs"][4], attn_l1norm) + self.assertIs(captured["inputs"][5], index_score) + self.assertIs(captured["inputs"][6], index_lse) + self.assertEqual( + captured["static_args"], + { + "heads": 64, + "head_dim": 128, + "block_i": 128, + "ratio": 1, + "sm_scale": 1.0, + "grad_scale": 1.0 / (2 * 64), + "max_seqlen_q": 64, + "max_seqlen_k": 256, + }, + ) + + def test_dense_launcher_orders_stages_and_uses_runtime_grad_loss(self): + placeholders = [object() for _ in range(13)] + with self._optional_modules(include_kernel=True): + self.module._launch_dense( + *placeholders, + heads=64, + head_dim=128, + block_i=128, + ratio=2, + sm_scale=0.125, + grad_scale=0.25, + max_seqlen_q=64, + max_seqlen_k=256, + ) + + score_grad = _DenseScoreGradKernel.instances[-1] + backward = _DenseBackwardKernel.instances[-1] + self.assertEqual(score_grad.configuration, (2, 128)) + self.assertEqual(backward.configuration, (128, 64, 128, 2)) + + score_args = score_grad.calls[-1] + self.assertEqual( + score_args[:9], + ( + placeholders[6], + placeholders[12], + placeholders[4], + placeholders[7], + placeholders[5], + None, + None, + None, + placeholders[8], + ), + ) + self.assertEqual(score_args[9].value, 0.25) + self.assertEqual(score_args[10].value, 64) + self.assertEqual(score_args[11].value, 256) + self.assertIs(score_args[12], placeholders[0]) + + backward_args = backward.calls[-1] + self.assertEqual( + backward_args[:10], + ( + placeholders[1], + placeholders[2], + placeholders[3], + placeholders[9], + placeholders[10], + placeholders[11], + placeholders[12], + None, + None, + None, + ), + ) + self.assertEqual(backward_args[10].value, 0.125) + self.assertEqual(backward_args[11].value, 64) + self.assertEqual(backward_args[12].value, 256) + self.assertIs(backward_args[13], placeholders[0]) + + def test_dense_rejects_unsupported_shape_and_configuration(self): + q, weights, k, attn_score, attn_l1norm, index_score, index_lse = ( + self._dense_inputs() + ) + with self.assertRaisesRegex(ValueError, "heads=64 and head_dim=128"): + self.module.dense_indexer_backward_wrapper( + _Array((2, 64, 32, 128), self.bfloat16), + _Array((2, 64, 32), self.bfloat16), + k, + attn_score, + attn_l1norm, + index_score, + index_lse, + ) + + with self.assertRaisesRegex(ValueError, "block_I must be 128"): + self.module.dense_indexer_backward_wrapper( + q, + weights, + k, + attn_score, + attn_l1norm, + index_score, + index_lse, + block_I=64, + ) + + def test_rejects_unsupported_dtype_shape_and_topk(self): + q, weights, k, attn_score, index_score, topk_indices = self._inputs() + with self.assertRaisesRegex(ValueError, "index_q.dtype"): + self.module.indexer_backward_wrapper( + _Array(q.shape, self.float32), + weights, + k, + attn_score, + index_score, + topk_indices, + ) + + with self.assertRaisesRegex(ValueError, "heads=64 and head_dim=128"): + self.module.indexer_backward_wrapper( + _Array((2, 64, 32, 128), self.bfloat16), + _Array((2, 64, 32), self.bfloat16), + k, + attn_score, + index_score, + topk_indices, + ) + + inputs = self._inputs(topk=64) + with self.assertRaisesRegex(ValueError, "must be divisible"): + self.module.indexer_backward_wrapper(*inputs) + + def test_child_package_declares_literal_api_surface(self): + init_path = _CUDNN_ROOT / "deepseek_sparse_attention" / "indexer_backward" / "__init__.py" + tree = ast.parse(init_path.read_text()) + exports = next( + node.value + for node in tree.body + if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == "_API_EXPORTS" for target in node.targets) + ) + self.assertEqual( + ast.literal_eval(exports), + ( + "DenseIndexerBackward", + "IndexerBackward", + "dense_indexer_backward_wrapper", + "indexer_backward_wrapper", + ), + ) + + +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..290a3b7b8 --- /dev/null +++ b/test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py @@ -0,0 +1,367 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for JAX DSA sparse-attention backward.""" + +from __future__ import annotations + +from importlib import import_module +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 + + +def _identity_jit(fn=None, **_kwargs): + return (lambda decorated_fn: decorated_fn) if fn is None else fn + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CUDNN_ROOT = _REPO_ROOT / "python" / "cudnn" +_TEST_PACKAGE = "cudnn_frontend_jax_dsa_sparse_bwd_contract_test" + + +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 + + +class _TensorSpec: + def __init__( + self, + *, + layout=None, + mode=None, + static=None, + ptr_assumed_align=256, + divisibility=None, + ): + self.layout = layout + self.mode = mode + self.static = static + self.ptr_assumed_align = ptr_assumed_align + self.divisibility = divisibility + + +class _Float32: + width = 32 + + def __init__(self, value=0.0): + self.value = value + + +class _Kernel: + instances = [] + + def __init__(self, *, head_dim, head_dim_v, block_tile): + self.configuration = (head_dim, head_dim_v, block_tile) + self.calls = [] + self.instances.append(self) + + @staticmethod + def get_workspace_size_lse_odo(q, d, h, b, acc_dtype): + assert acc_dtype is _Float32 + return (b, h, ((q + 7) // 8) * 8, 2 * acc_dtype.width // 8) + + @staticmethod + def get_workspace_size_dkv(k, d, b, acc_dtype): + assert acc_dtype is _Float32 + return ( + b, + 1, + ((k + 7) // 8) * 8, + ((d + 7) // 8) * 8 * acc_dtype.width // 8, + ) + + def __call__(self, *args): + self.calls.append(args) + + +class JaxDsaSparseAttentionBackwardContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.bfloat16 = _DType("bfloat16", 2) + cls.float32 = _DType("float32", 4) + cls.int32 = _DType("int32", 4) + cls.uint8 = _DType("uint8", 1) + + cls.fake_jnp = types.ModuleType("jax.numpy") + cls.fake_jnp.bfloat16 = cls.bfloat16 + cls.fake_jnp.float32 = cls.float32 + cls.fake_jnp.int32 = cls.int32 + cls.fake_jnp.uint8 = cls.uint8 + cls.fake_jnp.dtype = lambda value: value + + cls.fake_jax = types.ModuleType("jax") + cls.fake_jax.__path__ = [] + cls.fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + cls.fake_jax.numpy = cls.fake_jnp + cls.fake_jax.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, + ) + cls.fake_jax.ShapeDtypeStruct = lambda shape, dtype: (shape, dtype) + + cls.fake_cutlass = types.ModuleType("cutlass") + cls.fake_cutlass.__path__ = [] + cls.fake_cutlass.Constexpr = object + cls.fake_cutlass.Float32 = _Float32 + cls.fake_cutlass.Int32 = int + cls.fake_cutlass_cute = types.ModuleType("cutlass.cute") + cls.fake_cutlass_cute.jit = _identity_jit + cls.fake_cutlass.cute = cls.fake_cutlass_cute + cls.fake_cutlass_jax = types.ModuleType("cutlass.jax") + cls.fake_cutlass_jax.TensorSpec = _TensorSpec + cls.fake_cutlass_jax.cutlass_call = None + cls.fake_cutlass.jax = cls.fake_cutlass_jax + + cls.kernel_module_name = f"{_TEST_PACKAGE}.deepseek_sparse_attention.sparse_attention_backward.dsa_bwd_sm100" + cls.kernel_module = types.ModuleType(cls.kernel_module_name) + cls.kernel_module.FlashAttentionDSABackwardSm100 = _Kernel + + package_paths = { + _TEST_PACKAGE: _CUDNN_ROOT, + f"{_TEST_PACKAGE}.deepseek_sparse_attention": _CUDNN_ROOT / "deepseek_sparse_attention", + f"{_TEST_PACKAGE}.deepseek_sparse_attention.sparse_attention_backward": _CUDNN_ROOT / "deepseek_sparse_attention" / "sparse_attention_backward", + } + for package_name, package_path in package_paths.items(): + package = types.ModuleType(package_name) + package.__path__ = [str(package_path)] + package.__package__ = package_name + sys.modules[package_name] = package + + with cls._optional_modules(): + cls.module = import_module(f"{_TEST_PACKAGE}.deepseek_sparse_attention.sparse_attention_backward.jax") + + @classmethod + def tearDownClass(cls): + for module_name in tuple(sys.modules): + if module_name == _TEST_PACKAGE or module_name.startswith(f"{_TEST_PACKAGE}."): + sys.modules.pop(module_name, None) + + @classmethod + def _optional_modules(cls, *, include_kernel=False): + modules = { + "jax": cls.fake_jax, + "jax.numpy": cls.fake_jnp, + "cutlass": cls.fake_cutlass, + "cutlass.cute": cls.fake_cutlass_cute, + "cutlass.jax": cls.fake_cutlass_jax, + } + if include_kernel: + modules[cls.kernel_module_name] = cls.kernel_module + return mock.patch.dict(sys.modules, modules) + + @staticmethod + def _inputs(bfloat16, float32, int32, *, with_length): + q = _Array((65, 64, 512), bfloat16) + kv = _Array((130, 512), bfloat16) + out = _Array(q.shape, bfloat16) + dout = _Array(q.shape, bfloat16) + lse = _Array((65, 64), float32) + attn_sink = _Array((64,), float32) + topk_idxs = _Array((65, 32), int32) + topk_length = _Array((65,), int32) if with_length else None + return q, kv, out, dout, lse, attn_sink, topk_idxs, topk_length + + @staticmethod + def _fake_call(captured): + def call(launcher, inputs, **options): + captured.update(launcher=launcher, inputs=inputs, **options) + return tuple(_Array(spec.shape, spec.dtype) for spec in options["outputs"]) + + return call + + def test_kernel_module_is_lazy(self): + self.assertNotIn(self.kernel_module_name, sys.modules) + + def test_functional_outputs_and_zero_initialized_workspaces(self): + captured = {} + launcher = object() + inputs = self._inputs( + self.bfloat16, + self.float32, + self.int32, + with_length=False, + ) + q, kv, out, dout, lse, attn_sink, topk_idxs, _ = inputs + + with ( + self._optional_modules(include_kernel=True), + mock.patch.object(self.module, "_launch_without_topk_length", new=launcher), + mock.patch.object( + self.module, + "call_cutedsl", + side_effect=self._fake_call(captured), + ), + ): + result = self.module.sparse_attention_backward_wrapper( + q, + kv, + out, + dout, + lse, + attn_sink, + topk_idxs, + ) + + self.assertEqual(result["dq"].shape, q.shape) + self.assertEqual(result["dkv"].shape, kv.shape) + self.assertEqual(result["d_sink"].shape, attn_sink.shape) + self.assertEqual( + captured["inputs"], + (q, kv, out, dout, lse, attn_sink, topk_idxs), + ) + self.assertIs(captured["launcher"], launcher) + self.assertEqual( + [(spec.name, spec.shape, spec.dtype, spec.fill_value) for spec in captured["outputs"]], + [ + ("dq", q.shape, self.bfloat16, None), + ("dkv", kv.shape, self.bfloat16, 0), + ("d_sink", attn_sink.shape, self.float32, 0.0), + ], + ) + self.assertEqual( + [(spec.name, spec.shape, spec.dtype, spec.fill_value) for spec in captured["workspaces"]], + [ + ("workspace_lse_odo", (1, 64, 72, 8), self.uint8, 0), + ("workspace_dkv", (1, 1, 136, 2048), self.uint8, 0), + ], + ) + self.assertEqual(captured["outputs"][0].tensor_spec.divisibility, 512) + self.assertEqual(len(captured["input_specs"]), 7) + self.assertEqual( + captured["static_args"], + { + "total_seqlen_q": 65, + "total_seqlen_kv": 130, + "num_heads": 64, + "head_dim": 512, + "block_tile": 64, + "softmax_scale": 1.0 / (512**0.5), + }, + ) + + def test_optional_topk_length_is_a_real_kernel_input(self): + captured = {} + launcher = object() + inputs = self._inputs( + self.bfloat16, + self.float32, + self.int32, + with_length=True, + ) + + with ( + self._optional_modules(include_kernel=True), + mock.patch.object(self.module, "_launch_with_topk_length", new=launcher), + mock.patch.object( + self.module, + "call_cutedsl", + side_effect=self._fake_call(captured), + ), + ): + self.module.sparse_attention_backward_wrapper( + *inputs[:-1], + topk_length=inputs[-1], + ) + + self.assertIs(captured["inputs"][-1], inputs[-1]) + self.assertEqual(len(captured["input_specs"]), 8) + self.assertEqual(captured["static_args"]["softmax_scale"], 1.0 / (512**0.5)) + + def test_launchers_preserve_kernel_argument_order(self): + for launcher_name, num_args, topk_length_idx, dq_idx, workspace_idx in ( + ("_launch_with_topk_length", 14, 8, 9, 13), + ("_launch_without_topk_length", 13, None, 8, 12), + ): + with self.subTest(launcher=launcher_name): + _Kernel.instances.clear() + placeholders = [object() for _ in range(num_args)] + with self._optional_modules(include_kernel=True): + getattr(self.module, launcher_name)( + *placeholders, + total_seqlen_q=64, + total_seqlen_kv=128, + num_heads=64, + head_dim=512, + block_tile=64, + softmax_scale=0.125, + ) + + kernel = _Kernel.instances[-1] + self.assertEqual(kernel.configuration, (512, 512, 64)) + args = kernel.calls[-1] + self.assertEqual(args[0], (64, 128, 512, (64, 1))) + if topk_length_idx is None: + self.assertIsNone(args[8]) + else: + self.assertIs(args[8], placeholders[topk_length_idx]) + self.assertIs(args[9], placeholders[dq_idx]) + self.assertIs(args[13], placeholders[workspace_idx]) + self.assertEqual(args[14].value, 0.125) + self.assertIs(args[15], placeholders[0]) + + def test_rejects_non_bf16_or_non_multiple_of_64_heads(self): + inputs = self._inputs( + self.bfloat16, + self.float32, + self.int32, + with_length=False, + ) + q, kv, out, dout, lse, attn_sink, topk_idxs, _ = inputs + bad_q = _Array(q.shape, self.float32) + with ( + self._optional_modules(), + self.assertRaisesRegex(ValueError, "q.dtype"), + ): + self.module.sparse_attention_backward_wrapper( + bad_q, + kv, + out, + dout, + lse, + attn_sink, + topk_idxs, + ) + + bad_q = _Array((65, 32, 512), self.bfloat16) + bad_out = _Array(bad_q.shape, self.bfloat16) + bad_lse = _Array((65, 32), self.float32) + bad_sink = _Array((32,), self.float32) + with ( + self._optional_modules(), + self.assertRaisesRegex(ValueError, "divisible by 64"), + ): + self.module.sparse_attention_backward_wrapper( + bad_q, + kv, + bad_out, + bad_out, + bad_lse, + bad_sink, + topk_idxs, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_gemm.py b/test/python/fe_api/test_jax_gemm.py new file mode 100644 index 000000000..0c326b366 --- /dev/null +++ b/test/python/fe_api/test_jax_gemm.py @@ -0,0 +1,600 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX shape and GPU integration tests for dense GEMM fusions.""" + +from __future__ import annotations + +import pytest + + +def _jax_dependencies(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + pytest.importorskip("cutlass.jax") + return jax, jnp + + +def _sm100_device(jax): + for device in jax.local_devices(): + if device.platform != "gpu": + continue + capability = getattr(device, "compute_capability", None) + if capability is None: + continue + if isinstance(capability, (tuple, list)): + major, minor = (int(value) for value in capability[:2]) + capability_number = major * 10 + minor + elif "." in str(capability): + major, minor = (int(value) for value in str(capability).split(".", 1)) + capability_number = major * 10 + minor + else: + capability_number = int(capability) + if capability_number < 10: + capability_number *= 10 + if capability_number >= 100: + return device + pytest.skip("JAX SM100+ device is not available") + + +def _swiglu_reference(jax, jnp, a, b, alpha): + ab12 = alpha * jnp.einsum( + "lmk,lnk->lmn", + a.astype(jnp.float32), + b.astype(jnp.float32), + ) + batch, m, n = ab12.shape + blocks = 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) + return ab12, input_blocks * gate_blocks * jax.nn.sigmoid(gate_blocks) + + +@pytest.mark.L0 +def test_jax_dense_gemm_abstract_shapes(): + jax, jnp = _jax_dependencies() + + from cudnn.jax import ( + TupleDict, + gemm_amax_wrapper_sm100, + gemm_dsrelu_wrapper_sm100, + gemm_srelu_wrapper_sm100, + gemm_swiglu_wrapper_sm100, + ) + + a_bf16 = jax.ShapeDtypeStruct((1, 128, 128), jnp.bfloat16) + b_bf16 = jax.ShapeDtypeStruct((1, 128, 128), jnp.bfloat16) + swiglu = jax.eval_shape(gemm_swiglu_wrapper_sm100, a_bf16, b_bf16) + assert isinstance(swiglu, TupleDict) + assert tuple(swiglu.keys()) == ("ab12_tensor", "c_tensor", "sfc_tensor", "amax_tensor") + assert swiglu["ab12_tensor"].shape == (1, 128, 128) + assert swiglu["ab12_tensor"].dtype == jnp.float32 + assert swiglu["c_tensor"].shape == (1, 128, 64) + assert swiglu["c_tensor"].dtype == jnp.float16 + assert swiglu["sfc_tensor"] is None + assert swiglu["amax_tensor"] is None + + a_fp8 = jax.ShapeDtypeStruct((2, 256, 512), jnp.float8_e4m3fn) + b_fp8 = jax.ShapeDtypeStruct((2, 256, 512), jnp.float8_e4m3fn) + scales = jax.ShapeDtypeStruct((32, 4, 2, 4, 4, 2), jnp.float8_e8m0fnu) + prob = jax.ShapeDtypeStruct((256, 1, 2), jnp.float32) + c = jax.ShapeDtypeStruct((2, 256, 256), jnp.bfloat16) + + swiglu_mxfp8 = jax.eval_shape( + lambda a, b, sfa, sfb: gemm_swiglu_wrapper_sm100( + a, + b, + sfa_tensor=sfa, + sfb_tensor=sfb, + sf_vec_size=32, + ab12_dtype=jnp.bfloat16, + c_dtype=jnp.bfloat16, + ), + a_fp8, + b_fp8, + scales, + scales, + ) + assert swiglu_mxfp8["ab12_tensor"].shape == (2, 256, 256) + assert swiglu_mxfp8["ab12_tensor"].dtype == jnp.bfloat16 + assert swiglu_mxfp8["c_tensor"].shape == (2, 256, 128) + assert swiglu_mxfp8["c_tensor"].dtype == jnp.bfloat16 + assert swiglu_mxfp8["sfc_tensor"] is None + assert swiglu_mxfp8["amax_tensor"] is None + + a_fp4 = jax.ShapeDtypeStruct((1, 128, 128), jnp.float4_e2m1fn) + b_fp4 = jax.ShapeDtypeStruct((1, 128, 128), jnp.float4_e2m1fn) + fp4_scales = jax.ShapeDtypeStruct((32, 4, 1, 4, 2, 1), jnp.float8_e4m3fn) + swiglu_fp4 = jax.eval_shape( + lambda a, b, sfa, sfb: gemm_swiglu_wrapper_sm100( + a, + b, + sfa_tensor=sfa, + sfb_tensor=sfb, + sf_vec_size=16, + c_dtype=jnp.bfloat16, + ), + a_fp4, + b_fp4, + fp4_scales, + fp4_scales, + ) + assert swiglu_fp4["ab12_tensor"].shape == (1, 128, 128) + assert swiglu_fp4["c_tensor"].shape == (1, 128, 64) + assert swiglu_fp4["amax_tensor"].shape == (1, 1, 1) + assert swiglu_fp4["amax_tensor"].dtype == jnp.float32 + assert swiglu_fp4["sfc_tensor"] is None + + amax = jax.eval_shape( + gemm_amax_wrapper_sm100, + a_fp8, + b_fp8, + scales, + scales, + ) + assert amax["c_tensor"].shape == (2, 256, 256) + assert amax["c_tensor"].dtype == jnp.float32 + assert amax["amax_tensor"].shape == (1, 1, 1) + assert amax["amax_tensor"].dtype == jnp.float32 + + srelu = jax.eval_shape( + gemm_srelu_wrapper_sm100, + a_fp8, + b_fp8, + scales, + scales, + prob, + ) + assert srelu["c_tensor"].shape == (2, 256, 256) + assert srelu["d_tensor"].shape == (2, 256, 256) + assert srelu["c_tensor"].dtype == jnp.bfloat16 + assert srelu["d_tensor"].dtype == jnp.bfloat16 + assert srelu["amax_tensor"] is None + assert srelu["sfd_tensor"] is None + + dsrelu = jax.eval_shape( + gemm_dsrelu_wrapper_sm100, + a_fp8, + b_fp8, + c, + scales, + scales, + prob, + ) + assert dsrelu["d_tensor"].shape == (2, 256, 256) + assert dsrelu["d_tensor"].dtype == jnp.bfloat16 + assert dsrelu["dprob_tensor"].shape == (256, 1, 2) + assert dsrelu["dprob_tensor"].dtype == jnp.float32 + assert dsrelu["amax_tensor"] is None + assert dsrelu["sfd_tensor"] is None + + +@pytest.mark.L0 +def test_jax_gemm_swiglu_jit(): + jax, jnp = _jax_dependencies() + device = _sm100_device(jax) + + from cudnn.jax import gemm_swiglu_wrapper_sm100 + + m = n = k = 128 + batch = 1 + alpha = 0.5 + a = jax.device_put( + jax.random.normal(jax.random.key(0), (batch, k, m), dtype=jnp.bfloat16), + device, + ) + b = jax.device_put( + jax.random.normal(jax.random.key(1), (batch, k, n), dtype=jnp.bfloat16), + device, + ) + + @jax.jit + def run(a, b): + return gemm_swiglu_wrapper_sm100( + a, + b, + alpha=alpha, + ab12_dtype=jnp.float32, + c_dtype=jnp.bfloat16, + a_layout="LKM", + b_layout="LKN", + c_layout="LNM", + ) + + lowered = run.lower(a, b) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + compiled = lowered.compile() + + def reference(a_value, b_value): + ab12_lmn = alpha * jnp.einsum( + "lmk,lnk->lmn", + a_value.transpose(0, 2, 1).astype(jnp.float32), + b_value.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 * gate_blocks * jax.nn.sigmoid(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() + expected_ab12, expected_c = reference(a_value, b) + assert jnp.allclose(result["ab12_tensor"], expected_ab12, atol=5e-2, rtol=2e-2) + assert jnp.allclose( + result["c_tensor"].astype(jnp.float32), + expected_c, + atol=8e-2, + rtol=3e-2, + ) + + +@pytest.mark.L0 +def test_jax_gemm_swiglu_mxfp8_jit(): + jax, jnp = _jax_dependencies() + device = _sm100_device(jax) + + from cudnn.jax import gemm_swiglu_wrapper_sm100 + + m = n = k = 128 + batch = 1 + a = jax.device_put( + jax.random.uniform( + jax.random.key(8), + (batch, m, k), + dtype=jnp.float32, + minval=-0.25, + maxval=0.25, + ).astype(jnp.float8_e4m3fn), + device, + ) + b = jax.device_put( + jax.random.uniform( + jax.random.key(9), + (batch, n, k), + dtype=jnp.float32, + minval=-0.25, + maxval=0.25, + ).astype(jnp.float8_e4m3fn), + device, + ) + scales = jax.device_put( + jnp.ones((32, 4, 1, 4, 1, batch), dtype=jnp.float32).astype(jnp.float8_e8m0fnu), + device, + ) + + @jax.jit + def run(a, b, sfa, sfb): + return gemm_swiglu_wrapper_sm100( + a, + b, + sfa_tensor=sfa, + sfb_tensor=sfb, + sf_vec_size=32, + ab12_dtype=jnp.bfloat16, + c_dtype=jnp.bfloat16, + vector_f32=True, + ab12_stages=3, + ) + + lowered = run.lower(a, b, scales, scales) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + result = lowered.compile()(a, b, scales, scales) + result["c_tensor"].block_until_ready() + + expected_ab12, expected_c = _swiglu_reference(jax, jnp, a, b, 1.0) + assert jnp.allclose( + result["ab12_tensor"].astype(jnp.float32), + expected_ab12, + atol=3e-1, + rtol=5e-2, + ) + assert jnp.allclose( + result["c_tensor"].astype(jnp.float32), + expected_c, + atol=5e-1, + rtol=8e-2, + ) + assert result["sfc_tensor"] is None + assert result["amax_tensor"] is None + + +@pytest.mark.L0 +def test_jax_gemm_swiglu_native_fp4_amax_jit(): + jax, jnp = _jax_dependencies() + device = _sm100_device(jax) + + from cudnn.jax import gemm_swiglu_wrapper_sm100 + + m = n = k = 128 + batch = 1 + a = jax.device_put( + jax.random.uniform( + jax.random.key(10), + (batch, m, k), + dtype=jnp.float32, + minval=-1.0, + maxval=1.0, + ).astype(jnp.float4_e2m1fn), + device, + ) + b = jax.device_put( + jax.random.uniform( + jax.random.key(11), + (batch, n, k), + dtype=jnp.float32, + minval=-1.0, + maxval=1.0, + ).astype(jnp.float4_e2m1fn), + device, + ) + scales = jax.device_put( + jnp.ones((32, 4, 1, 4, 2, batch), dtype=jnp.float32).astype(jnp.float8_e4m3fn), + device, + ) + + @jax.jit + def run(a, b, sfa, sfb): + return gemm_swiglu_wrapper_sm100( + a, + b, + sfa_tensor=sfa, + sfb_tensor=sfb, + sf_vec_size=16, + ab12_dtype=jnp.float32, + c_dtype=jnp.bfloat16, + ) + + lowered = run.lower(a, b, scales, scales) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + compiled = lowered.compile() + + result = compiled(a, b, scales, scales) + result["amax_tensor"].block_until_ready() + expected_ab12, expected_c = _swiglu_reference(jax, jnp, a, b, 1.0) + expected_amax = jnp.max(jnp.abs(expected_c)).reshape(1, 1, 1) + assert jnp.allclose(result["ab12_tensor"], expected_ab12, atol=2e-1, rtol=3e-2) + assert jnp.allclose( + result["c_tensor"].astype(jnp.float32), + expected_c, + atol=5e-1, + rtol=5e-2, + ) + assert jnp.allclose(result["amax_tensor"], expected_amax, atol=5e-1, rtol=5e-2) + assert result["sfc_tensor"] is None + + zero_result = compiled(a, jnp.zeros_like(b), scales, scales) + zero_result["amax_tensor"].block_until_ready() + assert jnp.array_equal(zero_result["c_tensor"], jnp.zeros_like(zero_result["c_tensor"])) + assert jnp.array_equal(zero_result["amax_tensor"], jnp.zeros_like(zero_result["amax_tensor"])) + + +@pytest.mark.L0 +def test_jax_gemm_amax_jit_reinitializes_reduction(): + jax, jnp = _jax_dependencies() + device = _sm100_device(jax) + + from cudnn.jax import gemm_amax_wrapper_sm100 + + m = n = k = 128 + batch = 1 + fp8_dtype = jnp.float8_e4m3fn + a = jax.device_put( + jax.random.uniform( + jax.random.key(2), + (batch, m, k), + dtype=jnp.float32, + minval=-1.0, + maxval=1.0, + ).astype(fp8_dtype), + device, + ) + b = jax.device_put( + jax.random.uniform( + jax.random.key(3), + (batch, n, k), + dtype=jnp.float32, + minval=-1.0, + maxval=1.0, + ).astype(fp8_dtype), + device, + ) + sf_k = k // 32 + + def pack_scales(canonical): + return canonical.reshape(1, 4, 32, 1, 4, batch).transpose(2, 1, 0, 4, 3, 5).astype(jnp.float8_e8m0fnu) + + row = jnp.arange(m, dtype=jnp.int32)[:, None, None] + scale_column = jnp.arange(sf_k, dtype=jnp.int32)[None, :, None] + sfa_canonical = jnp.where((row + scale_column) % 2, 2.0, 1.0) + sfb_canonical = jnp.where((row + 2 * scale_column) % 2, 1.0, 0.5) + sfa = jax.device_put(pack_scales(sfa_canonical), device) + sfb = jax.device_put(pack_scales(sfb_canonical), device) + + @jax.jit + def run(a, b, sfa, sfb): + return gemm_amax_wrapper_sm100(a, b, sfa, sfb) + + lowered = run.lower(a, b, sfa, sfb) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + compiled = lowered.compile() + + result = compiled(a, b, sfa, sfb) + result["c_tensor"].block_until_ready() + expected = jnp.einsum( + "lmk,lnk->lmn", + a.astype(jnp.float32) * jnp.repeat(sfa_canonical, 32, axis=1).transpose(2, 0, 1), + b.astype(jnp.float32) * jnp.repeat(sfb_canonical, 32, axis=1).transpose(2, 0, 1), + ) + expected_amax = jnp.max(jnp.abs(expected)).reshape(1, 1, 1) + assert jnp.allclose(result["c_tensor"], expected, atol=2e-1, rtol=3e-2) + assert jnp.allclose(result["amax_tensor"], expected_amax, atol=5e-1, rtol=3e-2) + + zero_result = compiled(a, jnp.zeros_like(b), sfa, sfb) + zero_result["amax_tensor"].block_until_ready() + assert jnp.array_equal(zero_result["c_tensor"], jnp.zeros_like(zero_result["c_tensor"])) + assert jnp.array_equal( + zero_result["amax_tensor"], + jnp.zeros_like(zero_result["amax_tensor"]), + ) + + +@pytest.mark.L0 +def test_jax_gemm_squared_relu_forward_and_backward_jit(): + jax, jnp = _jax_dependencies() + device = _sm100_device(jax) + + from cudnn.jax import gemm_dsrelu_wrapper_sm100, gemm_srelu_wrapper_sm100 + + m = n = k = 128 + batch = 1 + fp8_dtype = jnp.float8_e4m3fn + a = jax.device_put( + jax.random.uniform( + jax.random.key(4), + (batch, m, k), + dtype=jnp.float32, + minval=-0.5, + maxval=0.5, + ).astype(fp8_dtype), + device, + ) + b = jax.device_put( + jax.random.uniform( + jax.random.key(5), + (batch, n, k), + dtype=jnp.float32, + minval=-0.5, + maxval=0.5, + ).astype(fp8_dtype), + device, + ) + c_input = jax.device_put( + jax.random.uniform( + jax.random.key(6), + (batch, m, n), + dtype=jnp.float32, + minval=-0.5, + maxval=0.5, + ).astype(jnp.bfloat16), + device, + ) + prob = jax.device_put( + jax.random.uniform( + jax.random.key(7), + (m, 1, batch), + dtype=jnp.float32, + minval=0.25, + maxval=1.0, + ), + device, + ) + scales = jax.device_put( + jnp.ones((32, 4, 1, 4, 1, batch), dtype=jnp.float32).astype(jnp.float8_e8m0fnu), + device, + ) + + @jax.jit + def forward(a, b, sfa, sfb, prob): + return gemm_srelu_wrapper_sm100( + a, + b, + sfa, + sfb, + prob, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + + @jax.jit + def backward(a, b, c, sfa, sfb, prob): + return gemm_dsrelu_wrapper_sm100( + a, + b, + c, + sfa, + sfb, + prob, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + + forward_lowered = forward.lower(a, b, scales, scales, prob) + backward_lowered = backward.lower(a, b, c_input, scales, scales, prob) + for lowered in (forward_lowered, backward_lowered): + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + x = jnp.einsum( + "lmk,lnk->lmn", + a.astype(jnp.float32), + b.astype(jnp.float32), + ) + relu_x = jnp.maximum(x, 0.0) + prob_lm1 = prob.transpose(2, 0, 1) + + forward_result = forward_lowered.compile()(a, b, scales, scales, prob) + forward_result["d_tensor"].block_until_ready() + expected_forward = jnp.square(relu_x) * prob_lm1 + assert jnp.allclose( + forward_result["c_tensor"].astype(jnp.float32), + x, + atol=2e-1, + rtol=3e-2, + ) + assert jnp.allclose( + forward_result["d_tensor"].astype(jnp.float32), + expected_forward, + atol=4e-1, + rtol=5e-2, + ) + + backward_compiled = backward_lowered.compile() + backward_result = backward_compiled( + a, + b, + c_input, + scales, + scales, + prob, + ) + backward_result["dprob_tensor"].block_until_ready() + c_f32 = c_input.astype(jnp.float32) + expected_d = c_f32 * prob_lm1 * 2 * relu_x + expected_dprob = jnp.sum(c_f32 * jnp.square(relu_x), axis=2, keepdims=True).transpose(1, 2, 0) + 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_input), + 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"]), + ) diff --git a/test/python/fe_api/test_jax_gemm_contract.py b/test/python/fe_api/test_jax_gemm_contract.py new file mode 100644 index 000000000..0b0824bdf --- /dev/null +++ b/test/python/fe_api/test_jax_gemm_contract.py @@ -0,0 +1,1116 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for the JAX dense GEMM wrappers.""" + +from __future__ import annotations + +import ast +import importlib +import sys +import types +import unittest +from importlib.machinery import ModuleSpec +from pathlib import Path +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +def _identity_jit(fn=None, **_kwargs): + return (lambda decorated_fn: decorated_fn) if fn is None else fn + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CUDNN_ROOT = _REPO_ROOT / "python" / "cudnn" +_TEST_PACKAGE = "cudnn_frontend_jax_gemm_contract_test" +_COMMON_KERNEL_CAPABILITIES = { + "MMA_TILER_M", + "MMA_TILER_N", + "TWO_CTA_MMA_TILER_M", + "MAX_CLUSTER_CTAS", + "MAX_CLUSTER_DIMENSION", +} +_COMMON_KERNEL_METHODS = {"require_mma_tiler", "require_cluster_shape"} + + +def _kernel_capabilities( + path: Path, + class_name: str, + field_names: set[str], + method_names: set[str], +) -> type: + tree = ast.parse(path.read_text(), filename=str(path)) + kernel_class = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == class_name) + field_nodes = [] + found_fields = set() + method_nodes = [] + found_methods = set() + for node in kernel_class.body: + if isinstance(node, ast.Assign) and len(node.targets) == 1: + target = node.targets[0] + elif isinstance(node, ast.AnnAssign): + target = node.target + else: + target = None + if isinstance(target, ast.Name) and target.id in field_names: + field_nodes.append(node) + found_fields.add(target.id) + elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in method_names: + method_nodes.append(node) + found_methods.add(node.name) + + missing = field_names - found_fields + if missing: + raise AssertionError(f"{class_name} is missing public kernel capabilities: {sorted(missing)}") + missing = method_names - found_methods + if missing: + raise AssertionError(f"{class_name} is missing public kernel methods: {sorted(missing)}") + + def validation_helper(name): + def call(*args, **kwargs): + validation = importlib.import_module(f"{_TEST_PACKAGE}.gemm_validation") + return getattr(validation, name)(*args, **kwargs) + + return call + + surface_class = ast.ClassDef( + name=class_name, + bases=[], + keywords=[], + body=[*field_nodes, *method_nodes], + decorator_list=[], + ) + surface_module = ast.fix_missing_locations( + ast.Module( + body=[ + ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), + surface_class, + ], + type_ignores=[], + ) + ) + namespace = { + "__name__": __name__, + "_require_mma_tiler": validation_helper("require_mma_tiler"), + "_require_cluster_shape": validation_helper("require_cluster_shape"), + } + exec(compile(surface_module, str(path), "exec"), namespace) + return namespace[class_name] + + +class _DType: + def __init__(self, name, itemsize): + 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 + + +class _TensorSpec: + def __init__( + self, + *, + layout=None, + mode=None, + static=None, + ptr_assumed_align=256, + divisibility=None, + ): + self.layout = layout + self.mode = mode + self.static = static + self.ptr_assumed_align = ptr_assumed_align + self.divisibility = divisibility + + +class JaxGemmContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.bfloat16 = _DType("bfloat16", 2) + cls.float16 = _DType("float16", 2) + cls.float32 = _DType("float32", 4) + cls.float8_e4m3fn = _DType("float8_e4m3fn", 1) + cls.float8_e5m2 = _DType("float8_e5m2", 1) + cls.float8_e8m0fnu = _DType("float8_e8m0fnu", 1) + cls.float4_e2m1fn = _DType("float4_e2m1fn", 1) + cls.uint8 = _DType("uint8", 1) + + cls.fake_jnp = types.ModuleType("jax.numpy") + for name in ( + "bfloat16", + "float16", + "float32", + "float8_e4m3fn", + "float8_e5m2", + "float8_e8m0fnu", + "float4_e2m1fn", + "uint8", + ): + setattr(cls.fake_jnp, name, getattr(cls, name)) + cls.fake_jnp.dtype = lambda value: value + + cls.fake_jax = types.ModuleType("jax") + cls.fake_jax.__path__ = [] + cls.fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + cls.fake_jax.numpy = cls.fake_jnp + cls.fake_jax.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, + ) + + cls.fake_cutlass = types.ModuleType("cutlass") + cls.fake_cutlass.__path__ = [] + cls.fake_cutlass.Constexpr = object + cls.fake_cutlass.Float32 = float + cls.fake_cutlass.utils = types.SimpleNamespace(HardwareInfo=lambda: types.SimpleNamespace(get_max_active_clusters=lambda _cluster_size: 8)) + cls.fake_cutlass_jax = types.ModuleType("cutlass.jax") + cls.fake_cutlass_jax.is_available = lambda: True + cls.fake_cutlass_jax.TensorSpec = _TensorSpec + cls.fake_cutlass_jax.jax_to_cutlass_dtype = lambda dtype: f"cutlass.{dtype.name}" + cls.fake_cutlass.jax = cls.fake_cutlass_jax + cls.fake_cutlass_cute = types.ModuleType("cutlass.cute") + cls.fake_cutlass_cute.jit = _identity_jit + cls.fake_cutlass_cute.where = lambda condition, yes, no: (yes if condition else no) + cls.fake_cutlass_cute.full_like = lambda value, fill: fill + cls.fake_cutlass.cute = cls.fake_cutlass_cute + + cls.swiglu_capabilities = _kernel_capabilities( + _CUDNN_ROOT / "gemm_swiglu" / "dense_gemm_persistent_swiglu.py", + "PersistentDenseGemmKernel", + _COMMON_KERNEL_CAPABILITIES + | { + "SINGLE_CTA_CLUSTER_SHAPE", + "SWIGLU_BLOCK_COLUMNS", + "SWIGLU_BLOCKS_PER_PAIR", + }, + _COMMON_KERNEL_METHODS | {"get_output_n"}, + ) + cls.quantized_swiglu_capabilities = _kernel_capabilities( + _CUDNN_ROOT / "gemm_swiglu" / "dense_blockscaled_gemm_persistent_swiglu_interleaved_quant.py", + "Sm100BlockScaledPersistentDenseGemmKernel", + _COMMON_KERNEL_CAPABILITIES | {"SF_VEC_SIZES", "SWIGLU_BLOCK_COLUMNS", "SWIGLU_BLOCKS_PER_PAIR"}, + _COMMON_KERNEL_METHODS | {"get_output_n"}, + ) + cls.quantized_swiglu_capabilities.calls = [] + + def init_quantized_kernel(instance, **kwargs): + instance.config = kwargs + + def call_quantized_kernel(instance, *args): + type(instance).calls.append((instance.config, args)) + + cls.quantized_swiglu_capabilities.__init__ = init_quantized_kernel + cls.quantized_swiglu_capabilities.__call__ = call_quantized_kernel + cls.srelu_capabilities = _kernel_capabilities( + _CUDNN_ROOT / "gemm_srelu" / "dense_blockscaled_gemm_persistent_srelu_quant.py", + "Sm100BlockScaledPersistentDenseGemmKernel", + _COMMON_KERNEL_CAPABILITIES, + _COMMON_KERNEL_METHODS, + ) + cls.dsrelu_capabilities = _kernel_capabilities( + _CUDNN_ROOT / "gemm_dsrelu" / "dense_blockscaled_gemm_persistent_dsrelu_quant.py", + "Sm100BlockScaledPersistentDenseGemmKernel", + _COMMON_KERNEL_CAPABILITIES, + _COMMON_KERNEL_METHODS, + ) + cls.amax_capabilities = _kernel_capabilities( + _CUDNN_ROOT / "gemm_amax" / "dense_blockscaled_gemm_persistent_amax.py", + "Sm100BlockScaledPersistentDenseGemmKernel", + _COMMON_KERNEL_CAPABILITIES | {"KNOWN_HANG_MMA_TILER_M", "SF_VEC_SIZES"}, + _COMMON_KERNEL_METHODS, + ) + cls.fake_kernel_modules = {} + for module_suffix, class_name, capabilities in ( + ( + "gemm_swiglu.dense_gemm_persistent_swiglu", + "PersistentDenseGemmKernel", + cls.swiglu_capabilities, + ), + ( + "gemm_swiglu.dense_blockscaled_gemm_persistent_swiglu_interleaved_quant", + "Sm100BlockScaledPersistentDenseGemmKernel", + cls.quantized_swiglu_capabilities, + ), + ( + "gemm_amax.dense_blockscaled_gemm_persistent_amax", + "Sm100BlockScaledPersistentDenseGemmKernel", + cls.amax_capabilities, + ), + ( + "gemm_srelu.dense_blockscaled_gemm_persistent_srelu_quant", + "Sm100BlockScaledPersistentDenseGemmKernel", + cls.srelu_capabilities, + ), + ( + "gemm_dsrelu.dense_blockscaled_gemm_persistent_dsrelu_quant", + "Sm100BlockScaledPersistentDenseGemmKernel", + cls.dsrelu_capabilities, + ), + ): + module_name = f"{_TEST_PACKAGE}.{module_suffix}" + module = types.ModuleType(module_name) + setattr(module, class_name, capabilities) + cls.fake_kernel_modules[module_name] = module + + parent = types.ModuleType(_TEST_PACKAGE) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = _TEST_PACKAGE + sys.modules[_TEST_PACKAGE] = parent + with cls._optional_modules(): + importlib.import_module(f"{_TEST_PACKAGE}.jax") + cls.gemm = importlib.import_module(f"{_TEST_PACKAGE}._jax.gemm") + cls.swiglu = importlib.import_module(f"{_TEST_PACKAGE}.gemm_swiglu.jax") + cls.amax = importlib.import_module(f"{_TEST_PACKAGE}.gemm_amax.jax") + cls.srelu = importlib.import_module(f"{_TEST_PACKAGE}.gemm_srelu.jax") + cls.dsrelu = importlib.import_module(f"{_TEST_PACKAGE}.gemm_dsrelu.jax") + + @classmethod + def tearDownClass(cls): + for module_name in tuple(sys.modules): + if module_name == _TEST_PACKAGE or module_name.startswith(f"{_TEST_PACKAGE}."): + sys.modules.pop(module_name, None) + + @classmethod + def _optional_modules(cls): + modules = { + "jax": cls.fake_jax, + "jax.numpy": cls.fake_jnp, + "cutlass": cls.fake_cutlass, + "cutlass.cute": cls.fake_cutlass_cute, + "cutlass.jax": cls.fake_cutlass_jax, + } + modules.update(cls.fake_kernel_modules) + return mock.patch.dict( + sys.modules, + modules, + ) + + @staticmethod + def _fake_call(captured): + def call(launcher, inputs, **options): + captured.update(launcher=launcher, inputs=inputs, **options) + return tuple(_Array(spec.shape, spec.dtype) for spec in options["outputs"]) + + return call + + def test_swiglu_declares_logical_outputs_and_physical_layouts(self): + captured = {} + a = _Array((2, 128, 256), self.bfloat16) + b = _Array((2, 128, 128), self.bfloat16) + with ( + self._optional_modules(), + mock.patch.object( + self.swiglu, + "call_cutedsl", + side_effect=self._fake_call(captured), + ), + ): + result = self.swiglu.gemm_swiglu_wrapper_sm100( + a, + b, + c_layout="LNM", + ab12_dtype=_Array((), self.float32), + c_dtype=_Array((), self.bfloat16), + acc_dtype=_Array((), self.float32), + a_layout="LKM", + b_layout="LKN", + ) + + self.assertIs(captured["launcher"], self.swiglu._launch) + self.assertEqual(captured["inputs"], (a, b)) + self.assertEqual( + captured["static_args"], + { + "alpha": 1.0, + "acc_dtype": self.float32, + "mma_tiler_mn": (128, 128), + "cluster_shape_mn": (1, 1), + "cluster_overlap_margin": 0, + }, + ) + self.assertNotIn("workspaces", captured) + self.assertEqual( + [(spec.name, spec.shape, spec.dtype, spec.fill_value) for spec in captured["outputs"]], + [ + ("ab12_tensor", (2, 128, 256), self.float32, None), + ("c_tensor", (2, 64, 256), self.bfloat16, None), + ], + ) + self.assertEqual([spec.layout for spec in captured["input_specs"]], [(2, 1, 0), (2, 1, 0)]) + self.assertEqual([spec.mode for spec in captured["input_specs"]], [(2, 1, 0), (2, 1, 0)]) + self.assertEqual( + [spec.tensor_spec.layout for spec in captured["outputs"]], + [(2, 1, 0), (2, 1, 0)], + ) + self.assertEqual([spec.tensor_spec.mode for spec in captured["outputs"]], [(2, 1, 0), (2, 1, 0)]) + self.assertEqual(tuple(result.keys()), ("ab12_tensor", "c_tensor", "sfc_tensor", "amax_tensor")) + self.assertIsNone(result["sfc_tensor"]) + self.assertIsNone(result["amax_tensor"]) + + def test_layout_strings_map_public_shapes_to_canonical_kernel_metadata(self): + cases = ( + (self.swiglu.gemm_a_tensor_spec, "lmk", (2, 3, 5), (3, 5, 2), (5, 1, 15), (1, 2, 0)), + (self.swiglu.gemm_a_tensor_spec, "LKM", (2, 5, 3), (3, 5, 2), (1, 3, 15), (2, 1, 0)), + (self.swiglu.gemm_b_tensor_spec, "LNK", (2, 7, 5), (7, 5, 2), (5, 1, 35), (1, 2, 0)), + (self.swiglu.gemm_b_tensor_spec, "LKN", (2, 5, 7), (7, 5, 2), (1, 7, 35), (2, 1, 0)), + (self.swiglu.gemm_c_tensor_spec, "LMN", (2, 3, 7), (3, 7, 2), (7, 1, 21), (1, 2, 0)), + (self.swiglu.gemm_c_tensor_spec, "LNM", (2, 7, 3), (3, 7, 2), (1, 3, 21), (2, 1, 0)), + ) + + with self._optional_modules(): + for factory, layout, array_shape, kernel_shape, stride, mode in cases: + with self.subTest(layout=layout): + spec = factory(layout) + desc = self.swiglu.JaxTensorDesc.from_value( + _Array(array_shape, self.bfloat16), + tensor_spec=spec, + ) + self.assertEqual(spec.layout, (2, 1, 0)) + self.assertEqual(spec.mode, mode) + self.assertEqual(desc.array_shape, array_shape) + self.assertEqual(desc.shape, kernel_shape) + self.assertEqual(desc.stride, stride) + + for factory, layout in ( + (self.swiglu.gemm_a_tensor_spec, "MKL"), + (self.swiglu.gemm_b_tensor_spec, "NKL"), + (self.swiglu.gemm_c_tensor_spec, "MNL"), + ): + with self.subTest(layout=layout), self.assertRaisesRegex(ValueError, "must be one of"): + factory(layout) + + a_desc = self.swiglu.JaxTensorDesc.from_value( + _Array((2, 3, 5), self.bfloat16), + tensor_spec=self.gemm.gemm_a_tensor_spec("LMK"), + ) + with self.assertRaisesRegex(ValueError, "descriptor layout does not match"): + self.gemm.as_gemm_tensor_desc( + "a_tensor", + a_desc, + self.gemm.gemm_a_tensor_spec("LKM"), + ) + + def test_kernel_capability_fields_are_public(self): + for capabilities in ( + self.swiglu_capabilities, + self.quantized_swiglu_capabilities, + self.amax_capabilities, + self.srelu_capabilities, + self.dsrelu_capabilities, + ): + with self.subTest(capabilities=capabilities): + self.assertLessEqual(_COMMON_KERNEL_CAPABILITIES, vars(capabilities).keys()) + for method_name in _COMMON_KERNEL_METHODS: + self.assertTrue(callable(getattr(capabilities, method_name))) + + def test_swiglu_kernel_validates_output_n(self): + for kernel in (self.swiglu_capabilities, self.quantized_swiglu_capabilities): + with self.subTest(kernel=kernel): + self.assertEqual(kernel.get_output_n(128), 64) + with self.assertRaisesRegex(ValueError, "32-column SwiGLU block pairs"): + kernel.get_output_n(96) + + def test_kernel_surface_applies_kernel_specific_configuration_rules(self): + with self.assertRaisesRegex(NotImplementedError, "currently hangs"): + self.amax_capabilities.require_mma_tiler((256, 128)) + with self.assertRaisesRegex(ValueError, "single-CTA MMA tile"): + self.swiglu_capabilities.require_cluster_shape((2, 2), mma_tiler_mn=(128, 128)) + + def test_amax_declares_scale_layout_and_initialized_reduction(self): + captured = {} + a = _Array((1, 128, 128), self.float8_e4m3fn) + b = _Array((1, 128, 128), self.float8_e4m3fn) + sfa = _Array((32, 4, 1, 4, 1, 1), self.float8_e8m0fnu) + sfb = _Array((32, 4, 1, 4, 1, 1), self.float8_e8m0fnu) + with ( + self._optional_modules(), + mock.patch.object( + self.amax, + "call_cutedsl", + side_effect=self._fake_call(captured), + ), + ): + result = self.amax.gemm_amax_wrapper_sm100( + a, + b, + sfa, + sfb, + c_layout="LMN", + c_dtype=self.float32, + ) + + self.assertEqual(captured["inputs"], (a, b, sfa, sfb)) + self.assertIs(captured["launcher"], self.amax._launch) + self.assertEqual( + captured["static_args"], + { + "sf_vec_size": 32, + "mma_tiler_mn": (128, 128), + "cluster_shape_mn": (1, 1), + "cluster_overlap_margin": 0, + }, + ) + self.assertEqual( + [spec.layout for spec in captured["input_specs"]], + [(2, 1, 0), (2, 1, 0), (2, 1, 4, 0, 3, 5), (2, 1, 4, 0, 3, 5)], + ) + c_spec, amax_spec = captured["outputs"] + self.assertEqual( + (c_spec.name, c_spec.shape, c_spec.dtype), + ("c_tensor", (1, 128, 128), self.float32), + ) + self.assertEqual(c_spec.tensor_spec.layout, (2, 1, 0)) + self.assertEqual(c_spec.tensor_spec.mode, (1, 2, 0)) + self.assertEqual( + (amax_spec.name, amax_spec.shape, amax_spec.dtype, amax_spec.fill_value), + ("amax_tensor", (1, 1, 1), self.float32, float("-inf")), + ) + self.assertEqual(tuple(result.keys()), ("c_tensor", "amax_tensor")) + + def test_swiglu_declares_mxfp8_quantized_call(self): + captured = {} + a = _Array((1, 128, 128), self.float8_e4m3fn) + b = _Array((1, 128, 128), self.float8_e4m3fn) + sfa = _Array((32, 4, 1, 4, 1, 1), self.float8_e8m0fnu) + sfb = _Array((32, 4, 1, 4, 1, 1), self.float8_e8m0fnu) + with ( + self._optional_modules(), + mock.patch.object( + self.swiglu, + "call_cutedsl", + side_effect=self._fake_call(captured), + ), + ): + result = self.swiglu.gemm_swiglu_wrapper_sm100( + a, + b, + alpha=0.5, + ab12_dtype=self.bfloat16, + c_dtype=self.bfloat16, + sfa_tensor=sfa, + sfb_tensor=sfb, + sf_vec_size=32, + vector_f32=True, + ab12_stages=3, + ) + + self.assertIs(captured["launcher"], self.swiglu._launch_quantized) + self.assertEqual(captured["inputs"], (a, b, sfa, sfb)) + self.assertEqual( + captured["static_args"], + { + "alpha": 0.5, + "sf_vec_size": 32, + "mma_tiler_mn": (128, 128), + "cluster_shape_mn": (1, 1), + "vector_f32": True, + "ab12_stages": 3, + "has_amax": False, + "cluster_overlap_margin": 0, + }, + ) + self.assertEqual( + [spec.layout for spec in captured["input_specs"]], + [ + (2, 1, 0), + (2, 1, 0), + (2, 1, 4, 0, 3, 5), + (2, 1, 4, 0, 3, 5), + ], + ) + self.assertEqual( + [(spec.name, spec.shape, spec.dtype, spec.fill_value) for spec in captured["outputs"]], + [ + ("ab12_tensor", (1, 128, 128), self.bfloat16, None), + ("c_tensor", (1, 128, 64), self.bfloat16, None), + ], + ) + self.assertEqual(tuple(result.keys()), ("ab12_tensor", "c_tensor", "sfc_tensor", "amax_tensor")) + self.assertIsNone(result["sfc_tensor"]) + self.assertIsNone(result["amax_tensor"]) + + def test_swiglu_declares_native_fp4_amax_call(self): + captured = {} + a = _Array((1, 128, 128), self.float4_e2m1fn) + b = _Array((1, 128, 128), self.float4_e2m1fn) + sfa = _Array((32, 4, 1, 4, 2, 1), self.float8_e4m3fn) + sfb = _Array((32, 4, 1, 4, 2, 1), self.float8_e4m3fn) + with ( + self._optional_modules(), + mock.patch.object( + self.swiglu, + "call_cutedsl", + side_effect=self._fake_call(captured), + ), + ): + result = self.swiglu.gemm_swiglu_wrapper_sm100( + a, + b, + ab12_dtype=self.bfloat16, + c_dtype=self.bfloat16, + sfa_tensor=sfa, + sfb_tensor=sfb, + sf_vec_size=16, + a_layout="LMK", + b_layout="LNK", + c_layout="LMN", + ) + + self.assertIs(captured["launcher"], self.swiglu._launch_quantized) + self.assertEqual(captured["inputs"], (a, b, sfa, sfb)) + self.assertEqual( + captured["static_args"], + { + "alpha": 1.0, + "sf_vec_size": 16, + "mma_tiler_mn": (128, 128), + "cluster_shape_mn": (1, 1), + "vector_f32": False, + "ab12_stages": 4, + "has_amax": True, + "cluster_overlap_margin": 0, + }, + ) + self.assertEqual( + [(spec.name, spec.shape, spec.dtype, spec.fill_value) for spec in captured["outputs"]], + [ + ("ab12_tensor", (1, 128, 128), self.bfloat16, None), + ("c_tensor", (1, 128, 64), self.bfloat16, None), + ("amax_tensor", (1, 1, 1), self.float32, float("-inf")), + ], + ) + self.assertEqual(tuple(result.keys()), ("ab12_tensor", "c_tensor", "sfc_tensor", "amax_tensor")) + self.assertIsNone(result["sfc_tensor"]) + self.assertEqual(result["amax_tensor"].shape, (1, 1, 1)) + + def test_swiglu_revalidation_rebuilds_conditional_outputs(self): + a = _Array((1, 128, 128), self.float4_e2m1fn) + b = _Array((1, 128, 128), self.float4_e2m1fn) + sfa = _Array((32, 4, 1, 4, 2, 1), self.float8_e4m3fn) + sfb = _Array((32, 4, 1, 4, 2, 1), self.float8_e4m3fn) + + with self._optional_modules(): + operation = self.swiglu.GemmSwigluSm100( + a, + b, + sample_sfa=sfa, + sample_sfb=sfb, + sf_vec_size=16, + ab12_dtype=self.bfloat16, + c_dtype=self.bfloat16, + ) + self.assertTrue(operation.check_support()) + self.assertIsNotNone(operation.amax_desc) + + operation._c_dtype = self.float16 + self.assertTrue(operation.check_support()) + self.assertIsNone(operation.amax_desc) + + def test_swiglu_quantized_launcher_uses_kernel_argument_order(self): + kernel = self.quantized_swiglu_capabilities + with self._optional_modules(): + for optional_outputs, has_amax in (((), False), (("amax",), True)): + with self.subTest(has_amax=has_amax): + kernel.calls.clear() + self.swiglu._launch_quantized( + "stream", + "a", + "b", + "sfa", + "sfb", + "ab12", + "c", + *optional_outputs, + alpha=0.5, + sf_vec_size=16, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + vector_f32=True, + ab12_stages=3, + has_amax=has_amax, + cluster_overlap_margin=1, + ) + + self.assertEqual(len(kernel.calls), 1) + config, args = kernel.calls[0] + self.assertEqual( + config, + { + "sf_vec_size": 16, + "mma_tiler_mn": (128, 128), + "cluster_shape_mn": (1, 1), + "vector_f32": True, + "ab12_stages": 3, + }, + ) + self.assertEqual( + args, + ( + "a", + "b", + "sfa", + "sfb", + "c", + "ab12", + "amax" if has_amax else None, + None, + None, + 0.5, + 7, + "stream", + ), + ) + + def test_swiglu_rejects_unused_nondefault_configuration(self): + a = _Array((1, 128, 128), self.bfloat16) + b = _Array((1, 128, 128), self.bfloat16) + + with self._optional_modules(): + for option, value, expected in ( + ("sf_vec_size", 32, "sf_vec_size applies only"), + ("vector_f32", True, "vector_f32 applies only"), + ("ab12_stages", 3, "ab12_stages applies only"), + ): + with ( + self.subTest(option=option), + self.assertRaisesRegex(NotImplementedError, expected), + mock.patch.object(self.swiglu, "call_cutedsl") as lower, + ): + self.swiglu.gemm_swiglu_wrapper_sm100(a, b, **{option: value}) + lower.assert_not_called() + + def test_swiglu_rejects_unsupported_epilogue_and_incomplete_2cta_tiles(self): + a = _Array((1, 128, 128), self.bfloat16) + b = _Array((1, 64, 128), self.bfloat16) + + with self._optional_modules(): + for mma_tiler_mn, message in ( + ((128, 32), r"N in \{64, 128, 192, 256\}"), + ((256, 64), "M must be divisible by 256"), + ): + with ( + self.subTest(mma_tiler_mn=mma_tiler_mn), + self.assertRaisesRegex(ValueError, message), + mock.patch.object(self.swiglu, "call_cutedsl") as lower, + ): + self.swiglu.gemm_swiglu_wrapper_sm100(a, b, mma_tiler_mn=mma_tiler_mn) + lower.assert_not_called() + + def test_amax_validates_scale_shapes_before_lowering(self): + a = _Array((1, 128, 128), self.float8_e5m2) + b = _Array((1, 128, 128), self.float8_e5m2) + bad_sfa = _Array((32, 4, 1, 4, 2, 1), self.float8_e8m0fnu) + sfb = _Array((32, 4, 1, 4, 1, 1), self.float8_e8m0fnu) + with ( + self._optional_modules(), + self.assertRaisesRegex(ValueError, "sfa_tensor must have shape"), + mock.patch.object(self.amax, "call_cutedsl") as lower, + ): + self.amax.gemm_amax_wrapper_sm100(a, b, bad_sfa, sfb) + lower.assert_not_called() + + def test_srelu_and_dsrelu_declare_functional_outputs(self): + a = _Array((2, 384, 512), self.float8_e4m3fn) + b = _Array((2, 256, 512), self.float8_e4m3fn) + c = _Array((2, 384, 256), self.bfloat16) + sfa = _Array((32, 4, 3, 4, 4, 2), self.float8_e8m0fnu) + sfb = _Array((32, 4, 2, 4, 4, 2), self.float8_e8m0fnu) + prob = _Array((384, 1, 2), self.float32) + srelu_captured = {} + dsrelu_captured = {} + + with ( + self._optional_modules(), + mock.patch.object( + self.srelu, + "call_cutedsl", + side_effect=self._fake_call(srelu_captured), + ), + mock.patch.object( + self.dsrelu, + "call_cutedsl", + side_effect=self._fake_call(dsrelu_captured), + ), + ): + srelu_result = self.srelu.gemm_srelu_wrapper_sm100( + a, + b, + sfa, + sfb, + prob, + ) + dsrelu_result = self.dsrelu.gemm_dsrelu_wrapper_sm100( + a, + b, + c, + sfa, + sfb, + prob, + ) + + self.assertEqual(srelu_captured["inputs"], (a, b, sfa, sfb, prob)) + self.assertIs(srelu_captured["launcher"], self.srelu._launch) + self.assertEqual( + set(srelu_captured["static_args"]), + {"alpha", "sf_vec_size", "mma_tiler_mn", "cluster_shape_mn", "vector_f32", "cluster_overlap_margin"}, + ) + self.assertEqual( + [(spec.name, spec.shape, spec.dtype, spec.fill_value) for spec in srelu_captured["outputs"]], + [ + ("c_tensor", (2, 384, 256), self.bfloat16, None), + ("d_tensor", (2, 384, 256), self.bfloat16, None), + ], + ) + self.assertEqual( + [spec.layout for spec in srelu_captured["input_specs"]], + [ + (2, 1, 0), + (2, 1, 0), + (2, 1, 4, 0, 3, 5), + (2, 1, 4, 0, 3, 5), + (0, 1, 2), + ], + ) + self.assertIsNone(srelu_result["amax_tensor"]) + self.assertIsNone(srelu_result["sfd_tensor"]) + + self.assertEqual(dsrelu_captured["inputs"], (a, b, c, sfa, sfb, prob)) + self.assertIs(dsrelu_captured["launcher"], self.dsrelu._launch) + self.assertEqual( + set(dsrelu_captured["static_args"]), + {"alpha", "sf_vec_size", "mma_tiler_mn", "cluster_shape_mn", "vector_f32", "cluster_overlap_margin"}, + ) + d_spec, dprob_spec = dsrelu_captured["outputs"] + self.assertEqual( + (d_spec.name, d_spec.shape, d_spec.dtype), + ("d_tensor", (2, 384, 256), self.bfloat16), + ) + self.assertEqual( + ( + dprob_spec.name, + dprob_spec.shape, + dprob_spec.dtype, + dprob_spec.fill_value, + ), + ("dprob_tensor", (384, 1, 2), self.float32, 0.0), + ) + self.assertEqual(dprob_spec.tensor_spec.layout, (0, 1, 2)) + self.assertEqual(dprob_spec.tensor_spec.mode, (0, 1, 2)) + self.assertIsNone(dsrelu_result["amax_tensor"]) + self.assertIsNone(dsrelu_result["sfd_tensor"]) + + def test_relu_wrappers_reject_unpredicated_partial_mma_tile(self): + with self._optional_modules(): + for m, mma_tiler_mn in ((129, (128, 128)), (130, (256, 256))): + a = _Array((1, m, 512), self.float8_e4m3fn) + b = _Array((1, 256, 512), self.float8_e4m3fn) + c = _Array((1, m, 256), self.bfloat16) + scale_m = (m + 127) // 128 + sfa = _Array((32, 4, scale_m, 4, 4, 1), self.float8_e8m0fnu) + sfb = _Array((32, 4, 2, 4, 4, 1), self.float8_e8m0fnu) + prob = _Array((m, 1, 1), self.float32) + + for module, call_args in ( + (self.srelu, (a, b, sfa, sfb, prob)), + (self.dsrelu, (a, b, c, sfa, sfb, prob)), + ): + wrapper_name = "gemm_srelu_wrapper_sm100" if module is self.srelu else "gemm_dsrelu_wrapper_sm100" + with ( + self.subTest(module=module.__name__, mma_tiler_mn=mma_tiler_mn), + self.assertRaisesRegex(ValueError, f"TILE_M={mma_tiler_mn[0]}"), + mock.patch.object(module, "call_cutedsl") as lower, + ): + getattr(module, wrapper_name)(*call_args, mma_tiler_mn=mma_tiler_mn) + lower.assert_not_called() + + def test_dense_class_uses_sample_metadata_without_retaining_arrays(self): + sample_a = _Array((1, 128, 128), self.float8_e4m3fn) + sample_b = _Array((1, 128, 128), self.float8_e4m3fn) + sample_sfa = _Array((32, 4, 1, 4, 1, 1), self.float8_e8m0fnu) + sample_sfb = _Array((32, 4, 1, 4, 1, 1), self.float8_e8m0fnu) + samples = (sample_a, sample_b, sample_sfa, sample_sfb) + dtype_carrier = _Array((1,), self.float32) + + with self._optional_modules(): + operation = self.amax.GemmAmaxSm100(*samples, c_dtype=dtype_carrier) + self.assertTrue(operation.check_support()) + + self.assertIs(operation.get_jax_callable(), operation) + self.assertEqual(operation.a_desc.shape, (128, 128, 1)) + self.assertEqual(operation.a_desc.array_shape, sample_a.shape) + self.assertEqual(operation.sfa_desc.shape, sample_sfa.shape) + self.assertFalse(any(name.endswith("_spec") for name in vars(operation))) + self.assertTrue(all(stored is not sample for stored in vars(operation).values() for sample in samples)) + self.assertTrue(all(stored is not dtype_carrier for stored in vars(operation).values())) + + actual_a = _Array(sample_a.shape, sample_a.dtype) + actual_b = _Array(sample_b.shape, sample_b.dtype) + actual_sfa = _Array(sample_sfa.shape, sample_sfa.dtype) + actual_sfb = _Array(sample_sfb.shape, sample_sfb.dtype) + captured = {} + with ( + self._optional_modules(), + mock.patch.object( + self.amax, + "call_cutedsl", + side_effect=self._fake_call(captured), + ), + ): + operation(actual_a, actual_b, actual_sfa, actual_sfb) + self.assertEqual(captured["inputs"], (actual_a, actual_b, actual_sfa, actual_sfb)) + self.assertEqual( + captured["input_specs"], + ( + operation.a_desc.tensor_spec, + operation.b_desc.tensor_spec, + operation.sfa_desc.tensor_spec, + operation.sfb_desc.tensor_spec, + ), + ) + self.assertIs(captured["outputs"][0].tensor_spec, operation.c_desc.tensor_spec) + + with ( + self._optional_modules(), + self.assertRaisesRegex(ValueError, "A tensor shape mismatch"), + mock.patch.object(self.amax, "call_cutedsl") as lower, + ): + operation( + _Array((1, 64, 128), self.float8_e4m3fn), + actual_b, + actual_sfa, + actual_sfb, + ) + lower.assert_not_called() + + def test_launchers_preserve_native_argument_order(self): + swiglu_calls = [] + amax_calls = [] + + class FakeHardwareInfo: + def get_max_active_clusters(self, cluster_size): + self.cluster_size = cluster_size + return 11 + + class FakeSwigluKernel: + TWO_CTA_MMA_TILER_M = 256 + + def __init__(self, **options): + self.options = options + + def __call__(self, *args): + swiglu_calls.append(args) + + class FakeAmaxKernel: + def __init__(self, **options): + self.options = options + + def __call__(self, *args): + amax_calls.append(args) + + swiglu_module_name = f"{_TEST_PACKAGE}.gemm_swiglu.dense_gemm_persistent_swiglu" + swiglu_module = types.ModuleType(swiglu_module_name) + swiglu_module.PersistentDenseGemmKernel = FakeSwigluKernel + amax_module_name = f"{_TEST_PACKAGE}.gemm_amax.dense_blockscaled_gemm_persistent_amax" + amax_module = types.ModuleType(amax_module_name) + amax_module.Sm100BlockScaledPersistentDenseGemmKernel = FakeAmaxKernel + + with ( + self._optional_modules(), + mock.patch.dict( + sys.modules, + { + swiglu_module_name: swiglu_module, + amax_module_name: amax_module, + }, + ), + mock.patch.object( + self.fake_cutlass, + "utils", + types.SimpleNamespace(HardwareInfo=FakeHardwareInfo), + create=True, + ), + mock.patch.object( + self.fake_cutlass, + "Float32", + lambda value: ("Float32", value), + create=True, + ), + ): + self.swiglu._launch( + "stream", + "a", + "b", + "ab12", + "c", + alpha=0.5, + acc_dtype=self.float32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + cluster_overlap_margin=1, + ) + self.amax._launch( + "stream", + "a", + "b", + "sfa", + "sfb", + "c", + "amax", + sf_vec_size=32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(2, 2), + cluster_overlap_margin=2, + ) + + self.assertEqual( + swiglu_calls, + [("a", "b", "ab12", "c", ("Float32", 0.5), 10, "stream")], + ) + self.assertEqual( + amax_calls, + [("a", "b", "sfa", "sfb", "c", "amax", 9, "stream")], + ) + + def test_relu_launchers_preserve_optional_native_slots(self): + srelu_calls = [] + dsrelu_calls = [] + + class FakeHardwareInfo: + def get_max_active_clusters(self, cluster_size): + return 8 + cluster_size + + class FakeSreluKernel: + def __init__(self, **options): + self.options = options + + def __call__(self, *args, **kwargs): + srelu_calls.append((args, kwargs)) + + class FakeDsreluKernel: + def __init__(self, **options): + self.options = options + + def __call__(self, *args, **kwargs): + dsrelu_calls.append((args, kwargs)) + + srelu_module_name = f"{_TEST_PACKAGE}.gemm_srelu.dense_blockscaled_gemm_persistent_srelu_quant" + srelu_module = types.ModuleType(srelu_module_name) + srelu_module.Sm100BlockScaledPersistentDenseGemmKernel = FakeSreluKernel + dsrelu_module_name = f"{_TEST_PACKAGE}.gemm_dsrelu.dense_blockscaled_gemm_persistent_dsrelu_quant" + dsrelu_module = types.ModuleType(dsrelu_module_name) + dsrelu_module.Sm100BlockScaledPersistentDenseGemmKernel = FakeDsreluKernel + + with ( + self._optional_modules(), + mock.patch.dict( + sys.modules, + { + srelu_module_name: srelu_module, + dsrelu_module_name: dsrelu_module, + }, + ), + mock.patch.object( + self.fake_cutlass, + "utils", + types.SimpleNamespace(HardwareInfo=FakeHardwareInfo), + create=True, + ), + mock.patch.object( + self.fake_cutlass, + "Float32", + lambda value: ("Float32", value), + create=True, + ), + ): + self.srelu._launch( + "stream", + "a", + "b", + "sfa", + "sfb", + "prob", + "c", + "d", + alpha=0.25, + sf_vec_size=32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + vector_f32=False, + cluster_overlap_margin=1, + ) + self.dsrelu._launch( + "stream", + "a", + "b", + "c", + "sfa", + "sfb", + "prob", + "d", + "dprob", + alpha=0.5, + sf_vec_size=32, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + vector_f32=True, + cluster_overlap_margin=2, + ) + + srelu_args, srelu_kwargs = srelu_calls[0] + self.assertEqual( + srelu_args, + ( + "a", + "b", + "sfa", + "sfb", + "c", + "d", + "prob", + None, + None, + None, + ("Float32", 0.25), + 8, + "stream", + ), + ) + self.assertEqual(set(srelu_kwargs), {"epilogue_op"}) + dsrelu_args, dsrelu_kwargs = dsrelu_calls[0] + self.assertEqual( + dsrelu_args, + ( + "a", + "b", + "sfa", + "sfb", + "c", + "d", + "prob", + "dprob", + None, + None, + None, + ("Float32", 0.5), + 8, + "stream", + ), + ) + self.assertEqual(set(dsrelu_kwargs), {"epilogue_op"}) + + +if __name__ == "__main__": + unittest.main() 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..696b143a0 --- /dev/null +++ b/test/python/fe_api/test_jax_grouped_gemm_contract.py @@ -0,0 +1,968 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for the JAX contiguous grouped GEMM wrappers.""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +import ast +import importlib +import importlib.util +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 + + +def _identity_jit(fn=None, **_kwargs): + return (lambda decorated_fn: decorated_fn) if fn is None else fn + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CUDNN_ROOT = _REPO_ROOT / "python" / "cudnn" +_TEST_PACKAGE = "cudnn_frontend_jax_grouped_gemm_contract_test" + + +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 + + +class _TensorSpec: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + +@dataclass(frozen=True) +class _BufferSpec: + name: str + shape: tuple + dtype: object + tensor_spec: object = None + fill_value: object = None + + +class _KernelSurface: + MMA_TILER_M = (64, 128, 256) + MMA_TILER_N = (128, 256) + TWO_CTA_MMA_TILER_M = 256 + MAX_CLUSTER_CTAS = 16 + MAX_CLUSTER_DIMENSION = 4 + CLUSTER_TILER_M = (128, 256) + FP8_SF_VEC_SIZE = 32 + SF_VEC_SIZES = (16, 32) + HADAMARD_SIZE = 16 + FIX_PAD_SIZE = 256 + MAX_EXPERTS = 1024 + DYNAMIC_SCHED_WORKSPACE_BYTES = 4 + calls = [] + + @classmethod + def require_mma_tiler(cls, value): + value = tuple(value) + if value[0] not in cls.MMA_TILER_M or value[1] not in cls.MMA_TILER_N: + raise ValueError(value) + return value + + @classmethod + def require_cluster_shape(cls, value, *, mma_tiler_mn): + del mma_tiler_mn + return tuple(value) + + @classmethod + def get_dense_workspace_bytes(cls, use_dynamic_sched): + return cls.DYNAMIC_SCHED_WORKSPACE_BYTES if use_dynamic_sched else 0 + + @classmethod + def can_implement(cls, *args): + del args + return True + + def __init__(self, **kwargs): + self.kwargs = kwargs + + def __call__(self, *args, **kwargs): + type(self).calls.append((self.kwargs, args, kwargs)) + + +class JaxGroupedGemmContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.bfloat16 = _DType("bfloat16", 2) + cls.float16 = _DType("float16", 2) + cls.float32 = _DType("float32", 4) + cls.float8_e4m3fn = _DType("float8_e4m3fn", 1) + cls.float8_e5m2 = _DType("float8_e5m2", 1) + cls.float8_e8m0fnu = _DType("float8_e8m0fnu", 1) + cls.float4_e2m1fn = _DType("float4_e2m1fn", 1) + cls.int32 = _DType("int32", 4) + cls.uint8 = _DType("uint8", 1) + + cls.fake_jnp = types.ModuleType("jax.numpy") + for name in ( + "bfloat16", + "float16", + "float32", + "float8_e4m3fn", + "float8_e5m2", + "float8_e8m0fnu", + "float4_e2m1fn", + "int32", + "uint8", + ): + setattr(cls.fake_jnp, name, getattr(cls, name)) + cls.fake_jnp.dtype = lambda value: value + cls.fake_jnp.ones = lambda shape, dtype: _Array(shape, dtype) + cls.fake_jnp.asarray = lambda value, dtype: _Array((len(value), len(value[0])), dtype) + + cls.fake_jax = types.ModuleType("jax") + cls.fake_jax.__path__ = [] + cls.fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + cls.fake_jax.numpy = cls.fake_jnp + cls.fake_jax.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, + ) + + cls.fake_cutlass = types.ModuleType("cutlass") + cls.fake_cutlass.__path__ = [] + cls.fake_cutlass.Constexpr = object + cls.fake_cutlass.utils = types.SimpleNamespace(HardwareInfo=lambda: types.SimpleNamespace(get_max_active_clusters=lambda _: 8)) + cls.fake_cutlass.Int32 = int + cls.fake_cutlass.Int64 = int + cls.fake_cutlass.Float32 = float + cls.fake_cutlass_cute = types.ModuleType("cutlass.cute") + cls.fake_cutlass_cute.jit = _identity_jit + cls.fake_cutlass_cute.where = lambda condition, yes, no: (yes if condition else no) + cls.fake_cutlass_cute.full_like = lambda _, fill: fill + cls.fake_cutlass.cute = cls.fake_cutlass_cute + cls.fake_cutlass_jax = types.ModuleType("cutlass.jax") + cls.fake_cutlass_jax.TensorSpec = _TensorSpec + cls.fake_cutlass_jax.jax_to_cutlass_dtype = lambda dtype: (f"cutlass.{dtype.name}") + cls.fake_cutlass.jax = cls.fake_cutlass_jax + cls.fake_cutlass_nvgpu = types.ModuleType("cutlass.cute.nvgpu") + cls.fake_cutlass_nvgpu.OperandMajorMode = types.SimpleNamespace(K="K") + + cls.calls = [] + + def call_cutedsl(fn, inputs, *, outputs, input_specs, **kwargs): + result_buffers = tuple(types.SimpleNamespace(name=spec.name) for spec in outputs) + workspace_buffers = tuple(types.SimpleNamespace(name=spec.name, iterator=f"{spec.name}_ptr") for spec in kwargs.get("workspaces", ())) + fn( + "stream", + *inputs, + *result_buffers, + *workspace_buffers, + **kwargs.get("static_args", {}), + ) + cls.calls.append((tuple(outputs), {**kwargs, "launcher": fn, "input_specs": tuple(input_specs)})) + return result_buffers + + cls.fake_gemm = types.ModuleType(f"{_TEST_PACKAGE}._jax.gemm") + cls.fake_gemm.require_16_byte_extent = lambda *_: None + + def require_layout(name, value, supported): + value = value.upper() + if value not in supported: + raise ValueError(f"{name} must be one of {supported}, got {value!r}") + return value + + def matrix_spec(name, value, supported, kernel_modes): + value = require_layout(name, value, supported) + return _TensorSpec( + layout=(2, 1, 0), + mode=tuple(value.index(dim) for dim in kernel_modes), + ) + + cls.fake_gemm.require_layout = require_layout + cls.fake_gemm.block_scale_tensor_spec = lambda: _TensorSpec(layout=(2, 1, 4, 0, 3, 5), mode=(0, 1, 2, 3, 4, 5)) + cls.fake_gemm.gemm_a_tensor_spec = lambda value: matrix_spec("a_layout", value, ("LMK", "LKM"), "MKL") + cls.fake_gemm.gemm_b_tensor_spec = lambda value: matrix_spec("b_layout", value, ("LNK", "LKN"), "NKL") + cls.fake_gemm.gemm_c_tensor_spec = lambda value, *, name="c_layout": matrix_spec(name, value, ("LMN", "LNM"), "MNL") + cls.fake_gemm.probability_tensor_spec = lambda: _TensorSpec(layout=(0, 1, 2), mode=(0, 1, 2)) + + cls.fake_kernel_modules = {} + for suffix, class_name in ( + ( + "grouped_gemm.grouped_gemm_swiglu.grouped_gemm_swiglu_quant", + "BlockScaledContiguousGroupedGemmKernel", + ), + ( + "grouped_gemm.grouped_gemm_dswiglu.grouped_gemm_dswiglu_quant", + "BlockScaledContiguousGroupedGemmKernel", + ), + ( + "grouped_gemm.grouped_gemm_quant.grouped_gemm_quant", + "BlockScaledMoEGroupedGemmQuantKernel", + ), + ( + "grouped_gemm.grouped_gemm_srelu.moe_blockscaled_grouped_gemm_srelu_quant", + "BlockScaledMoEGroupedGemmQuantKernel", + ), + ( + "grouped_gemm.grouped_gemm_glu.moe_blockscaled_grouped_gemm_glu_bias", + "BlockScaledMoEGroupedGemmGluBiasKernel", + ), + ( + "grouped_gemm.grouped_gemm_dglu.moe_blockscaled_grouped_gemm_dglu_dbias", + "BlockScaledMoEGroupedGemmDgluDbiasKernel", + ), + ( + "grouped_gemm.grouped_gemm_glu_hadamard.moe_blockscaled_grouped_gemm_glu_hadamard", + "BlockScaledMoEGroupedGemmGluHadamardKernel", + ), + ( + "grouped_gemm.grouped_gemm_dsrelu.moe_blockscaled_grouped_gemm_dsrelu_quant", + "BlockScaledMoEGroupedGemmQuantBwdKernel", + ), + ): + module = types.ModuleType(f"{_TEST_PACKAGE}.{suffix}") + setattr(module, class_name, _KernelSurface) + if "grouped_gemm_srelu" in suffix: + module.EpilogueType = types.SimpleNamespace(SRELU=types.SimpleNamespace(value=1)) + if "grouped_gemm_dsrelu" in suffix: + module.EpilogueType = types.SimpleNamespace(DSRELU=types.SimpleNamespace(value=1)) + cls.fake_kernel_modules[module.__name__] = module + + cls.fake_moe_utils = types.ModuleType(f"{_TEST_PACKAGE}.grouped_gemm.moe_utils") + cls.fake_moe_utils.MoEWeightMode = types.SimpleNamespace(DENSE="DENSE") + + with cls._modules(): + cls._make_package(_TEST_PACKAGE, _CUDNN_ROOT) + cls._make_package(f"{_TEST_PACKAGE}._jax", _CUDNN_ROOT / "_jax") + cls._make_package(f"{_TEST_PACKAGE}.grouped_gemm", _CUDNN_ROOT / "grouped_gemm") + for operation in ( + "grouped_gemm_swiglu", + "grouped_gemm_dswiglu", + "grouped_gemm_quant", + "grouped_gemm_srelu", + "grouped_gemm_glu", + "grouped_gemm_dglu", + "grouped_gemm_glu_hadamard", + "grouped_gemm_dsrelu", + ): + cls._make_package( + f"{_TEST_PACKAGE}.grouped_gemm.{operation}", + _CUDNN_ROOT / "grouped_gemm" / operation, + ) + cls._load_source(f"{_TEST_PACKAGE}.gemm_validation", _CUDNN_ROOT / "gemm_validation.py") + jax_api_base = cls._load_source( + f"{_TEST_PACKAGE}._jax.api_base", + _CUDNN_ROOT / "_jax" / "api_base.py", + ) + jax_api_base.BufferSpec = _BufferSpec + jax_api_base.call_cutedsl = call_cutedsl + cls.fake_gemm.as_gemm_tensor_desc = lambda name, value, tensor_spec: ( + value if isinstance(value, jax_api_base.JaxTensorDesc) else jax_api_base.JaxTensorDesc.from_value(value, tensor_spec=tensor_spec, name=name) + ) + cls._load_source( + f"{_TEST_PACKAGE}._jax.grouped_gemm", + _CUDNN_ROOT / "_jax" / "grouped_gemm.py", + ) + cls.swiglu = importlib.import_module(f"{_TEST_PACKAGE}.grouped_gemm.grouped_gemm_swiglu.jax") + cls.dswiglu = importlib.import_module(f"{_TEST_PACKAGE}.grouped_gemm.grouped_gemm_dswiglu.jax") + cls.quant = importlib.import_module(f"{_TEST_PACKAGE}.grouped_gemm.grouped_gemm_quant.jax") + cls.srelu = importlib.import_module(f"{_TEST_PACKAGE}.grouped_gemm.grouped_gemm_srelu.jax") + cls.glu = importlib.import_module(f"{_TEST_PACKAGE}.grouped_gemm.grouped_gemm_glu.jax") + cls.dglu = importlib.import_module(f"{_TEST_PACKAGE}.grouped_gemm.grouped_gemm_dglu.jax") + cls.glu_hadamard = importlib.import_module(f"{_TEST_PACKAGE}.grouped_gemm.grouped_gemm_glu_hadamard.jax") + cls.dsrelu = importlib.import_module(f"{_TEST_PACKAGE}.grouped_gemm.grouped_gemm_dsrelu.jax") + + @classmethod + def tearDownClass(cls): + for module_name in tuple(sys.modules): + if module_name == _TEST_PACKAGE or module_name.startswith(f"{_TEST_PACKAGE}."): + sys.modules.pop(module_name, None) + + @staticmethod + def _make_package(name, path): + module = types.ModuleType(name) + module.__path__ = [str(path)] + module.__package__ = name + sys.modules[name] = module + + @staticmethod + def _load_source(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + @classmethod + @contextmanager + def _modules(cls): + modules = { + "jax": cls.fake_jax, + "jax.numpy": cls.fake_jnp, + "cutlass": cls.fake_cutlass, + "cutlass.cute": cls.fake_cutlass_cute, + "cutlass.cute.nvgpu": cls.fake_cutlass_nvgpu, + "cutlass.jax": cls.fake_cutlass_jax, + "torch": None, + f"{_TEST_PACKAGE}._jax.gemm": cls.fake_gemm, + f"{_TEST_PACKAGE}.grouped_gemm.moe_utils": cls.fake_moe_utils, + **cls.fake_kernel_modules, + } + with mock.patch.dict(sys.modules, modules): + yield + + def setUp(self): + self.calls.clear() + _KernelSurface.calls.clear() + + def _inputs(self): + m, n, k, experts = 256, 256, 64, 4 + return { + "a": _Array((1, m, k), self.float8_e4m3fn), + "b": _Array((experts, n, k), self.float8_e4m3fn), + "sfa": _Array((32, 4, 2, 4, 1, 1), self.float8_e8m0fnu), + "sfb": _Array((32, 4, 2, 4, 1, experts), self.float8_e8m0fnu), + "offsets": _Array((experts,), self.int32), + "alpha": _Array((experts,), self.float32), + "norm": _Array((1,), self.float32), + "prob": _Array((m, 1, 1), self.float32), + } + + def test_forward_functional_outputs(self): + values = self._inputs() + with self._modules(): + result = self.swiglu.grouped_gemm_swiglu_wrapper_sm100( + values["a"], + values["b"], + values["sfa"], + values["sfb"], + values["offsets"], + values["alpha"], + values["norm"], + values["prob"], + ) + + self.assertEqual( + [item.name if item is not None else None for item in result], + [ + "c_tensor", + "d_tensor", + "d_col_tensor", + "amax_tensor", + "sfd_row_tensor", + "sfd_col_tensor", + ], + ) + specs, options = self.calls[-1] + self.assertEqual(next(spec.shape for spec in specs if spec.name == "c_tensor"), (1, 256, 256)) + self.assertEqual(next(spec.shape for spec in specs if spec.name == "d_tensor"), (1, 256, 128)) + self.assertEqual(options["input_specs"][0].mode, (1, 2, 0)) + self.assertEqual(options["input_specs"][1].mode, (1, 2, 0)) + self.assertEqual( + next(spec.fill_value for spec in specs if spec.name == "amax_tensor"), + -float("inf"), + ) + self.assertEqual( + next(spec.shape for spec in specs if spec.name == "sfd_row_tensor"), + (32, 4, 2, 4, 1, 1), + ) + self.assertIs(options["launcher"], self.swiglu._launch) + self.assertEqual(options["static_args"]["expert_cnt"], 4) + self.assertTrue(options["static_args"]["has_prob"]) + self.assertEqual(_KernelSurface.calls[-1][0]["expert_cnt"], 4) + + def test_class_configuration_and_sample_descriptors_are_immutable(self): + values = self._inputs() + with self._modules(): + operation = self.swiglu.GroupedGemmSwigluSm100( + values["a"], + values["b"], + values["sfa"], + values["sfb"], + values["offsets"], + values["alpha"], + values["norm"], + values["prob"], + ) + with self.assertRaises(TypeError): + operation._config["mma_tiler_mn"] = (128, 128) + with self.assertRaises(TypeError): + operation._sample_descs["a_tensor"] = None + + self.assertTrue(operation.check_support()) + operation( + values["a"], + values["b"], + values["sfa"], + values["sfb"], + values["offsets"], + values["alpha"], + values["norm"], + values["prob"], + ) + self.assertTrue(operation.check_support()) + with self.assertRaisesRegex(AttributeError, "immutable after its first call"): + operation._config = {} + + def test_environment_config_is_captured_at_construction(self): + values = self._inputs() + with self._modules(), mock.patch.object(self.swiglu.os, "getenv", return_value="3") as getenv: + operation = self.swiglu.GroupedGemmSwigluSm100( + values["a"], + values["b"], + values["sfa"], + values["sfb"], + values["offsets"], + values["alpha"], + values["norm"], + values["prob"], + ) + getenv.return_value = "9" + operation( + values["a"], + values["b"], + values["sfa"], + values["sfb"], + values["offsets"], + values["alpha"], + values["norm"], + values["prob"], + ) + + self.assertEqual(getenv.call_count, 1) + self.assertEqual(self.calls[-1][1]["static_args"]["cluster_overlap_margin"], 3) + + def test_backward_zero_initializes_dprob(self): + values = self._inputs() + c_tensor = _Array((1, 256, 512), self.bfloat16) + with self._modules(): + result = self.dswiglu.grouped_gemm_dswiglu_wrapper_sm100( + values["a"], + values["b"], + c_tensor, + values["sfa"], + values["sfb"], + values["offsets"], + values["alpha"], + None, + values["prob"], + values["norm"], + ) + + self.assertEqual( + [item.name if item is not None else None for item in result], + [ + "d_row_tensor", + "d_col_tensor", + "dprob_tensor", + None, + "sfd_row_tensor", + "sfd_col_tensor", + ], + ) + specs, _ = self.calls[-1] + dprob_spec = next(spec for spec in specs if spec.name == "dprob_tensor") + self.assertEqual(dprob_spec.fill_value, 0.0) + self.assertEqual(dprob_spec.tensor_spec.layout, (0, 1, 2)) + self.assertEqual( + next(spec.dtype for spec in specs if spec.name == "d_row_tensor"), + self.float8_e4m3fn, + ) + + def test_quantized_output_uses_hidden_dynamic_workspace(self): + values = self._inputs() + bias = _Array((256, 4), self.float32) + row_scale = _Array((256,), self.float32) + with self._modules(): + result = self.quant.grouped_gemm_quant_wrapper_sm100( + values["a"], + values["sfa"], + values["offsets"], + values["alpha"], + values["b"], + values["sfb"], + bias_tensor=bias, + norm_const_tensor=values["norm"], + prob_tensor=values["prob"], + row_scale_tensor=row_scale, + d_dtype=self.float8_e4m3fn, + use_dynamic_sched=True, + ) + + self.assertEqual( + [item.name if item is not None else None for item in result], + [ + "d_tensor", + "d_col_tensor", + None, + "sfd_row_tensor", + "sfd_col_tensor", + ], + ) + specs, options = self.calls[-1] + self.assertNotIn("amax_tensor", {spec.name for spec in specs}) + workspace = options["workspaces"][0] + self.assertEqual(workspace.shape, (4,)) + self.assertEqual(workspace.dtype, self.uint8) + self.assertEqual(workspace.tensor_spec.ptr_assumed_align, 128) + self.assertTrue(_KernelSurface.calls[-1][0]["use_dynamic_sched"]) + + def test_quant_accepts_public_lkn_weight_layout(self): + values = self._inputs() + values["b"] = _Array((4, 64, 128), self.float8_e4m3fn) + values["sfb"] = _Array((32, 4, 1, 4, 1, 4), self.float8_e8m0fnu) + with self._modules(): + self.quant.grouped_gemm_quant_wrapper_sm100( + values["a"], + values["sfa"], + values["offsets"], + values["alpha"], + values["b"], + values["sfb"], + prob_tensor=values["prob"], + b_layout="lkn", + mma_tiler_mn=(256, 128), + ) + + specs, options = self.calls[-1] + self.assertEqual(next(spec.shape for spec in specs if spec.name == "d_tensor"), (1, 256, 128)) + self.assertEqual(options["input_specs"][1].layout, (2, 1, 0)) + self.assertEqual(options["input_specs"][1].mode, (2, 1, 0)) + + def test_grouped_layouts_are_limited_by_kernel_capabilities(self): + values = self._inputs() + with self._modules(): + with self.assertRaisesRegex(ValueError, "output_layout must be one of"): + self.quant.grouped_gemm_quant_wrapper_sm100( + values["a"], + values["sfa"], + values["offsets"], + values["alpha"], + values["b"], + values["sfb"], + prob_tensor=values["prob"], + output_layout="LNM", + ) + + with self.assertRaisesRegex(ValueError, "b_layout must be one of"): + self.glu.grouped_gemm_glu_wrapper_sm100( + values["a"], + values["sfa"], + values["offsets"], + values["alpha"], + values["b"], + values["sfb"], + prob_tensor=values["prob"], + b_layout="LKN", + ) + + self.assertFalse(self.calls) + + def test_srelu_nonquantized_output_initializes_amax(self): + values = self._inputs() + with self._modules(): + result = self.srelu.grouped_gemm_srelu_wrapper_sm100( + values["a"], + values["b"], + values["sfa"], + values["sfb"], + values["offsets"], + values["alpha"], + prob_tensor=values["prob"], + ) + + self.assertEqual( + [item.name if item is not None else None for item in result], + ["c_tensor", "d_tensor", None, "amax_tensor", None, None], + ) + specs, options = self.calls[-1] + self.assertEqual( + next(spec.fill_value for spec in specs if spec.name == "amax_tensor"), + -float("inf"), + ) + self.assertEqual(options["workspaces"][0].shape, (1,)) + self.assertEqual(_KernelSurface.calls[-1][0]["epilogue_type"], 1) + + def test_glu_nonquantized_output_uses_hidden_d_col_and_amax(self): + values = self._inputs() + with self._modules(): + result = self.glu.grouped_gemm_glu_wrapper_sm100( + values["a"], + values["sfa"], + values["offsets"], + values["alpha"], + values["b"], + values["sfb"], + prob_tensor=values["prob"], + ) + + self.assertEqual( + [item.name if item is not None else None for item in result], + ["c_tensor", "d_tensor", None, "amax_tensor", None, None], + ) + specs, options = self.calls[-1] + self.assertEqual( + next(spec.fill_value for spec in specs if spec.name == "amax_tensor"), + -float("inf"), + ) + self.assertEqual( + [spec.name for spec in options["workspaces"]], + ["d_col_scratch", "workspace"], + ) + kernel_config, kernel_args, _ = _KernelSurface.calls[-1] + self.assertEqual(kernel_config["weight_mode"], "DENSE") + self.assertEqual(kernel_args[10].name, "d_col_scratch") + + def test_glu_quantized_output_supports_bias_and_dynamic_workspace(self): + values = self._inputs() + bias = _Array((256, 4), self.float32) + with self._modules(): + result = self.glu.grouped_gemm_glu_wrapper_sm100( + values["a"], + values["sfa"], + values["offsets"], + values["alpha"], + values["b"], + values["sfb"], + bias_tensor=bias, + norm_const_tensor=values["norm"], + prob_tensor=values["prob"], + d_dtype=self.float8_e4m3fn, + act_func="geglu", + use_dynamic_sched=True, + ) + + self.assertEqual( + [item.name if item is not None else None for item in result], + [ + "c_tensor", + "d_tensor", + "d_col_tensor", + None, + "sfd_row_tensor", + "sfd_col_tensor", + ], + ) + _, options = self.calls[-1] + self.assertEqual([spec.name for spec in options["workspaces"]], ["workspace"]) + self.assertEqual(options["workspaces"][0].shape, (4,)) + kernel_config, _, launch_kwargs = _KernelSurface.calls[-1] + self.assertTrue(kernel_config["enable_bias"]) + self.assertTrue(kernel_config["generate_sfd"]) + self.assertEqual(launch_kwargs["linear_offset"], 1.0) + + def test_dglu_zero_initializes_functional_reductions(self): + values = self._inputs() + c_tensor = _Array((1, 256, 512), self.bfloat16) + beta = _Array((4,), self.float32) + with self._modules(): + result = self.dglu.grouped_gemm_dglu_wrapper_sm100( + values["a"], + c_tensor, + values["sfa"], + values["offsets"], + values["alpha"], + beta, + values["prob"], + values["b"], + values["sfb"], + generate_dbias=True, + norm_const_tensor=values["norm"], + use_dynamic_sched=True, + ) + + self.assertEqual( + [item.name if item is not None else None for item in result], + [ + "d_row_tensor", + "d_col_tensor", + "dprob_tensor", + "dbias_tensor", + None, + "sfd_row_tensor", + "sfd_col_tensor", + ], + ) + specs, options = self.calls[-1] + fills = {spec.name: spec.fill_value for spec in specs} + self.assertEqual(fills["dprob_tensor"], 0.0) + self.assertEqual( + next(spec.tensor_spec.layout for spec in specs if spec.name == "dprob_tensor"), + (0, 1, 2), + ) + self.assertEqual(fills["dbias_tensor"], 0.0) + self.assertEqual(options["workspaces"][0].shape, (4,)) + kernel_config, _, _ = _KernelSurface.calls[-1] + self.assertEqual(kernel_config["weight_mode"], "DENSE") + self.assertEqual(kernel_config["act_func"], "dswiglu") + + def test_glu_hadamard_native_fp4_outputs_and_workspace_are_functional(self): + values = self._inputs() + values["a"] = _Array(values["a"].shape, self.float4_e2m1fn) + values["b"] = _Array(values["b"].shape, self.float4_e2m1fn) + bias = _Array((256, 4), self.float32) + with self._modules(): + result = self.glu_hadamard.grouped_gemm_glu_hadamard_wrapper_sm100( + values["a"], + values["b"], + values["sfa"], + values["sfb"], + values["offsets"], + values["alpha"], + values["prob"], + bias_tensor=bias, + act_func="geglu", + use_dynamic_sched=True, + use_tmem_post_rht_amax=True, + ) + + self.assertEqual( + [item.name for item in result], + ["c_tensor", "d_tensor", "amax_tensor", "post_rht_amax_tensor"], + ) + specs, options = self.calls[-1] + fills = {spec.name: spec.fill_value for spec in specs} + self.assertEqual( + fills, + { + "c_tensor": 0, + "d_tensor": 0, + "amax_tensor": 0.0, + "post_rht_amax_tensor": 0.0, + }, + ) + workspace = options["workspaces"][0] + self.assertEqual(workspace.shape, (4,)) + self.assertEqual(workspace.fill_value, 0) + self.assertEqual(workspace.tensor_spec.ptr_assumed_align, 128) + + kernel_config, kernel_args, launch_kwargs = _KernelSurface.calls[-1] + self.assertEqual(kernel_config["weight_mode"], "DENSE") + self.assertTrue(kernel_config["enable_bias"]) + self.assertTrue(kernel_config["use_tmem_post_rht_amax"]) + self.assertEqual(kernel_args[8], "workspace_ptr") + self.assertEqual(kernel_args[16].shape, (16, 16)) + self.assertIs(kernel_args[17], bias) + self.assertEqual(launch_kwargs["linear_offset"], 1.0) + + def test_glu_hadamard_rejects_raw_uint8_fp4(self): + values = self._inputs() + values["a"] = _Array(values["a"].shape, self.uint8) + values["b"] = _Array(values["b"].shape, self.uint8) + with self._modules(), self.assertRaisesRegex(ValueError, "a_tensor.dtype"): + self.glu_hadamard.grouped_gemm_glu_hadamard_wrapper_sm100( + values["a"], + values["b"], + values["sfa"], + values["sfb"], + values["offsets"], + values["alpha"], + values["prob"], + ) + + def test_glu_hadamard_srelu_supports_e4m3_scales_and_static_workspace(self): + values = self._inputs() + values["a"] = _Array(values["a"].shape, self.float4_e2m1fn) + values["b"] = _Array(values["b"].shape, self.float4_e2m1fn) + values["sfa"] = _Array(values["sfa"].shape, self.float8_e4m3fn) + values["sfb"] = _Array(values["sfb"].shape, self.float8_e4m3fn) + with self._modules(): + self.glu_hadamard.grouped_gemm_glu_hadamard_wrapper_sm100( + values["a"], + values["b"], + values["sfa"], + values["sfb"], + values["offsets"], + values["alpha"], + values["prob"], + act_func="srelu", + sf_vec_size=16, + ) + + specs, options = self.calls[-1] + self.assertEqual( + next(spec.shape for spec in specs if spec.name == "d_tensor"), + (1, 256, 256), + ) + self.assertEqual(options["workspaces"][0].shape, (1,)) + self.assertEqual(options["workspaces"][0].fill_value, 0) + + def test_glu_hadamard_lazy_surface_and_kernel_imports_are_torch_optional(self): + init_tree = ast.parse((_CUDNN_ROOT / "grouped_gemm/grouped_gemm_glu_hadamard/__init__.py").read_text()) + exports = next( + node for node in init_tree.body if isinstance(node, ast.Assign) and isinstance(node.targets[0], ast.Name) and node.targets[0].id == "_API_EXPORTS" + ) + self.assertIsInstance(exports.value, ast.Tuple) + + for relative_path in ( + "grouped_gemm/grouped_gemm_glu_hadamard/hadamard_utils.py", + "grouped_gemm/moe_kernel_helpers.py", + ): + tree = ast.parse((_CUDNN_ROOT / relative_path).read_text()) + required_torch_imports = [ + node + for node in tree.body + if (isinstance(node, ast.Import) and any(alias.name == "torch" for alias in node.names)) + or (isinstance(node, ast.ImportFrom) and node.module == "torch") + ] + self.assertEqual(required_torch_imports, []) + + def test_dsrelu_exposes_functional_outputs_and_hidden_workspace(self): + values = self._inputs() + c_tensor = _Array((1, 256, 256), self.bfloat16) + with self._modules(): + result = self.dsrelu.grouped_gemm_dsrelu_wrapper_sm100( + values["a"], + values["b"], + c_tensor, + values["sfa"], + values["sfb"], + values["offsets"], + values["alpha"], + values["prob"], + generate_dbias=True, + norm_const_tensor=values["norm"], + use_dynamic_sched=True, + use_dsrelu_reuse=True, + ) + + self.assertEqual( + [item.name if item is not None else None for item in result], + [ + "d_row_tensor", + "d_col_tensor", + "d_srelu_tensor", + "dprob_tensor", + "dbias_tensor", + None, + "sfd_row_tensor", + "sfd_col_tensor", + "sfd_col_d_srelu_tensor", + ], + ) + specs, options = self.calls[-1] + fills = {spec.name: spec.fill_value for spec in specs} + self.assertEqual(fills["dprob_tensor"], 0.0) + self.assertEqual(fills["dbias_tensor"], 0.0) + self.assertEqual(options["workspaces"][0].shape, (4,)) + kernel_config, _, _ = _KernelSurface.calls[-1] + self.assertEqual(kernel_config["weight_mode"], "DENSE") + self.assertEqual(kernel_config["epilogue_type"], 1) + self.assertTrue(kernel_config["generate_d_srelu"]) + self.assertTrue(kernel_config["use_dsrelu_reuse"]) + + def test_kernel_capabilities_are_public(self): + required_fields = { + "MMA_TILER_M", + "MMA_TILER_N", + "TWO_CTA_MMA_TILER_M", + "MAX_CLUSTER_CTAS", + "MAX_CLUSTER_DIMENSION", + "CLUSTER_TILER_M", + "FP8_SF_VEC_SIZE", + "FIX_PAD_SIZE", + "MAX_EXPERTS", + } + for relative_path in ( + "grouped_gemm/grouped_gemm_swiglu/grouped_gemm_swiglu_quant.py", + "grouped_gemm/grouped_gemm_dswiglu/grouped_gemm_dswiglu_quant.py", + ): + tree = ast.parse((_CUDNN_ROOT / relative_path).read_text()) + kernel = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "BlockScaledContiguousGroupedGemmKernel") + fields = { + node.targets[0].id for node in kernel.body if isinstance(node, ast.Assign) and len(node.targets) == 1 and isinstance(node.targets[0], ast.Name) + } + methods = {node.name for node in kernel.body if isinstance(node, ast.FunctionDef)} + self.assertTrue(required_fields <= fields) + self.assertTrue({"require_mma_tiler", "require_cluster_shape"} <= methods) + + workspace_fields = required_fields | {"DYNAMIC_SCHED_WORKSPACE_BYTES"} + for relative_path, class_name in ( + ( + "grouped_gemm/grouped_gemm_quant/grouped_gemm_quant.py", + "BlockScaledMoEGroupedGemmQuantKernel", + ), + ( + "grouped_gemm/grouped_gemm_srelu/moe_blockscaled_grouped_gemm_srelu_quant.py", + "BlockScaledMoEGroupedGemmQuantKernel", + ), + ( + "grouped_gemm/grouped_gemm_glu/moe_blockscaled_grouped_gemm_glu_bias.py", + "BlockScaledMoEGroupedGemmGluBiasKernel", + ), + ( + "grouped_gemm/grouped_gemm_dglu/moe_blockscaled_grouped_gemm_dglu_dbias.py", + "BlockScaledMoEGroupedGemmDgluDbiasKernel", + ), + ( + "grouped_gemm/grouped_gemm_dsrelu/moe_blockscaled_grouped_gemm_dsrelu_quant.py", + "BlockScaledMoEGroupedGemmQuantBwdKernel", + ), + ): + tree = ast.parse((_CUDNN_ROOT / relative_path).read_text()) + kernel = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == class_name) + fields = { + node.targets[0].id for node in kernel.body if isinstance(node, ast.Assign) and len(node.targets) == 1 and isinstance(node.targets[0], ast.Name) + } + methods = {node.name for node in kernel.body if isinstance(node, ast.FunctionDef)} + self.assertTrue(workspace_fields <= fields) + self.assertTrue( + { + "require_mma_tiler", + "require_cluster_shape", + "get_dense_workspace_bytes", + } + <= methods + ) + + tree = ast.parse((_CUDNN_ROOT / "grouped_gemm/grouped_gemm_glu_hadamard/moe_blockscaled_grouped_gemm_glu_hadamard.py").read_text()) + kernel = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "BlockScaledMoEGroupedGemmGluHadamardKernel") + fields = { + node.targets[0].id for node in kernel.body if isinstance(node, ast.Assign) and len(node.targets) == 1 and isinstance(node.targets[0], ast.Name) + } + methods = {node.name for node in kernel.body if isinstance(node, ast.FunctionDef)} + self.assertTrue( + { + "MMA_TILER_M", + "MMA_TILER_N", + "TWO_CTA_MMA_TILER_M", + "MAX_CLUSTER_CTAS", + "MAX_CLUSTER_DIMENSION", + "CLUSTER_TILER_M", + "SF_VEC_SIZES", + "MAX_EXPERTS", + "DYNAMIC_SCHED_WORKSPACE_BYTES", + "HADAMARD_SIZE", + "FIX_PAD_SIZE", + } + <= fields + ) + self.assertTrue( + { + "require_mma_tiler", + "require_cluster_shape", + "get_dense_workspace_bytes", + } + <= methods + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_grouped_gemm_wgrad_contract.py b/test/python/fe_api/test_jax_grouped_gemm_wgrad_contract.py new file mode 100644 index 000000000..5333a9890 --- /dev/null +++ b/test/python/fe_api/test_jax_grouped_gemm_wgrad_contract.py @@ -0,0 +1,470 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for JAX grouped GEMM weight gradients.""" + +from __future__ import annotations + +import ast +from enum import Enum +from importlib import import_module +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 + + +def _identity_jit(fn=None, **_kwargs): + return (lambda decorated_fn: decorated_fn) if fn is None else fn + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CUDNN_ROOT = _REPO_ROOT / "python" / "cudnn" +_TEST_PACKAGE = "cudnn_frontend_jax_grouped_wgrad_contract_test" + + +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 + + +class _TensorSpec: + def __init__( + self, + *, + layout=None, + mode=None, + static=None, + ptr_assumed_align=256, + divisibility=None, + ): + self.layout = layout + self.mode = mode + self.static = static + self.ptr_assumed_align = ptr_assumed_align + self.divisibility = divisibility + + +class _WeightMode(Enum): + DENSE = "dense" + DISCRETE = "discrete" + + +class _InputOrder(Enum): + Tensor2D = "tensor2d" + TensorRagged = "tensor_ragged" + + +class _HardwareInfo: + def get_max_active_clusters(self, cluster_size): + self.cluster_size = cluster_size + return 12 + + +class _Kernel: + MMA_TILER_M = (128, 256) + MMA_TILER_N = (128, 256) + TWO_CTA_MMA_TILER_M = 256 + FP8_SF_VEC_SIZE = 32 + instances = [] + + @classmethod + def require_mma_tiler(cls, mma_tiler_mn): + value = tuple(mma_tiler_mn) + if value[0] not in cls.MMA_TILER_M or value[1] not in cls.MMA_TILER_N: + raise ValueError("unsupported mma_tiler_mn") + return value + + @classmethod + def require_cluster_shape(cls, cluster_shape_mn, *, mma_tiler_mn): + value = tuple(cluster_shape_mn) + if mma_tiler_mn[0] == cls.TWO_CTA_MMA_TILER_M and value[0] % 2: + raise ValueError("cluster_shape_mn[0] must be divisible by 2") + return value + + @staticmethod + def require_input_order(input_order): + if isinstance(input_order, _InputOrder): + return input_order + return _InputOrder(input_order) + + @classmethod + def get_dense_workspace_bytes(cls, expert_cnt, input_order="tensor2d"): + order = cls.require_input_order(input_order) + slots = 4 if order is _InputOrder.TensorRagged else 2 + return slots * expert_cnt * 128 + + def __init__(self, **configuration): + self.configuration = configuration + self.calls = [] + self.instances.append(self) + + def __call__(self, *args): + self.calls.append(args) + + +class JaxGroupedGemmWgradContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.float8_e4m3fn = _DType("float8_e4m3fn", 1) + cls.float8_e5m2 = _DType("float8_e5m2", 1) + cls.float8_e8m0fnu = _DType("float8_e8m0fnu", 1) + cls.bfloat16 = _DType("bfloat16", 2) + cls.float16 = _DType("float16", 2) + cls.float32 = _DType("float32", 4) + cls.int32 = _DType("int32", 4) + cls.uint8 = _DType("uint8", 1) + + cls.fake_jnp = types.ModuleType("jax.numpy") + for name in ( + "float8_e4m3fn", + "float8_e5m2", + "float8_e8m0fnu", + "bfloat16", + "float16", + "float32", + "int32", + "uint8", + ): + setattr(cls.fake_jnp, name, getattr(cls, name)) + cls.fake_jnp.dtype = lambda value: value + + cls.fake_jax = types.ModuleType("jax") + cls.fake_jax.__path__ = [] + cls.fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + cls.fake_jax.numpy = cls.fake_jnp + cls.fake_jax.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, + ) + cls.fake_jax.ShapeDtypeStruct = lambda shape, dtype: (shape, dtype) + + cls.fake_cutlass = types.ModuleType("cutlass") + cls.fake_cutlass.__path__ = [] + cls.fake_cutlass.Constexpr = object + cls.fake_cutlass.utils = types.SimpleNamespace(HardwareInfo=_HardwareInfo) + cls.fake_cutlass_cute = types.ModuleType("cutlass.cute") + cls.fake_cutlass_cute.jit = _identity_jit + cls.fake_cutlass.cute = cls.fake_cutlass_cute + cls.fake_cutlass_jax = types.ModuleType("cutlass.jax") + cls.fake_cutlass_jax.TensorSpec = _TensorSpec + cls.fake_cutlass_jax.cutlass_call = None + cls.fake_cutlass_jax.jax_to_cutlass_dtype = lambda dtype: f"cutlass.{dtype.name}" + cls.fake_cutlass.jax = cls.fake_cutlass_jax + + cls.kernel_module_name = f"{_TEST_PACKAGE}.grouped_gemm.grouped_gemm_wgrad." "moe_blockscaled_grouped_gemm_wgrad" + cls.kernel_module = types.ModuleType(cls.kernel_module_name) + cls.kernel_module.BlockScaledMoEGroupedGemmWgradKernel = _Kernel + + cls.moe_utils_module_name = f"{_TEST_PACKAGE}.grouped_gemm.moe_utils" + cls.moe_utils_module = types.ModuleType(cls.moe_utils_module_name) + cls.moe_utils_module.MoEWeightMode = _WeightMode + cls.moe_utils_module.WGradInputOrder = _InputOrder + + package_paths = { + _TEST_PACKAGE: _CUDNN_ROOT, + f"{_TEST_PACKAGE}.grouped_gemm": _CUDNN_ROOT / "grouped_gemm", + f"{_TEST_PACKAGE}.grouped_gemm.grouped_gemm_wgrad": (_CUDNN_ROOT / "grouped_gemm" / "grouped_gemm_wgrad"), + } + for package_name, package_path in package_paths.items(): + package = types.ModuleType(package_name) + package.__path__ = [str(package_path)] + package.__package__ = package_name + sys.modules[package_name] = package + + with cls._optional_modules(): + cls.module = import_module(f"{_TEST_PACKAGE}.grouped_gemm.grouped_gemm_wgrad.jax") + + @classmethod + def tearDownClass(cls): + for module_name in tuple(sys.modules): + if module_name == _TEST_PACKAGE or module_name.startswith(f"{_TEST_PACKAGE}."): + sys.modules.pop(module_name, None) + + def setUp(self): + _Kernel.instances.clear() + + @classmethod + def _optional_modules(cls, *, include_kernel=False): + modules = { + "jax": cls.fake_jax, + "jax.numpy": cls.fake_jnp, + "cutlass": cls.fake_cutlass, + "cutlass.cute": cls.fake_cutlass_cute, + "cutlass.jax": cls.fake_cutlass_jax, + } + if include_kernel: + modules[cls.kernel_module_name] = cls.kernel_module + modules[cls.moe_utils_module_name] = cls.moe_utils_module + return mock.patch.dict(sys.modules, modules) + + @classmethod + def _inputs(cls): + return ( + _Array((384, 640), cls.float8_e4m3fn), + _Array((640, 640), cls.float8_e4m3fn), + _Array((384, 20), cls.float8_e8m0fnu), + _Array((640, 20), cls.float8_e8m0fnu), + _Array((2,), cls.int32), + ) + + @staticmethod + def _fake_call(captured): + def call(launcher, inputs, **options): + captured.update(launcher=launcher, inputs=tuple(inputs), **options) + return tuple(_Array(spec.shape, spec.dtype) for spec in options["outputs"]) + + return call + + def test_kernel_module_is_lazy(self): + self.assertNotIn(self.kernel_module_name, sys.modules) + + def test_sample_descriptors_use_the_lowering_tensor_specs(self): + with self._optional_modules(include_kernel=True): + operation = self.module.GroupedGemmWgradSm100(*self._inputs()) + + self.assertEqual(operation._sample_descs["a_tensor"].tensor_spec.layout, (1, 0)) + self.assertEqual(operation._sample_descs["b_tensor"].tensor_spec.layout, (0, 1)) + self.assertEqual(operation._sample_descs["sfa_tensor"].tensor_spec.layout, (1, 0)) + self.assertEqual(operation._sample_descs["sfb_tensor"].tensor_spec.layout, (1, 0)) + + def test_dense_output_and_workspace_are_xla_owned(self): + captured = {} + inputs = self._inputs() + with ( + self._optional_modules(include_kernel=True), + mock.patch.object( + self.module, + "call_cutedsl", + side_effect=self._fake_call(captured), + ), + ): + result = self.module.grouped_gemm_wgrad_wrapper_sm100(*inputs) + + self.assertEqual(result["wgrad_tensor"].shape, (2, 384, 640)) + self.assertIs(result["wgrad_tensor"].dtype, self.bfloat16) + self.assertEqual(captured["inputs"], inputs) + self.assertIs(captured["launcher"], self.module._launch) + self.assertEqual( + captured["static_args"], + { + "acc_dtype": self.float32, + "mma_tiler_mn": (256, 256), + "cluster_shape_mn": (2, 1), + "sf_vec_size": 32, + "accumulate_on_output": False, + "expert_cnt": 2, + "input_order": "tensor2d", + "has_global_scale": False, + "cluster_overlap_margin": 0, + }, + ) + + (output,) = captured["outputs"] + self.assertEqual( + (output.name, output.shape, output.dtype, output.fill_value), + ("wgrad_tensor", (2, 384, 640), self.bfloat16, None), + ) + self.assertEqual(output.tensor_spec.layout, (2, 1, 0)) + (workspace,) = captured["workspaces"] + self.assertEqual( + (workspace.name, workspace.shape, workspace.dtype), + ("workspace", (512,), self.uint8), + ) + self.assertEqual(workspace.tensor_spec.ptr_assumed_align, 128) + self.assertEqual( + [spec.layout if spec is not None else None for spec in captured["input_specs"]], + [(1, 0), (0, 1), (1, 0), (1, 0), None], + ) + + def test_global_scales_remain_runtime_inputs(self): + captured = {} + global_a = _Array((2,), self.float32) + global_b = _Array((2,), self.float32) + with ( + self._optional_modules(include_kernel=True), + mock.patch.object( + self.module, + "call_cutedsl", + side_effect=self._fake_call(captured), + ), + ): + self.module.grouped_gemm_wgrad_wrapper_sm100( + *self._inputs(), + global_scale_a=global_a, + global_scale_b=global_b, + ) + + self.assertEqual(captured["inputs"][-2:], (global_a, global_b)) + self.assertEqual(len(captured["input_specs"]), 7) + self.assertTrue(captured["static_args"]["has_global_scale"]) + + def test_accumulating_and_ragged_modes_use_initialized_storage(self): + captured = {} + with ( + self._optional_modules(include_kernel=True), + mock.patch.object( + self.module, + "call_cutedsl", + side_effect=self._fake_call(captured), + ), + ): + self.module.grouped_gemm_wgrad_wrapper_sm100( + *self._inputs(), + accumulate_on_output=True, + input_order="tensor_ragged", + ) + + self.assertEqual(captured["outputs"][0].fill_value, 0.0) + self.assertEqual(captured["workspaces"][0].shape, (1024,)) + self.assertTrue(captured["static_args"]["accumulate_on_output"]) + self.assertEqual(captured["static_args"]["input_order"], "tensor_ragged") + + def test_launcher_preserves_kernel_abi(self): + placeholders = [object() for _ in range(10)] + with self._optional_modules(include_kernel=True): + self.module._launch( + *placeholders, + acc_dtype=self.float32, + mma_tiler_mn=(256, 256), + cluster_shape_mn=(2, 1), + sf_vec_size=32, + accumulate_on_output=False, + expert_cnt=2, + input_order="tensor2d", + has_global_scale=True, + cluster_overlap_margin=1, + ) + + kernel = _Kernel.instances[-1] + self.assertEqual( + kernel.configuration, + { + "sf_vec_size": 32, + "acc_dtype": "cutlass.float32", + "use_2cta_instrs": True, + "mma_tiler_mn": (256, 256), + "cluster_shape_mn": (2, 1), + "accumulate_on_output": False, + "expert_cnt": 2, + "weight_mode": _WeightMode.DENSE, + "input_order": _InputOrder.Tensor2D, + }, + ) + args = kernel.calls[-1] + self.assertEqual( + args[:7], + ( + placeholders[1], + placeholders[2], + placeholders[3], + placeholders[4], + placeholders[8], + placeholders[5], + placeholders[9], + ), + ) + self.assertEqual(args[7], 11) + self.assertIs(args[8], placeholders[0]) + self.assertIs(args[9], placeholders[6]) + self.assertIs(args[10], placeholders[7]) + self.assertIsNone(args[11]) + + def test_rejects_non_fp8_bad_scales_and_partial_global_scale(self): + a, b, sfa, sfb, offsets = self._inputs() + with ( + self._optional_modules(include_kernel=True), + self.assertRaisesRegex(ValueError, "a_tensor.dtype"), + ): + self.module.grouped_gemm_wgrad_wrapper_sm100( + _Array(a.shape, self.bfloat16), + b, + sfa, + sfb, + offsets, + ) + + with ( + self._optional_modules(include_kernel=True), + self.assertRaisesRegex(ValueError, "sfa_tensor must have shape"), + ): + self.module.grouped_gemm_wgrad_wrapper_sm100( + a, + b, + _Array((384, 16), self.float8_e8m0fnu), + sfb, + offsets, + ) + + with ( + self._optional_modules(include_kernel=True), + self.assertRaisesRegex(ValueError, "must be provided together"), + ): + self.module.grouped_gemm_wgrad_wrapper_sm100( + a, + b, + sfa, + sfb, + offsets, + global_scale_a=_Array((2,), self.float32), + ) + + def test_kernel_capabilities_and_child_exports_are_literal(self): + kernel_path = _CUDNN_ROOT / "grouped_gemm" / "grouped_gemm_wgrad" / "moe_blockscaled_grouped_gemm_wgrad.py" + tree = ast.parse(kernel_path.read_text()) + kernel_class = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "BlockScaledMoEGroupedGemmWgradKernel") + constants = { + target.id: ast.literal_eval(node.value) + for node in kernel_class.body + if isinstance(node, ast.Assign) + for target in node.targets + if isinstance(target, ast.Name) and target.id in {"MMA_TILER_M", "MMA_TILER_N", "FP8_SF_VEC_SIZE"} + } + self.assertEqual( + constants, + { + "MMA_TILER_M": (128, 256), + "MMA_TILER_N": (128, 256), + "FP8_SF_VEC_SIZE": 32, + }, + ) + methods = {node.name for node in kernel_class.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))} + self.assertTrue({"require_mma_tiler", "require_cluster_shape", "get_dense_workspace_bytes"} <= methods) + + init_path = kernel_path.with_name("__init__.py") + init_tree = ast.parse(init_path.read_text()) + exports = next( + node.value + for node in init_tree.body + if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == "_API_EXPORTS" for target in node.targets) + ) + self.assertEqual( + ast.literal_eval(exports), + ( + "GroupedGemmWgradSm100", + "grouped_gemm_wgrad_wrapper_sm100", + ), + ) + + +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..0eaf14903 --- /dev/null +++ b/test/python/fe_api/test_jax_import_contract.py @@ -0,0 +1,215 @@ +# 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 code executed during import so a + # shared kernel file may still contain 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 " + f"{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_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_nsa.py b/test/python/fe_api/test_jax_nsa.py new file mode 100644 index 000000000..ee4496ab3 --- /dev/null +++ b/test/python/fe_api/test_jax_nsa.py @@ -0,0 +1,146 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX shape and GPU integration tests for NSA wrappers.""" + +from __future__ import annotations + +import pytest + + +def _jax_dependencies(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + pytest.importorskip("cutlass.jax") + return jax, jnp + + +def _gpu_device(jax, minimum_compute_capability): + for device in jax.local_devices(): + if device.platform != "gpu": + continue + capability = getattr(device, "compute_capability", None) + if capability is None: + continue + if isinstance(capability, (tuple, list)): + major, minor = (int(value) for value in capability[:2]) + capability_number = major * 10 + minor + elif "." in str(capability): + major, minor = (int(value) for value in str(capability).split(".", 1)) + capability_number = major * 10 + minor + else: + capability_number = int(capability) + if capability_number < 10: + capability_number *= 10 + if capability_number >= minimum_compute_capability: + return device + pytest.skip(f"JAX SM{minimum_compute_capability}+ device is not available") + + +@pytest.mark.L0 +def test_jax_sliding_window_attention_abstract_shape(): + jax, jnp = _jax_dependencies() + + from cudnn.jax import TupleDict, sliding_window_attention_wrapper + + q = jax.ShapeDtypeStruct((2, 4, 128, 64), jnp.bfloat16) + k = jax.ShapeDtypeStruct((2, 2, 128, 64), jnp.bfloat16) + v = jax.ShapeDtypeStruct((2, 2, 128, 64), jnp.bfloat16) + result = jax.eval_shape( + lambda q, k, v: sliding_window_attention_wrapper( + q, + k, + v, + left_bound=16, + is_infer=True, + ), + q, + k, + v, + ) + assert isinstance(result, TupleDict) + assert result["o_tensor"].shape == q.shape + assert result["o_tensor"].dtype == jnp.bfloat16 + assert result["stats_tensor"] is None + + +@pytest.mark.L0 +def test_jax_sliding_window_attention_jit(): + jax, jnp = _jax_dependencies() + device = _gpu_device(jax, 80) + + from cudnn.jax import sliding_window_attention_wrapper + + q_shape = (1, 4, 128, 64) + kv_shape = (1, 2, 128, 64) + q = jax.device_put( + jax.random.normal(jax.random.key(0), q_shape, dtype=jnp.bfloat16), + device, + ) + k = jax.device_put( + jax.random.normal(jax.random.key(1), kv_shape, dtype=jnp.bfloat16), + device, + ) + v = jax.device_put( + jax.random.normal(jax.random.key(2), kv_shape, dtype=jnp.bfloat16), + device, + ) + left_bound = 16 + scale = 0.125 + + @jax.jit + def run(q, k, v): + return sliding_window_attention_wrapper( + q, + k, + v, + left_bound=left_bound, + right_bound=0, + is_infer=True, + attn_scale=scale, + ) + + lowered = run.lower(q, k, v) + stablehlo = lowered.as_text("stablehlo") + assert "stablehlo.custom_call" in stablehlo + assert "__cudnn$fmha" in stablehlo + + result = lowered.compile()(q, k, v) + result["o_tensor"].block_until_ready() + assert result["stats_tensor"] is None + + k_expanded = jnp.repeat( + k.astype(jnp.float32), + q.shape[1] // k.shape[1], + axis=1, + ) + v_expanded = jnp.repeat( + v.astype(jnp.float32), + q.shape[1] // v.shape[1], + axis=1, + ) + logits = ( + jnp.einsum( + "bhqd,bhkd->bhqk", + q.astype(jnp.float32), + k_expanded, + ) + * scale + ) + q_position = jnp.arange(q.shape[2])[:, None] + k_position = jnp.arange(k.shape[2])[None, :] + valid = (k_position > q_position - left_bound) & (k_position <= q_position) + probability = jax.nn.softmax( + jnp.where(valid[None, None, :, :], logits, -jnp.inf), + axis=-1, + ) + reference = jnp.einsum("bhqk,bhkd->bhqd", probability, v_expanded) + max_error = jnp.max( + jnp.abs(result["o_tensor"].astype(jnp.float32) - reference.astype(jnp.float32)) + ) + assert jnp.allclose( + result["o_tensor"], + reference, + atol=3e-2, + rtol=3e-2, + ), f"maximum absolute error: {max_error}" diff --git a/test/python/fe_api/test_jax_nsa_contract.py b/test/python/fe_api/test_jax_nsa_contract.py new file mode 100644 index 000000000..9f0070766 --- /dev/null +++ b/test/python/fe_api/test_jax_nsa_contract.py @@ -0,0 +1,654 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contract tests for the JAX NSA wrappers.""" + +from __future__ import annotations + +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 + + +def _identity_jit(fn=None, **_kwargs): + return (lambda decorated_fn: decorated_fn) if fn is None else fn + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CUDNN_ROOT = _REPO_ROOT / "python" / "cudnn" +_TEST_PACKAGE = "cudnn_frontend_jax_nsa_contract_test" + + +class _DType: + def __init__(self, name, itemsize): + 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 + + def astype(self, dtype): + return _Array(self.shape, dtype) + + +class _TensorSpec: + def __init__( + self, + *, + layout=None, + mode=None, + static=None, + ptr_assumed_align=256, + divisibility=None, + ): + self.layout = layout + self.mode = mode + self.static = static + self.ptr_assumed_align = ptr_assumed_align + self.divisibility = divisibility + + +class JaxNsaContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.bfloat16 = _DType("bfloat16", 2) + cls.float16 = _DType("float16", 2) + cls.float32 = _DType("float32", 4) + cls.int32 = _DType("int32", 4) + cls.int64 = _DType("int64", 8) + + cls.fake_jnp = types.ModuleType("jax.numpy") + cls.fake_jnp.bfloat16 = cls.bfloat16 + cls.fake_jnp.float16 = cls.float16 + cls.fake_jnp.float32 = cls.float32 + cls.fake_jnp.int32 = cls.int32 + cls.fake_jnp.int64 = cls.int64 + cls.fake_jnp.dtype = lambda value: value + cls.fake_jnp.reshape = lambda value, shape: _Array(shape, value.dtype) + cls.fake_jnp.transpose = lambda value, axes: _Array( + tuple(value.shape[axis] for axis in axes), + value.dtype, + ) + + cls.fake_jax = types.ModuleType("jax") + cls.fake_jax.__path__ = [] + cls.fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + cls.fake_jax.numpy = cls.fake_jnp + cls.fake_jax.nn = types.SimpleNamespace() + cls.fake_jax.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, + ) + + cls.fake_cutlass = types.ModuleType("cutlass") + cls.fake_cutlass.__path__ = [] + cls.fake_cutlass.Constexpr = object + cls.fake_cutlass_cute = types.ModuleType("cutlass.cute") + cls.fake_cutlass_cute.jit = _identity_jit + cls.fake_cutlass.cute = cls.fake_cutlass_cute + cls.fake_cutlass_jax = types.ModuleType("cutlass.jax") + cls.fake_cutlass_jax.TensorSpec = _TensorSpec + cls.fake_cutlass_jax.jax_to_cutlass_dtype = lambda dtype: f"cutlass.{dtype.name}" + cls.fake_cutlass.jax = cls.fake_cutlass_jax + + parent = types.ModuleType(_TEST_PACKAGE) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = _TEST_PACKAGE + sys.modules[_TEST_PACKAGE] = parent + with mock.patch.dict( + sys.modules, + { + "jax": cls.fake_jax, + "jax.numpy": cls.fake_jnp, + "cutlass": cls.fake_cutlass, + "cutlass.cute": cls.fake_cutlass_cute, + "cutlass.jax": cls.fake_cutlass_jax, + "torch": None, + }, + ): + cls.nsa = importlib.import_module(f"{_TEST_PACKAGE}.native_sparse_attention") + cls.compression_package = importlib.import_module(f"{_TEST_PACKAGE}.native_sparse_attention.compression") + cls.selection_package = importlib.import_module(f"{_TEST_PACKAGE}.native_sparse_attention.selection") + cls.sliding_package = importlib.import_module( + f"{_TEST_PACKAGE}.native_sparse_attention.sliding_window_attention" + ) + cls.topk_package = importlib.import_module(f"{_TEST_PACKAGE}.native_sparse_attention.top_k") + cls.compression = importlib.import_module(f"{_TEST_PACKAGE}.native_sparse_attention.compression.jax") + cls.selection = importlib.import_module(f"{_TEST_PACKAGE}.native_sparse_attention.selection.jax") + cls.sliding = importlib.import_module( + f"{_TEST_PACKAGE}.native_sparse_attention.sliding_window_attention.jax" + ) + cls.topk = importlib.import_module(f"{_TEST_PACKAGE}.native_sparse_attention.top_k.jax") + + @classmethod + def tearDownClass(cls): + for module_name in tuple(sys.modules): + if module_name == _TEST_PACKAGE or module_name.startswith(f"{_TEST_PACKAGE}."): + sys.modules.pop(module_name, None) + + def _optional_modules(self): + return mock.patch.dict( + sys.modules, + { + "jax": self.fake_jax, + "jax.numpy": self.fake_jnp, + "cutlass": self.fake_cutlass, + "cutlass.cute": self.fake_cutlass_cute, + "cutlass.jax": self.fake_cutlass_jax, + }, + ) + + def test_jax_submodules_do_not_load_torch_apis(self): + self.assertNotIn( + f"{_TEST_PACKAGE}.native_sparse_attention.compression.api", + sys.modules, + ) + self.assertNotIn( + f"{_TEST_PACKAGE}.native_sparse_attention.top_k.api", + sys.modules, + ) + self.assertNotIn( + f"{_TEST_PACKAGE}.native_sparse_attention.selection.api", + sys.modules, + ) + self.assertNotIn( + f"{_TEST_PACKAGE}.native_sparse_attention.sliding_window_attention.api", + sys.modules, + ) + self.assertNotIn( + f"{_TEST_PACKAGE}.native_sparse_attention.compression.fmha", + sys.modules, + ) + self.assertNotIn( + f"{_TEST_PACKAGE}.native_sparse_attention.top_k.nsa_top_k_reduction_fwd", + sys.modules, + ) + self.assertNotIn( + f"{_TEST_PACKAGE}.native_sparse_attention.selection.NSA_select_attn_fwd_hmma", + sys.modules, + ) + self.assertIs(self.compression_package.jax, self.compression) + self.assertIs(self.selection_package.jax, self.selection) + self.assertIs(self.sliding_package.jax, self.sliding) + self.assertIs(self.topk_package.jax, self.topk) + + def test_sliding_window_attention_uses_public_cudnn_lowering(self): + calls = [] + + def dot_product_attention(q, k, v, **options): + calls.append((q, k, v, options)) + return _Array(q.shape, q.dtype) + + q = _Array((2, 4, 16, 64), self.bfloat16) + k = _Array((2, 2, 16, 64), self.bfloat16) + v = _Array((2, 2, 16, 64), self.bfloat16) + with ( + self._optional_modules(), + mock.patch.object( + self.fake_jax.nn, + "dot_product_attention", + side_effect=dot_product_attention, + create=True, + ), + ): + result = self.sliding.sliding_window_attention_wrapper( + q, + k, + v, + left_bound=8, + right_bound=0, + is_infer=True, + attn_scale=0.125, + ) + + self.assertEqual( + (result["o_tensor"].shape, result["o_tensor"].dtype), + (q.shape, self.bfloat16), + ) + self.assertIsNone(result["stats_tensor"]) + q_btnh, k_btnh, v_btnh, options = calls[0] + self.assertEqual(q_btnh.shape, (2, 16, 4, 64)) + self.assertEqual(k_btnh.shape, (2, 16, 2, 64)) + self.assertEqual(v_btnh.shape, (2, 16, 2, 64)) + self.assertEqual( + options, + { + "scale": 0.125, + "is_causal": True, + "local_window_size": (7, 0), + "implementation": "cudnn", + }, + ) + + def test_sliding_window_attention_rejects_unsupported_training_and_window(self): + q = _Array((2, 4, 16, 64), self.bfloat16) + k = _Array((2, 2, 16, 64), self.bfloat16) + v = _Array((2, 2, 16, 64), self.bfloat16) + with ( + self._optional_modules(), + mock.patch.object( + self.fake_jax.nn, + "dot_product_attention", + create=True, + ) as attention, + ): + with self.assertRaisesRegex(NotImplementedError, "inference only"): + self.sliding.sliding_window_attention_wrapper( + q, + k, + v, + is_infer=False, + ) + with self.assertRaisesRegex(NotImplementedError, "right_bound=0"): + self.sliding.sliding_window_attention_wrapper( + q, + k, + v, + right_bound=1, + ) + with self.assertRaisesRegex(NotImplementedError, "S_q == S_k"): + self.sliding.sliding_window_attention_wrapper( + q, + _Array((2, 2, 8, 64), self.bfloat16), + _Array((2, 2, 8, 64), self.bfloat16), + ) + with self.assertRaisesRegex(ValueError, "at least 1"): + self.sliding.sliding_window_attention_wrapper( + q, + k, + v, + left_bound=0, + ) + attention.assert_not_called() + + def test_selection_attention_declares_runtime_offsets_and_safe_outputs(self): + captured = {} + + def fake_call(launcher, inputs, **options): + captured.update(launcher=launcher, inputs=inputs, **options) + return tuple(_Array(spec.shape, spec.dtype) for spec in options["outputs"]) + + q = _Array((16, 4, 128), self.float16) + k = _Array((16, 1, 128), self.float16) + v = _Array((16, 1, 64), self.float16) + block_indices = _Array((16, 1, 8), self.int32) + block_counts = _Array((16, 1), self.int32) + cum_q = _Array((3,), self.int32) + cum_k = _Array((3,), self.int32) + with ( + self._optional_modules(), + mock.patch.object( + self.selection, + "call_cutedsl", + side_effect=fake_call, + ), + ): + result = self.selection.selection_attention_wrapper( + q, + k, + v, + block_indices, + block_counts, + cum_q, + cum_k, + max_s_q=8, + max_s_k=8, + ) + + self.assertEqual((result["o_tensor"].shape, result["o_tensor"].dtype), ((16, 4, 64), self.float16)) + self.assertEqual((result["l_tensor"].shape, result["l_tensor"].dtype), ((16, 4, 1), self.float32)) + self.assertEqual((result["m_tensor"].shape, result["m_tensor"].dtype), ((16, 4, 1), self.float32)) + self.assertIs(captured["launcher"], self.selection._launch) + self.assertEqual(tuple(value.shape for value in captured["inputs"][:3]), ((1, 16, 4, 128), (1, 16, 1, 128), (1, 16, 1, 64))) + self.assertEqual(captured["inputs"][3:], (block_indices, block_counts, cum_q, cum_k)) + + output, lse_sum, row_max = captured["outputs"] + self.assertEqual(output.shape, (1, 16, 4, 64)) + self.assertEqual(output.fill_value, 0) + self.assertEqual(lse_sum.shape, (1, 16, 4)) + self.assertEqual(lse_sum.fill_value, 0.0) + self.assertEqual(row_max.shape, (1, 16, 4)) + self.assertEqual(row_max.fill_value, float("-inf")) + + config = captured["static_args"] + self.assertEqual(config["element_dtype"], "cutlass.float16") + self.assertEqual(config["gqa_group_size"], 4) + self.assertEqual(config["max_s_q"], 8) + self.assertAlmostEqual(config["scale_softmax"], 1.0 / (128**0.5)) + + def test_selection_launcher_preserves_native_argument_order(self): + calls = [] + kernel_options = [] + + class FakeKernel: + def __init__(self, **options): + kernel_options.append(options) + + def __call__(self, *args): + calls.append(args) + + kernel_name = f"{_TEST_PACKAGE}.native_sparse_attention.selection.NSA_select_attn_fwd_hmma" + kernel_module = types.ModuleType(kernel_name) + kernel_module.HopperSelectAttentionFwd = FakeKernel + + def float32(value=None): + return ("Float32", value) + + with ( + self._optional_modules(), + mock.patch.dict(sys.modules, {kernel_name: kernel_module}), + mock.patch.object(self.fake_cutlass, "Float32", float32, create=True), + ): + self.selection._launch( + "stream", + "q", + "k", + "v", + "indices", + "counts", + "cum_q", + "cum_k", + "out", + "l", + "m", + element_dtype="float16", + head_dim=128, + value_dim=64, + gqa_group_size=4, + block_size=64, + max_s_q=8, + scale_softmax=0.125, + ) + + self.assertEqual(kernel_options[0]["head_dim"], 128) + self.assertEqual(kernel_options[0]["GQA_group_size"], 4) + self.assertEqual( + calls[0], + ( + "q", + "k", + "v", + "out", + "l", + "m", + "indices", + "counts", + 8, + "cum_q", + ("Float32", 0.125), + "stream", + ), + ) + + def test_selection_validation_fails_before_launch(self): + q = _Array((16, 4, 128), self.float16) + k = _Array((16, 1, 128), self.float16) + v = _Array((16, 1, 128), self.float16) + indices = _Array((16, 1, 8), self.int32) + counts = _Array((16, 1), self.int32) + offsets = _Array((3,), self.int32) + with ( + self._optional_modules(), + self.assertRaisesRegex(ValueError, "must be identical"), + mock.patch.object(self.selection, "call_cutedsl") as call_cutedsl, + ): + self.selection.selection_attention_wrapper( + q, + k, + v, + indices, + counts, + offsets, + offsets, + max_s_q=8, + max_s_k=16, + ) + call_cutedsl.assert_not_called() + + def test_compression_attention_declares_layouts_and_optional_lse(self): + captured = {} + + def fake_call(launcher, inputs, **options): + captured.update(launcher=launcher, inputs=inputs, **options) + return tuple(_Array(spec.shape, spec.dtype) for spec in options["outputs"]) + + q = _Array((2, 4, 16, 64), self.float16) + k = _Array((2, 2, 8, 64), self.float16) + v = _Array((2, 2, 8, 64), self.float16) + with ( + self._optional_modules(), + mock.patch.object( + self.compression, + "call_cutedsl", + side_effect=fake_call, + ), + ): + result = self.compression.compression_attention_wrapper( + q, + k, + v, + enable_lse=True, + o_dtype=self.bfloat16, + scale_q=2.0, + scale_k=0.5, + scale_v=3.0, + inv_scale_o=0.25, + ) + + self.assertEqual((result["o_tensor"].shape, result["o_tensor"].dtype), ((2, 4, 16, 64), self.bfloat16)) + self.assertEqual((result["lse_tensor"].shape, result["lse_tensor"].dtype), ((2, 4, 16), self.float32)) + self.assertIs(captured["launcher"], self.compression._launch) + self.assertEqual(captured["inputs"], (q, k, v)) + self.assertEqual(len(captured["input_specs"]), 3) + for spec in captured["input_specs"]: + self.assertEqual(spec.layout, (3, 1, 2, 0)) + self.assertEqual(spec.mode, (0, 2, 1, 3)) + + output, lse = captured["outputs"] + self.assertEqual(output.name, "o_tensor") + self.assertIsNone(output.fill_value) + self.assertEqual(output.tensor_spec.layout, (3, 1, 2, 0)) + self.assertEqual(output.tensor_spec.mode, (0, 2, 1, 3)) + self.assertEqual(lse.name, "lse_tensor") + self.assertEqual(lse.tensor_spec.layout, (2, 1, 0)) + self.assertEqual(lse.tensor_spec.mode, (0, 2, 1)) + + config = captured["static_args"] + self.assertTrue(config["enable_lse"]) + self.assertEqual(config["scale_softmax"], 0.125) + self.assertEqual(config["scale_output"], 0.75) + + def test_compression_launcher_preserves_native_argument_order(self): + calls = [] + kernel_options = [] + + class FakeKernel: + def __init__(self, **options): + kernel_options.append(options) + + def __call__(self, *args): + calls.append(args) + + package = f"{_TEST_PACKAGE}.native_sparse_attention.compression" + kernel_module = types.ModuleType(f"{package}.fmha") + kernel_module.BlackwellFusedMultiHeadAttentionForward = FakeKernel + helpers_module = types.ModuleType(f"{package}.fmha_helpers") + helpers_module.MaskType = types.SimpleNamespace(COMPRESSED_CAUSAL_MASK="compressed") + + def float32(value=None): + return ("Float32", value) + + def int32(value): + return ("Int32", value) + + with ( + self._optional_modules(), + mock.patch.dict( + sys.modules, + { + kernel_module.__name__: kernel_module, + helpers_module.__name__: helpers_module, + }, + ), + mock.patch.object(self.fake_cutlass, "Float32", float32, create=True), + mock.patch.object(self.fake_cutlass, "Int32", int32, create=True), + ): + self.compression._launch( + "stream", + "q", + "k", + "v", + "out", + "lse", + batch=2, + seqlen_q=16, + seqlen_k=8, + num_query_heads=4, + num_kv_heads=2, + head_dim=64, + enable_lse=True, + is_persistent=False, + scale_softmax=0.125, + scale_output=0.75, + ) + self.compression._launch( + "stream", + "q", + "k", + "v", + "out", + batch=2, + seqlen_q=16, + seqlen_k=8, + num_query_heads=4, + num_kv_heads=2, + head_dim=64, + enable_lse=False, + is_persistent=False, + scale_softmax=0.125, + scale_output=0.75, + ) + + self.assertEqual(kernel_options[0]["qk_acc_dtype"], float32) + self.assertEqual(kernel_options[0]["mask_type"], "compressed") + args = calls[0] + self.assertEqual(args[:4], ("q", "k", "v", "out")) + self.assertEqual(args[4][0], ("Int32", 2)) + self.assertEqual(args[4][1:4], (("Int32", 16), ("Int32", 16), ("Int32", 8))) + self.assertEqual(args[7], "lse") + self.assertEqual(args[-3:], (None, ("Int32", 0), "stream")) + self.assertIsNone(calls[1][7]) + + def test_topk_reduction_declares_initialized_outputs(self): + captured = {} + + def fake_call(launcher, inputs, **options): + captured.update(launcher=launcher, inputs=inputs, **options) + return tuple(_Array(spec.shape, spec.dtype) for spec in options["outputs"]) + + q = _Array((2, 4, 16, 64), self.bfloat16) + k = _Array((2, 2, 8, 64), self.bfloat16) + lse = _Array((2, 4, 16), self.float32) + with ( + self._optional_modules(), + mock.patch.object(self.topk, "call_cutedsl", side_effect=fake_call), + ): + result = self.topk.topk_reduction_wrapper(q, k, lse) + + self.assertEqual((result["topk_scores_tensor"].shape, result["topk_scores_tensor"].dtype), ((2, 2, 16, 16), self.float32)) + self.assertEqual((result["topk_indices_tensor"].shape, result["topk_indices_tensor"].dtype), ((2, 2, 16, 16), self.int32)) + self.assertEqual(captured["inputs"], (q, k, lse)) + self.assertIs(captured["launcher"], self.topk._launch) + scores, indices = captured["outputs"] + self.assertEqual(scores.fill_value, float("-inf")) + self.assertEqual(indices.fill_value, -1) + for spec in (scores.tensor_spec, indices.tensor_spec): + self.assertEqual(spec.layout, (3, 1, 2, 0)) + self.assertEqual(spec.mode, (0, 1, 2, 3)) + self.assertEqual(captured["input_specs"][2].layout, (2, 1, 0)) + self.assertEqual(captured["static_args"]["element_dtype"], "cutlass.bfloat16") + + def test_topk_launcher_preserves_native_argument_order(self): + calls = [] + + class FakeKernel: + def __init__(self, **options): + self.options = options + + def __call__(self, *args): + calls.append(args) + + kernel_name = f"{_TEST_PACKAGE}.native_sparse_attention.top_k.nsa_top_k_reduction_fwd" + kernel_module = types.ModuleType(kernel_name) + kernel_module.FineGrainedReductionQK = FakeKernel + + def float32(value=None): + return ("Float32", value) + + def int32(value): + return ("Int32", value) + + with ( + self._optional_modules(), + mock.patch.dict(sys.modules, {kernel_name: kernel_module}), + mock.patch.object(self.fake_cutlass, "Float32", float32, create=True), + mock.patch.object(self.fake_cutlass, "Int32", int32, create=True), + ): + self.topk._launch( + "stream", + "q", + "k", + "lse", + "scores", + "indices", + element_dtype="float16", + batch=2, + seqlen_q=16, + seqlen_k=8, + num_query_heads=4, + num_kv_heads=2, + head_dim=64, + k_value=16, + selection_block_size=64, + compress_stride=32, + is_causal=True, + scale_softmax=0.125, + ) + + args = calls[0] + self.assertEqual(args[1:6], ("q", "k", "lse", "scores", "indices")) + self.assertEqual(args[0][0], ("Int32", 2)) + self.assertEqual(args[-3:], (None, None, "stream")) + + def test_topk_validation_fails_before_launch(self): + q = _Array((1, 4, 16, 64), self.float16) + k = _Array((1, 2, 8, 64), self.float16) + lse = _Array((1, 4, 16), self.float32) + with ( + self._optional_modules(), + self.assertRaisesRegex(ValueError, "positive multiple of 4"), + mock.patch.object(self.topk, "call_cutedsl") as call_cutedsl, + ): + self.topk.topk_reduction_wrapper(q, k, lse, k_value=7) + call_cutedsl.assert_not_called() + + +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..2445d2e81 --- /dev/null +++ b/test/python/fe_api/test_jax_rmsnorm_rht_amax.py @@ -0,0 +1,144 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""GPU integration tests for the JAX RMSNorm + RHT + amax wrapper.""" + +import math + +import pytest + + +def _hadamard_16(jnp): + values = [[1.0 if ((row & col).bit_count() % 2 == 0) else -1.0 for col 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_shape(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + pytest.importorskip("cutlass.jax") + + from cudnn.jax import rmsnorm_rht_amax_sm100 + + x = jax.ShapeDtypeStruct((256, 2048), jnp.bfloat16) + weight = jax.ShapeDtypeStruct((2048,), jnp.bfloat16) + output, amax = jax.eval_shape(rmsnorm_rht_amax_sm100, x, weight) + + assert output.shape == (256, 2048) + assert output.dtype == jnp.bfloat16 + assert amax.shape == (128,) + assert amax.dtype == jnp.float32 + + +@pytest.mark.L0 +def test_jax_rmsnorm_rht_amax_class_uses_sample_metadata_only(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + pytest.importorskip("cutlass.jax") + + from cudnn.jax import RmsNormRhtAmaxSm100 + + sample_x = jax.ShapeDtypeStruct((256, 2048), jnp.bfloat16) + sample_weight = jax.ShapeDtypeStruct((2048,), jnp.bfloat16) + api = RmsNormRhtAmaxSm100(sample_x, sample_weight) + + assert not hasattr(api, "sample_x") + assert not hasattr(api, "sample_w") + assert all(value is not sample_x and value is not sample_weight for value in vars(api).values()) + assert api.check_support() + assert api.num_threads == 128 + assert api.rows_per_cta == 2 + + 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 + + wrong_x = jax.ShapeDtypeStruct((128, 2048), jnp.bfloat16) + with pytest.raises(ValueError, match="X tensor shape mismatch"): + jax.eval_shape(api, wrong_x, sample_weight) + + wrong_dtype_x = jax.ShapeDtypeStruct((256, 2048), jnp.float16) + with pytest.raises(ValueError, match="X dtype mismatch"): + jax.eval_shape(api, wrong_dtype_x, sample_weight) + + +@pytest.mark.L0 +def test_jax_rmsnorm_rht_amax_jit(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + pytest.importorskip("cutlass.jax") + + gpu_devices = [device for device in jax.local_devices() if device.platform == "gpu"] + if not gpu_devices: + pytest.skip("JAX CUDA device is not available") + + capable_devices = [] + reported_capabilities = [] + 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)) + reported_capabilities.append(f"SM{major}{minor}") + if major * 10 + minor >= 100: + capable_devices.append(device) + if not capable_devices: + reported = ", ".join(reported_capabilities) or "unknown capabilities" + pytest.skip(f"RMSNorm + RHT requires SM100+; local GPUs report {reported}") + device = capable_devices[0] + + from cudnn.jax import RmsNormRhtAmaxSm100 + + m, n = 256, 2048 + rows_per_cta = 2 + eps = 1e-5 + 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, + ) + + api = RmsNormRhtAmaxSm100( + jax.ShapeDtypeStruct((m, n), jnp.bfloat16), + jax.ShapeDtypeStruct((n,), jnp.bfloat16), + eps=eps, + num_threads=128, + rows_per_cta=rows_per_cta, + ) + assert api.check_support() + run = jax.jit(api) + + lowered = run.lower(x, weight) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + compiled = lowered.compile() + output, amax = compiled(x, weight) + output.block_until_ready() + second_output, second_amax = compiled(x, weight) + second_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 output.shape == (m, n) + assert output.dtype == jnp.bfloat16 + assert amax.shape == (m // rows_per_cta,) + assert amax.dtype == jnp.float32 + 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) + assert jnp.array_equal(second_output, output) + assert jnp.array_equal(second_amax, amax) diff --git a/test/python/fe_api/test_jax_sdpa_contract.py b/test/python/fe_api/test_jax_sdpa_contract.py new file mode 100644 index 000000000..9d13a36e0 --- /dev/null +++ b/test/python/fe_api/test_jax_sdpa_contract.py @@ -0,0 +1,443 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for the JAX SDPA wrappers.""" + +from __future__ import annotations + +from importlib import import_module +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 + + +def _identity_jit(fn=None, **_kwargs): + return (lambda decorated_fn: decorated_fn) if fn is None else fn + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CUDNN_ROOT = _REPO_ROOT / "python" / "cudnn" +_TEST_PACKAGE = "cudnn_frontend_jax_sdpa_contract_test" + + +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 + + +class _TensorSpec: + def __init__( + self, + *, + layout=None, + mode=None, + static=None, + ptr_assumed_align=256, + divisibility=None, + ): + self.layout = layout + self.mode = mode + self.static = static + self.ptr_assumed_align = ptr_assumed_align + self.divisibility = divisibility + + +class _Float32: + width = 32 + + +class _BackwardKernel: + @staticmethod + def get_workspace_size(seqlen_q, head_dim, num_query_heads, batch, acc_dtype): + assert acc_dtype is _Float32 + return ( + batch, + ((seqlen_q + 7) // 8) * 8, + num_query_heads, + 2 * (acc_dtype.width // 8) + head_dim * (acc_dtype.width // 8), + ) + + +class JaxSdpaContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.float16 = _DType("float16", 2) + cls.bfloat16 = _DType("bfloat16", 2) + cls.float32 = _DType("float32", 4) + cls.uint8 = _DType("uint8", 1) + + cls.fake_jnp = types.ModuleType("jax.numpy") + cls.fake_jnp.float16 = cls.float16 + cls.fake_jnp.bfloat16 = cls.bfloat16 + cls.fake_jnp.float32 = cls.float32 + cls.fake_jnp.uint8 = cls.uint8 + cls.fake_jnp.dtype = lambda value: value + + cls.fake_jax = types.ModuleType("jax") + cls.fake_jax.__path__ = [] + cls.fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + cls.fake_jax.numpy = cls.fake_jnp + cls.fake_jax.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, + ) + cls.fake_jax.ShapeDtypeStruct = lambda shape, dtype: (shape, dtype) + + cls.fake_cutlass = types.ModuleType("cutlass") + cls.fake_cutlass.__path__ = [] + cls.fake_cutlass.Constexpr = object + cls.fake_cutlass.Float32 = _Float32 + cls.fake_cutlass.Int32 = int + cls.fake_cutlass_cute = types.ModuleType("cutlass.cute") + cls.fake_cutlass_cute.jit = _identity_jit + cls.fake_cutlass.cute = cls.fake_cutlass_cute + cls.fake_cutlass_jax = types.ModuleType("cutlass.jax") + cls.fake_cutlass_jax.TensorSpec = _TensorSpec + cls.fake_cutlass_jax.jax_to_cutlass_dtype = lambda dtype: f"cutlass.{dtype.name}" + cls.fake_cutlass_jax.cutlass_call = None + cls.fake_cutlass.jax = cls.fake_cutlass_jax + + cls.backward_kernel_module_name = f"{_TEST_PACKAGE}.sdpa.bwd.fmha_backward_sm100_2kernel" + cls.backward_kernel_module = types.ModuleType(cls.backward_kernel_module_name) + cls.backward_kernel_module.BlackwellFusedMultiHeadAttentionBackward = _BackwardKernel + + package_paths = { + _TEST_PACKAGE: _CUDNN_ROOT, + f"{_TEST_PACKAGE}.sdpa": _CUDNN_ROOT / "sdpa", + f"{_TEST_PACKAGE}.sdpa.fwd": _CUDNN_ROOT / "sdpa" / "fwd", + f"{_TEST_PACKAGE}.sdpa.bwd": _CUDNN_ROOT / "sdpa" / "bwd", + } + for package_name, package_path in package_paths.items(): + package = types.ModuleType(package_name) + package.__path__ = [str(package_path)] + package.__package__ = package_name + sys.modules[package_name] = package + + with cls._optional_modules(): + cls.forward = import_module(f"{_TEST_PACKAGE}.sdpa.fwd.jax") + cls.backward = import_module(f"{_TEST_PACKAGE}.sdpa.bwd.jax") + + @classmethod + def tearDownClass(cls): + for module_name in tuple(sys.modules): + if module_name == _TEST_PACKAGE or module_name.startswith(f"{_TEST_PACKAGE}."): + sys.modules.pop(module_name, None) + + @classmethod + def _optional_modules(cls, *, include_backward_kernel=False): + modules = { + "jax": cls.fake_jax, + "jax.numpy": cls.fake_jnp, + "cutlass": cls.fake_cutlass, + "cutlass.cute": cls.fake_cutlass_cute, + "cutlass.jax": cls.fake_cutlass_jax, + } + if include_backward_kernel: + modules[cls.backward_kernel_module_name] = cls.backward_kernel_module + return mock.patch.dict(sys.modules, modules) + + @staticmethod + def _fake_call(captured): + def call(launcher, inputs, **options): + captured.update(launcher=launcher, inputs=inputs, **options) + return tuple(_Array(spec.shape, spec.dtype) for spec in options["outputs"]) + + return call + + def test_kernel_implementations_are_deferred_until_wrapper_call(self): + self.assertNotIn( + f"{_TEST_PACKAGE}.sdpa.fwd.fmha_forward_sm100_d256", + sys.modules, + ) + self.assertNotIn(self.backward_kernel_module_name, sys.modules) + + def test_forward_declares_bhsd_layout_and_functional_outputs(self): + captured = {} + q = _Array((2, 8, 64, 256), self.float16) + k = _Array((2, 2, 96, 256), self.float16) + v = _Array((2, 2, 96, 256), self.float16) + + with ( + self._optional_modules(), + mock.patch.object( + self.forward, + "call_cutedsl", + side_effect=self._fake_call(captured), + ), + ): + result = self.forward.sdpa_fwd_wrapper_sm100_d256(q, k, v) + + self.assertEqual(result["o_tensor"].shape, q.shape) + self.assertEqual(result["lse_tensor"].shape, (2, 8, 64)) + self.assertEqual(captured["inputs"], (q, k, v)) + self.assertNotIn("workspaces", captured) + self.assertIs(captured["launcher"], self.forward._launch) + self.assertEqual( + captured["static_args"], + { + "batch": 2, + "seqlen_q": 64, + "seqlen_k": 96, + "num_query_heads": 8, + "num_kv_heads": 2, + "scale_softmax": 0.0625, + "scale_output": 1.0, + "window_size_left": -1, + "window_size_right": -1, + "mask_kind": "residual", + }, + ) + + output, lse = captured["outputs"] + self.assertEqual( + (output.name, output.shape, output.dtype), + ("o_tensor", q.shape, self.float16), + ) + self.assertEqual( + (lse.name, lse.shape, lse.dtype), + ("lse_tensor", (2, 8, 64), self.float32), + ) + self.assertEqual(output.tensor_spec.layout, (3, 1, 2, 0)) + self.assertEqual(output.tensor_spec.mode, (0, 2, 1, 3)) + self.assertTrue(all(spec is output.tensor_spec for spec in captured["input_specs"])) + + def test_backward_declares_zero_initialized_hidden_workspace(self): + captured = {} + q = _Array((2, 8, 65, 256), self.bfloat16) + k = _Array((2, 2, 96, 256), self.bfloat16) + v = _Array((2, 2, 96, 256), self.bfloat16) + output = _Array(q.shape, self.bfloat16) + doutput = _Array(q.shape, self.bfloat16) + lse = _Array((2, 8, 65), self.float32) + + with ( + self._optional_modules(include_backward_kernel=True), + mock.patch.object( + self.backward, + "call_cutedsl", + side_effect=self._fake_call(captured), + ), + ): + result = self.backward.sdpa_bwd_wrapper_sm100_d256( + q, + k, + v, + output, + doutput, + lse, + ) + + self.assertEqual(result["dq_tensor"].shape, q.shape) + self.assertEqual(result["dk_tensor"].shape, k.shape) + self.assertEqual(result["dv_tensor"].shape, v.shape) + self.assertEqual(captured["inputs"], (q, k, v, output, doutput, lse)) + self.assertIs(captured["launcher"], self.backward._launch) + self.assertEqual( + [(spec.name, spec.shape, spec.dtype) for spec in captured["outputs"]], + [ + ("dq_tensor", q.shape, self.bfloat16), + ("dk_tensor", k.shape, self.bfloat16), + ("dv_tensor", v.shape, self.bfloat16), + ], + ) + (workspace,) = captured["workspaces"] + self.assertEqual(workspace.name, "workspace") + self.assertEqual(workspace.shape, (2, 72, 8, 1032)) + self.assertEqual(workspace.dtype, self.uint8) + self.assertEqual(workspace.fill_value, 0) + self.assertEqual(captured["static_args"]["element_dtype"], "cutlass.bfloat16") + self.assertEqual(captured["static_args"]["mask_kind"], "residual") + + def test_forward_launcher_preserves_native_argument_order(self): + calls = [] + kernel_options = [] + + class FakeKernel: + def __init__(self, **options): + kernel_options.append(options) + + def __call__(self, *args): + calls.append(args) + + mask_module_name = f"{_TEST_PACKAGE}.sdpa.fmha_utils" + mask_module = types.ModuleType(mask_module_name) + mask_module.MaskEnum = types.SimpleNamespace( + RESIDUAL_MASK="residual", + WINDOW_MASK_INFERENCE="window", + ) + kernel_module_name = f"{_TEST_PACKAGE}.sdpa.fwd.fmha_forward_sm100_d256" + kernel_module = types.ModuleType(kernel_module_name) + kernel_module.BlackwellFusedMultiHeadAttentionForward = FakeKernel + + def float32(value=None): + return ("Float32", value) + + def int32(value): + return ("Int32", value) + + with ( + self._optional_modules(), + mock.patch.dict( + sys.modules, + { + mask_module_name: mask_module, + kernel_module_name: kernel_module, + }, + ), + mock.patch.object(self.fake_cutlass, "Float32", float32), + mock.patch.object(self.fake_cutlass, "Int32", int32), + ): + self.forward._launch( + "stream", + "q", + "k", + "v", + "out", + "lse", + batch=2, + seqlen_q=64, + seqlen_k=96, + num_query_heads=8, + num_kv_heads=2, + scale_softmax=0.0625, + scale_output=0.5, + window_size_left=32, + window_size_right=0, + mask_kind="window", + ) + + self.assertEqual(kernel_options[0]["mask_type"], "window") + args = calls[0] + self.assertEqual(args[:4], ("q", "k", "v", "out")) + self.assertEqual(args[4][:3], (("Int32", 2), ("Int32", 64), ("Int32", 96))) + self.assertEqual(args[7], "lse") + self.assertEqual(args[-3:], (("Int32", 32), ("Int32", 0), "stream")) + + def test_backward_launcher_preserves_native_argument_order(self): + calls = [] + kernel_options = [] + + class FakeKernel: + def __init__(self, **options): + kernel_options.append(options) + + def __call__(self, *args): + calls.append(args) + + mask_module_name = f"{_TEST_PACKAGE}.sdpa.fmha_utils" + mask_module = types.ModuleType(mask_module_name) + mask_module.MaskEnum = types.SimpleNamespace( + RESIDUAL_MASK="residual", + WINDOW_MASK_INFERENCE="window", + ) + kernel_module = types.ModuleType(self.backward_kernel_module_name) + kernel_module.BlackwellFusedMultiHeadAttentionBackward = FakeKernel + + def float32(value=None): + return ("Float32", value) + + def int32(value): + return ("Int32", value) + + with ( + self._optional_modules(), + mock.patch.dict( + sys.modules, + { + mask_module_name: mask_module, + self.backward_kernel_module_name: kernel_module, + }, + ), + mock.patch.object(self.fake_cutlass, "Float32", float32), + mock.patch.object(self.fake_cutlass, "Int32", int32), + ): + self.backward._launch( + "stream", + "q", + "k", + "v", + "out", + "dout", + "lse", + "dq", + "dk", + "dv", + "workspace", + batch=2, + seqlen_q=64, + seqlen_k=96, + num_query_heads=8, + num_kv_heads=2, + element_dtype="bfloat16", + scale_softmax=0.0625, + is_causal=True, + window_size_left=32, + window_size_right=0, + mask_kind="window", + ) + + self.assertEqual(kernel_options[0]["element_dtype"], "bfloat16") + self.assertEqual(kernel_options[0]["mask_type"], "window") + args = calls[0] + self.assertEqual(args[1:9], ("q", "k", "v", "out", "dq", "dk", "dv", "dout")) + self.assertEqual(args[9], "lse") + self.assertEqual(args[-2:], ("workspace", "stream")) + + def test_noncausal_sliding_window_is_rejected(self): + q = _Array((1, 4, 64, 256), self.float16) + k = _Array((1, 1, 64, 256), self.float16) + v = _Array((1, 1, 64, 256), self.float16) + with ( + self._optional_modules(), + self.assertRaisesRegex(NotImplementedError, "non-causal"), + ): + self.forward.sdpa_fwd_wrapper_sm100_d256( + q, + k, + v, + window_size=(32, 0), + ) + + def test_backward_validates_lse_contract_before_loading_kernel(self): + q = _Array((1, 4, 64, 256), self.float16) + k = _Array((1, 1, 64, 256), self.float16) + v = _Array((1, 1, 64, 256), self.float16) + output = _Array(q.shape, self.float16) + doutput = _Array(q.shape, self.float16) + bad_lse = _Array((1, 64, 4), self.float32) + with ( + self._optional_modules(), + self.assertRaisesRegex(ValueError, "lse_tensor must have shape"), + ): + self.backward.sdpa_bwd_wrapper_sm100_d256( + q, + k, + v, + output, + doutput, + bad_lse, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_validation.py b/test/python/fe_api/test_jax_validation.py new file mode 100644 index 000000000..3e6718df2 --- /dev/null +++ b/test/python/fe_api/test_jax_validation.py @@ -0,0 +1,177 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free contracts for JAX API-base validation helpers.""" + +from __future__ import annotations + +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 + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CUDNN_ROOT = _REPO_ROOT / "python" / "cudnn" +_TEST_PACKAGE = "cudnn_frontend_jax_validation_test" + + +class _DType: + def __init__(self, name): + self.name = name + + def __repr__(self): + return self.name + + +class _Float16: + dtype = object() + + +class _Float32: + dtype = object() + + +class _ArrayMetadata: + def __init__(self, dtype, *, name=None): + self.dtype = dtype + if name is not None: + self.name = name + + +class JaxApiBaseValidationTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.float16 = _DType("float16") + cls.float32 = _DType("float32") + dtype_map = { + _Float16: cls.float16, + _Float32: cls.float32, + "float16": cls.float16, + "float32": cls.float32, + } + + fake_jnp = types.ModuleType("jax.numpy") + + def dtype(value): + if isinstance(value, _DType): + return value + try: + return dtype_map[value] + except (KeyError, TypeError): + raise TypeError(f"Cannot interpret {value!r} as a dtype") from None + + fake_jnp.dtype = 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.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, + ) + + def identity_jit(fn=None, **_kwargs): + return (lambda decorated_fn: decorated_fn) if fn is None else fn + + fake_cutlass_jax = types.ModuleType("cutlass.jax") + fake_cutlass_jax.TensorSpec = type("TensorSpec", (), {}) + fake_cute = types.ModuleType("cutlass.cute") + fake_cute.jit = identity_jit + fake_cutlass = types.ModuleType("cutlass") + fake_cutlass.__path__ = [] + fake_cutlass.Constexpr = object + fake_cutlass.cute = fake_cute + fake_cutlass.jax = fake_cutlass_jax + + parent = types.ModuleType(_TEST_PACKAGE) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = _TEST_PACKAGE + sys.modules[_TEST_PACKAGE] = parent + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "cutlass": fake_cutlass, + "cutlass.cute": fake_cute, + "cutlass.jax": fake_cutlass_jax, + }, + ): + cls.api_base = importlib.import_module(f"{_TEST_PACKAGE}._jax.api_base") + + @classmethod + def tearDownClass(cls): + for module_name in tuple(sys.modules): + if module_name == _TEST_PACKAGE or module_name.startswith(f"{_TEST_PACKAGE}."): + sys.modules.pop(module_name, None) + + def test_accepts_dtype_like_values_and_returns_normalized_dtype(self): + self.assertIs( + self.api_base.require_dtype("float32", (_Float16, _Float32)), + self.float32, + ) + + def test_accepts_objects_with_dtype_metadata(self): + value = _ArrayMetadata(self.float16) + self.assertIs( + self.api_base.require_dtype(value, (_Float16, _Float32)), + self.float16, + ) + self.assertIs(self.api_base.as_dtype(value), self.float16) + self.assertIs(self.api_base.as_optional_dtype(value), self.float16) + self.assertIsNone(self.api_base.as_optional_dtype(None)) + + def test_does_not_unwrap_scalar_dtype_classes(self): + self.assertIs( + self.api_base.require_dtype(_Float32, (_Float32,)), + self.float32, + ) + + def test_applies_default_only_to_none(self): + self.assertIs( + self.api_base.require_dtype( + None, + (_Float32,), + name="output_dtype", + default=_Float32, + ), + self.float32, + ) + with self.assertRaisesRegex(ValueError, "output_dtype must not be None"): + self.api_base.require_dtype(None, (_Float32,), name="output_dtype") + + def test_infers_or_accepts_diagnostic_name(self): + named_value = _ArrayMetadata(self.float16, name="sample") + with self.assertRaisesRegex( + ValueError, + r"sample\.dtype must be one of \{float32\}, got float16", + ): + self.api_base.require_dtype(named_value, (_Float32,)) + + with self.assertRaisesRegex(ValueError, r"dtype must be one of \{float32\}, got float16"): + self.api_base.require_dtype(_Float16, (_Float32,)) + + with self.assertRaisesRegex( + ValueError, + r"compute_dtype must be one of \{float32\}, got float16", + ): + self.api_base.require_dtype(_Float16, (_Float32,), name="compute_dtype") + + def test_rejects_invalid_dtype_values(self): + with self.assertRaisesRegex(TypeError, "Cannot interpret"): + self.api_base.require_dtype(object(), (_Float32,)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_optional_jax_test_support.py b/test/python/fe_api/test_optional_jax_test_support.py new file mode 100644 index 000000000..ee9f91b33 --- /dev/null +++ b/test/python/fe_api/test_optional_jax_test_support.py @@ -0,0 +1,146 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Dependency-free tests for optional JAX test selection.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path +import types +import unittest +from unittest import mock + +import pytest + +pytestmark = pytest.mark.L0 + +_CONFTEST_PATH = Path(__file__).with_name("conftest.py") + + +def _load_support_module(): + spec = importlib.util.spec_from_file_location( + "cudnn_frontend_jax_test_support", + _CONFTEST_PATH, + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"Unable to load {_CONFTEST_PATH}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _missing_module(name): + return ModuleNotFoundError(f"No module named {name!r}", name=name) + + +class _Item: + def __init__(self, filename): + self.path = Path(filename) + self.added_markers = [] + + def add_marker(self, marker): + self.added_markers.append(marker) + + +class JaxTestSupportTest(unittest.TestCase): + def setUp(self): + self.support = _load_support_module() + + def test_missing_jax_skips_jax_tests(self): + with mock.patch.object( + self.support, + "import_module", + side_effect=_missing_module("jax"), + ): + reason = self.support._jax_test_skip_reason() + + self.assertIn("require JAX", reason) + self.assertIn("nvidia-cudnn-frontend[jax]", reason) + + def test_missing_cutlass_jax_skips_jax_tests(self): + jax = types.SimpleNamespace(__version__="0.9.1") + + def import_module(name): + if name == "jax": + return jax + raise _missing_module("cutlass.jax") + + with mock.patch.object( + self.support, "import_module", side_effect=import_module + ): + reason = self.support._jax_test_skip_reason() + + self.assertIn("require CUTLASS JAX support", reason) + self.assertIn("nvidia-cudnn-frontend[jax]", reason) + + def test_unsupported_jax_version_skips_jax_tests(self): + jax = types.SimpleNamespace(__version__="0.8.0") + cutlass_jax = types.SimpleNamespace( + CUTE_DSL_MIN_SUPPORTED_JAX_VERSION=(0, 9, 1), + is_available=lambda: False, + ) + + with mock.patch.object( + self.support, + "import_module", + side_effect=lambda name: { + "jax": jax, + "cutlass.jax": cutlass_jax, + }[name], + ): + reason = self.support._jax_test_skip_reason() + + self.assertEqual( + reason, + "CUTLASS JAX support is unavailable with JAX 0.8.0; " + "the minimum supported JAX version is 0.9.1.", + ) + + def test_supported_runtime_does_not_skip(self): + jax = types.SimpleNamespace(__version__="0.9.1") + cutlass_jax = types.SimpleNamespace(is_available=lambda: True) + + with mock.patch.object( + self.support, + "import_module", + side_effect=lambda name: { + "jax": jax, + "cutlass.jax": cutlass_jax, + }[name], + ): + self.assertIsNone(self.support._jax_test_skip_reason()) + + def test_transitive_import_failure_is_not_hidden(self): + error = _missing_module("ml_dtypes") + with mock.patch.object( + self.support, + "import_module", + side_effect=error, + ): + with self.assertRaises(ModuleNotFoundError) as raised: + self.support._jax_test_skip_reason() + + self.assertIs(raised.exception, error) + + def test_collection_hook_selects_jax_test_filenames(self): + runtime_item = _Item("test_jax_gemm.py") + contract_item = _Item("test_jax_gemm_contract.py") + non_jax_item = _Item("test_api_base_framework_split.py") + with mock.patch.object( + self.support, + "_jax_test_skip_reason", + return_value="JAX is unavailable", + ): + self.support.pytest_collection_modifyitems( + None, + [runtime_item, contract_item, non_jax_item], + ) + + self.assertEqual(len(runtime_item.added_markers), 1) + self.assertEqual(len(contract_item.added_markers), 1) + self.assertEqual(non_jax_item.added_markers, []) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_rmsnorm_rht_amax.py b/test/python/fe_api/test_rmsnorm_rht_amax.py index c00d04b02..2e6eb093a 100644 --- a/test/python/fe_api/test_rmsnorm_rht_amax.py +++ b/test/python/fe_api/test_rmsnorm_rht_amax.py @@ -122,3 +122,34 @@ 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_aligned_api(request): + try: + from cudnn import rmsnorm_rht_amax_sm100 + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + skip_ref = request.config.getoption("--skip-ref", default=False) + eps = 1e-5 + m, n = 256, 2048 + rows_per_cta = 2 + x, weight = _make_inputs(m=m, n=n) + + try: + result = rmsnorm_rht_amax_sm100( + x, + weight, + eps=eps, + num_threads=128, + rows_per_cta=rows_per_cta, + ) + except (ValueError, RuntimeError) as exc: + pytest.skip(f"Unsupported testcase: {exc}") + + assert tuple(result.keys()) == ("output", "amax") + assert result["output"].shape == (m, n) + assert result["amax"].shape == (m // rows_per_cta,) + _assert_ref_close(x, weight, result["output"], result["amax"], eps=eps, rows_per_cta=rows_per_cta, skip_ref=skip_ref) diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_config.py b/test/python/fe_api/test_rmsnorm_rht_amax_config.py new file mode 100644 index 000000000..aad4deb67 --- /dev/null +++ b/test/python/fe_api/test_rmsnorm_rht_amax_config.py @@ -0,0 +1,174 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""CPU-only tests for shared RMSNorm + RHT launch configuration.""" + +import importlib.util +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + # Keep this contract test runnable with the standard library alone. + pass +else: + pytestmark = pytest.mark.L0 + + +_MODULE_PATH = Path(__file__).resolve().parents[3] / "python" / "cudnn" / "rmsnorm_rht_amax" / "config.py" +_SPEC = importlib.util.spec_from_file_location("cudnn_rmsnorm_rht_amax_config", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_MODULE = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _MODULE +_SPEC.loader.exec_module(_MODULE) + + +def _desc(shape, dtype_name): + shape = tuple(shape) + return types.SimpleNamespace( + shape=shape, + ndim=len(shape), + dtype_name=dtype_name, + ) + + +class RmsNormRhtAmaxLaunchConfigTest(unittest.TestCase): + def test_package_does_not_eagerly_import_torch_api(self): + package_dir = _MODULE_PATH.parent + root_name = "cudnn_frontend_rmsnorm_lazy_import_test" + module_name = f"{root_name}.rmsnorm_rht_amax" + root = types.ModuleType(root_name) + root.__path__ = [str(package_dir.parent)] + root.__package__ = root_name + sys.modules[root_name] = root + package_spec = importlib.util.spec_from_file_location( + module_name, + package_dir / "__init__.py", + submodule_search_locations=[str(package_dir)], + ) + assert package_spec is not None and package_spec.loader is not None + package = importlib.util.module_from_spec(package_spec) + sys.modules[module_name] = package + try: + package_spec.loader.exec_module(package) + self.assertNotIn(f"{module_name}.api", sys.modules) + finally: + for loaded_name in tuple(sys.modules): + if loaded_name == root_name or loaded_name.startswith(f"{root_name}."): + sys.modules.pop(loaded_name, None) + + def test_resolves_tuned_defaults(self): + self.assertEqual( + _MODULE.resolve_launch_config(256, 2048), + (128, 2), + ) + + def test_preserves_valid_overrides(self): + self.assertEqual( + _MODULE.resolve_launch_config( + 256, + 4096, + num_threads=256, + rows_per_cta=4, + ), + (256, 4), + ) + + def test_validates_tensor_metadata_and_infers_outputs(self): + plan = _MODULE.validate_rmsnorm_rht_amax( + _desc((256, 2048), "bfloat16"), + _desc((2048,), "bfloat16"), + output=_desc((256, 2048), "bfloat16"), + amax=_desc((64,), "float32"), + num_threads=128, + rows_per_cta=4, + ) + + self.assertEqual(plan.m, 256) + self.assertEqual(plan.n, 2048) + self.assertEqual(plan.num_threads, 128) + self.assertEqual(plan.rows_per_cta, 4) + self.assertEqual(plan.output_shape, (256, 2048)) + self.assertEqual(plan.amax_shape, (64,)) + + def test_rejects_invalid_tensor_metadata(self): + valid_x = _desc((256, 2048), "bfloat16") + valid_weight = _desc((2048,), "bfloat16") + cases = ( + ( + (_desc((256, 2048, 1), "bfloat16"), valid_weight), + {}, + "X must have rank 2", + ), + ( + (valid_x, _desc((2048, 1), "bfloat16")), + {}, + "W must have rank 1", + ), + ( + (_desc((256, 2048), "float16"), valid_weight), + {}, + "X must have dtype bfloat16", + ), + ( + (valid_x, _desc((2048,), "float32")), + {}, + "W must have dtype bfloat16", + ), + ( + (valid_x, _desc((1024,), "bfloat16")), + {}, + r"W must have shape \(2048,\)", + ), + ( + (valid_x, valid_weight), + {"output": _desc((256, 1024), "bfloat16")}, + "O must have shape", + ), + ( + (valid_x, valid_weight), + {"output": _desc((256, 2048), "float16")}, + "O must have dtype bfloat16", + ), + ( + (valid_x, valid_weight), + {"amax": _desc((128,), "float32"), "rows_per_cta": 4}, + "Amax must have shape", + ), + ( + (valid_x, valid_weight), + {"amax": _desc((64,), "bfloat16"), "rows_per_cta": 4}, + "Amax must have dtype float32", + ), + ) + + for args, kwargs, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + _MODULE.validate_rmsnorm_rht_amax(*args, **kwargs) + + def test_rejects_invalid_dimensions_and_launch_parameters(self): + cases = ( + ({"m": 0, "n": 2048}, "M must be positive"), + ({"m": 256, "n": 0}, "N must be positive"), + ({"m": 256, "n": 2047}, "Hadamard block size"), + ( + {"m": 256, "n": 2048, "num_threads": 512}, + "EPT=4 must be >= 8 and divisible by 8", + ), + ( + {"m": 255, "n": 2048, "rows_per_cta": 2}, + "M must be divisible", + ), + ) + for kwargs, message in cases: + with self.subTest(kwargs=kwargs): + with self.assertRaisesRegex(ValueError, message): + _MODULE.resolve_launch_config(**kwargs) + + +if __name__ == "__main__": + unittest.main()