diff --git a/docs/fe-oss-apis/dsa.md b/docs/fe-oss-apis/dsa.md index 9ec5f3400..7e543d52a 100644 --- a/docs/fe-oss-apis/dsa.md +++ b/docs/fe-oss-apis/dsa.md @@ -6,10 +6,9 @@ The DeepSeek Sparse Attention (DSA) module integrates a set of CuTe-DSL kernels that support the sparse-attention path used by DeepSeek-style models. -Most kernels target Hopper (SM90) and Blackwell (SM100+) GPUs; Indexer -Forward and Indexer Top-K remain SM100+ only. The kernels are -delivered as Python classes / wrappers that follow the same `APIBase` -pattern as other cuDNN Frontend operations. +The kernels target Hopper (SM90) and Blackwell (SM100+) GPUs and are exposed +through PyTorch and JAX classes and convenience wrappers. See the JAX support +matrix below for operation-specific layout restrictions. **Scope:** this module ships CuTe-DSL kernels for DSA backward, indexer scores/top-K, sparse/dense score recompute, and sparse/dense indexer @@ -24,7 +23,7 @@ The module packages the following operations: 2. **Indexer Forward** – CuTe-DSL score kernel (Q @ K^T, ReLU, head reduce, ratio causal mask). Non-fused; pair with **Indexer Top-K** for the top-K step. -3. **Indexer Top-K** – SM100 CuTe-DSL radix top-K kernel with per-row +3. **Indexer Top-K** – CuTe-DSL radix top-K kernel with per-row ``seq_lens``. 4. **Sparse Indexer / Attention Score Recompute** – sparse (top-K) recompute of indexer and attention scores for training loss. @@ -59,15 +58,23 @@ Training-score loss path: ## Installation +PyTorch: + +```bash +pip install 'nvidia-cudnn-frontend[cutedsl]' +``` + +JAX: + ```bash -pip install nvidia-cudnn-frontend[cutedsl] +pip install 'nvidia-cudnn-frontend[jax]' ``` --- ## API Usage -### DSA Namespace +### PyTorch namespace ```python from cudnn import DSA @@ -100,6 +107,113 @@ DSA.DenseIndexerBackward DSA.dense_indexer_backward_wrapper ``` +### JAX namespace and wrappers + +Importing `cudnn.jax` validates the optional JAX and CUTLASS JAX +dependencies. DSA operation classes and wrappers are then loaded lazily. Each +public symbol is available both directly and through `cudnn.jax.DSA`: + +```python +import cudnn.jax as cudnn_jax + +cudnn_jax.IndexerForward +cudnn_jax.indexer_forward_wrapper + +cudnn_jax.DSA.IndexerForward +cudnn_jax.DSA.indexer_forward_wrapper +``` + +The function wrappers are already decorated with `jax.jit`; their tuning and +target arguments, including maximum sequence lengths, are static compilation +options. They infer output metadata and return a `TupleDict` of JAX arrays: + +```python +result = cudnn_jax.indexer_forward_wrapper(q, k, w, ratio=4) +scores = result["scores"] +``` + +Layout strings describe the public JAX axis order and are static compilation +arguments. `IndexerForward` accepts batch-major or sequence-major fixed +inputs while keeping the head dimension contiguous: + +```python +result = cudnn_jax.indexer_forward_wrapper( + q_sbhd, + k_sbhd, + w_sbh, + q_layout="SBHD", + k_layout="SBHD", + w_layout="SBH", + output_layout="SBK", + ratio=4, +) +scores_sbk = result["scores"] +``` + +The defaults are `q_layout="BSHD"`, `k_layout="BSHD"`, +`w_layout="BSH"`, and `output_layout="BSK"`. Backward and score-recompute +adapters use the same convention and additionally expose layout selectors for +their score, denominator, and gradient tensors. Packed inputs use `THD` and +`TH`, with a `TK` output. The adapters map these public axes to the kernel's +canonical axes; callers do not need to transpose arrays into canonical order. + +Use the class API when the caller needs to choose the enclosing JAX +transformation. Classes specialize from array-like exemplars that expose +`shape` and `dtype`; `jax.ShapeDtypeStruct` makes that intent explicit: + +```python +import jax + +q_spec = jax.ShapeDtypeStruct(q.shape, q.dtype) +k_spec = jax.ShapeDtypeStruct(k.shape, k.dtype) +w_spec = jax.ShapeDtypeStruct(w.shape, w.dtype) + +op = cudnn_jax.IndexerForward(q_spec, k_spec, w_spec, ratio=4) +scores = jax.jit(op)(q, k, w)["scores"] +``` + +The JAX adapters allocate outputs and private workspaces through XLA. Where a +class accepts optional `sample_*` output arguments, those exemplars validate +an explicit output signature; otherwise the adapter infers it. Backward APIs +are functional: they do not update caller-owned gradient buffers in place. +For `IndexerForward`, an explicit `sample_out` describes the physical FP32 +buffer whose last extent is rounded up to four; the returned `scores` view is +sliced back to the logical K extent. +`SparseAttentionBackward` returns `dq`, `dkv`, and `d_sink`; +`IndexerBackward` and `DenseIndexerBackward` return `d_index_q`, `d_weights`, +and `d_index_k`. + +### JAX target selection + +`target_compute_capability` is the exact compilation target, such as `90`, +`100`, `103`, or `107`. When it is omitted, the adapter infers it from a +homogeneous set of local JAX GPUs. An explicit target must match every local +GPU. Device-free AOT and remote compilation are not currently supported. + +### JAX support matrix + +The SM100+ column covers the implemented SM100-family targets: SM100, SM103, +and SM107. + +| Operation | JAX layout | SM90 | SM100+ | +| --- | --- | --- | --- | +| Sparse Attention Backward | Flat MQA tensors | Yes | Yes; inputs must be BF16 | +| Indexer Forward | Fixed BSHD/SBHD or packed THD | Yes; `H_kv=1` and default tuning | Yes | +| Indexer Top-K and index utilities | Flattened rows; fixed or packed index conversion | Yes | Yes | +| Sparse score recompute | Fixed BSHD/SBHD outer layouts, MQA | Yes | Yes | +| Dense score recompute | Fixed BSHD/SBHD outer layouts, MQA | Yes | Yes | +| Dense score recompute | Packed THD/MQA | No | Yes | +| Indexer Backward | Fixed BSHD/SBHD; JAX-only packed THD with global indices | Yes | Yes | +| Dense Indexer Backward | Fixed BSHD/SBHD or packed THD | Yes | Yes | + +SM90 packed THD dense score recompute is rejected because that backend needs +host-side cumulative-length reads, which cannot be represented during JAX +tracing. Sparse score recompute has no packed THD signature and its SM90 +kernel requires a top-K extent divisible by 128. Both indexer +backward variants require `H >= 64`, `D == 128`, and `block_I == 128`; sparse +Indexer Backward also requires its top-K extent to be divisible by 128. Its +packed THD signature requires `topk_indices_global=True` on JAX. + --- ## Components @@ -118,7 +232,12 @@ Backward pass for DeepSeek Sparse Attention. Expects the forward outputs - `topk_idxs`: `(total_S_q, topk_max)` INT32 (global) - `topk_length` (optional): `(total_S_q,)` INT32 — per-query valid count - **Outputs** — tuple `(dq, dkv, d_sink)` -- **Constraints** — SM90 or SM100; SM90 supports the FlashMLA DSA shape with `head_dim ∈ {512, 576}` +- **Constraints** — SM90 or SM100+, `head_dim ∈ {512, 576}`. SM100+ + currently requires BF16 inputs; SM90 also supports FP16. Every + `topk_length` value is interpreted on device and clamped to + `[0, topk_max]`; zero-length rows are supported. When `topk_length` is + supplied, every index before it must be in `[0, total_S_kv)`. When it is + omitted, negative entries are padding sentinels. ```python result = DSA.sparse_attention_backward_wrapper( @@ -145,13 +264,17 @@ The K columns are assumed to be a compressed-KV prefix starting at global compressed column 0. - **Inputs** - - `q`: `(B, S_q, H_q, D)` BF16 - - `k`: `(B, S_k, H_kv, D)` BF16 - - `w`: `(B, S_q, H_q)` BF16 - - `q_causal_offsets` (optional): CUDA INT32 tensor with one entry per + - Fixed: `q` `(B, S_q, H_q, D)`, `k` `(B, S_k, H_kv, D)`, and + `w` `(B, S_q, H_q)`, all BF16. + - Packed THD: `q` `(T_q, H_q, D)`, `k` `(T_k, H_kv, D)`, and + `w` `(T_q, H_q)`, with `cu_seqlens_q`, `cu_seqlens_k`, + `max_seqlen_q`, and `max_seqlen_k`. + - `q_causal_offsets` (optional): INT32 device array with one entry per batch/THD segment, on the same device as `q`. -- **Output** — `scores`: `(B, S_q, S_k)` FP32 -- **Constraints** — SM100+, `head_dim == 128`, `qhead_per_kv_head ∈ {32, 64}` +- **Output** — `scores`: `(B, S_q, S_k)` or `(T_q, max_seqlen_k)` FP32. +- **Constraints** — SM90 or SM100+, `head_dim == 128`, + `qhead_per_kv_head ∈ {32, 64}`. SM90 requires `H_kv == 1` and the + default tuning configuration. ```python result = DSA.indexer_forward_wrapper( @@ -170,7 +293,8 @@ with variable per-row effective length. - `seq_lens`: `(batch_size,)` INT32 (per-batch effective column count) - **Outputs** — tuple `(indices, values)` (values is `None` when `return_val=False`) -- **Constraints** — SM100+, `top_k ≤ 2048` +- **Constraints** — SM90 or SM100+, `top_k ≤ min(2048, num_cols)` and + `top_k + next_n - 1 ≤ seq_lens[b] ≤ num_cols` for every batch item. ```python result = DSA.indexer_top_k_wrapper( @@ -180,6 +304,10 @@ result = DSA.indexer_top_k_wrapper( indices, values = result["indices"], result["values"] ``` +JAX also exposes `local_to_global_wrapper` for fixed or packed local index +conversion and `compactify_wrapper` for packing nonnegative indices and +returning each row's valid `topk_length`. + ### 4. Sparse Indexer Score Recompute Computes softmax over top-K entries of the indexer score: @@ -207,21 +335,34 @@ L1-normalised head-summed softmax over top-K entries: Full-KV (no top-K) analogues of §4 and §5. Each returns `{'out', 'denom'}`. They apply the same ratio-causal mask as Indexer Forward; masked positions are -written as zero and excluded from `denom`. Pass the same `q_causal_offsets` to -all dense score tensors that feed the same loss path. +excluded from `denom`. JAX initializes skipped output positions to `-inf`. +Pass the same `q_causal_offsets` to all dense score tensors that feed the same +loss path. ### 7. Indexer Backward Three-stage sparse top-K pipeline that produces the training gradients for the indexer tower: -1. `ScoreGradSm90` / `ScoreGradSm100` (kernel 1) — in-place score-grad precompute from - `attn_score` (target) and `index_score` (predict). +1. `ScoreGradSm90` / `ScoreGradSm100` (kernel 1) — score-grad precompute from + `attn_score` (target) and `index_score` (predict) into a private temporary. 2. `IndexerBackwardSm90` / `IndexerBackwardSm100` (kernel 2) — three warp-specialised GEMMs produce `d_index_q`, `d_weights`, and a `dIndexK_f32` accumulator. -3. Pure-torch dtype cast (kernel 3) converts `dIndexK_f32` to the output - dtype. +3. A dtype cast converts `dIndexK_f32` to the output dtype. + +The JAX wrapper treats `grad_loss` as a runtime array operand and returns new +gradient arrays; no input or caller-owned output buffer is mutated. + +The common fixed signature uses `index_q` `(B, S_q, H, 128)`, `weights` +`(B, S_q, H)`, `index_k` `(B, S_k, 128)`, and score/index tensors +`(B, S_q, topk)`. JAX additionally accepts a packed signature with +`index_q` `(T_q, H, 128)`, `weights` `(T_q, H)`, `index_k` `(T_k, 128)`, and +score/index tensors `(T_q, topk)`. The packed adapter presents these tensors +to the kernel as a synthetic batch of one, so it requires +`topk_indices_global=True`; each top-K entry must address the flattened packed +K range directly. It does not infer original batch-local IDs from cumulative +sequence lengths. The PyTorch sparse wrapper remains fixed BSHD. **The TileLang fallback present in the upstream repo is dropped here (CuTe-DSL only).** If the CuTe-DSL path fails the wrapper raises @@ -278,11 +419,13 @@ result = DSA.dense_indexer_backward_wrapper( ## Limitations - **Architecture support** — Sparse Attention Backward, Score Recompute, and - Indexer Backward support SM90 and SM100; Indexer Forward and Indexer Top-K - remain SM100+ only. + Indexer operations support SM90 and SM100+ subject to the layout restrictions + in the JAX support matrix. - **No fused forward** — the production forward is FlashMLA (C++); this module ships only the CuTe-DSL kernels. - **Indexer Forward only supports `head_dim = 128`** and `qhead_per_kv_head ∈ {32, 64}`. - **Top-K only up to 2048**; `top_k > 2048` is not supported by the underlying radix top-K kernel. +- **JAX Top-K does not use the PyTorch row-chunking fallback** when its private + workspace would exceed the INT32 indexing range. diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_amax.md b/docs/fe-oss-apis/gemm_fusions/gemm_amax.md index 12f14f11e..2022d9037 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_amax.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_amax.md @@ -103,6 +103,66 @@ op.compile() op.execute(a, b, sfa, sfb, c, amax, current_stream=None) ``` +### JAX API + +The JAX wrapper is decorated with `jax.jit`. GEMM arrays are compact row-major +arrays whose public axis order is selected explicitly: + +```python +import jax +import jax.numpy as jnp + +from cudnn.jax import gemm_amax_wrapper_sm100 + +# A: (L, M, K), B: (L, N, K) +# SFA/SFB: (L, ceil(rows/128), K_tiles, 32, 4, 4) +result = gemm_amax_wrapper_sm100( + a, + b, + sfa, + sfb, + a_layout="LMK", + b_layout="LNK", + c_layout="LMN", + c_dtype=jnp.float32, + acc_dtype=jnp.float32, + sf_vec_size=32, +) + +c = result["c_tensor"] # (L, M, N) +amax = result["amax_tensor"] # (1, 1, 1) +``` + +The scale arrays use ordinary row-major public storage with shape +`(L, ceil(rows/128), ceil(ceil(K/sf_vec_size)/4), 32, 4, 4)`. The adapter +maps that public representation to the kernel's canonical packed descriptor +`(32, 4, row_tiles, 4, k_tiles, L)`. + +The advanced class specializes from shape/dtype metadata and optionally accepts +explicit output exemplars: + +```python +from cudnn.jax import GemmAmaxSm100 + +operation = GemmAmaxSm100( + jax.ShapeDtypeStruct(a.shape, a.dtype), + jax.ShapeDtypeStruct(b.shape, b.dtype), + jax.ShapeDtypeStruct(sfa.shape, sfa.dtype), + jax.ShapeDtypeStruct(sfb.shape, sfb.dtype), +) +result = operation(a, b, sfa, sfb) +``` + +Supported public layouts are: + +- A: `LMK` (K-major) or `LKM` (M-major) +- B: `LNK` (K-major) or `LKN` (N-major) +- C: `LMN` (N-major) or `LNM` (M-major) + +JAX uses native FP4/FP8 dtypes. The Torch-only `uint8` packed-FP4 and `int8` +E8M0 storage aliases are not accepted by the JAX adapter. The JAX API requires +a homogeneous local GPU with an exact supported target in the SM100 family. + --- ## Parameters @@ -131,7 +191,7 @@ op.execute(a, b, sfa, sfb, c, amax, current_stream=None) - Dtype: `float32` ### Common parameters -- `acc_dtype: torch.dtype` +- `acc_dtype`: a Torch dtype or JAX dtype-like value - Accumulator dtype. Default: `torch.float32` (only supported value) - `mma_tiler_mn: Tuple[int, int]` - Kernel tile size `(TILE_M, TILE_N)`. Default: `(128, 128)` @@ -143,12 +203,17 @@ op.execute(a, b, sfa, sfb, c, amax, current_stream=None) - `sf_vec_size: int` - Size of K-group per scale factor: `{16, 32}`. Default: `32` - CUDA stream (`current_stream` in class API, `stream` in wrapper) + JAX receives its stream from lowering and does not expose a stream argument. ### Wrapper-specific parameters: `gemm_amax_wrapper_sm100` - `a_tensor`, `b_tensor`, `sfa_tensor`, `sfb_tensor`: see Input/Output tensors - `c_major: str`: see Input/Output tensors. Default: `"n"` - `c_dtype: torch.dtype`: see Input/Output tensors. Default: `torch.float32` +For JAX, `a_layout`, `b_layout`, and `c_layout` select the public axis orders +listed above. `c_dtype` defaults to `jnp.float32`; output allocation and amax +initialization are owned by the JAX wrapper. + ### Wrapper return values Returns a `TupleDict` with keys: diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md b/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md index a18656744..915506cdf 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_dsrelu.md @@ -7,7 +7,7 @@ **Block-scaled GEMM + dsReLU backward fusion**: A persistent, batched dense GEMM on NVIDIA Blackwell GPUs (SM100+) that supports block-scaled FP4 and FP8 inputs and produces both the backward output `D` and the probability gradient `dprob` in a single kernel launch. - **Inputs**: quantized `A` and `B`, the forward/intermediate tensor `C`, scale-factor tensors `SFA` and `SFB`, and a per-row probability tensor `prob` -- **Outputs**: backward output `D`, probability gradient `dprob`, and optional output scale factors `SFD` / `Amax` +- **Outputs**: backward output `D`, probability gradient `dprob`, and optional `Amax` ### Shapes @@ -22,7 +22,6 @@ - **Outputs** - `D`: shape `(M, N, L)` - `dprob`: shape `(M, 1, L)` - - `SFD`: shape `(32, 4, ceil_div(M, 128), 4, ceil_div(ceil_div(N, sf_vec_size), 4), L)` when `D` is FP8 - `Amax`: shape `(1,)` when FP4 input is written to fp16/bf16/fp32 output `L` is the batch dimension. @@ -43,7 +42,9 @@ $$ \mathrm{dprob}[m, 0, l] = \sum_n C[m, n, l] \cdot \mathrm{relu}(G[m, n, l])^2 $$ -As with the forward `srelu` kernel, FP8 `D` also emits `SFD` using the provided `norm_const_tensor`, and FP4 input written to fp16/bf16/fp32 `D` can emit `Amax`. +As with the forward `srelu` kernel, FP4 input written to fp16/bf16/fp32 `D` +can emit `Amax`. The FP8 `D`/SFD epilogue is reserved by the API but is not +implemented by the current kernel. ### Diagram @@ -63,16 +64,18 @@ A (MxKxL), SFA B (NxKxL), SFB | +--> dprob (Mx1xL) | - +-----------+-----------+ - | | - v v - SFD Amax + v + Amax ``` --- ## API Usage +The existing APIs in this section use Torch tensors. The JAX API is described +separately below because its public arrays are row-major and name their logical +axis order explicitly. + ### High-level wrapper ```python @@ -142,6 +145,47 @@ op.execute( ) ``` +### JAX API + +```python +import jax.numpy as jnp +from cudnn.jax import gemm_dsrelu_wrapper_sm100 + +result = gemm_dsrelu_wrapper_sm100( + a, # (L, M, K), float4_e2m1fn or FP8 + b, # (L, N, K), same dtype as A + c, # (L, M, N) + sfa, # (L, ceil(M/128), rest_k, 32, 4, 4) + sfb, # (L, ceil(N/128), rest_k, 32, 4, 4) + prob, # (L, 1, M), float32 + d_dtype=jnp.bfloat16, + sf_vec_size=16, # use 32 for MXFP8 + E8M0 scales +) +``` + +`GemmDsreluSm100` provides the corresponding class API and accepts optional +`sample_d` and `sample_dprob` output exemplars. The high-level wrapper is JIT +compiled and infers its outputs. + +JAX arrays use compact row-major storage: + +- `a_layout="LMK"` and `"LKM"` map public A shapes `(L,M,K)` and `(L,K,M)` + to kernel axes `(M,K,L)`. +- `b_layout="LNK"` and `"LKN"` map public B shapes `(L,N,K)` and `(L,K,N)` + to kernel axes `(N,K,L)`. +- `d_layout="LMN"` consumes C and produces D as `(L,M,N)`; `"LNM"` uses + `(L,N,M)`. +- SFA/SFB use public shape `(L, tiles, rest, 32, 4, 4)`. The adapter maps + this row-major representation to the packed six-dimensional kernel ABI. +- `prob` and the inferred `dprob_tensor` use public shape `(L,1,M)`. + +Here `rest_k = ceil(ceil(K / sf_vec_size) / 4)`. + +Native JAX `float4_e2m1fn` inputs return an initialized `amax_tensor` with +shape `(1,)`. JAX currently rejects FP8 `D`, `norm_const_tensor`, and SFD +generation because the kernel's SFD epilogue is not implemented. `D` must use +`float16`, `bfloat16`, or `float32`. + --- ## Parameters @@ -169,14 +213,14 @@ op.execute( - Dtype: `float32` - Output tensor **D**: `result["d_tensor"]` (wrapper) or `sample_d` / `d_tensor` (class) - Shape: `(M, N, L)` - - Dtype: `{float16, bfloat16, float32, float8_e4m3fn, float8_e5m2}` + - Functional dtypes: `{float16, bfloat16, float32}` - Output tensor **dprob**: `result["dprob_tensor"]` (wrapper) or `sample_dprob` / `dprob_tensor` (class) - Shape: `(M, 1, L)` - Dtype: `float32` - Output tensor **SFD**: `result["sfd_tensor"]` (wrapper) or `sample_sfd` / `sfd_tensor` (class) - Shape: `(32, 4, ceil_div(M, 128), 4, ceil_div(ceil_div(N, sf_vec_size), 4), L)` - Dtype: Must match `SFA` - - Required when `D` is FP8 + - Reserved for FP8 `D`; generation is not implemented - Output tensor **Amax**: `result["amax_tensor"]` (wrapper) or `sample_amax` / `amax_tensor` (class) - Shape: `(1,)` - Dtype: `float32` @@ -184,7 +228,7 @@ op.execute( - Input tensor **Norm Const**: `norm_const_tensor` (wrapper) or `sample_norm_const` / `norm_const_tensor` (class) - Shape: `(1,)` - Dtype: `float32` - - Required when `D` is FP8 + - Reserved for FP8 `D`; currently unsupported ### Common parameters @@ -237,11 +281,13 @@ Tuple unpacking order is: `(d_tensor, dprob_tensor, amax_tensor, sfd_tensor)`. - `sf_vec_size == 32` is unsupported with `sf_dtype == float8_e4m3fn` - FP8 input requires `sf_vec_size == 32` - FP4 input with FP8 `D` is unsupported -- FP8 `D` requires both `SFD` and `norm_const_tensor` +- FP8 `D` and SFD generation are not implemented ### Environment - Requires CUDA with SM100+ compute capability +- The JAX API requires a homogeneous local SM100-family GPU and resolves + occupancy before CuTe lowering. Device-free compilation is not supported. --- @@ -251,3 +297,4 @@ For end-to-end usage and regression coverage, see: - `test/python/fe_api/test_gemm_dsrelu.py` - `test/python/fe_api/test_gemm_dsrelu_utils.py` +- `test/python/fe_api/test_jax_gemm_relu.py` diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md b/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md index 963165e3f..242f10117 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_srelu.md @@ -7,7 +7,7 @@ **Block-scaled GEMM + sReLU fusion**: A persistent, batched dense GEMM on NVIDIA Blackwell GPUs (SM100+) that supports block-scaled FP4 and FP8 inputs and produces both the full GEMM result `C` and a probability-gated squared-ReLU output `D` in a single kernel launch. - **Inputs**: quantized `A` and `B`, scale-factor tensors `SFA` and `SFB`, and a per-row probability tensor `prob` -- **Outputs**: full GEMM result `C`, squared-ReLU output `D`, and optional output scale factors `SFD` / `Amax` +- **Outputs**: full GEMM result `C`, squared-ReLU output `D`, and optional `Amax` ### Shapes @@ -21,7 +21,6 @@ - **Outputs** - `C`: shape `(M, N, L)` - `D`: shape `(M, N, L)` - - `SFD`: shape `(32, 4, ceil_div(M, 128), 4, ceil_div(ceil_div(N, sf_vec_size), 4), L)` when `D` is FP8 - `Amax`: shape `(1,)` when FP4 input is written to fp16/bf16/fp32 output `L` is the batch dimension. @@ -38,7 +37,9 @@ $$ D[m, n, l] = \mathrm{prob}[m, 0, l] \cdot \mathrm{relu}(C[m, n, l])^2 $$ -When `D` is FP8, the kernel also emits output scale factors `SFD` using the provided `norm_const_tensor`. When FP4 input is written to a higher-precision `D`, the kernel can also emit `Amax`. +When FP4 input is written to a higher-precision `D`, the kernel can also emit +`Amax`. The FP8 `D`/SFD epilogue is reserved by the API but is not implemented +by the current kernel. ### Diagram @@ -56,16 +57,18 @@ A (MxKxL), SFA B (NxKxL), SFB v D (MxNxL) | - +-----------+-----------+ - | | - v v - SFD Amax + v + Amax ``` --- ## API Usage +The existing APIs in this section use Torch tensors. The JAX API is described +separately below because its public arrays are row-major and name their logical +axis order explicitly. + ### High-level wrapper ```python @@ -133,6 +136,46 @@ op.execute( ) ``` +### JAX API + +```python +import jax.numpy as jnp +from cudnn.jax import gemm_srelu_wrapper_sm100 + +result = gemm_srelu_wrapper_sm100( + a, # (L, M, K), float4_e2m1fn or FP8 + b, # (L, N, K), same dtype as A + sfa, # (L, ceil(M/128), rest_k, 32, 4, 4) + sfb, # (L, ceil(N/128), rest_k, 32, 4, 4) + prob, # (L, 1, M), float32 + c_dtype=jnp.bfloat16, + d_dtype=jnp.bfloat16, + sf_vec_size=16, # use 32 for MXFP8 + E8M0 scales +) +``` + +`GemmSreluSm100` provides the corresponding class API and accepts optional +`sample_c` and `sample_d` output exemplars. The high-level wrapper is JIT +compiled and infers its outputs. + +JAX arrays use compact row-major storage: + +- `a_layout="LMK"` and `"LKM"` map public A shapes `(L,M,K)` and `(L,K,M)` + to kernel axes `(M,K,L)`. +- `b_layout="LNK"` and `"LKN"` map public B shapes `(L,N,K)` and `(L,K,N)` + to kernel axes `(N,K,L)`. +- `c_layout="LMN"` produces `(L,M,N)`; `"LNM"` produces `(L,N,M)`. +- SFA/SFB use public shape `(L, tiles, rest, 32, 4, 4)`. The adapter maps + this row-major representation to the packed six-dimensional kernel ABI. +- `prob` uses public shape `(L,1,M)`. + +Here `rest_k = ceil(ceil(K / sf_vec_size) / 4)`. + +Native JAX `float4_e2m1fn` inputs return an initialized `amax_tensor` with +shape `(1,)`. JAX currently rejects FP8 `D`, `norm_const_tensor`, and SFD +generation because the kernel's SFD epilogue is not implemented. `D` must use +`float16`, `bfloat16`, or `float32`. + --- ## Parameters @@ -160,11 +203,11 @@ op.execute( - Dtype: `{float16, bfloat16, float32, float8_e4m3fn, float8_e5m2}` - Output tensor **D**: `result["d_tensor"]` (wrapper) or `sample_d` / `d_tensor` (class) - Shape: `(M, N, L)` - - Dtype: `{float16, bfloat16, float32, float8_e4m3fn, float8_e5m2}` + - Functional dtypes: `{float16, bfloat16, float32}` - Output tensor **SFD**: `result["sfd_tensor"]` (wrapper) or `sample_sfd` / `sfd_tensor` (class) - Shape: `(32, 4, ceil_div(M, 128), 4, ceil_div(ceil_div(N, sf_vec_size), 4), L)` - Dtype: Must match `SFA` - - Required when `D` is FP8 + - Reserved for FP8 `D`; generation is not implemented - Output tensor **Amax**: `result["amax_tensor"]` (wrapper) or `sample_amax` / `amax_tensor` (class) - Shape: `(1,)` - Dtype: `float32` @@ -172,7 +215,7 @@ op.execute( - Input tensor **Norm Const**: `norm_const_tensor` (wrapper) or `sample_norm_const` / `norm_const_tensor` (class) - Shape: `(1,)` - Dtype: `float32` - - Required when `D` is FP8 + - Reserved for FP8 `D`; currently unsupported ### Common parameters @@ -222,11 +265,13 @@ Tuple unpacking order is: `(c_tensor, d_tensor, amax_tensor, sfd_tensor)`. - `sf_vec_size == 32` is unsupported with `sf_dtype == float8_e4m3fn` - FP8 input requires `sf_vec_size == 32` - FP4 input with FP8 `D` is unsupported -- FP8 `D` requires both `SFD` and `norm_const_tensor` +- FP8 `D` and SFD generation are not implemented ### Environment - Requires CUDA with SM100+ compute capability +- The JAX API requires a homogeneous local SM100-family GPU and resolves + occupancy before CuTe lowering. Device-free compilation is not supported. --- @@ -236,3 +281,4 @@ For end-to-end usage and regression coverage, see: - `test/python/fe_api/test_gemm_srelu.py` - `test/python/fe_api/test_gemm_srelu_utils.py` +- `test/python/fe_api/test_jax_gemm_relu.py` diff --git a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md index e03c041e4..6c5976232 100644 --- a/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md +++ b/docs/fe-oss-apis/gemm_fusions/gemm_swiglu.md @@ -10,6 +10,10 @@ This API supports two modes: 1. **Standard mode**: High-precision GEMM with SwiGLU epilogue 2. **Quantized mode** (block-scaled): Low-precision GEMM using block scaling supporting FP4 and FP8 data types +The Torch and JAX APIs support both modes. JAX scale-factor arrays use an +ordinary row-major public shape and are mapped to the kernel's packed ABI by +the adapter. + ### Shapes - Inputs: @@ -178,12 +182,108 @@ gemm.execute( ) ``` +### JAX API (Standard Mode) + +The JAX wrapper is already decorated with `jax.jit`. Public arrays are compact +row-major arrays whose axis order is described explicitly: + +```python +import jax +import jax.numpy as jnp + +from cudnn.jax import gemm_swiglu_wrapper_sm100 + +# A: (L, M, K), B: (L, N, K) +result = gemm_swiglu_wrapper_sm100( + a, + b, + alpha=1.0, + a_layout="LMK", + b_layout="LNK", + c_layout="LMN", + ab12_dtype=jnp.float32, + c_dtype=jnp.bfloat16, + acc_dtype=jnp.float32, +) + +# AB12: (L, M, N), C: (L, M, N/2) +ab12 = result["ab12_tensor"] +c = result["c_tensor"] +``` + +The advanced class API specializes from shape/dtype metadata and may optionally +accept explicit output exemplars: + +```python +from cudnn.jax import GemmSwigluSm100 + +operation = GemmSwigluSm100( + jax.ShapeDtypeStruct(a.shape, a.dtype), + jax.ShapeDtypeStruct(b.shape, b.dtype), + a_layout="LMK", + b_layout="LNK", + c_layout="LMN", +) +result = operation(a, b) +``` + +### JAX API (Block-Scaled Mode) + +Providing both `sfa_tensor` and `sfb_tensor` selects the block-scaled kernel. +The output dtypes remain explicit because the wrapper defaults are unchanged +from standard mode; MXFP8 callers must request a supported AB12 dtype. + +```python +from cudnn.jax import gemm_swiglu_wrapper_sm100 + +# A: (L, M, K), B: (L, N, K) +# SFA/SFB: (L, row_tiles, k_tiles, 32, 4, 4) +result = gemm_swiglu_wrapper_sm100( + a_fp8, + b_fp8, + sfa_tensor=sfa_e8m0, + sfb_tensor=sfb_e8m0, + sf_vec_size=32, + ab12_dtype=jnp.bfloat16, + c_dtype=jnp.bfloat16, + a_layout="LMK", + b_layout="LNK", + c_layout="LMN", +) +``` + +For a scale tensor associated with `rows` and reduction extent `K`, the public +JAX shape is +`(L, ceil(rows/128), ceil(ceil(K/sf_vec_size)/4), 32, 4, 4)`. +Its fixed mode `(3, 4, 1, 5, 2, 0)` binds that row-major array to the canonical +kernel shape `(32, 4, ceil(rows/128), 4, k_tiles, L)`. + +The current JAX surface supports `C` in `float16`, `bfloat16`, or `float32`. +FP8 `C`, SFC generation, and `norm_const_tensor` are disabled because none of +the currently supported input configurations can use the kernel's FP8-C path. +`result["sfc_tensor"]` is therefore `None`. + +Supported public layouts are: + +- A: `LMK` (K-major) or `LKM` (M-major) +- B: `LNK` (K-major) or `LKN` (N-major) +- AB12/C: `LMN` (N-major) or `LNM` (M-major) + +The JAX API requires a homogeneous, supported local SM100-family GPU. It +resolves the exact local compute capability for CuTe compilation and queries +that GPU's occupancy once per specialized adapter before CuTe lowering. +Device-free AOT and remote compilation are not currently supported. + --- ## Parameters ### Input/Output tensors +The shapes and strides below use the kernel's canonical axis order, which is +also the Torch API order. JAX arrays use the public axis orders selected by +`a_layout`, `b_layout`, and `c_layout` above. + - Input tensor **A**: `a_tensor` (wrapper) or `sample_a`, `a_tensor` (class) - Shape: `(M, K, L)` - Stride: `(1, M, M·K)` for `m`-major or `(K, 1, M·K)` for `k`-major @@ -202,7 +302,7 @@ gemm.execute( - Quantized mode: Must be `n`-major for FP4 outputs - Dtype (`ab12_dtype`, provided as `ab12_dtype` argument for wrapper): - Standard mode: `{float32, float16, bfloat16}` if `acc_dtype == float32`, `{float16, bfloat16}` if `acc_dtype == float16` - - Quantized mode: `{float32, float16, bfloat16, float8_e4m3fn, float8_e5m2}` + - Quantized mode: `{float32, float16, bfloat16}` - Output tensor **C**: `result["c_tensor"]` (wrapper) or `sample_c`, `c_tensor` (class) - Shape: `(M, N/2, L)` - Stride: `(1, M, M·N/2)` for `m`-major or `(N/2, 1, M·N/2)` for `n`-major. Must match with `AB12` @@ -216,10 +316,12 @@ gemm.execute( - Input tensor **SFB** (B scale factor): `sfb_tensor` (wrapper) or `sample_sfb`, `sfb_tensor` (class) - Shape: `(32, 4, ceil(N/128), 4, ceil(ceil(K/sf_vec_size)/4), L)` - Dtype: Must match `SFA` + - JAX public shape for SFA/SFB is + `(L, row_tiles, k_tiles, 32, 4, 4)` as described above. - Output tensor **SFC** (C scale factor, **Optional**): `result["sfc_tensor"]` (wrapper) or `sample_sfc`, `sfc_tensor` (class) - Shape: `(32, 4, ceil(M/128), 4, ceil(ceil((N/2)/sf_vec_size)/4), L)` - Dtype: Must match `SFA` - - **Required when**: `c_dtype ∈ {float8_e4m3fn, float8_e5m2}` + - Reserved for FP8 `C`; current supported input combinations reject FP8 `C` - Output tensor **AMAX** (**Optional**): `result["amax_tensor"]` (wrapper) or `sample_amax`, `amax_tensor` (class) - Shape: `(1,)` - Dtype: `float32` @@ -227,27 +329,27 @@ gemm.execute( - Input tensor **Norm Const** (**Optional**): `norm_const_tensor` (wrapper) or `sample_norm_const`, `norm_const_tensor` (class) - Shape: `(1,)` - Dtype: `float32` - - **Required when**: `c_dtype ∈ {float8_e4m3fn, float8_e5m2}` + - Reserved for FP8 `C`; currently unsupported ### Common parameters - `alpha: float` - Scalar multiplier applied to the GEMM result before SwiGLU. - Default: `1.0` -- `acc_dtype: torch.dtype` +- `acc_dtype`: `torch.dtype` for Torch or a JAX dtype-like value for JAX - Accumulator dtype. - - Standard mode: `{float32, float16}`. Default: `torch.float32` + - Standard mode: `{float32, float16}`. Default: `float32` - Quantized mode: Must be `float32` - `mma_tiler_mn: Tuple[int, int]` - Kernel tile size `(TILE_M, TILE_N)`. Default: `(128, 128)` - `TILE_M ∈ {128, 256}` - - Standard mode: `TILE_N ∈ {32, 64, ..., 224, 256}` + - Standard mode: `TILE_N ∈ {64, 128, 192, 256}` - Quantized mode: `TILE_N ∈ {64, 128, 192, 256}` - `cluster_shape_mn: Tuple[int, int] | None` - Thread Block cluster shape `(CLUSTER_M, CLUSTER_N)` - Constraints: positive powers of 2, `CLUSTER_M*CLUSTER_N ≤ 16`. - Default: `(1,1)` if `mma_tiler_mn[0] != 256` else `(2,2)`. -- CUDA stream (`current_stream` in class API, `stream` in wrapper) +- CUDA stream (`current_stream` in the Torch class API, `stream` in the Torch wrapper). JAX supplies its stream during lowering. - **Quantization-specific parameters** - `sf_vec_size: int` - Scale factor vector size (number of elements per scale factor) @@ -262,14 +364,20 @@ gemm.execute( - Number of pipeline stages for AB12 output - Default: `4` -### Wrapper-specific parameters: `gemm_swiglu_wrapper_sm100` +### Wrapper-specific parameters + +Torch and JAX both expose `gemm_swiglu_wrapper_sm100`. The JAX wrapper resolves +the exact local SM100-family compilation target while tracing. - `a_tensor`, `b_tensor`: see Input/Output tensors -- `c_major: str`: see Input/Output tensors. Default: `"n"` -- `ab12_dtype: torch.dtype`: see Input/Output tensors. Default: `torch.float32` -- `c_dtype: torch.dtype`: see Input/Output tensors. Default: `torch.float16` -- `sfa_tensor`, `sfb_tensor`, `norm_const_tensor`: see Quantization-specific tensors -- `sf_vec_size`, `vector_f32`, `ab12_stages`: see Quantization-specific parameters +- Torch: `c_major` selects `"m"` or `"n"`; `ab12_dtype` defaults to `torch.float32`; `c_dtype` defaults to `torch.float16`. +- JAX: `a_layout`, `b_layout`, and `c_layout` select public axis order; `ab12_dtype` defaults to `jnp.float32`; `c_dtype` defaults to `jnp.float16`. +- Torch quantized mode: `sfa_tensor`, `sfb_tensor`, `norm_const_tensor`, + `sf_vec_size`, `vector_f32`, and `ab12_stages` use the quantization rules + above. +- JAX block-scaled mode: `sfa_tensor`, `sfb_tensor`, `sf_vec_size`, + `vector_f32`, and `ab12_stages` are supported. `norm_const_tensor` is + rejected while FP8 `C` remains disabled. ### Wrapper return values @@ -290,8 +398,11 @@ Tuple unpacking order is always: #### `GemmSwigluSm100` (constructor) -- `sample_a`, `sample_b`, `sample_ab12`, `sample_c` – see Input/Output tensors -- `sample_sfa`, `sample_sfb`, `sample_sfc`, `sample_amax`, `sample_norm_const` – see Scale factor tensors (quantized mode) +- `sample_a`, `sample_b`, `sample_ab12`, `sample_c` – see Input/Output tensors. JAX requires only the input samples and accepts the output samples optionally. +- `sample_sfa`, `sample_sfb`, `sample_sfc`, `sample_amax`, + `sample_norm_const` – quantized-mode samples. JAX supports SFA/SFB and can + infer AMAX; it rejects SFC and norm-constant samples while FP8 `C` remains + disabled. #### `GemmSwigluSm100.execute` @@ -305,7 +416,7 @@ Tuple unpacking order is always: ### Layouts and strides - `AB12` and `C` must have the same major order. -- `A`, `B`, `AB12` must be 16-byte aligned along the contiguous dimension. +- `A`, `B`, `AB12`, and `C` must be 16-byte aligned along the contiguous dimension. - For FP4 inputs (quantized mode): `A` and `B` must be `k`-major, `AB12` must be `n`-major. ### Dtypes @@ -327,11 +438,14 @@ The quantized kernel supports the following configurations: | **MXFP4** | `float4_e2m1fn_x2` or `uint8` | `float8_e4m3fn` | 16 | NVF4 variant | | **MXFP8** | `float8_e4m3fn` or `float8_e5m2` | `float8_e8m0fnu` | 32 | Standard MX FP8 | +JAX uses its native logical `float4_e2m1fn` dtype for FP4 when that dtype is +available in the installed JAX runtime. The raw `uint8` packed-FP4 alias is a +Torch compatibility path and is not reinterpreted by the JAX adapter. + Additional constraints: - `acc_dtype` must be `float32` - Not compatible with FP8 c_dtype. BF16 `c_dtype` is expected. -- For MXFP8 inputs, ab12_dtype` should be float16 or bfloat16. -- When `c_dtype ∈ {float8_e4m3fn, float8_e5m2}`: `sfc_tensor` and `norm_const_tensor` are required +- For MXFP8 inputs, `ab12_dtype` should be float16 or bfloat16. - When `ab_dtype` is FP4 and `c_dtype == bfloat16`: `amax_tensor` is required - `c_dtype` and `ab12_dtype` cannot both be `float32` diff --git a/docs/fe-oss-apis/nsa.md b/docs/fe-oss-apis/nsa.md index 10a2458a6..559c52318 100644 --- a/docs/fe-oss-apis/nsa.md +++ b/docs/fe-oss-apis/nsa.md @@ -74,6 +74,46 @@ NSA.TopKReduction NSA.topk_reduction_wrapper ``` +### JAX namespace + +Install the optional integration with `nvidia-cudnn-frontend[jax]`, then use +the lazy `cudnn.jax.NSA` namespace (the symbols are also available directly +from `cudnn.jax`): + +```python +import cudnn.jax as cudnn_jax + +result = cudnn_jax.NSA.compression_attention_wrapper( + q, + k, + v, + layout="BSHD", + enable_lse=True, +) +output = result["o_tensor"] +lse = result["lse_tensor"] +``` + +Functional JAX wrappers are decorated with `jax.jit`. Their layout, tuning, +maximum-sequence-length, and target arguments are static; cumulative sequence +lengths remain dynamic array operands. Classes accept array-like exemplars and +can be placed under a caller-owned JAX transformation. + +| Component | JAX public layouts | +| --- | --- | +| Selection | Packed `THD` | +| Compression | Fixed `BHSD`/`BSHD` or packed `THD` | +| Sliding Window | Fixed `BHSD`/`BSHD` or fully packed `THD` | +| Top-K | Fixed `BHSD`/`BSHD` or packed `THD` | + +Packed calls require cumulative Q/K sequence lengths and static `max_s_q` and +`max_s_k` bounds. The arrays must start at zero, be monotonic, end at the +corresponding packed token count, and respect those bounds. Sliding Window +uses those offsets to form padded cuDNN attention operands in the lowered +computation; it does not accept arbitrary ragged byte-offset tensors. Fixed +data outputs follow the selected layout, while LSE/statistics remain in +`(B, H, S[, 1])` order. + --- ## Components diff --git a/docs/fe-oss-apis/rmsnorm_rht_amax.md b/docs/fe-oss-apis/rmsnorm_rht_amax.md index ba875fef8..69493cd5c 100644 --- a/docs/fe-oss-apis/rmsnorm_rht_amax.md +++ b/docs/fe-oss-apis/rmsnorm_rht_amax.md @@ -9,6 +9,8 @@ This frontend integration exposes the kernel as a standard FE-OSS Python API with: - a class API (`RmsNormRhtAmaxSm100`) - a wrapper API (`rmsnorm_rht_amax_wrapper_sm100`) +- an explicit JAX API under `cudnn.jax` +- a framework-neutral operation signature in `op.py` and a CUTE DSL implementation in `kernel.py` - grouped-gemm-style regression coverage for compile/execute, wrapper use, and cache reuse ## Shapes @@ -97,9 +99,60 @@ op.execute( ) ``` +The PyTorch class preserves caller-owned output buffers. During +`check_support()`, the framework-neutral operation validates the complete +input and output exemplar signature and resolves the launch configuration. +`compile()` then constructs the CUTE DSL kernel from that resolved +configuration. + +### JAX API + +The CUDA 13 JAX packages require Python 3.11 or newer. Install the JAX integration explicitly: + +```bash +pip install 'nvidia-cudnn-frontend[jax]' +``` + +```python +import jax +from cudnn.jax import RmsNormRhtAmaxSm100, rmsnorm_rht_amax_sm100 + +o, amax = rmsnorm_rht_amax_sm100( + x, + w, + eps=1e-5, + num_threads=128, + rows_per_cta=2, +) + +# The class API specializes from abstract placeholders, then accepts arrays +# when called under JAX transformations. +x_spec = jax.ShapeDtypeStruct(x.shape, x.dtype) +w_spec = jax.ShapeDtypeStruct(w.shape, w.dtype) +op = RmsNormRhtAmaxSm100(x_spec, w_spec, rows_per_cta=2) +o, amax = jax.jit(op)(x, w) +``` + +The function API is JIT-compiled and treats `eps`, `num_threads`, and +`rows_per_cta` as static compilation options. Use `RmsNormRhtAmaxSm100` +directly when the caller needs to control JAX transformations explicitly. A +JAX array or another object exposing shape and dtype may be used instead of a +`ShapeDtypeStruct` exemplar; only its metadata is retained. Unlike PyTorch, +JAX derives output descriptors in its adapter and lets XLA allocate the +results. Callers may instead provide both `sample_o` and `sample_amax` as +explicit output placeholders. In either case, the common operation validates +the complete signature through `check_support()`. + +JAX support validation requires every local GPU to meet the operation's SM100 +minimum. Device-free AOT and remote compilation are not currently supported. + +The current JAX adapter models its buffers as row-major with the public and +kernel axes in the same order. Additional JAX-specific axis mappings can be +added without changing the framework-neutral operation signature. + ## Parameters -### Input and output tensors +### PyTorch input and output tensors - `x_tensor` / `sample_x` - Shape: `(M, N)` @@ -117,6 +170,9 @@ op.execute( - Shape: `(M / rows_per_cta,)` - Dtype: `torch.float32` +The JAX API uses the same shapes and layouts with `jax.numpy.bfloat16` for +`X`, `W`, and `O`, and `jax.numpy.float32` for `Amax`. + ### Common parameters - `eps: float` @@ -125,16 +181,28 @@ 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`) -### Wrapper return values +The PyTorch `execute()` and wrapper APIs additionally accept the optional +CUDA stream argument `current_stream`. JAX execution uses the stream supplied +by XLA. + +### Return values + +The PyTorch wrapper returns a `TupleDict` with keys: -Returns a `TupleDict` with keys: - `o_tensor` + - Shape: `(M, N)` + - Layout: row-major contiguous + - Dtype: `torch.bfloat16` - `amax_tensor` + - Shape: `(M / rows_per_cta,)` + - Dtype: `torch.float32` Tuple unpacking order is `(o_tensor, amax_tensor)`. +The class API writes into its caller-provided output tensors and returns +`None`. + ## Support surface and constraints - Requires SM100+. @@ -142,10 +210,13 @@ Tuple unpacking order is `(o_tensor, amax_tensor)`. - `N` must be divisible by the resolved `num_threads`. - `EPT = N / num_threads` must be at least `8` and divisible by `8`. - `M` must be divisible by `rows_per_cta`. -- Inputs and output are currently bf16 only. +- `X`, `W`, and `O` use bf16; `amax` uses float32. +- PyTorch input and output tensors must be 16-byte aligned. - The frontend integration matches the upstream RMSNorm kernel semantics; it does not expose full LayerNorm mean/bias behavior. ## Verification Focused correctness and cache coverage live in: + - `test/python/fe_api/test_rmsnorm_rht_amax.py` +- `test/python/fe_api/test_jax_rmsnorm_rht_amax.py` diff --git a/pyproject.toml b/pyproject.toml index ca8243c38..2c5a14d30 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -76,3 +76,8 @@ cutedsl = [ "apache-tvm-ffi", "torch-c-dlpack-ext", ] +jax = [ + "nvidia-cutlass-dsl[cu13]==4.5.0", + "cuda-python", + "jax[cuda13]>=0.9.1", +] diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index af48c763f..7465c8d76 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -51,6 +51,8 @@ def is_windows(): if hasattr(_pybind_module, _optional_symbol): globals()[_optional_symbol] = getattr(_pybind_module, _optional_symbol) +from .common.op import Op +from .common.tensor_desc import TensorDesc from .datatypes import _library_type, _is_torch_tensor __version__ = "1.26.0" @@ -334,6 +336,11 @@ def __getattr__(name: str) -> Any: globals()["ops"] = _ops return _ops + if name == "jax": + _jax = importlib.import_module(".jax", __name__) + globals()["jax"] = _jax + return _jax + if name == "experimental": from . import experimental as _experimental diff --git a/python/cudnn/_jax/__init__.py b/python/cudnn/_jax/__init__.py new file mode 100644 index 000000000..25f9cc8a7 --- /dev/null +++ b/python/cudnn/_jax/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Internal implementation helpers for optional JAX operation adapters.""" + +from .api_base import JaxApiBase, JaxTensorDesc +from .result import TupleDict + +__all__ = ["JaxApiBase", "JaxTensorDesc", "TupleDict"] diff --git a/python/cudnn/_jax/api_base.py b/python/cudnn/_jax/api_base.py new file mode 100644 index 000000000..e708f9a6d --- /dev/null +++ b/python/cudnn/_jax/api_base.py @@ -0,0 +1,685 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Base class, tensor metadata, and CuTe binding for optional JAX adapters.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable +from dataclasses import dataclass, field, replace +from operator import index +from typing import Any + +from .. import data_type +from ..common.tensor_desc import TensorDesc +from .layout import compact_stride, normalize_mode, to_canonical_axes, to_cutlass_layout, to_public_axes + + +@dataclass(frozen=True) +class JaxTensorDesc(TensorDesc[Any]): + """JAX tensor metadata and CUTLASS lowering constraints.""" + + mode: tuple[int, ...] = field(default_factory=tuple) + divisibility: tuple[int | None, ...] | None = None + ptr_assumed_align: int | None = None + + def __post_init__(self) -> None: + super().__post_init__() + object.__setattr__( + self, + "mode", + normalize_mode(self.ndim, None if not self.mode else self.mode), + ) + object.__setattr__( + self, + "divisibility", + _normalize_divisibility(self.ndim, self.divisibility), + ) + object.__setattr__( + self, + "ptr_assumed_align", + _normalize_ptr_assumed_align(self.ptr_assumed_align), + ) + + @property + def array_shape(self) -> tuple[int, ...]: + """Return the shape in the axis order used by the JAX array.""" + + return to_public_axes(self.shape, self.mode) + + @classmethod + def from_array( + cls, + value: Any, + *, + name: str = "", + mode: tuple[int, ...] | None = None, + public_stride_order: tuple[int, ...] | None = None, + init_value: bool | int | float | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, + ) -> "JaxTensorDesc": + """Describe a JAX array-like value using its shape and dtype metadata.""" + + shape = _require_array_metadata(value, name or "value") + return cls.from_shape( + shape, + value.dtype, + name=name, + mode=mode, + public_stride_order=public_stride_order, + init_value=init_value, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + ) + + @classmethod + def from_shape( + cls, + shape: tuple[int, ...], + dtype: Any, + *, + name: str = "", + mode: tuple[int, ...] | None = None, + public_stride_order: tuple[int, ...] | None = None, + init_value: bool | int | float | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, + ) -> "JaxTensorDesc": + """Describe a compact JAX tensor from its user-facing shape and dtype.""" + + shape = tuple(shape) + rank = len(shape) + mode = normalize_mode(rank, mode) + public_stride_order = _normalize_stride_order(rank, public_stride_order) + public_stride = compact_stride(shape, public_stride_order) + canonical_axis_by_public_axis = to_public_axes(tuple(range(rank)), mode) + + return cls( + dtype=dtype, + shape=to_canonical_axes(shape, mode), + stride=to_canonical_axes(public_stride, mode), + stride_order=tuple( + canonical_axis_by_public_axis[axis] + for axis in public_stride_order + ), + name=name, + init_value=init_value, + mode=mode, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + ) + + @property + def cudnn_dtype(self) -> data_type: + from .datatypes import jax_to_cudnn_dtype + + return jax_to_cudnn_dtype(self.dtype) + + def compact_like( + self, + *, + cudnn_dtype: data_type, + shape: tuple[int, ...], + stride_order: tuple[int, ...] | None = None, + name: str = "", + init_value: bool | int | float | None = None, + mode: tuple[int, ...] | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, + ) -> "JaxTensorDesc": + """Create a compact JAX descriptor from canonical output metadata.""" + + canonical = super().compact_like( + cudnn_dtype=cudnn_dtype, + shape=shape, + stride_order=stride_order, + name=name, + init_value=init_value, + ) + + from .datatypes import cudnn_to_jax_dtype + + return JaxTensorDesc( + dtype=cudnn_to_jax_dtype(cudnn_dtype), + shape=canonical.shape, + stride=canonical.stride, + stride_order=canonical.stride_order, + name=canonical.name, + init_value=canonical.init_value, + mode=normalize_mode(canonical.ndim, mode), + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + ) + + def with_divisibility( + self, + divisibility: tuple[int | None, ...] | None, + ) -> "JaxTensorDesc": + """Return this descriptor with canonical-axis divisibility constraints.""" + + return replace(self, divisibility=divisibility) + + +def _require_array_metadata(value: Any, name: str) -> tuple[Any, ...]: + if not hasattr(value, "shape") or not hasattr(value, "dtype"): + raise TypeError(f"{name} must have shape and dtype metadata") + return tuple(value.shape) + + +def _normalize_stride_order(rank: int, stride_order: tuple[int, ...] | None) -> tuple[int, ...]: + if stride_order is None: + return tuple(reversed(range(rank))) + + normalized = [] + for dimension in stride_order: + if isinstance(dimension, bool): + raise TypeError(f"public_stride_order entries must be integers, got {dimension!r}") + try: + normalized.append(index(dimension)) + except TypeError as error: + raise TypeError(f"public_stride_order entries must be integers, got {dimension!r}") from error + + normalized_order = tuple(normalized) + if len(normalized_order) != rank: + raise ValueError(f"public_stride_order rank mismatch: expected {rank}, got {len(normalized_order)}") + if tuple(sorted(normalized_order)) != tuple(range(rank)): + raise ValueError(f"public_stride_order must be a permutation of [0, {rank - 1}], got {normalized_order}") + return normalized_order + + +def _normalize_divisibility( + rank: int, + divisibility: tuple[int | None, ...] | None, +) -> tuple[int | None, ...] | None: + if divisibility is None: + return None + divisibility = tuple(divisibility) + if len(divisibility) != rank: + raise ValueError( + f"divisibility rank mismatch: expected {rank}, got {len(divisibility)}" + ) + + normalized = [] + for value in divisibility: + if value is None: + normalized.append(None) + continue + if isinstance(value, bool): + raise TypeError( + f"divisibility entries must be positive integers or None, got {value!r}" + ) + try: + value = index(value) + except TypeError as error: + raise TypeError( + f"divisibility entries must be positive integers or None, got {value!r}" + ) from error + if value <= 0: + raise ValueError( + f"divisibility entries must be positive integers or None, got {value}" + ) + normalized.append(value) + return tuple(normalized) + + +def _normalize_ptr_assumed_align(value: int | None) -> int | None: + if value is None: + return None + if isinstance(value, bool): + raise TypeError( + f"ptr_assumed_align must be a positive integer or None, got {value!r}" + ) + try: + value = index(value) + except TypeError as error: + raise TypeError( + f"ptr_assumed_align must be a positive integer or None, got {value!r}" + ) from error + if value <= 0: + raise ValueError( + f"ptr_assumed_align must be a positive integer or None, got {value}" + ) + return value + + +class JaxApiBase(ABC): + """Common tensor metadata, validation, and CuTe binding for JAX adapters.""" + + @staticmethod + def _device_compute_capability(device: Any, operation_name: str) -> int: + """Normalize JAX's device capability metadata to ``major * 10 + minor``.""" + + reported = getattr(device, "compute_capability", None) + try: + if isinstance(reported, (tuple, list)): + if len(reported) < 2: + raise ValueError + major, minor = int(reported[0]), int(reported[1]) + else: + text = str(reported) + if "." in text: + major_text, minor_text = text.split(".", 1) + major, minor = int(major_text), int(minor_text) + else: + capability = int(text) + major, minor = divmod(capability, 10) + if major < 0 or minor < 0 or minor > 9: + raise ValueError + except (TypeError, ValueError) as error: + raise RuntimeError(f"{operation_name}: JAX reported an invalid compute capability {reported!r} for {device}") from error + return major * 10 + minor + + @staticmethod + def _local_gpu_capabilities(operation_name: str) -> tuple[tuple[Any, int], ...]: + import jax + + try: + devices = tuple(jax.local_devices(backend="gpu")) + except RuntimeError as error: + raise RuntimeError(f"{operation_name}: JAX could not discover a local GPU") from error + return tuple((device, JaxApiBase._device_compute_capability(device, operation_name)) for device in devices) + + @staticmethod + def _compute_capability_family( + compute_capability: int, + supported_compute_capabilities: tuple[int, ...], + ) -> int | None: + compatible = tuple(capability for capability in supported_compute_capabilities if capability <= compute_capability) + return max(compatible, default=None) + + @staticmethod + def _resolve_compute_capability( + target_compute_capability: int | None, + supported_compute_capabilities: tuple[int, ...], + operation_name: str, + ) -> int: + """Resolve an exact JAX compilation target for a multi-arch adapter. + + ``supported_compute_capabilities`` lists the exact CuTe compilation + targets implemented by the adapter. The returned value remains exact + (for example ``103``) so the compiler selects the matching target. + An explicit target must match every local GPU exactly. + """ + + supported = tuple(sorted(set(supported_compute_capabilities))) + if not supported or any(isinstance(value, bool) or not isinstance(value, int) or value <= 0 for value in supported): + raise ValueError(f"supported_compute_capabilities must contain positive integers, got {supported_compute_capabilities!r}") + + if target_compute_capability is not None: + if isinstance(target_compute_capability, bool) or not isinstance(target_compute_capability, int): + raise TypeError(f"target_compute_capability must be an int or None, got {type(target_compute_capability).__name__}") + if target_compute_capability <= 0: + raise ValueError(f"target_compute_capability must be positive, got {target_compute_capability}") + if target_compute_capability not in supported: + supported_text = ", ".join(f"SM{value}" for value in supported) + raise ValueError(f"{operation_name} has no kernel for SM{target_compute_capability}; supported targets are {supported_text}") + local = JaxApiBase._local_gpu_capabilities(operation_name) + if not local: + raise RuntimeError(f"{operation_name} targets SM{target_compute_capability}, but no local JAX GPU is available") + mismatched = tuple((device, capability) for device, capability in local if capability != target_compute_capability) + if mismatched: + found = ", ".join(f"{device} (SM{capability})" for device, capability in mismatched) + raise RuntimeError(f"{operation_name} targets SM{target_compute_capability}, but found {found}") + return target_compute_capability + + local = JaxApiBase._local_gpu_capabilities(operation_name) + if not local: + raise RuntimeError(f"{operation_name}: no local JAX GPU is available") + + capabilities = tuple(sorted({capability for _, capability in local})) + if len(capabilities) != 1: + found = ", ".join(f"SM{capability}" for capability in capabilities) + raise RuntimeError(f"{operation_name}: local JAX GPUs have heterogeneous targets ({found})") + + resolved = capabilities[0] + if resolved not in supported: + supported_text = ", ".join(f"SM{value}" for value in supported) + raise RuntimeError(f"{operation_name} has no kernel for SM{resolved}; supported targets are {supported_text}") + return resolved + + @staticmethod + def _check_device_compatibility( + *, + minimum_compute_capability: int, + operation_name: str, + ) -> None: + """Require every local JAX GPU to satisfy an operation's minimum SM.""" + + def compatibility_error(reason: str) -> RuntimeError: + return RuntimeError(f"{operation_name} requires SM{minimum_compute_capability}+, {reason}") + + try: + devices = JaxApiBase._local_gpu_capabilities(operation_name) + except RuntimeError as error: + reason = str(error) + prefix = f"{operation_name}: " + if reason.startswith(prefix): + reason = reason[len(prefix) :] + raise compatibility_error(f"but {reason}") from error + + if not devices: + raise compatibility_error("but no local JAX GPU is available") + + incompatible = [] + for device, compute_capability in devices: + if compute_capability < minimum_compute_capability: + incompatible.append((device, compute_capability)) + + if incompatible: + found = ", ".join(f"{device} (SM{compute_capability})" for device, compute_capability in incompatible) + raise compatibility_error(f"found {found}") + + def _get_max_active_clusters( + self, + cluster_size: int, + *, + overlap_margin: int = 0, + ) -> int: + """Return cached occupancy for this adapter and cluster size. + + The hardware query performs its own CuTe compilation, so adapters must + call this before entering ``cutlass.jax.cutlass_call`` lowering. + """ + + cache = getattr(self, "_max_active_clusters_cache", None) + if cache is None: + cache = {} + self._max_active_clusters_cache = cache + + if cluster_size not in cache: + import cutlass + + cache[cluster_size] = cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_size) + + max_active_clusters = cache[cluster_size] - overlap_margin + if max_active_clusters <= 0: + raise ValueError("max_active_clusters must be positive after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") + return max_active_clusters + + def _get_device_multiprocessor_count(self) -> int: + """Return a cached local SM count before entering custom-call lowering.""" + + count = getattr(self, "_device_multiprocessor_count", None) + if count is None: + import cutlass + + count = int(cutlass.utils.HardwareInfo().get_device_multiprocessor_count()) + if count <= 0: + raise RuntimeError(f"CUDA reported an invalid multiprocessor count {count}") + self._device_multiprocessor_count = count + return count + + @staticmethod + def _to_tensor_desc( + value: Any, + name: str, + *, + mode: tuple[int, ...] | None = None, + public_stride_order: tuple[int, ...] | None = None, + init_value: bool | int | float | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, + ) -> JaxTensorDesc: + """Describe a public JAX array in canonical kernel-axis order. + + ``mode[kernel_axis]`` selects the corresponding public array axis. + ``public_stride_order`` lists the public array dimensions from fastest + to slowest. It defaults to compact row-major storage. + """ + + return JaxTensorDesc.from_array( + value, + name=name, + mode=mode, + public_stride_order=public_stride_order, + init_value=init_value, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + ) + + @staticmethod + def _check_tensor_signature( + value: Any, + expected: JaxTensorDesc, + ) -> None: + """Validate a public JAX value against a canonical descriptor.""" + + if not isinstance(expected, JaxTensorDesc): + raise TypeError( + f"expected must be a JaxTensorDesc, got {type(expected).__name__}" + ) + name = expected.name or "value" + public_shape = _require_array_metadata(value, name) + if len(public_shape) != expected.ndim: + raise ValueError(f"{name} tensor shape mismatch: expected {expected.shape}, got public shape {public_shape}") + actual_shape = to_canonical_axes(public_shape, expected.mode) + if actual_shape != expected.shape: + raise ValueError(f"{name} tensor shape mismatch: expected {expected.shape}, got {actual_shape}") + + from .datatypes import jax_to_cudnn_dtype + + actual_dtype = jax_to_cudnn_dtype(value.dtype) + if actual_dtype != expected.cudnn_dtype: + raise ValueError(f"{name} tensor dtype mismatch: expected {expected.cudnn_dtype}, got {actual_dtype}") + + @staticmethod + def _to_tensor_spec( + desc: JaxTensorDesc, + ) -> Any: + """Build a CUTLASS TensorSpec indexed by public JAX array axes. + + ``desc.divisibility`` uses canonical kernel axes. TensorSpec + layout and divisibility use public axes, while TensorSpec ``mode`` + records the canonical-to-public binding. + """ + + if not isinstance(desc, JaxTensorDesc): + raise TypeError( + f"desc must be a JaxTensorDesc, got {type(desc).__name__}" + ) + public_layout = to_cutlass_layout( + desc.shape, + desc.stride, + desc.stride_order, + mode=desc.mode, + name=desc.name or "tensor", + ) + default_mode = tuple(range(desc.ndim)) + default_layout = tuple(reversed(range(desc.ndim))) + if ( + desc.mode == default_mode + and public_layout == default_layout + and desc.divisibility is None + and desc.ptr_assumed_align is None + ): + return None + + if desc.divisibility is None: + public_divisibility = None + else: + public_divisibility = to_public_axes( + desc.divisibility, + desc.mode, + ) + + from cutlass.jax import TensorSpec + + options: dict[str, Any] = dict( + layout=public_layout, + mode=desc.mode, + ) + if public_divisibility is not None: + options["divisibility"] = public_divisibility + if desc.ptr_assumed_align is not None: + options["ptr_assumed_align"] = desc.ptr_assumed_align + return TensorSpec(**options) + + @staticmethod + def _to_shape_dtype_struct( + desc: JaxTensorDesc, + ) -> Any: + """Convert a canonical descriptor to abstract JAX output metadata.""" + + if not isinstance(desc, JaxTensorDesc): + raise TypeError(f"desc must be a JaxTensorDesc, got {type(desc).__name__}") + + from jax import ShapeDtypeStruct + + from .datatypes import cudnn_to_jax_dtype + + return ShapeDtypeStruct( + desc.array_shape, + cudnn_to_jax_dtype(desc.cudnn_dtype), + ) + + def _call_kernel( + self, + inputs: tuple[Any, ...], + *, + launch: Callable[..., None], + input_descs: tuple[JaxTensorDesc, ...], + output_descs: tuple[JaxTensorDesc, ...], + workspace_descs: tuple[JaxTensorDesc, ...] = (), + allow_cuda_graph: bool = True, + compile_options: Any = None, + use_static_tensors: bool = True, + ) -> tuple[Any, ...]: + """Bind an explicit kernel launcher to JAX and return public outputs. + + CUTLASS JAX supplies the stream first; the launcher receives + ``stream, *inputs, *outputs, *workspaces`` in that order. Workspaces + are declared as custom-call results so XLA owns their lifetime, then + omitted from this method's return value. Descriptors with a non-``None`` + ``init_value`` are materialized as JAX inputs and aliased to their + corresponding custom-call results. Outputs without an initial value + are allocated as custom-call results by XLA. + Descriptors validate input signatures and provide all CUTLASS + ``TensorSpec`` metadata. + """ + + import cutlass.cute as cute + import cutlass.jax as cutlass_jax + import jax.numpy as jnp + + if not callable(launch): + raise TypeError(f"launch must be callable, got {type(launch).__name__}") + + inputs = tuple(inputs) + for input_index, value in enumerate(inputs): + _require_array_metadata(value, f"input #{input_index}") + outputs = tuple(output_descs) + workspaces = tuple(workspace_descs) + if not outputs: + raise ValueError("A JAX operation must provide at least one output descriptor") + + input_descs = tuple(input_descs) + if len(input_descs) != len(inputs): + raise ValueError( + f"Expected {len(inputs)} input descriptors, got {len(input_descs)}" + ) + + for label, descs in ( + ("input", input_descs), + ("output", outputs), + ("workspace", workspaces), + ): + for desc in descs: + if not isinstance(desc, JaxTensorDesc): + raise TypeError( + f"{label}_descs must contain JaxTensorDesc values, " + f"got {type(desc).__name__}" + ) + + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc) + + input_specs = tuple(self._to_tensor_spec(desc) for desc in input_descs) + output_specs = tuple(self._to_tensor_spec(desc) for desc in outputs) + workspace_specs = tuple(self._to_tensor_spec(desc) for desc in workspaces) + buffers = outputs + workspaces + buffer_specs = output_specs + workspace_specs + buffer_metadata = tuple( + self._to_shape_dtype_struct(desc) for desc in buffers + ) + + initialized_buffer_by_index: dict[int, Any] = {} + for buffer_index, desc in enumerate(buffers): + if desc.init_value is not None: + initialized_buffer_by_index[buffer_index] = jnp.full( + buffer_metadata[buffer_index].shape, + desc.init_value, + dtype=buffer_metadata[buffer_index].dtype, + ) + + initialized = tuple(sorted(initialized_buffer_by_index)) + uninitialized = tuple( + buffer_index + for buffer_index in range(len(buffers)) + if buffer_index not in initialized_buffer_by_index + ) + initialized_inputs = tuple( + initialized_buffer_by_index[buffer_index] + for buffer_index in initialized + ) + aliases = { + len(inputs) + initialized_input_index: result_index + for initialized_input_index, result_index in enumerate(initialized) + } + call_input_specs = input_specs + tuple( + buffer_specs[buffer_index] for buffer_index in initialized + ) + + input_count = len(inputs) + initialized_count = len(initialized_inputs) + + @cute.jit(preprocess=False) + def launcher(stream, *args): + kernel_inputs = args[:input_count] + initialized_buffers = args[ + input_count : input_count + initialized_count + ] + allocated_buffers = args[input_count + initialized_count :] + if len(allocated_buffers) != len(uninitialized): + raise RuntimeError(f"Kernel received {len(allocated_buffers)} allocated buffers; expected {len(uninitialized)}") + + ordered_buffers: list[Any] = [None] * len(buffers) + for buffer_index, value in zip(initialized, initialized_buffers): + ordered_buffers[buffer_index] = value + for buffer_index, value in zip(uninitialized, allocated_buffers): + ordered_buffers[buffer_index] = value + output_buffers = tuple(ordered_buffers[: len(outputs)]) + workspace_buffers = tuple(ordered_buffers[len(outputs) :]) + launch(stream, *kernel_inputs, *output_buffers, *workspace_buffers) + + call = cutlass_jax.cutlass_call( + launcher, + output_shape_dtype=buffer_metadata, + input_spec=call_input_specs, + output_spec=buffer_specs, + input_output_aliases=aliases, + allow_cuda_graph=allow_cuda_graph, + compile_options=compile_options, + use_static_tensors=use_static_tensors, + ) + results = call(*inputs, *initialized_inputs) + if not isinstance(results, (tuple, list)): + results = (results,) + return tuple(results[: len(outputs)]) + + @abstractmethod + def check_support(self) -> bool: + """Validate the operation signature and static configuration.""" + + @abstractmethod + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """Trace or execute the operation with JAX array arguments.""" + + def get_jax_callable(self) -> Callable[..., Any]: + """Return this callable without imposing a JIT policy.""" + + return self + + +__all__ = ["JaxApiBase", "JaxTensorDesc"] diff --git a/python/cudnn/_jax/compiler.py b/python/cudnn/_jax/compiler.py new file mode 100644 index 000000000..eda074de9 --- /dev/null +++ b/python/cudnn/_jax/compiler.py @@ -0,0 +1,20 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""CUTLASS JAX compilation options.""" + +from __future__ import annotations + +from ..common.cute_arch import gpu_arch_flag_for_compute_capability + + +def compile_options_for_target(compute_capability: int, extra: str = "") -> str: + """Build CUTLASS JAX compile options for an explicit target.""" + + parts = [f"--gpu-arch {gpu_arch_flag_for_compute_capability(compute_capability)}"] + if extra: + parts.append(extra) + return " ".join(parts) + + +__all__ = ["compile_options_for_target"] diff --git a/python/cudnn/_jax/datatypes.py b/python/cudnn/_jax/datatypes.py new file mode 100644 index 000000000..005a23b19 --- /dev/null +++ b/python/cudnn/_jax/datatypes.py @@ -0,0 +1,73 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Conversions between JAX and canonical cuDNN data types.""" + +from __future__ import annotations + +from typing import Any + +import jax.numpy as jnp + +from .. import data_type + +_DTYPE_NAMES = ( + ("bool_", "BOOLEAN"), + ("float16", "HALF"), + ("bfloat16", "BFLOAT16"), + ("float32", "FLOAT"), + ("float64", "DOUBLE"), + ("int8", "INT8"), + ("int32", "INT32"), + ("int64", "INT64"), + ("uint8", "UINT8"), + ("float8_e4m3fn", "FP8_E4M3"), + ("float8_e5m2", "FP8_E5M2"), + ("float8_e8m0fnu", "FP8_E8M0"), + ("float4_e2m1fn", "FP4_E2M1"), + ("int4", "INT4"), +) + + +def _make_mappings(): + jax_to_cudnn = {} + cudnn_to_jax = {} + for jax_name, cudnn_name in _DTYPE_NAMES: + jax_scalar_type = getattr(jnp, jax_name, None) + cudnn_dtype = getattr(data_type, cudnn_name, None) + if jax_scalar_type is None or cudnn_dtype is None: + continue + jax_dtype = jnp.dtype(jax_scalar_type) + jax_to_cudnn[jax_dtype] = cudnn_dtype + cudnn_to_jax[cudnn_dtype] = jax_dtype + return jax_to_cudnn, cudnn_to_jax + + +_JAX_TO_CUDNN_DATA_TYPE, _CUDNN_TO_JAX_DATA_TYPE = _make_mappings() + + +def normalize_jax_dtype(value: Any | None, default: Any, name: str) -> Any: + """Normalize a JAX dtype-like value and identify invalid arguments by name.""" + + try: + return jnp.dtype(default if value is None else value) + except TypeError as error: + raise TypeError(f"{name} must be a JAX dtype, got {value!r}") from error + + +def jax_to_cudnn_dtype(dtype: Any) -> data_type: + """Return the canonical cuDNN type for a JAX dtype-like value.""" + + return _JAX_TO_CUDNN_DATA_TYPE.get(jnp.dtype(dtype), data_type.NOT_SET) + + +def cudnn_to_jax_dtype(dtype: data_type) -> Any: + """Return the JAX dtype corresponding to a canonical cuDNN type.""" + + try: + return _CUDNN_TO_JAX_DATA_TYPE[dtype] + except KeyError as error: + raise ValueError(f"Unsupported JAX data type {dtype}") from error + + +__all__ = ["cudnn_to_jax_dtype", "jax_to_cudnn_dtype", "normalize_jax_dtype"] diff --git a/python/cudnn/_jax/gemm.py b/python/cudnn/_jax/gemm.py new file mode 100644 index 000000000..f0f359420 --- /dev/null +++ b/python/cudnn/_jax/gemm.py @@ -0,0 +1,68 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Public-axis bindings for dense JAX GEMMs.""" + +from __future__ import annotations + +from .layout import mode_from_layout + +GEMM_A_LAYOUTS = ("LMK", "LKM") +GEMM_B_LAYOUTS = ("LNK", "LKN") +GEMM_OUTPUT_LAYOUTS = ("LMN", "LNM") + +ROW_MAJOR_STRIDE_ORDER_3D = (2, 1, 0) +# Public row-major shape: (L, tiles_M_or_N, rest_K_or_N, 32, 4, 4). +BLOCK_SCALE_MODE = (3, 4, 1, 5, 2, 0) +BLOCK_SCALE_STRIDE_ORDER = (3, 1, 0, 4, 2, 5) +# Public row-major shape: (L, 1, M). +PROBABILITY_MODE = (2, 1, 0) +PROBABILITY_STRIDE_ORDER = (0, 1, 2) + + +def require_layout(name: str, layout: str, supported: tuple[str, ...]) -> str: + """Validate an explicit, case-sensitive public axis-order string.""" + + if not isinstance(layout, str): + raise TypeError(f"{name} must be a string, got {type(layout).__name__}") + if layout not in supported: + choices = ", ".join(repr(value) for value in supported) + raise ValueError(f"{name} must be one of ({choices}), got {layout!r}") + return layout + + +def gemm_a_mode(layout: str) -> tuple[int, ...]: + """Map public A layout ``LMK`` or ``LKM`` to canonical ``MKL`` axes.""" + + layout = require_layout("a_layout", layout, GEMM_A_LAYOUTS) + return mode_from_layout(layout, kernel_axes="MKL") + + +def gemm_b_mode(layout: str) -> tuple[int, ...]: + """Map public B layout ``LNK`` or ``LKN`` to canonical ``NKL`` axes.""" + + layout = require_layout("b_layout", layout, GEMM_B_LAYOUTS) + return mode_from_layout(layout, kernel_axes="NKL") + + +def gemm_output_mode(layout: str, *, name: str = "c_layout") -> tuple[int, ...]: + """Map public output layout ``LMN`` or ``LNM`` to canonical ``MNL`` axes.""" + + layout = require_layout(name, layout, GEMM_OUTPUT_LAYOUTS) + return mode_from_layout(layout, kernel_axes="MNL") + + +__all__ = [ + "BLOCK_SCALE_MODE", + "BLOCK_SCALE_STRIDE_ORDER", + "GEMM_A_LAYOUTS", + "GEMM_B_LAYOUTS", + "GEMM_OUTPUT_LAYOUTS", + "PROBABILITY_MODE", + "PROBABILITY_STRIDE_ORDER", + "ROW_MAJOR_STRIDE_ORDER_3D", + "gemm_a_mode", + "gemm_b_mode", + "gemm_output_mode", + "require_layout", +] diff --git a/python/cudnn/_jax/layout.py b/python/cudnn/_jax/layout.py new file mode 100644 index 000000000..9cdc93b13 --- /dev/null +++ b/python/cudnn/_jax/layout.py @@ -0,0 +1,167 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Axis and layout conversion for CUTLASS JAX calls.""" + +from __future__ import annotations + +from operator import index +from typing import TypeVar + +AxisValueT = TypeVar("AxisValueT") + + +def mode_from_layout(layout: str, *, kernel_axes: str) -> tuple[int, ...]: + """Map named kernel axes to their positions in a public JAX layout. + + Axis names are single, case-sensitive characters. The returned mapping + follows ``mode[kernel_axis] = public_axis``. + """ + + if not isinstance(layout, str): + raise TypeError(f"layout must be a string, got {layout!r}") + if not isinstance(kernel_axes, str): + raise TypeError(f"kernel_axes must be a string, got {kernel_axes!r}") + if len(layout) != len(kernel_axes): + raise ValueError( + f"layout rank must match kernel_axes: got {layout!r} and {kernel_axes!r}" + ) + if len(set(layout)) != len(layout): + raise ValueError(f"layout axes must be unique, got {layout!r}") + if len(set(kernel_axes)) != len(kernel_axes): + raise ValueError(f"kernel_axes must be unique, got {kernel_axes!r}") + if set(layout) != set(kernel_axes): + raise ValueError( + f"layout must contain exactly the axes in {kernel_axes!r}, got {layout!r}" + ) + return tuple(layout.index(axis) for axis in kernel_axes) + + +def normalize_mode(rank: int, mode: tuple[int, ...] | None = None) -> tuple[int, ...]: + """Return and validate a kernel-axis to public-axis mapping. + + ``mode[kernel_axis]`` is the corresponding axis in the public JAX array. + Omitting ``mode`` keeps the public and kernel axis orders identical. + """ + + if isinstance(rank, bool): + raise TypeError(f"rank must be an integer, got {rank!r}") + try: + rank = index(rank) + except TypeError as error: + raise TypeError(f"rank must be an integer, got {rank!r}") from error + if rank < 0: + raise ValueError(f"rank must be non-negative, got {rank}") + + if mode is None: + return tuple(range(rank)) + + normalized = [] + for axis in mode: + if isinstance(axis, bool): + raise TypeError(f"mode entries must be integers, got {axis!r}") + try: + normalized.append(index(axis)) + except TypeError as error: + raise TypeError(f"mode entries must be integers, got {axis!r}") from error + normalized_mode = tuple(normalized) + if tuple(sorted(normalized_mode)) != tuple(range(rank)): + raise ValueError( + f"mode must be a permutation of [0, {rank - 1}], got {normalized_mode}" + ) + return normalized_mode + + +def to_canonical_axes( + public_values: tuple[AxisValueT, ...], + mode: tuple[int, ...] | None = None, +) -> tuple[AxisValueT, ...]: + """Reorder public-axis values into canonical kernel-axis order.""" + + public_values = tuple(public_values) + mode = normalize_mode(len(public_values), mode) + return tuple(public_values[public_axis] for public_axis in mode) + + +def to_public_axes( + canonical_values: tuple[AxisValueT, ...], + mode: tuple[int, ...] | None = None, +) -> tuple[AxisValueT, ...]: + """Reorder canonical kernel-axis values into public JAX axis order.""" + + canonical_values = tuple(canonical_values) + mode = normalize_mode(len(canonical_values), mode) + canonical_axis_by_public_axis = [0] * len(mode) + for canonical_axis, public_axis in enumerate(mode): + canonical_axis_by_public_axis[public_axis] = canonical_axis + return tuple( + canonical_values[canonical_axis_by_public_axis[public_axis]] + for public_axis in range(len(mode)) + ) + + +def compact_stride( + shape: tuple[int, ...], stride_order: tuple[int, ...] +) -> tuple[int, ...]: + """Return compact strides for dimensions ordered fastest to slowest.""" + + stride = [0] * len(shape) + running = 1 + for dimension in stride_order: + stride[dimension] = running + running *= max(shape[dimension], 1) + return tuple(stride) + + +def stride_order_to_public( + canonical_stride_order: tuple[int, ...], + mode: tuple[int, ...] | None = None, +) -> tuple[int, ...]: + """Map a fastest-to-slowest canonical axis order to public array axes.""" + + canonical_stride_order = tuple(canonical_stride_order) + mode = normalize_mode(len(canonical_stride_order), mode) + if tuple(sorted(canonical_stride_order)) != tuple( + range(len(canonical_stride_order)) + ): + raise ValueError( + "canonical_stride_order must be a permutation of " + f"[0, {len(canonical_stride_order) - 1}], got {canonical_stride_order}" + ) + return tuple(mode[canonical_axis] for canonical_axis in canonical_stride_order) + + +def to_cutlass_layout( + shape: tuple[int, ...], + stride: tuple[int, ...], + stride_order: tuple[int, ...], + *, + mode: tuple[int, ...] | None = None, + name: str = "tensor", +) -> tuple[int, ...]: + """Return CUTLASS stride ranks indexed by public JAX array axes. + + The descriptor arguments are in canonical kernel-axis order. ``mode`` + maps those axes back to the public array axes consumed by TensorSpec. + """ + + if stride != compact_stride(shape, stride_order): + raise ValueError( + f"JAX TensorSpec cannot represent non-compact stride {stride} for {name}" + ) + + layout = [0] * len(stride_order) + for rank, dimension in enumerate(stride_order): + layout[dimension] = rank + return to_public_axes(tuple(layout), mode) + + +__all__ = [ + "compact_stride", + "mode_from_layout", + "normalize_mode", + "stride_order_to_public", + "to_canonical_axes", + "to_cutlass_layout", + "to_public_axes", +] diff --git a/python/cudnn/_jax/result.py b/python/cudnn/_jax/result.py new file mode 100644 index 000000000..7f302ede2 --- /dev/null +++ b/python/cudnn/_jax/result.py @@ -0,0 +1,40 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX pytree registration for framework-neutral frontend results.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import jax + +from ..common.result import TupleDict + + +def _flatten(value: TupleDict) -> tuple[tuple[Any, ...], tuple[Any, ...]]: + keys = tuple(dict.keys(value)) + return tuple(dict.__getitem__(value, key) for key in keys), keys + + +def _flatten_with_keys(value: TupleDict): + children, keys = _flatten(value) + return tuple((jax.tree_util.DictKey(key), child) for key, child in zip(keys, children)), keys + + +def _unflatten(keys: tuple[Any, ...], children: Iterable[Any]) -> TupleDict: + return TupleDict(zip(keys, children)) + + +if not getattr(TupleDict, "_jax_pytree_registered", False): + jax.tree_util.register_pytree_with_keys( + TupleDict, + _flatten_with_keys, + _unflatten, + _flatten, + ) + TupleDict._jax_pytree_registered = True + + +__all__ = ["TupleDict"] diff --git a/python/cudnn/api_base.py b/python/cudnn/api_base.py index 2a2a1741a..ed382369a 100644 --- a/python/cudnn/api_base.py +++ b/python/cudnn/api_base.py @@ -11,7 +11,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any, List, Tuple, Optional import logging import threading @@ -20,8 +20,15 @@ import torch import cutlass.cute as cute +from cudnn import data_type from cudnn._experimental_warnings import warn_experimental_api_once -from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.common.result import TupleDict as TupleDict +from cudnn.common.tensor_desc import TensorDesc as _TensorDesc +from cudnn.datatypes import ( + _convert_to_cutlass_data_type, + _cudnn_to_torch_data_type, + _torch_to_cudnn_data_type, +) def ceil_div(a: int, b: int) -> int: @@ -53,23 +60,37 @@ def _reset_experimental_api_warning_registry() -> None: _experimental_api_warnings_emitted.clear() -@dataclass(frozen=True) -class TensorDesc: - """Metadata needed to validate/compile tensor signatures without storage.""" +@dataclass(frozen=True, init=False) +class TensorDesc(_TensorDesc[torch.dtype]): + """Torch tensor signature with device and packed-storage metadata.""" - dtype: torch.dtype - shape: Tuple[int, ...] - stride: Tuple[int, ...] - stride_order: Tuple[int, ...] device: torch.device interpret_uint8_as_fp4x2: bool = False - ndim: int = field(init=False) - name: str = "" + + def __init__( + self, + dtype: torch.dtype, + shape: Tuple[int, ...], + stride: Tuple[int, ...], + stride_order: Tuple[int, ...], + device: torch.device, + interpret_uint8_as_fp4x2: bool = False, + name: str = "", + init_value: bool | int | float | None = None, + ) -> None: + object.__setattr__(self, "device", device) + object.__setattr__(self, "interpret_uint8_as_fp4x2", interpret_uint8_as_fp4x2) + super().__init__( + dtype=dtype, + shape=shape, + stride=stride, + stride_order=stride_order, + name=name, + init_value=init_value, + ) def __post_init__(self): - shape = tuple(self.shape) - stride = tuple(self.stride) - stride_order = tuple(self.stride_order) + super().__post_init__() device = self.device if not isinstance(device, torch.device): try: @@ -77,19 +98,94 @@ def __post_init__(self): except (TypeError, ValueError, RuntimeError) as exc: raise TypeError(f"Invalid device for TensorDesc: {self.device!r}") from exc - ndim = len(shape) - if len(stride) != ndim: - raise ValueError(f"Stride rank mismatch: expected {ndim}, got {len(stride)}") - if len(stride_order) != ndim: - raise ValueError(f"Stride order rank mismatch: expected {ndim}, got {len(stride_order)}") - if tuple(sorted(stride_order)) != tuple(range(ndim)): - raise ValueError(f"Stride order must be a permutation of [0, {ndim - 1}], got {stride_order}") - - object.__setattr__(self, "shape", shape) - object.__setattr__(self, "stride", stride) - object.__setattr__(self, "stride_order", stride_order) object.__setattr__(self, "device", device) - object.__setattr__(self, "ndim", ndim) + + @classmethod + def from_tensor( + cls, + tensor: torch.Tensor, + name: str = "", + *, + shape: Optional[Tuple[int, ...]] = None, + stride: Optional[Tuple[int, ...]] = None, + interpret_uint8_as_fp4x2: bool = False, + ) -> "TensorDesc": + """Capture a Torch tensor's physical or supplied logical layout.""" + + tensor_shape = tuple(tensor.shape) if shape is None else tuple(shape) + tensor_stride = tuple(tensor.stride()) if stride is None else tuple(stride) + tensor_stride_order = cls._compute_stride_order(tensor_shape, tensor_stride) + return cls( + dtype=tensor.dtype, + shape=tensor_shape, + stride=tensor_stride, + stride_order=tensor_stride_order, + device=tensor.device, + interpret_uint8_as_fp4x2=interpret_uint8_as_fp4x2, + name=name, + ) + + @property + def cudnn_dtype(self) -> data_type: + if self.interpret_uint8_as_fp4x2 and self.dtype == torch.uint8: + return data_type.FP4_E2M1 + converted = _torch_to_cudnn_data_type(self.dtype) + return data_type.NOT_SET if converted is None else converted + + def compact_like( + self, + *, + cudnn_dtype: data_type, + shape: tuple[int, ...], + stride_order: tuple[int, ...] | None = None, + name: str = "", + init_value: bool | int | float | None = None, + ) -> "TensorDesc": + """Create a compact Torch descriptor on this descriptor's device.""" + + canonical = super().compact_like( + cudnn_dtype=cudnn_dtype, + shape=shape, + stride_order=stride_order, + name=name, + init_value=init_value, + ) + dtype = _cudnn_to_torch_data_type(cudnn_dtype) + if dtype is None: + raise ValueError(f"Unsupported Torch output dtype {cudnn_dtype}") + return TensorDesc( + dtype=dtype, + shape=canonical.shape, + stride=canonical.stride, + stride_order=canonical.stride_order, + device=self.device, + name=canonical.name, + init_value=canonical.init_value, + ) + + def materialize(self, stream: Optional[cuda.CUstream] = None) -> torch.Tensor: + """Allocate and optionally initialize a tensor described by this object.""" + + if self.interpret_uint8_as_fp4x2 and self.dtype == torch.uint8: + raise ValueError("Materializing a logical FP4 descriptor backed by uint8 storage is unsupported") + + def allocate() -> torch.Tensor: + tensor = torch.empty_strided( + self.shape, + self.stride, + dtype=self.dtype, + device=self.device, + ) + if self.init_value is not None: + tensor.fill_(self.init_value) + return tensor + + if stream is None: + return allocate() + if self.device.type != "cuda": + raise ValueError("A CUDA stream can only be used to materialize a CUDA tensor") + with torch.cuda.stream(torch.cuda.ExternalStream(int(stream), device=self.device)): + return allocate() @staticmethod def _normalize_dim(dim: int, ndim: int, *, allow_new_dim: bool = False) -> int: @@ -181,6 +277,7 @@ def _with_layout(self, shape: Tuple[int, ...], stride: Tuple[int, ...]) -> "Tens device=self.device, interpret_uint8_as_fp4x2=self.interpret_uint8_as_fp4x2, name=self.name, + init_value=self.init_value, ) def __len__(self) -> int: @@ -910,6 +1007,58 @@ def _make_fake_cute_tensor_like( assumed_align=assumed_align, ) + @staticmethod + def _to_tensor_desc( + tensor: torch.Tensor, + name: str = "", + *, + shape: Optional[Tuple[int, ...]] = None, + stride: Optional[Tuple[int, ...]] = None, + interpret_uint8_as_fp4x2: bool = False, + ) -> TensorDesc: + """Create a Torch descriptor from tensor and optional logical layout metadata.""" + + return TensorDesc.from_tensor( + tensor, + name, + shape=shape, + stride=stride, + interpret_uint8_as_fp4x2=interpret_uint8_as_fp4x2, + ) + + @staticmethod + def _materialize_tensor_desc( + desc: _TensorDesc[Any], + *, + device: torch.device, + stream: Optional[cuda.CUstream] = None, + ) -> torch.Tensor: + """Allocate a Torch tensor from a canonical tensor descriptor.""" + + if not isinstance(desc, _TensorDesc): + raise TypeError(f"desc must be a TensorDesc, got {type(desc).__name__}") + dtype = _cudnn_to_torch_data_type(desc.cudnn_dtype) + if dtype is None: + raise ValueError(f"Unsupported Torch output dtype {desc.cudnn_dtype}") + + def allocate() -> torch.Tensor: + tensor = torch.empty_strided( + desc.shape, + desc.stride, + dtype=dtype, + device=device, + ) + if desc.init_value is not None: + tensor.fill_(desc.init_value) + return tensor + + if stream is None: + return allocate() + if torch.device(device).type != "cuda": + raise ValueError("A CUDA stream can only be used to materialize a CUDA tensor") + with torch.cuda.stream(torch.cuda.ExternalStream(int(stream), device=device)): + return allocate() + def _make_tensor_desc( self, tensor: Optional[torch.Tensor], @@ -928,15 +1077,12 @@ def _make_tensor_desc( tensor_stride = self._tensor_stride(tensor, name=name) finally: self._interpret_uint8_as_fp4x2 = prev_interpret - tensor_stride_order = tuple(i for i, s in sorted(enumerate(tensor_stride), key=lambda x: (x[1], tensor_shape[x[0]]))) - return TensorDesc( - dtype=tensor.dtype, + return self._to_tensor_desc( + tensor, + name, shape=tensor_shape, stride=tensor_stride, - stride_order=tensor_stride_order, - device=tensor.device, interpret_uint8_as_fp4x2=interpret_uint8_as_fp4x2, - name=name, ) def _make_fake_cute_tensor_from_desc( @@ -1019,35 +1165,3 @@ def _make_fake_cute_compact_tensor( stride_order=stride_order, assumed_align=assumed_align, ) - - -class TupleDict(dict): - """A dictionary that supports tuple unpacking. - - This class extends dict to allow unpacking like a tuple while still - providing dictionary-style key access. The unpacking order is determined - by the _keys attribute which preserves insertion order. - - Example: - >>> result = TupleDict(a=1, b=2, c=3) - >>> x, y, z = result # Unpacks as (1, 2, 3) - >>> result['a'] # Returns 1 - >>> result[0] # Returns 1 (integer indexing) - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - # Store keys in order for tuple unpacking - self._keys = list(self.keys()) - - def __iter__(self): - """Iterate over values in insertion order for tuple unpacking.""" - return (self[k] for k in self._keys) - - def __getitem__(self, key): - """Support both string keys and integer indices.""" - if isinstance(key, int): - if key < 0 or key >= len(self._keys): - raise IndexError(f"index {key} out of range for TupleDict with {len(self._keys)} items") - return super().__getitem__(self._keys[key]) - return super().__getitem__(key) diff --git a/python/cudnn/block_sparse_attention/__init__.py b/python/cudnn/block_sparse_attention/__init__.py index 172f703e8..ab3f43dd8 100644 --- a/python/cudnn/block_sparse_attention/__init__.py +++ b/python/cudnn/block_sparse_attention/__init__.py @@ -6,8 +6,14 @@ _SYMBOLS = { "block_sparse_attention_forward": (".api", "block_sparse_attention_forward"), "block_sparse_attention_backward": (".api", "block_sparse_attention_backward"), + "BlockSparseAttentionForwardOp": (".op", "BlockSparseAttentionForwardOp"), + "BlockSparseAttentionBackwardOp": (".op", "BlockSparseAttentionBackwardOp"), } +_BSA_SYMBOLS = frozenset( + {"block_sparse_attention_forward", "block_sparse_attention_backward"} +) + def _load_symbol(name): module_name, symbol_name = _SYMBOLS[name] @@ -27,11 +33,11 @@ def __getattr__(name): class BSANamespace: def __getattr__(self, name): - if name in _SYMBOLS: + if name in _BSA_SYMBOLS: return _load_symbol(name) raise AttributeError(f"BSA has no attribute {name!r}") BSA = BSANamespace() -__all__ = ["BSA", *_SYMBOLS.keys()] +__all__ = ["BSA", *_SYMBOLS] diff --git a/python/cudnn/block_sparse_attention/csrc/bwd/bucketed_k2q_csr.py b/python/cudnn/block_sparse_attention/csrc/bwd/bucketed_k2q_csr.py index e11ed68df..5df63c40d 100644 --- a/python/cudnn/block_sparse_attention/csrc/bwd/bucketed_k2q_csr.py +++ b/python/cudnn/block_sparse_attention/csrc/bwd/bucketed_k2q_csr.py @@ -1,11 +1,14 @@ -from typing import Optional, Tuple +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional, Tuple + +if TYPE_CHECKING: + import torch import cutlass import cutlass.cute as cute from cutlass import Int32, const_expr -from cutlass.cute.runtime import from_dlpack import cuda.bindings.driver as cuda -import torch class BucketedK2QCsrUniversal: @@ -240,6 +243,8 @@ def _bucketed_k2q_csr_compile_key( def _to_cute_tensor(tensor: torch.Tensor) -> cute.Tensor: + from cutlass.cute.runtime import from_dlpack + return from_dlpack( tensor.detach(), assumed_align=4, @@ -256,6 +261,8 @@ def build_bucketed_k2q_csr_cutedsl( q2k_block_nums: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, int, int]: """Build bucketed K-to-Q CSR metadata with CuTe DSL kernels.""" + import torch + assert q2k_block_index.dtype == torch.int32 assert q2k_block_index.is_cuda assert q2k_block_index.ndim == 4 diff --git a/python/cudnn/block_sparse_attention/csrc/bwd/sm100_blk128/bsa_bwd_sm100.py b/python/cudnn/block_sparse_attention/csrc/bwd/sm100_blk128/bsa_bwd_sm100.py index 899ef48cb..dd3756a7a 100644 --- a/python/cudnn/block_sparse_attention/csrc/bwd/sm100_blk128/bsa_bwd_sm100.py +++ b/python/cudnn/block_sparse_attention/csrc/bwd/sm100_blk128/bsa_bwd_sm100.py @@ -1,11 +1,14 @@ # Copyright (c) 2025, Ted Zadouri, Markus Hoehnerbach, Jay Shah, Tri Dao. +from __future__ import annotations + import math -from typing import Callable, NamedTuple, Optional, Tuple +from typing import TYPE_CHECKING, Callable, NamedTuple, Optional, Tuple from functools import partial import cuda.bindings.driver as cuda -import torch +if TYPE_CHECKING: + import torch import cutlass import cutlass.cute as cute @@ -28,11 +31,6 @@ from cudnn.block_sparse_attention.csrc.utils.tcgen05_mma_helpers import gemm_w_idx, gemm_ptx_w_idx from cudnn.block_sparse_attention.csrc.utils.seqlen_info import SeqlenInfoQK from cudnn.block_sparse_attention.csrc.utils.block_info import BlockInfo -from cudnn.block_sparse_attention.csrc.bwd.bsa_bwd_prepost import ( - _bwd_postprocess_convert, - _bwd_preprocess, - _get_device_arch, -) from cudnn.block_sparse_attention.csrc.utils.cute_dsl_utils import ParamsBase, sub_packed_f32x2 from cudnn.block_sparse_attention.csrc.utils.tile_scheduler import ( TileSchedulerArguments, @@ -2193,6 +2191,14 @@ def bsa_sm100_blk128_bwd_bucketed_k2q_csr( dv: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """SM100/SM110 blk128 bwd entry receiving BSA bucketed k2q CSR.""" + import torch + + from cudnn.block_sparse_attention.csrc.bwd.bsa_bwd_prepost import ( + _bwd_postprocess_convert, + _bwd_preprocess, + _get_device_arch, + ) + assert q.dtype == torch.bfloat16, "SM100 blk128 bwd only supports bfloat16" assert q.dtype == k.dtype == v.dtype == out.dtype == dout.dtype assert lse.dtype == torch.float32 diff --git a/python/cudnn/block_sparse_attention/csrc/utils/cute_dsl_utils.py b/python/cudnn/block_sparse_attention/csrc/utils/cute_dsl_utils.py index a22e7c7a7..736115a6d 100644 --- a/python/cudnn/block_sparse_attention/csrc/utils/cute_dsl_utils.py +++ b/python/cudnn/block_sparse_attention/csrc/utils/cute_dsl_utils.py @@ -5,18 +5,25 @@ # this file are adapted from quack-kernels 0.4.1 (Apache-2.0) and modified for # cudnn-frontend's CUTLASS DSL 4.5 integration. +from __future__ import annotations + from dataclasses import dataclass, fields from functools import partial from typing import Tuple, get_origin -import torch - import cutlass import cutlass.cute as cute from cutlass._mlir.dialects import nvvm from cutlass.base_dsl.tvm_ffi_builder import spec from cutlass.cutlass_dsl import NumericMeta -from cutlass.cute.runtime import from_dlpack + +try: + import torch +except ModuleNotFoundError: + # Kernel-only users, including cudnn.jax, must not require PyTorch. The + # Torch conversion helpers below fail explicitly if they are invoked + # without the optional dependency. + torch = None # Python scalars and CuTe numeric types are compile-time values. Everything # else in a ParamsBase dataclass is flattened into MLIR values at JIT time. @@ -51,13 +58,17 @@ def convert_single_arg(arg, arg_name, arg_type, ctx): _install_constexpr_tvm_ffi_converter() -torch2cute_dtype_map = { - torch.float16: cutlass.Float16, - torch.bfloat16: cutlass.BFloat16, - torch.float32: cutlass.Float32, - torch.float8_e4m3fn: cutlass.Float8E4M3FN, - torch.float8_e5m2: cutlass.Float8E5M2, -} +torch2cute_dtype_map = ( + {} + if torch is None + else { + torch.float16: cutlass.Float16, + torch.bfloat16: cutlass.BFloat16, + torch.float32: cutlass.Float32, + torch.float8_e4m3fn: cutlass.Float8E4M3FN, + torch.float8_e5m2: cutlass.Float8E5M2, + } +) def _partition_param_fields(obj): @@ -139,6 +150,11 @@ def assume_tensor_aligned(t): def to_cute_tensor(t, assumed_align=16, leading_dim=-1, fully_dynamic=False, enable_tvm_ffi=True): """Convert torch tensor to cute tensor for TVM FFI. leading_dim=-1 defaults to t.ndim-1.""" + if torch is None: + raise ImportError("to_cute_tensor requires PyTorch") + + from cutlass.cute.runtime import from_dlpack + # NOTE: torch 2.9.1 doesn't support fp8 via DLPack but 2.11.0 nightly does # currently export raw bytes as uint8 and tell cutlass correct type # can directly export as fp8 when torch supports it diff --git a/python/cudnn/block_sparse_attention/jax.py b/python/cudnn/block_sparse_attention/jax.py new file mode 100644 index 000000000..53147ad82 --- /dev/null +++ b/python/cudnn/block_sparse_attention/jax.py @@ -0,0 +1,1389 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX adapters for the CuTe DSL block-sparse attention kernels.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +import jax +import jax.numpy as jnp + +from .. import data_type +from .._jax.compiler import compile_options_for_target +from .._jax import JaxApiBase, JaxTensorDesc, TupleDict +from .._jax.layout import mode_from_layout, to_public_axes +from .op import ( + BlockSparseAttentionBackwardOp, + BlockSparseAttentionForwardOp, + SUPPORTED_BACKWARD_COMPUTE_CAPABILITIES, + SUPPORTED_FORWARD_COMPUTE_CAPABILITIES, + ceil_div, + compute_capability_family, +) + + +_DATA_LAYOUTS = ("bhsd", "bshd") + + +def _data_mode(layout: str) -> tuple[int, ...]: + if not isinstance(layout, str): + raise TypeError(f"layout must be a string, got {type(layout).__name__}") + normalized = "".join( + character for character in layout.lower() if character.isalpha() + ) + if normalized not in _DATA_LAYOUTS: + raise ValueError(f"layout must be one of {_DATA_LAYOUTS}, got {layout!r}") + return mode_from_layout(normalized.upper(), kernel_axes="BHSD") + + +def _optional_signature( + value: Any | None, desc: JaxTensorDesc | None, name: str +) -> None: + if (value is None) != (desc is None): + expected = "omitted" if desc is None else "provided" + raise ValueError(f"{name} must be {expected} for this specialized callable") + + +def _default_forward_sparse_block_size(target: int) -> int: + return 64 if compute_capability_family(target) in {90, 120} else 128 + + +def _default_backward_sparse_block_size(target: int) -> int: + return 64 if compute_capability_family(target) == 90 else 128 + + +def _default_bucket_size( + target: int, sparse_block_size: int, num_q_blocks: int, num_heads: int +) -> int: + if compute_capability_family(target) == 90: + return 384 + if sparse_block_size == 64: + if num_q_blocks >= 3000: + return 1024 + return 1088 if num_q_blocks < 2048 or num_q_blocks >= 8192 else 1152 + if num_q_blocks >= 4096 and num_heads <= 1: + return 256 + if num_q_blocks >= 2048: + return 512 + return 384 + + +def _round_up(value: int, alignment: int) -> int: + return ceil_div(value, alignment) * alignment + + +def _sm100_blk64_use_clc( + *, + batch: int, + heads: int, + seqlen_q: int, + block_sparse_num: int, + has_variable_block_nums: bool, +) -> bool: + if has_variable_block_nums: + return True + num_m_blocks = ceil_div(seqlen_q, 64) + if num_m_blocks >= 8192 and block_sparse_num >= 512: + return True + if heads == 1: + return False + return ( + num_m_blocks >= 128 + and batch * heads * num_m_blocks >= 512 + and block_sparse_num <= (64 if heads == 2 else 128) + ) + + +def _resolve_sm100_blk64_use_clc( + *, + kv_splits: int, + requested: bool | None, + batch: int, + heads: int, + seqlen_q: int, + block_sparse_num: int, + has_variable_block_nums: bool, +) -> bool: + if kv_splits > 1: + return False + if requested is not None: + return requested + return _sm100_blk64_use_clc( + batch=batch, + heads=heads, + seqlen_q=seqlen_q, + block_sparse_num=block_sparse_num, + has_variable_block_nums=has_variable_block_nums, + ) + + +def _sm100_blk64_auto_kv_splits(block_sparse_num: int) -> int: + if block_sparse_num >= 900: + return 8 + if block_sparse_num >= 450: + return 4 + if block_sparse_num >= 256: + return 2 + return 1 + + +def _sm100_blk64_uses_int64_kv_strides(desc: JaxTensorDesc) -> bool: + coord_stride_limit = 1 << 27 + batch, heads, seqlen_k, _ = desc.shape + stride_b, stride_h, stride_s, stride_d = desc.stride + rank6_shape = (64, 64, 2, heads, ceil_div(seqlen_k, 64), batch) + rank6_stride = ( + stride_s, + stride_d, + 64 * stride_d, + stride_h, + 64 * stride_s, + stride_b, + ) + if any(stride < 0 or stride > jnp.iinfo(jnp.int32).max for stride in rank6_stride): + return True + return (rank6_shape[4] > 1 and rank6_stride[4] >= coord_stride_limit) or ( + rank6_shape[5] > 1 and rank6_stride[5] >= coord_stride_limit + ) + + +class BlockSparseAttentionForward(JaxApiBase): + """JAX callable specialized from block-sparse attention metadata.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_q2k_block_index: Any, + *, + sample_block_sizes: Any | None = None, + sample_q2k_block_nums: Any | None = None, + sample_o: Any | None = None, + sample_lse: Any | None = None, + block_sparse_num: int | None = None, + sparse_block_size: int | None = None, + allow_empty_block_nums: bool = False, + softmax_scale: float | None = None, + pack_gqa: bool | None = None, + layout: str = "bhsd", + kv_splits: int | str = 1, + use_clc: bool | None = None, + target_compute_capability: int | None = None, + ) -> None: + self.layout = layout + self.data_mode = _data_mode(layout) + self.compute_capability = self._resolve_compute_capability( + target_compute_capability, + SUPPORTED_FORWARD_COMPUTE_CAPABILITIES, + "BlockSparseAttentionForward", + ) + self.compute_capability_family = compute_capability_family( + self.compute_capability + ) + self.q_desc = self._to_tensor_desc(sample_q, "sample_q", mode=self.data_mode) + self.k_desc = self._to_tensor_desc(sample_k, "sample_k", mode=self.data_mode) + self.v_desc = self._to_tensor_desc(sample_v, "sample_v", mode=self.data_mode) + self.block_index_desc = self._to_tensor_desc( + sample_q2k_block_index, "sample_q2k_block_index" + ) + self.block_sizes_desc = ( + None + if sample_block_sizes is None + else self._to_tensor_desc(sample_block_sizes, "sample_block_sizes") + ) + self.block_nums_desc = ( + None + if sample_q2k_block_nums is None + else self._to_tensor_desc(sample_q2k_block_nums, "sample_q2k_block_nums") + ) + + if self.q_desc.ndim != 4 or self.v_desc.ndim != 4: + raise ValueError("sample_q and sample_v must be rank 4") + batch, num_q_heads, seqlen_q, head_dim = self.q_desc.shape + value_dim = self.v_desc.shape[3] + canonical_output_shape = (batch, num_q_heads, seqlen_q, value_dim) + if sample_o is None: + self.o_desc = JaxTensorDesc.from_shape( + to_public_axes(canonical_output_shape, self.data_mode), + self.q_desc.dtype, + name="sample_o", + mode=self.data_mode, + ) + else: + self.o_desc = self._to_tensor_desc( + sample_o, "sample_o", mode=self.data_mode + ) + if sample_lse is None: + self.lse_desc = JaxTensorDesc.from_shape( + (batch, num_q_heads, seqlen_q), + jnp.float32, + name="sample_lse", + ) + else: + self.lse_desc = self._to_tensor_desc(sample_lse, "sample_lse") + + self.sparse_block_size = ( + _default_forward_sparse_block_size(self.compute_capability) + if sparse_block_size is None + else sparse_block_size + ) + capacity = ( + self.block_index_desc.shape[3] if self.block_index_desc.ndim == 4 else 0 + ) + self.block_sparse_num = ( + capacity if block_sparse_num is None else block_sparse_num + ) + self.softmax_scale = ( + head_dim**-0.5 if softmax_scale is None else float(softmax_scale) + ) + if isinstance(kv_splits, str): + if kv_splits != "auto": + raise ValueError("kv_splits string value must be 'auto'") + if not ( + self.compute_capability_family == 100 and self.sparse_block_size == 64 + ): + raise ValueError( + "kv_splits='auto' is available only on the SM100-family blk64 path" + ) + auto_block_count = ( + self.block_sparse_num if self.block_sparse_num > 0 else capacity + ) + self.kv_splits = _sm100_blk64_auto_kv_splits(auto_block_count) + else: + self.kv_splits = kv_splits + self.use_clc = use_clc + self._op = BlockSparseAttentionForwardOp( + q=self.q_desc, + k=self.k_desc, + v=self.v_desc, + block_index=self.block_index_desc, + output=self.o_desc, + lse=self.lse_desc, + block_sizes=self.block_sizes_desc, + block_nums=self.block_nums_desc, + block_sparse_num=self.block_sparse_num, + sparse_block_size=self.sparse_block_size, + softmax_scale=self.softmax_scale, + pack_gqa=pack_gqa, + allow_empty_block_nums=allow_empty_block_nums, + kv_splits=kv_splits, + use_clc=use_clc, + target_compute_capability=self.compute_capability, + ) + self.check_support() + + def check_support(self) -> bool: + return self._op.check_support() + + def __call__( + self, + q: Any, + k: Any, + v: Any, + q2k_block_index: Any, + block_sizes: Any | None = None, + q2k_block_nums: Any | None = None, + ) -> TupleDict: + self.check_support() + _optional_signature(block_sizes, self.block_sizes_desc, "block_sizes") + _optional_signature(q2k_block_nums, self.block_nums_desc, "q2k_block_nums") + + inputs: list[Any] = [q, k, v, q2k_block_index] + input_descs: list[JaxTensorDesc] = [ + self.q_desc, + self.k_desc, + self.v_desc, + self.block_index_desc, + ] + if self.block_sizes_desc is not None: + inputs.append(block_sizes) + input_descs.append(self.block_sizes_desc) + + effective_block_nums = q2k_block_nums + effective_block_nums_desc = self.block_nums_desc + needs_materialized_block_nums = ( + self.compute_capability_family == 90 and effective_block_nums is None + ) + if needs_materialized_block_nums: + block_nums_shape = ( + self._op.batch, + self.block_index_desc.shape[1], + self._op.num_q_blocks, + ) + effective_block_nums = jnp.full( + block_nums_shape, self.block_sparse_num, dtype=jnp.int32 + ) + effective_block_nums_desc = JaxTensorDesc.from_shape( + block_nums_shape, + jnp.int32, + name="effective_q2k_block_nums", + ) + if effective_block_nums is not None: + inputs.append(effective_block_nums) + input_descs.append(effective_block_nums_desc) + + split_offsets_desc = None + if self.kv_splits > 1: + if effective_block_nums is None: + valid_kv = jnp.full( + ( + self._op.batch, + self.block_index_desc.shape[1], + self._op.num_q_blocks, + ), + self.block_sparse_num, + dtype=jnp.int32, + ) + else: + valid_kv = jnp.maximum(effective_block_nums.astype(jnp.int32), 0) + split_ids = jnp.arange(self.kv_splits + 1, dtype=jnp.int32) + quotient = valid_kv // self.kv_splits + remainder_even = valid_kv - quotient * self.kv_splits + even_offsets = ( + quotient[..., None] * split_ids + + (remainder_even[..., None] * split_ids + self.kv_splits - 1) + // self.kv_splits + ) + aligned_base = (quotient // 8) * 8 + aligned_remainder = valid_kv - aligned_base * self.kv_splits + aligned_offsets = aligned_base[..., None] * split_ids + jnp.minimum( + aligned_remainder[..., None], + split_ids * 8, + ) + split_offsets = jnp.where( + (aligned_base == 0)[..., None], + even_offsets, + jnp.minimum(aligned_offsets, valid_kv[..., None]), + ).astype(jnp.int32) + split_offsets_desc = self._to_tensor_desc( + split_offsets, + "split_offsets", + ) + inputs.append(split_offsets) + input_descs.append(split_offsets_desc) + + self._forward_has_effective_block_nums = effective_block_nums is not None + self._forward_has_split_offsets = split_offsets_desc is not None + self._forward_input_count = len(inputs) + + workspace_descs: tuple[JaxTensorDesc, ...] = () + if self.kv_splits > 1: + partial_shape = ( + self._op.batch, + self.kv_splits * self._op.num_q_heads, + self._op.seqlen_q, + self._op.value_dim, + ) + partial_lse_shape = partial_shape[:3] + partial_o_desc = self.q_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=partial_shape, + name="partial_o_workspace", + ) + partial_lse_desc = self.lse_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=partial_lse_shape, + name="partial_lse_workspace", + ) + workspace_descs = (partial_o_desc, partial_lse_desc) + + o, lse = self._call_kernel( + tuple(inputs), + launch=self._launch, + output_descs=(self.o_desc, self.lse_desc), + input_descs=tuple(input_descs), + workspace_descs=workspace_descs, + compile_options=compile_options_for_target(self.compute_capability), + ) + return TupleDict(o_tensor=o, lse_tensor=lse) + + def _launch(self, stream: Any, *arguments: Any) -> None: + import cutlass + import cutlass.cute as cute + + inputs = list(arguments[: self._forward_input_count]) + buffers = list(arguments[self._forward_input_count :]) + q, k, v, block_index = inputs[:4] + cursor = 4 + block_sizes = None + if self.block_sizes_desc is not None: + block_sizes = inputs[cursor] + cursor += 1 + block_nums = None + if self._forward_has_effective_block_nums: + block_nums = inputs[cursor] + cursor += 1 + split_offsets = None + if self._forward_has_split_offsets: + split_offsets = inputs[cursor] + cursor += 1 + if cursor != len(inputs): + raise RuntimeError("Unexpected BSA forward kernel inputs") + + output, lse = buffers[:2] + if self.kv_splits > 1: + partial_output, partial_lse = buffers[2:] + kernel_output, kernel_lse = partial_output, partial_lse + else: + kernel_output, kernel_lse = output, lse + + if self.compute_capability_family == 90: + from .csrc.fwd.sm90_blk64.bsa_fwd_sm90 import ( + BlockSparseAttnForwardSm90Blk64, + ) + + def select(tensor: Any, modes: tuple[int, ...]) -> Any: + return cute.make_tensor( + tensor.iterator, cute.select(tensor.layout, mode=modes) + ) + + q_kernel = select(q, (2, 3, 1, 0)) + k_kernel = select(k, (2, 3, 1, 0)) + v_kernel = select(v, (3, 2, 1, 0)) + o_kernel = select(kernel_output, (2, 3, 1, 0)) + lse_kernel = select(kernel_lse, (2, 1, 0)) + index_kernel = select(block_index, (3, 2, 1, 0)) + nums_kernel = select(block_nums, (2, 1, 0)) + if block_sizes is None: + sizes_kernel = nums_kernel + elif self.block_sizes_desc.ndim == 1: + sizes_kernel = cute.make_tensor( + block_sizes.iterator, + cute.make_layout( + ( + self._op.num_kv_blocks, + self.block_index_desc.shape[1], + self._op.batch, + ), + stride=(block_sizes.stride[0], 0, 0), + ), + ) + elif self.block_sizes_desc.ndim == 2: + sizes_kernel = cute.make_tensor( + block_sizes.iterator, + cute.make_layout( + ( + self._op.num_kv_blocks, + self.block_index_desc.shape[1], + self._op.batch, + ), + stride=(block_sizes.stride[1], 0, block_sizes.stride[0]), + ), + ) + else: + sizes_kernel = select(block_sizes, (2, 1, 0)) + split_kernel = ( + nums_kernel + if split_offsets is None + else select(split_offsets, (3, 2, 1, 0)) + ) + kernel = BlockSparseAttnForwardSm90Blk64( + gqa_ratio=self._op.gqa_ratio, + head_dim=self._op.head_dim, + value_dim=self._op.value_dim, + blocksparse_blocksize_q=64, + blocksparse_blocksize_k=64, + dtype=q.element_type, + acc_dtype=cutlass.Float32, + has_block_sizes=block_sizes is not None, + num_splits=self.kv_splits, + allow_empty_block_nums=self._op.allow_empty_block_nums + and self.block_nums_desc is not None, + ) + kernel( + q_kernel, + k_kernel, + v_kernel, + o_kernel, + lse_kernel, + index_kernel, + nums_kernel, + sizes_kernel, + split_kernel, + cutlass.Float32(self.softmax_scale), + stream, + ) + elif self.compute_capability_family == 120: + from .csrc.fwd.sm120_blk64.bsa_fwd_sm120 import ( + BlockSparseAttnForwardSm120Blk64, + ) + + def select(tensor: Any, modes: tuple[int, ...]) -> Any: + return cute.make_tensor( + tensor.iterator, cute.select(tensor.layout, mode=modes) + ) + + q_kernel = select(q, (2, 3, 1, 0)) + k_kernel = select(k, (2, 3, 1, 0)) + v_kernel = select(v, (3, 2, 1, 0)) + o_kernel = select(kernel_output, (2, 3, 1, 0)) + lse_kernel = select(kernel_lse, (2, 1, 0)) + index_kernel = select(block_index, (3, 2, 1, 0)) + nums_kernel = ( + index_kernel if block_nums is None else select(block_nums, (2, 1, 0)) + ) + block_sizes_mode = 0 + sizes_kernel = nums_kernel + if block_sizes is not None: + block_sizes_mode = self.block_sizes_desc.ndim + sizes_kernel = ( + block_sizes + if block_sizes_mode == 1 + else select(block_sizes, tuple(reversed(range(block_sizes_mode)))) + ) + kernel = BlockSparseAttnForwardSm120Blk64( + gqa_ratio=self._op.gqa_ratio, + head_dim=self._op.head_dim, + value_dim=self._op.value_dim, + blocksparse_blocksize_q=64, + blocksparse_blocksize_k=64, + dtype=q.element_type, + acc_dtype=cutlass.Float32, + has_block_sizes=block_sizes is not None, + has_block_nums=self.block_nums_desc is not None, + block_sizes_mode=block_sizes_mode, + ) + kernel( + q_kernel, + k_kernel, + v_kernel, + o_kernel, + lse_kernel, + index_kernel, + nums_kernel, + cutlass.Int32( + 0 if self.block_nums_desc is not None else self.block_sparse_num + ), + sizes_kernel, + cutlass.Float32(self.softmax_scale), + stream, + ) + elif self.sparse_block_size == 64: + from .csrc.fwd.sm100_blk64.bsa_fwd_sm100 import ( + BlockSparseAttnForwardSm100Blk64, + ) + + use_clc = _resolve_sm100_blk64_use_clc( + kv_splits=self.kv_splits, + requested=self.use_clc, + batch=self._op.batch, + heads=self._op.num_q_heads, + seqlen_q=self._op.seqlen_q, + block_sparse_num=self.block_sparse_num, + has_variable_block_nums=self.block_nums_desc is not None, + ) + kernel = BlockSparseAttnForwardSm100Blk64( + self._op.head_dim, + self._op.value_dim, + qhead_per_kvhead=1, + pack_gqa=False, + m_block_size=64, + n_block_size=256, + sparse_block_size=64, + is_persistent=use_clc, + use_clc_scheduler=use_clc, + allow_empty_block_nums=self._op.allow_empty_block_nums + and self.block_nums_desc is not None, + has_block_sizes=block_sizes is not None, + num_splits=self.kv_splits, + use_int64_kv_strides=( + _sm100_blk64_uses_int64_kv_strides(self.k_desc) + or _sm100_blk64_uses_int64_kv_strides(self.v_desc) + ), + ) + kernel( + q, + k, + v, + kernel_output, + kernel_lse, + cutlass.Float32(self.softmax_scale), + block_index, + block_sizes, + cutlass.Int32( + 0 if self.block_nums_desc is not None else self.block_sparse_num + ), + block_nums, + split_offsets, + stream, + ) + else: + from .csrc.fwd.sm100_blk128.bsa_fwd_sm100 import ( + BlockSparseAttnForwardSm100Blk128, + ) + + kernel = BlockSparseAttnForwardSm100Blk128( + self._op.head_dim, + self._op.value_dim, + qhead_per_kvhead=self._op.gqa_ratio, + pack_gqa=self._op.pack_gqa, + allow_empty_block_nums=self._op.allow_empty_block_nums + and self.block_nums_desc is not None, + has_block_sizes=block_sizes is not None, + ) + kernel( + q, + k, + v, + output, + lse, + cutlass.Float32(self.softmax_scale), + block_index, + block_sizes, + cutlass.Int32(self.block_sparse_num), + block_nums, + stream, + ) + + if self.kv_splits > 1: + from .csrc.fwd.sm100_blk64.bsa_fwd_combine import ( + BlockSparseAttnForwardCombine, + ) + + split_heads = self.kv_splits * self._op.num_q_heads + o_partial = cute.make_tensor( + kernel_output.iterator, + cute.make_layout( + ( + self.kv_splits, + self._op.batch, + self._op.seqlen_q, + self._op.num_q_heads, + self._op.value_dim, + ), + stride=( + self._op.num_q_heads * self._op.seqlen_q * self._op.value_dim, + self._op.seqlen_q * split_heads * self._op.value_dim, + self._op.value_dim, + self._op.seqlen_q * self._op.value_dim, + 1, + ), + ), + ) + lse_partial = cute.make_tensor( + kernel_lse.iterator, + cute.make_layout( + ( + self.kv_splits, + self._op.batch, + self._op.seqlen_q, + self._op.num_q_heads, + ), + stride=( + self._op.num_q_heads * self._op.seqlen_q, + self._op.seqlen_q * split_heads, + 1, + self._op.seqlen_q, + ), + ), + ) + output_bshd = cute.make_tensor( + output.iterator, cute.select(output.layout, mode=(0, 2, 1, 3)) + ) + lse_bsh = cute.make_tensor( + lse.iterator, cute.select(lse.layout, mode=(0, 2, 1)) + ) + combine = BlockSparseAttnForwardCombine( + dtype=output.element_type, + head_dim=self._op.value_dim, + tile_m=16, + k_block_size=64, + log_max_splits=(self.kv_splits - 1).bit_length(), + num_threads=128, + stages=4, + ) + combine( + o_partial, + lse_partial, + output_bshd, + lse_bsh, + None, + None, + None, + None, + None, + stream, + ) + + +class BlockSparseAttentionBackward(JaxApiBase): + """JAX callable specialized from BSA backward tensor metadata.""" + + def __init__( + self, + sample_do: Any, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_o: Any, + sample_lse: Any, + sample_q2k_block_index: Any, + *, + sample_block_sizes: Any | None = None, + sample_q2k_block_nums: Any | None = None, + sample_dq: Any | None = None, + sample_dk: Any | None = None, + sample_dv: Any | None = None, + block_sparse_num: int | None = None, + bucket_size_blocks: int | None = None, + sparse_block_size: int | None = None, + softmax_scale: float | None = None, + layout: str = "bhsd", + target_compute_capability: int | None = None, + ) -> None: + self.layout = layout + self.data_mode = _data_mode(layout) + self.compute_capability = self._resolve_compute_capability( + target_compute_capability, + SUPPORTED_BACKWARD_COMPUTE_CAPABILITIES, + "BlockSparseAttentionBackward", + ) + self.compute_capability_family = compute_capability_family( + self.compute_capability + ) + self.do_desc = self._to_tensor_desc(sample_do, "sample_do", mode=self.data_mode) + self.q_desc = self._to_tensor_desc(sample_q, "sample_q", mode=self.data_mode) + self.k_desc = self._to_tensor_desc(sample_k, "sample_k", mode=self.data_mode) + self.v_desc = self._to_tensor_desc(sample_v, "sample_v", mode=self.data_mode) + self.o_desc = self._to_tensor_desc(sample_o, "sample_o", mode=self.data_mode) + self.lse_desc = self._to_tensor_desc(sample_lse, "sample_lse") + self.block_index_desc = self._to_tensor_desc( + sample_q2k_block_index, "sample_q2k_block_index" + ) + self.block_sizes_desc = ( + None + if sample_block_sizes is None + else self._to_tensor_desc(sample_block_sizes, "sample_block_sizes") + ) + self.block_nums_desc = ( + None + if sample_q2k_block_nums is None + else self._to_tensor_desc(sample_q2k_block_nums, "sample_q2k_block_nums") + ) + + def output_desc( + sample: Any | None, source: JaxTensorDesc, name: str + ) -> JaxTensorDesc: + if sample is not None: + return self._to_tensor_desc(sample, name, mode=self.data_mode) + return JaxTensorDesc.from_shape( + to_public_axes(source.shape, self.data_mode), + source.dtype, + name=name, + mode=self.data_mode, + ) + + self.dq_desc = output_desc(sample_dq, self.q_desc, "sample_dq") + self.dk_desc = output_desc(sample_dk, self.k_desc, "sample_dk") + self.dv_desc = output_desc(sample_dv, self.v_desc, "sample_dv") + self.sparse_block_size = ( + _default_backward_sparse_block_size(self.compute_capability) + if sparse_block_size is None + else sparse_block_size + ) + capacity = ( + self.block_index_desc.shape[3] if self.block_index_desc.ndim == 4 else 0 + ) + self.block_sparse_num = ( + capacity if block_sparse_num is None else block_sparse_num + ) + head_dim = self.q_desc.shape[3] if self.q_desc.ndim == 4 else 1 + self.softmax_scale = ( + head_dim**-0.5 if softmax_scale is None else float(softmax_scale) + ) + num_q_blocks = ( + ceil_div(self.q_desc.shape[2], self.sparse_block_size) + if self.q_desc.ndim == 4 + else 0 + ) + self.bucket_size_blocks = ( + _default_bucket_size( + self.compute_capability, + self.sparse_block_size, + num_q_blocks, + self.q_desc.shape[1], + ) + if bucket_size_blocks is None + else bucket_size_blocks + ) + self._op = BlockSparseAttentionBackwardOp( + dout=self.do_desc, + q=self.q_desc, + k=self.k_desc, + v=self.v_desc, + output=self.o_desc, + lse=self.lse_desc, + block_index=self.block_index_desc, + dq=self.dq_desc, + dk=self.dk_desc, + dv=self.dv_desc, + block_sizes=self.block_sizes_desc, + block_nums=self.block_nums_desc, + block_sparse_num=self.block_sparse_num, + sparse_block_size=self.sparse_block_size, + softmax_scale=self.softmax_scale, + bucket_size_blocks=self.bucket_size_blocks, + target_compute_capability=self.compute_capability, + ) + self.check_support() + + def check_support(self) -> bool: + return self._op.check_support() + + def _workspace_descs(self) -> tuple[tuple[str, JaxTensorDesc], ...]: + int_source = self.block_index_desc + float_source = self.lse_desc + b, h = self._op.batch, self._op.num_heads + g, nk = self._op.num_q_groups, self._op.num_kv_blocks + workspaces: list[tuple[str, JaxTensorDesc]] = [ + ( + "counts", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, h, g, nk), + name="counts_workspace", + init_value=0, + ), + ), + ( + "local_offsets", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, h, g, nk + 1), + name="local_offsets_workspace", + ), + ), + ( + "group_totals", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, h, g), + name="group_totals_workspace", + ), + ), + ( + "bucket_offsets", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, h, g, nk + 1), + name="bucket_offsets_workspace", + ), + ), + ( + "cursors", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, h, g, nk), + name="cursors_workspace", + ), + ), + ( + "bucket_indices", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, h, self._op.max_edges), + name="bucket_indices_workspace", + ), + ), + ] + if self.block_nums_desc is None: + workspaces.append( + ( + "block_nums_placeholder", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, h, self._op.num_q_blocks), + name="block_nums_placeholder", + init_value=self.block_sparse_num, + ), + ) + ) + if self.block_sizes_desc is None: + workspaces.append( + ( + "block_sizes_placeholder", + int_source.compact_like( + cudnn_dtype=data_type.INT32, + shape=(b, nk), + name="block_sizes_placeholder", + init_value=0, + ), + ) + ) + + if self.compute_capability_family == 100 and self.sparse_block_size == 128: + q_rounded = self._op.num_q_blocks * 128 + k_rounded = self._op.num_kv_blocks * 128 + d_rounded = _round_up(self._op.head_dim, 32) + workspaces.extend( + ( + ( + "dpsum", + float_source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(b, h, q_rounded), + name="dpsum_workspace", + ), + ), + ( + "lse_log2", + float_source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(b, h, q_rounded), + name="lse_log2_workspace", + ), + ), + ( + "dq_accum", + float_source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(b, h, q_rounded * d_rounded), + name="dq_accum_workspace", + ), + ), + ) + ) + if self._op.num_q_groups > 1: + workspaces.extend( + ( + ( + "dk_accum", + float_source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(b, h, k_rounded * d_rounded), + name="dk_accum_workspace", + init_value=0.0, + ), + ), + ( + "dv_accum", + float_source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(b, h, k_rounded * d_rounded), + name="dv_accum_workspace", + init_value=0.0, + ), + ), + ) + ) + else: + q_round_to = 64 if self.compute_capability_family == 90 else 8 + k_round_to = 64 if self.compute_capability_family == 90 else 8 + d_round_to = 32 if self.compute_capability_family == 90 else 8 + q_rounded = _round_up(self._op.seqlen_q, q_round_to) + k_rounded = _round_up(self._op.seqlen_k, k_round_to) + d_rounded = _round_up(self._op.head_dim, d_round_to) + elems_per_bh = ( + 2 * q_rounded + q_rounded * d_rounded + 2 * k_rounded * d_rounded + ) + workspaces.append( + ( + "bwd_workspace", + float_source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(b, h, elems_per_bh), + name="bwd_workspace", + init_value=0.0, + ), + ) + ) + return tuple(workspaces) + + def __call__( + self, + do: Any, + q: Any, + k: Any, + v: Any, + o: Any, + lse: Any, + q2k_block_index: Any, + block_sizes: Any | None = None, + q2k_block_nums: Any | None = None, + ) -> TupleDict: + self.check_support() + _optional_signature(block_sizes, self.block_sizes_desc, "block_sizes") + _optional_signature(q2k_block_nums, self.block_nums_desc, "q2k_block_nums") + + inputs: list[Any] = [do, q, k, v, o, lse, q2k_block_index] + input_descs: list[JaxTensorDesc] = [ + self.do_desc, + self.q_desc, + self.k_desc, + self.v_desc, + self.o_desc, + self.lse_desc, + self.block_index_desc, + ] + if self.block_sizes_desc is not None: + inputs.append(block_sizes) + input_descs.append(self.block_sizes_desc) + if self.block_nums_desc is not None: + inputs.append(q2k_block_nums) + input_descs.append(self.block_nums_desc) + self._backward_input_count = len(inputs) + workspace_items = self._workspace_descs() + self._backward_workspace_names = tuple(name for name, _ in workspace_items) + workspace_descs = tuple(desc for _, desc in workspace_items) + + dq, dk, dv = self._call_kernel( + tuple(inputs), + launch=self._launch, + output_descs=(self.dq_desc, self.dk_desc, self.dv_desc), + input_descs=tuple(input_descs), + workspace_descs=workspace_descs, + compile_options=compile_options_for_target(self.compute_capability), + ) + return TupleDict(dq_tensor=dq, dk_tensor=dk, dv_tensor=dv) + + def _launch(self, stream: Any, *arguments: Any) -> None: + import cutlass + import cutlass.cute as cute + + inputs = list(arguments[: self._backward_input_count]) + dq, dk, dv = arguments[ + self._backward_input_count : self._backward_input_count + 3 + ] + workspaces = dict( + zip( + self._backward_workspace_names, + arguments[self._backward_input_count + 3 :], + ) + ) + do, q, k, v, o, lse, block_index = inputs[:7] + cursor = 7 + block_sizes = None + if self.block_sizes_desc is not None: + block_sizes = inputs[cursor] + cursor += 1 + block_nums = None + if self.block_nums_desc is not None: + block_nums = inputs[cursor] + cursor += 1 + if cursor != len(inputs): + raise RuntimeError("Unexpected BSA backward kernel inputs") + if block_nums is None: + block_nums_kernel = workspaces["block_nums_placeholder"] + else: + block_nums_kernel = block_nums + + from .csrc.bwd.bucketed_k2q_csr import BucketedK2QCsrUniversal + + builder = BucketedK2QCsrUniversal( + self.block_sparse_num, + self.bucket_size_blocks, + self.block_nums_desc is not None, + self.block_index_desc.shape[3], + ) + builder( + workspaces["counts"], + workspaces["local_offsets"], + workspaces["group_totals"], + workspaces["bucket_offsets"], + workspaces["cursors"], + workspaces["bucket_indices"], + block_index, + block_nums_kernel, + stream, + ) + + if block_sizes is None: + variable_block_sizes = workspaces["block_sizes_placeholder"] + elif self.block_sizes_desc.ndim == 1: + variable_block_sizes = cute.make_tensor( + block_sizes.iterator, + cute.make_layout( + (self._op.batch, self._op.num_kv_blocks), + stride=(0, block_sizes.stride[0]), + ), + ) + else: + variable_block_sizes = block_sizes + + if self.compute_capability_family == 90: + from .csrc.bwd.sm90_blk64.bsa_bwd_sm90 import ( + BlockSparseAttnBackwardSm90Blk64, + ) + + kernel = BlockSparseAttnBackwardSm90Blk64( + q.element_type, self._op.head_dim, self._op.head_dim + ) + kernel( + ( + cutlass.Int32(self._op.seqlen_q), + cutlass.Int32(self._op.seqlen_k), + cutlass.Int32(self._op.head_dim), + (cutlass.Int32(self._op.num_heads), cutlass.Int32(self._op.batch)), + ), + do, + o, + q, + k, + v, + lse, + dq, + dk, + dv, + workspaces["bucket_offsets"], + workspaces["bucket_indices"], + variable_block_sizes if self.block_sizes_desc is not None else None, + workspaces["bwd_workspace"], + cutlass.Float32(self.softmax_scale), + stream, + ) + return + + if self.sparse_block_size == 64: + from .csrc.bwd.sm100_blk64.bsa_bwd_sm100 import ( + BlockSparseAttnBackwardSm100Blk64, + ) + + kernel = BlockSparseAttnBackwardSm100Blk64( + sparse_block_size=64, + has_block_sizes=self.block_sizes_desc is not None, + ) + kernel( + ( + cutlass.Int32(self._op.seqlen_q), + cutlass.Int32(self._op.seqlen_k), + cutlass.Int32(self._op.head_dim), + (cutlass.Int32(self._op.num_heads), cutlass.Int32(self._op.batch)), + ), + do, + o, + q, + k, + v, + lse, + dq, + dk, + dv, + workspaces["bucket_offsets"], + workspaces["bucket_indices"], + variable_block_sizes, + workspaces["bwd_workspace"], + cutlass.Float32(self.softmax_scale), + stream, + ) + return + + from .csrc.bwd.bsa_bwd_postprocess import BlockSparseAttnBackwardPostprocess + from .csrc.bwd.bsa_bwd_preprocess import BlockSparseAttnBackwardPreprocess + from .csrc.bwd.sm100_blk128.bsa_bwd_sm100 import ( + BsaK2qCsrTensors, + BlockSparseAttnBackwardSm100Blk128, + ) + + def bhsd_to_bshd(tensor: Any) -> Any: + return cute.make_tensor( + tensor.iterator, cute.select(tensor.layout, mode=(0, 2, 1, 3)) + ) + + do_bshd, q_bshd, k_bshd, v_bshd, o_bshd = [ + bhsd_to_bshd(tensor) for tensor in (do, q, k, v, o) + ] + dq_bshd, dk_bshd, dv_bshd = [bhsd_to_bshd(tensor) for tensor in (dq, dk, dv)] + preprocess = BlockSparseAttnBackwardPreprocess( + q.element_type, + self._op.head_dim, + self._op.head_dim, + 128, + ) + preprocess( + o_bshd, + do_bshd, + workspaces["dpsum"], + lse, + workspaces["lse_log2"], + workspaces["dq_accum"], + None, + None, + None, + stream, + ) + use_dkv_postprocess = self._op.num_q_groups > 1 + core = BlockSparseAttnBackwardSm100Blk128( + self._op.head_dim, + force_dkv_postprocess=use_dkv_postprocess, + ) + core( + q_bshd, + k_bshd, + v_bshd, + do_bshd, + workspaces["lse_log2"], + workspaces["dpsum"], + workspaces["dq_accum"], + workspaces["dk_accum"] if use_dkv_postprocess else dk_bshd, + workspaces["dv_accum"] if use_dkv_postprocess else dv_bshd, + cutlass.Float32(self.softmax_scale), + BsaK2qCsrTensors( + workspaces["bucket_offsets"], + workspaces["bucket_indices"], + ), + stream, + ) + postprocess = BlockSparseAttnBackwardPostprocess( + q.element_type, + self._op.head_dim, + self.compute_capability, + 128, + False, + ) + postprocess( + workspaces["dq_accum"], + dq_bshd, + cutlass.Float32(self.softmax_scale), + None, + None, + stream, + ) + if use_dkv_postprocess: + postprocess( + workspaces["dk_accum"], + dk_bshd, + cutlass.Float32(self.softmax_scale), + None, + None, + stream, + ) + postprocess( + workspaces["dv_accum"], + dv_bshd, + cutlass.Float32(1.0), + None, + None, + stream, + ) + + +@partial( + jax.jit, + static_argnames=( + "block_sparse_num", + "sparse_block_size", + "allow_empty_block_nums", + "softmax_scale", + "pack_gqa", + "layout", + "kv_splits", + "use_clc", + "target_compute_capability", + ), +) +def block_sparse_attention_forward( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + q2k_block_index: Any, + block_sparse_num: int | None = None, + block_sizes: Any | None = None, + q2k_block_nums: Any | None = None, + *, + sparse_block_size: int | None = None, + allow_empty_block_nums: bool = False, + softmax_scale: float | None = None, + pack_gqa: bool | None = None, + layout: str = "bhsd", + kv_splits: int | str = 1, + use_clc: bool | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + """Run non-causal block-sparse attention from JAX arrays.""" + + return BlockSparseAttentionForward( + q_tensor, + k_tensor, + v_tensor, + q2k_block_index, + sample_block_sizes=block_sizes, + sample_q2k_block_nums=q2k_block_nums, + block_sparse_num=block_sparse_num, + sparse_block_size=sparse_block_size, + allow_empty_block_nums=allow_empty_block_nums, + softmax_scale=softmax_scale, + pack_gqa=pack_gqa, + layout=layout, + kv_splits=kv_splits, + use_clc=use_clc, + target_compute_capability=target_compute_capability, + )( + q_tensor, + k_tensor, + v_tensor, + q2k_block_index, + block_sizes, + q2k_block_nums, + ) + + +@partial( + jax.jit, + static_argnames=( + "block_sparse_num", + "bucket_size_blocks", + "sparse_block_size", + "softmax_scale", + "layout", + "target_compute_capability", + ), +) +def block_sparse_attention_backward( + do_tensor: Any, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + o_tensor: Any, + lse_tensor: Any, + q2k_block_index: Any, + block_sparse_num: int | None = None, + block_sizes: Any | None = None, + q2k_block_nums: Any | None = None, + *, + softmax_scale: float | None = None, + bucket_size_blocks: int | None = None, + sparse_block_size: int | None = None, + layout: str = "bhsd", + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute explicit block-sparse attention gradients from JAX arrays.""" + + return BlockSparseAttentionBackward( + do_tensor, + q_tensor, + k_tensor, + v_tensor, + o_tensor, + lse_tensor, + q2k_block_index, + sample_block_sizes=block_sizes, + sample_q2k_block_nums=q2k_block_nums, + block_sparse_num=block_sparse_num, + bucket_size_blocks=bucket_size_blocks, + sparse_block_size=sparse_block_size, + softmax_scale=softmax_scale, + layout=layout, + target_compute_capability=target_compute_capability, + )( + do_tensor, + q_tensor, + k_tensor, + v_tensor, + o_tensor, + lse_tensor, + q2k_block_index, + block_sizes, + q2k_block_nums, + ) + + +__all__ = [ + "BlockSparseAttentionBackward", + "BlockSparseAttentionForward", + "block_sparse_attention_backward", + "block_sparse_attention_forward", +] diff --git a/python/cudnn/block_sparse_attention/op.py b/python/cudnn/block_sparse_attention/op.py new file mode 100644 index 000000000..673f2cc78 --- /dev/null +++ b/python/cudnn/block_sparse_attention/op.py @@ -0,0 +1,481 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral contracts for block-sparse attention kernels.""" + +from __future__ import annotations + +from typing import Any + +from .. import data_type +from ..common.op import Op +from ..common.tensor_desc import TensorDesc + + +SUPPORTED_FORWARD_COMPUTE_CAPABILITIES = (90, 100, 103, 107, 110, 120) +SUPPORTED_BACKWARD_COMPUTE_CAPABILITIES = (90, 100, 103, 107, 110) + + +def ceil_div(value: int, divisor: int) -> int: + return (value + divisor - 1) // divisor + + +def compute_capability_family(compute_capability: int) -> int: + if compute_capability == 90: + return 90 + if 100 <= compute_capability < 120: + return 100 + if compute_capability == 120: + return 120 + raise ValueError( + f"Unsupported block-sparse attention target SM{compute_capability}" + ) + + +def _require_shape(desc: TensorDesc[Any], expected: tuple[int, ...]) -> None: + if desc.shape != expected: + raise ValueError( + f"{desc.name or 'tensor'} shape must be {expected}, got {desc.shape}" + ) + + +def _require_dtype(desc: TensorDesc[Any], expected: data_type) -> None: + if desc.cudnn_dtype != expected: + raise ValueError( + f"{desc.name or 'tensor'} dtype must be {expected}, got {desc.cudnn_dtype}" + ) + + +class BlockSparseAttentionForwardOp(Op): + """Complete, framework-neutral signature for BSA forward.""" + + def __init__( + self, + *, + q: TensorDesc[Any], + k: TensorDesc[Any], + v: TensorDesc[Any], + block_index: TensorDesc[Any], + output: TensorDesc[Any], + lse: TensorDesc[Any], + block_sizes: TensorDesc[Any] | None, + block_nums: TensorDesc[Any] | None, + block_sparse_num: int, + sparse_block_size: int, + softmax_scale: float, + pack_gqa: bool | None, + allow_empty_block_nums: bool, + kv_splits: int, + use_clc: bool | None, + target_compute_capability: int, + ) -> None: + self.q = q + self.k = k + self.v = v + self.block_index = block_index + self.output = output + self.lse = lse + self.block_sizes = block_sizes + self.block_nums = block_nums + self.block_sparse_num = block_sparse_num + self.sparse_block_size = sparse_block_size + self.softmax_scale = softmax_scale + self.pack_gqa = pack_gqa + self.allow_empty_block_nums = allow_empty_block_nums + self.kv_splits = kv_splits + self.use_clc = use_clc + self.target_compute_capability = target_compute_capability + + self.batch = 0 + self.num_q_heads = 0 + self.num_kv_heads = 0 + self.seqlen_q = 0 + self.seqlen_k = 0 + self.head_dim = 0 + self.value_dim = 0 + self.num_q_blocks = 0 + self.num_kv_blocks = 0 + self.gqa_ratio = 0 + self.pack_gqa_effective = False + + def check_support(self) -> bool: + if self.target_compute_capability not in SUPPORTED_FORWARD_COMPUTE_CAPABILITIES: + raise ValueError( + f"BlockSparseAttentionForward has no kernel for SM{self.target_compute_capability}; " + f"supported targets are {SUPPORTED_FORWARD_COMPUTE_CAPABILITIES}" + ) + family = compute_capability_family(self.target_compute_capability) + if self.q.ndim != 4 or self.k.ndim != 4 or self.v.ndim != 4: + raise ValueError("q, k, and v must all be rank-4 tensors") + + batch, num_q_heads, seqlen_q, head_dim = self.q.shape + batch_k, num_kv_heads, seqlen_k, head_dim_k = self.k.shape + batch_v, num_v_heads, seqlen_v, value_dim = self.v.shape + if ( + min( + batch, + num_q_heads, + num_kv_heads, + seqlen_q, + seqlen_k, + head_dim, + value_dim, + ) + <= 0 + ): + raise ValueError("q, k, and v dimensions must all be positive") + if batch_k != batch or batch_v != batch: + raise ValueError("q, k, and v batch dimensions must match") + if num_v_heads != num_kv_heads or seqlen_v != seqlen_k: + raise ValueError("k and v head and sequence dimensions must match") + if head_dim_k != head_dim: + raise ValueError("q and k head dimensions must match") + if num_q_heads % num_kv_heads: + raise ValueError( + "the number of query heads must be divisible by the number of KV heads" + ) + if self.q.cudnn_dtype not in {data_type.HALF, data_type.BFLOAT16}: + raise ValueError("q, k, and v must use float16 or bfloat16") + if ( + self.k.cudnn_dtype != self.q.cudnn_dtype + or self.v.cudnn_dtype != self.q.cudnn_dtype + ): + raise ValueError("q, k, and v must have the same dtype") + if self.q.stride[3] != 1 or self.k.stride[3] != 1 or self.v.stride[3] != 1: + raise ValueError("q, k, and v must have a contiguous head dimension") + + if self.sparse_block_size not in {64, 128}: + raise ValueError("sparse_block_size must be 64 or 128") + if family in {90, 120} and self.sparse_block_size != 64: + raise ValueError( + f"SM{self.target_compute_capability} forward requires sparse_block_size=64" + ) + if family == 90: + if head_dim not in {64, 96, 128} or value_dim not in {64, 96, 128}: + raise ValueError( + "SM90 forward supports QK and V dimensions 64, 96, or 128" + ) + if seqlen_q % 64: + raise ValueError( + "SM90 forward requires seqlen_q to be a multiple of 64" + ) + elif family == 120: + if (head_dim, value_dim) != (128, 128): + raise ValueError("SM120 forward requires QK and V dimensions of 128") + elif self.sparse_block_size == 64: + if self.q.cudnn_dtype != data_type.BFLOAT16 or (head_dim, value_dim) != ( + 128, + 128, + ): + raise ValueError( + "SM100-family blk64 forward requires BF16 and QK=V=128" + ) + if num_q_heads != num_kv_heads: + raise ValueError("SM100-family blk64 forward supports MHA only") + elif (head_dim, value_dim) not in {(64, 64), (96, 96), (128, 128)}: + raise ValueError( + "SM100-family blk128 forward supports (QK, V) dimensions " + f"(64, 64), (96, 96), or (128, 128); got ({head_dim}, {value_dim})" + ) + + gqa_ratio = num_q_heads // num_kv_heads + if self.pack_gqa is True and not ( + family == 100 and self.sparse_block_size == 128 + ): + raise ValueError( + "pack_gqa is available only on the SM100-family blk128 path" + ) + if self.pack_gqa is True and 128 % gqa_ratio: + raise ValueError("pack_gqa=True requires the GQA ratio to divide 128") + pack_gqa_effective = ( + family == 100 + and self.sparse_block_size == 128 + and (gqa_ratio > 1 if self.pack_gqa is None else self.pack_gqa) + and 128 % gqa_ratio == 0 + ) + metadata_heads = num_kv_heads if pack_gqa_effective else num_q_heads + metadata_tokens = seqlen_q * gqa_ratio if pack_gqa_effective else seqlen_q + num_q_blocks = ceil_div(metadata_tokens, self.sparse_block_size) + num_kv_blocks = ceil_div(seqlen_k, self.sparse_block_size) + + if self.block_index.ndim != 4: + raise ValueError("block_index must be rank 4") + _require_dtype(self.block_index, data_type.INT32) + if self.block_index.shape[:3] != (batch, metadata_heads, num_q_blocks): + raise ValueError( + "block_index shape prefix must be " + f"{(batch, metadata_heads, num_q_blocks)}, got {self.block_index.shape[:3]}" + ) + capacity = self.block_index.shape[3] + if capacity <= 0: + raise ValueError("block_index must have non-empty KV-block capacity") + + if self.block_nums is None: + minimum = 2 if family == 100 and self.sparse_block_size == 128 else 1 + if not minimum <= self.block_sparse_num <= capacity: + raise ValueError( + f"block_sparse_num must be in [{minimum}, {capacity}], got {self.block_sparse_num}" + ) + if minimum == 2 and self.block_sparse_num % 2: + raise ValueError( + "block_sparse_num must be even for the SM100-family blk128 kernel" + ) + else: + _require_dtype(self.block_nums, data_type.INT32) + _require_shape(self.block_nums, (batch, metadata_heads, num_q_blocks)) + + if self.block_sizes is not None: + _require_dtype(self.block_sizes, data_type.INT32) + allowed_ranks = {1, 2, 3} if family in {90, 120} else {1} + if self.block_sizes.ndim not in allowed_ranks: + raise ValueError( + f"block_sizes rank must be one of {tuple(sorted(allowed_ranks))}" + ) + expected = { + 1: (num_kv_blocks,), + 2: (batch, num_kv_blocks), + 3: (batch, metadata_heads, num_kv_blocks), + }[self.block_sizes.ndim] + _require_shape(self.block_sizes, expected) + + if ( + isinstance(self.kv_splits, bool) + or not isinstance(self.kv_splits, int) + or not 1 <= self.kv_splits <= 256 + ): + raise ValueError( + f"kv_splits must be an integer in [1, 256], got {self.kv_splits!r}" + ) + if self.use_clc is not None and not isinstance(self.use_clc, bool): + raise TypeError( + f"use_clc must be a bool or None, got {type(self.use_clc).__name__}" + ) + if self.use_clc is not None and not ( + family == 100 and self.sparse_block_size == 64 + ): + raise ValueError("use_clc is available only on the SM100-family blk64 path") + if family == 120 and self.kv_splits != 1: + raise ValueError("SM120 forward does not support kv_splits") + if family == 100 and self.sparse_block_size == 128 and self.kv_splits != 1: + raise ValueError("SM100-family blk128 forward does not support kv_splits") + if self.kv_splits > 1 and self.use_clc: + raise ValueError("kv_splits > 1 does not support use_clc=True") + + _require_shape(self.output, (batch, num_q_heads, seqlen_q, value_dim)) + _require_dtype(self.output, self.q.cudnn_dtype) + _require_shape(self.lse, (batch, num_q_heads, seqlen_q)) + _require_dtype(self.lse, data_type.FLOAT) + + self.batch = batch + self.num_q_heads = num_q_heads + self.num_kv_heads = num_kv_heads + self.seqlen_q = seqlen_q + self.seqlen_k = seqlen_k + self.head_dim = head_dim + self.value_dim = value_dim + self.num_q_blocks = num_q_blocks + self.num_kv_blocks = num_kv_blocks + self.gqa_ratio = gqa_ratio + self.pack_gqa_effective = pack_gqa_effective + return True + + +class BlockSparseAttentionBackwardOp(Op): + """Complete, framework-neutral signature for BSA backward.""" + + def __init__( + self, + *, + dout: TensorDesc[Any], + q: TensorDesc[Any], + k: TensorDesc[Any], + v: TensorDesc[Any], + output: TensorDesc[Any], + lse: TensorDesc[Any], + block_index: TensorDesc[Any], + dq: TensorDesc[Any], + dk: TensorDesc[Any], + dv: TensorDesc[Any], + block_sizes: TensorDesc[Any] | None, + block_nums: TensorDesc[Any] | None, + block_sparse_num: int, + sparse_block_size: int, + softmax_scale: float, + bucket_size_blocks: int, + target_compute_capability: int, + ) -> None: + self.dout = dout + self.q = q + self.k = k + self.v = v + self.output = output + self.lse = lse + self.block_index = block_index + self.dq = dq + self.dk = dk + self.dv = dv + self.block_sizes = block_sizes + self.block_nums = block_nums + self.block_sparse_num = block_sparse_num + self.sparse_block_size = sparse_block_size + self.softmax_scale = softmax_scale + self.bucket_size_blocks = bucket_size_blocks + self.target_compute_capability = target_compute_capability + + self.batch = 0 + self.num_heads = 0 + self.seqlen_q = 0 + self.seqlen_k = 0 + self.head_dim = 0 + self.num_q_blocks = 0 + self.num_kv_blocks = 0 + self.num_q_groups = 0 + self.max_edges = 0 + + def check_support(self) -> bool: + if ( + self.target_compute_capability + not in SUPPORTED_BACKWARD_COMPUTE_CAPABILITIES + ): + raise ValueError( + f"BlockSparseAttentionBackward has no kernel for SM{self.target_compute_capability}; " + f"supported targets are {SUPPORTED_BACKWARD_COMPUTE_CAPABILITIES}" + ) + family = compute_capability_family(self.target_compute_capability) + if any( + desc.ndim != 4 + for desc in ( + self.dout, + self.q, + self.k, + self.v, + self.output, + self.dq, + self.dk, + self.dv, + ) + ): + raise ValueError("BSA backward data tensors must all be rank 4") + batch, num_heads, seqlen_q, head_dim = self.q.shape + batch_k, num_kv_heads, seqlen_k, head_dim_k = self.k.shape + if min(batch, num_heads, seqlen_q, seqlen_k, head_dim) <= 0: + raise ValueError("BSA backward dimensions must be positive") + if batch_k != batch or num_kv_heads != num_heads or head_dim_k != head_dim: + raise ValueError( + "BSA backward supports MHA with matching q and k dimensions" + ) + _require_shape(self.v, self.k.shape) + _require_shape(self.dout, self.q.shape) + _require_shape(self.output, self.q.shape) + _require_shape(self.dq, self.q.shape) + _require_shape(self.dk, self.k.shape) + _require_shape(self.dv, self.v.shape) + if self.q.cudnn_dtype != data_type.BFLOAT16: + raise ValueError("BSA backward requires bfloat16 data tensors") + for desc in (self.dout, self.k, self.v, self.output, self.dq, self.dk, self.dv): + _require_dtype(desc, data_type.BFLOAT16) + if any( + desc.stride[3] != 1 + for desc in ( + self.dout, + self.q, + self.k, + self.v, + self.output, + self.dq, + self.dk, + self.dv, + ) + ): + raise ValueError( + "BSA backward data tensors must have a contiguous head dimension" + ) + _require_shape(self.lse, (batch, num_heads, seqlen_q)) + _require_dtype(self.lse, data_type.FLOAT) + + if family == 90: + if self.sparse_block_size != 64 or head_dim != 128: + raise ValueError( + "SM90 backward requires sparse_block_size=64 and head_dim=128" + ) + elif self.sparse_block_size == 64: + if head_dim != 128: + raise ValueError("SM100-family blk64 backward requires head_dim=128") + elif self.sparse_block_size == 128: + if head_dim not in {64, 128}: + raise ValueError( + "SM100-family blk128 backward requires head_dim=64 or 128" + ) + if self.block_sizes is not None: + raise ValueError( + "SM100-family blk128 backward does not support block_sizes" + ) + else: + raise ValueError("sparse_block_size must be 64 or 128") + + num_q_blocks = ceil_div(seqlen_q, self.sparse_block_size) + num_kv_blocks = ceil_div(seqlen_k, self.sparse_block_size) + if self.block_index.ndim != 4: + raise ValueError("block_index must be rank 4") + _require_dtype(self.block_index, data_type.INT32) + if self.block_index.shape[:3] != (batch, num_heads, num_q_blocks): + raise ValueError( + f"block_index shape prefix must be {(batch, num_heads, num_q_blocks)}, got {self.block_index.shape[:3]}" + ) + capacity = self.block_index.shape[3] + if capacity <= 0: + raise ValueError("block_index must have non-empty KV-block capacity") + if self.block_nums is None: + if not 1 <= self.block_sparse_num <= capacity: + raise ValueError( + f"block_sparse_num must be in [1, {capacity}], got {self.block_sparse_num}" + ) + if self.sparse_block_size == 128 and self.block_sparse_num % 2: + raise ValueError( + "block_sparse_num must be even for the blk128 backward kernel" + ) + max_edges = num_q_blocks * self.block_sparse_num + else: + _require_dtype(self.block_nums, data_type.INT32) + _require_shape(self.block_nums, (batch, num_heads, num_q_blocks)) + max_edges = num_q_blocks * capacity + + if self.block_sizes is not None: + _require_dtype(self.block_sizes, data_type.INT32) + if self.block_sizes.ndim not in {1, 2}: + raise ValueError("block_sizes must have rank 1 or 2 for BSA backward") + expected = ( + (num_kv_blocks,) + if self.block_sizes.ndim == 1 + else (batch, num_kv_blocks) + ) + _require_shape(self.block_sizes, expected) + if ( + isinstance(self.bucket_size_blocks, bool) + or not isinstance(self.bucket_size_blocks, int) + or self.bucket_size_blocks <= 0 + ): + raise ValueError( + f"bucket_size_blocks must be positive, got {self.bucket_size_blocks!r}" + ) + + self.batch = batch + self.num_heads = num_heads + self.seqlen_q = seqlen_q + self.seqlen_k = seqlen_k + self.head_dim = head_dim + self.num_q_blocks = num_q_blocks + self.num_kv_blocks = num_kv_blocks + self.num_q_groups = ceil_div(num_q_blocks, self.bucket_size_blocks) + self.max_edges = max(1, max_edges) + return True + + +__all__ = [ + "BlockSparseAttentionBackwardOp", + "BlockSparseAttentionForwardOp", + "SUPPORTED_BACKWARD_COMPUTE_CAPABILITIES", + "SUPPORTED_FORWARD_COMPUTE_CAPABILITIES", + "ceil_div", + "compute_capability_family", +] diff --git a/python/cudnn/common/__init__.py b/python/cudnn/common/__init__.py new file mode 100644 index 000000000..c833f57f0 --- /dev/null +++ b/python/cudnn/common/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral metadata and API infrastructure.""" diff --git a/python/cudnn/common/cute_arch.py b/python/cudnn/common/cute_arch.py new file mode 100644 index 000000000..cbc678310 --- /dev/null +++ b/python/cudnn/common/cute_arch.py @@ -0,0 +1,30 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""CuTe architecture naming shared by framework adapters.""" + +from __future__ import annotations + +_ARCH_MAP = { + 90: "sm_90a", + 100: "sm_100a", + 103: "sm_103a", + 107: "sm_100f", + 110: "sm_110a", + 120: "sm_120a", +} + + +def gpu_arch_flag_for_compute_capability(compute_capability: int) -> str: + """Return the explicit CuTe architecture flag for a compilation target.""" + + if isinstance(compute_capability, bool) or not isinstance(compute_capability, int): + raise TypeError(f"compute_capability must be an int, got {type(compute_capability).__name__}") + try: + return _ARCH_MAP[compute_capability] + except KeyError as error: + supported = ", ".join(f"SM{value}" for value in sorted(_ARCH_MAP)) + raise RuntimeError(f"Unsupported GPU compute capability SM{compute_capability} for CuTe kernels; supported targets are {supported}") from error + + +__all__ = ["gpu_arch_flag_for_compute_capability"] diff --git a/python/cudnn/common/op.py b/python/cudnn/common/op.py new file mode 100644 index 000000000..4c5f97712 --- /dev/null +++ b/python/cudnn/common/op.py @@ -0,0 +1,17 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral operation contract.""" + +from abc import ABC, abstractmethod + + +class Op(ABC): + """Logical operation specialized from a complete tensor signature.""" + + @abstractmethod + def check_support(self) -> bool: + """Validate the operation signature and resolve static configuration.""" + + +__all__ = ["Op"] diff --git a/python/cudnn/common/operation_api.py b/python/cudnn/common/operation_api.py new file mode 100644 index 000000000..e2a395bd5 --- /dev/null +++ b/python/cudnn/common/operation_api.py @@ -0,0 +1,54 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Reusable lazy exports for framework-specific operation packages.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, MutableMapping, Sequence +from importlib import import_module +from typing import Any + + +def make_operation_api( + module_globals: MutableMapping[str, Any], + *, + exports: Mapping[str, Sequence[str]], + submodules: Sequence[str] = (), +) -> tuple[list[str], Callable[[str], Any], Callable[[], list[str]]]: + """Create ``__all__``, ``__getattr__``, and ``__dir__`` for an operation. + + ``exports`` maps sibling module names to the public symbols provided by + each module. ``submodules`` lists sibling modules that are themselves + available as lazy attributes. Resolved values are cached in the package's + globals, following Python's normal module import behavior. + """ + + package_name = module_globals["__name__"] + export_routes: dict[str, str] = {} + for module_name, names in exports.items(): + for name in names: + if name in export_routes: + raise ValueError(f"Operation export {name!r} is provided by more than one module") + export_routes[name] = module_name + + lazy_submodules = frozenset(submodules) + + def get_attribute(name: str) -> Any: + if name in lazy_submodules: + value = import_module(f".{name}", package_name) + elif name in export_routes: + value = getattr(import_module(f".{export_routes[name]}", package_name), name) + else: + raise AttributeError(f"module {package_name!r} has no attribute {name!r}") + + module_globals[name] = value + return value + + def list_attributes() -> list[str]: + return sorted((*module_globals, *export_routes, *lazy_submodules)) + + return sorted(export_routes), get_attribute, list_attributes + + +__all__ = ["make_operation_api"] diff --git a/python/cudnn/common/result.py b/python/cudnn/common/result.py new file mode 100644 index 000000000..eca466092 --- /dev/null +++ b/python/cudnn/common/result.py @@ -0,0 +1,34 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral result containers shared by frontend APIs.""" + +from __future__ import annotations + +from typing import Any + + +class TupleDict(dict): + """Dictionary result that also unpacks and indexes like a tuple. + + Values follow insertion order, matching the result order documented by + the operation wrapper. String-key access remains ordinary dictionary + access. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._keys = tuple(self.keys()) + + def __iter__(self): + return (dict.__getitem__(self, key) for key in self._keys) + + def __getitem__(self, key: Any) -> Any: + if isinstance(key, int): + if key < 0 or key >= len(self._keys): + raise IndexError(f"index {key} out of range for TupleDict with {len(self._keys)} items") + key = self._keys[key] + return dict.__getitem__(self, key) + + +__all__ = ["TupleDict"] diff --git a/python/cudnn/common/tensor_desc.py b/python/cudnn/common/tensor_desc.py new file mode 100644 index 000000000..6a998d059 --- /dev/null +++ b/python/cudnn/common/tensor_desc.py @@ -0,0 +1,182 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral tensor metadata for operation kernels.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from numbers import Real +from operator import index +from typing import Any, Generic, TypeVar + +from .. import data_type + +DTypeT = TypeVar("DTypeT") + + +@dataclass(frozen=True) +class TensorDesc(Generic[DTypeT]): + """Framework-neutral tensor signature used by operation kernels. + + ``stride_order`` lists dimension indices from the smallest stride to the + largest. Framework adapters may extend this descriptor with storage or + lowering metadata, but operation kernels only consume the fields defined + here and :attr:`cudnn_dtype`. ``init_value`` is a scalar used to initialize + inferred output or workspace storage; ``None`` leaves the storage + uninitialized. + """ + + dtype: DTypeT + shape: tuple[int, ...] + stride: tuple[int, ...] + stride_order: tuple[int, ...] + ndim: int = field(init=False) + name: str = "" + init_value: bool | int | float | None = None + + def __post_init__(self) -> None: + shape = self._integer_tuple(self.shape, "shape") + stride = self._integer_tuple(self.stride, "stride") + stride_order = self._integer_tuple(self.stride_order, "stride_order") + ndim = len(shape) + + if self.init_value is not None and not isinstance(self.init_value, (bool, Real)): + raise TypeError(f"TensorDesc.init_value must be a real scalar or None, got {self.init_value!r}") + + if any(size < 0 for size in shape): + raise ValueError(f"TensorDesc.shape entries must be non-negative, got {shape}") + if any(value < 0 for value in stride): + raise ValueError(f"TensorDesc.stride entries must be non-negative, got {stride}") + if len(stride) != ndim: + raise ValueError(f"Stride rank mismatch: expected {ndim}, got {len(stride)}") + if len(stride_order) != ndim: + raise ValueError(f"Stride order rank mismatch: expected {ndim}, got {len(stride_order)}") + if tuple(sorted(stride_order)) != tuple(range(ndim)): + raise ValueError(f"Stride order must be a permutation of [0, {ndim - 1}], got {stride_order}") + ordered_strides = tuple(stride[dimension] for dimension in stride_order) + if any(left > right for left, right in zip(ordered_strides, ordered_strides[1:])): + raise ValueError(f"Stride order {stride_order} is inconsistent with stride {stride}") + + object.__setattr__(self, "shape", shape) + object.__setattr__(self, "stride", stride) + object.__setattr__(self, "stride_order", stride_order) + object.__setattr__(self, "ndim", ndim) + + @staticmethod + def _integer_tuple(values, field_name: str) -> tuple[int, ...]: + normalized = [] + for value in values: + if isinstance(value, bool): + raise TypeError(f"TensorDesc.{field_name} entries must be integers, got {value!r}") + try: + normalized.append(index(value)) + except TypeError as error: + raise TypeError(f"TensorDesc.{field_name} entries must be integers, got {value!r}") from error + return tuple(normalized) + + @property + def cudnn_dtype(self) -> data_type: + """Return the canonical dtype consumed by operation kernels. + + Framework-specific subclasses override this property when ``dtype`` is + native to that framework. + """ + + if not isinstance(self.dtype, data_type): + raise TypeError(f"{type(self).__name__}.dtype must provide a cudnn.data_type mapping, " f"got {type(self.dtype).__name__}") + return self.dtype + + def compact_like( + self, + *, + cudnn_dtype: data_type, + shape: tuple[int, ...], + stride_order: tuple[int, ...] | None = None, + name: str = "", + init_value: bool | int | float | None = None, + ) -> "TensorDesc[Any]": + """Create a compact descriptor derived from this descriptor. + + The framework-neutral implementation returns a canonical descriptor. + Framework-specific descriptors override this method to preserve their + native dtype and allocation metadata. + """ + + return make_compact_tensor_desc( + dtype=cudnn_dtype, + shape=shape, + stride_order=stride_order, + name=name, + init_value=init_value, + ) + + def is_compact(self, stride_order: tuple[int, ...] | None = None) -> bool: + """Return whether the descriptor is contiguous in the given order. + + Size-one dimensions may carry any non-negative stride because they do + not change the addressed storage. When no order is supplied, the + descriptor's canonical ``stride_order`` is used. + """ + + order = self.stride_order if stride_order is None else self._integer_tuple(stride_order, "stride_order") + if len(order) != self.ndim or tuple(sorted(order)) != tuple(range(self.ndim)): + raise ValueError(f"Stride order must be a permutation of [0, {self.ndim - 1}], got {order}") + + expected_stride = 1 + for dimension in order: + size = self.shape[dimension] + if size != 1 and self.stride[dimension] != expected_stride: + return False + expected_stride *= max(size, 1) + return True + + +def make_compact_tensor_desc( + *, + dtype: data_type, + shape: tuple[int, ...], + stride_order: tuple[int, ...] | None = None, + name: str = "", + init_value: bool | int | float | None = None, +) -> TensorDesc[data_type]: + """Construct a framework-neutral descriptor for a compact tensor. + + ``stride_order`` lists dimensions from fastest varying to slowest varying. + It defaults to the conventional row-major order. + """ + + if not isinstance(dtype, data_type): + raise TypeError(f"dtype must be a cudnn.data_type, got {type(dtype).__name__}") + + shape = TensorDesc._integer_tuple(shape, "shape") + if any(size < 0 for size in shape): + raise ValueError(f"TensorDesc.shape entries must be non-negative, got {shape}") + + if stride_order is None: + stride_order = tuple(reversed(range(len(shape)))) + else: + stride_order = TensorDesc._integer_tuple(stride_order, "stride_order") + + if len(stride_order) != len(shape): + raise ValueError(f"Stride order rank mismatch: expected {len(shape)}, got {len(stride_order)}") + if tuple(sorted(stride_order)) != tuple(range(len(shape))): + raise ValueError(f"Stride order must be a permutation of [0, {len(shape) - 1}], got {stride_order}") + + stride = [0] * len(shape) + running = 1 + for dimension in stride_order: + stride[dimension] = running + running *= max(shape[dimension], 1) + + return TensorDesc( + dtype=dtype, + shape=shape, + stride=tuple(stride), + stride_order=stride_order, + name=name, + init_value=init_value, + ) + + +__all__ = ["TensorDesc", "make_compact_tensor_desc"] diff --git a/python/cudnn/deepseek_sparse_attention/README.md b/python/cudnn/deepseek_sparse_attention/README.md index fc892210a..46b1a77ec 100644 --- a/python/cudnn/deepseek_sparse_attention/README.md +++ b/python/cudnn/deepseek_sparse_attention/README.md @@ -1,11 +1,16 @@ ## DSA module -- **Indexer Forward**: CuTe-DSL score kernel (Q @ Kᵗ, ReLU, head reduce, ratio causal mask). Supports the SM90 MQA path through `indexer_forward_wrapper` and SM100+ through `IndexerForward`; non-fused, so pair with **Indexer Top-K** for the top-K stage. +Install the JAX integration with +`pip install 'nvidia-cudnn-frontend[jax]'`. JAX classes and JIT-compiled +wrappers are available as direct `cudnn.jax` exports and through +`cudnn.jax.DSA`. + +- **Indexer Forward**: CuTe-DSL score kernel (Q @ Kᵗ, ReLU, head reduce, ratio causal mask). Supports SM90 MQA and SM100+ with fixed BSHD/SBHD or packed THD JAX inputs; non-fused, so pair with **Indexer Top-K** for the top-K stage. - **Indexer Top-K**: SM90+ CuTe-DSL radix top-K kernel with per-row ``seq_lens``. - **Sparse Attention Backward**: DSA backward (FlashMLA-shape, SM90/SM100). - **Sparse Indexer / Attention Score Recompute**: Sparse (top-K) recomputation of indexer and attention scores for training loss. - **Dense Indexer / Attention Score Recompute**: Dense (full-KV) analogues of the above. -- **Indexer Backward**: Three-stage pipeline (score-grad, three GEMMs, dtype cast) for sparse top-K score tensors. +- **Indexer Backward**: Three-stage pipeline (score-grad, three GEMMs, dtype cast) for sparse top-K score tensors. JAX additionally supports packed THD tensors when `topk_indices_global=True`; the PyTorch sparse wrapper remains fixed BSHD. - **Dense Indexer Backward**: Full-KV counterpart of Indexer Backward. ## Acknowledgements diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py index fbae12785..787e3868e 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/__init__.py @@ -1,13 +1,20 @@ -from .api import ( - DenseIndexerBackward, - IndexerBackward, - dense_indexer_backward_wrapper, - indexer_backward_wrapper, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy Torch API and framework-neutral indexer-backward exports.""" -__all__ = [ - "DenseIndexerBackward", - "IndexerBackward", - "dense_indexer_backward_wrapper", - "indexer_backward_wrapper", -] +from ...common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ("DenseIndexerBackwardOp", "IndexerBackwardOp"), + "api": ( + "DenseIndexerBackward", + "IndexerBackward", + "dense_indexer_backward_wrapper", + "indexer_backward_wrapper", + ), + }, + submodules=("api", "op"), +) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py index 023c0e20a..b9440c255 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/api.py @@ -17,6 +17,8 @@ import torch import cuda.bindings.driver as cuda +from cudnn import data_type +from cudnn.common.tensor_desc import make_compact_tensor_desc from cudnn.api_base import APIBase, TupleDict from cudnn.deepseek_sparse_attention.utils.runtime import ( torch_stream_context as _torch_stream_context, @@ -27,6 +29,7 @@ from .dense_indexer_backward_sm90 import dense_indexer_backward_sm90 from .indexer_backward_sm100 import indexer_backward_sm100 from .indexer_backward_sm90 import indexer_backward_sm90 +from .op import DenseIndexerBackwardOp, IndexerBackwardOp def _as_grad_loss_tensor(grad_loss: Union[float, torch.Tensor], device: torch.device) -> torch.Tensor: @@ -161,18 +164,20 @@ def __init__( self.idx_score_desc = self._make_tensor_desc(sample_index_score, name="sample_index_score") self.topk_desc = self._make_tensor_desc(sample_topk_indices, name="sample_topk_indices") - b, s_q, h, d = sample_index_q.shape - s_k = sample_index_k.shape[1] - topk = sample_topk_indices.shape[-1] - self.batch = b - self.seqlen = s_q - self.seqlen_k = s_k - self.heads = h - self.head_dim = d - self.topk = topk - self.sm_scale = float(sm_scale) - self.block_I = int(block_I) - self.topk_indices_global = bool(topk_indices_global) + self._op = IndexerBackwardOp( + index_q=self.iq_desc, + weights=self.w_desc, + index_k=self.ik_desc, + d_index_q=self.diq_desc, + d_weights=self.dw_desc, + d_index_k=self.dik_desc, + attn_score=self.attn_desc, + index_score=self.idx_score_desc, + topk_indices=self.topk_desc, + sm_scale=sm_scale, + block_i=block_I, + topk_indices_global=topk_indices_global, + ) def check_support(self) -> bool: major, _ = torch.cuda.get_device_capability() @@ -180,12 +185,7 @@ def check_support(self) -> bool: major != 9 and major < 10, f"IndexerBackward requires SM90 or SM100+, found SM{major}", ) - self._check_dtype(self.iq_desc, torch.bfloat16, name="index_q") - self._check_dtype(self.w_desc, torch.bfloat16, name="weights") - self._check_dtype(self.ik_desc, torch.bfloat16, name="index_k") - self._check_dtype(self.attn_desc, torch.float32, name="attn_score") - self._check_dtype(self.idx_score_desc, torch.float32, name="index_score") - self._check_dtype(self.topk_desc, torch.int32, name="topk_indices") + self._op.check_support() self._is_supported = True return True @@ -201,15 +201,15 @@ def compile(self) -> None: major, _ = torch.cuda.get_device_capability() kernel_factory = indexer_backward_sm90 if major == 9 else indexer_backward_sm100 self._compiled_kernel = kernel_factory( - self.batch, - self.seqlen, - self.seqlen_k, - self.heads, - self.head_dim, - self.topk, - sm_scale=self.sm_scale, - block_I=self.block_I, - topk_indices_global=self.topk_indices_global, + self._op.batch, + self._op.seqlen_q, + self._op.seqlen_k, + self._op.heads, + self._op.head_dim, + self._op.topk, + sm_scale=self._op.sm_scale, + block_I=self._op.block_i, + topk_indices_global=self._op.topk_indices_global, ) def execute( @@ -231,7 +231,7 @@ def execute( # ``grad_scale`` is forwarded as a runtime ``Float32`` arg; the # compiled kernel is reused when loss_coeff changes for the same # cached tensor shape. - grad_scale = float(loss_coeff) / (int(self.batch) * int(self.seqlen)) + grad_scale = float(loss_coeff) / (int(self._op.batch) * int(self._op.seqlen_q)) grad_loss_tensor = _as_grad_loss_tensor(grad_loss, index_q.device) self._compiled_kernel( index_q, @@ -290,31 +290,35 @@ def __init__( self.idx_score_desc = self._make_tensor_desc(sample_index_score, name="sample_index_score") self.idx_lse_desc = self._make_tensor_desc(sample_index_lse, name="sample_index_lse") - self.sm_scale = float(sm_scale) - self.block_I = int(block_I) - self.ratio = int(ratio) - self.is_thd = bool(is_thd) - self.has_q_causal_offsets = bool(has_q_causal_offsets) - - if self.is_thd: - total_q, heads, head_dim = sample_index_q.shape - total_k = sample_index_k.shape[0] - if batch is None or max_seqlen_q is None or max_seqlen_k is None: - raise ValueError("THD dense indexer backward requires batch and max_seqlen_q/k") - self.batch = int(batch) - self.normalization_tokens = int(total_q) - self.total_k = int(total_k) - self.max_seqlen_q = int(max_seqlen_q) - self.max_seqlen_k = int(max_seqlen_k) - else: - b, seqlen_q, heads, head_dim = sample_index_q.shape - self.batch = int(b) - self.normalization_tokens = int(b * seqlen_q) - self.total_k = int(b * sample_index_k.shape[1]) - self.max_seqlen_q = int(seqlen_q) - self.max_seqlen_k = int(sample_index_k.shape[1]) - self.heads = int(heads) - self.head_dim = int(head_dim) + is_thd = bool(is_thd) + if is_thd and batch is None: + raise ValueError("THD dense indexer backward requires batch") + inferred_batch = int(batch) if is_thd else int(sample_index_q.shape[0]) + cu_seqlens_q_desc = make_compact_tensor_desc(dtype=data_type.INT32, shape=(inferred_batch + 1,), name="sample_cu_seqlens_q") if is_thd else None + cu_seqlens_k_desc = make_compact_tensor_desc(dtype=data_type.INT32, shape=(inferred_batch + 1,), name="sample_cu_seqlens_k") if is_thd else None + q_causal_offsets_desc = ( + make_compact_tensor_desc(dtype=data_type.INT32, shape=(inferred_batch,), name="sample_q_causal_offsets") if has_q_causal_offsets else None + ) + self._op = DenseIndexerBackwardOp( + index_q=self.iq_desc, + weights=self.w_desc, + index_k=self.ik_desc, + d_index_q=self.diq_desc, + d_weights=self.dw_desc, + d_index_k=self.dik_desc, + attn_score=self.attn_desc, + attn_l1norm=self.attn_denom_desc, + index_score=self.idx_score_desc, + index_lse=self.idx_lse_desc, + cu_seqlens_q=cu_seqlens_q_desc, + cu_seqlens_k=cu_seqlens_k_desc, + q_causal_offsets=q_causal_offsets_desc, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + sm_scale=sm_scale, + block_i=block_I, + ratio=ratio, + ) self._uses_current_stream_pipeline = False def check_support(self) -> bool: @@ -323,19 +327,7 @@ def check_support(self) -> bool: major != 9 and major < 10, f"DenseIndexerBackward requires SM90 or SM100+, found SM{major}", ) - self._check_dtype(self.iq_desc, torch.bfloat16, name="index_q") - self._check_dtype(self.w_desc, torch.bfloat16, name="weights") - self._check_dtype(self.ik_desc, torch.bfloat16, name="index_k") - self._check_dtype(self.diq_desc, torch.bfloat16, name="d_index_q") - self._check_dtype(self.dw_desc, torch.bfloat16, name="d_weights") - self._check_dtype(self.dik_desc, [torch.bfloat16, torch.float32], name="d_index_k") - self._check_dtype(self.attn_desc, torch.float32, name="attn_score") - self._check_dtype(self.attn_denom_desc, torch.float32, name="attn_l1norm") - self._check_dtype(self.idx_score_desc, torch.float32, name="index_score") - self._check_dtype(self.idx_lse_desc, torch.float32, name="index_lse") - self._value_error_if(self.block_I <= 0, f"block_I must be positive, got {self.block_I}") - self._value_error_if(self.ratio < 1, f"ratio must be >= 1, got {self.ratio}") - self._value_error_if(self.heads < 64, f"DenseIndexerBackward requires heads >= 64, got {self.heads}") + self._op.check_support() self._is_supported = True return True @@ -347,16 +339,16 @@ def compile(self) -> None: kernel_factory = dense_indexer_backward_sm90 if major == 9 else dense_indexer_backward_sm100 self._uses_current_stream_pipeline = major == 9 self._compiled_kernel = kernel_factory( - self.batch, - self.max_seqlen_q, - self.max_seqlen_k, - self.heads, - self.head_dim, - sm_scale=self.sm_scale, - block_I=self.block_I, - ratio=self.ratio, - is_varlen=self.is_thd, - has_q_causal_offsets=self.has_q_causal_offsets, + self._op.batch, + self._op.max_seqlen_q, + self._op.max_seqlen_k, + self._op.heads, + self._op.head_dim, + sm_scale=self._op.sm_scale, + block_I=self._op.block_i, + ratio=self._op.ratio, + is_varlen=self._op.is_thd, + has_q_causal_offsets=self._op.q_causal_offsets is not None, ) def execute( @@ -381,7 +373,7 @@ def execute( backend_stream = None if self._uses_current_stream_pipeline else current_stream with _torch_stream_context(backend_stream): grad_loss_value = float(_as_grad_loss_tensor(grad_loss, index_q.device).item()) - grad_scale = float(loss_coeff) * grad_loss_value / max(int(self.normalization_tokens), 1) + grad_scale = float(loss_coeff) * grad_loss_value / max(int(self._op.normalization_tokens), 1) # Dense backward's dK path uses atomic/bulk reductions into fp32. d_index_k_target = d_index_k @@ -500,7 +492,7 @@ def indexer_backward_wrapper( block_I=block_I, topk_indices_global=topk_indices_global, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_IndexerBackwardObjects[key] = obj @@ -647,7 +639,7 @@ def dense_indexer_backward_wrapper( max_seqlen_k=max_k, has_q_causal_offsets=q_causal_offsets is not None, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_DenseIndexerBackwardObjects[key] = obj diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py index 7e8112211..c55059124 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm100.py @@ -53,7 +53,6 @@ import math from functools import partial -import torch import cuda.bindings.driver as cuda import cutlass @@ -73,12 +72,7 @@ import cutlass.utils.blackwell_helpers as sm100_utils_basic -from cudnn.deepseek_sparse_attention.utils.compiler import compile_options - from cudnn.deepseek_sparse_attention.utils.copy import cpasync_reduce_bulk_add_f32 -from cudnn.deepseek_sparse_attention.utils.runtime import ( - resolve_stream as _resolve_stream, -) from cudnn.deepseek_sparse_attention.utils.seqlen import seqlen_info as _seqlen_info mul_packed_f32x2 = partial(cute.arch.mul_packed_f32x2, rnd="rn") @@ -151,7 +145,12 @@ def __call__( max_seqlen_q: Int32, max_seqlen_k: Int32, stream, + mGradSignal=None, + mGradLoss=None, ): + if const_expr(mGradSignal is None): + mGradSignal = mPredictRaw + # BSHD: mCuSeqlensQ/K = None; tensors carry batch dim. # PredictRaw/TargetRaw: (B, S_q, S_k) LSE/Denom: (B, S_q) # THD: mCuSeqlensQ/K provided; tensors are packed. @@ -167,6 +166,10 @@ def __call__( mPredictRaw.iterator, cute.select(mPredictRaw.layout, mode=[1, 2, 0]), ) + mGradSignal = cute.make_tensor( + mGradSignal.iterator, + cute.select(mGradSignal.layout, mode=[1, 2, 0]), + ) mTargetRaw = cute.make_tensor( mTargetRaw.iterator, cute.select(mTargetRaw.layout, mode=[1, 2, 0]), @@ -186,12 +189,14 @@ def __call__( self.kernel_score_grad( mPredictRaw, + mGradSignal, mTargetRaw, mLSE, mDenomTarget, mCuSeqlensQ, mCuSeqlensK, mQCausalOffsets, + mGradLoss, grad_scale, seqlen_k_pad, max_seqlen_q, @@ -206,12 +211,14 @@ def __call__( def kernel_score_grad( self, mPredictRaw, + mGradSignal, mTargetRaw, mLSE, mDenomTarget, mCuSeqlensQ, mCuSeqlensK, mQCausalOffsets, + mGradLoss, grad_scale: Float32, seqlen_k_pad: Int32, seqlen_q_static: Int32, @@ -252,11 +259,13 @@ class SharedStorage: # THD : (T_q, max_K) -> domain-offset by q_offset along T_q; LSE/Denom (T_q,) similarly. if const_expr(is_varlen): mPredict_b = cute.domain_offset((q_offset, Int32(0)), mPredictRaw) + mGradSignal_b = cute.domain_offset((q_offset, Int32(0)), mGradSignal) mTarget_b = cute.domain_offset((q_offset, Int32(0)), mTargetRaw) mLSE_b = cute.domain_offset((q_offset,), mLSE) mDenom_b = cute.domain_offset((q_offset,), mDenomTarget) else: mPredict_b = mPredictRaw[None, None, batch_idx] + mGradSignal_b = mGradSignal[None, None, batch_idx] mTarget_b = mTargetRaw[None, None, batch_idx] mLSE_b = mLSE[None, batch_idx] mDenom_b = mDenomTarget[None, batch_idx] @@ -264,6 +273,7 @@ class SharedStorage: # Scalar loads (shared across all threads in CTA) lse_val = mLSE_b[seq_local] denom_val = mDenom_b[seq_local] + effective_grad_scale = grad_scale if const_expr(mGradLoss is None) else grad_scale * mGradLoss[0] LOG2E = Float32(1.4426950408889634) @@ -284,7 +294,7 @@ class SharedStorage: target = tr / (denom_val + Float32(DENOM_EPS)) target_eff = target if target >= Float32(CLIP_PROB_MIN) else Float32(CLIP_PROB_MIN) log_clip_mask = Float32(1.0) if score_minus_lse >= Float32(CLIP_LOG_MIN) else Float32(0.0) - g = -target_eff * log_clip_mask * grad_scale + g = -target_eff * log_clip_mask * effective_grad_scale local_sum = local_sum + g pos = pos + Int32(128) @@ -324,14 +334,14 @@ class SharedStorage: target = tr / (denom_val + Float32(DENOM_EPS)) target_eff = target if target >= Float32(CLIP_PROB_MIN) else Float32(CLIP_PROB_MIN) log_clip_mask = Float32(1.0) if score_minus_lse >= Float32(CLIP_LOG_MIN) else Float32(0.0) - g = -target_eff * log_clip_mask * grad_scale + g = -target_eff * log_clip_mask * effective_grad_scale grad_signal = g - predict * sum_grad - mPredict_b[seq_local, pos] = grad_signal + mGradSignal_b[seq_local, pos] = grad_signal else: # masked / padding: grad_signal must be 0 — downstream # GEMM multiplies by this directly and must not contaminate dQ/dK/dW. - mPredict_b[seq_local, pos] = Float32(0.0) + mGradSignal_b[seq_local, pos] = Float32(0.0) pos = pos + Int32(128) @@ -1981,6 +1991,10 @@ def dense_indexer_backward_sm100( def _build_cute_dsl_kernel(batch, max_seqlen_q, max_seqlen_k, heads, dim, sm_scale, block_I, ratio, is_varlen, has_q_causal_offsets): + import torch + + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor if torch.cuda.get_device_capability()[0] < 10: @@ -2024,7 +2038,7 @@ def _ensure_compiled_score_grad( the runtime ``grad_scale`` differs from the value seen at compile. """ if compile_key not in _score_grad_compile_cache: - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) cuq_arg = to_cute_tensor(CuSeqlensQ) if CuSeqlensQ is not None else None cuk_arg = to_cute_tensor(CuSeqlensK) if CuSeqlensK is not None else None q_offsets_arg = to_cute_tensor(QCausalOffsets) if QCausalOffsets is not None else None @@ -2063,7 +2077,7 @@ def _ensure_compiled_gemm( int32 cu_seqlens tensors of shape (B+1,). """ if compile_key not in _gemm_compile_cache: - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) cuq_arg = to_cute_tensor(CuSeqlensQ) if CuSeqlensQ is not None else None cuk_arg = to_cute_tensor(CuSeqlensK) if CuSeqlensK is not None else None q_offsets_arg = to_cute_tensor(QCausalOffsets) if QCausalOffsets is not None else None @@ -2110,7 +2124,7 @@ def _run_gemm_only( assert QCausalOffsets is not None, "offset-compiled kernel requires q_causal_offsets at runtime" else: assert QCausalOffsets is None, "non-offset compiled kernel must not receive q_causal_offsets" - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) _ensure_compiled_gemm( IndexQ, Weights, @@ -2177,7 +2191,7 @@ def _run( assert QCausalOffsets is not None, "offset-compiled kernel requires q_causal_offsets at runtime" else: assert QCausalOffsets is None, "non-offset compiled kernel must not receive q_causal_offsets" - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) # Kernel 1 (CuTe DSL): in-place overwrites IdxScoreRaw with grad_signal. # Grid = (max_seqlen_q, B, 1); CTAs past per-batch seqlen exit early. diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm90.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm90.py index 5e7a18696..506bd8c76 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm90.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/dense_indexer_backward_sm90.py @@ -2,20 +2,12 @@ from __future__ import annotations -import torch import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute from cutlass import Float32, Int32, const_expr -from cudnn.deepseek_sparse_attention.utils.compiler import compile_options -from cudnn.deepseek_sparse_attention.utils.runtime import ( - resolve_stream as _resolve_stream, - torch_stream_context as _torch_stream_context, -) -from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor - from .indexer_backward_sm90 import ( CLIP_LOG_MIN, CLIP_PROB_MIN, @@ -71,7 +63,12 @@ def __call__( max_seqlen_q: Int32, max_seqlen_k: Int32, stream, + mGradSignal=None, + mGradLoss=None, ): + if const_expr(mGradSignal is None): + mGradSignal = mIdxScoreRaw + # BSHD: # IdxScoreRaw/AttnScoreRaw: (B, S_q, S_k), IdxLSE/AttnL1Norm: (B, S_q) # THD: @@ -85,6 +82,10 @@ def __call__( mIdxScoreRaw.iterator, cute.select(mIdxScoreRaw.layout, mode=[1, 2, 0]), ) + mGradSignal = cute.make_tensor( + mGradSignal.iterator, + cute.select(mGradSignal.layout, mode=[1, 2, 0]), + ) mAttnScoreRaw = cute.make_tensor( mAttnScoreRaw.iterator, cute.select(mAttnScoreRaw.layout, mode=[1, 2, 0]), @@ -103,12 +104,14 @@ def __call__( self.kernel_score_grad( mIdxScoreRaw, + mGradSignal, mAttnScoreRaw, mIdxLSE, mAttnL1Norm, mCuSeqlensQ, mCuSeqlensK, mQCausalOffsets, + mGradLoss, grad_scale, seqlen_k_pad, max_seqlen_q, @@ -123,12 +126,14 @@ def __call__( def kernel_score_grad( self, mIdxScoreRaw, + mGradSignal, mAttnScoreRaw, mIdxLSE, mAttnL1Norm, mCuSeqlensQ, mCuSeqlensK, mQCausalOffsets, + mGradLoss, grad_scale: Float32, seqlen_k_pad: Int32, seqlen_q_static: Int32, @@ -164,17 +169,20 @@ class SharedStorage: if const_expr(is_varlen): mIdxScore_b = cute.domain_offset((q_offset, Int32(0)), mIdxScoreRaw) + mGradSignal_b = cute.domain_offset((q_offset, Int32(0)), mGradSignal) mAttnScore_b = cute.domain_offset((q_offset, Int32(0)), mAttnScoreRaw) mIdxLSE_b = cute.domain_offset((q_offset,), mIdxLSE) mAttnL1_b = cute.domain_offset((q_offset,), mAttnL1Norm) else: mIdxScore_b = mIdxScoreRaw[None, None, batch_idx] + mGradSignal_b = mGradSignal[None, None, batch_idx] mAttnScore_b = mAttnScoreRaw[None, None, batch_idx] mIdxLSE_b = mIdxLSE[None, batch_idx] mAttnL1_b = mAttnL1Norm[None, batch_idx] idx_lse_val = mIdxLSE_b[seq_local] attn_l1_val = mAttnL1_b[seq_local] + effective_grad_scale = grad_scale if const_expr(mGradLoss is None) else grad_scale * mGradLoss[0] LOG2E = Float32(1.4426950408889634) ratio = Int32(self.ratio) @@ -195,7 +203,7 @@ class SharedStorage: target = attn_raw / (attn_l1_val + Float32(EPS)) target_eff = target if target >= Float32(CLIP_PROB_MIN) else Float32(CLIP_PROB_MIN) log_clip_mask = Float32(1.0) if score_minus_lse >= Float32(CLIP_LOG_MIN) else Float32(0.0) - local_sum = local_sum + (-target_eff * log_clip_mask * grad_scale) + local_sum = local_sum + (-target_eff * log_clip_mask * effective_grad_scale) pos = pos + Int32(128) warp_sum = cute.arch.warp_reduction_sum(local_sum) @@ -221,10 +229,10 @@ class SharedStorage: target = attn_raw / (attn_l1_val + Float32(EPS)) target_eff = target if target >= Float32(CLIP_PROB_MIN) else Float32(CLIP_PROB_MIN) log_clip_mask = Float32(1.0) if score_minus_lse >= Float32(CLIP_LOG_MIN) else Float32(0.0) - g = -target_eff * log_clip_mask * grad_scale - mIdxScore_b[seq_local, pos] = g - predict * sum_grad + g = -target_eff * log_clip_mask * effective_grad_scale + mGradSignal_b[seq_local, pos] = g - predict * sum_grad else: - mIdxScore_b[seq_local, pos] = Float32(0.0) + mGradSignal_b[seq_local, pos] = Float32(0.0) pos = pos + Int32(128) @@ -270,6 +278,12 @@ def _build_cute_dsl_dense_kernel( is_varlen, has_q_causal_offsets, ): + import torch + + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream, torch_stream_context + from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor + cap = torch.cuda.get_device_capability()[0] if cap < 9: raise RuntimeError(f"Requires SM90+ (got SM{cap}0)") @@ -295,7 +309,7 @@ def _build_cute_dsl_dense_kernel( def _get_dummy_topk(device, current_stream=None): if dummy_topk_holder[0] is None or dummy_topk_holder[0].device != device: - with _torch_stream_context(current_stream): + with torch_stream_context(current_stream): dummy_topk_holder[0] = torch.zeros(batch, seqlen, seqlen_k, device=device, dtype=torch.int32) return dummy_topk_holder[0] @@ -312,7 +326,7 @@ def _ensure_compiled( QCausalOffsets, current_stream=None, ): - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) if gemm_key not in _dense_compile_cache: dummy_topk = _get_dummy_topk(IndexQ.device, current_stream=current_stream) cuq_arg = to_cute_tensor(CuSeqlensQ) if CuSeqlensQ is not None else None @@ -356,7 +370,7 @@ def _ensure_compiled_score_grad( current_stream=None, ): if score_grad_key not in _dense_score_grad_compile_cache: - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) cuq_arg = to_cute_tensor(CuSeqlensQ) if CuSeqlensQ is not None else None cuk_arg = to_cute_tensor(CuSeqlensK) if CuSeqlensK is not None else None q_offsets_arg = to_cute_tensor(QCausalOffsets) if QCausalOffsets is not None else None @@ -395,7 +409,7 @@ def _run_score_grad_only( assert QCausalOffsets is not None, "offset-compiled score-grad kernel requires q_causal_offsets at runtime" else: assert QCausalOffsets is None, "non-offset compiled score-grad kernel must not receive q_causal_offsets" - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) _ensure_compiled_score_grad( IdxScoreRaw, IdxLSE, @@ -444,7 +458,7 @@ def _run_gemm_only( else: assert QCausalOffsets is None, "non-offset compiled kernel must not receive q_causal_offsets" dummy_topk = _get_dummy_topk(IndexQ.device, current_stream=current_stream) - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) _ensure_compiled( IndexQ, @@ -526,7 +540,7 @@ def _run( current_stream=current_stream, ) else: - with _torch_stream_context(current_stream): + with torch_stream_context(current_stream): dIndexK_f32 = torch.zeros_like(dIndexK, dtype=torch.float32) _run_gemm_only( IndexQ, @@ -541,7 +555,7 @@ def _run( QCausalOffsets, current_stream=current_stream, ) - with _torch_stream_context(current_stream): + with torch_stream_context(current_stream): dIndexK.copy_(dIndexK_f32) _run.score_grad = _run_score_grad_only diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py index 740c3578d..b0b96abc9 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm100.py @@ -51,7 +51,8 @@ import math from functools import partial -import torch +from typing import Optional + import cuda.bindings.driver as cuda import cutlass @@ -71,12 +72,6 @@ import cutlass.utils.blackwell_helpers as sm100_utils_basic -from cudnn.deepseek_sparse_attention.utils.compiler import compile_options -from cudnn.deepseek_sparse_attention.utils.runtime import ( - resolve_stream as _resolve_stream, - torch_stream_context as _torch_stream_context, -) - mul_packed_f32x2 = partial(cute.arch.mul_packed_f32x2, rnd="rn") fma_packed_f32x2 = partial(cute.arch.fma_packed_f32x2, rnd="rn") @@ -1453,7 +1448,7 @@ def indexer_backward_sm100( class ScoreGradSm100: - """CuTe DSL kernel for in-place score_grad precompute.""" + """CuTe DSL sparse score-gradient kernel with optional outputs.""" THREADS_PER_CTA = 128 WARP_SIZE = 32 @@ -1470,14 +1465,24 @@ def __call__( mGradLoss: cute.Tensor, grad_scale: Float32 | float, stream: cuda.CUstream, + mGradSignal: Optional[cute.Tensor] = None, + mSumGrad: Optional[cute.Tensor] = None, ): + write_sum_grad = mSumGrad is not None or mGradSignal is None + if cutlass.const_expr(mGradSignal is None): + mGradSignal = mAttnScore + if cutlass.const_expr(mSumGrad is None): + mSumGrad = mIndexScore + # (b, s, t) -> (s, t, b): topk dim contiguous for per-CTA strided loops. mAttnScore = cute.make_tensor(mAttnScore.iterator, cute.select(mAttnScore.layout, mode=[1, 2, 0])) mIndexScore = cute.make_tensor(mIndexScore.iterator, cute.select(mIndexScore.layout, mode=[1, 2, 0])) + mGradSignal = cute.make_tensor(mGradSignal.iterator, cute.select(mGradSignal.layout, mode=[1, 2, 0])) + mSumGrad = cute.make_tensor(mSumGrad.iterator, cute.select(mSumGrad.layout, mode=[1, 2, 0])) seqlen = cute.size(mAttnScore.shape[0]) batch_size = cute.size(mAttnScore.shape[2]) if cute.rank(mAttnScore.shape) > 2 else 1 - self.kernel_score_grad(mAttnScore, mIndexScore, mGradLoss, grad_scale).launch( + self.kernel_score_grad(mAttnScore, mIndexScore, mGradLoss, mGradSignal, mSumGrad, grad_scale, write_sum_grad).launch( grid=(seqlen, batch_size, 1), block=[self.THREADS_PER_CTA, 1, 1], cluster=[1, 1, 1], @@ -1486,7 +1491,9 @@ def __call__( ) @cute.kernel - def kernel_score_grad(self, mAttnScore, mIndexScore, mGradLoss, grad_scale: Float32 | float): + def kernel_score_grad( + self, mAttnScore, mIndexScore, mGradLoss, mGradSignal, mSumGrad, grad_scale: Float32 | float, write_sum_grad: cutlass.Constexpr[bool] + ): tidx = cute.arch.thread_idx()[0] seq_idx = cute.arch.block_idx()[0] batch_idx = cute.arch.block_idx()[1] @@ -1528,11 +1535,14 @@ class SharedStorage: target_eff = cute.arch.fmax(target, Float32(CLIP_PROB_MIN)) log_clip_mask = Float32(1.0) if predict >= Float32(CLIP_PROB_MIN) else Float32(0.0) g_i = -target_eff * log_clip_mask * grad_scale_f32 - mAttnScore[seq_idx, pos, batch_idx] = g_i - predict * sum_grad - mIndexScore[seq_idx, pos, batch_idx] = sum_grad + mGradSignal[seq_idx, pos, batch_idx] = g_i - predict * sum_grad + if cutlass.const_expr(write_sum_grad): + mSumGrad[seq_idx, pos, batch_idx] = sum_grad def _score_grad_inplace_cute(AttnScore, IndexScore, GradLoss, grad_scale, current_stream=None): + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor # Kernel reads ``mGradLoss[0]`` so it must be at least 1-D. ``to_cute_tensor`` @@ -1544,7 +1554,7 @@ def _score_grad_inplace_cute(AttnScore, IndexScore, GradLoss, grad_scale, curren _, _, topk = AttnScore.shape compile_key = (topk,) - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) if compile_key not in _score_grad_cute_cache: kernel_obj = ScoreGradSm100(topk=topk) _score_grad_cute_cache[compile_key] = cute.compile( @@ -1582,6 +1592,8 @@ def _score_grad_inplace(AttnScore, IndexScore, GradLoss, grad_scale, block_I=128 # target = clip(log_target) # dL/dlog_predict = -exp(target) * I(log_predict >= -100) # dL/dlogits = g - predict * sum(g) + import torch + can_use_cute = ( AttnScore.is_cuda and IndexScore.is_cuda @@ -1598,6 +1610,10 @@ def _score_grad_inplace(AttnScore, IndexScore, GradLoss, grad_scale, block_I=128 def _build_cute_dsl_kernel(heads, dim, topk, sm_scale, block_I, topk_indices_global: bool = True): + import torch + + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream, torch_stream_context from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor if torch.cuda.get_device_capability()[0] < 10: @@ -1617,7 +1633,7 @@ def _build_cute_dsl_kernel(heads, dim, topk, sm_scale, block_I, topk_indices_glo def _ensure_compiled(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, AttnScore, TopkIndices, current_stream=None): """Lazy-compile the GEMM kernel (kernel 2) on first execute (needs real tensors).""" if compile_key not in _compile_cache: - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) cute_args = [to_cute_tensor(t) for t in [IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, AttnScore, TopkIndices]] _compile_cache[compile_key] = cute.compile( kernel_obj, @@ -1629,7 +1645,7 @@ def _ensure_compiled(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, At def _run_gemm_only(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, GradSignal, TopkIndices, current_stream=None): """Run only kernel 2 (GEMM). Caller must have run kernel 1 and zeroed dIndexK_f32.""" - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) _ensure_compiled(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, GradSignal, TopkIndices, current_stream=current_stream) with torch.cuda.nvtx.range("indexer_backward_dsl_gemm"): _compile_cache[compile_key]( @@ -1660,10 +1676,10 @@ def _run(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK, AttnScore, IndexSc _run_gemm_only(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK, AttnScore, TopkIndices, current_stream=current_stream) else: # Need a separate f32 buffer for atomicAdd, then cast back to output dtype. - with _torch_stream_context(current_stream): + with torch_stream_context(current_stream): dIndexK_f32 = torch.zeros_like(dIndexK, dtype=torch.float32) _run_gemm_only(IndexQ, Weights, IndexK, dIndexQ, dWeights, dIndexK_f32, AttnScore, TopkIndices, current_stream=current_stream) - with _torch_stream_context(current_stream): + with torch_stream_context(current_stream): dIndexK.copy_(dIndexK_f32) _run.score_grad = partial(_score_grad_inplace, block_I=block_I) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm90.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm90.py index 558e8d614..23f402c2e 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm90.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/indexer_backward_sm90.py @@ -44,7 +44,6 @@ from typing import Optional import cuda.bindings.driver as cuda -import torch import cutlass import cutlass.cute as cute @@ -55,13 +54,7 @@ from cutlass.utils import LayoutEnum from cudnn.deepseek_sparse_attention.utils import copy as copy_ops -from cudnn.deepseek_sparse_attention.utils.compiler import compile_options -from cudnn.deepseek_sparse_attention.utils.runtime import ( - resolve_stream as _resolve_stream, - torch_stream_context as _torch_stream_context, -) from cudnn.deepseek_sparse_attention.utils.seqlen import seqlen_info as _seqlen_info -from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor from cudnn.deepseek_sparse_attention.utils.sm90.mma import ( gemm, gemm_zero_init, @@ -1385,7 +1378,7 @@ def indexer_backward_sm90( class ScoreGradSm90: - """CuTe DSL kernel for in-place score_grad precompute (SM90).""" + """CuTe DSL sparse score-gradient kernel with an optional output.""" THREADS_PER_CTA = 128 @@ -1401,14 +1394,19 @@ def __call__( mGradLoss: cute.Tensor, grad_scale: Float32 | float, stream: cuda.CUstream, + mGradSignal: Optional[cute.Tensor] = None, ): + if cutlass.const_expr(mGradSignal is None): + mGradSignal = mAttnScore + # (b, s, t) -> (s, t, b): contiguous topk traversal per CTA. mAttnScore = cute.make_tensor(mAttnScore.iterator, cute.select(mAttnScore.layout, mode=[1, 2, 0])) mIndexScore = cute.make_tensor(mIndexScore.iterator, cute.select(mIndexScore.layout, mode=[1, 2, 0])) + mGradSignal = cute.make_tensor(mGradSignal.iterator, cute.select(mGradSignal.layout, mode=[1, 2, 0])) seqlen = cute.size(mAttnScore.shape[0]) batch_size = cute.size(mAttnScore.shape[2]) if cute.rank(mAttnScore.shape) > 2 else 1 - self.kernel_score_grad(mAttnScore, mIndexScore, mGradLoss, grad_scale).launch( + self.kernel_score_grad(mAttnScore, mIndexScore, mGradLoss, mGradSignal, grad_scale).launch( grid=(seqlen, batch_size, 1), block=[self.THREADS_PER_CTA, 1, 1], cluster=[1, 1, 1], @@ -1417,7 +1415,7 @@ def __call__( ) @cute.kernel - def kernel_score_grad(self, mAttnScore, mIndexScore, mGradLoss, grad_scale: Float32 | float): + def kernel_score_grad(self, mAttnScore, mIndexScore, mGradLoss, mGradSignal, grad_scale: Float32 | float): tidx = cute.arch.thread_idx()[0] seq_idx = cute.arch.block_idx()[0] batch_idx = cute.arch.block_idx()[1] @@ -1473,7 +1471,7 @@ class SharedStorage: predict = Float32(mIndexScore[seq_idx, pos, batch_idx]) log_clip_mask = Float32(1.0) if predict >= Float32(CLIP_PROB_MIN) else Float32(0.0) g_i = -target_eff * log_clip_mask * grad_scale_f32 - mAttnScore[seq_idx, pos, batch_idx] = g_i - predict * sum_grad + mGradSignal[seq_idx, pos, batch_idx] = g_i - predict * sum_grad def _score_grad_inplace_cute( @@ -1484,6 +1482,10 @@ def _score_grad_inplace_cute( index_is_log: bool = False, current_stream=None, ): + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream + from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor + # Kernel reads ``mGradLoss[0]`` so it must be at least 1-D; defend direct # factory callers passing a 0-D scalar (the public wrapper reshapes already). if GradLoss.ndim == 0: @@ -1491,7 +1493,7 @@ def _score_grad_inplace_cute( _, _, topk = AttnScore.shape compile_key = (topk, bool(index_is_log)) - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) if compile_key not in _score_grad_cute_cache: kernel_obj = ScoreGradSm90( topk=topk, @@ -1538,6 +1540,8 @@ def _score_grad_inplace( # target = clip(log_target) # dL/dlog_predict = -exp(target) * I(log_predict >= -100) # dL/dlogits = g - predict * sum(g) + import torch + can_use_cute = ( AttnScore.is_cuda and IndexScore.is_cuda @@ -1571,6 +1575,12 @@ def _build_cute_dsl_kernel( score_input_is_log=True, topk_indices_global: bool = True, ): + import torch + + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream, torch_stream_context + from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor + cap = torch.cuda.get_device_capability()[0] if cap < 9: raise RuntimeError(f"Requires SM90+ (got SM{cap}0)") @@ -1603,7 +1613,7 @@ def _ensure_compiled( current_stream=None, ): """Lazy-compile the GEMM kernel (kernel 2) on first execute (needs real tensors).""" - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) if compile_key not in _compile_cache: cute_args = [ to_cute_tensor(t) @@ -1648,7 +1658,7 @@ def _run_gemm_only( current_stream=None, ): """Run only kernel 2 (GEMM). Caller must have run kernel 1 and zeroed dIndexK_f32.""" - s = _resolve_stream(current_stream) + s = resolve_stream(current_stream) _ensure_compiled( IndexQ, Weights, @@ -1719,7 +1729,7 @@ def _run( ) else: # Need separate f32 buffer for atomicAdd, then convert back - with _torch_stream_context(current_stream): + with torch_stream_context(current_stream): dIndexK_f32 = torch.zeros_like(dIndexK, dtype=torch.float32) _run_gemm_only( IndexQ, @@ -1732,7 +1742,7 @@ def _run( TopkIndices, current_stream=current_stream, ) - with _torch_stream_context(current_stream): + with torch_stream_context(current_stream): dIndexK.copy_(dIndexK_f32) _run.score_grad = partial(_score_grad_inplace, index_is_log=score_input_is_log) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py new file mode 100644 index 000000000..ec1e58cf6 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/jax.py @@ -0,0 +1,952 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX APIs for sparse and dense DeepSeek indexer backward.""" + +from __future__ import annotations + +import math +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..._jax.compiler import compile_options_for_target +from ..._jax.layout import mode_from_layout +from .op import DEFAULT_BLOCK_I, DenseIndexerBackwardOp, IndexerBackwardOp + +_SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) +_SUPPORTED_COMPUTE_CAPABILITY_FAMILIES = (90, 100) + + +def _resolve_layout( + name: str, + layout: str | None, + *, + default: str, + kernel_axes: str, + supported: tuple[str, ...], +) -> tuple[str, tuple[int, ...]]: + layout = default if layout is None else layout + if layout not in supported: + choices = ", ".join(repr(value) for value in supported) + raise ValueError(f"{name} must be one of ({choices}), got {layout!r}") + return layout, mode_from_layout(layout, kernel_axes=kernel_axes) + + +def _grad_loss_operand(value: Any) -> Any: + value = jnp.asarray(value, dtype=jnp.float32) + if value.shape == (): + return jnp.reshape(value, (1,)) + if value.shape != (1,): + raise ValueError(f"grad_loss must be scalar or shape (1,), got {value.shape}") + return value + + +def _finite_float(value: float, name: str) -> float: + try: + result = float(value) + except (TypeError, ValueError) as error: + raise TypeError(f"{name} must be a real scalar, got {value!r}") from error + if not math.isfinite(result): + raise ValueError(f"{name} must be finite, got {result}") + return result + + +class _IndexerJaxBase(JaxApiBase): + target_compute_capability: int | None + compute_capability: int | None + + def _resolve_target(self, operation_name: str) -> None: + self.compute_capability = self._resolve_compute_capability( + self.target_compute_capability, + _SUPPORTED_COMPUTE_CAPABILITIES, + operation_name, + ) + + @property + def _architecture_family(self) -> int: + if self.compute_capability is None: + raise RuntimeError( + "check_support() must resolve the compute capability before lowering" + ) + family = self._compute_capability_family( + self.compute_capability, _SUPPORTED_COMPUTE_CAPABILITY_FAMILIES + ) + if family is None: + raise RuntimeError( + f"No indexer-backward kernel for SM{self.compute_capability}" + ) + return family + + @staticmethod + def _output_desc( + sample: Any | None, + *, + source: JaxTensorDesc, + cudnn_dtype: data_type, + shape: tuple[int, ...], + name: str, + mode: tuple[int, ...], + ) -> JaxTensorDesc: + if sample is not None: + return JaxApiBase._to_tensor_desc(sample, name, mode=mode) + return source.compact_like( + cudnn_dtype=cudnn_dtype, + shape=shape, + stride_order=source.stride_order, + name=name, + mode=mode, + ) + + +class IndexerBackward(_IndexerJaxBase): + """Sparse indexer backward for BSHD or packed THD JAX arrays. + + Packed THD uses ``Q=(T_q,H,D)``, ``W=(T_q,H)``, ``K=(T_k,D)``, and + score/index tensors shaped ``(T_q,topk)``. Packed top-K indices must be + global flat K indices because this API intentionally has no host-visible + sequence metadata from which to reconstruct per-batch offsets. + """ + + def __init__( + self, + sample_index_q: Any, + sample_weights: Any, + sample_index_k: Any, + sample_attn_score: Any, + sample_index_score: Any, + sample_topk_indices: Any, + sample_d_index_q: Any | None = None, + sample_d_weights: Any | None = None, + sample_d_index_k: Any | None = None, + sm_scale: float = 1.0, + loss_coeff: float = 1.0, + block_I: int = DEFAULT_BLOCK_I, + topk_indices_global: bool = False, + q_layout: str | None = None, + w_layout: str | None = None, + k_layout: str | None = None, + score_layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + is_thd = len(tuple(sample_index_q.shape)) == 3 + if is_thd: + q_defaults = ("THD", "THD", ("THD",)) + w_defaults = ("TH", "TH", ("TH",)) + k_defaults = ("TD", "TD", ("TD",)) + score_defaults = ("TK", "TK", ("TK",)) + else: + q_defaults = ("BSHD", "BSHD", ("BSHD", "SBHD")) + w_defaults = ("BSH", "BSH", ("BSH", "SBH")) + k_defaults = ("BSD", "BSD", ("BSD", "SBD")) + score_defaults = ("BSK", "BSK", ("BSK", "SBK")) + self.q_layout, self.q_mode = _resolve_layout( + "q_layout", + q_layout, + default=q_defaults[0], + kernel_axes=q_defaults[1], + supported=q_defaults[2], + ) + self.w_layout, self.w_mode = _resolve_layout( + "w_layout", + w_layout, + default=w_defaults[0], + kernel_axes=w_defaults[1], + supported=w_defaults[2], + ) + self.k_layout, self.k_mode = _resolve_layout( + "k_layout", + k_layout, + default=k_defaults[0], + kernel_axes=k_defaults[1], + supported=k_defaults[2], + ) + self.score_layout, self.score_mode = _resolve_layout( + "score_layout", + score_layout, + default=score_defaults[0], + kernel_axes=score_defaults[1], + supported=score_defaults[2], + ) + + self.iq_desc = self._to_tensor_desc( + sample_index_q, "sample_index_q", mode=self.q_mode + ) + self.w_desc = self._to_tensor_desc( + sample_weights, "sample_weights", mode=self.w_mode + ) + self.ik_desc = self._to_tensor_desc( + sample_index_k, "sample_index_k", mode=self.k_mode + ) + self.attn_desc = self._to_tensor_desc( + sample_attn_score, "sample_attn_score", mode=self.score_mode + ) + self.idx_score_desc = self._to_tensor_desc( + sample_index_score, "sample_index_score", mode=self.score_mode + ) + self.topk_desc = self._to_tensor_desc( + sample_topk_indices, "sample_topk_indices", mode=self.score_mode + ) + self.diq_desc = self._output_desc( + sample_d_index_q, + source=self.iq_desc, + cudnn_dtype=data_type.BFLOAT16, + shape=self.iq_desc.shape, + name="sample_d_index_q", + mode=self.q_mode, + ) + self.dw_desc = self._output_desc( + sample_d_weights, + source=self.w_desc, + cudnn_dtype=data_type.BFLOAT16, + shape=self.w_desc.shape, + name="sample_d_weights", + mode=self.w_mode, + ) + self.dik_desc = self._output_desc( + sample_d_index_k, + source=self.ik_desc, + cudnn_dtype=data_type.BFLOAT16, + shape=self.ik_desc.shape, + name="sample_d_index_k", + mode=self.k_mode, + ) + self.grad_loss_desc = self.iq_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(1,), + name="grad_loss", + ) + self._op = IndexerBackwardOp( + index_q=self.iq_desc, + weights=self.w_desc, + index_k=self.ik_desc, + d_index_q=self.diq_desc, + d_weights=self.dw_desc, + d_index_k=self.dik_desc, + attn_score=self.attn_desc, + index_score=self.idx_score_desc, + topk_indices=self.topk_desc, + sm_scale=sm_scale, + block_i=block_I, + topk_indices_global=topk_indices_global, + ) + self.loss_coeff = _finite_float(loss_coeff, "loss_coeff") + self.target_compute_capability = target_compute_capability + self.compute_capability = None + + def check_support(self) -> bool: + self._op.check_support() + self._resolve_target("IndexerBackward") + return True + + def __call__( + self, + index_q: Any, + weights: Any, + index_k: Any, + attn_score: Any, + index_score: Any, + topk_indices: Any, + grad_loss: Any = 1.0, + ) -> TupleDict: + self.check_support() + grad_loss = _grad_loss_operand(grad_loss) + + dk_accum_desc = self.ik_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=self.ik_desc.shape, + stride_order=self.ik_desc.stride_order, + name="d_index_k_accum", + init_value=0.0, + mode=self.k_mode, + ) + grad_signal_desc = self.attn_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=self.attn_desc.shape, + stride_order=self.attn_desc.stride_order, + name="grad_signal", + mode=self.score_mode, + ) + head_dim = self._op.head_dim + iq_desc = self.iq_desc.with_divisibility( + (None,) * (self.iq_desc.ndim - 1) + (head_dim,) + ) + ik_desc = self.ik_desc.with_divisibility( + (None,) * (self.ik_desc.ndim - 1) + (head_dim,) + ) + diq_desc = self.diq_desc.with_divisibility( + (None,) * (self.diq_desc.ndim - 1) + (head_dim,) + ) + dk_accum_desc = dk_accum_desc.with_divisibility( + (None,) * (dk_accum_desc.ndim - 1) + (head_dim,) + ) + d_index_q, d_weights, d_index_k_accum = self._call_kernel( + ( + index_q, + weights, + index_k, + attn_score, + index_score, + topk_indices, + grad_loss, + ), + launch=self._launch_kernel, + output_descs=(diq_desc, self.dw_desc, dk_accum_desc), + input_descs=( + iq_desc, + self.w_desc, + ik_desc, + self.attn_desc, + self.idx_score_desc, + self.topk_desc, + self.grad_loss_desc, + ), + workspace_descs=(grad_signal_desc,), + compile_options=compile_options_for_target( + self.compute_capability, "--opt-level 3" + ), + ) + return TupleDict( + d_index_q=d_index_q, + d_weights=d_weights, + d_index_k=d_index_k_accum.astype(self.dik_desc.dtype), + ) + + def _launch_kernel( + self, + stream, + index_q, + weights, + index_k, + attn_score, + index_score, + topk_indices, + grad_loss, + d_index_q, + d_weights, + d_index_k_accum, + grad_signal, + ) -> None: + import cutlass + + grad_scale = self.loss_coeff / (self._op.batch * self._op.seqlen_q) + + if self._op.is_thd: + index_q = _prepend_unit_batch(index_q) + weights = _prepend_unit_batch(weights) + index_k = _prepend_unit_batch(index_k) + attn_score = _prepend_unit_batch(attn_score) + index_score = _prepend_unit_batch(index_score) + topk_indices = _prepend_unit_batch(topk_indices) + d_index_q = _prepend_unit_batch(d_index_q) + d_weights = _prepend_unit_batch(d_weights) + d_index_k_accum = _prepend_unit_batch(d_index_k_accum) + grad_signal = _prepend_unit_batch(grad_signal) + + if self._architecture_family == 100: + from .indexer_backward_sm100 import IndexerBackwardSm100, ScoreGradSm100 + + ScoreGradSm100(topk=self._op.topk)( + attn_score, + index_score, + grad_loss, + cutlass.Float32(grad_scale), + stream, + grad_signal, + None, + ) + IndexerBackwardSm100( + head_dim=self._op.head_dim, + heads=self._op.heads, + block_I=self._op.block_i, + topk=self._op.topk, + topk_indices_global=self._op.topk_indices_global, + )( + index_q, + weights, + index_k, + d_index_q, + d_weights, + d_index_k_accum, + grad_signal, + topk_indices, + cutlass.Float32(self._op.sm_scale), + stream, + ) + return + + from .indexer_backward_sm90 import IndexerBackwardSm90, ScoreGradSm90 + + ScoreGradSm90(topk=self._op.topk, index_is_log=False)( + attn_score, + index_score, + grad_loss, + cutlass.Float32(grad_scale), + stream, + grad_signal, + ) + IndexerBackwardSm90( + head_dim=self._op.head_dim, + heads=self._op.heads, + block_I=self._op.block_i, + topk=self._op.topk, + topk_indices_global=self._op.topk_indices_global, + )( + index_q, + weights, + index_k, + d_index_q, + d_weights, + d_index_k_accum, + grad_signal, + topk_indices, + cutlass.Float32(self._op.sm_scale), + stream, + None, + None, + cutlass.Int32(self._op.seqlen_q), + cutlass.Int32(self._op.seqlen_k), + None, + ) + + +def _prepend_unit_batch(tensor: Any) -> Any: + """Expose a compact packed tensor to the sparse kernels as batch one.""" + + import cutlass.cute as cute + + leading_stride = tensor.shape[0] * tensor.stride[0] + return cute.make_tensor( + tensor.iterator, + cute.make_layout( + (1, *tensor.shape), + stride=(leading_stride, *tensor.stride), + ), + ) + + +class DenseIndexerBackward(_IndexerJaxBase): + """Sample-signature-bound dense BSHD or packed-THD backward callable.""" + + def __init__( + self, + sample_index_q: Any, + sample_weights: Any, + sample_index_k: Any, + sample_attn_score: Any, + sample_attn_l1norm: Any, + sample_index_score: Any, + sample_index_lse: Any, + sample_d_index_q: Any | None = None, + sample_d_weights: Any | None = None, + sample_d_index_k: Any | None = None, + sample_cu_seqlens_q: Any | None = None, + sample_cu_seqlens_k: Any | None = None, + sample_q_causal_offsets: Any | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + sm_scale: float = 1.0, + loss_coeff: float = 1.0, + block_I: int = DEFAULT_BLOCK_I, + ratio: int = 1, + q_layout: str | None = None, + w_layout: str | None = None, + k_layout: str | None = None, + score_layout: str | None = None, + denom_layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + is_thd = len(tuple(sample_index_q.shape)) == 3 + if is_thd: + q_defaults = ("THD", "THD", ("THD",)) + w_defaults = ("TH", "TH", ("TH",)) + k_defaults = ("TD", "TD", ("TD",)) + score_defaults = ("TK", "TK", ("TK",)) + denom_defaults = ("T", "T", ("T",)) + else: + q_defaults = ("BSHD", "BSHD", ("BSHD", "SBHD")) + w_defaults = ("BSH", "BSH", ("BSH", "SBH")) + k_defaults = ("BSD", "BSD", ("BSD", "SBD")) + score_defaults = ("BSK", "BSK", ("BSK", "SBK")) + denom_defaults = ("BS", "BS", ("BS", "SB")) + self.q_layout, self.q_mode = _resolve_layout( + "q_layout", + q_layout, + default=q_defaults[0], + kernel_axes=q_defaults[1], + supported=q_defaults[2], + ) + self.w_layout, self.w_mode = _resolve_layout( + "w_layout", + w_layout, + default=w_defaults[0], + kernel_axes=w_defaults[1], + supported=w_defaults[2], + ) + self.k_layout, self.k_mode = _resolve_layout( + "k_layout", + k_layout, + default=k_defaults[0], + kernel_axes=k_defaults[1], + supported=k_defaults[2], + ) + self.score_layout, self.score_mode = _resolve_layout( + "score_layout", + score_layout, + default=score_defaults[0], + kernel_axes=score_defaults[1], + supported=score_defaults[2], + ) + self.denom_layout, self.denom_mode = _resolve_layout( + "denom_layout", + denom_layout, + default=denom_defaults[0], + kernel_axes=denom_defaults[1], + supported=denom_defaults[2], + ) + + self.iq_desc = self._to_tensor_desc( + sample_index_q, "sample_index_q", mode=self.q_mode + ) + self.w_desc = self._to_tensor_desc( + sample_weights, "sample_weights", mode=self.w_mode + ) + self.ik_desc = self._to_tensor_desc( + sample_index_k, "sample_index_k", mode=self.k_mode + ) + self.attn_desc = self._to_tensor_desc( + sample_attn_score, "sample_attn_score", mode=self.score_mode + ) + self.attn_denom_desc = self._to_tensor_desc( + sample_attn_l1norm, + "sample_attn_l1norm", + mode=self.denom_mode, + ) + self.idx_score_desc = self._to_tensor_desc( + sample_index_score, "sample_index_score", mode=self.score_mode + ) + self.idx_lse_desc = self._to_tensor_desc( + sample_index_lse, "sample_index_lse", mode=self.denom_mode + ) + self.diq_desc = self._output_desc( + sample_d_index_q, + source=self.iq_desc, + cudnn_dtype=data_type.BFLOAT16, + shape=self.iq_desc.shape, + name="sample_d_index_q", + mode=self.q_mode, + ) + self.dw_desc = self._output_desc( + sample_d_weights, + source=self.w_desc, + cudnn_dtype=data_type.BFLOAT16, + shape=self.w_desc.shape, + name="sample_d_weights", + mode=self.w_mode, + ) + self.dik_desc = self._output_desc( + sample_d_index_k, + source=self.ik_desc, + cudnn_dtype=data_type.BFLOAT16, + shape=self.ik_desc.shape, + name="sample_d_index_k", + mode=self.k_mode, + ) + self.grad_loss_desc = self.iq_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(1,), + name="grad_loss", + ) + self.cuq_desc = ( + None + if sample_cu_seqlens_q is None + else self._to_tensor_desc(sample_cu_seqlens_q, "sample_cu_seqlens_q") + ) + self.cuk_desc = ( + None + if sample_cu_seqlens_k is None + else self._to_tensor_desc(sample_cu_seqlens_k, "sample_cu_seqlens_k") + ) + self.q_offsets_desc = ( + None + if sample_q_causal_offsets is None + else self._to_tensor_desc( + sample_q_causal_offsets, "sample_q_causal_offsets" + ) + ) + self._op = DenseIndexerBackwardOp( + index_q=self.iq_desc, + weights=self.w_desc, + index_k=self.ik_desc, + d_index_q=self.diq_desc, + d_weights=self.dw_desc, + d_index_k=self.dik_desc, + attn_score=self.attn_desc, + attn_l1norm=self.attn_denom_desc, + index_score=self.idx_score_desc, + index_lse=self.idx_lse_desc, + cu_seqlens_q=self.cuq_desc, + cu_seqlens_k=self.cuk_desc, + q_causal_offsets=self.q_offsets_desc, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + sm_scale=sm_scale, + block_i=block_I, + ratio=ratio, + ) + self.loss_coeff = _finite_float(loss_coeff, "loss_coeff") + self.target_compute_capability = target_compute_capability + self.compute_capability = None + + def check_support(self) -> bool: + self._op.check_support() + self._resolve_target("DenseIndexerBackward") + return True + + def __call__( + self, + index_q: Any, + weights: Any, + index_k: Any, + attn_score: Any, + attn_l1norm: Any, + index_score: Any, + index_lse: Any, + grad_loss: Any = 1.0, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + q_causal_offsets: Any | None = None, + ) -> TupleDict: + self.check_support() + for value, expected, name in ( + (cu_seqlens_q, self.cuq_desc, "cu_seqlens_q"), + (cu_seqlens_k, self.cuk_desc, "cu_seqlens_k"), + (q_causal_offsets, self.q_offsets_desc, "q_causal_offsets"), + ): + if (value is None) != (expected is None): + raise ValueError(f"{name} presence must match its sample") + grad_loss = _grad_loss_operand(grad_loss) + + inputs = ( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + grad_loss, + ) + input_descs = ( + self.iq_desc, + self.w_desc, + self.ik_desc, + self.attn_desc, + self.attn_denom_desc, + self.idx_score_desc, + self.idx_lse_desc, + self.grad_loss_desc, + ) + for optional, desc in ( + (cu_seqlens_q, self.cuq_desc), + (cu_seqlens_k, self.cuk_desc), + (q_causal_offsets, self.q_offsets_desc), + ): + if desc is not None: + inputs += (optional,) + input_descs += (desc,) + dk_accum_desc = self.ik_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=self.ik_desc.shape, + stride_order=self.ik_desc.stride_order, + name="d_index_k_accum", + init_value=0.0, + mode=self.k_mode, + ) + grad_signal_desc = self.idx_score_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=self.idx_score_desc.shape, + stride_order=self.idx_score_desc.stride_order, + name="grad_signal", + mode=self.score_mode, + ) + head_dim = self._op.head_dim + iq_desc = self.iq_desc.with_divisibility( + (None,) * (self.iq_desc.ndim - 1) + (head_dim,) + ) + ik_desc = self.ik_desc.with_divisibility( + (None,) * (self.ik_desc.ndim - 1) + (head_dim,) + ) + diq_desc = self.diq_desc.with_divisibility( + (None,) * (self.diq_desc.ndim - 1) + (head_dim,) + ) + dk_accum_desc = dk_accum_desc.with_divisibility( + (None,) * (dk_accum_desc.ndim - 1) + (head_dim,) + ) + input_descs = ( + iq_desc, + self.w_desc, + ik_desc, + *input_descs[3:], + ) + d_index_q, d_weights, d_index_k_accum = self._call_kernel( + inputs, + launch=self._launch_kernel, + output_descs=(diq_desc, self.dw_desc, dk_accum_desc), + input_descs=input_descs, + workspace_descs=(grad_signal_desc,), + compile_options=compile_options_for_target( + self.compute_capability, "--opt-level 3" + ), + ) + return TupleDict( + d_index_q=d_index_q, + d_weights=d_weights, + d_index_k=d_index_k_accum.astype(self.dik_desc.dtype), + ) + + def _launch_kernel(self, stream, *arguments) -> None: + import cutlass + + *inputs, d_index_q, d_weights, d_index_k_accum, grad_signal = arguments + ( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + grad_loss, + *optional, + ) = inputs + cursor = 0 + cu_seqlens_q = cu_seqlens_k = q_causal_offsets = None + if self._op.is_thd: + cu_seqlens_q, cu_seqlens_k = optional[:2] + cursor = 2 + if self._op.q_causal_offsets is not None: + q_causal_offsets = optional[cursor] + + grad_scale = self.loss_coeff / max(self._op.normalization_tokens, 1) + + if self._architecture_family == 100: + from .dense_indexer_backward_sm100 import ( + DenseIndexerBackward2QGemmSm100, + ScoreGradDense, + ) + + ScoreGradDense(ratio=self._op.ratio, block_I=self._op.block_i)( + index_score, + attn_score, + index_lse, + attn_l1norm, + cu_seqlens_q, + cu_seqlens_k, + q_causal_offsets, + cutlass.Float32(grad_scale), + cutlass.Int32(self._op.max_seqlen_q), + cutlass.Int32(self._op.max_seqlen_k), + stream, + grad_signal, + grad_loss, + ) + DenseIndexerBackward2QGemmSm100( + head_dim=self._op.head_dim, + heads=self._op.heads, + block_I=self._op.block_i, + ratio=self._op.ratio, + )( + index_q, + weights, + index_k, + d_index_q, + d_weights, + d_index_k_accum, + grad_signal, + cu_seqlens_q, + cu_seqlens_k, + q_causal_offsets, + cutlass.Float32(self._op.sm_scale), + cutlass.Int32(self._op.max_seqlen_q), + cutlass.Int32(self._op.max_seqlen_k), + stream, + ) + return + + from .dense_indexer_backward_sm90 import ScoreGradDenseSm90 + from .indexer_backward_sm90 import IndexerBackwardSm90 + + ScoreGradDenseSm90(ratio=self._op.ratio, block_I=self._op.block_i)( + index_score, + attn_score, + index_lse, + attn_l1norm, + cu_seqlens_q, + cu_seqlens_k, + q_causal_offsets, + cutlass.Float32(grad_scale), + cutlass.Int32(self._op.max_seqlen_q), + cutlass.Int32(self._op.max_seqlen_k), + stream, + grad_signal, + grad_loss, + ) + IndexerBackwardSm90( + head_dim=self._op.head_dim, + heads=self._op.heads, + block_I=self._op.block_i, + topk=self._op.max_seqlen_k, + is_dense=True, + ratio=self._op.ratio, + )( + index_q, + weights, + index_k, + d_index_q, + d_weights, + d_index_k_accum, + grad_signal, + None, + cutlass.Float32(self._op.sm_scale), + stream, + cu_seqlens_q, + cu_seqlens_k, + cutlass.Int32(self._op.max_seqlen_q), + cutlass.Int32(self._op.max_seqlen_k), + q_causal_offsets, + ) + + +@jax.jit( + static_argnames=( + "sm_scale", + "loss_coeff", + "block_I", + "topk_indices_global", + "q_layout", + "w_layout", + "k_layout", + "score_layout", + "target_compute_capability", + ) +) +def indexer_backward_wrapper( + index_q: Any, + weights: Any, + index_k: Any, + attn_score: Any, + index_score: Any, + topk_indices: Any, + sm_scale: float = 1.0, + loss_coeff: float = 1.0, + grad_loss: Any = 1.0, + block_I: int = DEFAULT_BLOCK_I, + topk_indices_global: bool = False, + q_layout: str | None = None, + w_layout: str | None = None, + k_layout: str | None = None, + score_layout: str | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute sparse indexer gradients for fixed BSHD or packed THD inputs. + + Packed THD callers must pass ``topk_indices_global=True``. + """ + + return IndexerBackward( + index_q, + weights, + index_k, + attn_score, + index_score, + topk_indices, + sm_scale=sm_scale, + loss_coeff=loss_coeff, + block_I=block_I, + topk_indices_global=topk_indices_global, + q_layout=q_layout, + w_layout=w_layout, + k_layout=k_layout, + score_layout=score_layout, + target_compute_capability=target_compute_capability, + )(index_q, weights, index_k, attn_score, index_score, topk_indices, grad_loss) + + +@jax.jit( + static_argnames=( + "sm_scale", + "loss_coeff", + "block_I", + "ratio", + "max_seqlen_q", + "max_seqlen_k", + "q_layout", + "w_layout", + "k_layout", + "score_layout", + "denom_layout", + "target_compute_capability", + ) +) +def dense_indexer_backward_wrapper( + index_q: Any, + weights: Any, + index_k: Any, + attn_score: Any, + attn_l1norm: Any, + index_score: Any, + index_lse: Any, + sm_scale: float = 1.0, + loss_coeff: float = 1.0, + grad_loss: Any = 1.0, + block_I: int = DEFAULT_BLOCK_I, + ratio: int = 1, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + q_causal_offsets: Any | None = None, + q_layout: str | None = None, + w_layout: str | None = None, + k_layout: str | None = None, + score_layout: str | None = None, + denom_layout: str | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + core = (index_q, weights, index_k, attn_score, attn_l1norm, index_score, index_lse) + return DenseIndexerBackward( + *core, + sample_cu_seqlens_q=cu_seqlens_q, + sample_cu_seqlens_k=cu_seqlens_k, + sample_q_causal_offsets=q_causal_offsets, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + sm_scale=sm_scale, + loss_coeff=loss_coeff, + block_I=block_I, + ratio=ratio, + q_layout=q_layout, + w_layout=w_layout, + k_layout=k_layout, + score_layout=score_layout, + denom_layout=denom_layout, + target_compute_capability=target_compute_capability, + )( + *core, + grad_loss=grad_loss, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + q_causal_offsets=q_causal_offsets, + ) + + +__all__ = [ + "DenseIndexerBackward", + "IndexerBackward", + "dense_indexer_backward_wrapper", + "indexer_backward_wrapper", +] diff --git a/python/cudnn/deepseek_sparse_attention/indexer_backward/op.py b/python/cudnn/deepseek_sparse_attention/indexer_backward/op.py new file mode 100644 index 000000000..a9ba6f9cf --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_backward/op.py @@ -0,0 +1,536 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral DeepSeek sparse and dense indexer backward operations.""" + +from __future__ import annotations + +import math +from typing import Any, Optional + +from ... import data_type +from ...common.op import Op +from ...common.tensor_desc import TensorDesc + +DEFAULT_BLOCK_I = 128 +SUPPORTED_HEAD_DIM = 128 +MIN_HEADS = 64 + + +def _require_desc( + name: str, desc: TensorDesc[Any] | None, *, optional: bool = False +) -> None: + if desc is None and optional: + return + if not isinstance(desc, TensorDesc): + expected = "a TensorDesc or None" if optional else "a TensorDesc" + raise TypeError(f"{name} must be {expected}, got {type(desc).__name__}") + + +def _require_rank(desc: TensorDesc[Any], rank: int, name: str) -> None: + if desc.ndim != rank: + raise ValueError(f"{name} must have rank {rank}, got shape {desc.shape}") + + +def _require_dtype( + desc: TensorDesc[Any], expected: data_type | tuple[data_type, ...], name: str +) -> None: + expected_values = expected if isinstance(expected, tuple) else (expected,) + if desc.cudnn_dtype not in expected_values: + expected_text = " or ".join(value.name.lower() for value in expected_values) + raise ValueError(f"{name} must have dtype {expected_text}, got {desc.dtype}") + + +def _require_shape(desc: TensorDesc[Any], expected: tuple[int, ...], name: str) -> None: + if desc.shape != expected: + raise ValueError(f"{name} must have shape {expected}, got {desc.shape}") + + +def _require_compact(desc: TensorDesc[Any]) -> None: + if not desc.is_compact(): + raise ValueError( + f"{desc.name or 'tensor'} must be compact, got stride {desc.stride}" + ) + + +def _require_contiguous_tail(desc: TensorDesc[Any], rank: int) -> None: + expected_stride = 1 + for axis in range(desc.ndim - 1, desc.ndim - rank - 1, -1): + if desc.shape[axis] != 1 and desc.stride[axis] != expected_stride: + raise ValueError( + f"{desc.name or 'tensor'} must have its final {rank} axes contiguous, " + f"got shape {desc.shape} and stride {desc.stride}" + ) + expected_stride *= max(desc.shape[axis], 1) + + +def _resolve_scale(value: float, name: str) -> float: + try: + resolved = float(value) + except (TypeError, ValueError) as error: + raise TypeError(f"{name} must be a real scalar, got {value!r}") from error + if not math.isfinite(resolved): + raise ValueError(f"{name} must be finite, got {resolved}") + return resolved + + +class IndexerBackwardOp(Op): + """Complete fixed BSHD or packed THD sparse backward signature.""" + + def __init__( + self, + *, + index_q: TensorDesc[Any], + weights: TensorDesc[Any], + index_k: TensorDesc[Any], + d_index_q: TensorDesc[Any], + d_weights: TensorDesc[Any], + d_index_k: TensorDesc[Any], + attn_score: TensorDesc[Any], + index_score: TensorDesc[Any], + topk_indices: TensorDesc[Any], + sm_scale: float = 1.0, + block_i: int = DEFAULT_BLOCK_I, + topk_indices_global: bool = False, + ) -> None: + descriptors = ( + ("index_q", index_q), + ("weights", weights), + ("index_k", index_k), + ("d_index_q", d_index_q), + ("d_weights", d_weights), + ("d_index_k", d_index_k), + ("attn_score", attn_score), + ("index_score", index_score), + ("topk_indices", topk_indices), + ) + for name, desc in descriptors: + _require_desc(name, desc) + + self.index_q = index_q + self.weights = weights + self.index_k = index_k + self.d_index_q = d_index_q + self.d_weights = d_weights + self.d_index_k = d_index_k + self.attn_score = attn_score + self.index_score = index_score + self.topk_indices = topk_indices + self.requested_sm_scale = sm_scale + self.block_i = int(block_i) + self.topk_indices_global = bool(topk_indices_global) + + self.is_thd: Optional[bool] = None + self.batch: Optional[int] = None + self.seqlen_q: Optional[int] = None + self.seqlen_k: Optional[int] = None + self.heads: Optional[int] = None + self.head_dim: Optional[int] = None + self.topk: Optional[int] = None + self.sm_scale: Optional[float] = None + + def check_support(self) -> bool: + self.is_thd = None + self.batch = self.seqlen_q = self.seqlen_k = self.heads = self.head_dim = ( + self.topk + ) = None + self.sm_scale = None + + if self.index_q.ndim not in (3, 4): + raise ValueError( + f"index_q must use BSHD rank 4 or packed THD rank 3, got shape {self.index_q.shape}" + ) + is_thd = self.index_q.ndim == 3 + q_rank = 3 if is_thd else 4 + auxiliary_rank = q_rank - 1 + for desc, rank, name in ( + (self.index_q, q_rank, "index_q"), + (self.weights, auxiliary_rank, "weights"), + (self.index_k, auxiliary_rank, "index_k"), + (self.d_index_q, q_rank, "d_index_q"), + (self.d_weights, auxiliary_rank, "d_weights"), + (self.d_index_k, auxiliary_rank, "d_index_k"), + (self.attn_score, auxiliary_rank, "attn_score"), + (self.index_score, auxiliary_rank, "index_score"), + (self.topk_indices, auxiliary_rank, "topk_indices"), + ): + _require_rank(desc, rank, name) + + for desc, name in ( + (self.index_q, "index_q"), + (self.weights, "weights"), + (self.index_k, "index_k"), + ): + _require_dtype(desc, data_type.BFLOAT16, name) + _require_dtype(self.d_index_q, data_type.BFLOAT16, "d_index_q") + _require_dtype(self.d_weights, data_type.BFLOAT16, "d_weights") + _require_dtype( + self.d_index_k, (data_type.BFLOAT16, data_type.FLOAT), "d_index_k" + ) + _require_dtype(self.attn_score, data_type.FLOAT, "attn_score") + _require_dtype(self.index_score, data_type.FLOAT, "index_score") + _require_dtype(self.topk_indices, data_type.INT32, "topk_indices") + + if is_thd: + seqlen_q, heads, head_dim = self.index_q.shape + seqlen_k, head_dim_k = self.index_k.shape + batch = batch_k = 1 + if not self.topk_indices_global: + raise ValueError( + "Packed THD IndexerBackward requires topk_indices_global=True" + ) + else: + batch, seqlen_q, heads, head_dim = self.index_q.shape + batch_k, seqlen_k, head_dim_k = self.index_k.shape + topk = self.topk_indices.shape[-1] + dimensions = (batch, seqlen_q, seqlen_k, heads, head_dim, topk) + if any(value <= 0 for value in dimensions): + raise ValueError( + f"Indexer-backward dimensions must be positive, got {dimensions}" + ) + if batch_k != batch or head_dim_k != head_dim: + raise ValueError("index_k batch and head dimensions must match index_q") + if heads < MIN_HEADS: + raise ValueError( + f"IndexerBackward requires heads >= {MIN_HEADS}, got {heads}" + ) + if head_dim != SUPPORTED_HEAD_DIM: + raise ValueError( + f"IndexerBackward requires head_dim={SUPPORTED_HEAD_DIM}, got {head_dim}" + ) + if self.block_i != DEFAULT_BLOCK_I: + raise ValueError(f"block_i must be {DEFAULT_BLOCK_I}, got {self.block_i}") + if topk % self.block_i != 0: + raise ValueError( + f"topk ({topk}) must be divisible by block_i ({self.block_i})" + ) + + weights_shape = (seqlen_q, heads) if is_thd else (batch, seqlen_q, heads) + score_shape = (seqlen_q, topk) if is_thd else (batch, seqlen_q, topk) + _require_shape(self.weights, weights_shape, "weights") + _require_shape(self.d_index_q, self.index_q.shape, "d_index_q") + _require_shape(self.d_weights, self.weights.shape, "d_weights") + _require_shape(self.d_index_k, self.index_k.shape, "d_index_k") + _require_shape(self.attn_score, score_shape, "attn_score") + _require_shape(self.index_score, score_shape, "index_score") + _require_shape(self.topk_indices, score_shape, "topk_indices") + for desc in ( + self.index_q, + self.weights, + self.index_k, + self.d_index_q, + self.d_weights, + self.d_index_k, + self.attn_score, + self.index_score, + self.topk_indices, + ): + _require_compact(desc) + for desc in (self.index_q, self.d_index_q): + _require_contiguous_tail(desc, 2) + for desc in ( + self.weights, + self.index_k, + self.d_weights, + self.d_index_k, + self.attn_score, + self.index_score, + self.topk_indices, + ): + _require_contiguous_tail(desc, 1) + + self.is_thd = is_thd + self.batch = batch + self.seqlen_q = seqlen_q + self.seqlen_k = seqlen_k + self.heads = heads + self.head_dim = head_dim + self.topk = topk + self.sm_scale = _resolve_scale(self.requested_sm_scale, "sm_scale") + return True + + +class DenseIndexerBackwardOp(Op): + """Complete BSHD or packed-THD dense indexer-backward signature.""" + + def __init__( + self, + *, + index_q: TensorDesc[Any], + weights: TensorDesc[Any], + index_k: TensorDesc[Any], + d_index_q: TensorDesc[Any], + d_weights: TensorDesc[Any], + d_index_k: TensorDesc[Any], + attn_score: TensorDesc[Any], + attn_l1norm: TensorDesc[Any], + index_score: TensorDesc[Any], + index_lse: TensorDesc[Any], + cu_seqlens_q: TensorDesc[Any] | None = None, + cu_seqlens_k: TensorDesc[Any] | None = None, + q_causal_offsets: TensorDesc[Any] | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + sm_scale: float = 1.0, + block_i: int = DEFAULT_BLOCK_I, + ratio: int = 1, + ) -> None: + descriptors = ( + ("index_q", index_q), + ("weights", weights), + ("index_k", index_k), + ("d_index_q", d_index_q), + ("d_weights", d_weights), + ("d_index_k", d_index_k), + ("attn_score", attn_score), + ("attn_l1norm", attn_l1norm), + ("index_score", index_score), + ("index_lse", index_lse), + ) + for name, desc in descriptors: + _require_desc(name, desc) + for name, desc in ( + ("cu_seqlens_q", cu_seqlens_q), + ("cu_seqlens_k", cu_seqlens_k), + ("q_causal_offsets", q_causal_offsets), + ): + _require_desc(name, desc, optional=True) + + self.index_q = index_q + self.weights = weights + self.index_k = index_k + self.d_index_q = d_index_q + self.d_weights = d_weights + self.d_index_k = d_index_k + self.attn_score = attn_score + self.attn_l1norm = attn_l1norm + self.index_score = index_score + self.index_lse = index_lse + self.cu_seqlens_q = cu_seqlens_q + self.cu_seqlens_k = cu_seqlens_k + self.q_causal_offsets = q_causal_offsets + self.requested_max_seqlen_q = max_seqlen_q + self.requested_max_seqlen_k = max_seqlen_k + self.requested_sm_scale = sm_scale + self.block_i = int(block_i) + self.ratio = int(ratio) + + self.is_thd: Optional[bool] = None + self.batch: Optional[int] = None + self.normalization_tokens: Optional[int] = None + self.total_k: Optional[int] = None + self.heads: Optional[int] = None + self.head_dim: Optional[int] = None + self.max_seqlen_q: Optional[int] = None + self.max_seqlen_k: Optional[int] = None + self.sm_scale: Optional[float] = None + + def check_support(self) -> bool: + self.is_thd = None + self.batch = self.normalization_tokens = self.total_k = self.heads = ( + self.head_dim + ) = None + self.max_seqlen_q = self.max_seqlen_k = None + self.sm_scale = None + + if (self.cu_seqlens_q is None) != (self.cu_seqlens_k is None): + raise ValueError("cu_seqlens_q and cu_seqlens_k must be provided together") + is_thd = self.cu_seqlens_q is not None + if is_thd: + batch, normalization_tokens, total_k, heads, head_dim, max_q, max_k = ( + self._check_thd() + ) + else: + batch, normalization_tokens, total_k, heads, head_dim, max_q, max_k = ( + self._check_bshd() + ) + + for desc, name in ( + (self.index_q, "index_q"), + (self.weights, "weights"), + (self.index_k, "index_k"), + ): + _require_dtype(desc, data_type.BFLOAT16, name) + _require_dtype(self.d_index_q, data_type.BFLOAT16, "d_index_q") + _require_dtype(self.d_weights, data_type.BFLOAT16, "d_weights") + _require_dtype( + self.d_index_k, (data_type.BFLOAT16, data_type.FLOAT), "d_index_k" + ) + for desc, name in ( + (self.attn_score, "attn_score"), + (self.attn_l1norm, "attn_l1norm"), + (self.index_score, "index_score"), + (self.index_lse, "index_lse"), + ): + _require_dtype(desc, data_type.FLOAT, name) + for desc, name in ( + (self.cu_seqlens_q, "cu_seqlens_q"), + (self.cu_seqlens_k, "cu_seqlens_k"), + (self.q_causal_offsets, "q_causal_offsets"), + ): + if desc is not None: + _require_dtype(desc, data_type.INT32, name) + + if heads < MIN_HEADS: + raise ValueError( + f"DenseIndexerBackward requires heads >= {MIN_HEADS}, got {heads}" + ) + if head_dim != SUPPORTED_HEAD_DIM: + raise ValueError( + f"DenseIndexerBackward requires head_dim={SUPPORTED_HEAD_DIM}, got {head_dim}" + ) + if self.block_i != DEFAULT_BLOCK_I: + raise ValueError(f"block_i must be {DEFAULT_BLOCK_I}, got {self.block_i}") + if self.ratio < 1: + raise ValueError(f"ratio must be >= 1, got {self.ratio}") + + for desc in ( + self.index_q, + self.weights, + self.index_k, + self.d_index_q, + self.d_weights, + self.d_index_k, + self.attn_score, + self.attn_l1norm, + self.index_score, + self.index_lse, + self.cu_seqlens_q, + self.cu_seqlens_k, + self.q_causal_offsets, + ): + if desc is not None: + _require_compact(desc) + for desc in (self.index_q, self.d_index_q): + _require_contiguous_tail(desc, 2) + for desc in ( + self.weights, + self.index_k, + self.d_weights, + self.d_index_k, + self.attn_score, + self.index_score, + ): + _require_contiguous_tail(desc, 1) + + self.is_thd = is_thd + self.batch = batch + self.normalization_tokens = normalization_tokens + self.total_k = total_k + self.heads = heads + self.head_dim = head_dim + self.max_seqlen_q = max_q + self.max_seqlen_k = max_k + self.sm_scale = _resolve_scale(self.requested_sm_scale, "sm_scale") + return True + + def _check_bshd(self) -> tuple[int, int, int, int, int, int, int]: + for desc, rank, name in ( + (self.index_q, 4, "index_q"), + (self.weights, 3, "weights"), + (self.index_k, 3, "index_k"), + (self.d_index_q, 4, "d_index_q"), + (self.d_weights, 3, "d_weights"), + (self.d_index_k, 3, "d_index_k"), + (self.attn_score, 3, "attn_score"), + (self.attn_l1norm, 2, "attn_l1norm"), + (self.index_score, 3, "index_score"), + (self.index_lse, 2, "index_lse"), + ): + _require_rank(desc, rank, name) + + batch, seqlen_q, heads, head_dim = self.index_q.shape + batch_k, seqlen_k, head_dim_k = self.index_k.shape + if min(batch, seqlen_q, seqlen_k, heads, head_dim) <= 0: + raise ValueError("Dense indexer-backward dimensions must be positive") + if batch_k != batch or head_dim_k != head_dim: + raise ValueError("index_k batch and head dimensions must match index_q") + _require_shape(self.weights, (batch, seqlen_q, heads), "weights") + _require_shape(self.d_index_q, self.index_q.shape, "d_index_q") + _require_shape(self.d_weights, self.weights.shape, "d_weights") + _require_shape(self.d_index_k, self.index_k.shape, "d_index_k") + score_shape = (batch, seqlen_q, seqlen_k) + denom_shape = (batch, seqlen_q) + _require_shape(self.attn_score, score_shape, "attn_score") + _require_shape(self.index_score, score_shape, "index_score") + _require_shape(self.attn_l1norm, denom_shape, "attn_l1norm") + _require_shape(self.index_lse, denom_shape, "index_lse") + if self.q_causal_offsets is not None: + _require_rank(self.q_causal_offsets, 1, "q_causal_offsets") + _require_shape(self.q_causal_offsets, (batch,), "q_causal_offsets") + + max_q = ( + seqlen_q + if self.requested_max_seqlen_q is None + else int(self.requested_max_seqlen_q) + ) + max_k = ( + seqlen_k + if self.requested_max_seqlen_k is None + else int(self.requested_max_seqlen_k) + ) + if (max_q, max_k) != (seqlen_q, seqlen_k): + raise ValueError( + f"BSHD max_seqlen_q/k must match tensor extents {(seqlen_q, seqlen_k)}, got {(max_q, max_k)}" + ) + return batch, batch * seqlen_q, batch * seqlen_k, heads, head_dim, max_q, max_k + + def _check_thd(self) -> tuple[int, int, int, int, int, int, int]: + for desc, rank, name in ( + (self.index_q, 3, "index_q"), + (self.weights, 2, "weights"), + (self.index_k, 2, "index_k"), + (self.d_index_q, 3, "d_index_q"), + (self.d_weights, 2, "d_weights"), + (self.d_index_k, 2, "d_index_k"), + (self.attn_score, 2, "attn_score"), + (self.attn_l1norm, 1, "attn_l1norm"), + (self.index_score, 2, "index_score"), + (self.index_lse, 1, "index_lse"), + (self.cu_seqlens_q, 1, "cu_seqlens_q"), + (self.cu_seqlens_k, 1, "cu_seqlens_k"), + ): + _require_rank(desc, rank, name) + + total_q, heads, head_dim = self.index_q.shape + total_k, head_dim_k = self.index_k.shape + batch = self.cu_seqlens_q.shape[0] - 1 + if min(total_q, total_k, batch, heads, head_dim) <= 0: + raise ValueError("Dense THD indexer-backward dimensions must be positive") + if self.cu_seqlens_k.shape != (batch + 1,): + raise ValueError( + "cu_seqlens_q and cu_seqlens_k must describe the same batch" + ) + if head_dim_k != head_dim: + raise ValueError("index_k head dimension must match index_q") + if self.requested_max_seqlen_q is None or self.requested_max_seqlen_k is None: + raise ValueError( + "THD dense indexer backward requires max_seqlen_q and max_seqlen_k" + ) + max_q = int(self.requested_max_seqlen_q) + max_k = int(self.requested_max_seqlen_k) + if max_q <= 0 or max_k <= 0: + raise ValueError(f"max_seqlen_q/k must be positive, got {(max_q, max_k)}") + + _require_shape(self.weights, (total_q, heads), "weights") + _require_shape(self.d_index_q, self.index_q.shape, "d_index_q") + _require_shape(self.d_weights, self.weights.shape, "d_weights") + _require_shape(self.d_index_k, self.index_k.shape, "d_index_k") + score_shape = (total_q, max_k) + _require_shape(self.attn_score, score_shape, "attn_score") + _require_shape(self.index_score, score_shape, "index_score") + _require_shape(self.attn_l1norm, (total_q,), "attn_l1norm") + _require_shape(self.index_lse, (total_q,), "index_lse") + if self.q_causal_offsets is not None: + _require_rank(self.q_causal_offsets, 1, "q_causal_offsets") + _require_shape(self.q_causal_offsets, (batch,), "q_causal_offsets") + return batch, total_q, total_k, heads, head_dim, max_q, max_k + + +__all__ = [ + "DEFAULT_BLOCK_I", + "DenseIndexerBackwardOp", + "IndexerBackwardOp", + "MIN_HEADS", + "SUPPORTED_HEAD_DIM", +] diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py index 4b8df4c14..49fbc38ac 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/__init__.py @@ -1,3 +1,19 @@ -from .api import IndexerForward, indexer_forward_wrapper +"""Lazy Torch API and framework-neutral indexer-forward exports. -__all__ = ["IndexerForward", "indexer_forward_wrapper"] +JAX APIs are exported through :mod:`cudnn.jax`. +""" + +from ...common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ( + "IndexerForwardOp", + "SUPPORTED_COMPUTE_CAPABILITIES", + "TMA_ALIGN_ELEMENTS", + ), + "api": ("IndexerForward", "indexer_forward_wrapper"), + }, + submodules=("api", "op"), +) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py index d35e0da8b..b66de8af3 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py @@ -26,8 +26,9 @@ from .indexer_fwd_sm100 import IndexerForwardSm100 from ._interface import indexer_fwd as indexer_fwd_sm100 from ._interface_sm90 import indexer_fwd as indexer_fwd_sm90 +from .op import IndexerForwardOp, TMA_ALIGN_ELEMENTS -TMA_ALIGN_ELEMS = 4 # FP32 output => seqlen_k padded to multiples of 4 (16 B) +TMA_ALIGN_ELEMS = TMA_ALIGN_ELEMENTS class IndexerForward(APIBase): @@ -66,6 +67,21 @@ def __init__( self.sm_scale = float(sm_scale) self.qhead_per_kv_head = qhead_per_kv_head + self._op = IndexerForwardOp( + q=self.q_desc, + k=self.k_desc, + weight=self.w_desc, + output=self.o_desc, + ratio=self.ratio, + qhead_per_kv_head=self.qhead_per_kv_head, + m_block_size=self.m_block_size, + n_block_size=self.n_block_size, + q_stage=self.q_stage, + kv_stage=self.kv_stage, + sm_scale=self.sm_scale, + target_compute_capability=100, + ) + self.batch_size = None self.s_q = None self.s_k = None @@ -76,55 +92,7 @@ def __init__( def check_support(self) -> bool: self._logger.debug("Entering check_support") - self._value_error_if( - self.q_desc.ndim != 4, - f"Q must be 4-D (B, S_q, H_q, D), got {self.q_desc.shape}", - ) - self._value_error_if( - self.k_desc.ndim != 4, - f"K must be 4-D (B, S_k, H_kv, D), got {self.k_desc.shape}", - ) - self._value_error_if( - self.w_desc.ndim != 3, - f"W must be 3-D (B, S_q, H_q), got {self.w_desc.shape}", - ) - self._value_error_if( - self.o_desc.ndim != 3, - f"Out must be 3-D (B, S_q, S_k_padded), got {self.o_desc.shape}", - ) - - b, s_q, h_q, d = self.q_desc.shape - b_k, s_k, h_kv, d_k = self.k_desc.shape - b_o, s_q_out, s_k_padded_from_out = self.o_desc.shape - self._value_error_if(b != b_k, f"Batch size mismatch Q={b} vs K={b_k}") - self._value_error_if(b != b_o, f"Batch size mismatch Q={b} vs Out={b_o}") - self._value_error_if(s_q != s_q_out, f"S_q mismatch Q={s_q} vs Out={s_q_out}") - self._value_error_if(d != d_k, f"Head dim mismatch Q={d} vs K={d_k}") - self._value_error_if( - d != 128, - f"IndexerForward is tuned for head_dim=128 only, got {d}", - ) - - qhpkv = self.qhead_per_kv_head if self.qhead_per_kv_head is not None else (h_q // h_kv) - self._value_error_if( - qhpkv * h_kv != h_q, - f"qhead_per_kv_head * h_kv != h_q ({qhpkv} * {h_kv} != {h_q})", - ) - self._value_error_if( - qhpkv not in (32, 64), - f"qhead_per_kv_head must be 32 or 64, got {qhpkv}", - ) - self.qhead_per_kv_head = qhpkv - - self._check_dtype(self.q_desc, torch.bfloat16, name="Q") - self._check_dtype(self.k_desc, torch.bfloat16, name="K") - self._check_dtype(self.w_desc, torch.bfloat16, name="W") - self._check_dtype(self.o_desc, torch.float32, name="Out") - - self._value_error_if( - s_k_padded_from_out % TMA_ALIGN_ELEMS != 0, - f"Out seqlen_k dim must be a multiple of {TMA_ALIGN_ELEMS}, got {s_k_padded_from_out}", - ) + self._op.check_support() major = device_major() self._runtime_error_if( @@ -132,13 +100,14 @@ def check_support(self) -> bool: f"IndexerForward requires SM100+ compute capability, found SM{major}", ) - self.batch_size = b - self.s_q = s_q - self.s_k = s_k - self.s_k_padded = s_k_padded_from_out - self.h_q = h_q - self.h_kv = h_kv - self.head_dim = d + self.batch_size = self._op.batch_size + self.s_q = self._op.s_q + self.s_k = self._op.s_k + self.s_k_padded = self.o_desc.shape[-1] + self.h_q = self._op.h_q + self.h_kv = self._op.h_kv + self.head_dim = self._op.head_dim + self.qhead_per_kv_head = self._op.qhead_per_kv_head self._is_supported = True return True @@ -172,6 +141,7 @@ def compile(self) -> None: cutlass.Float32(self.sm_scale), None, None, + None, fake_stream, options=compile_options(), ) @@ -191,6 +161,7 @@ def tensor_api(q, k, w, out, stream): cutlass.Float32(self.sm_scale), None, None, + None, stream, ) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py new file mode 100644 index 000000000..2746b8539 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/jax.py @@ -0,0 +1,401 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX API for DeepSeek indexer forward.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +import jax +import jax.numpy as jnp + +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..._jax.compiler import compile_options_for_target +from ..._jax.layout import mode_from_layout, to_public_axes +from .op import IndexerForwardOp, SUPPORTED_COMPUTE_CAPABILITIES, TMA_ALIGN_ELEMENTS + + +def _resolve_layout( + name: str, + layout: str | None, + *, + default: str, + kernel_axes: str, + supported: tuple[str, ...], +) -> tuple[str, tuple[int, ...]]: + layout = default if layout is None else layout + mode = mode_from_layout(layout, kernel_axes=kernel_axes) + if layout not in supported: + choices = ", ".join(repr(value) for value in supported) + raise ValueError(f"{name} must be one of ({choices}), got {layout!r}") + return layout, mode + + +class IndexerForward(JaxApiBase): + """JAX callable specialized from fixed BSHD or packed THD metadata. + + Fixed inputs may use batch-major ``BSHD``/``BSH`` or sequence-major + ``SBHD``/``SBH`` public axis orders. The output may similarly use ``BSK`` + or ``SBK``. Packed inputs currently use ``THD``/``TH`` and return ``TK``. + + ``sample_out`` describes the padded physical FP32 kernel buffer. When it + is omitted, the adapter creates that descriptor and returns a logical view + with the unpadded K extent. + """ + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_w: Any, + *, + sample_out: Any | None = None, + sample_cu_seqlens_q: Any | None = None, + sample_cu_seqlens_k: Any | None = None, + sample_q_causal_offsets: Any | None = None, + ratio: int = 4, + qhead_per_kv_head: int | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + m_block_size: int = 128, + n_block_size: int = 128, + q_stage: int = 2, + kv_stage: int = 4, + sm_scale: float = 1.0, + q_layout: str | None = None, + k_layout: str | None = None, + w_layout: str | None = None, + output_layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + if (sample_cu_seqlens_q is None) != (sample_cu_seqlens_k is None): + raise ValueError("THD input requires both sample_cu_seqlens_q and sample_cu_seqlens_k") + + self.is_varlen = sample_cu_seqlens_q is not None + if self.is_varlen: + self.q_layout, self.q_mode = _resolve_layout( + "q_layout", + q_layout, + default="THD", + kernel_axes="THD", + supported=("THD",), + ) + self.k_layout, self.k_mode = _resolve_layout( + "k_layout", + k_layout, + default="THD", + kernel_axes="THD", + supported=("THD",), + ) + self.w_layout, self.w_mode = _resolve_layout( + "w_layout", + w_layout, + default="TH", + kernel_axes="TH", + supported=("TH",), + ) + self.output_layout, self.output_mode = _resolve_layout( + "output_layout", + output_layout, + default="TK", + kernel_axes="TK", + supported=("TK",), + ) + else: + self.q_layout, self.q_mode = _resolve_layout( + "q_layout", + q_layout, + default="BSHD", + kernel_axes="BSHD", + supported=("BSHD", "SBHD"), + ) + self.k_layout, self.k_mode = _resolve_layout( + "k_layout", + k_layout, + default="BSHD", + kernel_axes="BSHD", + supported=("BSHD", "SBHD"), + ) + self.w_layout, self.w_mode = _resolve_layout( + "w_layout", + w_layout, + default="BSH", + kernel_axes="BSH", + supported=("BSH", "SBH"), + ) + self.output_layout, self.output_mode = _resolve_layout( + "output_layout", + output_layout, + default="BSK", + kernel_axes="BSK", + supported=("BSK", "SBK"), + ) + + self.target_compute_capability = self._resolve_compute_capability( + target_compute_capability=target_compute_capability, + supported_compute_capabilities=SUPPORTED_COMPUTE_CAPABILITIES, + operation_name="IndexerForward", + ) + self.q_desc = self._to_tensor_desc(sample_q, "sample_q", mode=self.q_mode) + self.k_desc = self._to_tensor_desc(sample_k, "sample_k", mode=self.k_mode) + self.w_desc = self._to_tensor_desc(sample_w, "sample_w", mode=self.w_mode) + self.cu_seqlens_q_desc = None if sample_cu_seqlens_q is None else self._to_tensor_desc(sample_cu_seqlens_q, "sample_cu_seqlens_q") + self.cu_seqlens_k_desc = None if sample_cu_seqlens_k is None else self._to_tensor_desc(sample_cu_seqlens_k, "sample_cu_seqlens_k") + self.q_causal_offsets_desc = ( + None + if sample_q_causal_offsets is None + else self._to_tensor_desc( + sample_q_causal_offsets, + "sample_q_causal_offsets", + ) + ) + self.max_seqlen_q = max_seqlen_q + self.max_seqlen_k = max_seqlen_k + + if sample_out is None: + self.o_desc = self._default_output_desc(max_seqlen_k) + else: + self.o_desc = self._to_tensor_desc( + sample_out, + "sample_out", + mode=self.output_mode, + init_value=float("-inf"), + ) + + self._op = IndexerForwardOp( + q=self.q_desc, + k=self.k_desc, + weight=self.w_desc, + output=self.o_desc, + cu_seqlens_q=self.cu_seqlens_q_desc, + cu_seqlens_k=self.cu_seqlens_k_desc, + q_causal_offsets=self.q_causal_offsets_desc, + ratio=ratio, + qhead_per_kv_head=qhead_per_kv_head, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + m_block_size=m_block_size, + n_block_size=n_block_size, + q_stage=q_stage, + kv_stage=kv_stage, + sm_scale=sm_scale, + target_compute_capability=self.target_compute_capability, + ) + + def _default_output_desc(self, max_seqlen_k: int | None) -> JaxTensorDesc: + if self.is_varlen: + if self.q_desc.ndim != 3: + raise ValueError(f"THD Q must have rank 3, got {self.q_desc.shape}") + if max_seqlen_k is None: + raise ValueError("THD input requires max_seqlen_k") + leading_shape = (self.q_desc.shape[0],) + logical_seqlen_k = int(max_seqlen_k) + else: + if self.q_desc.ndim != 4 or self.k_desc.ndim != 4: + raise ValueError("Fixed indexer forward requires rank-4 Q and K") + leading_shape = (self.q_desc.shape[0], self.q_desc.shape[1]) + logical_seqlen_k = self.k_desc.shape[1] + + padded_seqlen_k = ((logical_seqlen_k + TMA_ALIGN_ELEMENTS - 1) // TMA_ALIGN_ELEMENTS) * TMA_ALIGN_ELEMENTS + canonical_shape = (*leading_shape, padded_seqlen_k) + return JaxTensorDesc.from_shape( + to_public_axes(canonical_shape, self.output_mode), + jnp.float32, + name="sample_out", + mode=self.output_mode, + init_value=float("-inf"), + ) + + def check_support(self) -> bool: + return self._op.check_support() + + def __call__( + self, + q: Any, + k: Any, + w: Any, + *, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + q_causal_offsets: Any | None = None, + ) -> TupleDict: + self.check_support() + self._check_optional_signature(cu_seqlens_q, self.cu_seqlens_q_desc, "cu_seqlens_q") + self._check_optional_signature(cu_seqlens_k, self.cu_seqlens_k_desc, "cu_seqlens_k") + self._check_optional_signature(q_causal_offsets, self.q_causal_offsets_desc, "q_causal_offsets") + + inputs = [q, k, w] + input_descs = [self.q_desc, self.k_desc, self.w_desc] + for value, desc in ( + (cu_seqlens_q, self.cu_seqlens_q_desc), + (cu_seqlens_k, self.cu_seqlens_k_desc), + (q_causal_offsets, self.q_causal_offsets_desc), + ): + if desc is not None: + inputs.append(value) + input_descs.append(desc) + + output_desc = self.o_desc.with_divisibility( + (None,) * (self.o_desc.ndim - 1) + (TMA_ALIGN_ELEMENTS,) + ) + (scores_padded,) = self._call_kernel( + tuple(inputs), + launch=self._launch_kernel, + output_descs=(output_desc,), + input_descs=tuple(input_descs), + compile_options=compile_options_for_target(self.target_compute_capability), + ) + if self._op.s_k is None: + raise RuntimeError("IndexerForward output shape was not resolved by check_support()") + output_slice = [slice(None)] * self.o_desc.ndim + output_slice[self.output_mode[-1]] = slice(0, self._op.s_k) + return TupleDict(scores=scores_padded[tuple(output_slice)]) + + @staticmethod + def _check_optional_signature(value: Any | None, desc: JaxTensorDesc | None, name: str) -> None: + if (value is None) != (desc is None): + expected = "omitted" if desc is None else "provided" + raise ValueError(f"{name} must be {expected} for this specialized callable") + + def _launch_kernel( + self, + stream: Any, + *arguments: Any, + ) -> None: + from cutlass import BFloat16, Float32, Int32 + + *inputs, output = arguments + q, k, w, *optional_inputs = inputs + cu_seqlens_q = optional_inputs.pop(0) if self.cu_seqlens_q_desc is not None else None + cu_seqlens_k = optional_inputs.pop(0) if self.cu_seqlens_k_desc is not None else None + q_causal_offsets = optional_inputs.pop(0) if self.q_causal_offsets_desc is not None else None + if optional_inputs: + raise RuntimeError("Unexpected IndexerForward kernel inputs") + resolved = ( + self._op.head_dim, + self._op.qhead_per_kv_head, + self._op.h_kv, + self._op.s_q, + self._op.s_k, + ) + if any(value is None for value in resolved): + raise RuntimeError("IndexerForward launch configuration was not resolved by check_support()") + head_dim, qhead_per_kv_head, h_kv, s_q, s_k = resolved + if self.target_compute_capability < 100: + from .indexer_fwd_sm90 import IndexerForwardSm90 + + kernel = IndexerForwardSm90( + BFloat16, + head_dim=head_dim, + qhead_per_kvhead=qhead_per_kv_head, + ratio=self._op.ratio, + is_varlen=bool(self._op.is_varlen), + ) + else: + from .indexer_fwd_sm100 import IndexerForwardSm100 + + kernel = IndexerForwardSm100( + head_dim=head_dim, + qhead_per_kvhead=qhead_per_kv_head, + ratio=self._op.ratio, + is_varlen=bool(self._op.is_varlen), + m_block_size=self._op.m_block_size, + n_block_size=self._op.n_block_size, + q_stage=self._op.q_stage, + kv_stage=self._op.kv_stage, + ) + + kernel( + q, + k, + w, + output, + Int32(h_kv), + Int32(s_q), + Int32(s_k), + Float32(self._op.sm_scale), + cu_seqlens_q, + cu_seqlens_k, + q_causal_offsets, + stream, + ) + + +@partial( + jax.jit, + static_argnames=( + "ratio", + "qhead_per_kv_head", + "max_seqlen_q", + "max_seqlen_k", + "m_block_size", + "n_block_size", + "q_stage", + "kv_stage", + "sm_scale", + "q_layout", + "k_layout", + "w_layout", + "output_layout", + "target_compute_capability", + ), +) +def indexer_forward_wrapper( + q: Any, + k: Any, + w: Any, + *, + ratio: int = 4, + qhead_per_kv_head: int | None = None, + m_block_size: int = 128, + n_block_size: int = 128, + q_stage: int = 2, + kv_stage: int = 4, + sm_scale: float = 1.0, + q_layout: str | None = None, + k_layout: str | None = None, + w_layout: str | None = None, + output_layout: str | None = None, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + q_causal_offsets: Any | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute fixed batch/sequence-major or packed THD indexer scores.""" + + return IndexerForward( + q, + k, + w, + sample_cu_seqlens_q=cu_seqlens_q, + sample_cu_seqlens_k=cu_seqlens_k, + sample_q_causal_offsets=q_causal_offsets, + ratio=ratio, + qhead_per_kv_head=qhead_per_kv_head, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + m_block_size=m_block_size, + n_block_size=n_block_size, + q_stage=q_stage, + kv_stage=kv_stage, + sm_scale=sm_scale, + q_layout=q_layout, + k_layout=k_layout, + w_layout=w_layout, + output_layout=output_layout, + target_compute_capability=target_compute_capability, + )( + q, + k, + w, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + q_causal_offsets=q_causal_offsets, + ) + + +__all__ = ["IndexerForward", "indexer_forward_wrapper"] diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/op.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/op.py new file mode 100644 index 000000000..347b6a7a8 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/op.py @@ -0,0 +1,239 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral validation for DeepSeek indexer forward.""" + +from __future__ import annotations + +from typing import Any + +from ... import data_type +from ...common.op import Op +from ...common.tensor_desc import TensorDesc + +TMA_ALIGN_ELEMENTS = 4 +SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) + + +def _require_compact( + desc: TensorDesc[Any], + label: str, + contiguous_axis: str, +) -> None: + if not desc.is_compact() or desc.stride[-1] != 1: + raise ValueError( + f"{label} must be compact with its canonical {contiguous_axis} axis " f"contiguous, got stride {desc.stride} and stride order {desc.stride_order}" + ) + + +def _require_desc(name: str, desc: TensorDesc[Any] | None) -> None: + if desc is not None and not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc or None, got {type(desc).__name__}") + + +class IndexerForwardOp(Op): + """Complete tensor signature and static configuration for indexer forward.""" + + def __init__( + self, + *, + q: TensorDesc[Any], + k: TensorDesc[Any], + weight: TensorDesc[Any], + output: TensorDesc[Any], + cu_seqlens_q: TensorDesc[Any] | None = None, + cu_seqlens_k: TensorDesc[Any] | None = None, + q_causal_offsets: TensorDesc[Any] | None = None, + ratio: int = 4, + qhead_per_kv_head: int | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + m_block_size: int = 128, + n_block_size: int = 128, + q_stage: int = 2, + kv_stage: int = 4, + sm_scale: float = 1.0, + target_compute_capability: int = 100, + ) -> None: + for name, desc in (("q", q), ("k", k), ("weight", weight), ("output", output)): + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + for name, desc in ( + ("cu_seqlens_q", cu_seqlens_q), + ("cu_seqlens_k", cu_seqlens_k), + ("q_causal_offsets", q_causal_offsets), + ): + _require_desc(name, desc) + + self.q = q + self.k = k + self.weight = weight + self.output = output + self.cu_seqlens_q = cu_seqlens_q + self.cu_seqlens_k = cu_seqlens_k + self.q_causal_offsets = q_causal_offsets + self.ratio = int(ratio) + self.requested_qhead_per_kv_head = qhead_per_kv_head + self.requested_max_seqlen_q = max_seqlen_q + self.requested_max_seqlen_k = max_seqlen_k + self.m_block_size = int(m_block_size) + self.n_block_size = int(n_block_size) + self.q_stage = int(q_stage) + self.kv_stage = int(kv_stage) + self.sm_scale = float(sm_scale) + self.target_compute_capability = int(target_compute_capability) + + self.is_varlen: bool | None = None + self.batch_size: int | None = None + self.s_q: int | None = None + self.s_k: int | None = None + self.h_q: int | None = None + self.h_kv: int | None = None + self.head_dim: int | None = None + self.qhead_per_kv_head: int | None = None + + def check_support(self) -> bool: + is_varlen = self.cu_seqlens_q is not None or self.cu_seqlens_k is not None + if (self.cu_seqlens_q is None) != (self.cu_seqlens_k is None): + raise ValueError("THD input requires both cu_seqlens_q and cu_seqlens_k") + if self.target_compute_capability not in SUPPORTED_COMPUTE_CAPABILITIES: + raise ValueError("target_compute_capability must be one of " f"{SUPPORTED_COMPUTE_CAPABILITIES}, got {self.target_compute_capability}") + + for desc, label in ((self.q, "Q"), (self.k, "K"), (self.weight, "W")): + if desc.cudnn_dtype != data_type.BFLOAT16: + raise ValueError(f"{label} must have dtype bfloat16, got {desc.dtype}") + if self.output.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"Out must have dtype float32, got {self.output.dtype}") + + if is_varlen: + batch_size, s_q, s_k, h_q, h_kv, head_dim = self._check_varlen_signature() + else: + batch_size, s_q, s_k, h_q, h_kv, head_dim = self._check_fixed_signature() + for desc, label, contiguous_axis in ( + (self.q, "Q", "D"), + (self.k, "K", "D"), + (self.weight, "W", "H"), + (self.output, "Out", "K"), + ): + _require_compact(desc, label, contiguous_axis) + + if any(value <= 0 for value in (batch_size, s_q, s_k, h_q, h_kv, head_dim)): + raise ValueError("IndexerForward dimensions must be positive, got " f"B={batch_size}, S_q={s_q}, S_k={s_k}, H_q={h_q}, H_kv={h_kv}, D={head_dim}") + if head_dim != 128: + raise ValueError(f"IndexerForward requires head_dim=128, got {head_dim}") + + qhead_per_kv_head = self.requested_qhead_per_kv_head + if qhead_per_kv_head is None: + if h_q % h_kv != 0: + raise ValueError(f"H_q ({h_q}) must be divisible by H_kv ({h_kv})") + qhead_per_kv_head = h_q // h_kv + if qhead_per_kv_head * h_kv != h_q: + raise ValueError("qhead_per_kv_head * H_kv must equal H_q, got " f"{qhead_per_kv_head} * {h_kv} != {h_q}") + if qhead_per_kv_head not in (32, 64): + raise ValueError(f"qhead_per_kv_head must be 32 or 64, got {qhead_per_kv_head}") + if self.target_compute_capability < 100 and h_kv != 1: + raise ValueError(f"SM90 IndexerForward requires H_kv=1, got {h_kv}") + if self.ratio < 1: + raise ValueError(f"ratio must be at least 1, got {self.ratio}") + + tuning = (self.m_block_size, self.n_block_size, self.q_stage, self.kv_stage) + if any(value <= 0 for value in tuning): + raise ValueError(f"IndexerForward tuning values must be positive, got {tuning}") + if self.target_compute_capability < 100 and tuning != (128, 128, 2, 4): + raise ValueError("SM90 IndexerForward supports only m_block_size=128, n_block_size=128, " f"q_stage=2, kv_stage=4; got {tuning}") + if self.target_compute_capability >= 100: + packed_q = s_q * qhead_per_kv_head + minimum_packed_q = self.q_stage * self.m_block_size + if packed_q < minimum_packed_q: + raise ValueError( + "SM100-family IndexerForward requires at least one complete packed query tile: max_seqlen_q * " + f"qhead_per_kv_head >= q_stage * m_block_size, got " + f"{s_q} * {qhead_per_kv_head} < {self.q_stage} * " + f"{self.m_block_size}" + ) + + if self.q_causal_offsets is not None: + offsets = self.q_causal_offsets + if offsets.ndim != 1 or offsets.shape != (batch_size,): + raise ValueError(f"q_causal_offsets must have shape {(batch_size,)}, got {offsets.shape}") + if offsets.cudnn_dtype != data_type.INT32: + raise ValueError(f"q_causal_offsets must have dtype int32, got {offsets.dtype}") + _require_compact(offsets, "q_causal_offsets", "entry") + + expected_leading_shape = (self.q.shape[0],) if is_varlen else (batch_size, s_q) + output_seqlen_k = self.output.shape[-1] + if self.output.shape[:-1] != expected_leading_shape or output_seqlen_k < s_k or output_seqlen_k % TMA_ALIGN_ELEMENTS: + raise ValueError( + "Out must have leading shape " + f"{expected_leading_shape} and an S_k extent >= {s_k} divisible by " + f"{TMA_ALIGN_ELEMENTS}, got {self.output.shape}" + ) + + self.is_varlen = is_varlen + self.batch_size = batch_size + self.s_q = s_q + self.s_k = s_k + self.h_q = h_q + self.h_kv = h_kv + self.head_dim = head_dim + self.qhead_per_kv_head = qhead_per_kv_head + return True + + def _check_fixed_signature(self) -> tuple[int, int, int, int, int, int]: + if self.q.ndim != 4: + raise ValueError(f"Q must be 4-D (B, S_q, H_q, D), got {self.q.shape}") + if self.k.ndim != 4: + raise ValueError(f"K must be 4-D (B, S_k, H_kv, D), got {self.k.shape}") + if self.weight.ndim != 3: + raise ValueError(f"W must be 3-D (B, S_q, H_q), got {self.weight.shape}") + if self.output.ndim != 3: + raise ValueError(f"Out must be 3-D, got {self.output.shape}") + + batch_size, s_q, h_q, head_dim = self.q.shape + k_batch, s_k, h_kv, k_head_dim = self.k.shape + if k_batch != batch_size: + raise ValueError(f"Q and K batch dimensions must match, got {batch_size} and {k_batch}") + if k_head_dim != head_dim: + raise ValueError(f"Q and K head dimensions must match, got {head_dim} and {k_head_dim}") + if self.weight.shape != (batch_size, s_q, h_q): + raise ValueError(f"W must have shape {(batch_size, s_q, h_q)}, got {self.weight.shape}") + if self.requested_max_seqlen_q not in (None, s_q): + raise ValueError(f"max_seqlen_q must be omitted or equal S_q={s_q}") + if self.requested_max_seqlen_k not in (None, s_k): + raise ValueError(f"max_seqlen_k must be omitted or equal S_k={s_k}") + return batch_size, s_q, s_k, h_q, h_kv, head_dim + + def _check_varlen_signature(self) -> tuple[int, int, int, int, int, int]: + if self.q.ndim != 3: + raise ValueError(f"THD Q must be 3-D (total_q, H_q, D), got {self.q.shape}") + if self.k.ndim != 3: + raise ValueError(f"THD K must be 3-D (total_k, H_kv, D), got {self.k.shape}") + if self.weight.ndim != 2: + raise ValueError(f"THD W must be 2-D (total_q, H_q), got {self.weight.shape}") + if self.output.ndim != 2: + raise ValueError(f"THD Out must be 2-D, got {self.output.shape}") + + if self.cu_seqlens_q is None or self.cu_seqlens_k is None: + raise RuntimeError("THD cumulative sequence descriptors were not configured") + for desc, label in ((self.cu_seqlens_q, "cu_seqlens_q"), (self.cu_seqlens_k, "cu_seqlens_k")): + if desc.ndim != 1 or desc.cudnn_dtype != data_type.INT32: + raise ValueError(f"{label} must be a 1-D int32 tensor, got shape {desc.shape} and dtype {desc.dtype}") + _require_compact(desc, label, "entry") + if self.cu_seqlens_q.shape != self.cu_seqlens_k.shape: + raise ValueError("cu_seqlens_q and cu_seqlens_k must have the same shape") + if self.cu_seqlens_q.shape[0] < 2: + raise ValueError("cu_seqlens_q must contain at least two entries") + if self.requested_max_seqlen_q is None or self.requested_max_seqlen_k is None: + raise ValueError("THD input requires max_seqlen_q and max_seqlen_k") + + total_q, h_q, head_dim = self.q.shape + _, h_kv, k_head_dim = self.k.shape + if k_head_dim != head_dim: + raise ValueError(f"Q and K head dimensions must match, got {head_dim} and {k_head_dim}") + if self.weight.shape != (total_q, h_q): + raise ValueError(f"THD W must have shape {(total_q, h_q)}, got {self.weight.shape}") + batch_size = self.cu_seqlens_q.shape[0] - 1 + return batch_size, int(self.requested_max_seqlen_q), int(self.requested_max_seqlen_k), h_q, h_kv, head_dim + + +__all__ = ["IndexerForwardOp", "SUPPORTED_COMPUTE_CAPABILITIES", "TMA_ALIGN_ELEMENTS"] diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/__init__.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/__init__.py index 481ef694e..7bf0c95d1 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/__init__.py @@ -1,13 +1,20 @@ -from .api import ( - IndexerTopK, - indexer_top_k_wrapper, - local_to_global_wrapper, - compactify_wrapper, -) +"""Lazy Torch API and framework-neutral indexer top-K exports. + +JAX APIs are exported through :mod:`cudnn.jax`. +""" -__all__ = [ - "IndexerTopK", - "indexer_top_k_wrapper", - "local_to_global_wrapper", - "compactify_wrapper", -] +from ...common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ("IndexerTopKOp", "bucket_num_cols"), + "api": ( + "IndexerTopK", + "compactify_wrapper", + "indexer_top_k_wrapper", + "local_to_global_wrapper", + ), + }, + submodules=("api", "op"), +) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/api.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/api.py index fd5a107a6..b1502bffd 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/api.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/api.py @@ -7,10 +7,13 @@ import torch import cuda.bindings.driver as cuda +from cudnn import data_type from cudnn.api_base import APIBase, TupleDict +from cudnn.deepseek_sparse_attention.utils.runtime import torch_stream_context from .local_to_global_dsl import local_to_global as _local_to_global from .compactify import compactify as _compactify +from .op import IndexerTopKOp _SUPPORTED_DTYPES = (torch.float32, torch.float16, torch.bfloat16) @@ -77,37 +80,52 @@ def __init__( self.return_val = bool(return_val) self.num_copy_bits = int(num_copy_bits) - def check_support(self) -> bool: - self._logger.debug("Entering check_support") + if self.input_desc.ndim != 2: + raise ValueError(f"input_values must be 2-D, got {self.input_desc.shape}") + if self.seq_lens_desc.ndim != 1: + raise ValueError(f"seq_lens must be 1-D, got {self.seq_lens_desc.shape}") self._check_dtype(self.input_desc, list(_SUPPORTED_DTYPES), name="input_values") - self._check_dtype(self.seq_lens_desc, torch.int32, name="seq_lens") - self._value_error_if( - self.input_desc.ndim != 2, - f"input_values must be 2-D (n_rows, num_cols), got {self.input_desc.shape}", + if self.top_k <= 0: + raise ValueError(f"top_k must be positive, got {self.top_k}") + + output_shape = (self.input_desc.shape[0], self.top_k) + self.indices_desc = self.input_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=output_shape, + name="indices", ) - self._value_error_if( - self.seq_lens_desc.ndim != 1, - f"seq_lens must be 1-D, got {self.seq_lens_desc.shape}", + self.values_desc = ( + self.input_desc.compact_like( + cudnn_dtype=self.input_desc.cudnn_dtype, + shape=output_shape, + name="values", + ) + if self.return_val + else None ) - self._value_error_if( - self.top_k <= 0 or self.top_k > 2048, - f"top_k must be in (0, 2048], got {self.top_k}", + buffer_count = 2 if self.input_desc.cudnn_dtype == data_type.FLOAT else 1 + self.workspace_desc = self.input_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=(self.input_desc.shape[0], buffer_count, self.input_desc.shape[1]), + name="extra_buffer", ) - - # Enforce the kernel's n_rows == batch_size * next_n invariant - # up-front so misuse surfaces here rather than as silently-empty - # rows (the stagger formula produces non-positive lengths when - # next_n exceeds n_rows per batch). - n_rows = self.input_desc.shape[0] - batch_size = self.seq_lens_desc.shape[0] - self._value_error_if( - n_rows != batch_size * self.next_n, - f"n_rows ({n_rows}) must equal seq_lens.numel() * next_n " - f"({batch_size} * {self.next_n} = {batch_size * self.next_n}). " - f"For independent top-K over equal-length rows use " - f"next_n=1 and seq_lens of shape (n_rows,).", + self._op = IndexerTopKOp( + input_values=self.input_desc, + seq_lens=self.seq_lens_desc, + output_indices=self.indices_desc, + output_values=self.values_desc, + workspace=self.workspace_desc, + top_k=self.top_k, + next_n=self.next_n, + return_val=self.return_val, + num_copy_bits=self.num_copy_bits, + target_compute_capability=90, ) + def check_support(self) -> bool: + self._logger.debug("Entering check_support") + self._op.check_support() + major, _ = torch.cuda.get_device_capability() self._runtime_error_if( major < 9, @@ -136,14 +154,15 @@ def execute( self.compile() kernel = self._compiled_kernel or _get_cute_dsl_topk_wrapper() - indices, values = kernel( - input_values, - seq_lens, - self.top_k, - self.next_n, - return_val=self.return_val, - num_copy_bits=self.num_copy_bits, - ) + with torch_stream_context(current_stream): + indices, values = kernel( + input_values, + seq_lens, + self.top_k, + self.next_n, + return_val=self.return_val, + num_copy_bits=self.num_copy_bits, + ) return indices, values @@ -170,6 +189,8 @@ def indexer_top_k_wrapper( full details. ``values`` is ``None`` when ``return_val=False``. + Every runtime ``seq_lens[b]`` must satisfy + ``top_k + next_n - 1 <= seq_lens[b] <= input_values.shape[1]``. """ cache_key = ( input_values.dtype, @@ -191,7 +212,7 @@ def indexer_top_k_wrapper( return_val=return_val, num_copy_bits=num_copy_bits, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_IndexerTopKObjects[cache_key] = obj diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/compactify.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/compactify.py index ea1997711..48739b6f4 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/compactify.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/compactify.py @@ -9,27 +9,26 @@ from __future__ import annotations -from typing import Optional, Tuple +from typing import TYPE_CHECKING, Optional, Tuple -import torch import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute from cutlass import Int32, const_expr -from cudnn.deepseek_sparse_attention.utils.compiler import compile_options -from cudnn.deepseek_sparse_attention.utils.runtime import ( - device_major, - resolve_stream as _resolve_stream, -) -from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor as _to_cute +if TYPE_CHECKING: + import torch WARPS_PER_CTA = 4 ROWS_PER_CTA = WARPS_PER_CTA # One warp per row. def is_available() -> bool: + import torch + + from cudnn.deepseek_sparse_attention.utils.runtime import device_major + return torch.cuda.is_available() and device_major() >= 9 @@ -124,6 +123,8 @@ def kernel(self, mIn: cute.Tensor, mOut: cute.Tensor, mLen: cute.Tensor, rows: I def _compile_or_fetch(key, kernel_obj, *args): + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + if key not in _compile_cache: _compile_cache[key] = cute.compile(kernel_obj, *args, options=compile_options("--opt-level 3")) return _compile_cache[key] @@ -139,6 +140,11 @@ def compactify( batch-major to (B*S, K), matching DSA sparse attention's flat top-K layout. Returns (indices, topk_length). """ + import torch + + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream as _resolve_stream + from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor as _to_cute + if not is_available(): raise RuntimeError("compactify requires SM90+ CUDA device") if idxs.dtype != torch.int32: diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_decode_varlen.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_decode_varlen.py index 41ff88173..593b2dd59 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_decode_varlen.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_decode_varlen.py @@ -14,17 +14,12 @@ """SM90+ CuTe DSL indexer top-K decode kernel.""" -import math - import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute import cutlass.utils as utils -import torch from cutlass.utils.distributed import atomicAdd -from cudnn.deepseek_sparse_attention.utils.compiler import compile_options - from .block_scan import block_prefix_sum_kernel from .indexer_top_k_varlen_util import IndexerTopKKernelVarlen @@ -573,13 +568,6 @@ def _next_positive_power_of_2(x: int) -> int: return 1 << (x - 1).bit_length() -_TORCH_TO_CUTLASS_DTYPE = { - torch.float16: cutlass.Float16, - torch.bfloat16: cutlass.BFloat16, - torch.float32: cutlass.Float32, -} - - def _bucket_num_cols(num_cols: int) -> int: """Bucket num_cols to the next power of 2 for compilation caching. @@ -603,8 +591,19 @@ def cute_dsl_topk_wrapper( load_balance=False, num_copy_bits=256, ): + import math + + import torch + + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + + torch_to_cutlass_dtype = { + torch.float16: cutlass.Float16, + torch.bfloat16: cutlass.BFloat16, + torch.float32: cutlass.Float32, + } torch_dtype = input_values.dtype - dtype = _TORCH_TO_CUTLASS_DTYPE[torch_dtype] + dtype = torch_to_cutlass_dtype[torch_dtype] num_rows, num_cols = input_values.shape bucketed_num_cols = _bucket_num_cols(num_cols) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py index 57c7f9f37..ac72aa09c 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py @@ -12,15 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + +from typing import TYPE_CHECKING import cutlass import cutlass.cute as cute -import torch from cutlass._mlir.dialects import llvm from cutlass.utils.distributed import atomicAdd from .block_scan import block_prefix_sum_kernel, fence_acq_rel_cta +if TYPE_CHECKING: + import torch + """ top-k varlen utils. could be used by prefill and decode phase. """ @@ -1089,6 +1094,8 @@ def create_random_logits( Returns: Tensor of shape (num_rows, max_row_length) with random values and -inf padding """ + import torch + torch.manual_seed(seed) torch.cuda.manual_seed(seed) num_rows = row_starts.shape[0] @@ -1145,6 +1152,8 @@ def compare_top_k_results( Returns: True if results match within tolerance, False otherwise """ + import torch + num_rows = cuda_indices.shape[0] # Calculate valid lengths for each row (vectorized) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py new file mode 100644 index 000000000..83d62194c --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/jax.py @@ -0,0 +1,397 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX APIs for DeepSeek indexer top-K and index utilities.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._jax import JaxApiBase, TupleDict +from ..._jax.compiler import compile_options_for_target +from .op import INT32_MAX, IndexerTopKOp, SUPPORTED_COMPUTE_CAPABILITIES, SUPPORTED_INPUT_DTYPES + + +class IndexerTopK(JaxApiBase): + """JAX callable specialized from top-K input metadata.""" + + def __init__( + self, + sample_input_values: Any, + sample_seq_lens: Any, + top_k: int, + next_n: int = 1, + return_val: bool = True, + num_copy_bits: int = 256, + *, + target_compute_capability: int | None = None, + ) -> None: + self.target_compute_capability = self._resolve_compute_capability( + target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "IndexerTopK", + ) + self.input_desc = self._to_tensor_desc(sample_input_values, "sample_input_values") + self.seq_lens_desc = self._to_tensor_desc(sample_seq_lens, "sample_seq_lens") + self.top_k = int(top_k) + self.next_n = int(next_n) + self.return_val = bool(return_val) + self.num_copy_bits = int(num_copy_bits) + + if self.input_desc.ndim != 2: + raise ValueError(f"input_values must have rank 2, got {self.input_desc.shape}") + if self.seq_lens_desc.ndim != 1: + raise ValueError(f"seq_lens must have rank 1, got {self.seq_lens_desc.shape}") + if self.input_desc.cudnn_dtype not in SUPPORTED_INPUT_DTYPES: + raise ValueError(f"input_values must have dtype float16, bfloat16, or float32, got {self.input_desc.dtype}") + if self.top_k <= 0: + raise ValueError(f"top_k must be positive, got {self.top_k}") + output_shape = (self.input_desc.shape[0], self.top_k) + self.indices_desc = self.input_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=output_shape, + name="indices", + ) + self.values_desc = ( + self.input_desc.compact_like( + cudnn_dtype=self.input_desc.cudnn_dtype, + shape=output_shape, + name="values", + ) + if self.return_val + else None + ) + buffer_count = 2 if self.input_desc.cudnn_dtype == data_type.FLOAT else 1 + self.workspace_desc = self.input_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=(self.input_desc.shape[0], buffer_count, self.input_desc.shape[1]), + name="extra_buffer", + ) + self._op = IndexerTopKOp( + input_values=self.input_desc, + seq_lens=self.seq_lens_desc, + output_indices=self.indices_desc, + output_values=self.values_desc, + workspace=self.workspace_desc, + top_k=self.top_k, + next_n=self.next_n, + return_val=self.return_val, + num_copy_bits=self.num_copy_bits, + target_compute_capability=self.target_compute_capability, + ) + + def check_support(self) -> bool: + self._op.check_support() + workspace_elements = 1 + for extent in self.workspace_desc.shape: + workspace_elements *= extent + if workspace_elements > INT32_MAX: + raise NotImplementedError( + "JAX IndexerTopK does not support the Torch row-chunking fallback when " + f"the workspace exceeds {INT32_MAX} elements (got {workspace_elements})" + ) + return True + + def __call__(self, input_values: Any, seq_lens: Any) -> TupleDict: + self.check_support() + + output_descs = (self.indices_desc,) if self.values_desc is None else (self.indices_desc, self.values_desc) + results = self._call_kernel( + (input_values, seq_lens), + launch=self._launch_kernel, + output_descs=output_descs, + input_descs=(self.input_desc, self.seq_lens_desc), + workspace_descs=(self.workspace_desc,), + compile_options=compile_options_for_target(self.target_compute_capability), + ) + indices = results[0] + values = results[1] if self.return_val else None + return TupleDict(indices=indices, values=values) + + def _launch_kernel( + self, + stream: Any, + input_values: Any, + seq_lens: Any, + *buffers: Any, + ) -> None: + import cutlass + + from .indexer_top_k_decode_varlen import IndexerTopKKernelVarlenDecode + + dtype_by_cudnn = { + data_type.HALF: cutlass.Float16, + data_type.BFLOAT16: cutlass.BFloat16, + data_type.FLOAT: cutlass.Float32, + } + outputs = buffers[:-1] + extra_buffer = buffers[-1] + output_indices = outputs[0] + output_values = outputs[1] if self.return_val else None + resolved = ( + self._op.max_num_cols, + self._op.large_occupancy, + ) + if any(value is None for value in resolved): + raise RuntimeError("IndexerTopK launch configuration was not resolved by check_support()") + max_num_cols, large_occupancy = resolved + kernel = IndexerTopKKernelVarlenDecode( + dtype_by_cudnn[self.input_desc.cudnn_dtype], + max_num_cols, + self.top_k, + self.next_n, + num_copy_bits=self.num_copy_bits, + return_val=self.return_val, + large_occupancy=large_occupancy, + ) + kernel( + input_values, + None, + extra_buffer, + None, + seq_lens, + output_indices, + output_values, + stream, + enable_persistent_dynamic_scheduling=False, + min_blocks_per_mp=1, + ) + + +class _LocalToGlobal(JaxApiBase): + def __init__( + self, + sample_local_indices: Any, + seqlen_k: int, + *, + sample_cu_seqlens_q: Any | None = None, + sample_cu_seqlens_k: Any | None = None, + target_compute_capability: int | None = None, + ) -> None: + if (sample_cu_seqlens_q is None) != (sample_cu_seqlens_k is None): + raise ValueError("Packed local-to-global requires both cumulative sequence tensors") + self.target_compute_capability = self._resolve_compute_capability( + target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "local_to_global", + ) + self.local_desc = self._to_tensor_desc(sample_local_indices, "local_indices") + self.cu_q_desc = None if sample_cu_seqlens_q is None else self._to_tensor_desc(sample_cu_seqlens_q, "cu_seqlens_q") + self.cu_k_desc = None if sample_cu_seqlens_k is None else self._to_tensor_desc(sample_cu_seqlens_k, "cu_seqlens_k") + self.is_varlen = self.cu_q_desc is not None + self.seqlen_k = int(seqlen_k) + self.output_desc = self.local_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=self.local_desc.shape, + name="indices", + ) + + def check_support(self) -> bool: + expected_rank = 2 if self.is_varlen else 3 + if self.local_desc.ndim != expected_rank: + raise ValueError(f"local_indices must have rank {expected_rank}, got {self.local_desc.shape}") + if self.local_desc.cudnn_dtype not in (data_type.INT32, data_type.INT64): + raise ValueError(f"local_indices must have dtype int32 or int64, got {self.local_desc.dtype}") + if any(extent <= 0 for extent in self.local_desc.shape): + raise ValueError(f"local_indices dimensions must be positive, got {self.local_desc.shape}") + if not self.local_desc.is_compact(tuple(reversed(range(expected_rank)))): + raise ValueError("local_indices must be row-major contiguous") + if self.seqlen_k <= 0: + raise ValueError(f"seqlen_k must be positive, got {self.seqlen_k}") + if self.is_varlen: + if self.cu_q_desc is None or self.cu_k_desc is None: + raise RuntimeError("Packed cumulative sequence descriptors were not configured") + for desc, name in ((self.cu_q_desc, "cu_seqlens_q"), (self.cu_k_desc, "cu_seqlens_k")): + if desc.ndim != 1 or desc.cudnn_dtype != data_type.INT32 or not desc.is_compact((0,)): + raise ValueError(f"{name} must be a contiguous rank-1 int32 tensor") + if self.cu_q_desc.shape != self.cu_k_desc.shape: + raise ValueError("cu_seqlens_q and cu_seqlens_k must have the same shape") + if self.cu_q_desc.shape[0] < 2: + raise ValueError("Packed cumulative sequence tensors must contain at least two entries") + return True + + def __call__( + self, + local_indices: Any, + *, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + ) -> TupleDict: + self.check_support() + inputs = [local_indices] + input_descs = [self.local_desc] + if self.is_varlen: + if cu_seqlens_q is None or cu_seqlens_k is None: + raise ValueError("Packed local-to-global requires both cumulative sequence tensors") + if self.cu_q_desc is None or self.cu_k_desc is None: + raise RuntimeError("Packed cumulative sequence descriptors were not configured") + inputs.extend((cu_seqlens_q, cu_seqlens_k)) + input_descs.extend((self.cu_q_desc, self.cu_k_desc)) + elif cu_seqlens_q is not None or cu_seqlens_k is not None: + raise ValueError("Fixed local-to-global does not accept cumulative sequence tensors") + + (result,) = self._call_kernel( + tuple(inputs), + launch=self._launch_kernel, + output_descs=(self.output_desc,), + input_descs=tuple(input_descs), + compile_options=compile_options_for_target(self.target_compute_capability, "--opt-level 3"), + ) + return TupleDict(indices=result) + + def _launch_kernel(self, stream, *arguments) -> None: + from cutlass import Int32 + + from .local_to_global_dsl import LocalToGlobalTopK + + *inputs, output = arguments + local_indices, *optional = inputs + if self.is_varlen: + cu_q, cu_k = optional + else: + cu_q = cu_k = None + LocalToGlobalTopK(is_varlen=self.is_varlen)( + local_indices, + output, + Int32(self.seqlen_k), + cu_q, + cu_k, + stream, + ) + + +class _Compactify(JaxApiBase): + def __init__(self, sample_indices: Any, *, target_compute_capability: int | None = None) -> None: + self.target_compute_capability = self._resolve_compute_capability( + target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "compactify", + ) + self.input_desc = self._to_tensor_desc(sample_indices, "indices") + if self.input_desc.ndim != 2: + raise ValueError(f"Compactify expects flattened rank-2 input, got {self.input_desc.shape}") + rows, cols = self.input_desc.shape + self.rows = rows + self.cols = cols + self.output_desc = self.input_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=(rows, cols), + name="indices", + ) + self.length_desc = self.input_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=(rows,), + name="topk_length", + ) + + def check_support(self) -> bool: + if self.input_desc.cudnn_dtype != data_type.INT32: + raise ValueError(f"indices must have dtype int32, got {self.input_desc.dtype}") + if self.rows <= 0 or self.cols <= 0: + raise ValueError(f"indices dimensions must be positive, got {self.input_desc.shape}") + if not self.input_desc.is_compact((1, 0)): + raise ValueError("indices must be row-major contiguous") + return True + + def __call__(self, indices: Any) -> TupleDict: + self.check_support() + compact_indices, topk_length = self._call_kernel( + (indices,), + launch=self._launch_kernel, + output_descs=(self.output_desc, self.length_desc), + input_descs=(self.input_desc,), + compile_options=compile_options_for_target(self.target_compute_capability, "--opt-level 3"), + ) + return TupleDict(indices=compact_indices, topk_length=topk_length) + + def _launch_kernel(self, stream, indices, compact_indices, topk_length) -> None: + from cutlass import Int32 + + from .compactify import CompactifyKernel + + CompactifyKernel(cols=self.cols)(indices, compact_indices, topk_length, Int32(self.rows), stream) + + +@partial( + jax.jit, + static_argnames=("top_k", "next_n", "return_val", "num_copy_bits", "target_compute_capability"), +) +def indexer_top_k_wrapper( + input_values: Any, + seq_lens: Any, + top_k: int, + next_n: int = 1, + return_val: bool = True, + num_copy_bits: int = 256, + *, + target_compute_capability: int | None = None, +) -> TupleDict: + """Select per-row top-K indices and optional values from JAX. + + Runtime lengths are consumed on device and cannot be inspected while + tracing. Every ``seq_lens[b]`` must satisfy + ``top_k + next_n - 1 <= seq_lens[b] <= input_values.shape[1]``. + """ + + return IndexerTopK( + input_values, + seq_lens, + top_k, + next_n=next_n, + return_val=return_val, + num_copy_bits=num_copy_bits, + target_compute_capability=target_compute_capability, + )(input_values, seq_lens) + + +@partial(jax.jit, static_argnames=("seqlen_k", "target_compute_capability")) +def local_to_global_wrapper( + local_indices: Any, + seqlen_k: int, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + *, + target_compute_capability: int | None = None, +) -> TupleDict: + """Convert local top-K indices to the global flattened KV index space.""" + + return _LocalToGlobal( + local_indices, + seqlen_k, + sample_cu_seqlens_q=cu_seqlens_q, + sample_cu_seqlens_k=cu_seqlens_k, + target_compute_capability=target_compute_capability, + )( + local_indices, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + ) + + +@partial(jax.jit, static_argnames=("target_compute_capability",)) +def compactify_wrapper(indices: Any, *, target_compute_capability: int | None = None) -> TupleDict: + """Pack nonnegative indices to the front of every flattened row.""" + + if indices.ndim not in (2, 3): + raise ValueError(f"indices must have rank 2 or 3, got shape {indices.shape}") + rows = 1 + for extent in indices.shape[:-1]: + rows *= extent + flat_indices = jnp.reshape(indices, (rows, indices.shape[-1])) + return _Compactify( + flat_indices, + target_compute_capability=target_compute_capability, + )(flat_indices) + + +__all__ = [ + "IndexerTopK", + "compactify_wrapper", + "indexer_top_k_wrapper", + "local_to_global_wrapper", +] diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/local_to_global_dsl.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/local_to_global_dsl.py index bfa61b73f..eaa4ee637 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/local_to_global_dsl.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/local_to_global_dsl.py @@ -9,19 +9,16 @@ from __future__ import annotations -import torch +from typing import TYPE_CHECKING + import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute from cutlass import Int32, Int64, const_expr -from cudnn.deepseek_sparse_attention.utils.compiler import compile_options -from cudnn.deepseek_sparse_attention.utils.runtime import ( - device_major as _get_device_capability, - resolve_stream, -) -from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor as _to_cute_tensor +if TYPE_CHECKING: + import torch class LocalToGlobalTopK: @@ -123,7 +120,11 @@ def is_available() -> bool: # The kernel is arch-agnostic CuTe DSL: thread/block indexing, integer # ALU, and plain global mem load/store — no TMA / TMEM / tcgen05 / wgmma. # SM90+ is enough. - return torch.cuda.is_available() and _get_device_capability() >= 9 + import torch + + from cudnn.deepseek_sparse_attention.utils.runtime import device_major + + return torch.cuda.is_available() and device_major() >= 9 def local_to_global( @@ -133,6 +134,12 @@ def local_to_global( cu_seqlens_k: torch.Tensor | None = None, stream: cuda.CUstream | None = None, ) -> torch.Tensor: + import torch + + from cudnn.deepseek_sparse_attention.utils.compiler import compile_options + from cudnn.deepseek_sparse_attention.utils.runtime import resolve_stream + from cudnn.deepseek_sparse_attention.utils.tensor_conversion import to_cute_tensor as _to_cute_tensor + if not is_available(): raise RuntimeError("CuTe DSL local_to_global requires an SM90+ CUDA device") if not local_indices.is_cuda or not local_indices.is_contiguous(): diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/op.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/op.py new file mode 100644 index 000000000..543180ffe --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/op.py @@ -0,0 +1,194 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral validation for the DeepSeek indexer top-K operation.""" + +from __future__ import annotations + +from typing import Any + +from ... import data_type +from ...common.op import Op +from ...common.tensor_desc import TensorDesc + +INT32_MAX = (1 << 31) - 1 +SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) +SUPPORTED_INPUT_DTYPES = (data_type.HALF, data_type.BFLOAT16, data_type.FLOAT) + + +def bucket_num_cols(num_cols: int) -> int: + if num_cols <= 0: + return 1 + return 1 << (num_cols - 1).bit_length() + + +def _dtype_bits(dtype: data_type) -> int: + return 32 if dtype == data_type.FLOAT else 16 + + +def _require_compact(desc: TensorDesc[Any], label: str) -> None: + expected_order = tuple(reversed(range(desc.ndim))) + if not desc.is_compact(expected_order): + raise ValueError(f"{label} must be row-major contiguous, got stride {desc.stride}") + + +def _num_threads_per_cta( + *, + dtype: data_type, + max_num_cols: int, + num_copy_bits: int, + large_occupancy: bool, +) -> int: + if large_occupancy: + return 512 + vec_size = num_copy_bits // _dtype_bits(dtype) + if dtype == data_type.FLOAT: + if max_num_cols >= vec_size * 1024: + return 1024 + if 2048 < max_num_cols < 8192: + return 512 + return 256 + if max_num_cols >= 43008: + return 1024 + if 4096 < max_num_cols < 43008: + return 512 + return 256 + + +class IndexerTopKOp(Op): + """Complete input, output, and workspace signature for indexer top-K.""" + + def __init__( + self, + *, + input_values: TensorDesc[Any], + seq_lens: TensorDesc[Any], + output_indices: TensorDesc[Any], + output_values: TensorDesc[Any] | None, + workspace: TensorDesc[Any], + top_k: int, + next_n: int = 1, + return_val: bool = True, + num_copy_bits: int = 256, + target_compute_capability: int = 90, + ) -> None: + for name, desc in ( + ("input_values", input_values), + ("seq_lens", seq_lens), + ("output_indices", output_indices), + ("workspace", workspace), + ): + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + if output_values is not None and not isinstance(output_values, TensorDesc): + raise TypeError(f"output_values must be a TensorDesc or None, got {type(output_values).__name__}") + + self.input_values = input_values + self.seq_lens = seq_lens + self.output_indices = output_indices + self.output_values = output_values + self.workspace = workspace + self.top_k = int(top_k) + self.next_n = int(next_n) + self.return_val = bool(return_val) + self.num_copy_bits = int(num_copy_bits) + self.target_compute_capability = int(target_compute_capability) + + self.num_rows: int | None = None + self.num_cols: int | None = None + self.batch_size: int | None = None + self.buffer_count: int | None = None + self.dtype_bits: int | None = None + self.max_num_cols: int | None = None + self.large_occupancy: bool | None = None + self.num_threads_per_cta: int | None = None + + def check_support(self) -> bool: + if self.target_compute_capability not in SUPPORTED_COMPUTE_CAPABILITIES: + raise ValueError("target_compute_capability must be one of " f"{SUPPORTED_COMPUTE_CAPABILITIES}, got {self.target_compute_capability}") + if self.input_values.ndim != 2: + raise ValueError(f"input_values must be 2-D, got {self.input_values.shape}") + if self.seq_lens.ndim != 1: + raise ValueError(f"seq_lens must be 1-D, got {self.seq_lens.shape}") + if self.input_values.cudnn_dtype not in SUPPORTED_INPUT_DTYPES: + raise ValueError(f"input_values must have dtype float16, bfloat16, or float32, got {self.input_values.dtype}") + if self.seq_lens.cudnn_dtype != data_type.INT32: + raise ValueError(f"seq_lens must have dtype int32, got {self.seq_lens.dtype}") + _require_compact(self.input_values, "input_values") + _require_compact(self.seq_lens, "seq_lens") + + num_rows, num_cols = self.input_values.shape + (batch_size,) = self.seq_lens.shape + if num_rows <= 0 or num_cols <= 0: + raise ValueError(f"input_values dimensions must be positive, got {(num_rows, num_cols)}") + if batch_size <= 0: + raise ValueError("seq_lens must not be empty") + if self.next_n <= 0: + raise ValueError(f"next_n must be positive, got {self.next_n}") + if num_rows != batch_size * self.next_n: + raise ValueError(f"num_rows ({num_rows}) must equal seq_lens.size * next_n " f"({batch_size} * {self.next_n} = {batch_size * self.next_n})") + if self.top_k <= 0 or self.top_k > min(2048, num_cols): + raise ValueError(f"top_k must be in (0, min(2048, num_cols={num_cols})], got {self.top_k}") + if self.num_copy_bits <= 0 or self.num_copy_bits % 8: + raise ValueError(f"num_copy_bits must be a positive whole-byte width, got {self.num_copy_bits}") + copy_bytes = self.num_copy_bits // 8 + if copy_bytes & (copy_bytes - 1): + raise ValueError(f"num_copy_bits must describe a power-of-two byte alignment, got {self.num_copy_bits}") + + dtype_bits = _dtype_bits(self.input_values.cudnn_dtype) + if self.num_copy_bits % dtype_bits: + raise ValueError(f"num_copy_bits ({self.num_copy_bits}) must be divisible by input dtype width ({dtype_bits})") + + output_shape = (num_rows, self.top_k) + if self.output_indices.shape != output_shape or self.output_indices.cudnn_dtype != data_type.INT32: + raise ValueError(f"output_indices must be int32 with shape {output_shape}, got {self.output_indices.shape}") + _require_compact(self.output_indices, "output_indices") + if self.return_val: + if self.output_values is None: + raise ValueError("output_values is required when return_val=True") + if self.output_values.shape != output_shape or self.output_values.cudnn_dtype != self.input_values.cudnn_dtype: + raise ValueError(f"output_values must have shape {output_shape} and match input_values dtype") + _require_compact(self.output_values, "output_values") + elif self.output_values is not None: + raise ValueError("output_values must be None when return_val=False") + + buffer_count = 2 if self.input_values.cudnn_dtype == data_type.FLOAT else 1 + workspace_shape = (num_rows, buffer_count, num_cols) + if self.workspace.shape != workspace_shape or self.workspace.cudnn_dtype != data_type.INT32: + raise ValueError(f"workspace must be int32 with shape {workspace_shape}, got {self.workspace.shape}") + _require_compact(self.workspace, "workspace") + max_num_cols = bucket_num_cols(num_cols) + large_occupancy = num_rows > 148 + num_threads_per_cta = _num_threads_per_cta( + dtype=self.input_values.cudnn_dtype, + max_num_cols=max_num_cols, + num_copy_bits=self.num_copy_bits, + large_occupancy=large_occupancy, + ) + vector_size = min( + self.top_k, + (self.top_k + num_threads_per_cta - 1) // num_threads_per_cta, + self.num_copy_bits // dtype_bits, + 2, + ) + if self.top_k % vector_size: + raise ValueError(f"top_k ({self.top_k}) must be divisible by the selected output vector width ({vector_size})") + + self.num_rows = num_rows + self.num_cols = num_cols + self.batch_size = batch_size + self.buffer_count = buffer_count + self.dtype_bits = dtype_bits + self.max_num_cols = max_num_cols + self.large_occupancy = large_occupancy + self.num_threads_per_cta = num_threads_per_cta + return True + + +__all__ = [ + "INT32_MAX", + "IndexerTopKOp", + "SUPPORTED_COMPUTE_CAPABILITIES", + "SUPPORTED_INPUT_DTYPES", + "bucket_num_cols", +] diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/__init__.py b/python/cudnn/deepseek_sparse_attention/score_recompute/__init__.py index d5434b6ff..7838de720 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/__init__.py @@ -1,21 +1,38 @@ -from .api import ( - SparseIndexerScoreRecompute, - sparse_indexer_score_recompute_wrapper, - SparseAttnScoreRecompute, - sparse_attn_score_recompute_wrapper, - DenseIndexerScoreRecompute, - dense_indexer_score_recompute_wrapper, - DenseAttnScoreRecompute, - dense_attn_score_recompute_wrapper, -) +"""Lazy Torch API and framework-neutral score-recompute exports. + +JAX APIs are exported through :mod:`cudnn.jax`. +""" -__all__ = [ - "SparseIndexerScoreRecompute", - "sparse_indexer_score_recompute_wrapper", - "SparseAttnScoreRecompute", - "sparse_attn_score_recompute_wrapper", - "DenseIndexerScoreRecompute", - "dense_indexer_score_recompute_wrapper", - "DenseAttnScoreRecompute", - "dense_attn_score_recompute_wrapper", -] +from ...common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "config": ( + "DenseScoreKernelConfig", + "SparseScoreKernelConfig", + "dispatch_sparse_attn_tile_params", + "resolve_dense_score_kernel_config", + "resolve_dense_score_smem_config", + "resolve_sparse_score_kernel_config", + "resolve_sparse_score_smem_config", + ), + "op": ( + "DenseScoreRecomputeOp", + "DenseScoreSm90Config", + "SparseScoreRecomputeOp", + "SparseScoreSm90Config", + ), + "api": ( + "SparseIndexerScoreRecompute", + "sparse_indexer_score_recompute_wrapper", + "SparseAttnScoreRecompute", + "sparse_attn_score_recompute_wrapper", + "DenseIndexerScoreRecompute", + "dense_indexer_score_recompute_wrapper", + "DenseAttnScoreRecompute", + "dense_attn_score_recompute_wrapper", + ), + }, + submodules=("api", "config", "op"), +) diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py b/python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py index 1310cc690..7649f98be 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/_interface_sm100.py @@ -10,7 +10,6 @@ from __future__ import annotations -import math from typing import Optional import torch @@ -21,6 +20,12 @@ from .sparse_score_recompute_sm100 import SparseScoreRecomputeSm100 from .dense_score_recompute_sm100 import DenseScoreRecomputeSm100 +from .config import ( + dispatch_sparse_attn_tile_params as _shared_dispatch_sparse_attn_tile_params, + resolve_dense_score_kernel_config, + resolve_dense_score_smem_config, + resolve_sparse_score_smem_config, +) from cudnn.deepseek_sparse_attention.utils.compiler import compile_options from cudnn.deepseek_sparse_attention.utils.runtime import ( device_major as _get_device_capability, @@ -97,29 +102,14 @@ def _sparse_indexer_score_recompute( with _torch_stream_context(current_stream): out = torch.empty((bs, seqlen_q, topk), dtype=torch.float32, device=device) - # Compute kv_stage and topk_in_smem from SMEM budget (SM100: 228 KB) - SM100_SMEM_BYTES = 228 * 1024 - head_dim_padded = int(math.ceil(head_dim / 16) * 16) - sK_per_stage = n_block_size * head_dim_padded * 2 # BF16 - sQ_size = m_block_size * head_dim_padded * 2 # BF16 - sTopkIdx_bytes = topk * 2 * 4 # double-buffer, INT32 - sPerHead_bytes = m_block_size * 2 * 2 # double-buffer, BF16 - smem_fixed = sPerHead_bytes + 2048 # barriers, sScoreAll, alignment - - topk_in_smem = True - smem_overhead = sTopkIdx_bytes + smem_fixed - kv_stage = min(4, max(1, (SM100_SMEM_BYTES - sQ_size - smem_overhead) // sK_per_stage)) - total_smem_est = sQ_size + sK_per_stage * kv_stage + smem_overhead - if total_smem_est > SM100_SMEM_BYTES: - topk_in_smem = False - smem_overhead = smem_fixed - kv_stage = min(4, max(1, (SM100_SMEM_BYTES - sQ_size - smem_overhead) // sK_per_stage)) - total_smem_est = sQ_size + sK_per_stage * kv_stage + smem_overhead - assert total_smem_est <= SM100_SMEM_BYTES, ( - f"SMEM overflow ({total_smem_est} > {SM100_SMEM_BYTES}) even without sTopkIdx: " - f"topk={topk}, head_dim={head_dim}(padded={head_dim_padded}), " - f"m_block={m_block_size}, n_block={n_block_size}, kv_stage={kv_stage}." - ) + kv_stage, topk_in_smem = resolve_sparse_score_smem_config( + score_type="indexer", + head_dim=head_dim, + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=None, + topk=topk, + ) compute_capability = _get_device_capability() @@ -319,30 +309,14 @@ def _sparse_attn_score_recompute( with _torch_stream_context(current_stream): out = torch.empty((bs, seqlen_q, topk), dtype=torch.float32, device=device) - # Compute kv_stage and topk_in_smem from SMEM budget (SM100: 228 KB) - SM100_SMEM_BYTES = 228 * 1024 - head_dim_padded = int(math.ceil(head_dim / 16) * 16) - k_block_size_eff = k_block_size if k_block_size is not None else head_dim_padded - sK_per_stage = n_block_size * k_block_size_eff * 2 # BF16 - sQ_size = m_block_size * head_dim_padded * 2 # BF16 (Q total unchanged) - sTopkIdx_bytes = topk * 2 * 4 # double-buffer, INT32 - sPerHead_bytes = m_block_size * 2 * 4 # double-buffer, FP32 - smem_fixed = sPerHead_bytes + 2048 # barriers, sScoreAll, alignment - - topk_in_smem = True - smem_overhead = sTopkIdx_bytes + smem_fixed - kv_stage = min(4, max(1, (SM100_SMEM_BYTES - sQ_size - smem_overhead) // sK_per_stage)) - total_smem_est = sQ_size + sK_per_stage * kv_stage + smem_overhead - if total_smem_est > SM100_SMEM_BYTES: - topk_in_smem = False - smem_overhead = smem_fixed - kv_stage = min(4, max(1, (SM100_SMEM_BYTES - sQ_size - smem_overhead) // sK_per_stage)) - total_smem_est = sQ_size + sK_per_stage * kv_stage + smem_overhead - assert total_smem_est <= SM100_SMEM_BYTES, ( - f"SMEM overflow ({total_smem_est} > {SM100_SMEM_BYTES}) even without sTopkIdx: " - f"topk={topk}, head_dim={head_dim}(padded={head_dim_padded}), " - f"m_block={m_block_size}, n_block={n_block_size}, kv_stage={kv_stage}." - ) + kv_stage, topk_in_smem = resolve_sparse_score_smem_config( + score_type="attention", + head_dim=head_dim, + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=k_block_size, + topk=topk, + ) compute_capability = _get_device_capability() @@ -434,37 +408,12 @@ def _dispatch_sparse_attn_tile_params( Rules tuned on B200 via tests/sweep_tile_params.py (--compact / default). """ - m = qhead_per_kv_head - - if compact: - # --- Compact tuned configs --- - # hd=512 - if head_dim == 512 and qhead_per_kv_head == 64: - return m, 64, None # kvs=2, ~789 TFLOPS - if head_dim == 512 and qhead_per_kv_head == 128: - return m, 128, 128 # kvs=2, ~844/1004 TFLOPS - # hd=576 - if head_dim == 576 and qhead_per_kv_head == 64: - return m, 64, None # kvs=2, ~827 TFLOPS - if head_dim == 576 and qhead_per_kv_head == 128: - return m, 64, 192 # kvs=2~3, ~677/1054 TFLOPS - return m, 64, None - - # --- Non-compact tuned configs --- - # hd=512 - if head_dim == 512 and qhead_per_kv_head == 64: - return m, 128, 256 # kvs=2, ~800 TFLOPS - if head_dim == 512 and qhead_per_kv_head == 128 and topk <= 512: - return m, 128, 128 # kvs=2, ~1000 TFLOPS - if head_dim == 512 and qhead_per_kv_head == 128: - return m, 128, 256 # kvs=1, ~910 TFLOPS - # hd=576 - if head_dim == 576 and qhead_per_kv_head == 64: - return m, 128, 192 # kvs=3, ~770 TFLOPS - if head_dim == 576 and qhead_per_kv_head == 128: - return m, 128, 192 # kvs=1, ~830 TFLOPS - - return m, 64, None + return _shared_dispatch_sparse_attn_tile_params( + head_dim, + qhead_per_kv_head, + topk, + compact=compact, + ) def sparse_attn_score_recompute( @@ -550,130 +499,23 @@ def sparse_attn_score_recompute( # Dense backward: full KV via TMA, no topk # ============================================================================= -_SM100_SMEM_BYTES = 228 * 1024 - - -def _select_dense_k_block_size(head_dim_padded, m_block_size, n_block_size, per_head_elem_bytes): - """Auto-select k_block_size so that sQ + sK(1 stage) + overhead fits in SMEM. - - When head_dim_padded is too large (e.g. m_block=128, hd=512 -> sQ=128KB, - sK=128KB > 228KB), we reduce k_block_size to shrink sK per stage. - sQ total = m_block * head_dim_padded * 2 (fixed, since Q has num_k_chunks stages). - - k_block_size must be a multiple of 64 and divide head_dim_padded evenly. - """ - sQ_size = m_block_size * head_dim_padded * 2 - sPerHead_bytes = m_block_size * 2 * per_head_elem_bytes - smem_overhead = sPerHead_bytes + 2048 - - avail_for_sK = _SM100_SMEM_BYTES - sQ_size - smem_overhead - max_kbs = max(64, avail_for_sK // (n_block_size * 2)) - max_kbs = (max_kbs // 64) * 64 - - kbs = head_dim_padded - while kbs > max_kbs: - # Try common divisors of head_dim_padded in descending order - found = False - for candidate in range(kbs - 64, 63, -64): - if candidate > 0 and head_dim_padded % candidate == 0: - kbs = candidate - found = True - break - if not found: - kbs = 64 - break - - assert head_dim_padded % kbs == 0, f"Cannot find valid k_block_size: head_dim_padded={head_dim_padded}, " f"m_block={m_block_size}, n_block={n_block_size}" - return kbs - - -def _compute_dense_kv_stage(head_dim_padded, m_block_size, n_block_size, k_block_size_eff, per_head_elem_bytes): - sK_per_stage = n_block_size * k_block_size_eff * 2 - sQ_size = m_block_size * head_dim_padded * 2 - sPerHead_bytes = m_block_size * 2 * per_head_elem_bytes - smem_overhead = sPerHead_bytes + 2048 - return min(4, max(1, (_SM100_SMEM_BYTES - sQ_size - smem_overhead) // sK_per_stage)) - - -# ---- Dispatch helpers (mirror _dispatch_sparse_attn_tile_params) ------------- - - -def _dense_m_block_with_smem_check(qhead_per_kv_head, head_dim, per_head_elem_bytes): - """Return m=qhpkv*2 if SMEM fits, else fall back to m=qhpkv.""" - head_dim_padded = int(math.ceil(head_dim / 16) * 16) - n = 128 - min_kbs = 64 - smem_overhead = 4096 - - m2 = qhead_per_kv_head * 2 - sQ_m2 = m2 * head_dim_padded * 2 - sK_min = n * min_kbs * 2 - sPerHead_m2 = m2 * 2 * per_head_elem_bytes - if sQ_m2 + sK_min + sPerHead_m2 + smem_overhead <= _SM100_SMEM_BYTES: - return m2 - - return qhead_per_kv_head - def _dispatch_dense_attn_tile_params(head_dim, qhead_per_kv_head): - """Select (m_block_size, n_block_size, k_block_size) for dense attention backward. - - n_block_size is always 128 (required for Ld32x32bOp TMEM path). - m_block_size = qhpkv * 2 (2 q_tokens per tile) when SMEM allows, - falling back to qhpkv when it doesn't (e.g. nh_q=128, hd=512). - - Returns (m_block_size, n_block_size, k_block_size) where k_block_size=None - means no head_dim splitting. - - Rules tuned on B200 via tests/sweep_tile_params.py (dense mode, 2q sweep). - k=64 kvs=4 is consistently best or within 1-2% of best across all sizes. - 2q (m=2*qhpkv) gives 35-46% speedup over 1q for nh_q=64. - """ - m = _dense_m_block_with_smem_check(qhead_per_kv_head, head_dim, per_head_elem_bytes=4) - n = 128 - - # hd=512, nh_q=64: m=128(2q) k=64 kvs=4 ~1600 TFLOPS (16K best, +46% vs 1q) - # hd=512, nh_q=128: m=128(1q) k=64 kvs=4 ~1512 TFLOPS (m=256 doesn't fit SMEM) - # hd=576, nh_q=64: m=128(2q) k=64 kvs=4 ~1601 TFLOPS (16K best, +41% vs 1q) - # hd=576, nh_q=128: m=128(1q) k=64 kvs=4 ~1527 TFLOPS (m=256 doesn't fit SMEM) - if head_dim in (512, 576): - return m, n, 64 - - # Fallback: auto-select from SMEM budget - head_dim_padded = int(math.ceil(head_dim / 16) * 16) - k = _select_dense_k_block_size(head_dim_padded, m, n, per_head_elem_bytes=4) - if k == head_dim_padded: - k = None - return m, n, k + config = resolve_dense_score_kernel_config( + score_type="attention", + head_dim=head_dim, + qhead_per_kv_head=qhead_per_kv_head, + ) + return config.m_block_size, config.n_block_size, config.k_block_size def _dispatch_dense_indexer_tile_params(head_dim, qhead_per_kv_head): - """Select (m_block_size, n_block_size, k_block_size) for dense indexer backward. - - m_block_size = qhpkv * 2 (2 q_tokens per tile) when SMEM allows, - falling back to qhpkv when it doesn't. - - Returns (m_block_size, n_block_size, k_block_size) where k_block_size=None - means no head_dim splitting. - - Rules tuned on B200 via tests/sweep_tile_params.py (dense mode, 2q sweep). - 2q with k=full is consistently best for all tested configs. - nh_q=32: 2q gives ~68% speedup; nh_q=64: 2q gives ~54% speedup. - """ - m = _dense_m_block_with_smem_check(qhead_per_kv_head, head_dim, per_head_elem_bytes=2) - n = 128 - - # hd=128, nh_q=32: m=64(2q) k=full kvs=4 ~541 TFLOPS (16K best, +71% vs 1q) - # hd=128, nh_q=64: m=128(2q) k=full kvs=4 ~871 TFLOPS (16K best, +56% vs 1q) - if head_dim == 128: - return m, n, None - - # Fallback: auto-select from SMEM budget - head_dim_padded = int(math.ceil(head_dim / 16) * 16) - k = _select_dense_k_block_size(head_dim_padded, m, n, per_head_elem_bytes=2) - if k == head_dim_padded: - k = None - return m, n, k + config = resolve_dense_score_kernel_config( + score_type="indexer", + head_dim=head_dim, + qhead_per_kv_head=qhead_per_kv_head, + ) + return config.m_block_size, config.n_block_size, config.k_block_size # ---- Internal dense functions (explicit tile params) ------------------------ @@ -763,10 +605,13 @@ def _dense_indexer_score_recompute( with _torch_stream_context(current_stream): denom_out = torch.empty(denom_shape, dtype=torch.float32, device=device) - head_dim_padded = int(math.ceil(head_dim / 16) * 16) - if k_block_size is None: - k_block_size = _select_dense_k_block_size(head_dim_padded, m_block_size, n_block_size, per_head_elem_bytes=2) - kv_stage = _compute_dense_kv_stage(head_dim_padded, m_block_size, n_block_size, k_block_size, per_head_elem_bytes=2) + k_block_size, kv_stage = resolve_dense_score_smem_config( + score_type="indexer", + head_dim=head_dim, + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=k_block_size, + ) compute_capability = _get_device_capability() if compute_capability < 10: @@ -943,10 +788,13 @@ def _dense_attn_score_recompute( with _torch_stream_context(current_stream): denom_out = torch.empty(denom_shape, dtype=torch.float32, device=device) - head_dim_padded = int(math.ceil(head_dim / 16) * 16) - if k_block_size is None: - k_block_size = _select_dense_k_block_size(head_dim_padded, m_block_size, n_block_size, per_head_elem_bytes=4) - kv_stage = _compute_dense_kv_stage(head_dim_padded, m_block_size, n_block_size, k_block_size, per_head_elem_bytes=4) + k_block_size, kv_stage = resolve_dense_score_smem_config( + score_type="attention", + head_dim=head_dim, + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=k_block_size, + ) compute_capability = _get_device_capability() if compute_capability < 10: diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/api.py b/python/cudnn/deepseek_sparse_attention/score_recompute/api.py index b8d044063..6b31557d4 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/api.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/api.py @@ -17,18 +17,31 @@ from cudnn.api_base import APIBase, TupleDict from . import _interface_sm100 as _iface_sm100 +from .op import DenseScoreRecomputeOp, SparseScoreRecomputeOp # --------------------------------------------------------------------------- # Base helpers # --------------------------------------------------------------------------- -def _check_score_arch(api: APIBase) -> None: - major, _ = torch.cuda.get_device_capability() +def _check_score_arch(api: APIBase) -> int: + major, minor = torch.cuda.get_device_capability() api._runtime_error_if( major != 9 and major < 10, f"{type(api).__name__} requires SM90 or SM100+ compute capability, found SM{major}", ) + return major * 10 + minor + + +def _compact_for_kernel(desc): + if desc is None: + return None + return desc.compact_like( + cudnn_dtype=desc.cudnn_dtype, + shape=desc.shape, + name=desc.name, + init_value=desc.init_value, + ) class _ScoreRecomputeBase(APIBase): @@ -52,36 +65,6 @@ def compile(self) -> None: self._compiled_kernel = True -def _check_sparse_score_shapes( - api: APIBase, - q_desc, - k_desc, - aux_desc, - topk_desc, - out_desc, - topk_length_desc, - aux_name: str, - qhead_per_kv_head: Optional[int], -) -> None: - api._value_error_if(q_desc.ndim != 4, f"Q must be 4-D (B, S_q, H_q, D), got {q_desc.shape}") - api._value_error_if(k_desc.ndim != 3, f"K must be 3-D (B, S_k, D) MQA, got {k_desc.shape}") - api._value_error_if(aux_desc.ndim != 3, f"{aux_name} must be 3-D (B, S_q, H_q), got {aux_desc.shape}") - api._value_error_if(topk_desc.ndim != 3, f"topk_indices must be 3-D (B, S_q, topk), got {topk_desc.shape}") - api._value_error_if(out_desc.ndim != 3, f"out must be 3-D (B, S_q, topk), got {out_desc.shape}") - - b, s_q, h_q, d = q_desc.shape - api._value_error_if(k_desc.shape[0] != b, f"K batch dim {k_desc.shape[0]} must match Q batch dim {b}") - api._value_error_if(k_desc.shape[2] != d, f"K head dim {k_desc.shape[2]} must match Q head dim {d}") - api._value_error_if(aux_desc.shape != (b, s_q, h_q), f"{aux_name} shape must be {(b, s_q, h_q)}, got {aux_desc.shape}") - api._value_error_if(topk_desc.shape[:2] != (b, s_q), f"topk_indices leading dims must be {(b, s_q)}, got {topk_desc.shape[:2]}") - api._value_error_if(out_desc.shape != topk_desc.shape, f"out shape must match topk_indices shape {topk_desc.shape}, got {out_desc.shape}") - if qhead_per_kv_head is not None: - api._value_error_if(qhead_per_kv_head != h_q, f"qhead_per_kv_head must equal H_q ({h_q}) for MQA sparse score, got {qhead_per_kv_head}") - if topk_length_desc is not None: - api._check_dtype(topk_length_desc, torch.int32, name="topk_length") - api._value_error_if(topk_length_desc.shape != (b, s_q), f"topk_length must be shape {(b, s_q)}, got {topk_length_desc.shape}") - - def _check_dense_score_shapes( api: APIBase, q_desc, @@ -154,25 +137,25 @@ def __init__( self.topk_length_desc = self._make_tensor_desc(sample_topk_length, name="sample_topk_length") self.qhead_per_kv_head = qhead_per_kv_head self.topk_indices_global = bool(topk_indices_global) + self._op = None def check_support(self) -> bool: - _check_score_arch(self) - self._check_dtype(self.q_desc, torch.bfloat16, name="Q") - self._check_dtype(self.k_desc, torch.bfloat16, name="K") - self._check_dtype(self.w_desc, torch.bfloat16, name="W") - self._check_dtype(self.topk_desc, torch.int32, name="topk_indices") - self._check_dtype(self.out_desc, torch.float32, name="out") - _check_sparse_score_shapes( - self, - self.q_desc, - self.k_desc, - self.w_desc, - self.topk_desc, - self.out_desc, - self.topk_length_desc, - "W", - self.qhead_per_kv_head, - ) + target_compute_capability = _check_score_arch(self) + if self._op is None: + self._op = SparseScoreRecomputeOp( + q=_compact_for_kernel(self.q_desc), + k=_compact_for_kernel(self.k_desc), + per_head=_compact_for_kernel(self.w_desc), + topk_indices=_compact_for_kernel(self.topk_desc), + output=_compact_for_kernel(self.out_desc), + topk_length=_compact_for_kernel(self.topk_length_desc), + score_type="indexer", + softmax_scale=1.0, + qhead_per_kv_head=self.qhead_per_kv_head, + topk_indices_global=self.topk_indices_global, + target_compute_capability=target_compute_capability, + ) + self._op.check_support() self._is_supported = True return True @@ -269,7 +252,7 @@ def sparse_indexer_score_recompute_wrapper( qhead_per_kv_head=qhead_per_kv_head, topk_indices_global=topk_indices_global, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_SparseIndexerScoreRecomputeObjects[key] = obj @@ -320,25 +303,25 @@ def __init__( self.softmax_scale = float(softmax_scale) self.qhead_per_kv_head = qhead_per_kv_head self.topk_indices_global = bool(topk_indices_global) + self._op = None def check_support(self) -> bool: - _check_score_arch(self) - self._check_dtype(self.q_desc, torch.bfloat16, name="Q") - self._check_dtype(self.k_desc, torch.bfloat16, name="K") - self._check_dtype(self.lse_desc, torch.float32, name="LSE") - self._check_dtype(self.topk_desc, torch.int32, name="topk_indices") - self._check_dtype(self.out_desc, torch.float32, name="out") - _check_sparse_score_shapes( - self, - self.q_desc, - self.k_desc, - self.lse_desc, - self.topk_desc, - self.out_desc, - self.topk_length_desc, - "LSE", - self.qhead_per_kv_head, - ) + target_compute_capability = _check_score_arch(self) + if self._op is None: + self._op = SparseScoreRecomputeOp( + q=_compact_for_kernel(self.q_desc), + k=_compact_for_kernel(self.k_desc), + per_head=_compact_for_kernel(self.lse_desc), + topk_indices=_compact_for_kernel(self.topk_desc), + output=_compact_for_kernel(self.out_desc), + topk_length=_compact_for_kernel(self.topk_length_desc), + score_type="attention", + softmax_scale=self.softmax_scale, + qhead_per_kv_head=self.qhead_per_kv_head, + topk_indices_global=self.topk_indices_global, + target_compute_capability=target_compute_capability, + ) + self._op.check_support() self._is_supported = True return True @@ -442,7 +425,7 @@ def sparse_attn_score_recompute_wrapper( qhead_per_kv_head=qhead_per_kv_head, topk_indices_global=topk_indices_global, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_SparseAttnScoreRecomputeObjects[key] = obj @@ -524,9 +507,28 @@ def __init__( self.sm_scale = float(sm_scale) self.ratio = int(ratio) self.is_thd = bool(is_thd) + self._op = None def check_support(self) -> bool: - _check_score_arch(self) + target_compute_capability = _check_score_arch(self) + if not self.is_thd: + if self._op is None: + self._op = DenseScoreRecomputeOp( + q=_compact_for_kernel(self.q_desc), + k=_compact_for_kernel(self.k_desc), + per_head=_compact_for_kernel(self.w_desc), + output=_compact_for_kernel(self.out_desc), + denominator=_compact_for_kernel(self.denom_desc), + score_type="indexer", + scale=self.sm_scale, + ratio=self.ratio, + qhead_per_kv_head=self.qhead_per_kv_head, + is_thd=False, + target_compute_capability=target_compute_capability, + ) + self._op.check_support() + self._is_supported = True + return True self._check_dtype(self.q_desc, torch.bfloat16, name="Q") self._check_dtype(self.k_desc, torch.bfloat16, name="K") self._check_dtype(self.w_desc, torch.bfloat16, name="W") @@ -683,7 +685,7 @@ def dense_indexer_score_recompute_wrapper( ratio=ratio, is_thd=is_thd, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_DenseIndexerScoreRecomputeObjects[key] = obj @@ -735,9 +737,28 @@ def __init__( self.qhead_per_kv_head = qhead_per_kv_head self.ratio = int(ratio) self.is_thd = bool(is_thd) + self._op = None def check_support(self) -> bool: - _check_score_arch(self) + target_compute_capability = _check_score_arch(self) + if not self.is_thd: + if self._op is None: + self._op = DenseScoreRecomputeOp( + q=_compact_for_kernel(self.q_desc), + k=_compact_for_kernel(self.k_desc), + per_head=_compact_for_kernel(self.lse_desc), + output=_compact_for_kernel(self.out_desc), + denominator=_compact_for_kernel(self.denom_desc), + score_type="attention", + scale=self.softmax_scale, + ratio=self.ratio, + qhead_per_kv_head=self.qhead_per_kv_head, + is_thd=False, + target_compute_capability=target_compute_capability, + ) + self._op.check_support() + self._is_supported = True + return True self._check_dtype(self.q_desc, torch.bfloat16, name="Q") self._check_dtype(self.k_desc, torch.bfloat16, name="K") self._check_dtype(self.lse_desc, torch.float32, name="LSE") @@ -894,7 +915,7 @@ def dense_attn_score_recompute_wrapper( ratio=ratio, is_thd=is_thd, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_DenseAttnScoreRecomputeObjects[key] = obj diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/config.py b/python/cudnn/deepseek_sparse_attention/score_recompute/config.py new file mode 100644 index 000000000..f666747b7 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/config.py @@ -0,0 +1,292 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral SM100 configuration for score-recompute kernels.""" + +from __future__ import annotations + +from dataclasses import dataclass +import math + +_SM100_SMEM_BYTES = 228 * 1024 + + +@dataclass(frozen=True) +class SparseScoreKernelConfig: + """Static launch configuration for ``SparseScoreRecomputeSm100``.""" + + m_block_size: int + n_block_size: int + k_block_size: int | None + kv_stage: int + have_topk_length: bool + topk_in_smem: bool + + +@dataclass(frozen=True) +class DenseScoreKernelConfig: + """Static launch configuration for ``DenseScoreRecomputeSm100``.""" + + m_block_size: int + n_block_size: int + k_block_size: int | None + kv_stage: int + + +def dispatch_sparse_attn_tile_params( + head_dim: int, + qhead_per_kv_head: int, + topk: int, + *, + compact: bool, +) -> tuple[int, int, int | None]: + """Select the tuned SM100 sparse-attention tile.""" + + m = qhead_per_kv_head + if compact: + if head_dim == 512 and qhead_per_kv_head == 64: + return m, 64, None + if head_dim == 512 and qhead_per_kv_head == 128: + return m, 128, 128 + if head_dim == 576 and qhead_per_kv_head == 64: + return m, 64, None + if head_dim == 576 and qhead_per_kv_head == 128: + return m, 64, 192 + return m, 64, None + + if head_dim == 512 and qhead_per_kv_head == 64: + return m, 128, 256 + if head_dim == 512 and qhead_per_kv_head == 128 and topk <= 512: + return m, 128, 128 + if head_dim == 512 and qhead_per_kv_head == 128: + return m, 128, 256 + if head_dim == 576 and qhead_per_kv_head in (64, 128): + return m, 128, 192 + return m, 64, None + + +def resolve_sparse_score_smem_config( + *, + score_type: str, + head_dim: int, + m_block_size: int, + n_block_size: int, + k_block_size: int | None, + topk: int, +) -> tuple[int, bool]: + """Return ``(kv_stage, topk_in_smem)`` for an explicit tile choice.""" + + if score_type not in ("indexer", "attention"): + raise ValueError(f"score_type must be 'indexer' or 'attention', got {score_type!r}") + if head_dim <= 0: + raise ValueError(f"head dimension must be positive, got {head_dim}") + if m_block_size <= 0 or n_block_size <= 0: + raise ValueError(f"m_block_size and n_block_size must be positive, got {m_block_size} and {n_block_size}") + if topk <= 0 or topk % n_block_size: + raise ValueError(f"topk ({topk}) must be positive and a multiple of n_block_size ({n_block_size})") + + head_dim_padded = math.ceil(head_dim / 16) * 16 + k_block_size_eff = head_dim_padded if k_block_size is None else k_block_size + if k_block_size_eff <= 0: + raise ValueError(f"k_block_size must be positive, got {k_block_size_eff}") + if head_dim_padded % k_block_size_eff: + raise ValueError(f"padded head dimension ({head_dim_padded}) must be divisible by k_block_size ({k_block_size_eff})") + if k_block_size_eff % 64: + raise ValueError(f"k_block_size must be a multiple of 64, got {k_block_size_eff}") + + per_head_element_bytes = 2 if score_type == "indexer" else 4 + k_bytes_per_stage = n_block_size * k_block_size_eff * 2 + q_bytes = m_block_size * head_dim_padded * 2 + topk_bytes = topk * 2 * 4 + per_head_bytes = m_block_size * 2 * per_head_element_bytes + fixed_overhead = per_head_bytes + 2048 + + topk_in_smem = True + smem_overhead = topk_bytes + fixed_overhead + kv_stage = min(4, max(1, (_SM100_SMEM_BYTES - q_bytes - smem_overhead) // k_bytes_per_stage)) + total_smem = q_bytes + k_bytes_per_stage * kv_stage + smem_overhead + if total_smem > _SM100_SMEM_BYTES: + topk_in_smem = False + smem_overhead = fixed_overhead + kv_stage = min(4, max(1, (_SM100_SMEM_BYTES - q_bytes - smem_overhead) // k_bytes_per_stage)) + total_smem = q_bytes + k_bytes_per_stage * kv_stage + smem_overhead + if total_smem > _SM100_SMEM_BYTES: + raise ValueError( + "SM100 shared-memory requirement exceeds 228 KiB even with top-K indices in global memory: " + f"head_dim={head_dim}, m_block_size={m_block_size}, " + f"n_block_size={n_block_size}, k_block_size={k_block_size_eff}" + ) + + return kv_stage, topk_in_smem + + +def resolve_sparse_score_kernel_config( + *, + score_type: str, + head_dim: int, + qhead_per_kv_head: int, + topk: int, + have_topk_length: bool, +) -> SparseScoreKernelConfig: + """Resolve the SM100 sparse tile and shared-memory configuration.""" + + if score_type not in ("indexer", "attention"): + raise ValueError(f"score_type must be 'indexer' or 'attention', got {score_type!r}") + if head_dim <= 0 or head_dim % 64: + raise ValueError(f"head dimension must be a positive multiple of 64, got {head_dim}") + if qhead_per_kv_head not in (32, 64, 128): + raise ValueError(f"qhead_per_kv_head must be 32, 64, or 128 for the SM100 sparse score kernel, got {qhead_per_kv_head}") + if topk <= 0: + raise ValueError(f"topk must be positive, got {topk}") + + if score_type == "indexer": + m_block_size, n_block_size, k_block_size = qhead_per_kv_head, 128, None + else: + m_block_size, n_block_size, k_block_size = dispatch_sparse_attn_tile_params( + head_dim, + qhead_per_kv_head, + topk, + compact=have_topk_length, + ) + + if topk % n_block_size: + raise ValueError(f"topk ({topk}) must be a multiple of the selected n_block_size ({n_block_size})") + + kv_stage, topk_in_smem = resolve_sparse_score_smem_config( + score_type=score_type, + head_dim=head_dim, + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=k_block_size, + topk=topk, + ) + return SparseScoreKernelConfig( + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=k_block_size, + kv_stage=kv_stage, + have_topk_length=bool(have_topk_length), + topk_in_smem=topk_in_smem, + ) + + +def _select_dense_k_block_size( + head_dim_padded: int, + m_block_size: int, + n_block_size: int, + per_head_element_bytes: int, +) -> int: + q_bytes = m_block_size * head_dim_padded * 2 + per_head_bytes = m_block_size * 2 * per_head_element_bytes + available_k_bytes = _SM100_SMEM_BYTES - q_bytes - per_head_bytes - 2048 + max_k_block_size = max(64, available_k_bytes // (n_block_size * 2)) + max_k_block_size = (max_k_block_size // 64) * 64 + + for candidate in range(min(head_dim_padded, max_k_block_size), 63, -64): + if head_dim_padded % candidate == 0: + return candidate + raise ValueError( + "No SM100 dense-score K tile divides the padded head dimension: " + f"head_dim_padded={head_dim_padded}, m_block_size={m_block_size}, n_block_size={n_block_size}" + ) + + +def _dense_m_block_size(qhead_per_kv_head: int, head_dim: int, per_head_element_bytes: int) -> int: + head_dim_padded = math.ceil(head_dim / 16) * 16 + n_block_size = 128 + m_block_size = qhead_per_kv_head * 2 + q_bytes = m_block_size * head_dim_padded * 2 + minimum_k_bytes = n_block_size * 64 * 2 + per_head_bytes = m_block_size * 2 * per_head_element_bytes + if q_bytes + minimum_k_bytes + per_head_bytes + 4096 <= _SM100_SMEM_BYTES: + return m_block_size + return qhead_per_kv_head + + +def resolve_dense_score_smem_config( + *, + score_type: str, + head_dim: int, + m_block_size: int, + n_block_size: int, + k_block_size: int | None, +) -> tuple[int, int]: + """Return the effective K tile and stage count for a dense tile choice.""" + + if score_type not in ("indexer", "attention"): + raise ValueError(f"score_type must be 'indexer' or 'attention', got {score_type!r}") + if head_dim <= 0 or m_block_size <= 0 or n_block_size <= 0: + raise ValueError(f"head_dim, m_block_size, and n_block_size must be positive, got {head_dim}, {m_block_size}, and {n_block_size}") + head_dim_padded = math.ceil(head_dim / 16) * 16 + per_head_element_bytes = 2 if score_type == "indexer" else 4 + effective_k_block_size = k_block_size + if effective_k_block_size is None: + effective_k_block_size = _select_dense_k_block_size( + head_dim_padded, + m_block_size, + n_block_size, + per_head_element_bytes, + ) + if effective_k_block_size <= 0 or effective_k_block_size % 64: + raise ValueError(f"k_block_size must be a positive multiple of 64, got {effective_k_block_size}") + if head_dim_padded % effective_k_block_size: + raise ValueError(f"padded head dimension ({head_dim_padded}) must be divisible by k_block_size ({effective_k_block_size})") + + k_bytes_per_stage = n_block_size * effective_k_block_size * 2 + q_bytes = m_block_size * head_dim_padded * 2 + per_head_bytes = m_block_size * 2 * per_head_element_bytes + kv_stage = min(4, max(1, (_SM100_SMEM_BYTES - q_bytes - per_head_bytes - 2048) // k_bytes_per_stage)) + return effective_k_block_size, kv_stage + + +def resolve_dense_score_kernel_config(*, score_type: str, head_dim: int, qhead_per_kv_head: int) -> DenseScoreKernelConfig: + """Resolve the tuned SM100 dense-score tile.""" + + if score_type not in ("indexer", "attention"): + raise ValueError(f"score_type must be 'indexer' or 'attention', got {score_type!r}") + if head_dim <= 0 or head_dim % 64: + raise ValueError(f"head dimension must be a positive multiple of 64, got {head_dim}") + if qhead_per_kv_head not in (32, 64, 128): + raise ValueError(f"qhead_per_kv_head must be 32, 64, or 128 for the SM100 dense score kernel, got {qhead_per_kv_head}") + + per_head_element_bytes = 2 if score_type == "indexer" else 4 + m_block_size = _dense_m_block_size(qhead_per_kv_head, head_dim, per_head_element_bytes) + n_block_size = 128 + head_dim_padded = math.ceil(head_dim / 16) * 16 + + if score_type == "indexer" and head_dim == 128: + k_block_size = head_dim_padded + elif score_type == "attention" and head_dim in (512, 576): + k_block_size = 64 + else: + k_block_size = _select_dense_k_block_size( + head_dim_padded, + m_block_size, + n_block_size, + per_head_element_bytes, + ) + + k_block_size, kv_stage = resolve_dense_score_smem_config( + score_type=score_type, + head_dim=head_dim, + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=k_block_size, + ) + return DenseScoreKernelConfig( + m_block_size=m_block_size, + n_block_size=n_block_size, + k_block_size=None if k_block_size == head_dim_padded else k_block_size, + kv_stage=kv_stage, + ) + + +__all__ = [ + "DenseScoreKernelConfig", + "SparseScoreKernelConfig", + "dispatch_sparse_attn_tile_params", + "resolve_dense_score_kernel_config", + "resolve_dense_score_smem_config", + "resolve_sparse_score_kernel_config", + "resolve_sparse_score_smem_config", +] diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py b/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py new file mode 100644 index 000000000..f34e42a7e --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/jax.py @@ -0,0 +1,1081 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX APIs for DSA sparse and dense score recomputation.""" + +from __future__ import annotations + +from functools import partial +from typing import Any, cast + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..._jax.layout import mode_from_layout, to_public_axes +from .config import DenseScoreKernelConfig, SparseScoreKernelConfig +from .op import ( + DenseScoreRecomputeOp, + DenseScoreSm90Config, + SUPPORTED_COMPUTE_CAPABILITIES, + SparseScoreRecomputeOp, + SparseScoreSm90Config, +) + + +def _compile_options(target_compute_capability: int) -> str: + from ..._jax.compiler import compile_options_for_target + + return compile_options_for_target(target_compute_capability) + + +def _resolve_layout( + name: str, + layout: str | None, + *, + default: str, + kernel_axes: str, + supported: tuple[str, ...], +) -> tuple[str, tuple[int, ...]]: + layout = default if layout is None else layout + if layout not in supported: + choices = ", ".join(repr(value) for value in supported) + raise ValueError(f"{name} must be one of ({choices}), got {layout!r}") + return layout, mode_from_layout(layout, kernel_axes=kernel_axes) + + +def _optional_desc( + api: JaxApiBase, + value: Any | None, + name: str, + *, + mode: tuple[int, ...] | None = None, +) -> JaxTensorDesc | None: + return None if value is None else api._to_tensor_desc(value, name, mode=mode) + + +def _check_optional_signature( + api: JaxApiBase, + value: Any | None, + expected: JaxTensorDesc | None, + name: str, + *, + mode: tuple[int, ...] | None = None, +) -> None: + if expected is None: + if value is not None: + raise ValueError(f"{name} was not part of the sample signature") + return + if value is None: + raise ValueError(f"{name} is required by the sample signature") + + +class _SparseScoreRecompute(JaxApiBase): + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_per_head: Any, + sample_topk_indices: Any, + *, + score_type: str, + softmax_scale: float, + sample_out: Any | None, + sample_topk_length: Any | None, + qhead_per_kv_head: int | None, + topk_indices_global: bool, + q_layout: str | None, + k_layout: str | None, + per_head_layout: str | None, + score_layout: str | None, + output_layout: str | None, + target_compute_capability: int | None, + ) -> None: + self.q_layout, self.q_mode = _resolve_layout( + "q_layout", + q_layout, + default="BSHD", + kernel_axes="BSHD", + supported=("BSHD", "SBHD"), + ) + self.k_layout, self.k_mode = _resolve_layout( + "k_layout", + k_layout, + default="BSD", + kernel_axes="BSD", + supported=("BSD", "SBD"), + ) + self.per_head_layout, self.per_head_mode = _resolve_layout( + "per_head_layout", + per_head_layout, + default="BSH", + kernel_axes="BSH", + supported=("BSH", "SBH"), + ) + self.score_layout, self.score_mode = _resolve_layout( + "score_layout", + score_layout, + default="BSK", + kernel_axes="BSK", + supported=("BSK", "SBK"), + ) + self.output_layout, self.output_mode = _resolve_layout( + "output_layout", + output_layout, + default="BSK", + kernel_axes="BSK", + supported=("BSK", "SBK"), + ) + self.length_layout = self.score_layout.replace("K", "") + self.length_mode = mode_from_layout(self.length_layout, kernel_axes="BS") + + self.q_desc = self._to_tensor_desc(sample_q, "sample_q", mode=self.q_mode) + self.k_desc = self._to_tensor_desc(sample_k, "sample_k", mode=self.k_mode) + self.per_head_desc = self._to_tensor_desc( + sample_per_head, "sample_per_head", mode=self.per_head_mode + ) + self.topk_indices_desc = self._to_tensor_desc( + sample_topk_indices, + "sample_topk_indices", + mode=self.score_mode, + ) + self.topk_length_desc = _optional_desc( + self, + sample_topk_length, + "sample_topk_length", + mode=self.length_mode, + ) + self.out_desc = ( + self._default_output_desc() + if sample_out is None + else self._to_tensor_desc(sample_out, "sample_out", mode=self.output_mode) + ) + self.target_compute_capability = self._resolve_compute_capability( + target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + type(self).__name__, + ) + self._op = SparseScoreRecomputeOp( + q=self.q_desc, + k=self.k_desc, + per_head=self.per_head_desc, + topk_indices=self.topk_indices_desc, + output=self.out_desc, + topk_length=self.topk_length_desc, + score_type=score_type, + softmax_scale=softmax_scale, + qhead_per_kv_head=qhead_per_kv_head, + topk_indices_global=topk_indices_global, + target_compute_capability=self.target_compute_capability, + ) + + def _default_output_desc(self) -> JaxTensorDesc: + if self.topk_indices_desc.ndim != 3: + raise ValueError( + f"topk_indices must have rank 3, got shape {self.topk_indices_desc.shape}" + ) + return JaxTensorDesc.from_shape( + to_public_axes(self.topk_indices_desc.shape, self.output_mode), + jnp.float32, + name="sample_out", + mode=self.output_mode, + ) + + def check_support(self) -> bool: + return self._op.check_support() + + def _call( + self, q: Any, k: Any, per_head: Any, topk_indices: Any, topk_length: Any | None + ) -> Any: + self.check_support() + _check_optional_signature( + self, + topk_length, + self.topk_length_desc, + "topk_length", + mode=self.length_mode, + ) + + if self.target_compute_capability == 90: + kernel_k = jnp.expand_dims(k, axis=self.k_layout.index("D")) + kernel_k_layout = self.k_layout.replace("D", "HD") + kernel_k_mode = mode_from_layout(kernel_k_layout, kernel_axes="BSHD") + kernel_k_desc = self._to_tensor_desc( + kernel_k, + "kernel_k", + mode=kernel_k_mode, + ) + per_head_axes = tuple(self.per_head_layout.index(axis) for axis in "BHS") + kernel_per_head = jnp.transpose(per_head, per_head_axes) + kernel_per_head_desc = self._to_tensor_desc( + kernel_per_head, + "kernel_per_head", + ) + inputs = (q, kernel_k, kernel_per_head, topk_indices) + input_descs = ( + self.q_desc, + kernel_k_desc, + kernel_per_head_desc, + self.topk_indices_desc, + ) + else: + inputs = (q, k, per_head, topk_indices) + input_descs = ( + self.q_desc, + self.k_desc, + self.per_head_desc, + self.topk_indices_desc, + ) + + workspace_descs = () + if topk_length is not None: + inputs += (topk_length,) + input_descs += (self.topk_length_desc,) + elif self.target_compute_capability != 90: + workspace_descs = ( + self.topk_indices_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=(1, 1), + name="topk_length_workspace", + ), + ) + + (output,) = self._call_kernel( + inputs, + launch=self._launch_kernel, + output_descs=(self.out_desc,), + input_descs=input_descs, + workspace_descs=workspace_descs, + compile_options=_compile_options(self.target_compute_capability), + ) + return output + + def _launch_kernel( + self, + stream: Any, + *arguments: Any, + ) -> None: + uses_workspace = ( + self.topk_length_desc is None and self.target_compute_capability != 90 + ) + if uses_workspace: + *inputs, output, topk_length_workspace = arguments + else: + *inputs, output = arguments + topk_length_workspace = None + q, k, per_head, topk_indices, *optional_inputs = inputs + topk_length = optional_inputs[0] if optional_inputs else topk_length_workspace + + if self.target_compute_capability == 90: + import cutlass + + from .sparse_score_recompute_sm90 import SparseScoreRecomputeSm90 + + config = cast(SparseScoreSm90Config, self._op.config) + kernel = SparseScoreRecomputeSm90( + cutlass.BFloat16, + head_dim=self.q_desc.shape[-1], + qhead_per_kvhead=cast(int, self._op.qhead_per_kv_head), + tile_m=config.tile_m, + tile_n=config.tile_n, + KV_stage=config.kv_stage, + num_threads=config.num_threads, + swap_AB=True, + topk_max=self.topk_indices_desc.shape[-1], + is_index_scores=self._op.score_type == "indexer", + softmax_scale=self._op.softmax_scale, + has_topk_length=self.topk_length_desc is not None, + num_head_tiles=config.num_head_tiles, + is_sparse=True, + output_log_probs=False, + topk_indices_global=self._op.topk_indices_global, + ) + kernel(q, k, topk_indices, stream, output, per_head, topk_length, None) + return + + import cutlass + + from .sparse_score_recompute_sm100 import SparseScoreRecomputeSm100 + + config = cast(SparseScoreKernelConfig, self._op.config) + kernel = SparseScoreRecomputeSm100( + head_dim=self.q_desc.shape[-1], + qhead_per_kvhead=cast(int, self._op.qhead_per_kv_head), + m_block_size=config.m_block_size, + n_block_size=config.n_block_size, + k_block_size=config.k_block_size, + topk=self.topk_indices_desc.shape[-1], + kv_stage=config.kv_stage, + score_type=self._op.score_type, + have_topk_length=config.have_topk_length, + topk_in_smem=config.topk_in_smem, + topk_indices_global=self._op.topk_indices_global, + ) + kernel( + q, + k, + per_head, + topk_indices, + output, + topk_length, + cutlass.Float32(self._op.softmax_scale), + stream, + ) + + +class SparseIndexerScoreRecompute(_SparseScoreRecompute): + """JAX callable specialized for sparse indexer score recomputation.""" + + def __init__( + self, + sample_q_indexer: Any, + sample_k_indexer: Any, + sample_weights: Any, + sample_topk_indices: Any, + sample_out: Any | None = None, + sample_topk_length: Any | None = None, + qhead_per_kv_head: int | None = None, + topk_indices_global: bool = False, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + score_layout: str | None = None, + output_layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + super().__init__( + sample_q_indexer, + sample_k_indexer, + sample_weights, + sample_topk_indices, + score_type="indexer", + softmax_scale=1.0, + sample_out=sample_out, + sample_topk_length=sample_topk_length, + qhead_per_kv_head=qhead_per_kv_head, + topk_indices_global=topk_indices_global, + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + score_layout=score_layout, + output_layout=output_layout, + target_compute_capability=target_compute_capability, + ) + + def __call__( + self, + q_indexer: Any, + k_indexer: Any, + weights: Any, + topk_indices: Any, + topk_length: Any | None = None, + ) -> TupleDict: + return TupleDict( + predict=self._call(q_indexer, k_indexer, weights, topk_indices, topk_length) + ) + + +class SparseAttnScoreRecompute(_SparseScoreRecompute): + """JAX callable specialized for sparse attention score recomputation.""" + + def __init__( + self, + sample_q_attn: Any, + sample_k_attn: Any, + sample_lse: Any, + sample_topk_indices: Any, + softmax_scale: float, + sample_out: Any | None = None, + sample_topk_length: Any | None = None, + qhead_per_kv_head: int | None = None, + topk_indices_global: bool = False, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + score_layout: str | None = None, + output_layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + super().__init__( + sample_q_attn, + sample_k_attn, + sample_lse, + sample_topk_indices, + score_type="attention", + softmax_scale=softmax_scale, + sample_out=sample_out, + sample_topk_length=sample_topk_length, + qhead_per_kv_head=qhead_per_kv_head, + topk_indices_global=topk_indices_global, + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + score_layout=score_layout, + output_layout=output_layout, + target_compute_capability=target_compute_capability, + ) + + def __call__( + self, + q_attn: Any, + k_attn: Any, + lse: Any, + topk_indices: Any, + topk_length: Any | None = None, + ) -> TupleDict: + return TupleDict( + target=self._call(q_attn, k_attn, lse, topk_indices, topk_length) + ) + + +class _DenseScoreRecompute(JaxApiBase): + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_per_head: Any, + *, + score_type: str, + scale: float, + sample_out: Any | None, + sample_denom_out: Any | None, + qhead_per_kv_head: int | None, + ratio: int, + sample_cu_seqlens_q: Any | None, + sample_cu_seqlens_k: Any | None, + max_seqlen_q: int | None, + max_seqlen_k: int | None, + sample_q_causal_offsets: Any | None, + q_layout: str | None, + k_layout: str | None, + per_head_layout: str | None, + output_layout: str | None, + denom_layout: str | None, + target_compute_capability: int | None, + ) -> None: + self.is_thd = len(tuple(sample_q.shape)) == 3 + if self.is_thd: + q_defaults = ("THD", "THD", ("THD",)) + k_defaults = ("THD", "THD", ("THD",)) + per_head_defaults = ("TH", "TH", ("TH",)) + output_defaults = ("TK", "TK", ("TK",)) + denom_defaults = ("T", "T", ("T",)) + else: + q_defaults = ("BSHD", "BSHD", ("BSHD", "SBHD")) + k_defaults = ("BSHD", "BSHD", ("BSHD", "SBHD")) + per_head_defaults = ("BSH", "BSH", ("BSH", "SBH")) + output_defaults = ("BSK", "BSK", ("BSK", "SBK")) + denom_defaults = ("BS", "BS", ("BS", "SB")) + self.q_layout, self.q_mode = _resolve_layout( + "q_layout", + q_layout, + default=q_defaults[0], + kernel_axes=q_defaults[1], + supported=q_defaults[2], + ) + self.k_layout, self.k_mode = _resolve_layout( + "k_layout", + k_layout, + default=k_defaults[0], + kernel_axes=k_defaults[1], + supported=k_defaults[2], + ) + self.per_head_layout, self.per_head_mode = _resolve_layout( + "per_head_layout", + per_head_layout, + default=per_head_defaults[0], + kernel_axes=per_head_defaults[1], + supported=per_head_defaults[2], + ) + self.output_layout, self.output_mode = _resolve_layout( + "output_layout", + output_layout, + default=output_defaults[0], + kernel_axes=output_defaults[1], + supported=output_defaults[2], + ) + self.denom_layout, self.denom_mode = _resolve_layout( + "denom_layout", + denom_layout, + default=denom_defaults[0], + kernel_axes=denom_defaults[1], + supported=denom_defaults[2], + ) + + self.q_desc = self._to_tensor_desc(sample_q, "sample_q", mode=self.q_mode) + self.k_desc = self._to_tensor_desc(sample_k, "sample_k", mode=self.k_mode) + self.per_head_desc = self._to_tensor_desc( + sample_per_head, "sample_per_head", mode=self.per_head_mode + ) + self.cu_seqlens_q_desc = _optional_desc( + self, sample_cu_seqlens_q, "sample_cu_seqlens_q" + ) + self.cu_seqlens_k_desc = _optional_desc( + self, sample_cu_seqlens_k, "sample_cu_seqlens_k" + ) + self.q_causal_offsets_desc = _optional_desc( + self, sample_q_causal_offsets, "sample_q_causal_offsets" + ) + self.requested_max_seqlen_q = max_seqlen_q + self.requested_max_seqlen_k = max_seqlen_k + + default_out, default_denom = self._default_output_descs(max_seqlen_k) + self.out_desc = ( + default_out + if sample_out is None + else self._to_tensor_desc( + sample_out, + "sample_out", + mode=self.output_mode, + init_value=float("-inf"), + ) + ) + self.denom_desc = ( + default_denom + if sample_denom_out is None + else self._to_tensor_desc( + sample_denom_out, + "sample_denom_out", + mode=self.denom_mode, + ) + ) + self.target_compute_capability = self._resolve_compute_capability( + target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + type(self).__name__, + ) + self._op = DenseScoreRecomputeOp( + q=self.q_desc, + k=self.k_desc, + per_head=self.per_head_desc, + output=self.out_desc, + denominator=self.denom_desc, + score_type=score_type, + scale=scale, + ratio=ratio, + qhead_per_kv_head=qhead_per_kv_head, + is_thd=self.is_thd, + cu_seqlens_q=self.cu_seqlens_q_desc, + cu_seqlens_k=self.cu_seqlens_k_desc, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + q_causal_offsets=self.q_causal_offsets_desc, + target_compute_capability=self.target_compute_capability, + ) + + def _default_output_descs( + self, max_seqlen_k: int | None + ) -> tuple[JaxTensorDesc, JaxTensorDesc]: + if self.is_thd: + if max_seqlen_k is None: + raise ValueError( + "max_seqlen_k is required to infer THD dense score outputs" + ) + out_shape = (self.q_desc.shape[0], int(max_seqlen_k)) + denom_shape = (self.q_desc.shape[0],) + elif self.q_desc.ndim == 4 and self.k_desc.ndim == 4: + out_shape = ( + self.q_desc.shape[0], + self.q_desc.shape[1], + self.k_desc.shape[1], + ) + denom_shape = self.q_desc.shape[:2] + else: + raise ValueError( + f"Dense score inputs must both use BSHD or THD layout, got Q={self.q_desc.shape}, K={self.k_desc.shape}" + ) + return ( + JaxTensorDesc.from_shape( + to_public_axes(out_shape, self.output_mode), + jnp.float32, + name="sample_out", + mode=self.output_mode, + init_value=float("-inf"), + ), + JaxTensorDesc.from_shape( + to_public_axes(denom_shape, self.denom_mode), + jnp.float32, + name="sample_denom_out", + mode=self.denom_mode, + ), + ) + + def check_support(self) -> bool: + return self._op.check_support() + + def _call( + self, + q: Any, + k: Any, + per_head: Any, + cu_seqlens_q: Any | None, + cu_seqlens_k: Any | None, + q_causal_offsets: Any | None, + ) -> TupleDict: + self.check_support() + _check_optional_signature( + self, cu_seqlens_q, self.cu_seqlens_q_desc, "cu_seqlens_q" + ) + _check_optional_signature( + self, cu_seqlens_k, self.cu_seqlens_k_desc, "cu_seqlens_k" + ) + _check_optional_signature( + self, q_causal_offsets, self.q_causal_offsets_desc, "q_causal_offsets" + ) + + if self.target_compute_capability == 90: + per_head_axes = tuple(self.per_head_layout.index(axis) for axis in "BHS") + kernel_per_head = jnp.transpose(per_head, per_head_axes) + kernel_per_head_desc = self._to_tensor_desc( + kernel_per_head, + "kernel_per_head", + ) + inputs = (q, k, kernel_per_head) + input_descs = ( + self.q_desc, + self.k_desc, + kernel_per_head_desc, + ) + else: + inputs = (q, k, per_head) + input_descs = ( + self.q_desc, + self.k_desc, + self.per_head_desc, + ) + if self.is_thd: + inputs += (cu_seqlens_q, cu_seqlens_k) + input_descs += ( + self.cu_seqlens_q_desc, + self.cu_seqlens_k_desc, + ) + if q_causal_offsets is not None: + inputs += (q_causal_offsets,) + input_descs += (self.q_causal_offsets_desc,) + + output, denominator = self._call_kernel( + inputs, + launch=self._launch_kernel, + output_descs=(self.out_desc, self.denom_desc), + input_descs=input_descs, + compile_options=_compile_options(self.target_compute_capability), + ) + return TupleDict(out=output, denom=denominator) + + def _launch_kernel( + self, + stream: Any, + *arguments: Any, + ) -> None: + *inputs, output, denominator = arguments + q, k, per_head, *optional_inputs = inputs + + if self.target_compute_capability == 90: + import cutlass + + from .dense_score_recompute_sm90 import DenseScoreRecomputeSm90 + + q_causal_offsets = optional_inputs[0] if optional_inputs else None + config = cast(DenseScoreSm90Config, self._op.config) + kernel = DenseScoreRecomputeSm90( + cutlass.BFloat16, + head_dim=self.q_desc.shape[-1], + qhead_per_kvhead=cast(int, self._op.qhead_per_kv_head), + tile_m=config.tile_m, + tile_n=config.tile_n, + KV_stage=config.kv_stage, + num_threads=config.num_threads, + swap_AB=True, + topk_max=cast(int, self._op.max_seqlen_k), + is_index_scores=self._op.score_type == "indexer", + softmax_scale=self._op.scale, + has_topk_length=False, + num_head_tiles=config.num_head_tiles, + ratio=self._op.ratio, + ) + kernel( + q, + k, + None, + stream, + output, + per_head, + None, + denominator, + q_causal_offsets, + ) + return + + import cutlass + + from .dense_score_recompute_sm100 import DenseScoreRecomputeSm100 + + if self.is_thd: + cu_seqlens_q, cu_seqlens_k, *optional_inputs = optional_inputs + else: + cu_seqlens_q = cu_seqlens_k = None + q_causal_offsets = optional_inputs[0] if optional_inputs else None + config = cast(DenseScoreKernelConfig, self._op.config) + kernel = DenseScoreRecomputeSm100( + head_dim=self.q_desc.shape[-1], + qhead_per_kvhead=cast(int, self._op.qhead_per_kv_head), + m_block_size=config.m_block_size, + n_block_size=config.n_block_size, + k_block_size=config.k_block_size, + kv_stage=config.kv_stage, + score_type=self._op.score_type, + ratio=self._op.ratio, + is_varlen=self.is_thd, + ) + kernel( + q, + k, + per_head, + output, + denominator, + cutlass.Float32(self._op.scale), + cutlass.Int32(cast(int, self._op.max_seqlen_q)), + cutlass.Int32(cast(int, self._op.max_seqlen_k)), + cu_seqlens_q, + cu_seqlens_k, + q_causal_offsets, + stream, + ) + + +class DenseIndexerScoreRecompute(_DenseScoreRecompute): + """JAX callable specialized for dense indexer score recomputation.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_weights: Any, + sample_out: Any | None = None, + sample_denom_out: Any | None = None, + qhead_per_kv_head: int | None = None, + sm_scale: float = 1.0, + ratio: int = 1, + sample_cu_seqlens_q: Any | None = None, + sample_cu_seqlens_k: Any | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + sample_q_causal_offsets: Any | None = None, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + output_layout: str | None = None, + denom_layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + super().__init__( + sample_q, + sample_k, + sample_weights, + score_type="indexer", + scale=sm_scale, + sample_out=sample_out, + sample_denom_out=sample_denom_out, + qhead_per_kv_head=qhead_per_kv_head, + ratio=ratio, + sample_cu_seqlens_q=sample_cu_seqlens_q, + sample_cu_seqlens_k=sample_cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + sample_q_causal_offsets=sample_q_causal_offsets, + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + output_layout=output_layout, + denom_layout=denom_layout, + target_compute_capability=target_compute_capability, + ) + + def __call__( + self, + q: Any, + k: Any, + weights: Any, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + q_causal_offsets: Any | None = None, + ) -> TupleDict: + return self._call(q, k, weights, cu_seqlens_q, cu_seqlens_k, q_causal_offsets) + + +class DenseAttnScoreRecompute(_DenseScoreRecompute): + """JAX callable specialized for dense attention score recomputation.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_lse: Any, + softmax_scale: float, + sample_out: Any | None = None, + sample_denom_out: Any | None = None, + qhead_per_kv_head: int | None = None, + ratio: int = 1, + sample_cu_seqlens_q: Any | None = None, + sample_cu_seqlens_k: Any | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + sample_q_causal_offsets: Any | None = None, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + output_layout: str | None = None, + denom_layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + super().__init__( + sample_q, + sample_k, + sample_lse, + score_type="attention", + scale=softmax_scale, + sample_out=sample_out, + sample_denom_out=sample_denom_out, + qhead_per_kv_head=qhead_per_kv_head, + ratio=ratio, + sample_cu_seqlens_q=sample_cu_seqlens_q, + sample_cu_seqlens_k=sample_cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + sample_q_causal_offsets=sample_q_causal_offsets, + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + output_layout=output_layout, + denom_layout=denom_layout, + target_compute_capability=target_compute_capability, + ) + + def __call__( + self, + q: Any, + k: Any, + lse: Any, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + q_causal_offsets: Any | None = None, + ) -> TupleDict: + return self._call(q, k, lse, cu_seqlens_q, cu_seqlens_k, q_causal_offsets) + + +@partial( + jax.jit, + static_argnames=( + "qhead_per_kv_head", + "topk_indices_global", + "q_layout", + "k_layout", + "per_head_layout", + "score_layout", + "output_layout", + "target_compute_capability", + ), +) +def sparse_indexer_score_recompute_wrapper( + q_indexer: Any, + k_indexer: Any, + weights: Any, + topk_indices: Any, + qhead_per_kv_head: int | None = None, + topk_length: Any | None = None, + topk_indices_global: bool = False, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + score_layout: str | None = None, + output_layout: str | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + return SparseIndexerScoreRecompute( + q_indexer, + k_indexer, + weights, + topk_indices, + sample_topk_length=topk_length, + qhead_per_kv_head=qhead_per_kv_head, + topk_indices_global=topk_indices_global, + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + score_layout=score_layout, + output_layout=output_layout, + target_compute_capability=target_compute_capability, + )(q_indexer, k_indexer, weights, topk_indices, topk_length) + + +@partial( + jax.jit, + static_argnames=( + "softmax_scale", + "qhead_per_kv_head", + "topk_indices_global", + "q_layout", + "k_layout", + "per_head_layout", + "score_layout", + "output_layout", + "target_compute_capability", + ), +) +def sparse_attn_score_recompute_wrapper( + q_attn: Any, + k_attn: Any, + lse: Any, + topk_indices: Any, + softmax_scale: float, + qhead_per_kv_head: int | None = None, + topk_length: Any | None = None, + topk_indices_global: bool = False, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + score_layout: str | None = None, + output_layout: str | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + return SparseAttnScoreRecompute( + q_attn, + k_attn, + lse, + topk_indices, + softmax_scale, + sample_topk_length=topk_length, + qhead_per_kv_head=qhead_per_kv_head, + topk_indices_global=topk_indices_global, + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + score_layout=score_layout, + output_layout=output_layout, + target_compute_capability=target_compute_capability, + )(q_attn, k_attn, lse, topk_indices, topk_length) + + +@partial( + jax.jit, + static_argnames=( + "qhead_per_kv_head", + "sm_scale", + "ratio", + "max_seqlen_q", + "max_seqlen_k", + "q_layout", + "k_layout", + "per_head_layout", + "output_layout", + "denom_layout", + "target_compute_capability", + ), +) +def dense_indexer_score_recompute_wrapper( + q: Any, + k: Any, + weights: Any, + qhead_per_kv_head: int | None = None, + sm_scale: float = 1.0, + ratio: int = 1, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + q_causal_offsets: Any | None = None, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + output_layout: str | None = None, + denom_layout: str | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + return DenseIndexerScoreRecompute( + q, + k, + weights, + qhead_per_kv_head=qhead_per_kv_head, + sm_scale=sm_scale, + ratio=ratio, + sample_cu_seqlens_q=cu_seqlens_q, + sample_cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + sample_q_causal_offsets=q_causal_offsets, + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + output_layout=output_layout, + denom_layout=denom_layout, + target_compute_capability=target_compute_capability, + )(q, k, weights, cu_seqlens_q, cu_seqlens_k, q_causal_offsets) + + +@partial( + jax.jit, + static_argnames=( + "softmax_scale", + "qhead_per_kv_head", + "ratio", + "max_seqlen_q", + "max_seqlen_k", + "q_layout", + "k_layout", + "per_head_layout", + "output_layout", + "denom_layout", + "target_compute_capability", + ), +) +def dense_attn_score_recompute_wrapper( + q: Any, + k: Any, + lse: Any, + softmax_scale: float, + qhead_per_kv_head: int | None = None, + ratio: int = 1, + cu_seqlens_q: Any | None = None, + cu_seqlens_k: Any | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + q_causal_offsets: Any | None = None, + q_layout: str | None = None, + k_layout: str | None = None, + per_head_layout: str | None = None, + output_layout: str | None = None, + denom_layout: str | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + return DenseAttnScoreRecompute( + q, + k, + lse, + softmax_scale, + qhead_per_kv_head=qhead_per_kv_head, + ratio=ratio, + sample_cu_seqlens_q=cu_seqlens_q, + sample_cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + sample_q_causal_offsets=q_causal_offsets, + q_layout=q_layout, + k_layout=k_layout, + per_head_layout=per_head_layout, + output_layout=output_layout, + denom_layout=denom_layout, + target_compute_capability=target_compute_capability, + )(q, k, lse, cu_seqlens_q, cu_seqlens_k, q_causal_offsets) + + +__all__ = [ + "DenseAttnScoreRecompute", + "DenseIndexerScoreRecompute", + "SparseAttnScoreRecompute", + "SparseIndexerScoreRecompute", + "dense_attn_score_recompute_wrapper", + "dense_indexer_score_recompute_wrapper", + "sparse_attn_score_recompute_wrapper", + "sparse_indexer_score_recompute_wrapper", +] diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/op.py b/python/cudnn/deepseek_sparse_attention/score_recompute/op.py new file mode 100644 index 000000000..ede5f357d --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/op.py @@ -0,0 +1,562 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral signatures for DSA score-recompute operations.""" + +from __future__ import annotations + +from dataclasses import dataclass +from numbers import Real +from operator import index +from typing import Any + +from ... import data_type +from ...common.op import Op +from ...common.tensor_desc import TensorDesc +from .config import ( + DenseScoreKernelConfig, + SparseScoreKernelConfig, + resolve_dense_score_kernel_config, + resolve_sparse_score_kernel_config, +) + +SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) + + +@dataclass(frozen=True) +class SparseScoreSm90Config: + tile_m: int + tile_n: int + kv_stage: int + num_threads: int + num_head_tiles: int + + +@dataclass(frozen=True) +class DenseScoreSm90Config: + tile_m: int + tile_n: int + kv_stage: int + num_threads: int + num_head_tiles: int + + +def _require_desc( + value: Any, name: str, *, optional: bool = False +) -> TensorDesc[Any] | None: + if optional and value is None: + return None + if not isinstance(value, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(value).__name__}") + return value + + +def _require_dtype(desc: TensorDesc[Any], expected: data_type, name: str) -> None: + if desc.cudnn_dtype != expected: + raise ValueError(f"{name} must have dtype {expected}, got {desc.cudnn_dtype}") + + +def _require_rank(desc: TensorDesc[Any], rank: int, name: str) -> None: + if desc.ndim != rank: + raise ValueError(f"{name} must have rank {rank}, got shape {desc.shape}") + + +def _require_compact(desc: TensorDesc[Any], name: str) -> None: + if not desc.is_compact(): + raise ValueError(f"{name} must be compact, got stride {desc.stride}") + + +def _require_contiguous_tail(desc: TensorDesc[Any], rank: int, name: str) -> None: + expected_stride = 1 + for axis in range(desc.ndim - 1, desc.ndim - rank - 1, -1): + if desc.shape[axis] != 1 and desc.stride[axis] != expected_stride: + raise ValueError( + f"{name} must have its final {rank} axes contiguous, " + f"got shape {desc.shape} and stride {desc.stride}" + ) + expected_stride *= max(desc.shape[axis], 1) + + +def _positive_dimensions(named_dimensions: dict[str, int], operation: str) -> None: + invalid = ", ".join( + f"{name}={value}" for name, value in named_dimensions.items() if value <= 0 + ) + if invalid: + raise ValueError(f"{operation} dimensions must be positive, got {invalid}") + + +def _normalize_target(target_compute_capability: int) -> int: + if isinstance(target_compute_capability, bool): + raise TypeError("target_compute_capability must be an integer") + try: + target = index(target_compute_capability) + except TypeError as error: + raise TypeError("target_compute_capability must be an integer") from error + if target not in SUPPORTED_COMPUTE_CAPABILITIES: + supported = ", ".join(f"SM{value}" for value in SUPPORTED_COMPUTE_CAPABILITIES) + raise ValueError( + f"Unsupported score-recompute target SM{target}; supported targets are {supported}" + ) + return target + + +def _normalize_scale(scale: float, name: str) -> float: + if isinstance(scale, bool) or not isinstance(scale, Real): + raise TypeError(f"{name} must be a real number, got {type(scale).__name__}") + return float(scale) + + +class SparseScoreRecomputeOp(Op): + """Complete sparse score signature and architecture-specific configuration.""" + + def __init__( + self, + *, + q: TensorDesc[Any], + k: TensorDesc[Any], + per_head: TensorDesc[Any], + topk_indices: TensorDesc[Any], + output: TensorDesc[Any], + score_type: str, + softmax_scale: float, + target_compute_capability: int, + topk_length: TensorDesc[Any] | None = None, + qhead_per_kv_head: int | None = None, + topk_indices_global: bool = False, + ) -> None: + self.q = _require_desc(q, "q") + self.k = _require_desc(k, "k") + self.per_head = _require_desc(per_head, "per_head") + self.topk_indices = _require_desc(topk_indices, "topk_indices") + self.output = _require_desc(output, "output") + self.topk_length = _require_desc(topk_length, "topk_length", optional=True) + if score_type not in ("indexer", "attention"): + raise ValueError( + f"score_type must be 'indexer' or 'attention', got {score_type!r}" + ) + if not isinstance(topk_indices_global, bool): + raise TypeError( + f"topk_indices_global must be a bool, got {type(topk_indices_global).__name__}" + ) + + self.score_type = score_type + self.softmax_scale = _normalize_scale(softmax_scale, "softmax_scale") + self.target_compute_capability = _normalize_target(target_compute_capability) + self.requested_qhead_per_kv_head = qhead_per_kv_head + self.topk_indices_global = topk_indices_global + + self.qhead_per_kv_head: int | None = None + self.config: SparseScoreKernelConfig | SparseScoreSm90Config | None = None + + def check_support(self) -> bool: + self.qhead_per_kv_head = None + self.config = None + + _require_rank(self.q, 4, "Q") + _require_rank(self.k, 3, "K") + _require_rank( + self.per_head, 3, "weights" if self.score_type == "indexer" else "LSE" + ) + _require_rank(self.topk_indices, 3, "topk_indices") + _require_rank(self.output, 3, "output") + _require_dtype(self.q, data_type.BFLOAT16, "Q") + _require_dtype(self.k, data_type.BFLOAT16, "K") + _require_dtype( + self.per_head, + data_type.BFLOAT16 if self.score_type == "indexer" else data_type.FLOAT, + "weights" if self.score_type == "indexer" else "LSE", + ) + _require_dtype(self.topk_indices, data_type.INT32, "topk_indices") + _require_dtype(self.output, data_type.FLOAT, "output") + for desc, name in ( + (self.q, "Q"), + (self.k, "K"), + (self.per_head, "weights" if self.score_type == "indexer" else "LSE"), + (self.topk_indices, "topk_indices"), + (self.output, "output"), + ): + _require_compact(desc, name) + _require_contiguous_tail(self.q, 2, "Q") + for desc, name in ( + (self.k, "K"), + ( + self.per_head, + "weights" if self.score_type == "indexer" else "LSE", + ), + (self.topk_indices, "topk_indices"), + (self.output, "output"), + ): + _require_contiguous_tail(desc, 1, name) + + batch, seqlen_q, num_query_heads, head_dim = self.q.shape + k_batch, seqlen_k, k_head_dim = self.k.shape + topk_batch, topk_seqlen_q, topk = self.topk_indices.shape + _positive_dimensions( + { + "B": batch, + "S_q": seqlen_q, + "S_k": seqlen_k, + "H_q": num_query_heads, + "D": head_dim, + "topk": topk, + }, + "Sparse score-recompute", + ) + if k_batch != batch or k_head_dim != head_dim: + raise ValueError( + f"K shape must be {(batch, seqlen_k, head_dim)}, got {self.k.shape}" + ) + expected_per_head = (batch, seqlen_q, num_query_heads) + if self.per_head.shape != expected_per_head: + raise ValueError( + f"per-head tensor must have shape {expected_per_head}, got {self.per_head.shape}" + ) + if (topk_batch, topk_seqlen_q) != (batch, seqlen_q): + raise ValueError( + f"topk_indices leading dimensions must be {(batch, seqlen_q)}, got {self.topk_indices.shape[:2]}" + ) + if self.output.shape != self.topk_indices.shape: + raise ValueError( + f"output must have shape {self.topk_indices.shape}, got {self.output.shape}" + ) + + if self.topk_length is not None: + _require_rank(self.topk_length, 2, "topk_length") + _require_dtype(self.topk_length, data_type.INT32, "topk_length") + _require_compact(self.topk_length, "topk_length") + if self.topk_length.shape != (batch, seqlen_q): + raise ValueError( + f"topk_length must have shape {(batch, seqlen_q)}, got {self.topk_length.shape}" + ) + + qhead_per_kv_head = ( + num_query_heads + if self.requested_qhead_per_kv_head is None + else self.requested_qhead_per_kv_head + ) + if isinstance(qhead_per_kv_head, bool): + raise TypeError("qhead_per_kv_head must be an integer") + try: + qhead_per_kv_head = index(qhead_per_kv_head) + except TypeError as error: + raise TypeError("qhead_per_kv_head must be an integer") from error + if qhead_per_kv_head != num_query_heads: + raise ValueError( + f"qhead_per_kv_head must equal H_q ({num_query_heads}) for sparse MQA, got {qhead_per_kv_head}" + ) + + if self.target_compute_capability == 90: + if topk % 128: + raise ValueError( + f"SM90 sparse score recompute requires topk to be a multiple of 128, got {topk}" + ) + if qhead_per_kv_head <= 1: + raise ValueError( + "SM90 sparse score recompute requires qhead_per_kv_head > 1" + ) + tile_m = min(qhead_per_kv_head, 64) + if qhead_per_kv_head % tile_m: + raise ValueError( + f"qhead_per_kv_head ({qhead_per_kv_head}) must be divisible by tile_m ({tile_m})" + ) + config: SparseScoreKernelConfig | SparseScoreSm90Config = ( + SparseScoreSm90Config( + tile_m=tile_m, + tile_n=64, + kv_stage=2, + num_threads=256, + num_head_tiles=qhead_per_kv_head // tile_m, + ) + ) + else: + config = resolve_sparse_score_kernel_config( + score_type=self.score_type, + head_dim=head_dim, + qhead_per_kv_head=qhead_per_kv_head, + topk=topk, + have_topk_length=self.topk_length is not None, + ) + + self.qhead_per_kv_head = qhead_per_kv_head + self.config = config + return True + + +class DenseScoreRecomputeOp(Op): + """Complete dense score signature and architecture-specific configuration.""" + + def __init__( + self, + *, + q: TensorDesc[Any], + k: TensorDesc[Any], + per_head: TensorDesc[Any], + output: TensorDesc[Any], + denominator: TensorDesc[Any], + score_type: str, + scale: float, + ratio: int, + target_compute_capability: int, + qhead_per_kv_head: int | None = None, + is_thd: bool = False, + cu_seqlens_q: TensorDesc[Any] | None = None, + cu_seqlens_k: TensorDesc[Any] | None = None, + max_seqlen_q: int | None = None, + max_seqlen_k: int | None = None, + q_causal_offsets: TensorDesc[Any] | None = None, + ) -> None: + self.q = _require_desc(q, "q") + self.k = _require_desc(k, "k") + self.per_head = _require_desc(per_head, "per_head") + self.output = _require_desc(output, "output") + self.denominator = _require_desc(denominator, "denominator") + self.cu_seqlens_q = _require_desc(cu_seqlens_q, "cu_seqlens_q", optional=True) + self.cu_seqlens_k = _require_desc(cu_seqlens_k, "cu_seqlens_k", optional=True) + self.q_causal_offsets = _require_desc( + q_causal_offsets, "q_causal_offsets", optional=True + ) + if score_type not in ("indexer", "attention"): + raise ValueError( + f"score_type must be 'indexer' or 'attention', got {score_type!r}" + ) + if not isinstance(is_thd, bool): + raise TypeError(f"is_thd must be a bool, got {type(is_thd).__name__}") + if isinstance(ratio, bool): + raise TypeError("ratio must be an integer") + try: + ratio = index(ratio) + except TypeError as error: + raise TypeError("ratio must be an integer") from error + + self.score_type = score_type + self.scale = _normalize_scale(scale, "scale") + self.ratio = ratio + self.target_compute_capability = _normalize_target(target_compute_capability) + self.requested_qhead_per_kv_head = qhead_per_kv_head + self.is_thd = is_thd + self.requested_max_seqlen_q = max_seqlen_q + self.requested_max_seqlen_k = max_seqlen_k + + self.qhead_per_kv_head: int | None = None + self.max_seqlen_q: int | None = None + self.max_seqlen_k: int | None = None + self.config: DenseScoreKernelConfig | DenseScoreSm90Config | None = None + + @staticmethod + def _positive_static(value: int | None, name: str) -> int: + if value is None or isinstance(value, bool): + raise ValueError( + f"{name} must be provided as a positive integer for THD inputs" + ) + try: + value = index(value) + except TypeError as error: + raise TypeError(f"{name} must be an integer") from error + if value <= 0: + raise ValueError(f"{name} must be positive, got {value}") + return value + + def check_support(self) -> bool: + self.qhead_per_kv_head = None + self.max_seqlen_q = None + self.max_seqlen_k = None + self.config = None + if self.ratio < 1: + raise ValueError(f"ratio must be >= 1, got {self.ratio}") + + expected_rank = 3 if self.is_thd else 4 + _require_rank(self.q, expected_rank, "Q") + _require_rank(self.k, expected_rank, "K") + _require_rank( + self.per_head, + expected_rank - 1, + "weights" if self.score_type == "indexer" else "LSE", + ) + _require_rank(self.output, expected_rank - 1, "output") + _require_rank(self.denominator, expected_rank - 2, "denominator") + _require_dtype(self.q, data_type.BFLOAT16, "Q") + _require_dtype(self.k, data_type.BFLOAT16, "K") + _require_dtype( + self.per_head, + data_type.BFLOAT16 if self.score_type == "indexer" else data_type.FLOAT, + "weights" if self.score_type == "indexer" else "LSE", + ) + _require_dtype(self.output, data_type.FLOAT, "output") + _require_dtype(self.denominator, data_type.FLOAT, "denominator") + for desc, name in ( + (self.q, "Q"), + (self.k, "K"), + (self.per_head, "weights" if self.score_type == "indexer" else "LSE"), + (self.output, "output"), + (self.denominator, "denominator"), + ): + _require_compact(desc, name) + for desc, name in ((self.q, "Q"), (self.k, "K")): + _require_contiguous_tail(desc, 2, name) + for desc, name in ( + ( + self.per_head, + "weights" if self.score_type == "indexer" else "LSE", + ), + (self.output, "output"), + ): + _require_contiguous_tail(desc, 1, name) + + if self.is_thd: + total_q, num_query_heads, head_dim = self.q.shape + total_k, num_kv_heads, k_head_dim = self.k.shape + _positive_dimensions( + { + "total_q": total_q, + "total_k": total_k, + "H_q": num_query_heads, + "H_kv": num_kv_heads, + "D": head_dim, + }, + "Dense THD score-recompute", + ) + if self.per_head.shape != (total_q, num_query_heads): + raise ValueError( + f"THD per-head tensor must have shape {(total_q, num_query_heads)}, got {self.per_head.shape}" + ) + if self.cu_seqlens_q is None or self.cu_seqlens_k is None: + raise ValueError( + "THD dense score recompute requires both cu_seqlens_q and cu_seqlens_k" + ) + for desc, name in ( + (self.cu_seqlens_q, "cu_seqlens_q"), + (self.cu_seqlens_k, "cu_seqlens_k"), + ): + _require_rank(desc, 1, name) + _require_dtype(desc, data_type.INT32, name) + _require_compact(desc, name) + if self.cu_seqlens_q.shape != self.cu_seqlens_k.shape: + raise ValueError( + f"cu_seqlens_q and cu_seqlens_k shapes must match, got {self.cu_seqlens_q.shape} and {self.cu_seqlens_k.shape}" + ) + batch = self.cu_seqlens_q.shape[0] - 1 + if batch <= 0: + raise ValueError( + "cumulative sequence-length tensors must contain at least two entries" + ) + max_seqlen_q = self._positive_static( + self.requested_max_seqlen_q, "max_seqlen_q" + ) + max_seqlen_k = self._positive_static( + self.requested_max_seqlen_k, "max_seqlen_k" + ) + expected_output = (total_q, max_seqlen_k) + expected_denominator = (total_q,) + if self.target_compute_capability == 90: + raise NotImplementedError( + "SM90 THD score recompute uses host-side sequence-length reads and cannot be traced by JAX; use SM100+ or BSHD inputs" + ) + else: + if self.cu_seqlens_q is not None or self.cu_seqlens_k is not None: + raise ValueError( + "BSHD dense score recompute does not accept cu_seqlens_q or cu_seqlens_k" + ) + batch, seqlen_q, num_query_heads, head_dim = self.q.shape + k_batch, seqlen_k, num_kv_heads, k_head_dim = self.k.shape + _positive_dimensions( + { + "B": batch, + "S_q": seqlen_q, + "S_k": seqlen_k, + "H_q": num_query_heads, + "H_kv": num_kv_heads, + "D": head_dim, + }, + "Dense BSHD score-recompute", + ) + if k_batch != batch: + raise ValueError(f"K batch dimension must be {batch}, got {k_batch}") + if self.per_head.shape != (batch, seqlen_q, num_query_heads): + raise ValueError( + f"BSHD per-head tensor must have shape {(batch, seqlen_q, num_query_heads)}, got {self.per_head.shape}" + ) + max_seqlen_q, max_seqlen_k = seqlen_q, seqlen_k + expected_output = (batch, seqlen_q, seqlen_k) + expected_denominator = (batch, seqlen_q) + + if k_head_dim != head_dim: + raise ValueError(f"K head dimension must be {head_dim}, got {k_head_dim}") + if num_kv_heads != 1: + raise ValueError( + f"Dense score recompute currently requires MQA with H_kv=1, got {num_kv_heads}" + ) + if self.output.shape != expected_output: + raise ValueError( + f"output must have shape {expected_output}, got {self.output.shape}" + ) + if self.denominator.shape != expected_denominator: + raise ValueError( + f"denominator must have shape {expected_denominator}, got {self.denominator.shape}" + ) + + if self.q_causal_offsets is not None: + _require_rank(self.q_causal_offsets, 1, "q_causal_offsets") + _require_dtype(self.q_causal_offsets, data_type.INT32, "q_causal_offsets") + _require_compact(self.q_causal_offsets, "q_causal_offsets") + if self.q_causal_offsets.shape != (batch,): + raise ValueError( + f"q_causal_offsets must have shape {(batch,)}, got {self.q_causal_offsets.shape}" + ) + + inferred_qhead_per_kv_head = num_query_heads // num_kv_heads + qhead_per_kv_head = ( + inferred_qhead_per_kv_head + if self.requested_qhead_per_kv_head is None + else self.requested_qhead_per_kv_head + ) + if isinstance(qhead_per_kv_head, bool): + raise TypeError("qhead_per_kv_head must be an integer") + try: + qhead_per_kv_head = index(qhead_per_kv_head) + except TypeError as error: + raise TypeError("qhead_per_kv_head must be an integer") from error + if qhead_per_kv_head != inferred_qhead_per_kv_head: + raise ValueError( + f"qhead_per_kv_head must equal H_q / H_kv ({inferred_qhead_per_kv_head}), got {qhead_per_kv_head}" + ) + + if self.target_compute_capability == 90: + if qhead_per_kv_head <= 1: + raise ValueError( + "SM90 dense score recompute requires qhead_per_kv_head > 1" + ) + tile_m = min(qhead_per_kv_head, 64) + if qhead_per_kv_head % tile_m: + raise ValueError( + f"qhead_per_kv_head ({qhead_per_kv_head}) must be divisible by tile_m ({tile_m})" + ) + config: DenseScoreKernelConfig | DenseScoreSm90Config = ( + DenseScoreSm90Config( + tile_m=tile_m, + tile_n=64, + kv_stage=2, + num_threads=384, + num_head_tiles=qhead_per_kv_head // tile_m, + ) + ) + else: + config = resolve_dense_score_kernel_config( + score_type=self.score_type, + head_dim=head_dim, + qhead_per_kv_head=qhead_per_kv_head, + ) + + self.qhead_per_kv_head = qhead_per_kv_head + self.max_seqlen_q = max_seqlen_q + self.max_seqlen_k = max_seqlen_k + self.config = config + return True + + +__all__ = [ + "DenseScoreRecomputeOp", + "DenseScoreSm90Config", + "SUPPORTED_COMPUTE_CAPABILITIES", + "SparseScoreRecomputeOp", + "SparseScoreSm90Config", +] diff --git a/python/cudnn/deepseek_sparse_attention/score_recompute/sparse_score_recompute_sm100.py b/python/cudnn/deepseek_sparse_attention/score_recompute/sparse_score_recompute_sm100.py index bcb9a4890..eabd8853f 100644 --- a/python/cudnn/deepseek_sparse_attention/score_recompute/sparse_score_recompute_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/score_recompute/sparse_score_recompute_sm100.py @@ -348,6 +348,7 @@ def kernel( tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams, ): """Device-side kernel entry with CLC persistent scheduling.""" + mK = copy_utils.as_global_tensor(mK) seqlen_q_packed = cute.size(mQ.shape[0]) seqlen_k = cute.size(mK.shape[0]) warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/__init__.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/__init__.py index 039c5be6a..19498935a 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/__init__.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/__init__.py @@ -1,3 +1,20 @@ -from .api import SparseAttentionBackward, sparse_attention_backward_wrapper +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT -__all__ = ["SparseAttentionBackward", "sparse_attention_backward_wrapper"] +"""Lazy Torch API and framework-neutral sparse-attention backward exports.""" + +from ...common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ("SparseAttentionBackwardOp",), + "dsa_bwd_sm90": ("FlashAttentionDSABackwardSm90",), + "dsa_bwd_sm100": ("FlashAttentionDSABackwardSm100",), + "api": ( + "SparseAttentionBackward", + "sparse_attention_backward_wrapper", + ), + }, + submodules=("api", "dsa_bwd_sm90", "dsa_bwd_sm100", "op"), +) diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py index 3030f8198..db3a9603f 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/_interface_sm100.py @@ -59,7 +59,8 @@ def flash_attn_bwd_sm100( head_dim_v = 512 if head_dim == 576 else head_dim device = q.device - assert q.dtype in [torch.float16, torch.bfloat16] + # The SM100 kernel currently fixes element storage to BFloat16. + assert q.dtype == torch.bfloat16 assert q.dtype == kv.dtype == out.dtype == dout.dtype assert lse.dtype == torch.float32 assert attn_sink.dtype == torch.float32 @@ -99,7 +100,7 @@ def flash_attn_bwd_sm100( # Allocate workspace tensors acc_dtype = cutlass.Float32 - ws_lse_odo_shape = FlashAttentionDSABackwardSm100._get_workspace_size_LSE_OdO( + ws_lse_odo_shape = FlashAttentionDSABackwardSm100.get_workspace_size_lse_odo( total_S_q, head_dim, num_head, @@ -112,7 +113,7 @@ def flash_attn_bwd_sm100( device=device, ) - ws_dkv_shape = FlashAttentionDSABackwardSm100._get_workspace_size_dKV( + ws_dkv_shape = FlashAttentionDSABackwardSm100.get_workspace_size_dkv( total_S_kv, head_dim, batch_size, diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py index 7ce340394..7805d71ad 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/api.py @@ -12,9 +12,22 @@ import torch import cuda.bindings.driver as cuda +from cudnn import data_type from cudnn.api_base import APIBase, TupleDict from . import _interface_sm100 as _iface_sm100 +from .op import SparseAttentionBackwardOp + + +def _compact_for_kernel(desc): + """Describe tensors that the Torch backends normalize before lowering.""" + + return desc.compact_like( + cudnn_dtype=desc.cudnn_dtype, + shape=desc.shape, + name=desc.name, + init_value=desc.init_value, + ) class SparseAttentionBackward(APIBase): @@ -42,33 +55,76 @@ def __init__( self.attn_sink_desc = self._make_tensor_desc(sample_attn_sink, name="sample_attn_sink") self.topk_idxs_desc = self._make_tensor_desc(sample_topk_idxs, name="sample_topk_idxs") self.topk_length_desc = self._make_tensor_desc(sample_topk_length, name="sample_topk_length") - self.block_tile = int(block_tile) + self.dq_desc = ( + self._make_tensor_desc(sample_dq, name="sample_dq") + if sample_dq is not None + else self.q_desc.compact_like( + cudnn_dtype=self.q_desc.cudnn_dtype, + shape=self.q_desc.shape, + name="sample_dq", + ) + ) + self.dkv_desc = ( + self._make_tensor_desc(sample_dkv, name="sample_dkv") + if sample_dkv is not None + else self.kv_desc.compact_like( + cudnn_dtype=self.kv_desc.cudnn_dtype, + shape=self.kv_desc.shape, + name="sample_dkv", + ) + ) + self.d_sink_desc = self.attn_sink_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=self.attn_sink_desc.shape, + name="sample_d_sink", + ) + self._op = SparseAttentionBackwardOp( + q=_compact_for_kernel(self.q_desc), + kv=_compact_for_kernel(self.kv_desc), + output=_compact_for_kernel(self.out_desc), + doutput=_compact_for_kernel(self.dout_desc), + lse=_compact_for_kernel(self.lse_desc), + attn_sink=self.attn_sink_desc, + topk_idxs=self.topk_idxs_desc, + topk_length=self.topk_length_desc, + dq=self.dq_desc, + dkv=self.dkv_desc, + d_sink=self.d_sink_desc, + softmax_scale=softmax_scale, + block_tile=block_tile, + ) + self.block_tile = self._op.block_tile self.softmax_scale = softmax_scale def check_support(self) -> bool: - major, _ = torch.cuda.get_device_capability() + major, minor = torch.cuda.get_device_capability(self.q_desc.device) + compute_capability = major * 10 + minor self._runtime_error_if( - major < 9, - f"SparseAttentionBackward requires SM90+, found SM{major}", - ) - self._value_error_if( - self.q_desc.ndim != 3, - f"Q must be 3-D (total_S_q, H, D), got {self.q_desc.shape}", + compute_capability < 90, + f"SparseAttentionBackward requires SM90+, found SM{compute_capability}", ) - self._value_error_if( - self.kv_desc.ndim != 2, - f"KV must be 2-D (total_S_kv, D), got {self.kv_desc.shape}", - ) - self._check_dtype(self.q_desc, [torch.float16, torch.bfloat16], name="Q") - self._check_dtype( - self.kv_desc, - self.q_desc.dtype, - name="KV", - extra_error_msg="KV must have same dtype as Q", - ) - self._check_dtype(self.lse_desc, torch.float32, name="LSE") - self._check_dtype(self.attn_sink_desc, torch.float32, name="attn_sink") - self._check_dtype(self.topk_idxs_desc, torch.int32, name="topk_idxs") + self._op.check_support() + if compute_capability >= 100 and self.q_desc.cudnn_dtype != data_type.BFLOAT16: + raise ValueError("SparseAttentionBackward on SM100+ currently requires bfloat16 inputs") + + devices = { + desc.device + for desc in ( + self.q_desc, + self.kv_desc, + self.out_desc, + self.dout_desc, + self.lse_desc, + self.attn_sink_desc, + self.topk_idxs_desc, + self.topk_length_desc, + self.dq_desc, + self.dkv_desc, + self.d_sink_desc, + ) + if desc is not None + } + self._value_error_if(len(devices) != 1, f"All tensors must be on the same device, got {sorted(map(str, devices))}") self._is_supported = True return True @@ -94,8 +150,8 @@ def execute( softmax_scale: Optional[float] = None, current_stream: Optional[cuda.CUstream] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - major, _ = torch.cuda.get_device_capability() - scale = self.softmax_scale if softmax_scale is None else softmax_scale + major, _ = torch.cuda.get_device_capability(q.device) + scale = self._op.softmax_scale if softmax_scale is None else softmax_scale if major == 9: from . import _interface_sm90 as _iface_sm90 @@ -180,7 +236,7 @@ def sparse_attention_backward_wrapper( softmax_scale=softmax_scale, block_tile=block_tile, ) - assert obj.check_support() + obj.check_support() obj.compile() _cache_of_SparseAttentionBackwardObjects[key] = obj diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py index fd4c69cc6..44dd43119 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py @@ -148,7 +148,7 @@ def _setup_attributes(self): self.compute_tmastore_dQ_stage = 1 @staticmethod - def _get_workspace_size_LSE_OdO(q: int, d: int, h: int, b: int, acc_dtype: Type[cutlass.Numeric]): + def get_workspace_size_lse_odo(q: int, d: int, h: int, b: int, acc_dtype: Type[cutlass.Numeric]): # q is total seqlen, b=1 d = (d + 7) // 8 * 8 # round up to 8 q = (q + 7) // 8 * 8 # round up to 8 @@ -161,13 +161,17 @@ def _get_workspace_size_LSE_OdO(q: int, d: int, h: int, b: int, acc_dtype: Type[ return (b, h, q, workspace_bytes) @staticmethod - def _get_workspace_size_dKV(k: int, d: int, b: int, acc_dtype: Type[cutlass.Numeric]): + def get_workspace_size_dkv(k: int, d: int, b: int, acc_dtype: Type[cutlass.Numeric]): d = (d + 7) // 8 * 8 # round up to 8 k = (k + 7) // 8 * 8 # round up to 8 # FP32 versions of dKV workspace_bytes = d * acc_dtype.width // 8 return (b, 1, k, workspace_bytes) + # Compatibility aliases for the existing Torch interface. + _get_workspace_size_LSE_OdO = get_workspace_size_lse_odo + _get_workspace_size_dKV = get_workspace_size_dkv + def get_workspace_tensor( self, problem_shape: Tuple[Int32, Int32, Int32, Tuple[Int32, Int32]], @@ -563,25 +567,24 @@ class SharedStorage: stream=stream, ) - if cutlass.const_expr(self.same_hdim_kv): - dSink_grid = ( - cute.ceil_div(problem_shape[0], self.dSink_block_q), - problem_shape[3][0], - problem_shape[3][1], - ) - self.sum_dSink( - sum_OdO, - scaled_LSE, - mAttnSink, - mdSink, - problem_shape, - ).launch( - grid=dSink_grid, - block=[self.dSink_num_threads, 1, 1], - cluster=[1, 1, 1], - stream=stream, - min_blocks_per_mp=1, - ) + dSink_grid = ( + cute.ceil_div(problem_shape[0], self.dSink_block_q), + problem_shape[3][0], + problem_shape[3][1], + ) + self.sum_dSink( + sum_OdO, + scaled_LSE, + mAttnSink, + mdSink, + problem_shape, + ).launch( + grid=dSink_grid, + block=[self.dSink_num_threads, 1, 1], + cluster=[1, 1, 1], + stream=stream, + min_blocks_per_mp=1, + ) @cute.kernel def convert( @@ -665,21 +668,18 @@ def sum_OdO( lse_bhq = lse[bidy, (idx_q + offset, bidz)] sum_OdO_bhq = sum_OdO_scale * acc - if cutlass.const_expr(self.same_hdim_kv): - attn_sink_bh = attn_sink[bidy, (0, bidz)] + attn_sink_bh = attn_sink[bidy, (0, bidz)] - log2_e = -lse_scale - lse_log2 = lse_bhq * log2_e - sink_log2 = attn_sink_bh * log2_e - lse_max_log2 = cute.arch.fmax(lse_log2, sink_log2) - sum_exp2 = Float32(cute.math.exp2(lse_log2 - lse_max_log2) + cute.math.exp2(sink_log2 - lse_max_log2)) - lse_with_sink_log2 = lse_max_log2 + cute.math.log2(sum_exp2) - scaled_lse_bhq = -lse_with_sink_log2 + log2_e = -lse_scale + lse_log2 = lse_bhq * log2_e + sink_log2 = attn_sink_bh * log2_e + lse_max_log2 = cute.arch.fmax(lse_log2, sink_log2) + sum_exp2 = Float32(cute.math.exp2(lse_log2 - lse_max_log2) + cute.math.exp2(sink_log2 - lse_max_log2)) + lse_with_sink_log2 = lse_max_log2 + cute.math.log2(sum_exp2) + scaled_lse_bhq = -lse_with_sink_log2 - if lse_bhq == Float32(float("inf")): - scaled_lse_bhq = Float32(float("-inf")) - else: - scaled_lse_bhq = lse_scale * lse_bhq + if lse_bhq == Float32(float("inf")): + scaled_lse_bhq = Float32(float("-inf")) sum_OdO[bidy, (idx_q + offset, bidz)] = sum_OdO_bhq scaled_lse[bidy, (idx_q + offset, bidz)] = scaled_lse_bhq @@ -865,10 +865,14 @@ def bwd( if cutlass.const_expr(mTopkLength is not None): topk = mTopkLength[token_idx] + topk = cutlass.max(Int32(0), cutlass.min(topk, Int32(self.max_topk))) else: topk = mTopkIdxs.shape[0] - tile_count = cute.ceil_div(topk, self.block_tile) + # Keep every warp role on the same balanced one-tile pipeline when a + # query has no selected KV rows. The partial-tile predicates below + # turn that synthetic tile into all zeros. + tile_count = cutlass.max(Int32(1), cute.ceil_div(topk, self.block_tile)) if warp_idx == self.load_warp_id: cute.arch.setmaxregister_decrease(self.num_regs_load) @@ -1296,7 +1300,7 @@ def load_KV( load_mma_K_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.load_mma_K_stage) tile_index = tile_count - 1 - full_tiles = (topk % self.block_tile) == 0 + full_tiles = (topk > 0) and ((topk % self.block_tile) == 0) rows_per_warp = self.block_tile // self.num_load_KV_warps rTopkIdx = cute.make_rmem_tensor((rows_per_warp,), cutlass.Int32) @@ -2127,7 +2131,7 @@ def reduce_dKV( rTopkIdx = cute.make_rmem_tensor((8,), cutlass.Int32) if cutlass.const_expr(not self.same_hdim_kv): rTopkIdx_64 = cute.make_rmem_tensor((8,), cutlass.Int32) - full_tiles = (topk % self.block_tile) == 0 + full_tiles = (topk > 0) and ((topk % self.block_tile) == 0) while tile_index >= 0: # Preload topk indices into rmem (shared across all 4 store_dKV calls) for i in cutlass.range_constexpr(8): diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm90.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm90.py index 0b5aef54d..aad8bb9bc 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm90.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm90.py @@ -1174,9 +1174,12 @@ def mma_wg0( if const_expr(self.have_topk_length): topK = mTopkLength[batch_idx, seq_idx] + topK = cutlass.max(Int32(0), cutlass.min(topK, Int32(self.max_topk))) else: topK = self.max_topk - n_block_max = (topK + self.tile_n - 1) // self.tile_n + # A zero-length row still executes one fully predicated tile so + # both warp groups reach the same barriers and write zero dQ. + n_block_max = cutlass.max(Int32(1), (topK + self.tile_n - 1) // self.tile_n) topk_tail_rows = topK - (n_block_max - 1) * self.tile_n n_block = n_block_max - 1 @@ -1691,9 +1694,10 @@ def mma_wg1( if const_expr(self.have_topk_length): topK = mTopkLength[batch_idx, seq_idx] + topK = cutlass.max(Int32(0), cutlass.min(topK, Int32(self.max_topk))) else: topK = self.max_topk - n_block_max = (topK + self.tile_n - 1) // self.tile_n + n_block_max = cutlass.max(Int32(1), (topK + self.tile_n - 1) // self.tile_n) # scatter AtomicAdd — flat gmem dKVaccum for per-row addressing mdKVaccum_cur = mdKVaccum[None, head_idx_kv, batch_idx] diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py new file mode 100644 index 000000000..da79f0d53 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/jax.py @@ -0,0 +1,464 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for DeepSeek sparse-attention backward on SM90 and SM100+.""" + +from __future__ import annotations + +from typing import Any, Optional + +import jax + +from ... import data_type +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..._jax.compiler import compile_options_for_target +from .op import BLOCK_TILE, SparseAttentionBackwardOp + +_SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) +_SUPPORTED_COMPUTE_CAPABILITY_FAMILIES = (90, 100) + + +class SparseAttentionBackward(JaxApiBase): + """JAX callable specialized from sparse-attention tensor metadata.""" + + def __init__( + self, + sample_q: Any, + sample_kv: Any, + sample_out: Any, + sample_dout: Any, + sample_lse: Any, + sample_attn_sink: Any, + sample_topk_idxs: Any, + sample_dq: Any | None = None, + sample_dkv: Any | None = None, + sample_d_sink: Any | None = None, + sample_topk_length: Any | None = None, + softmax_scale: Optional[float] = None, + block_tile: int = BLOCK_TILE, + target_compute_capability: int | None = None, + ) -> None: + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.kv_desc = self._to_tensor_desc(sample_kv, "sample_kv") + self.out_desc = self._to_tensor_desc(sample_out, "sample_out") + self.dout_desc = self._to_tensor_desc(sample_dout, "sample_dout") + self.lse_desc = self._to_tensor_desc(sample_lse, "sample_lse") + self.attn_sink_desc = self._to_tensor_desc(sample_attn_sink, "sample_attn_sink") + self.topk_idxs_desc = self._to_tensor_desc(sample_topk_idxs, "sample_topk_idxs") + self.topk_length_desc = None if sample_topk_length is None else self._to_tensor_desc(sample_topk_length, "sample_topk_length") + + self.dq_desc = self._output_desc( + sample_dq, + source=self.q_desc, + cudnn_dtype=self.q_desc.cudnn_dtype, + shape=self.q_desc.shape, + name="sample_dq", + ) + self.dkv_desc = self._output_desc( + sample_dkv, + source=self.kv_desc, + cudnn_dtype=self.kv_desc.cudnn_dtype, + shape=self.kv_desc.shape, + name="sample_dkv", + ) + self.d_sink_desc = self._output_desc( + sample_d_sink, + source=self.attn_sink_desc, + cudnn_dtype=data_type.FLOAT, + shape=self.attn_sink_desc.shape, + name="sample_d_sink", + init_value=0.0, + ) + self._op = SparseAttentionBackwardOp( + q=self.q_desc, + kv=self.kv_desc, + output=self.out_desc, + doutput=self.dout_desc, + lse=self.lse_desc, + attn_sink=self.attn_sink_desc, + topk_idxs=self.topk_idxs_desc, + topk_length=self.topk_length_desc, + dq=self.dq_desc, + dkv=self.dkv_desc, + d_sink=self.d_sink_desc, + softmax_scale=softmax_scale, + block_tile=block_tile, + ) + self.target_compute_capability = target_compute_capability + self.compute_capability: int | None = None + + @staticmethod + def _output_desc( + sample: Any | None, + *, + source: JaxTensorDesc, + cudnn_dtype: data_type, + shape: tuple[int, ...], + name: str, + init_value: bool | int | float | None = None, + ) -> JaxTensorDesc: + if sample is not None: + desc = JaxApiBase._to_tensor_desc(sample, name, init_value=init_value) + else: + desc = source.compact_like( + cudnn_dtype=cudnn_dtype, + shape=shape, + name=name, + init_value=init_value, + ) + return desc + + def check_support(self) -> bool: + self._op.check_support() + compute_capability = self._resolve_compute_capability( + self.target_compute_capability, + _SUPPORTED_COMPUTE_CAPABILITIES, + "SparseAttentionBackward", + ) + family = self._compute_capability_family(compute_capability, _SUPPORTED_COMPUTE_CAPABILITY_FAMILIES) + if family == 100 and self.q_desc.cudnn_dtype != data_type.BFLOAT16: + raise ValueError("SparseAttentionBackward on SM100+ currently requires bfloat16 inputs") + self.compute_capability = compute_capability + return True + + def __call__( + self, + q: Any, + kv: Any, + out: Any, + dout: Any, + lse: Any, + attn_sink: Any, + topk_idxs: Any, + topk_length: Any | None = None, + ) -> TupleDict: + self.check_support() + if (topk_length is None) != (self.topk_length_desc is None): + raise ValueError("topk_length presence must match sample_topk_length") + + inputs = (q, kv, out, dout, lse, attn_sink, topk_idxs) + q_desc = self.q_desc.with_divisibility( + (None, None, self._op.head_dim) + ) + kv_desc = self.kv_desc.with_divisibility( + (None, self._op.head_dim) + ) + out_desc = self.out_desc.with_divisibility( + (None, None, self._op.head_dim_v) + ) + dout_desc = self.dout_desc.with_divisibility( + (None, None, self._op.head_dim_v) + ) + input_descs = ( + q_desc, + kv_desc, + out_desc, + dout_desc, + self.lse_desc, + self.attn_sink_desc, + self.topk_idxs_desc, + ) + if topk_length is not None: + inputs += (topk_length,) + input_descs += (self.topk_length_desc,) + + dq_desc = self.dq_desc.with_divisibility( + (None, None, self._op.head_dim) + ) + dkv_desc = self.dkv_desc.with_divisibility( + (None, self._op.head_dim) + ) + + dq, dkv, d_sink = self._call_kernel( + inputs, + launch=self._launch_kernel, + output_descs=(dq_desc, dkv_desc, self.d_sink_desc), + input_descs=input_descs, + workspace_descs=self._workspace_descs(), + compile_options=compile_options_for_target(self.compute_capability), + ) + return TupleDict(dq=dq, dkv=dkv, d_sink=d_sink) + + def _workspace_descs(self) -> tuple[JaxTensorDesc, ...]: + if self._architecture_family == 100: + q_rounded = _round_up(self._op.total_seqlen_q, 8) + kv_rounded = _round_up(self._op.total_seqlen_kv, 8) + head_dim_rounded = _round_up(self._op.head_dim, 8) + return ( + self.q_desc.compact_like( + cudnn_dtype=data_type.UINT8, + shape=(1, self._op.num_heads, q_rounded, 8), + name="workspace_lse_odo", + init_value=0, + ), + self.q_desc.compact_like( + cudnn_dtype=data_type.UINT8, + shape=(1, 1, kv_rounded, head_dim_rounded * 4), + name="workspace_dkv", + init_value=0, + ), + ) + + q_rounded = _round_up(self._op.total_seqlen_q, BLOCK_TILE) + kv_rounded = _round_up(self._op.total_seqlen_kv, BLOCK_TILE) + head_dim_rounded = _round_up(self._op.head_dim, 32) + workspaces = ( + self.q_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(1, q_rounded, self._op.num_heads), + name="workspace_dpsum", + ), + self.q_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(1, q_rounded, self._op.num_heads), + name="workspace_lse_log2", + ), + self.q_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(1, 1, kv_rounded * head_dim_rounded), + name="workspace_dkv_accum", + init_value=0.0, + ), + ) + if self.topk_length_desc is None: + workspaces += ( + self.q_desc.compact_like( + cudnn_dtype=data_type.INT32, + shape=(1,), + name="workspace_dummy_topk_length", + ), + ) + return workspaces + + @property + def _architecture_family(self) -> int: + if self.compute_capability is None: + raise RuntimeError("check_support() must resolve the compute capability before lowering") + family = self._compute_capability_family(self.compute_capability, _SUPPORTED_COMPUTE_CAPABILITY_FAMILIES) + if family is None: + raise RuntimeError(f"No sparse-attention backward kernel for SM{self.compute_capability}") + return family + + def _launch_kernel( + self, + stream: Any, + *arguments: Any, + ) -> None: + workspace_count = len(self._workspace_descs()) + input_count = len(arguments) - 3 - workspace_count + inputs = arguments[:input_count] + outputs = arguments[input_count : input_count + 3] + workspaces = arguments[input_count + 3 :] + if self._architecture_family == 90: + self._launch_sm90(inputs, outputs, workspaces, stream) + else: + self._launch_sm100(inputs, outputs, workspaces, stream) + + def _launch_sm100( + self, + inputs: tuple[Any, ...], + outputs: tuple[Any, ...], + workspaces: tuple[Any, ...], + stream: Any, + ) -> None: + import cutlass + + from .dsa_bwd_sm100 import FlashAttentionDSABackwardSm100 + + q, kv, out, dout, lse, attn_sink, topk_idxs, *optional = inputs + topk_length = optional[0] if optional else None + dq, dkv, d_sink = outputs + workspace_lse_odo, workspace_dkv = workspaces + kernel = FlashAttentionDSABackwardSm100( + head_dim=self._op.head_dim, + head_dim_v=self._op.head_dim_v, + block_tile=self._op.block_tile, + max_topk=self._op.max_topk, + ) + problem_shape = ( + cutlass.Int32(self._op.total_seqlen_q), + cutlass.Int32(self._op.total_seqlen_kv), + cutlass.Int32(self._op.head_dim), + (cutlass.Int32(self._op.num_heads), cutlass.Int32(1)), + ) + kernel( + problem_shape, + q, + kv, + out, + dout, + lse, + attn_sink, + topk_idxs, + topk_length, + dq, + dkv, + d_sink, + workspace_lse_odo, + workspace_dkv, + cutlass.Float32(self._op.softmax_scale), + stream, + ) + + def _launch_sm90( + self, + inputs: tuple[Any, ...], + outputs: tuple[Any, ...], + workspaces: tuple[Any, ...], + stream: Any, + ) -> None: + import cutlass + + from .dsa_bwd_sm90 import ( + FlashAttentionDSABackwardSm90, + _FlashAttentionDSABackwardPostprocessSm90, + _FlashAttentionDSABackwardPreprocessSm90, + ) + + q, kv, out, dout, lse, attn_sink, topk_idxs, *optional = inputs + dq, dkv, d_sink = outputs + workspace_dpsum, workspace_lse_log2, workspace_dkv_accum, *optional_workspaces = workspaces + topk_length = optional[0] if optional else optional_workspaces[0] + + q4 = _prepend_unit_dim(q) + kv4 = _as_bshd_kv(kv) + out4 = _prepend_unit_dim(out) + dout4 = _prepend_unit_dim(dout) + lse4 = _prepend_unit_dim(lse) + topk4 = _prepend_unit_dim(topk_idxs) + topk_length4 = _prepend_unit_dim(topk_length) + dq4 = _prepend_unit_dim(dq) + dkv4 = _as_bshd_kv(dkv) + + dtype = q.element_type + preprocess = _FlashAttentionDSABackwardPreprocessSm90( + dtype, + self._op.head_dim_v, + 90, + BLOCK_TILE, + num_threads=256, + ) + preprocess( + out4, + dout4, + workspace_dpsum, + lse4, + workspace_lse_log2, + attn_sink, + d_sink, + None, + None, + None, + stream, + ) + + backward = FlashAttentionDSABackwardSm90( + dtype, + self._op.head_dim, + self._op.head_dim_v, + self._op.num_heads, + tile_m=BLOCK_TILE, + tile_n=BLOCK_TILE, + KV_stage=1, + PdS_stage=1, + SdP_swapAB=False, + dKV_swapAB=False, + dQ_swapAB=False, + num_threads=256, + have_topk_length=self.topk_length_desc is not None, + max_topk=self._op.max_topk, + ) + backward( + q4, + kv4, + dout4, + workspace_lse_log2, + workspace_dpsum, + dq4, + workspace_dkv_accum, + topk4, + topk_length4, + cutlass.Float32(self._op.softmax_scale), + stream, + ) + + hdim_chunk = 64 if self._op.head_dim == 576 else min(128, self._op.head_dim) + postprocess = _FlashAttentionDSABackwardPostprocessSm90( + dtype, + hdim_chunk=hdim_chunk, + tile_n=BLOCK_TILE, + head_dim=self._op.head_dim, + num_threads=hdim_chunk, + N_hdim_chunks=self._op.head_dim // hdim_chunk, + ) + postprocess( + workspace_dkv_accum, + dkv4, + cutlass.Int32(self._op.total_seqlen_kv), + stream, + ) + + +def _prepend_unit_dim(tensor: Any) -> Any: + """Return a CuTe view with a leading batch dimension of one.""" + + import cutlass.cute as cute + + leading_stride = tensor.shape[0] * tensor.stride[0] + return cute.make_tensor( + tensor.iterator, + cute.make_layout((1, *tensor.shape), stride=(leading_stride, *tensor.stride)), + ) + + +def _as_bshd_kv(tensor: Any) -> Any: + """Return a CuTe ``(1, sequence, 1, head_dim)`` MQA view.""" + + import cutlass.cute as cute + + leading_stride = tensor.shape[0] * tensor.stride[0] + return cute.make_tensor( + tensor.iterator, + cute.make_layout( + (1, tensor.shape[0], 1, tensor.shape[1]), + stride=(leading_stride, tensor.stride[0], tensor.stride[0], tensor.stride[1]), + ), + ) + + +def _round_up(value: int, alignment: int) -> int: + return (value + alignment - 1) // alignment * alignment + + +@jax.jit(static_argnames=("softmax_scale", "block_tile", "target_compute_capability")) +def sparse_attention_backward_wrapper( + q: Any, + kv: Any, + out: Any, + dout: Any, + lse: Any, + attn_sink: Any, + topk_idxs: Any, + softmax_scale: Optional[float] = None, + topk_length: Any | None = None, + block_tile: int = BLOCK_TILE, + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute DeepSeek sparse-attention gradients from JAX arrays.""" + + return SparseAttentionBackward( + q, + kv, + out, + dout, + lse, + attn_sink, + topk_idxs, + sample_topk_length=topk_length, + softmax_scale=softmax_scale, + block_tile=block_tile, + target_compute_capability=target_compute_capability, + )(q, kv, out, dout, lse, attn_sink, topk_idxs, topk_length) + + +__all__ = ["SparseAttentionBackward", "sparse_attention_backward_wrapper"] diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/op.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/op.py new file mode 100644 index 000000000..cfc9bbd55 --- /dev/null +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/op.py @@ -0,0 +1,226 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral DeepSeek sparse-attention backward operation.""" + +from __future__ import annotations + +import math +from typing import Any, Optional + +from ... import data_type +from ...common.op import Op +from ...common.tensor_desc import TensorDesc + +SUPPORTED_HEAD_DIMS = (512, 576) +BLOCK_TILE = 64 + + +class SparseAttentionBackwardOp(Op): + """Complete logical signature for DeepSeek sparse-attention backward. + + The operation describes the flat MQA interface shared by the Hopper and + Blackwell implementations. Framework adapters own device selection, + allocation, and lowering. + """ + + def __init__( + self, + *, + q: TensorDesc[Any], + kv: TensorDesc[Any], + output: TensorDesc[Any], + doutput: TensorDesc[Any], + lse: TensorDesc[Any], + attn_sink: TensorDesc[Any], + topk_idxs: TensorDesc[Any], + dq: TensorDesc[Any], + dkv: TensorDesc[Any], + d_sink: TensorDesc[Any], + topk_length: Optional[TensorDesc[Any]] = None, + softmax_scale: Optional[float] = None, + block_tile: int = BLOCK_TILE, + ) -> None: + descriptors = ( + ("q", q), + ("kv", kv), + ("output", output), + ("doutput", doutput), + ("lse", lse), + ("attn_sink", attn_sink), + ("topk_idxs", topk_idxs), + ("dq", dq), + ("dkv", dkv), + ("d_sink", d_sink), + ) + for name, desc in descriptors: + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + if topk_length is not None and not isinstance(topk_length, TensorDesc): + raise TypeError(f"topk_length must be a TensorDesc or None, got {type(topk_length).__name__}") + + self.q = q + self.kv = kv + self.output = output + self.doutput = doutput + self.lse = lse + self.attn_sink = attn_sink + self.topk_idxs = topk_idxs + self.topk_length = topk_length + self.dq = dq + self.dkv = dkv + self.d_sink = d_sink + self.requested_softmax_scale = softmax_scale + self.block_tile = int(block_tile) + + self.total_seqlen_q: Optional[int] = None + self.total_seqlen_kv: Optional[int] = None + self.num_heads: Optional[int] = None + self.head_dim: Optional[int] = None + self.head_dim_v: Optional[int] = None + self.max_topk: Optional[int] = None + self.softmax_scale: Optional[float] = None + + def check_support(self) -> bool: + """Validate the complete tensor signature and static configuration.""" + + self.total_seqlen_q = None + self.total_seqlen_kv = None + self.num_heads = None + self.head_dim = None + self.head_dim_v = None + self.max_topk = None + self.softmax_scale = None + + self._check_ranks_and_dtypes() + + total_seqlen_q, num_heads, head_dim = self.q.shape + total_seqlen_kv, kv_head_dim = self.kv.shape + max_topk = self.topk_idxs.shape[1] + head_dim_v = 512 if head_dim == 576 else head_dim + + dimensions = { + "total_seqlen_q": total_seqlen_q, + "total_seqlen_kv": total_seqlen_kv, + "num_heads": num_heads, + "head_dim": head_dim, + "max_topk": max_topk, + } + invalid = ", ".join(f"{name}={value}" for name, value in dimensions.items() if value <= 0) + if invalid: + raise ValueError(f"Sparse-attention dimensions must be positive, got {invalid}") + if kv_head_dim != head_dim: + raise ValueError(f"KV head dimension must match Q ({head_dim}), got {kv_head_dim}") + if head_dim not in SUPPORTED_HEAD_DIMS: + raise ValueError(f"head_dim must be one of {SUPPORTED_HEAD_DIMS}, got {head_dim}") + if num_heads % BLOCK_TILE != 0: + raise ValueError(f"num_heads must be divisible by {BLOCK_TILE}, got {num_heads}") + if self.block_tile != BLOCK_TILE: + raise ValueError(f"block_tile must be {BLOCK_TILE}, got {self.block_tile}") + + self._check_shapes(total_seqlen_q, total_seqlen_kv, num_heads, head_dim, head_dim_v, max_topk) + self._check_compact_layouts() + + if self.requested_softmax_scale is None: + softmax_scale = 1.0 / math.sqrt(head_dim) + else: + try: + softmax_scale = float(self.requested_softmax_scale) + except (TypeError, ValueError) as error: + raise TypeError(f"softmax_scale must be a real scalar or None, got {self.requested_softmax_scale!r}") from error + if not math.isfinite(softmax_scale): + raise ValueError(f"softmax_scale must be finite, got {softmax_scale}") + + self.total_seqlen_q = total_seqlen_q + self.total_seqlen_kv = total_seqlen_kv + self.num_heads = num_heads + self.head_dim = head_dim + self.head_dim_v = head_dim_v + self.max_topk = max_topk + self.softmax_scale = softmax_scale + return True + + def _check_ranks_and_dtypes(self) -> None: + expected_ranks = ( + (self.q, 3, "Q"), + (self.kv, 2, "KV"), + (self.output, 3, "O"), + (self.doutput, 3, "dO"), + (self.lse, 2, "LSE"), + (self.attn_sink, 1, "attn_sink"), + (self.topk_idxs, 2, "topk_idxs"), + (self.dq, 3, "dQ"), + (self.dkv, 2, "dKV"), + (self.d_sink, 1, "d_sink"), + ) + if self.topk_length is not None: + expected_ranks += ((self.topk_length, 1, "topk_length"),) + for desc, rank, name in expected_ranks: + if desc.ndim != rank: + raise ValueError(f"{name} must have rank {rank}, got shape {desc.shape}") + + if self.q.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError(f"Q must have dtype float16 or bfloat16, got {self.q.dtype}") + for desc, name in ( + (self.kv, "KV"), + (self.output, "O"), + (self.doutput, "dO"), + (self.dq, "dQ"), + (self.dkv, "dKV"), + ): + if desc.cudnn_dtype != self.q.cudnn_dtype: + raise ValueError(f"{name} must have the same dtype as Q, got {desc.dtype}") + for desc, name in ((self.lse, "LSE"), (self.attn_sink, "attn_sink"), (self.d_sink, "d_sink")): + if desc.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"{name} must have dtype float32, got {desc.dtype}") + for desc, name in ((self.topk_idxs, "topk_idxs"), (self.topk_length, "topk_length")): + if desc is not None and desc.cudnn_dtype != data_type.INT32: + raise ValueError(f"{name} must have dtype int32, got {desc.dtype}") + + def _check_shapes( + self, + total_seqlen_q: int, + total_seqlen_kv: int, + num_heads: int, + head_dim: int, + head_dim_v: int, + max_topk: int, + ) -> None: + expected_shapes = ( + (self.output, (total_seqlen_q, num_heads, head_dim_v), "O"), + (self.doutput, (total_seqlen_q, num_heads, head_dim_v), "dO"), + (self.lse, (total_seqlen_q, num_heads), "LSE"), + (self.attn_sink, (num_heads,), "attn_sink"), + (self.topk_idxs, (total_seqlen_q, max_topk), "topk_idxs"), + (self.dq, (total_seqlen_q, num_heads, head_dim), "dQ"), + (self.dkv, (total_seqlen_kv, head_dim), "dKV"), + (self.d_sink, (num_heads,), "d_sink"), + ) + if self.topk_length is not None: + expected_shapes += ((self.topk_length, (total_seqlen_q,), "topk_length"),) + for desc, expected, name in expected_shapes: + if desc.shape != expected: + raise ValueError(f"{name} must have shape {expected}, got {desc.shape}") + + def _check_compact_layouts(self) -> None: + descriptors = ( + self.q, + self.kv, + self.output, + self.doutput, + self.lse, + self.attn_sink, + self.topk_idxs, + self.dq, + self.dkv, + self.d_sink, + ) + if self.topk_length is not None: + descriptors += (self.topk_length,) + for desc in descriptors: + if not desc.is_compact(tuple(reversed(range(desc.ndim)))): + name = desc.name or "tensor" + raise ValueError(f"{name} must be row-major contiguous, got stride {desc.stride}") + + +__all__ = ["BLOCK_TILE", "SUPPORTED_HEAD_DIMS", "SparseAttentionBackwardOp"] diff --git a/python/cudnn/deepseek_sparse_attention/utils/compiler.py b/python/cudnn/deepseek_sparse_attention/utils/compiler.py index 65b68e81f..31c9bfd5c 100644 --- a/python/cudnn/deepseek_sparse_attention/utils/compiler.py +++ b/python/cudnn/deepseek_sparse_attention/utils/compiler.py @@ -18,18 +18,7 @@ from functools import lru_cache -import torch - -# (compute_capability) → cute DSL --gpu-arch flag value. -# H100, B200/B300, and sm_100f require architecture-specific variants because -# the kernels use TMA / tcgen05 instructions that are only guaranteed to lower -# correctly under the matching SASS gencode. -_ARCH_MAP = { - (9, 0): "sm_90a", # Hopper H100 - (10, 0): "sm_100a", # Blackwell B200 - (10, 3): "sm_103a", # Blackwell Ultra B300 - (10, 7): "sm_100f", -} +from ...common.cute_arch import gpu_arch_flag_for_compute_capability @lru_cache(maxsize=None) @@ -39,15 +28,12 @@ def gpu_arch_flag() -> str: Cached because torch.cuda.get_device_capability() is cheap but the function gets called inside every cute.compile site. """ + import torch + if not torch.cuda.is_available(): raise RuntimeError("cute.compile requires CUDA; no GPU available") - cap = torch.cuda.get_device_capability() - arch = _ARCH_MAP.get(cap) - if arch is None: - raise RuntimeError( - f"Unsupported GPU compute capability {cap} for DSA CuTe kernels. " f"Add it to deepseek_sparse_attention/utils/compiler.py::_ARCH_MAP." - ) - return arch + major, minor = torch.cuda.get_device_capability() + return gpu_arch_flag_for_compute_capability(major * 10 + minor) def compile_options(extra: str = "") -> str: @@ -59,7 +45,8 @@ def compile_options(extra: str = "") -> str: Example: cute.compile(..., options=compile_options("--opt-level 3")) """ - parts = ["--enable-tvm-ffi", f"--gpu-arch {gpu_arch_flag()}"] + arch = gpu_arch_flag() + parts = ["--enable-tvm-ffi", f"--gpu-arch {arch}"] if extra: parts.append(extra) return " ".join(parts) diff --git a/python/cudnn/deepseek_sparse_attention/utils/copy.py b/python/cudnn/deepseek_sparse_attention/utils/copy.py index 82c50a477..b2526824a 100644 --- a/python/cudnn/deepseek_sparse_attention/utils/copy.py +++ b/python/cudnn/deepseek_sparse_attention/utils/copy.py @@ -19,6 +19,35 @@ def load_s2r(src: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor: return dst +@dsl_user_op +def as_global_tensor(tensor: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor: + """Cast a generic kernel tensor back to global memory for async copies.""" + + if tensor.memspace == cute.AddressSpace.gmem: + return tensor + if tensor.memspace != cute.AddressSpace.generic: + raise ValueError(f"Expected a global or generic tensor, got {tensor.memspace}") + + pointer = tensor.iterator + global_pointer = llvm.addrspacecast( + llvm.PointerType.get(cute.AddressSpace.gmem), + pointer.llvm_ptr, + loc=loc, + ip=ip, + ) + return cute.make_tensor( + cute.make_ptr( + tensor.element_type, + global_pointer, + cute.AddressSpace.gmem, + assumed_align=pointer.alignment, + loc=loc, + ip=ip, + ), + tensor.layout, + ) + + def tiled_copy_1d(dtype: Type[cutlass.Numeric], num_threads: int, num_copy_elems: int = 1, is_async: bool = False) -> cute.TiledCopy: num_copy_bits = min(128, num_copy_elems * dtype.width) copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp() diff --git a/python/cudnn/discrete_grouped_gemm/__init__.py b/python/cudnn/discrete_grouped_gemm/__init__.py index db09418bf..57c00cb64 100644 --- a/python/cudnn/discrete_grouped_gemm/__init__.py +++ b/python/cudnn/discrete_grouped_gemm/__init__.py @@ -1,26 +1,27 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -""" -Discrete-weight Grouped GEMM GLU Kernel Module +"""Lazy Torch exports for discrete-weight grouped GEMM operations. -This module provides forward and backward discrete-weight grouped GEMM with GLU -activation (SwiGLU/GeGLU) for MoE (Mixture of Experts) workloads on SM100+ GPUs. +JAX APIs are exported through :mod:`cudnn.jax`. """ -from .discrete_grouped_gemm_swiglu import ( - DiscreteGroupedGemmSwigluSm100, - discrete_grouped_gemm_swiglu_wrapper_sm100, -) +from ..common.operation_api import make_operation_api -from .discrete_grouped_gemm_dswiglu import ( - DiscreteGroupedGemmDswigluSm100, - discrete_grouped_gemm_dswiglu_wrapper_sm100, +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "discrete_grouped_gemm_swiglu": ( + "DiscreteGroupedGemmSwigluSm100", + "discrete_grouped_gemm_swiglu_wrapper_sm100", + ), + "discrete_grouped_gemm_dswiglu": ( + "DiscreteGroupedGemmDswigluSm100", + "discrete_grouped_gemm_dswiglu_wrapper_sm100", + ), + }, + submodules=( + "discrete_grouped_gemm_dswiglu", + "discrete_grouped_gemm_swiglu", + ), ) - -__all__ = [ - "DiscreteGroupedGemmSwigluSm100", - "discrete_grouped_gemm_swiglu_wrapper_sm100", - "DiscreteGroupedGemmDswigluSm100", - "discrete_grouped_gemm_dswiglu_wrapper_sm100", -] diff --git a/python/cudnn/discrete_grouped_gemm/_jax_common.py b/python/cudnn/discrete_grouped_gemm/_jax_common.py new file mode 100644 index 000000000..e1448123d --- /dev/null +++ b/python/cudnn/discrete_grouped_gemm/_jax_common.py @@ -0,0 +1,318 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared JAX metadata and validation for discrete grouped GEMMs.""" + +from __future__ import annotations + +import os +from typing import Any + +import jax.numpy as jnp + +from .. import data_type +from ..gemm.helpers import ( + block_scale_shape, + require_16_byte_alignment, + require_block_scale_layout, + require_cluster_shape, + require_compact_major, + require_mma_tiler, + require_tensor_shape, +) +from .._jax import JaxApiBase, JaxTensorDesc +from .._jax.datatypes import jax_to_cudnn_dtype, normalize_jax_dtype +from .._jax.gemm import ( + BLOCK_SCALE_MODE, + PROBABILITY_MODE, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, +) +from .._jax.layout import mode_from_layout, to_public_axes + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) +FIX_PAD_SIZE = 256 +MAX_EXPERTS = 1024 + +FP8_DTYPES = frozenset({data_type.FP8_E4M3, data_type.FP8_E5M2}) +AB_DTYPES = frozenset({data_type.FP4_E2M1, *FP8_DTYPES}) +SF_DTYPES = frozenset({data_type.FP8_E4M3, data_type.FP8_E8M0}) + + +def _require_dtype( + desc: JaxTensorDesc, allowed: frozenset[data_type], label: str +) -> data_type: + dtype = desc.cudnn_dtype + if dtype not in allowed: + expected = ", ".join(sorted(str(value) for value in allowed)) + raise ValueError( + f"{label} dtype must be one of {{{expected}}}, got {desc.dtype}" + ) + return dtype + + +class DiscreteGroupedGemmJaxBase(JaxApiBase): + """Common signature for the stacked-expert JAX binding. + + The Torch API accepts device pointer tables because Torch exposes stable + allocation addresses. JAX arrays do not expose such an ABI, and pointer + tables would hide the referenced buffers from XLA's liveness analysis. + This binding therefore accepts compact stacked B/SFB arrays as explicit + custom-call operands. The existing discrete kernel derives one pointer + per expert while initializing its TMA descriptors. + """ + + def _initialize_common( + self, + sample_a: Any, + sample_b: Any, + sample_sfa: Any, + sample_sfb: Any, + sample_padded_offsets: Any, + sample_alpha: Any, + *, + acc_dtype: Any | None, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int] | None, + sf_vec_size: int, + vector_f32: bool, + m_aligned: int, + use_dynamic_sched: bool, + a_layout: str, + b_layout: str, + output_layout: str, + ) -> None: + self.a_layout = a_layout + self.b_layout = b_layout + self.output_layout = output_layout + self.a_mode = gemm_a_mode(a_layout) + self.b_mode = gemm_b_mode(b_layout) + self.output_mode = gemm_output_mode(output_layout, name="output_layout") + self.scale_mode = BLOCK_SCALE_MODE + self.probability_mode = PROBABILITY_MODE + self.bias_mode = mode_from_layout("LN", kernel_axes="NL") + + self.compute_capability: int | None = None + self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=self.a_mode) + self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=self.b_mode) + self.sfa_desc = self._to_tensor_desc( + sample_sfa, "sample_sfa", mode=self.scale_mode + ) + self.sfb_desc = self._to_tensor_desc( + sample_sfb, "sample_sfb", mode=self.scale_mode + ) + self.padded_offsets_desc = self._to_tensor_desc( + sample_padded_offsets, "sample_padded_offsets" + ) + self.alpha_desc = self._to_tensor_desc(sample_alpha, "sample_alpha") + + self.acc_dtype = normalize_jax_dtype(acc_dtype, jnp.float32, "acc_dtype") + self.requested_mma_tiler_mn = mma_tiler_mn + self.requested_cluster_shape_mn = cluster_shape_mn + self.sf_vec_size = sf_vec_size + self.vector_f32 = vector_f32 + self.m_aligned = m_aligned + self.use_dynamic_sched = use_dynamic_sched + self.num_cluster_overlap_margin = int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ) + + self.m = self.n = self.k = self.expert_cnt = None + self.ab_dtype = self.sf_dtype = None + self.mma_tiler_mn = self.cluster_shape_mn = None + + def _check_common(self) -> None: + if self.a_desc.ndim != 3: + raise ValueError(f"A must have rank 3, got shape {self.a_desc.shape}") + if self.b_desc.ndim != 3: + raise ValueError(f"B must have rank 3, got shape {self.b_desc.shape}") + + m, k, a_groups = self.a_desc.shape + n, b_k, experts = self.b_desc.shape + if m < 0 or min(n, k, experts) <= 0: + raise ValueError( + "M must be non-negative and N, K, and expert count must be " + f"positive, got {(m, n, k, experts)}" + ) + if a_groups != 1: + raise ValueError( + f"A must flatten all expert rows into a singleton group dimension, got {a_groups}" + ) + if b_k != k: + raise ValueError(f"B K dimension must match A, got {b_k} and {k}") + if experts > MAX_EXPERTS: + raise ValueError( + f"expert count must not exceed {MAX_EXPERTS}, got {experts}" + ) + if m % FIX_PAD_SIZE: + raise ValueError( + f"A M dimension must be divisible by {FIX_PAD_SIZE}, got {m}" + ) + if n % 64: + raise ValueError(f"B N dimension must be divisible by 64, got {n}") + + require_tensor_shape( + self.padded_offsets_desc, (experts,), label="padded_offsets" + ) + require_tensor_shape(self.alpha_desc, (experts,), label="alpha") + if self.padded_offsets_desc.cudnn_dtype != data_type.INT32: + raise ValueError( + f"padded_offsets must have int32 dtype, got {self.padded_offsets_desc.dtype}" + ) + if self.alpha_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"alpha must have float32 dtype, got {self.alpha_desc.dtype}" + ) + + ab_dtype = _require_dtype(self.a_desc, AB_DTYPES, "A") + if self.b_desc.cudnn_dtype != ab_dtype: + raise ValueError( + f"A and B must have the same dtype, got {self.a_desc.dtype} and {self.b_desc.dtype}" + ) + sf_dtype = _require_dtype(self.sfa_desc, SF_DTYPES, "SFA") + if self.sfb_desc.cudnn_dtype != sf_dtype: + raise ValueError( + f"SFA and SFB must have the same dtype, got {self.sfa_desc.dtype} and {self.sfb_desc.dtype}" + ) + if self.sf_vec_size not in (16, 32): + raise ValueError(f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}") + if sf_dtype == data_type.FP8_E4M3 and self.sf_vec_size == 32: + raise ValueError("FP8_E4M3 scale factors require sf_vec_size=16") + if ab_dtype in FP8_DTYPES and self.sf_vec_size != 32: + raise ValueError("FP8 A and B require sf_vec_size=32") + + require_tensor_shape( + self.sfa_desc, block_scale_shape(m, k, 1, self.sf_vec_size), label="SFA" + ) + require_tensor_shape( + self.sfb_desc, + block_scale_shape(n, k, experts, self.sf_vec_size), + label="SFB", + ) + require_block_scale_layout(self.sfa_desc, "SFA") + require_block_scale_layout(self.sfb_desc, "SFB") + + a_major = require_compact_major(self.a_desc, "m", "k") + b_major = require_compact_major(self.b_desc, "n", "k") + if a_major != "k" or b_major != "k": + raise ValueError( + "The discrete grouped kernels currently require K-major A and B layouts" + ) + require_16_byte_alignment(self.a_desc) + require_16_byte_alignment(self.b_desc) + + if jax_to_cudnn_dtype(self.acc_dtype) != data_type.FLOAT: + raise ValueError(f"acc_dtype must be float32, got {self.acc_dtype}") + if not isinstance(self.vector_f32, bool): + raise TypeError( + f"vector_f32 must be a bool, got {type(self.vector_f32).__name__}" + ) + if not isinstance(self.use_dynamic_sched, bool): + raise TypeError( + f"use_dynamic_sched must be a bool, got {type(self.use_dynamic_sched).__name__}" + ) + if self.m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {self.m_aligned}") + + mma_tiler_mn = require_mma_tiler( + self.requested_mma_tiler_mn, + allowed_m=(128, 256), + allowed_n=(256,), + ) + cta_group_size = 2 if mma_tiler_mn[0] == 256 else 1 + default_cluster = (2, 1) if cta_group_size == 2 else (1, 1) + cluster_shape_mn = require_cluster_shape( + default_cluster + if self.requested_cluster_shape_mn is None + else self.requested_cluster_shape_mn, + cta_group_size=cta_group_size, + ) + if cluster_shape_mn[0] > 4 or cluster_shape_mn[1] > 4: + raise ValueError( + f"cluster_shape_mn entries must not exceed 4, got {cluster_shape_mn}" + ) + cluster_tile_m = cluster_shape_mn[0] // cta_group_size * mma_tiler_mn[0] + if cluster_tile_m not in (128, 256): + raise ValueError(f"cluster M tile must be 128 or 256, got {cluster_tile_m}") + + self.m, self.n, self.k, self.expert_cnt = m, n, k, experts + self.ab_dtype, self.sf_dtype = ab_dtype, sf_dtype + self.mma_tiler_mn, self.cluster_shape_mn = mma_tiler_mn, cluster_shape_mn + if m > 0: + self.compute_capability = self._resolve_compute_capability( + None, + SUPPORTED_COMPUTE_CAPABILITIES, + type(self).__name__, + ) + + def _canonical_desc( + self, + shape: tuple[int, ...], + dtype: Any, + name: str, + *, + mode: tuple[int, ...] | None = None, + init_value: bool | int | float | None = None, + ptr_assumed_align: int | None = None, + ) -> JaxTensorDesc: + return JaxTensorDesc.from_shape( + to_public_axes(shape, mode), + dtype, + name=name, + mode=mode, + init_value=init_value, + ptr_assumed_align=ptr_assumed_align, + ) + + def _workspace_desc(self, workspace_bytes: int) -> JaxTensorDesc: + if workspace_bytes <= 0: + raise ValueError( + f"kernel workspace size must be positive, got {workspace_bytes}" + ) + return self._canonical_desc( + (workspace_bytes,), + jnp.uint8, + "workspace", + ptr_assumed_align=128, + ) + + def _materialize_output_desc( + self, + desc: JaxTensorDesc | None, + ) -> Any | None: + """Materialize an inferred result without launching an empty GEMM.""" + + if desc is None: + return None + metadata = self._to_shape_dtype_struct(desc) + if desc.init_value is None: + return jnp.empty(metadata.shape, dtype=metadata.dtype) + return jnp.full(metadata.shape, desc.init_value, dtype=metadata.dtype) + + def _check_runtime_common( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + ) -> None: + self._check_tensor_signature(a_tensor, self.a_desc) + self._check_tensor_signature(b_tensor, self.b_desc) + self._check_tensor_signature(sfa_tensor, self.sfa_desc) + self._check_tensor_signature(sfb_tensor, self.sfb_desc) + self._check_tensor_signature(padded_offsets, self.padded_offsets_desc) + self._check_tensor_signature(alpha_tensor, self.alpha_desc) + + +__all__ = [ + "AB_DTYPES", + "DiscreteGroupedGemmJaxBase", + "FIX_PAD_SIZE", + "FP8_DTYPES", + "MAX_EXPERTS", + "SF_DTYPES", + "SUPPORTED_COMPUTE_CAPABILITIES", +] diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/__init__.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/__init__.py index e0355406a..2fcfd2688 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/__init__.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/__init__.py @@ -1,12 +1,17 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - DiscreteGroupedGemmDswigluSm100, - discrete_grouped_gemm_dswiglu_wrapper_sm100, -) +"""Lazy Torch API exports for discrete grouped GEMM dSwiGLU.""" + +from ...common.operation_api import make_operation_api -__all__ = [ - "DiscreteGroupedGemmDswigluSm100", - "discrete_grouped_gemm_dswiglu_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "api": ( + "DiscreteGroupedGemmDswigluSm100", + "discrete_grouped_gemm_dswiglu_wrapper_sm100", + ), + }, + submodules=("api",), +) diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/discrete_B_blockscaled_grouped_gemm_dglu_dbias.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/discrete_B_blockscaled_grouped_gemm_dglu_dbias.py index d52cb859c..e24c2a59b 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/discrete_B_blockscaled_grouped_gemm_dglu_dbias.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/discrete_B_blockscaled_grouped_gemm_dglu_dbias.py @@ -270,6 +270,7 @@ def __init__( expert_cnt: int, use_dynamic_sched: bool = False, act_func: str = "dswiglu", + stacked_expert_inputs: bool = False, ): """Initializes the configuration for a Blackwell blockscaled dense GEMM kernel. @@ -371,6 +372,10 @@ def __init__( self.vectorized_f32 = vectorized_f32 self.use_dynamic_sched = use_dynamic_sched + # See the forward kernel for the JAX liveness rationale. Torch keeps + # using device pointer tables; JAX passes stacked B/SFB operands whose + # expert addresses are derived here. + self.stacked_expert_inputs = stacked_expert_inputs # Amax reduction configuration self.num_epilog_warps = len(self.epilog_warp_id) @@ -638,10 +643,15 @@ def desc_init_kernel_device_ptrs( expert_idx = cute.arch.block_idx()[0] - b_ptr_tensor = cute.make_tensor(cute.make_ptr(cutlass.Int64, ptrs_b.toint(), AddressSpace.gmem, assumed_align=8), cute.make_layout((self.expert_cnt,))) - sfb_ptr_tensor = cute.make_tensor( - cute.make_ptr(cutlass.Int64, ptrs_sfb.toint(), AddressSpace.gmem, assumed_align=8), cute.make_layout((self.expert_cnt,)) - ) + if cutlass.const_expr(not self.stacked_expert_inputs): + b_ptr_tensor = cute.make_tensor( + cute.make_ptr(cutlass.Int64, ptrs_b.toint(), AddressSpace.gmem, assumed_align=8), + cute.make_layout((self.expert_cnt,)), + ) + sfb_ptr_tensor = cute.make_tensor( + cute.make_ptr(cutlass.Int64, ptrs_sfb.toint(), AddressSpace.gmem, assumed_align=8), + cute.make_layout((self.expert_cnt,)), + ) c1 = cutlass.Int32(1) c0 = cutlass.Int64(0) @@ -653,8 +663,12 @@ def desc_init_kernel_device_ptrs( stride_n = c1_64 stride_k = b_stride_size - b_ptr_val = b_ptr_tensor[expert_idx] - b_ptr = cute.make_ptr(self.b_dtype, b_ptr_val, AddressSpace.gmem) + if cutlass.const_expr(self.stacked_expert_inputs): + b_base_ptr = cute.make_ptr(self.b_dtype, ptrs_b.toint(), AddressSpace.gmem, assumed_align=16) + b_ptr = b_base_ptr + cutlass.Int64(expert_idx) * cutlass.Int64(n) * cutlass.Int64(k) + else: + b_ptr_val = b_ptr_tensor[expert_idx] + b_ptr = cute.make_ptr(self.b_dtype, b_ptr_val, AddressSpace.gmem) b_tensor_i = cute.make_tensor( b_ptr, cute.make_layout((n, k, c1), stride=(stride_n, stride_k, c0)), @@ -671,9 +685,13 @@ def desc_init_kernel_device_ptrs( workspace = TensormapWorkspace(workspace_ptr, ["b", "sfb"]) store_tma_desc(tma_atom_b, workspace.get_ptr("b", expert_idx)) - sfb_ptr_val = sfb_ptr_tensor[expert_idx] - sfb_ptr = cute.make_ptr(self.sf_dtype, sfb_ptr_val, AddressSpace.gmem) sfb_layout = blockscaled_utils.tile_atom_to_shape_SF((n, k, c1), self.sf_vec_size) + if cutlass.const_expr(self.stacked_expert_inputs): + sfb_base_ptr = cute.make_ptr(self.sf_dtype, ptrs_sfb.toint(), AddressSpace.gmem, assumed_align=16) + sfb_ptr = sfb_base_ptr + cutlass.Int64(expert_idx) * cutlass.Int64(cute.cosize(sfb_layout)) + else: + sfb_ptr_val = sfb_ptr_tensor[expert_idx] + sfb_ptr = cute.make_ptr(self.sf_dtype, sfb_ptr_val, AddressSpace.gmem) sfb_tensor_i = cute.make_tensor(sfb_ptr, sfb_layout) tma_atom_sfb, _ = cute.nvgpu.make_tiled_tma_atom_B( sfb_tma_op_arg, diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/jax.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/jax.py new file mode 100644 index 000000000..f8ecc0a6e --- /dev/null +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu/jax.py @@ -0,0 +1,577 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for the discrete-weight grouped GEMM dSwiGLU kernel.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._jax.compiler import compile_options_for_target +from ...gemm.helpers import ( + block_scale_shape, + require_16_byte_alignment, + require_compact_major, + require_tensor_shape, +) +from ..._jax import TupleDict +from ..._jax.datatypes import normalize_jax_dtype +from .._jax_common import ( + DiscreteGroupedGemmJaxBase, + FP8_DTYPES, + SUPPORTED_COMPUTE_CAPABILITIES, +) + + +class DiscreteGroupedGemmDswigluSm100(DiscreteGroupedGemmJaxBase): + """JAX callable for the discrete grouped dSwiGLU/dGeGLU kernel. + + B and SFB are compact stacked expert arrays rather than opaque device + pointer tables. Keeping them as custom-call operands gives XLA complete + buffer-liveness information while the kernel still initializes one TMA + descriptor pair per expert. + """ + + def __init__( + self, + sample_a: Any, + sample_b: Any, + sample_c: Any, + sample_sfa: Any, + sample_sfb: Any, + sample_padded_offsets: Any, + sample_alpha: Any, + sample_beta: Any, + sample_prob: Any, + *, + sample_norm_const: Any | None = None, + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + generate_dbias: bool = False, + act_func: str = "dswiglu", + epilogue_op: str | None = None, + use_dynamic_sched: bool = False, + a_layout: str = "LMK", + b_layout: str = "LNK", + output_layout: str = "LMN", + ) -> None: + self._initialize_common( + sample_a, + sample_b, + sample_sfa, + sample_sfb, + sample_padded_offsets, + sample_alpha, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + use_dynamic_sched=use_dynamic_sched, + a_layout=a_layout, + b_layout=b_layout, + output_layout=output_layout, + ) + self.c_desc = self._to_tensor_desc(sample_c, "sample_c", mode=self.output_mode) + self.beta_desc = self._to_tensor_desc(sample_beta, "sample_beta") + self.prob_desc = self._to_tensor_desc( + sample_prob, "sample_prob", mode=self.probability_mode + ) + self.norm_const_desc = ( + None + if sample_norm_const is None + else self._to_tensor_desc(sample_norm_const, "sample_norm_const") + ) + self.d_dtype = normalize_jax_dtype(d_dtype, jnp.bfloat16, "d_dtype") + self.discrete_col_sfd = discrete_col_sfd + self.generate_dbias = generate_dbias + self.act_func = act_func + self.epilogue_op = epilogue_op + + self.generate_sfd = False + self.d_row_desc = self.d_col_desc = self.dprob_desc = None + self.dbias_desc = self.sfd_row_desc = self.sfd_col_desc = self.amax_desc = None + + def check_support(self) -> bool: + self._check_common() + if self.act_func not in ("dswiglu", "dgeglu"): + raise ValueError( + f"act_func must be 'dswiglu' or 'dgeglu', got {self.act_func!r}" + ) + if self.epilogue_op not in (None, "none", "identity", "relu", "srelu"): + raise ValueError( + f"epilogue_op must be None, 'identity', 'relu', or 'srelu', got {self.epilogue_op!r}" + ) + for name, value in ( + ("discrete_col_sfd", self.discrete_col_sfd), + ("generate_dbias", self.generate_dbias), + ): + if not isinstance(value, bool): + raise TypeError(f"{name} must be a bool, got {type(value).__name__}") + + n_out = 2 * self.n + require_tensor_shape(self.c_desc, (self.m, n_out, 1), label="C") + if self.c_desc.cudnn_dtype not in ( + data_type.FLOAT, + data_type.HALF, + data_type.BFLOAT16, + ): + raise ValueError( + f"C must have float32, float16, or bfloat16 dtype, got {self.c_desc.dtype}" + ) + if require_compact_major(self.c_desc, "m", "n") != "n": + raise ValueError("C must use an N-major output layout") + require_16_byte_alignment(self.c_desc) + + require_tensor_shape(self.beta_desc, (self.expert_cnt,), label="beta") + if self.beta_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"beta must have float32 dtype, got {self.beta_desc.dtype}" + ) + require_tensor_shape(self.prob_desc, (self.m, 1, 1), label="prob") + if self.prob_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"prob must have float32 dtype, got {self.prob_desc.dtype}" + ) + if self.norm_const_desc is not None: + require_tensor_shape(self.norm_const_desc, (1,), label="norm_const") + if self.norm_const_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"norm_const must have float32 dtype, got {self.norm_const_desc.dtype}" + ) + + d_cudnn_dtype = self._output_cudnn_dtype(self.d_dtype) + if self.ab_dtype == data_type.FP4_E2M1: + if d_cudnn_dtype not in ( + data_type.HALF, + data_type.BFLOAT16, + data_type.FLOAT, + ): + raise ValueError(f"FP4 A and B do not support D dtype {self.d_dtype}") + elif d_cudnn_dtype not in ( + data_type.HALF, + data_type.BFLOAT16, + *FP8_DTYPES, + data_type.FP4_E2M1, + ): + raise ValueError(f"FP8 A and B do not support D dtype {self.d_dtype}") + + self.generate_sfd = ( + self.ab_dtype in FP8_DTYPES + and self.sf_dtype == data_type.FP8_E8M0 + and d_cudnn_dtype in FP8_DTYPES + ) + if self.generate_sfd and self.norm_const_desc is None: + raise ValueError( + "FP8 A/B and FP8 D require sample_norm_const so the kernel can generate SFD" + ) + if not self.generate_sfd and self.norm_const_desc is not None: + raise ValueError( + "sample_norm_const is only used by FP8 A/B to FP8 D configurations" + ) + if self.discrete_col_sfd and not self.generate_sfd: + raise ValueError("discrete_col_sfd requires generated SFD outputs") + + self.d_row_desc = self._canonical_desc( + (self.m, n_out, 1), self.d_dtype, "d_row_tensor", mode=self.output_mode + ) + self.d_col_desc = self._canonical_desc( + (self.m, n_out, 1), self.d_dtype, "d_col_tensor", mode=self.output_mode + ) + for desc, label in ((self.d_row_desc, "D_row"), (self.d_col_desc, "D_col")): + if require_compact_major(desc, "m", "n") != "n": + raise ValueError(f"{label} must use an N-major output layout") + require_16_byte_alignment(desc) + + self.dprob_desc = self._canonical_desc( + (self.m, 1, 1), + jnp.float32, + "dprob_tensor", + mode=self.probability_mode, + init_value=0.0, + ) + self.dbias_desc = ( + self._canonical_desc( + (self.expert_cnt, n_out, 1), + jnp.bfloat16, + "dbias_tensor", + init_value=0.0, + ) + if self.generate_dbias + else None + ) + if self.generate_sfd: + self.sfd_row_desc = self._canonical_desc( + block_scale_shape(self.m, n_out, 1, self.sf_vec_size), + self.sfa_desc.dtype, + "sfd_row_tensor", + mode=self.scale_mode, + ) + self.sfd_col_desc = self._canonical_desc( + block_scale_shape(n_out, self.m, 1, self.sf_vec_size), + self.sfa_desc.dtype, + "sfd_col_tensor", + mode=self.scale_mode, + ) + else: + self.sfd_row_desc = self.sfd_col_desc = None + self.amax_desc = ( + self._canonical_desc( + (self.expert_cnt, 2, 1), + jnp.float32, + "amax_tensor", + init_value=float("-inf"), + ) + if d_cudnn_dtype in (data_type.HALF, data_type.BFLOAT16) + else None + ) + return True + + @staticmethod + def _output_cudnn_dtype(dtype: Any) -> data_type: + from ..._jax.datatypes import jax_to_cudnn_dtype + + resolved = jax_to_cudnn_dtype(dtype) + allowed = { + data_type.FLOAT, + data_type.HALF, + data_type.BFLOAT16, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + data_type.FP4_E2M1, + } + if resolved not in allowed: + raise ValueError(f"d_dtype has unsupported dtype {dtype}") + return resolved + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Any, + prob_tensor: Any, + norm_const_tensor: Any | None = None, + *, + linear_offset: float | None = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + ) -> TupleDict: + self.check_support() + self._check_runtime_common( + a_tensor, b_tensor, sfa_tensor, sfb_tensor, padded_offsets, alpha_tensor + ) + self._check_tensor_signature(c_tensor, self.c_desc) + self._check_tensor_signature(beta_tensor, self.beta_desc) + self._check_tensor_signature(prob_tensor, self.prob_desc) + if (norm_const_tensor is None) != (self.norm_const_desc is None): + raise ValueError("norm_const_tensor presence must match sample_norm_const") + if self.norm_const_desc is not None: + self._check_tensor_signature(norm_const_tensor, self.norm_const_desc) + if linear_offset is None: + linear_offset = 1.0 if self.act_func == "dgeglu" else 0.0 + + if self.m == 0: + return TupleDict( + d_row_tensor=self._materialize_output_desc(self.d_row_desc), + d_col_tensor=self._materialize_output_desc(self.d_col_desc), + dprob_tensor=self._materialize_output_desc(self.dprob_desc), + dbias_tensor=self._materialize_output_desc(self.dbias_desc), + amax_tensor=self._materialize_output_desc(self.amax_desc), + sfd_row_tensor=self._materialize_output_desc(self.sfd_row_desc), + sfd_col_tensor=self._materialize_output_desc(self.sfd_col_desc), + ) + + import cutlass + import cutlass.cute as cute + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from .discrete_B_blockscaled_grouped_gemm_dglu_dbias import ( + BlockScaledDiscreteWeightDgluDbiasGroupedGemmKernel, + ) + + kernel = BlockScaledDiscreteWeightDgluDbiasGroupedGemmKernel( + sf_vec_size=self.sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(self.acc_dtype), + use_2cta_instrs=self.mma_tiler_mn[0] == 256, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=self.cluster_shape_mn, + vectorized_f32=self.vector_f32, + discrete_col_sfd=self.discrete_col_sfd, + expert_cnt=self.expert_cnt, + use_dynamic_sched=self.use_dynamic_sched, + act_func=self.act_func, + stacked_expert_inputs=True, + ) + max_active_clusters = self._get_max_active_clusters( + self.cluster_shape_mn[0] * self.cluster_shape_mn[1], + overlap_margin=self.num_cluster_overlap_margin, + ) + workspace_desc = self._workspace_desc(kernel.get_workspace_bytes()) + + inputs = [ + a_tensor, + b_tensor, + sfb_tensor, + c_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + ] + input_descs = [ + self.a_desc, + self.b_desc, + self.sfb_desc, + self.c_desc, + self.sfa_desc, + self.padded_offsets_desc, + self.alpha_desc, + self.beta_desc, + self.prob_desc, + ] + if self.norm_const_desc is not None: + inputs.append(norm_const_tensor) + input_descs.append(self.norm_const_desc) + + output_descs = [self.d_row_desc, self.d_col_desc, self.dprob_desc] + for desc in ( + self.dbias_desc, + self.sfd_row_desc, + self.sfd_col_desc, + self.amax_desc, + ): + if desc is not None: + output_descs.append(desc) + + has_norm = self.norm_const_desc is not None + has_dbias = self.dbias_desc is not None + has_sfd = self.sfd_row_desc is not None + has_amax = self.amax_desc is not None + + def launch(stream: Any, *args: Any) -> None: + arg_index = 0 + + def take() -> Any: + nonlocal arg_index + value = args[arg_index] + arg_index += 1 + return value + + a, b, sfb, c, sfa, offsets, alpha, beta, prob = (take() for _ in range(9)) + norm_const = take() if has_norm else None + d_row, d_col, dprob = (take() for _ in range(3)) + dbias = take() if has_dbias else None + sfd_row = take() if has_sfd else None + sfd_col = take() if has_sfd else None + amax = take() if has_amax else None + workspace = take() + if arg_index != len(args): + raise RuntimeError( + f"Unexpected discrete grouped dSwiGLU buffer count: consumed {arg_index}, received {len(args)}" + ) + + if self.epilogue_op in (None, "none", "identity"): + + def epilogue(x): + return x + elif self.epilogue_op == "relu": + + def epilogue(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) + else: + + def epilogue(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) ** 2 + + kernel( + a, + b.iterator, + sfb.iterator, + cutlass.Int32(self.n), + cutlass.Int32(self.k), + cutlass.Int64(self.k), + OperandMajorMode.K, + workspace.iterator, + c, + d_row, + d_col, + sfa, + sfd_row, + sfd_col, + amax, + norm_const, + offsets, + alpha, + beta, + prob, + dprob, + cutlass.Float32(linear_offset), + dbias, + max_active_clusters, + stream, + epilogue, + cutlass.Float32(geglu_alpha), + cutlass.Float32(glu_clamp_max), + cutlass.Float32(glu_clamp_min), + ) + + results = self._call_kernel( + tuple(inputs), + launch=launch, + output_descs=tuple(output_descs), + input_descs=tuple(input_descs), + workspace_descs=(workspace_desc,), + compile_options=compile_options_for_target(self.compute_capability), + ) + result_index = 0 + + def result() -> Any: + nonlocal result_index + value = results[result_index] + result_index += 1 + return value + + d_row_result, d_col_result, dprob_result = result(), result(), result() + dbias_result = result() if has_dbias else None + sfd_row_result = result() if has_sfd else None + sfd_col_result = result() if has_sfd else None + amax_result = result() if has_amax else None + return TupleDict( + d_row_tensor=d_row_result, + d_col_tensor=d_col_result, + dprob_tensor=dprob_result, + dbias_tensor=dbias_result, + amax_tensor=amax_result, + sfd_row_tensor=sfd_row_result, + sfd_col_tensor=sfd_col_result, + ) + + +@partial( + jax.jit, + static_argnames=( + "d_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "generate_dbias", + "act_func", + "epilogue_op", + "use_dynamic_sched", + "a_layout", + "b_layout", + "output_layout", + "linear_offset", + "geglu_alpha", + "glu_clamp_max", + "glu_clamp_min", + ), +) +def discrete_grouped_gemm_dswiglu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Any, + prob_tensor: Any, + norm_const_tensor: Any | None = None, + *, + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + generate_dbias: bool = False, + act_func: str = "dswiglu", + epilogue_op: str | None = None, + use_dynamic_sched: bool = False, + a_layout: str = "LMK", + b_layout: str = "LNK", + output_layout: str = "LMN", + linear_offset: float | None = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, +) -> TupleDict: + """Run discrete dSwiGLU with XLA-owned stacked expert operands.""" + + operation = DiscreteGroupedGemmDswigluSm100( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + sample_norm_const=norm_const_tensor, + d_dtype=d_dtype, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + generate_dbias=generate_dbias, + act_func=act_func, + epilogue_op=epilogue_op, + use_dynamic_sched=use_dynamic_sched, + a_layout=a_layout, + b_layout=b_layout, + output_layout=output_layout, + ) + return operation( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + norm_const_tensor, + linear_offset=linear_offset, + geglu_alpha=geglu_alpha, + glu_clamp_max=glu_clamp_max, + glu_clamp_min=glu_clamp_min, + ) + + +__all__ = [ + "DiscreteGroupedGemmDswigluSm100", + "SUPPORTED_COMPUTE_CAPABILITIES", + "discrete_grouped_gemm_dswiglu_wrapper_sm100", +] diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/__init__.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/__init__.py index dafdfc1c2..1d2297f71 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/__init__.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/__init__.py @@ -1,12 +1,17 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - DiscreteGroupedGemmSwigluSm100, - discrete_grouped_gemm_swiglu_wrapper_sm100, -) +"""Lazy Torch API exports for discrete grouped GEMM + SwiGLU.""" + +from ...common.operation_api import make_operation_api -__all__ = [ - "DiscreteGroupedGemmSwigluSm100", - "discrete_grouped_gemm_swiglu_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "api": ( + "DiscreteGroupedGemmSwigluSm100", + "discrete_grouped_gemm_swiglu_wrapper_sm100", + ), + }, + submodules=("api",), +) diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/discrete_B_blockscaled_grouped_gemm_glu_bias.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/discrete_B_blockscaled_grouped_gemm_glu_bias.py index 316cf3228..ca55a97a2 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/discrete_B_blockscaled_grouped_gemm_glu_bias.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/discrete_B_blockscaled_grouped_gemm_glu_bias.py @@ -277,6 +277,7 @@ def __init__( use_dynamic_sched: bool = False, act_func: str = "swiglu", enable_bias: bool = False, + stacked_expert_inputs: bool = False, ): """Initializes the configuration for a Blackwell blockscaled dense GEMM kernel. @@ -331,6 +332,14 @@ def __init__( self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE self.enable_bias = enable_bias + # JAX cannot safely pass an integer table containing addresses of + # otherwise-unreferenced buffers: XLA would not know that those + # buffers must remain live. In stacked mode, ``b_ptrs`` and + # ``sfb_ptrs`` are instead the base pointers of explicit B/SFB custom + # call operands. The descriptor initializer derives each expert's + # address from the compact outer expert dimension. Torch retains the + # original pointer-table ABI by using the default value. + self.stacked_expert_inputs = stacked_expert_inputs self.occupancy = 1 self.epilog_warp_id = (0, 1, 2, 3) self.mma_warp_id = 4 @@ -641,10 +650,15 @@ def desc_init_kernel_device_ptrs( expert_idx = cute.arch.block_idx()[0] - b_ptr_tensor = cute.make_tensor(cute.make_ptr(cutlass.Int64, ptrs_b.toint(), AddressSpace.gmem, assumed_align=8), cute.make_layout((self.expert_cnt,))) - sfb_ptr_tensor = cute.make_tensor( - cute.make_ptr(cutlass.Int64, ptrs_sfb.toint(), AddressSpace.gmem, assumed_align=8), cute.make_layout((self.expert_cnt,)) - ) + if cutlass.const_expr(not self.stacked_expert_inputs): + b_ptr_tensor = cute.make_tensor( + cute.make_ptr(cutlass.Int64, ptrs_b.toint(), AddressSpace.gmem, assumed_align=8), + cute.make_layout((self.expert_cnt,)), + ) + sfb_ptr_tensor = cute.make_tensor( + cute.make_ptr(cutlass.Int64, ptrs_sfb.toint(), AddressSpace.gmem, assumed_align=8), + cute.make_layout((self.expert_cnt,)), + ) c1 = cutlass.Int32(1) c0 = cutlass.Int64(0) @@ -657,8 +671,12 @@ def desc_init_kernel_device_ptrs( stride_k = b_stride_size # --- B TMA descriptor --- - b_ptr_val = b_ptr_tensor[expert_idx] - b_ptr = cute.make_ptr(self.b_dtype, b_ptr_val, AddressSpace.gmem) + if cutlass.const_expr(self.stacked_expert_inputs): + b_base_ptr = cute.make_ptr(self.b_dtype, ptrs_b.toint(), AddressSpace.gmem, assumed_align=16) + b_ptr = b_base_ptr + cutlass.Int64(expert_idx) * cutlass.Int64(n) * cutlass.Int64(k) + else: + b_ptr_val = b_ptr_tensor[expert_idx] + b_ptr = cute.make_ptr(self.b_dtype, b_ptr_val, AddressSpace.gmem) b_tensor_i = cute.make_tensor( b_ptr, cute.make_layout((n, k, c1), stride=(stride_n, stride_k, c0)), @@ -676,9 +694,13 @@ def desc_init_kernel_device_ptrs( store_tma_desc(tma_atom_b, workspace.get_ptr("b", expert_idx)) # --- SFB TMA descriptor --- - sfb_ptr_val = sfb_ptr_tensor[expert_idx] - sfb_ptr = cute.make_ptr(self.sf_dtype, sfb_ptr_val, AddressSpace.gmem) sfb_layout = blockscaled_utils.tile_atom_to_shape_SF((n, k, c1), self.sf_vec_size) + if cutlass.const_expr(self.stacked_expert_inputs): + sfb_base_ptr = cute.make_ptr(self.sf_dtype, ptrs_sfb.toint(), AddressSpace.gmem, assumed_align=16) + sfb_ptr = sfb_base_ptr + cutlass.Int64(expert_idx) * cutlass.Int64(cute.cosize(sfb_layout)) + else: + sfb_ptr_val = sfb_ptr_tensor[expert_idx] + sfb_ptr = cute.make_ptr(self.sf_dtype, sfb_ptr_val, AddressSpace.gmem) sfb_tensor_i = cute.make_tensor(sfb_ptr, sfb_layout) tma_atom_sfb, _ = cute.nvgpu.make_tiled_tma_atom_B( sfb_tma_op_arg, diff --git a/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/jax.py b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/jax.py new file mode 100644 index 000000000..669b6b6f0 --- /dev/null +++ b/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_swiglu/jax.py @@ -0,0 +1,548 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for the discrete-weight grouped GEMM + SwiGLU kernel.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._jax.compiler import compile_options_for_target +from ...gemm.helpers import ( + block_scale_shape, + require_16_byte_alignment, + require_compact_major, + require_tensor_shape, +) +from ..._jax import JaxTensorDesc, TupleDict +from ..._jax.datatypes import normalize_jax_dtype +from .._jax_common import ( + DiscreteGroupedGemmJaxBase, + FP8_DTYPES, + SUPPORTED_COMPUTE_CAPABILITIES, +) + + +class DiscreteGroupedGemmSwigluSm100(DiscreteGroupedGemmJaxBase): + """JAX callable for grouped GEMM + SwiGLU using stacked expert weights. + + ``sample_b`` and ``sample_sfb`` retain an explicit expert dimension. They + are passed to XLA as ordinary operands, so their storage remains live for + the custom call. During lowering the discrete kernel derives per-expert + addresses and builds the same TMA descriptor workspace used by the Torch + pointer-table API. + """ + + def __init__( + self, + sample_a: Any, + sample_b: Any, + sample_sfa: Any, + sample_sfb: Any, + sample_padded_offsets: Any, + sample_alpha: Any, + *, + sample_norm_const: Any | None = None, + sample_prob: Any | None = None, + sample_bias: Any | None = None, + c_dtype: Any | None = None, + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "swiglu", + use_dynamic_sched: bool = False, + a_layout: str = "LMK", + b_layout: str = "LNK", + output_layout: str = "LMN", + ) -> None: + self._initialize_common( + sample_a, + sample_b, + sample_sfa, + sample_sfb, + sample_padded_offsets, + sample_alpha, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + use_dynamic_sched=use_dynamic_sched, + a_layout=a_layout, + b_layout=b_layout, + output_layout=output_layout, + ) + self.norm_const_desc = ( + None + if sample_norm_const is None + else self._to_tensor_desc(sample_norm_const, "sample_norm_const") + ) + self._uses_implicit_prob = sample_prob is None + if sample_prob is None: + self.prob_desc = JaxTensorDesc.from_shape( + (1, 1, self.a_desc.shape[0]), + jnp.float32, + name="sample_prob", + mode=self.probability_mode, + ) + else: + self.prob_desc = self._to_tensor_desc( + sample_prob, "sample_prob", mode=self.probability_mode + ) + self.bias_desc = ( + None + if sample_bias is None + else self._to_tensor_desc(sample_bias, "sample_bias", mode=self.bias_mode) + ) + self.c_dtype = normalize_jax_dtype(c_dtype, jnp.bfloat16, "c_dtype") + self.d_dtype = normalize_jax_dtype(d_dtype, jnp.bfloat16, "d_dtype") + self.discrete_col_sfd = discrete_col_sfd + self.act_func = act_func + + self.generate_sfd = self.norm_const_desc is not None + self.c_desc = self.d_desc = self.d_col_desc = None + self.sfd_row_desc = self.sfd_col_desc = self.amax_desc = None + + def check_support(self) -> bool: + self._check_common() + if self.act_func not in ("swiglu", "geglu"): + raise ValueError( + f"act_func must be 'swiglu' or 'geglu', got {self.act_func!r}" + ) + if not isinstance(self.discrete_col_sfd, bool): + raise TypeError( + f"discrete_col_sfd must be a bool, got {type(self.discrete_col_sfd).__name__}" + ) + if self.discrete_col_sfd and not self.generate_sfd: + raise ValueError( + "discrete_col_sfd requires sample_norm_const and generated SFD outputs" + ) + + if self.norm_const_desc is not None: + require_tensor_shape(self.norm_const_desc, (1,), label="norm_const") + if self.norm_const_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"norm_const must have float32 dtype, got {self.norm_const_desc.dtype}" + ) + require_tensor_shape(self.prob_desc, (self.m, 1, 1), label="prob") + if self.prob_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"prob must have float32 dtype, got {self.prob_desc.dtype}" + ) + if self.bias_desc is not None: + require_tensor_shape( + self.bias_desc, (self.n, self.expert_cnt), label="bias" + ) + if self.bias_desc.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError( + f"bias must have float16 or bfloat16 dtype, got {self.bias_desc.dtype}" + ) + + c_cudnn_dtype = self._output_cudnn_dtype(self.c_dtype, "c_dtype") + d_cudnn_dtype = self._output_cudnn_dtype(self.d_dtype, "d_dtype") + if self.ab_dtype == data_type.FP4_E2M1: + if c_cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise NotImplementedError("FP4 A and B require float16 or bfloat16 C") + if d_cudnn_dtype not in ( + data_type.HALF, + data_type.BFLOAT16, + data_type.FLOAT, + ): + raise ValueError(f"FP4 A and B do not support D dtype {self.d_dtype}") + if self.sf_vec_size == 16 and d_cudnn_dtype == data_type.FLOAT: + raise NotImplementedError( + "FP4 A and B with sf_vec_size=16 do not support float32 D" + ) + elif d_cudnn_dtype not in ( + data_type.HALF, + data_type.BFLOAT16, + *FP8_DTYPES, + data_type.FP4_E2M1, + ): + raise ValueError(f"FP8 A and B do not support D dtype {self.d_dtype}") + + output_n = self.n // 2 + self.c_desc = self._canonical_desc( + (self.m, self.n, 1), self.c_dtype, "c_tensor", mode=self.output_mode + ) + self.d_desc = self._canonical_desc( + (self.m, output_n, 1), self.d_dtype, "d_tensor", mode=self.output_mode + ) + self.d_col_desc = self._canonical_desc( + (self.m, output_n, 1), self.d_dtype, "d_col_tensor", mode=self.output_mode + ) + for desc, label in ( + (self.c_desc, "C"), + (self.d_desc, "D"), + (self.d_col_desc, "D_col"), + ): + if require_compact_major(desc, "m", "n") != "n": + raise ValueError(f"{label} must use an N-major output layout") + require_16_byte_alignment(desc) + + if self.generate_sfd: + self.sfd_row_desc = self._canonical_desc( + block_scale_shape(self.m, output_n, 1, self.sf_vec_size), + self.sfa_desc.dtype, + "sfd_row_tensor", + mode=self.scale_mode, + ) + self.sfd_col_desc = self._canonical_desc( + block_scale_shape(output_n, self.m, 1, self.sf_vec_size), + self.sfa_desc.dtype, + "sfd_col_tensor", + mode=self.scale_mode, + ) + else: + self.sfd_row_desc = self.sfd_col_desc = None + self.amax_desc = ( + self._canonical_desc( + (self.expert_cnt, 1), + jnp.float32, + "amax_tensor", + init_value=float("-inf"), + ) + if d_cudnn_dtype in (data_type.HALF, data_type.BFLOAT16) + else None + ) + return True + + @staticmethod + def _output_cudnn_dtype(dtype: Any, name: str) -> data_type: + from ..._jax.datatypes import jax_to_cudnn_dtype + + resolved = jax_to_cudnn_dtype(dtype) + allowed = { + data_type.FLOAT, + data_type.HALF, + data_type.BFLOAT16, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + data_type.FP4_E2M1, + } + if resolved not in allowed: + raise ValueError(f"{name} has unsupported dtype {dtype}") + return resolved + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + norm_const_tensor: Any | None = None, + prob_tensor: Any | None = None, + bias_tensor: Any | None = None, + *, + linear_offset: float | None = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + ) -> TupleDict: + self.check_support() + self._check_runtime_common( + a_tensor, b_tensor, sfa_tensor, sfb_tensor, padded_offsets, alpha_tensor + ) + self._check_optional_input(norm_const_tensor, self.norm_const_desc) + if self._uses_implicit_prob: + if prob_tensor is not None: + raise ValueError( + "prob_tensor must be None when sample_prob was omitted" + ) + prob_tensor = jnp.ones((1, 1, self.m), dtype=jnp.float32) + else: + self._check_tensor_signature(prob_tensor, self.prob_desc) + self._check_optional_input(bias_tensor, self.bias_desc) + if linear_offset is None: + linear_offset = 1.0 if self.act_func == "geglu" else 0.0 + + if self.m == 0: + return TupleDict( + c_tensor=self._materialize_output_desc(self.c_desc), + d_tensor=self._materialize_output_desc(self.d_desc), + d_col_tensor=self._materialize_output_desc(self.d_col_desc), + amax_tensor=self._materialize_output_desc(self.amax_desc), + sfd_row_tensor=self._materialize_output_desc(self.sfd_row_desc), + sfd_col_tensor=self._materialize_output_desc(self.sfd_col_desc), + ) + + import cutlass + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from .discrete_B_blockscaled_grouped_gemm_glu_bias import ( + BlockScaledDiscreteWeightGroupedGemmBiasKernel, + ) + + kernel = BlockScaledDiscreteWeightGroupedGemmBiasKernel( + sf_vec_size=self.sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(self.acc_dtype), + use_2cta_instrs=self.mma_tiler_mn[0] == 256, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=self.cluster_shape_mn, + vectorized_f32=self.vector_f32, + generate_sfd=self.generate_sfd, + discrete_col_sfd=self.discrete_col_sfd, + expert_cnt=self.expert_cnt, + use_dynamic_sched=self.use_dynamic_sched, + act_func=self.act_func, + enable_bias=self.bias_desc is not None, + stacked_expert_inputs=True, + ) + max_active_clusters = self._get_max_active_clusters( + self.cluster_shape_mn[0] * self.cluster_shape_mn[1], + overlap_margin=self.num_cluster_overlap_margin, + ) + workspace_desc = self._workspace_desc(kernel.get_workspace_bytes()) + + inputs = [ + a_tensor, + b_tensor, + sfb_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + ] + input_descs = [ + self.a_desc, + self.b_desc, + self.sfb_desc, + self.sfa_desc, + self.padded_offsets_desc, + self.alpha_desc, + ] + for value, desc in ( + (norm_const_tensor, self.norm_const_desc), + (prob_tensor, self.prob_desc), + (bias_tensor, self.bias_desc), + ): + if desc is not None: + inputs.append(value) + input_descs.append(desc) + + output_descs = [self.c_desc, self.d_desc, self.d_col_desc] + for desc in ( + self.sfd_row_desc, + self.sfd_col_desc, + self.amax_desc, + ): + if desc is not None: + output_descs.append(desc) + + has_norm = self.norm_const_desc is not None + has_prob = True + has_bias = self.bias_desc is not None + has_sfd = self.sfd_row_desc is not None + has_amax = self.amax_desc is not None + + def launch(stream: Any, *args: Any) -> None: + arg_index = 0 + + def take() -> Any: + nonlocal arg_index + value = args[arg_index] + arg_index += 1 + return value + + a, b, sfb, sfa, offsets, alpha = (take() for _ in range(6)) + norm_const = take() if has_norm else None + prob = take() if has_prob else None + bias = take() if has_bias else None + c, d, d_col = (take() for _ in range(3)) + sfd_row = take() if has_sfd else None + sfd_col = take() if has_sfd else None + amax = take() if has_amax else None + workspace = take() + if arg_index != len(args): + raise RuntimeError( + f"Unexpected discrete grouped SwiGLU buffer count: consumed {arg_index}, received {len(args)}" + ) + + kernel( + a, + b.iterator, + sfb.iterator, + cutlass.Int32(self.n), + cutlass.Int32(self.k), + cutlass.Int64(self.k), + OperandMajorMode.K, + workspace.iterator, + c, + d, + d_col, + sfa, + sfd_row, + sfd_col, + amax, + norm_const, + offsets, + alpha, + prob, + bias, + max_active_clusters, + stream, + lambda x: x, + cutlass.Float32(linear_offset), + cutlass.Float32(geglu_alpha), + cutlass.Float32(glu_clamp_max), + cutlass.Float32(glu_clamp_min), + ) + + results = self._call_kernel( + tuple(inputs), + launch=launch, + output_descs=tuple(output_descs), + input_descs=tuple(input_descs), + workspace_descs=(workspace_desc,), + compile_options=compile_options_for_target(self.compute_capability), + ) + result_index = 0 + + def result() -> Any: + nonlocal result_index + value = results[result_index] + result_index += 1 + return value + + c_result, d_result, d_col_result = result(), result(), result() + sfd_row_result = result() if has_sfd else None + sfd_col_result = result() if has_sfd else None + amax_result = result() if has_amax else None + return TupleDict( + c_tensor=c_result, + d_tensor=d_result, + d_col_tensor=d_col_result, + amax_tensor=amax_result, + sfd_row_tensor=sfd_row_result, + sfd_col_tensor=sfd_col_result, + ) + + def _check_optional_input( + self, + value: Any | None, + desc: JaxTensorDesc | None, + ) -> None: + if (value is None) != (desc is None): + name = "optional input" if desc is None else desc.name + raise ValueError( + f"{name} presence must match the sample passed to the constructor" + ) + if desc is not None: + self._check_tensor_signature(value, desc) + + +@partial( + jax.jit, + static_argnames=( + "c_dtype", + "d_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "act_func", + "use_dynamic_sched", + "a_layout", + "b_layout", + "output_layout", + "linear_offset", + "geglu_alpha", + "glu_clamp_max", + "glu_clamp_min", + ), +) +def discrete_grouped_gemm_swiglu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + norm_const_tensor: Any | None = None, + prob_tensor: Any | None = None, + bias_tensor: Any | None = None, + *, + c_dtype: Any | None = None, + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "swiglu", + use_dynamic_sched: bool = False, + a_layout: str = "LMK", + b_layout: str = "LNK", + output_layout: str = "LMN", + linear_offset: float | None = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, +) -> TupleDict: + """Run the discrete kernel with XLA-owned stacked expert operands.""" + + operation = DiscreteGroupedGemmSwigluSm100( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + sample_norm_const=norm_const_tensor, + sample_prob=prob_tensor, + sample_bias=bias_tensor, + c_dtype=c_dtype, + d_dtype=d_dtype, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + act_func=act_func, + use_dynamic_sched=use_dynamic_sched, + a_layout=a_layout, + b_layout=b_layout, + output_layout=output_layout, + ) + return operation( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + norm_const_tensor, + prob_tensor, + bias_tensor, + linear_offset=linear_offset, + geglu_alpha=geglu_alpha, + glu_clamp_max=glu_clamp_max, + glu_clamp_min=glu_clamp_min, + ) + + +__all__ = [ + "DiscreteGroupedGemmSwigluSm100", + "SUPPORTED_COMPUTE_CAPABILITIES", + "discrete_grouped_gemm_swiglu_wrapper_sm100", +] diff --git a/python/cudnn/discrete_grouped_gemm/discrete_kernel_utils.py b/python/cudnn/discrete_grouped_gemm/discrete_kernel_utils.py index 992747d0c..28d14bb01 100644 --- a/python/cudnn/discrete_grouped_gemm/discrete_kernel_utils.py +++ b/python/cudnn/discrete_grouped_gemm/discrete_kernel_utils.py @@ -38,9 +38,12 @@ - Kernel helper functions that don't depend on kernel instance state """ -from typing import Type, Tuple, Union +from __future__ import annotations -import torch +from typing import TYPE_CHECKING, Any, Type, Tuple, Union + +if TYPE_CHECKING: + import torch import cutlass import cutlass.cute as cute @@ -69,7 +72,9 @@ # --------------------------------------------------------------------------- -def _require_pointer_tensor(ptrs: torch.Tensor, name: str, expected_len: int | None = None) -> None: +def _require_pointer_tensor(ptrs: Any, name: str, expected_len: int | None = None) -> None: + import torch + if ptrs.dtype != torch.int64: raise ValueError(f"{name} must be int64, got {ptrs.dtype}") if ptrs.ndim != 1: @@ -399,6 +404,8 @@ def silu_f32_geglu_scaled(a: Union[float, Float32], fastmath: bool = False) -> U def sigmoid(x): """PyTorch reference sigmoid using exp2 for numerical consistency.""" + import torch + LOG2_E = 1.4426950408889634 exp_x = torch.exp2(x * (-LOG2_E)) ret = 1.0 / (exp_x + 1.0) @@ -415,6 +422,8 @@ def compute_reference_amax(output_tensor: torch.Tensor) -> float: Returns: float: reference amax value """ + import torch + if output_tensor.dtype != torch.float32: output_fp32 = output_tensor.float() else: @@ -444,6 +453,8 @@ def compare_and_report_mismatches( rtol: Relative tolerance max_mismatches: Maximum number of mismatches to report """ + import torch + if gpu_tensor.is_cuda: gpu_data = gpu_tensor.cpu() else: diff --git a/python/cudnn/gemm/__init__.py b/python/cudnn/gemm/__init__.py new file mode 100644 index 000000000..acccf6a26 --- /dev/null +++ b/python/cudnn/gemm/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral implementation shared by dense GEMM operations.""" diff --git a/python/cudnn/gemm/helpers.py b/python/cudnn/gemm/helpers.py new file mode 100644 index 000000000..c3bd54f1c --- /dev/null +++ b/python/cudnn/gemm/helpers.py @@ -0,0 +1,253 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral validation helpers for dense GEMM operations.""" + +from __future__ import annotations + +from operator import index +from typing import Iterable + +from .. import data_type +from ..common.tensor_desc import TensorDesc + +_STANDARD_MMA_M = (128, 256) +_STANDARD_MMA_N = tuple(range(32, 257, 32)) +BLOCK_SCALE_STRIDE_ORDER = (3, 1, 0, 4, 2, 5) + +_DATA_TYPE_BITS_BY_NAME = { + "BOOLEAN": 8, + "INT4": 4, + "INT8": 8, + "UINT8": 8, + "HALF": 16, + "BFLOAT16": 16, + "FLOAT": 32, + "DOUBLE": 64, + "INT32": 32, + "INT64": 64, + "FP4_E2M1": 4, + "FP8_E4M3": 8, + "FP8_E5M2": 8, + "FP8_E8M0": 8, + "FAST_FLOAT_FOR_FP8": 32, +} + + +def _require_tensor_desc(tensor: TensorDesc, label: str) -> None: + if not isinstance(tensor, TensorDesc): + raise TypeError(f"{label} must be a TensorDesc, got {type(tensor).__name__}") + + +def _integer_pair(value: Iterable[int], label: str) -> tuple[int, int]: + try: + values = tuple(value) + except TypeError as error: + raise TypeError(f"{label} must contain two integers, got {value!r}") from error + + if len(values) != 2: + raise ValueError(f"{label} must contain two integers, got {values}") + + normalized = [] + for entry in values: + if isinstance(entry, bool): + raise TypeError(f"{label} must contain integers, got {values}") + try: + normalized.append(index(entry)) + except TypeError as error: + raise TypeError(f"{label} must contain integers, got {values}") from error + return normalized[0], normalized[1] + + +def _format_values(values: Iterable[int]) -> str: + return "{" + ", ".join(str(value) for value in values) + "}" + + +def require_gemm_inputs(a: TensorDesc, b: TensorDesc) -> tuple[int, int, int, int]: + """Validate canonical ``A[M,K,L]`` and ``B[N,K,L]`` descriptors. + + Returns ``(M, N, K, L)``. Layout and dtype validation are deliberately + separate because their supported combinations differ between kernels. + """ + + _require_tensor_desc(a, "A") + _require_tensor_desc(b, "B") + if a.ndim != 3: + raise ValueError(f"A must have rank 3, got shape {a.shape}") + if b.ndim != 3: + raise ValueError(f"B must have rank 3, got shape {b.shape}") + + m, k, l = a.shape + n, b_k, b_l = b.shape + for axis, extent in (("M", m), ("N", n), ("K", k), ("L", l)): + if extent <= 0: + raise ValueError(f"{axis} must be positive, got {extent}") + + if (b_k, b_l) != (k, l): + raise ValueError(f"B shape mismatch: expected (N, {k}, {l}), got {b.shape}") + return m, n, k, l + + +def block_scale_shape( + rows: int, + k: int, + batch: int, + sf_vec_size: int, +) -> tuple[int, int, int, int, int, int]: + """Return the canonical packed scale-factor shape for a dense GEMM.""" + + if isinstance(sf_vec_size, bool): + raise TypeError(f"sf_vec_size must be an integer, got {sf_vec_size!r}") + try: + sf_vec_size = index(sf_vec_size) + except TypeError as error: + raise TypeError(f"sf_vec_size must be an integer, got {sf_vec_size!r}") from error + if sf_vec_size <= 0: + raise ValueError(f"sf_vec_size must be positive, got {sf_vec_size}") + row_tiles = (rows + 127) // 128 + scale_k = (k + sf_vec_size - 1) // sf_vec_size + k_tiles = (scale_k + 3) // 4 + return (32, 4, row_tiles, 4, k_tiles, batch) + + +def require_block_scale_layout(tensor: TensorDesc, label: str) -> None: + """Require the packed scale-factor layout used by dense SM100 GEMMs.""" + + _require_tensor_desc(tensor, label) + if not tensor.is_compact(BLOCK_SCALE_STRIDE_ORDER): + raise ValueError( + f"{label} must use the packed block-scale layout with stride order " + f"{BLOCK_SCALE_STRIDE_ORDER}, got stride {tensor.stride} and " + f"stride order {tensor.stride_order}" + ) + + +def require_tensor_shape( + tensor: TensorDesc, + expected: tuple[int, ...], + *, + label: str | None = None, +) -> None: + """Require a descriptor to have an exact logical shape.""" + + tensor_label = label or getattr(tensor, "name", "") or "Tensor" + _require_tensor_desc(tensor, tensor_label) + expected = tuple(expected) + if tensor.shape != expected: + raise ValueError(f"{tensor_label} must have shape {expected}, got {tensor.shape}") + + +def require_compact_major( + tensor: TensorDesc, + mode0_label: str, + mode1_label: str, +) -> str: + """Resolve the compact major mode of a canonical rank-three GEMM tensor. + + Canonical GEMM tensors may make mode 0 or mode 1 contiguous. The batch + mode (mode 2) must remain outermost. The returned label is supplied by the + caller, for example ``("m", "k")`` for ``A[M,K,L]``. + """ + + tensor_label = getattr(tensor, "name", "") or "Tensor" + _require_tensor_desc(tensor, tensor_label) + if tensor.ndim != 3: + raise ValueError(f"{tensor_label} must have rank 3, got shape {tensor.shape}") + + labels_by_order = { + (0, 1, 2): mode0_label, + (1, 0, 2): mode1_label, + } + major = labels_by_order.get(tensor.stride_order) + if major is None or not tensor.is_compact(tensor.stride_order): + raise ValueError(f"{tensor_label} must be compact {mode0_label}-major or " f"{mode1_label}-major, got shape {tensor.shape} and stride {tensor.stride}") + return major + + +def data_type_bits(dtype: data_type) -> int: + """Return the storage width of a canonical cuDNN data type.""" + + if not isinstance(dtype, data_type): + raise TypeError(f"dtype must be a cudnn.data_type, got {type(dtype).__name__}") + + for name, bits in _DATA_TYPE_BITS_BY_NAME.items(): + member = getattr(data_type, name, None) + if member is not None and dtype == member: + return bits + raise ValueError(f"Unsupported cuDNN data type {dtype}") + + +def require_16_byte_alignment(tensor: TensorDesc) -> None: + """Require the tensor's contiguous extent to cover whole 16-byte units. + + This is the logical-extent requirement used by dense GEMM TMA operands; + framework adapters remain responsible for the base-pointer alignment. + ``require_compact_major`` should be called first for GEMM matrices. + """ + + tensor_label = getattr(tensor, "name", "") or "Tensor" + _require_tensor_desc(tensor, tensor_label) + if tensor.ndim == 0: + raise ValueError(f"{tensor_label} must have at least one dimension") + + contiguous_dimension = tensor.stride_order[0] + bits = data_type_bits(tensor.cudnn_dtype) + extent = tensor.shape[contiguous_dimension] + if extent * bits % 128 != 0: + required_multiple = 128 // bits + raise ValueError(f"{tensor_label} contiguous extent must be a multiple of " f"{required_multiple} elements for 16-byte alignment, got {extent}") + + +def require_mma_tiler( + mma_tiler_mn: Iterable[int], + *, + allowed_m: Iterable[int] = _STANDARD_MMA_M, + allowed_n: Iterable[int] = _STANDARD_MMA_N, +) -> tuple[int, int]: + """Validate and normalize an ``(M, N)`` MMA tile.""" + + mma_m, mma_n = _integer_pair(mma_tiler_mn, "mma_tiler_mn") + allowed_m = tuple(allowed_m) + allowed_n = tuple(allowed_n) + if mma_m not in allowed_m: + raise ValueError(f"mma_tiler_mn[0] must be in {_format_values(allowed_m)}, got {mma_m}") + if mma_n not in allowed_n: + raise ValueError(f"mma_tiler_mn[1] must be in {_format_values(allowed_n)}, got {mma_n}") + return mma_m, mma_n + + +def require_cluster_shape( + cluster_shape_mn: Iterable[int], + *, + cta_group_size: int = 1, + max_cluster_size: int = 16, +) -> tuple[int, int]: + """Validate and normalize an ``(M, N)`` Blackwell cluster shape.""" + + cluster_m, cluster_n = _integer_pair(cluster_shape_mn, "cluster_shape_mn") + if isinstance(cta_group_size, bool) or not isinstance(cta_group_size, int) or cta_group_size <= 0: + raise ValueError(f"cta_group_size must be a positive integer, got {cta_group_size!r}") + if isinstance(max_cluster_size, bool) or not isinstance(max_cluster_size, int) or max_cluster_size <= 0: + raise ValueError(f"max_cluster_size must be a positive integer, got {max_cluster_size!r}") + + if cluster_m <= 0 or cluster_n <= 0 or cluster_m & (cluster_m - 1) or cluster_n & (cluster_n - 1): + raise ValueError("cluster_shape_mn entries must be positive powers of two, " f"got {(cluster_m, cluster_n)}") + if cluster_m * cluster_n > max_cluster_size: + raise ValueError(f"cluster_shape_mn product must not exceed {max_cluster_size}, " f"got {(cluster_m, cluster_n)}") + if cluster_m % cta_group_size != 0: + raise ValueError(f"cluster_shape_mn[0] must be divisible by CTA group size " f"{cta_group_size}, got {cluster_m}") + return cluster_m, cluster_n + + +__all__ = [ + "BLOCK_SCALE_STRIDE_ORDER", + "block_scale_shape", + "data_type_bits", + "require_16_byte_alignment", + "require_block_scale_layout", + "require_cluster_shape", + "require_compact_major", + "require_gemm_inputs", + "require_mma_tiler", + "require_tensor_shape", +] diff --git a/python/cudnn/gemm/srelu.py b/python/cudnn/gemm/srelu.py new file mode 100644 index 000000000..5bdd04858 --- /dev/null +++ b/python/cudnn/gemm/srelu.py @@ -0,0 +1,217 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral validation shared by dense sReLU fusions.""" + +from __future__ import annotations + +from typing import Any + +from .. import data_type +from ..common.op import Op +from ..common.tensor_desc import TensorDesc +from .helpers import ( + block_scale_shape, + require_16_byte_alignment, + require_block_scale_layout, + require_cluster_shape, + require_compact_major, + require_gemm_inputs, + require_mma_tiler, + require_tensor_shape, +) + +_AB_DTYPES = frozenset( + { + data_type.FP4_E2M1, + data_type.UINT8, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + } +) +_FP4_DTYPES = frozenset({data_type.FP4_E2M1, data_type.UINT8}) +_FP8_DTYPES = frozenset({data_type.FP8_E4M3, data_type.FP8_E5M2}) +_WIDE_OUTPUT_DTYPES = frozenset({data_type.HALF, data_type.BFLOAT16, data_type.FLOAT}) +_C_DTYPES = _WIDE_OUTPUT_DTYPES | _FP8_DTYPES + + +class BlockScaledGemmSreluSm100OpBase(Op): + """Logical tensor signature shared by sReLU forward and backward.""" + + def __init__( + self, + *, + a: TensorDesc[Any], + b: TensorDesc[Any], + c: TensorDesc[Any], + d: TensorDesc[Any], + sfa: TensorDesc[Any], + sfb: TensorDesc[Any], + prob: TensorDesc[Any], + dprob: TensorDesc[Any] | None = None, + sfd: TensorDesc[Any] | None = None, + amax: TensorDesc[Any] | None = None, + norm_const: TensorDesc[Any] | None = None, + alpha: float = 1.0, + acc_dtype: data_type = data_type.FLOAT, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 16, + ) -> None: + for name, desc in ( + ("a", a), + ("b", b), + ("c", c), + ("d", d), + ("sfa", sfa), + ("sfb", sfb), + ("prob", prob), + ): + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + for name, desc in ( + ("dprob", dprob), + ("sfd", sfd), + ("amax", amax), + ("norm_const", norm_const), + ): + if desc is not None and not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc or None, got {type(desc).__name__}") + if not isinstance(acc_dtype, data_type): + raise TypeError(f"acc_dtype must be a cudnn.data_type, got {type(acc_dtype).__name__}") + if isinstance(sf_vec_size, bool) or not isinstance(sf_vec_size, int): + raise TypeError(f"sf_vec_size must be an int, got {type(sf_vec_size).__name__}") + + self.a = a + self.b = b + self.c = c + self.d = d + self.sfa = sfa + self.sfb = sfb + self.prob = prob + self.dprob = dprob + self.sfd = sfd + self.amax = amax + self.norm_const = norm_const + self.alpha = alpha + self.acc_dtype = acc_dtype + self.mma_tiler_mn = tuple(mma_tiler_mn) + self.cluster_shape_mn = None if cluster_shape_mn is None else tuple(cluster_shape_mn) + self.sf_vec_size = sf_vec_size + + self.m: int | None = None + self.n: int | None = None + self.k: int | None = None + self.l: int | None = None + self.a_major: str | None = None + self.b_major: str | None = None + self.output_major: str | None = None + + def check_support(self) -> bool: + """Validate shapes, dtypes, compact layouts, and static tiling.""" + + m, n, k, batch = require_gemm_inputs(self.a, self.b) + require_tensor_shape(self.c, (m, n, batch), label="C") + require_tensor_shape(self.d, (m, n, batch), label="D") + require_tensor_shape(self.prob, (m, 1, batch), label="prob") + if self.dprob is not None: + require_tensor_shape(self.dprob, (m, 1, batch), label="dprob") + + if self.sf_vec_size not in (16, 32): + raise ValueError(f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}") + require_tensor_shape( + self.sfa, + block_scale_shape(m, k, batch, self.sf_vec_size), + label="SFA", + ) + require_tensor_shape( + self.sfb, + block_scale_shape(n, k, batch, self.sf_vec_size), + label="SFB", + ) + for label, desc in (("SFA", self.sfa), ("SFB", self.sfb)): + require_block_scale_layout(desc, label) + + if self.sfd is not None: + raise NotImplementedError("SFD generation is not implemented by the current sReLU kernels") + if self.amax is not None: + require_tensor_shape(self.amax, (1,), label="amax") + if self.norm_const is not None: + raise ValueError("norm_const is only used with FP8 D output, which is not implemented") + + a_storage_dtype = self.a.cudnn_dtype + b_storage_dtype = self.b.cudnn_dtype + if b_storage_dtype != a_storage_dtype: + raise ValueError(f"A and B must have the same dtype, got {self.a.dtype} and {self.b.dtype}") + ab_dtype = data_type.FP4_E2M1 if a_storage_dtype == data_type.UINT8 else a_storage_dtype + if ab_dtype not in _AB_DTYPES: + raise ValueError(f"A has unsupported block-scaled GEMM dtype {self.a.dtype}") + + sf_dtype = self.sfa.cudnn_dtype + if self.sfb.cudnn_dtype != sf_dtype: + raise ValueError(f"SFA and SFB must have the same dtype, got {self.sfa.dtype} and {self.sfb.dtype}") + if sf_dtype not in {data_type.FP8_E8M0, data_type.FP8_E4M3}: + raise ValueError(f"SFA has unsupported scale-factor dtype {self.sfa.dtype}") + if ab_dtype in _FP8_DTYPES and (sf_dtype, self.sf_vec_size) != (data_type.FP8_E8M0, 32): + raise ValueError("FP8 inputs require FP8_E8M0 scales with sf_vec_size=32") + if ab_dtype in _FP4_DTYPES and self.sf_vec_size == 32 and sf_dtype != data_type.FP8_E8M0: + raise ValueError("FP4 inputs with sf_vec_size=32 require FP8_E8M0 scales") + + if self.c.cudnn_dtype not in _C_DTYPES: + raise ValueError(f"C has unsupported dtype {self.c.dtype}") + if self.d.cudnn_dtype in _FP8_DTYPES: + raise NotImplementedError("FP8 D output is unavailable because SFD generation is not implemented") + if self.d.cudnn_dtype not in _WIDE_OUTPUT_DTYPES: + raise ValueError(f"D has unsupported dtype {self.d.dtype}") + if self.prob.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"prob must have float32 dtype, got {self.prob.dtype}") + if self.dprob is not None and self.dprob.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"dprob must have float32 dtype, got {self.dprob.dtype}") + if self.acc_dtype != data_type.FLOAT: + raise ValueError(f"Accumulator dtype must be float32, got {self.acc_dtype}") + if self.amax is not None and self.amax.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"amax must have float32 dtype, got {self.amax.dtype}") + + a_major = require_compact_major(self.a, "m", "k") + b_major = require_compact_major(self.b, "n", "k") + c_major = require_compact_major(self.c, "m", "n") + d_major = require_compact_major(self.d, "m", "n") + if c_major != d_major: + raise ValueError(f"C and D must use the same major mode, got {c_major}-major and {d_major}-major") + if ab_dtype in _FP4_DTYPES and (a_major != "k" or b_major != "k"): + raise ValueError("FP4 A and B must use K-major layouts") + + self.mma_tiler_mn = require_mma_tiler( + self.mma_tiler_mn, + allowed_n=(64, 128, 192, 256), + ) + cta_group_size = 2 if self.mma_tiler_mn[0] == 256 else 1 + cta_tile_m = self.mma_tiler_mn[0] // cta_group_size + if m % cta_tile_m: + raise ValueError(f"M must be divisible by CTA_TILE_M={cta_tile_m} because the probability load is not predicated, got {m}") + if self.cluster_shape_mn is None: + self.cluster_shape_mn = (2, 1) if cta_group_size == 2 else (1, 1) + self.cluster_shape_mn = require_cluster_shape( + self.cluster_shape_mn, + cta_group_size=cta_group_size, + ) + if self.cluster_shape_mn[0] > 4 or self.cluster_shape_mn[1] > 4: + raise ValueError("cluster_shape_mn entries must be at most 4 for scale-factor multicast, " f"got {self.cluster_shape_mn}") + + for tensor in (self.c, self.d): + require_16_byte_alignment(tensor) + for tensor in (self.a, self.b): + if tensor.cudnn_dtype == data_type.UINT8: + contiguous_extent = tensor.shape[tensor.stride_order[0]] + if contiguous_extent * 4 % 128 != 0: + raise ValueError( + f"{tensor.name or 'Tensor'} contiguous extent must be a multiple of 32 logical FP4 elements " + f"for 16-byte alignment, got {contiguous_extent}" + ) + else: + require_16_byte_alignment(tensor) + + self.m, self.n, self.k, self.l = m, n, k, batch + self.ab_dtype = ab_dtype + self.a_major, self.b_major, self.output_major = a_major, b_major, c_major + return True diff --git a/python/cudnn/gemm_amax/__init__.py b/python/cudnn/gemm_amax/__init__.py index b65cf491e..70648275c 100644 --- a/python/cudnn/gemm_amax/__init__.py +++ b/python/cudnn/gemm_amax/__init__.py @@ -1,9 +1,18 @@ -from .api import ( - GemmAmaxSm100, - gemm_amax_wrapper_sm100, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy Torch API and framework-neutral GEMM + amax operation exports.""" -__all__ = [ - "GemmAmaxSm100", - "gemm_amax_wrapper_sm100", -] +from ..common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ("GemmAmaxSm100Op",), + "api": ( + "GemmAmaxSm100", + "gemm_amax_wrapper_sm100", + ), + }, + submodules=("api", "op"), +) diff --git a/python/cudnn/gemm_amax/api.py b/python/cudnn/gemm_amax/api.py index 62fc475bb..a43eaabd9 100644 --- a/python/cudnn/gemm_amax/api.py +++ b/python/cudnn/gemm_amax/api.py @@ -1,18 +1,26 @@ -from .dense_blockscaled_gemm_persistent_amax import ( - Sm100BlockScaledPersistentDenseGemmKernel, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT -from cuda.bindings import driver as cuda +"""Torch API for block-scaled dense GEMM + amax on SM100.""" + +import logging import os -import torch -from typing import Tuple, Optional +from typing import Optional, Tuple +from cuda.bindings import driver as cuda import cutlass import cutlass.cute as cute from cutlass.cute.runtime import make_fake_stream +import torch -from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.api_base import APIBase, TensorDesc, TupleDict, is_power_of_2, ceil_div +from cudnn import data_type +from cudnn.api_base import APIBase, TupleDict +from cudnn.datatypes import _torch_to_cudnn_data_type + +from .dense_blockscaled_gemm_persistent_amax import ( + Sm100BlockScaledPersistentDenseGemmKernel, +) +from .op import GemmAmaxSm100Op class GemmAmaxSm100(APIBase): @@ -39,216 +47,75 @@ def __init__( self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa") self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb") self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") - self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax") + self.amax_desc = self._pad_tensor_to_ndim( + self._make_tensor_desc(sample_amax, name="sample_amax"), + 3, + "sample_amax", + ) self.acc_dtype = acc_dtype self.mma_tiler_mn = mma_tiler_mn self.cluster_shape_mn = cluster_shape_mn self.sf_vec_size = sf_vec_size - self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) - self._logger.debug(f"setting num_cluster_overlap_margin: {self.num_cluster_overlap_margin}") - - # used to reshape sfa/sfb tensors to atom layout - self.atom_m = (32, 4) - self.atom_k = 4 - - self._interpret_uint8_as_fp4x2 = True self._logger.debug( - f"__init__ completed with args: sample_a {self.a_desc.shape}, sample_b {self.b_desc.shape}, sample_sfa {self.sfa_desc.shape}, sample_sfb {self.sfb_desc.shape}, sample_c {self.c_desc.shape}, sample_amax {self.amax_desc.shape}, acc_dtype {acc_dtype}, mma_tiler_mn {mma_tiler_mn}, cluster_shape_mn {cluster_shape_mn}, sf_vec_size {sf_vec_size}" + "setting num_cluster_overlap_margin: %s", + self.num_cluster_overlap_margin, ) - def check_support(self) -> bool: - self._logger.debug("Entering check_support") - - self._logger.debug("Checking dtypes and sf_vec_size") - ab_dtype = self._check_dtype( - self.a_desc, - dtype=[torch.float4_e2m1fn_x2, torch.uint8, torch.float8_e5m2, torch.float8_e4m3fn], - name="A", - ) - self._check_dtype( - self.b_desc, - dtype=ab_dtype, - name="B", - extra_error_msg="A and B tensor dtypes must match", + # Torch uint8 is the legacy packed-FP4 storage alias for this API. + self._interpret_uint8_as_fp4x2 = True + canonical_acc_dtype = _torch_to_cudnn_data_type(self.acc_dtype) + if canonical_acc_dtype is None: + canonical_acc_dtype = data_type.NOT_SET + self._op = GemmAmaxSm100Op( + a=self.a_desc, + b=self.b_desc, + sfa=self.sfa_desc, + sfb=self.sfb_desc, + c=self.c_desc, + amax=self.amax_desc, + acc_dtype=canonical_acc_dtype, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=self.cluster_shape_mn, + sf_vec_size=self.sf_vec_size, ) - if ab_dtype == torch.uint8: - self._logger.warning("Uint8 ab_dtype will be interpreted as packed fp4, not as native uint8") + self._kernel = Sm100BlockScaledPersistentDenseGemmKernel - self._value_error_if( - self.sf_vec_size not in {16, 32}, - f"Unsupported sf_vec_size: received {self.sf_vec_size}, expected {{16, 32}}", + self._logger.debug( + "__init__ completed with args: sample_a %s, sample_b %s, " + "sample_sfa %s, sample_sfb %s, sample_c %s, sample_amax %s, " + "acc_dtype %s, mma_tiler_mn %s, cluster_shape_mn %s, sf_vec_size %s", + self.a_desc.shape, + self.b_desc.shape, + self.sfa_desc.shape, + self.sfb_desc.shape, + self.c_desc.shape, + self.amax_desc.shape, + acc_dtype, + mma_tiler_mn, + cluster_shape_mn, + sf_vec_size, ) - sf_dtype = self._check_dtype( - self.sfa_desc, - dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn, torch.int8], - name="sfa", - ) - self._check_dtype( - self.sfb_desc, - dtype=sf_dtype, - name="sfb", - extra_error_msg="sfa and sfb tensor dtypes must match", - ) - if sf_dtype == torch.int8: + def check_support(self) -> bool: + """Validate the shared operation contract and Torch device support.""" + + self._is_supported = False + self._op.check_support() + self.mma_tiler_mn = self._op.mma_tiler_mn + self.cluster_shape_mn = self._op.cluster_shape_mn + self.a_major = self._op.a_major + self.b_major = self._op.b_major + self.c_major = self._op.c_major + self.ab_dtype = self.a_desc.dtype + self.c_dtype = self.c_desc.dtype + + if self.a_desc.dtype == torch.uint8: + self._logger.warning("Uint8 ab_dtype will be interpreted as packed fp4, not as native uint8") + if self.sfa_desc.dtype == torch.int8: self._logger.warning("Int8 sf_dtype will be interpreted as float8_e8m0fnu, not as native int8") - self._value_error_if( - sf_dtype == torch.float8_e4m3fn and self.sf_vec_size == 32, - "Unsupported sf_dtype and sf_vec_size combination: float8_e4m3fn and 32 is not supported", - ) - self._value_error_if( - ab_dtype in {torch.float8_e5m2, torch.float8_e4m3fn} and self.sf_vec_size == 16, - f"Unsupported ab_dtype and sf_vec_size combination: {{float8_e5m2, float8_e4m3fn}} and 16 is not supported", - ) - - c_dtype = self._check_dtype( - self.c_desc, - dtype=[torch.float32, torch.float16, torch.bfloat16, torch.float8_e5m2, torch.float8_e4m3fn, torch.float4_e2m1fn_x2, torch.uint8], - name="C", - ) - self._value_error_if( - self._is_fp4x2(c_dtype) and not self._is_fp4x2(ab_dtype), - f"Unsupported c_dtype and ab_dtype combination: fp4 c_dtype requires fp4 ab_dtype, got {ab_dtype}", - ) - self._not_implemented_error_if( - self._is_fp8(c_dtype) and self._is_fp8(ab_dtype), - "Unsupported c_dtype and ab_dtype combination: fp8 ab_dtype and fp8 c_dtype (fails to launch)", - ) - self._check_dtype( - self.acc_dtype, - dtype=torch.float32, - name="Accumulator", - extra_error_msg="Accumulator must be float32", - ) - - self.ab_dtype = ab_dtype - self.c_dtype = c_dtype - - self._logger.debug("Checking tensor layout") - m, k, l = self.a_desc.shape - n, _, _ = self.b_desc.shape - _, _, m_div_atom_m0_m1, _, sf_k_div_atom_k, _ = self.sfa_desc.shape - _, _, n_div_atom_m0_m1, _, sf_k_div_atom_k, _ = self.sfb_desc.shape - - self._check_tensor_shape(self.a_desc, (m, k, l), "A") - self._check_tensor_shape(self.b_desc, (n, k, l), "B") - self._check_tensor_shape(self.c_desc, (m, n, l), "C") - self._check_tensor_shape( - self.sfa_desc, - (self.atom_m[0], self.atom_m[1], m_div_atom_m0_m1, self.atom_k, sf_k_div_atom_k, l), - "sfa", - ) - self._check_tensor_shape( - self.sfb_desc, - (self.atom_m[0], self.atom_m[1], n_div_atom_m0_m1, self.atom_k, sf_k_div_atom_k, l), - "sfb", - ) - self.amax_desc = self._pad_tensor_to_ndim(self.amax_desc, 3, "amax") - self._check_tensor_shape(self.amax_desc, (1, 1, 1), "amax") - - expected_m_div_atom = ceil_div(m, self.atom_m[0] * self.atom_m[1]) - expected_n_div_atom = ceil_div(n, self.atom_m[0] * self.atom_m[1]) - self._value_error_if( - m_div_atom_m0_m1 != expected_m_div_atom, - f"Input/Output shape mismatch: expected m_div_atom_m0_m1 (sfa.shape[2]) = {expected_m_div_atom}, got {m_div_atom_m0_m1}", - ) - self._value_error_if( - n_div_atom_m0_m1 != expected_n_div_atom, - f"Input/Output shape mismatch: expected n_div_atom_m0_m1 (sfb.shape[2]) = {expected_n_div_atom}, got {n_div_atom_m0_m1}", - ) - - # Check tensor strides - _ = self._check_tensor_stride( - self.a_desc, - stride=[(1, m, m * k), (k, 1, m * k)], - name="A", - ) - _ = self._check_tensor_stride( - self.b_desc, - stride=[(1, n, n * k), (k, 1, n * k)], - name="B", - ) - _ = self._check_tensor_stride( - self.c_desc, - stride=[(1, m, m * n), (n, 1, m * n)], - name="C", - ) - - # Derive major mode from stride order - self.a_major = "m" if self.a_desc.stride_order == (0, 1, 2) else "k" - self.b_major = "n" if self.b_desc.stride_order == (0, 1, 2) else "k" - self.c_major = "m" if self.c_desc.stride_order == (0, 1, 2) else "n" - - self._value_error_if( - self._is_fp4x2(ab_dtype) and not (self.a_major == "k" and self.b_major == "k"), - f"Unsupported A or B tensor stride: Float4 tensors require k-major layout for hardware efficiency, got {self.a_major} and {self.b_major}", - ) - self._value_error_if( - self._is_fp4x2(c_dtype) and self.c_major == "m", - f"Unsupported C tensor stride: Float4 tensors require n-major layout for hardware efficiency, got {self.c_major}", - ) - - self._logger.debug("Checking mma tiler and cluster shape") - self._value_error_if( - self.mma_tiler_mn[0] not in [128, 256], - f"Unsupported mma_tiler_mn[0]: expected {{128, 256}}, got {self.mma_tiler_mn[0]}", - ) - self._value_error_if( - self.mma_tiler_mn[1] not in [128, 256], - f"Unsupported mma_tiler_mn[1]: expected {{128, 256}}, got {self.mma_tiler_mn[1]}", - ) - self._not_implemented_error_if( - self.mma_tiler_mn[0] == 256, - "mma_tiler_mn[0] == 256 currently hangs", - ) - self._value_error_if( - self._is_fp4x2(self.ab_dtype) and self.mma_tiler_mn[1] == 256 and k <= 128, - f"mma_tiler_mn (X, 256) requires k > 128 (packed x2), got {k}", - ) - self._value_error_if( - not (self.cluster_shape_mn[0] % (2 if self.mma_tiler_mn[0] == 256 else 1) == 0), - "Illegal cluster shape", - ) - self._not_implemented_error_if( - self.mma_tiler_mn == (128, 256) and self.sf_vec_size == 16 and c_dtype in {torch.float32, torch.float16, torch.bfloat16}, - "mma_tiler_mn (128, 256), sf_vec_size 16, c_dtype {torch.float32, torch.float16, torch.bfloat16} fails to launch", - ) - - # Special cluster shape check for scale factor multicasts. - # Due to limited size of scale factors, we can't multicast among more than 4 CTAs. - self._value_error_if( - not ( - self.cluster_shape_mn[0] <= 4 - and self.cluster_shape_mn[1] <= 4 - and self.cluster_shape_mn[0] > 0 - and self.cluster_shape_mn[1] > 0 - and is_power_of_2(self.cluster_shape_mn[0]) - and is_power_of_2(self.cluster_shape_mn[1]) - ), - f"Invalid cluster shape: expected cluster_shape_mn values in {{1, 2, 4}}, got {self.cluster_shape_mn}", - ) - - self._logger.debug("Checking tensor alignment") - - def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): - major_mode_idx = 0 if is_mode0_major else 1 - num_major_elements = tensor_shape[major_mode_idx] - num_contiguous_elements = 16 * 8 // (_convert_to_cutlass_data_type(dtype).width) - return num_major_elements % num_contiguous_elements == 0 - - self._value_error_if( - not ( - check_contigous_16B_alignment(ab_dtype, self.a_major == "m", (m, k, l)) - and check_contigous_16B_alignment(ab_dtype, self.b_major == "n", (n, k, l)) - and check_contigous_16B_alignment(c_dtype, self.c_major == "m", (m, n, l)) - ), - "Unsupported tensor alignment: tensors must be 16B aligned", - ) - - self._logger.debug("Checking environment") self._runtime_error_if(not torch.cuda.is_available(), "CUDA is not available") device = torch.cuda.current_device() major, minor = torch.cuda.get_device_capability(device) @@ -258,8 +125,6 @@ def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): f"GemmAmax requires SM100+ compute capability, but found SM{compute_capability} on device {device}", ) - self._kernel = Sm100BlockScaledPersistentDenseGemmKernel - self._is_supported = True self._logger.debug("check_support completed successfully") return True @@ -281,7 +146,7 @@ def compile(self) -> None: max_active_clusters -= self.num_cluster_overlap_margin self._value_error_if( max_active_clusters <= 0, - "max_active_clusters must be > 0 after applying overlap margin; reduce CUDNNFE_CLUSTER_OVERLAP_MARGIN", + "max_active_clusters must be > 0 after applying overlap margin; " "reduce CUDNNFE_CLUSTER_OVERLAP_MARGIN", ) self._logger.debug("Compiling gemm_amax") @@ -293,7 +158,7 @@ def compile(self) -> None: amax_cute = self._make_fake_cute_tensor_from_desc(self.amax_desc, assumed_align=16) fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) - _compiled_kernel = cute.compile( + compiled_kernel = cute.compile( gemm_amax, a_tensor=a_cute, b_tensor=b_cute, @@ -316,8 +181,7 @@ def tensor_api( stream: cuda.CUstream, ): amax_tensor = self._pad_tensor_to_ndim(amax_tensor, 3, "amax") - - return _compiled_kernel( + return compiled_kernel( a_tensor, b_tensor, sfa_tensor, @@ -328,11 +192,6 @@ def tensor_api( ) self._compiled_kernel = tensor_api - - for attr_name in tuple(vars(self)): - if attr_name.startswith("sample_"): - setattr(self, attr_name, None) - self._logger.debug("Kernel compiled successfully") def execute( @@ -353,7 +212,6 @@ def execute( "GemmAmaxSm100 kernel not compiled; call compile() first", ) self._logger.debug("Executing with compiled kernel") - self._compiled_kernel( a_tensor=a_tensor, b_tensor=b_tensor, @@ -366,8 +224,6 @@ def execute( self._logger.debug("Executed with compiled kernel successfully") -import logging - _logger = logging.getLogger(__name__) _cache_of_GemmAmaxSm100Objects = {} @@ -385,19 +241,32 @@ def gemm_amax_wrapper_sm100( sf_vec_size: int = 32, stream: Optional[cuda.CUstream] = None, ) -> TupleDict: - _logger.debug("gemm_amax_wrapper_sm100: Creating empty output tensors c and amax") m, _, l = a_tensor.shape n, _, l = b_tensor.shape - c_tensor = None if c_major == "m": - c_tensor = torch.empty_strided((m, n, l), (1, m, m * n), dtype=c_dtype, device=a_tensor.device) + c_tensor = torch.empty_strided( + (m, n, l), + (1, m, m * n), + dtype=c_dtype, + device=a_tensor.device, + ) elif c_major == "n": - c_tensor = torch.empty_strided((m, n, l), (n, 1, m * n), dtype=c_dtype, device=a_tensor.device) + c_tensor = torch.empty_strided( + (m, n, l), + (n, 1, m * n), + dtype=c_dtype, + device=a_tensor.device, + ) else: raise ValueError(f"c_major must be either 'm' or 'n', got {c_major}") - amax_tensor = torch.full((1, 1, 1), -float("inf"), device=a_tensor.device, dtype=torch.float32) + amax_tensor = torch.full( + (1, 1, 1), + -float("inf"), + device=a_tensor.device, + dtype=torch.float32, + ) cache_key = ( a_tensor.shape, @@ -432,7 +301,7 @@ def gemm_amax_wrapper_sm100( current_stream=stream, ) else: - _logger.debug("gemm_amax_wrapper_sm100: No previously cached GemmAmaxSm100 object found, creating new GemmAmaxSm100 object") + _logger.debug("gemm_amax_wrapper_sm100: No previously cached GemmAmaxSm100 object " "found, creating new GemmAmaxSm100 object") gemm_amax = GemmAmaxSm100( sample_a=a_tensor, sample_b=b_tensor, @@ -445,7 +314,7 @@ def gemm_amax_wrapper_sm100( cluster_shape_mn=cluster_shape_mn, sf_vec_size=sf_vec_size, ) - assert gemm_amax.check_support() + gemm_amax.check_support() gemm_amax.compile() gemm_amax.execute( a_tensor=a_tensor, diff --git a/python/cudnn/gemm_amax/jax.py b/python/cudnn/gemm_amax/jax.py new file mode 100644 index 000000000..89e52c5f0 --- /dev/null +++ b/python/cudnn/gemm_amax/jax.py @@ -0,0 +1,249 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX API for block-scaled dense GEMM + amax on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any + +import jax +import jax.numpy as jnp + +from .. import data_type +from .._jax.compiler import compile_options_for_target +from ..gemm.helpers import require_gemm_inputs +from .._jax import JaxApiBase, JaxTensorDesc, TupleDict +from .._jax.datatypes import jax_to_cudnn_dtype, normalize_jax_dtype +from .._jax.gemm import BLOCK_SCALE_MODE, gemm_a_mode, gemm_b_mode, gemm_output_mode +from .._jax.layout import to_public_axes +from .op import GemmAmaxSm100Op + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) + + +class GemmAmaxSm100(JaxApiBase): + """JAX callable specialized from block-scaled GEMM input metadata.""" + + def __init__( + self, + sample_a: Any, + sample_b: Any, + sample_sfa: Any, + sample_sfb: Any, + *, + sample_c: Any | None = None, + sample_amax: Any | None = None, + c_layout: str = "LMN", + c_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: tuple[int, int] = (1, 1), + sf_vec_size: int = 32, + a_layout: str = "LMK", + b_layout: str = "LNK", + ) -> None: + a_mode = gemm_a_mode(a_layout) + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(c_layout) + + self.compute_capability = self._resolve_compute_capability( + None, + SUPPORTED_COMPUTE_CAPABILITIES, + "GemmAmaxSm100", + ) + self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=a_mode) + self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=b_mode) + self.sfa_desc = self._to_tensor_desc( + sample_sfa, + "sample_sfa", + mode=BLOCK_SCALE_MODE, + ) + self.sfb_desc = self._to_tensor_desc( + sample_sfb, + "sample_sfb", + mode=BLOCK_SCALE_MODE, + ) + self.acc_dtype = normalize_jax_dtype(acc_dtype, jnp.float32, "acc_dtype") + + if (sample_c is None) != (sample_amax is None): + raise ValueError("sample_c and sample_amax must be provided together") + if sample_c is None: + resolved_c_dtype = normalize_jax_dtype(c_dtype, jnp.float32, "c_dtype") + self.c_desc, self.amax_desc = self._default_output_descs(resolved_c_dtype, output_mode) + else: + if c_dtype is not None: + raise ValueError("c_dtype cannot be specified with sample_c and sample_amax") + self.c_desc = self._to_tensor_desc( + sample_c, + "sample_c", + mode=output_mode, + ) + self.amax_desc = self._to_tensor_desc( + sample_amax, + "sample_amax", + init_value=float("-inf"), + ) + + acc_cudnn_dtype = jax_to_cudnn_dtype(self.acc_dtype) + if acc_cudnn_dtype == data_type.NOT_SET: + raise ValueError(f"Unsupported JAX accumulator dtype {self.acc_dtype}") + self._op = GemmAmaxSm100Op( + a=self.a_desc, + b=self.b_desc, + sfa=self.sfa_desc, + sfb=self.sfb_desc, + c=self.c_desc, + amax=self.amax_desc, + acc_dtype=acc_cudnn_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + ) + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + + def _default_output_descs( + self, + c_dtype: Any, + output_mode: tuple[int, ...], + ) -> tuple[JaxTensorDesc, JaxTensorDesc]: + m, n, _, batch = require_gemm_inputs(self.a_desc, self.b_desc) + return ( + JaxTensorDesc.from_shape( + to_public_axes((m, n, batch), output_mode), + c_dtype, + name="sample_c", + mode=output_mode, + ), + JaxTensorDesc.from_shape( + (1, 1, 1), + jnp.float32, + name="sample_amax", + init_value=float("-inf"), + ), + ) + + def check_support(self) -> bool: + """Validate the JAX dtype surface and common operation contract.""" + + if self.a_desc.cudnn_dtype not in { + data_type.FP4_E2M1, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + }: + raise ValueError(f"The JAX GEMM + amax API requires FP4 or FP8 A and B, got " f"{self.a_desc.dtype}") + if self.b_desc.cudnn_dtype not in { + data_type.FP4_E2M1, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + }: + raise ValueError(f"The JAX GEMM + amax API requires FP4 or FP8 A and B, got " f"{self.b_desc.dtype}") + if self.sfa_desc.cudnn_dtype not in { + data_type.FP8_E8M0, + data_type.FP8_E4M3, + }: + raise ValueError(f"The JAX GEMM + amax API requires E8M0 or E4M3 scale factors, " f"got {self.sfa_desc.dtype}") + if self.sfb_desc.cudnn_dtype not in { + data_type.FP8_E8M0, + data_type.FP8_E4M3, + }: + raise ValueError(f"The JAX GEMM + amax API requires E8M0 or E4M3 scale factors, " f"got {self.sfb_desc.dtype}") + return self._op.check_support() + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + ) -> TupleDict: + self.check_support() + max_active_clusters = self._get_max_active_clusters( + self._op.cluster_shape_mn[0] * self._op.cluster_shape_mn[1], + overlap_margin=self.num_cluster_overlap_margin, + ) + + def launch(stream, a, b, sfa, sfb, c, amax): + from .dense_blockscaled_gemm_persistent_amax import ( + Sm100BlockScaledPersistentDenseGemmKernel, + ) + + kernel = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size=self._op.sf_vec_size, + mma_tiler_mn=self._op.mma_tiler_mn, + cluster_shape_mn=self._op.cluster_shape_mn, + ) + kernel( + a, + b, + sfa, + sfb, + c, + amax, + max_active_clusters, + stream, + ) + + c_tensor, amax_tensor = self._call_kernel( + (a_tensor, b_tensor, sfa_tensor, sfb_tensor), + launch=launch, + input_descs=(self.a_desc, self.b_desc, self.sfa_desc, self.sfb_desc), + output_descs=(self.c_desc, self.amax_desc), + compile_options=compile_options_for_target(self.compute_capability), + ) + return TupleDict(c_tensor=c_tensor, amax_tensor=amax_tensor) + + +@partial( + jax.jit, + static_argnames=( + "c_layout", + "c_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "a_layout", + "b_layout", + ), +) +def gemm_amax_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + c_layout: str = "LMN", + c_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: tuple[int, int] = (1, 1), + sf_vec_size: int = 32, + *, + a_layout: str = "LMK", + b_layout: str = "LNK", +) -> TupleDict: + """Compute a block-scaled GEMM and global max-absolute reduction.""" + + return GemmAmaxSm100( + 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", + "SUPPORTED_COMPUTE_CAPABILITIES", + "gemm_amax_wrapper_sm100", +] diff --git a/python/cudnn/gemm_amax/op.py b/python/cudnn/gemm_amax/op.py new file mode 100644 index 000000000..4ef637851 --- /dev/null +++ b/python/cudnn/gemm_amax/op.py @@ -0,0 +1,236 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral block-scaled dense GEMM + amax operation.""" + +from __future__ import annotations + +from typing import Any, Iterable, Optional + +from .. import data_type +from ..gemm.helpers import ( + block_scale_shape, + data_type_bits, + require_block_scale_layout, + require_cluster_shape, + require_compact_major, + require_gemm_inputs, + require_mma_tiler, + require_tensor_shape, +) +from ..common.op import Op +from ..common.tensor_desc import TensorDesc + +_AB_DTYPES = frozenset( + { + data_type.FP4_E2M1, + data_type.UINT8, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + } +) +_SCALE_DTYPES = frozenset( + { + data_type.FP8_E4M3, + data_type.FP8_E8M0, + data_type.INT8, + } +) +_C_DTYPES = frozenset( + { + data_type.FLOAT, + data_type.HALF, + data_type.BFLOAT16, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + data_type.FP4_E2M1, + data_type.UINT8, + } +) +_WIDE_C_DTYPES = frozenset( + { + data_type.FLOAT, + data_type.HALF, + data_type.BFLOAT16, + } +) +_FP8_DTYPES = frozenset({data_type.FP8_E4M3, data_type.FP8_E5M2}) +_FP4_STORAGE_DTYPES = frozenset({data_type.FP4_E2M1, data_type.UINT8}) +_E8M0_STORAGE_DTYPES = frozenset({data_type.FP8_E8M0, data_type.INT8}) + + +def _require_supported_dtype( + label: str, + dtype: data_type, + supported: Iterable[data_type], +) -> None: + if dtype not in supported: + raise ValueError(f"{label} has unsupported dtype {dtype}") + + +def _require_logical_alignment( + tensor: TensorDesc[Any], + logical_dtype: data_type, +) -> None: + bits = 4 if logical_dtype == data_type.FP4_E2M1 else data_type_bits(logical_dtype) + extent = tensor.shape[tensor.stride_order[0]] + if extent * bits % 128: + required_multiple = 128 // bits + raise ValueError( + f"{tensor.name or 'Tensor'} contiguous extent must be a multiple of " f"{required_multiple} elements for 16-byte alignment, got {extent}" + ) + + +class GemmAmaxSm100Op(Op): + """Complete logical signature and static configuration for GEMM + amax.""" + + def __init__( + self, + *, + a: TensorDesc[Any], + b: TensorDesc[Any], + sfa: TensorDesc[Any], + sfb: TensorDesc[Any], + c: TensorDesc[Any], + amax: TensorDesc[Any], + acc_dtype: data_type = data_type.FLOAT, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: tuple[int, int] = (1, 1), + sf_vec_size: int = 32, + ) -> None: + for name, desc in ( + ("a", a), + ("b", b), + ("sfa", sfa), + ("sfb", sfb), + ("c", c), + ("amax", amax), + ): + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + if not isinstance(acc_dtype, data_type): + raise TypeError(f"acc_dtype must be a cudnn.data_type, got {type(acc_dtype).__name__}") + + self.a = a + self.b = b + self.sfa = sfa + self.sfb = sfb + self.c = c + self.amax = amax + self.acc_dtype = acc_dtype + self.requested_mma_tiler_mn = mma_tiler_mn + self.requested_cluster_shape_mn = cluster_shape_mn + self.sf_vec_size = sf_vec_size + + self.m: Optional[int] = None + self.n: Optional[int] = None + self.k: Optional[int] = None + self.l: Optional[int] = None + self.a_major: Optional[str] = None + self.b_major: Optional[str] = None + self.c_major: Optional[str] = None + self.ab_dtype: Optional[data_type] = None + self.scale_dtype: Optional[data_type] = None + self.c_dtype: Optional[data_type] = None + self.mma_tiler_mn: Optional[tuple[int, int]] = None + self.cluster_shape_mn: Optional[tuple[int, int]] = None + + def check_support(self) -> bool: + """Validate the complete tensor signature and static configuration.""" + + self.m = self.n = self.k = self.l = None + self.a_major = self.b_major = self.c_major = None + self.ab_dtype = self.scale_dtype = self.c_dtype = None + self.mma_tiler_mn = self.cluster_shape_mn = None + + m, n, k, batch = require_gemm_inputs(self.a, self.b) + require_tensor_shape(self.c, (m, n, batch), label="C") + require_tensor_shape(self.amax, (1, 1, 1), label="Amax") + + if self.a.cudnn_dtype != self.b.cudnn_dtype: + raise ValueError(f"A and B must have the same dtype, got {self.a.dtype} and {self.b.dtype}") + ab_storage_dtype = self.a.cudnn_dtype + _require_supported_dtype("A and B", ab_storage_dtype, _AB_DTYPES) + ab_dtype = data_type.FP4_E2M1 if ab_storage_dtype in _FP4_STORAGE_DTYPES else ab_storage_dtype + + if self.sfa.cudnn_dtype != self.sfb.cudnn_dtype: + raise ValueError(f"SFA and SFB must have the same dtype, got {self.sfa.dtype} and {self.sfb.dtype}") + scale_storage_dtype = self.sfa.cudnn_dtype + _require_supported_dtype("SFA and SFB", scale_storage_dtype, _SCALE_DTYPES) + scale_dtype = data_type.FP8_E8M0 if scale_storage_dtype in _E8M0_STORAGE_DTYPES else scale_storage_dtype + + if self.sf_vec_size not in (16, 32): + raise ValueError(f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}") + if scale_dtype == data_type.FP8_E4M3 and self.sf_vec_size == 32: + raise ValueError("FP8 E4M3 scale factors do not support sf_vec_size=32") + if ab_dtype in _FP8_DTYPES and self.sf_vec_size == 16: + raise ValueError("FP8 A and B do not support sf_vec_size=16") + require_tensor_shape( + self.sfa, + block_scale_shape(m, k, batch, self.sf_vec_size), + label="SFA", + ) + require_tensor_shape( + self.sfb, + block_scale_shape(n, k, batch, self.sf_vec_size), + label="SFB", + ) + require_block_scale_layout(self.sfa, "SFA") + require_block_scale_layout(self.sfb, "SFB") + + c_storage_dtype = self.c.cudnn_dtype + _require_supported_dtype("C", c_storage_dtype, _C_DTYPES) + c_dtype = data_type.FP4_E2M1 if c_storage_dtype in _FP4_STORAGE_DTYPES else c_storage_dtype + if self.amax.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"Amax must have dtype float32, got {self.amax.dtype}") + if self.acc_dtype != data_type.FLOAT: + raise ValueError(f"Accumulator dtype must be float32, got {self.acc_dtype}") + + a_major = require_compact_major(self.a, "m", "k") + b_major = require_compact_major(self.b, "n", "k") + c_major = require_compact_major(self.c, "m", "n") + + if ab_dtype == data_type.FP4_E2M1 and (a_major, b_major) != ("k", "k"): + raise ValueError("FP4 A and B require k-major layouts, got " f"{a_major}-major and {b_major}-major") + if c_dtype == data_type.FP4_E2M1 and c_major != "n": + raise ValueError(f"FP4 C requires n-major layout, got {c_major}-major") + if c_dtype == data_type.FP4_E2M1 and ab_dtype != data_type.FP4_E2M1: + raise ValueError("FP4 C requires FP4 A and B") + if ab_dtype in _FP8_DTYPES and c_dtype in _FP8_DTYPES: + raise NotImplementedError("FP8 A and B with FP8 C are unsupported because the kernel fails to launch") + + mma_tiler_mn = require_mma_tiler( + self.requested_mma_tiler_mn, + allowed_m=(128, 256), + allowed_n=(128, 256), + ) + if mma_tiler_mn[0] == 256: + raise NotImplementedError("mma_tiler_mn[0] == 256 currently hangs") + cluster_shape_mn = require_cluster_shape( + self.requested_cluster_shape_mn, + cta_group_size=2 if mma_tiler_mn[0] == 256 else 1, + max_cluster_size=16, + ) + if cluster_shape_mn[0] > 4 or cluster_shape_mn[1] > 4: + raise ValueError("cluster_shape_mn entries must be at most 4 for scale-factor multicast, " f"got {cluster_shape_mn}") + + if ab_dtype == data_type.FP4_E2M1 and mma_tiler_mn[1] == 256 and k <= 128: + raise ValueError(f"mma_tiler_mn (X, 256) requires K > 128 for FP4, got {k}") + if mma_tiler_mn == (128, 256) and self.sf_vec_size == 16 and c_dtype in _WIDE_C_DTYPES: + raise NotImplementedError("mma_tiler_mn (128, 256), sf_vec_size=16, and a 16/32-bit C " "dtype are unsupported because the kernel fails to launch") + + _require_logical_alignment(self.a, ab_dtype) + _require_logical_alignment(self.b, ab_dtype) + _require_logical_alignment(self.c, c_dtype) + + self.m, self.n, self.k, self.l = m, n, k, batch + self.a_major, self.b_major, self.c_major = a_major, b_major, c_major + self.ab_dtype = ab_dtype + self.scale_dtype = scale_dtype + self.c_dtype = c_dtype + self.mma_tiler_mn = mma_tiler_mn + self.cluster_shape_mn = cluster_shape_mn + return True + + +__all__ = ["GemmAmaxSm100Op", "block_scale_shape"] diff --git a/python/cudnn/gemm_dsrelu/__init__.py b/python/cudnn/gemm_dsrelu/__init__.py index ade25458b..b0be8c758 100644 --- a/python/cudnn/gemm_dsrelu/__init__.py +++ b/python/cudnn/gemm_dsrelu/__init__.py @@ -1,9 +1,18 @@ -from .api import ( - GemmDsreluSm100, - gemm_dsrelu_wrapper_sm100, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy Torch API and framework-neutral GEMM + dsReLU exports.""" -__all__ = [ - "GemmDsreluSm100", - "gemm_dsrelu_wrapper_sm100", -] +from ..common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ("GemmDsreluSm100Op",), + "api": ( + "GemmDsreluSm100", + "gemm_dsrelu_wrapper_sm100", + ), + }, + submodules=("api", "op"), +) diff --git a/python/cudnn/gemm_dsrelu/api.py b/python/cudnn/gemm_dsrelu/api.py index fb02fa027..d2e71490c 100644 --- a/python/cudnn/gemm_dsrelu/api.py +++ b/python/cudnn/gemm_dsrelu/api.py @@ -118,6 +118,10 @@ def check_support(self) -> bool: dtype=[torch.float16, torch.bfloat16, torch.float32, torch.float8_e4m3fn, torch.float8_e5m2], name="D", ) + self._not_implemented_error_if( + self._is_fp8(self.d_dtype), + "FP8 D output is unavailable because the current dsReLU kernel does not implement SFD generation", + ) self._check_dtype(self.prob_desc, dtype=torch.float32, name="prob") self._check_dtype(self.dprob_desc, dtype=torch.float32, name="dprob") @@ -128,9 +132,6 @@ def check_support(self) -> bool: self._check_dtype(self.acc_dtype, dtype=torch.float32, name="Accumulator") self._value_error_if(self.sf_vec_size not in {16, 32}, f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}") - self._value_error_if( - self._is_fp8(self.d_desc) and (self.sfd_desc is None or self.norm_const_desc is None), "sfd and norm_const are required when D is FP8" - ) self._value_error_if( self._is_fp4x2(self.ab_dtype) and self.d_dtype in {torch.float8_e4m3fn, torch.float8_e5m2}, "FP4 input with FP8 output is not supported" ) diff --git a/python/cudnn/gemm_dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.py b/python/cudnn/gemm_dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.py index 66f7d6ddb..9712c71bd 100644 --- a/python/cudnn/gemm_dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.py +++ b/python/cudnn/gemm_dsrelu/dense_blockscaled_gemm_persistent_dsrelu_quant.py @@ -29,8 +29,6 @@ from typing import Type, Tuple, Union, Optional import cuda.bindings.driver as cuda -import torch - import cutlass import cutlass.cute as cute from cutlass.cute.nvgpu import cpasync, tcgen05 diff --git a/python/cudnn/gemm_dsrelu/jax.py b/python/cudnn/gemm_dsrelu/jax.py new file mode 100644 index 000000000..f761fe8c7 --- /dev/null +++ b/python/cudnn/gemm_dsrelu/jax.py @@ -0,0 +1,285 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX API for block-scaled dense GEMM + squared-ReLU backward.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any + +import jax +import jax.numpy as jnp + +from .. import data_type +from .._jax.compiler import compile_options_for_target +from ..gemm.helpers import require_gemm_inputs +from .._jax import JaxApiBase, JaxTensorDesc, TupleDict +from .._jax.datatypes import jax_to_cudnn_dtype, normalize_jax_dtype +from .._jax.gemm import BLOCK_SCALE_MODE, PROBABILITY_MODE, gemm_a_mode, gemm_b_mode, gemm_output_mode +from .._jax.layout import to_public_axes +from .op import GemmDsreluSm100Op + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) +_JAX_INPUT_DTYPES = frozenset({data_type.FP4_E2M1, data_type.FP8_E4M3, data_type.FP8_E5M2}) +_WIDE_OUTPUT_DTYPES = frozenset({data_type.HALF, data_type.BFLOAT16, data_type.FLOAT}) + + +class GemmDsreluSm100(JaxApiBase): + """JAX callable specialized from a block-scaled squared-ReLU backward signature.""" + + def __init__( + self, + sample_a: Any, + sample_b: Any, + sample_c: Any, + sample_sfa: Any, + sample_sfb: Any, + sample_prob: Any, + *, + sample_d: Any | None = None, + sample_dprob: Any | None = None, + alpha: float = 1.0, + d_layout: str = "LMN", + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + a_layout: str = "LMK", + b_layout: str = "LNK", + ) -> None: + a_mode = gemm_a_mode(a_layout) + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(d_layout, name="d_layout") + + self.compute_capability = self._resolve_compute_capability( + None, + SUPPORTED_COMPUTE_CAPABILITIES, + "GemmDsreluSm100", + ) + self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=a_mode) + self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=b_mode) + self.c_desc = self._to_tensor_desc(sample_c, "sample_c", mode=output_mode) + self.sfa_desc = self._to_tensor_desc(sample_sfa, "sample_sfa", mode=BLOCK_SCALE_MODE) + self.sfb_desc = self._to_tensor_desc(sample_sfb, "sample_sfb", mode=BLOCK_SCALE_MODE) + self.prob_desc = self._to_tensor_desc(sample_prob, "sample_prob", mode=PROBABILITY_MODE) + + self.acc_dtype = normalize_jax_dtype(acc_dtype, jnp.float32, "acc_dtype") + if (sample_d is None) != (sample_dprob is None): + raise ValueError("sample_d and sample_dprob must be provided together") + if sample_d is None: + resolved_d_dtype = normalize_jax_dtype(d_dtype, jnp.bfloat16, "d_dtype") + self.d_desc, self.dprob_desc = self._default_output_descs(resolved_d_dtype, output_mode) + else: + if d_dtype is not None: + raise ValueError("d_dtype cannot be specified with sample_d and sample_dprob") + self.d_desc = self._to_tensor_desc(sample_d, "sample_d", mode=output_mode) + self.dprob_desc = self._to_tensor_desc( + sample_dprob, + "sample_dprob", + mode=PROBABILITY_MODE, + init_value=0.0, + ) + + self.amax_desc = None + if self.a_desc.cudnn_dtype == data_type.FP4_E2M1 and self.d_desc.cudnn_dtype in _WIDE_OUTPUT_DTYPES: + self.amax_desc = self.d_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(1,), + name="amax_tensor", + init_value=float("-inf"), + ) + acc_cudnn_dtype = jax_to_cudnn_dtype(self.acc_dtype) + if acc_cudnn_dtype == data_type.NOT_SET: + raise ValueError(f"Unsupported JAX accumulator dtype {self.acc_dtype}") + if not isinstance(vector_f32, bool): + raise TypeError(f"vector_f32 must be a bool, got {type(vector_f32).__name__}") + + self._op = GemmDsreluSm100Op( + a=self.a_desc, + b=self.b_desc, + c=self.c_desc, + d=self.d_desc, + sfa=self.sfa_desc, + sfb=self.sfb_desc, + prob=self.prob_desc, + dprob=self.dprob_desc, + sfd=None, + amax=self.amax_desc, + norm_const=None, + alpha=alpha, + acc_dtype=acc_cudnn_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + ) + self.vector_f32 = vector_f32 + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + + def _default_output_descs( + self, + d_dtype: Any, + output_mode: tuple[int, ...], + ) -> tuple[JaxTensorDesc, JaxTensorDesc]: + m, n, _, batch = require_gemm_inputs(self.a_desc, self.b_desc) + d_shape = to_public_axes((m, n, batch), output_mode) + dprob_shape = to_public_axes((m, 1, batch), PROBABILITY_MODE) + return ( + JaxTensorDesc.from_shape( + d_shape, + d_dtype, + name="sample_d", + mode=output_mode, + ), + JaxTensorDesc.from_shape( + dprob_shape, + jnp.float32, + name="sample_dprob", + mode=PROBABILITY_MODE, + init_value=0.0, + ), + ) + + def check_support(self) -> bool: + self._op.check_support() + if self.a_desc.cudnn_dtype not in _JAX_INPUT_DTYPES: + raise NotImplementedError("The JAX GEMM + dsReLU API supports native FP4 and FP8 inputs") + return True + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + prob_tensor: Any, + ) -> TupleDict: + self.check_support() + + max_active_clusters = self._get_max_active_clusters( + self._op.cluster_shape_mn[0] * self._op.cluster_shape_mn[1], + overlap_margin=self.num_cluster_overlap_margin, + ) + + def launch(stream, a, b, c, sfa, sfb, prob, d, dprob, *optional_outputs): + expected_optional_outputs = int(self.amax_desc is not None) + if len(optional_outputs) != expected_optional_outputs: + raise RuntimeError(f"GemmDsreluSm100 received {len(optional_outputs)} optional output buffers; expected {expected_optional_outputs}") + amax = optional_outputs[0] if self.amax_desc is not None else None + + import cutlass + import cutlass.cute as cute + + from .dense_blockscaled_gemm_persistent_dsrelu_quant import Sm100BlockScaledPersistentDenseGemmKernel + + kernel = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size=self._op.sf_vec_size, + mma_tiler_mn=self._op.mma_tiler_mn, + cluster_shape_mn=self._op.cluster_shape_mn, + vector_f32=self.vector_f32, + ) + + def squared_relu_backward(x, upstream): + return cute.where(x > 0, x, cute.full_like(x, 0)) * 2 * upstream + + kernel( + a, + b, + sfa, + sfb, + c, + d, + prob, + dprob, + amax, + None, + None, + cutlass.Float32(self._op.alpha), + max_active_clusters, + stream, + epilogue_op=squared_relu_backward, + ) + + optional_output_descs = () if self.amax_desc is None else (self.amax_desc,) + output_descs = (self.d_desc, self.dprob_desc) + optional_output_descs + results = self._call_kernel( + (a_tensor, b_tensor, c_tensor, sfa_tensor, sfb_tensor, prob_tensor), + launch=launch, + input_descs=(self.a_desc, self.b_desc, self.c_desc, self.sfa_desc, self.sfb_desc, self.prob_desc), + output_descs=output_descs, + compile_options=compile_options_for_target(self.compute_capability), + ) + d_tensor, dprob_tensor, *optional_results = results + return TupleDict( + d_tensor=d_tensor, + dprob_tensor=dprob_tensor, + amax_tensor=optional_results[0] if self.amax_desc is not None else None, + sfd_tensor=None, + ) + + +@partial( + jax.jit, + static_argnames=( + "alpha", + "d_layout", + "d_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "a_layout", + "b_layout", + ), +) +def gemm_dsrelu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + prob_tensor: Any, + alpha: float = 1.0, + d_layout: str = "LMN", + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + *, + a_layout: str = "LMK", + b_layout: str = "LNK", +) -> TupleDict: + """Compute the block-scaled squared-ReLU backward fusion.""" + + return GemmDsreluSm100( + 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, + 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) + + +__all__ = [ + "GemmDsreluSm100", + "SUPPORTED_COMPUTE_CAPABILITIES", + "gemm_dsrelu_wrapper_sm100", +] diff --git a/python/cudnn/gemm_dsrelu/op.py b/python/cudnn/gemm_dsrelu/op.py new file mode 100644 index 000000000..608c5f134 --- /dev/null +++ b/python/cudnn/gemm_dsrelu/op.py @@ -0,0 +1,21 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral dense GEMM + squared-ReLU backward operation.""" + +from .. import data_type +from ..gemm.srelu import BlockScaledGemmSreluSm100OpBase + + +class GemmDsreluSm100Op(BlockScaledGemmSreluSm100OpBase): + """Logical signature for block-scaled GEMM + squared-ReLU backward.""" + + def check_support(self) -> bool: + if self.dprob is None: + raise ValueError("dprob is required by the dsReLU backward signature") + if self.c.cudnn_dtype not in {data_type.HALF, data_type.BFLOAT16, data_type.FLOAT}: + raise ValueError(f"dsReLU C must use float16, bfloat16, or float32, got {self.c.dtype}") + return super().check_support() + + +__all__ = ["GemmDsreluSm100Op"] diff --git a/python/cudnn/gemm_srelu/__init__.py b/python/cudnn/gemm_srelu/__init__.py index 29324f369..59aa5689f 100644 --- a/python/cudnn/gemm_srelu/__init__.py +++ b/python/cudnn/gemm_srelu/__init__.py @@ -1,9 +1,18 @@ -from .api import ( - GemmSreluSm100, - gemm_srelu_wrapper_sm100, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy Torch API and framework-neutral GEMM + sReLU exports.""" -__all__ = [ - "GemmSreluSm100", - "gemm_srelu_wrapper_sm100", -] +from ..common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ("GemmSreluSm100Op",), + "api": ( + "GemmSreluSm100", + "gemm_srelu_wrapper_sm100", + ), + }, + submodules=("api", "op"), +) diff --git a/python/cudnn/gemm_srelu/api.py b/python/cudnn/gemm_srelu/api.py index 3bf21f95f..90ed43f44 100644 --- a/python/cudnn/gemm_srelu/api.py +++ b/python/cudnn/gemm_srelu/api.py @@ -117,6 +117,10 @@ def check_support(self) -> bool: dtype=[torch.float16, torch.bfloat16, torch.float32, torch.float8_e4m3fn, torch.float8_e5m2], name="D", ) + self._not_implemented_error_if( + self._is_fp8(self.d_dtype), + "FP8 D output is unavailable because the current sReLU kernel does not implement SFD generation", + ) self._check_dtype(self.prob_desc, dtype=torch.float32, name="prob") self.sf_dtype = self._check_dtype(self.sfa_desc, dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn], name="SFA") @@ -126,9 +130,6 @@ def check_support(self) -> bool: self._check_dtype(self.acc_dtype, dtype=torch.float32, name="Accumulator") self._value_error_if(self.sf_vec_size not in {16, 32}, f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}") - self._value_error_if( - self._is_fp8(self.d_desc) and (self.sfd_desc is None or self.norm_const_desc is None), "sfd and norm_const are required when D is FP8" - ) self._value_error_if( self._is_fp4x2(self.ab_dtype) and self.d_dtype in {torch.float8_e4m3fn, torch.float8_e5m2}, "FP4 input with FP8 output is not supported" ) diff --git a/python/cudnn/gemm_srelu/dense_blockscaled_gemm_persistent_srelu_quant.py b/python/cudnn/gemm_srelu/dense_blockscaled_gemm_persistent_srelu_quant.py index 1b338ba70..1ef5eb36a 100644 --- a/python/cudnn/gemm_srelu/dense_blockscaled_gemm_persistent_srelu_quant.py +++ b/python/cudnn/gemm_srelu/dense_blockscaled_gemm_persistent_srelu_quant.py @@ -29,8 +29,6 @@ from typing import Type, Tuple, Union, Optional import cuda.bindings.driver as cuda -import torch - import cutlass import cutlass.cute as cute from cutlass.cute.nvgpu import cpasync, tcgen05 @@ -1375,15 +1373,18 @@ def kernel( self.epilog_sync_barrier.arrive_and_wait() + # + # Apply squared ReLU and the row probability to D. + # + acc_values = acc_vec.load() + acc_values = epilogue_op(acc_values) + acc_values = acc_values * mProb + acc_vec.store(acc_values) + # # Generate amax # if cutlass.const_expr(self.generate_amax): - acc_values = acc_vec.load() - acc_values = epilogue_op(acc_values) - acc_values = acc_values * mProb - acc_vec.store(acc_values) - # Apply element-wise absolute value using math.absf (supports vectors) abs_acc_values_ir = cutlass._mlir.dialects.math.absf(acc_values.ir_value()) # operand (positional) abs_acc_values = type(acc_values)(abs_acc_values_ir, acc_values.shape, acc_values.dtype) diff --git a/python/cudnn/gemm_srelu/jax.py b/python/cudnn/gemm_srelu/jax.py new file mode 100644 index 000000000..efab492d8 --- /dev/null +++ b/python/cudnn/gemm_srelu/jax.py @@ -0,0 +1,281 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX API for block-scaled dense GEMM + squared ReLU.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any + +import jax +import jax.numpy as jnp + +from .. import data_type +from .._jax.compiler import compile_options_for_target +from ..gemm.helpers import require_gemm_inputs +from .._jax import JaxApiBase, JaxTensorDesc, TupleDict +from .._jax.datatypes import jax_to_cudnn_dtype, normalize_jax_dtype +from .._jax.gemm import BLOCK_SCALE_MODE, PROBABILITY_MODE, gemm_a_mode, gemm_b_mode, gemm_output_mode +from .._jax.layout import to_public_axes +from .op import GemmSreluSm100Op + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) +_JAX_INPUT_DTYPES = frozenset({data_type.FP4_E2M1, data_type.FP8_E4M3, data_type.FP8_E5M2}) +_WIDE_OUTPUT_DTYPES = frozenset({data_type.HALF, data_type.BFLOAT16, data_type.FLOAT}) + + +class GemmSreluSm100(JaxApiBase): + """JAX callable specialized from a block-scaled GEMM + sReLU signature.""" + + def __init__( + self, + sample_a: Any, + sample_b: Any, + sample_sfa: Any, + sample_sfb: Any, + sample_prob: Any, + *, + sample_c: Any | None = None, + sample_d: Any | None = None, + alpha: float = 1.0, + c_layout: str = "LMN", + c_dtype: Any | None = None, + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + a_layout: str = "LMK", + b_layout: str = "LNK", + ) -> None: + a_mode = gemm_a_mode(a_layout) + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(c_layout) + + self.compute_capability = self._resolve_compute_capability( + None, + SUPPORTED_COMPUTE_CAPABILITIES, + "GemmSreluSm100", + ) + self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=a_mode) + self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=b_mode) + self.sfa_desc = self._to_tensor_desc(sample_sfa, "sample_sfa", mode=BLOCK_SCALE_MODE) + self.sfb_desc = self._to_tensor_desc(sample_sfb, "sample_sfb", mode=BLOCK_SCALE_MODE) + self.prob_desc = self._to_tensor_desc(sample_prob, "sample_prob", mode=PROBABILITY_MODE) + + self.acc_dtype = normalize_jax_dtype(acc_dtype, jnp.float32, "acc_dtype") + if (sample_c is None) != (sample_d is None): + raise ValueError("sample_c and sample_d must be provided together") + if sample_c is None: + resolved_c_dtype = normalize_jax_dtype(c_dtype, jnp.bfloat16, "c_dtype") + resolved_d_dtype = normalize_jax_dtype(d_dtype, jnp.bfloat16, "d_dtype") + self.c_desc, self.d_desc = self._default_output_descs( + resolved_c_dtype, + resolved_d_dtype, + output_mode, + ) + else: + if c_dtype is not None or d_dtype is not None: + raise ValueError("c_dtype and d_dtype cannot be specified with sample_c and sample_d") + self.c_desc = self._to_tensor_desc(sample_c, "sample_c", mode=output_mode) + self.d_desc = self._to_tensor_desc(sample_d, "sample_d", mode=output_mode) + + self.amax_desc = None + if self.a_desc.cudnn_dtype == data_type.FP4_E2M1 and self.d_desc.cudnn_dtype in _WIDE_OUTPUT_DTYPES: + self.amax_desc = self.d_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(1,), + name="amax_tensor", + init_value=float("-inf"), + ) + acc_cudnn_dtype = jax_to_cudnn_dtype(self.acc_dtype) + if acc_cudnn_dtype == data_type.NOT_SET: + raise ValueError(f"Unsupported JAX accumulator dtype {self.acc_dtype}") + if not isinstance(vector_f32, bool): + raise TypeError(f"vector_f32 must be a bool, got {type(vector_f32).__name__}") + + self._op = GemmSreluSm100Op( + a=self.a_desc, + b=self.b_desc, + c=self.c_desc, + d=self.d_desc, + sfa=self.sfa_desc, + sfb=self.sfb_desc, + prob=self.prob_desc, + sfd=None, + amax=self.amax_desc, + norm_const=None, + alpha=alpha, + acc_dtype=acc_cudnn_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + ) + self.vector_f32 = vector_f32 + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + + def _default_output_descs( + self, + c_dtype: Any, + d_dtype: Any, + output_mode: tuple[int, ...], + ) -> tuple[JaxTensorDesc, JaxTensorDesc]: + m, n, _, batch = require_gemm_inputs(self.a_desc, self.b_desc) + public_shape = to_public_axes((m, n, batch), output_mode) + return ( + JaxTensorDesc.from_shape( + public_shape, + c_dtype, + name="sample_c", + mode=output_mode, + ), + JaxTensorDesc.from_shape( + public_shape, + d_dtype, + name="sample_d", + mode=output_mode, + ), + ) + + def check_support(self) -> bool: + self._op.check_support() + if self.a_desc.cudnn_dtype not in _JAX_INPUT_DTYPES: + raise NotImplementedError("The JAX GEMM + sReLU API supports native FP4 and FP8 inputs") + return True + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + prob_tensor: Any, + ) -> TupleDict: + self.check_support() + + max_active_clusters = self._get_max_active_clusters( + self._op.cluster_shape_mn[0] * self._op.cluster_shape_mn[1], + overlap_margin=self.num_cluster_overlap_margin, + ) + + def launch(stream, a, b, sfa, sfb, prob, c, d, *optional_outputs): + expected_optional_outputs = int(self.amax_desc is not None) + if len(optional_outputs) != expected_optional_outputs: + raise RuntimeError(f"GemmSreluSm100 received {len(optional_outputs)} optional output buffers; expected {expected_optional_outputs}") + amax = optional_outputs[0] if self.amax_desc is not None else None + + import cutlass + import cutlass.cute as cute + + from .dense_blockscaled_gemm_persistent_srelu_quant import Sm100BlockScaledPersistentDenseGemmKernel + + kernel = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size=self._op.sf_vec_size, + mma_tiler_mn=self._op.mma_tiler_mn, + cluster_shape_mn=self._op.cluster_shape_mn, + vector_f32=self.vector_f32, + ) + + def squared_relu(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) ** 2 + + kernel( + a, + b, + sfa, + sfb, + c, + d, + prob, + amax, + None, + None, + cutlass.Float32(self._op.alpha), + max_active_clusters, + stream, + epilogue_op=squared_relu, + ) + + optional_output_descs = () if self.amax_desc is None else (self.amax_desc,) + output_descs = (self.c_desc, self.d_desc) + optional_output_descs + results = self._call_kernel( + (a_tensor, b_tensor, sfa_tensor, sfb_tensor, prob_tensor), + launch=launch, + input_descs=(self.a_desc, self.b_desc, self.sfa_desc, self.sfb_desc, self.prob_desc), + output_descs=output_descs, + compile_options=compile_options_for_target(self.compute_capability), + ) + c_tensor, d_tensor, *optional_results = results + return TupleDict( + c_tensor=c_tensor, + d_tensor=d_tensor, + amax_tensor=optional_results[0] if self.amax_desc is not None else None, + sfd_tensor=None, + ) + + +@partial( + jax.jit, + static_argnames=( + "alpha", + "c_layout", + "c_dtype", + "d_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "a_layout", + "b_layout", + ), +) +def gemm_srelu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + prob_tensor: Any, + alpha: float = 1.0, + c_layout: str = "LMN", + c_dtype: Any | None = None, + d_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + *, + a_layout: str = "LMK", + b_layout: str = "LNK", +) -> TupleDict: + """Compute block-scaled GEMM and ``D = prob * relu(C) ** 2``.""" + + return GemmSreluSm100( + 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, + 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) + + +__all__ = [ + "GemmSreluSm100", + "SUPPORTED_COMPUTE_CAPABILITIES", + "gemm_srelu_wrapper_sm100", +] diff --git a/python/cudnn/gemm_srelu/op.py b/python/cudnn/gemm_srelu/op.py new file mode 100644 index 000000000..b8f3194a4 --- /dev/null +++ b/python/cudnn/gemm_srelu/op.py @@ -0,0 +1,18 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral dense GEMM + squared-ReLU operation.""" + +from ..gemm.srelu import BlockScaledGemmSreluSm100OpBase + + +class GemmSreluSm100Op(BlockScaledGemmSreluSm100OpBase): + """Logical signature for block-scaled GEMM + squared ReLU.""" + + def check_support(self) -> bool: + if self.dprob is not None: + raise ValueError("dprob is only part of the dsReLU backward signature") + return super().check_support() + + +__all__ = ["GemmSreluSm100Op"] diff --git a/python/cudnn/gemm_swiglu/__init__.py b/python/cudnn/gemm_swiglu/__init__.py index cf2170c32..239e6c8b2 100644 --- a/python/cudnn/gemm_swiglu/__init__.py +++ b/python/cudnn/gemm_swiglu/__init__.py @@ -1,9 +1,21 @@ -from .api import ( - GemmSwigluSm100, - gemm_swiglu_wrapper_sm100, -) +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Lazy Torch API and framework-neutral GEMM + SwiGLU operation exports.""" -__all__ = [ - "GemmSwigluSm100", - "gemm_swiglu_wrapper_sm100", -] +from ..common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ( + "BlockScaledGemmSwigluSm100Op", + "GemmSwigluSm100Op", + ), + "api": ( + "GemmSwigluSm100", + "gemm_swiglu_wrapper_sm100", + ), + }, + submodules=("api", "op"), +) diff --git a/python/cudnn/gemm_swiglu/api.py b/python/cudnn/gemm_swiglu/api.py index 2daeaa740..5a49b9034 100644 --- a/python/cudnn/gemm_swiglu/api.py +++ b/python/cudnn/gemm_swiglu/api.py @@ -41,8 +41,10 @@ import cutlass.cute as cute from cutlass.cute.runtime import make_fake_stream -from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn import data_type +from cudnn.datatypes import _convert_to_cutlass_data_type, _torch_to_cudnn_data_type from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +from .op import GemmSwigluSm100Op import os @@ -104,6 +106,24 @@ def __init__( self._logger.debug("Quantization arguments provided, using quantized GEMM swiglu kernel") self._kernel = Sm100BlockScaledPersistentDenseGemmKernel + self._op = None + if self._kernel is PersistentDenseGemmKernel: + canonical_acc_dtype = _torch_to_cudnn_data_type(self.acc_dtype) + if canonical_acc_dtype is None: + canonical_acc_dtype = data_type.NOT_SET + self._op = GemmSwigluSm100Op( + a=self.a_desc, + b=self.b_desc, + ab12=self.ab12_desc, + c=self.c_desc, + alpha=self.alpha, + acc_dtype=canonical_acc_dtype, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + ) + self.mma_tiler_mn = self._op.mma_tiler_mn + self.cluster_shape_mn = self._op.cluster_shape_mn + self._logger.debug( f"__init__ completed with args: sample_a {self.a_desc.shape}, sample_b {self.b_desc.shape}, sample_ab12 {self.ab12_desc.shape}, sample_c {self.c_desc.shape}, alpha {alpha}, acc_dtype {acc_dtype}, mma_tiler_mn {mma_tiler_mn}, cluster_shape_mn {cluster_shape_mn}, sample_sfa {self.sfa_desc.shape if self.sfa_desc is not None else None}, sample_sfb {self.sfb_desc.shape if self.sfb_desc is not None else None}, sample_amax {self.amax_desc.shape if self.amax_desc is not None else None}, sample_sfc {self.sfc_desc.shape if self.sfc_desc is not None else None}, sample_norm_const {self.norm_const_desc.shape if self.norm_const_desc is not None else None}, sf_vec_size {sf_vec_size}, vector_f32 {vector_f32}, ab12_stages {ab12_stages}" ) @@ -111,6 +131,38 @@ def __init__( self._interpret_uint8_as_fp4x2 = True def check_support(self) -> bool: + """Validate the shared operation contract and Torch device support.""" + + self._is_supported = False + if self._kernel is Sm100BlockScaledPersistentDenseGemmKernel: + return self._check_quantized_support() + + if self._op is None: + raise RuntimeError("Standard GEMM + SwiGLU operation was not initialized") + self._op.check_support() + self.mma_tiler_mn = self._op.mma_tiler_mn + self.cluster_shape_mn = self._op.cluster_shape_mn + + self.ab_dtype = self.a_desc.dtype + self.ab12_dtype = self.ab12_desc.dtype + self.c_dtype = self.c_desc.dtype + + self._logger.debug("Checking environment") + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available") + device = torch.cuda.current_device() + major, minor = torch.cuda.get_device_capability(device) + compute_capability = major * 10 + minor + if compute_capability < 100: + raise RuntimeError(f"GemmSwiglu requires SM100+ compute capability, but found SM{compute_capability} on device {device}") + + self._is_supported = True + self._logger.debug("check_support completed successfully") + return True + + def _check_quantized_support(self) -> bool: + """Preserve the existing Torch-only block-scaled validation path.""" + self._logger.debug("Entering check_support") self._logger.debug("Checking tensor shapes, strides, and dtypes") diff --git a/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py b/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py index 28f6aff86..5b6769395 100644 --- a/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py +++ b/python/cudnn/gemm_swiglu/dense_gemm_persistent_swiglu.py @@ -78,7 +78,7 @@ see detailed valid dtype combinations in below PersistentDenseGemmKernel class documentation * A/B tensor must have the same data type * Mma tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) -* Mma tiler N must be 32-256, step 32 +* Mma tiler N must be 64/128/192/256 so each tile contains complete SwiGLU block pairs * Cluster shape M/N must be positive and power of 2, total cluster size <= 16 * Cluster shape M must be multiple of 2 if use_2cta_instrs=True * The contiguous dimension of A/B/C tensors must be at least 16 bytes aligned, @@ -125,7 +125,7 @@ class PersistentDenseGemmKernel: :note: Constraints: - MMA tiler M must be 64/128 (use_2cta_instrs=False) or 128/256 (use_2cta_instrs=True) - - MMA tiler N must be 32-256, step 32 + - MMA tiler N must be 64/128/192/256 so each tile contains complete SwiGLU block pairs - Cluster shape M must be multiple of 2 if use_2cta_instrs=True - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 diff --git a/python/cudnn/gemm_swiglu/jax.py b/python/cudnn/gemm_swiglu/jax.py new file mode 100644 index 000000000..47d998071 --- /dev/null +++ b/python/cudnn/gemm_swiglu/jax.py @@ -0,0 +1,389 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Optional JAX API for standard and block-scaled dense GEMM + SwiGLU.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any + +import jax +import jax.numpy as jnp + +from .. import data_type +from .._jax.compiler import compile_options_for_target +from ..gemm.helpers import require_gemm_inputs +from .._jax import JaxApiBase, JaxTensorDesc, TupleDict +from .._jax.datatypes import jax_to_cudnn_dtype, normalize_jax_dtype +from .._jax.gemm import BLOCK_SCALE_MODE, gemm_a_mode, gemm_b_mode, gemm_output_mode +from .._jax.layout import to_public_axes +from .op import BlockScaledGemmSwigluSm100Op, GemmSwigluSm100Op + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) + + +class GemmSwigluSm100(JaxApiBase): + """JAX callable specialized from dense GEMM input metadata. + + ``a_layout``, ``b_layout``, and ``c_layout`` describe the public JAX axis + order. Descriptors stored by the operation remain in the kernel's + canonical ``MKL``, ``NKL``, and ``MNL`` orders. + + Providing SFA and SFB samples selects the block-scaled kernel. Public scale + arrays are compact row-major ``[L, row_tiles, k_tiles, 32, 4, 4]`` arrays; + the adapter maps them to the kernel's packed scale-factor layout. + """ + + def __init__( + self, + sample_a: Any, + sample_b: Any, + *, + sample_ab12: Any | None = None, + sample_c: Any | None = None, + sample_sfa: Any | None = None, + sample_sfb: Any | None = None, + sample_amax: Any | None = None, + alpha: float = 1.0, + c_layout: str = "LMN", + ab12_dtype: Any | None = None, + c_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: tuple[int, int] | None = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + ab12_stages: int = 4, + a_layout: str = "LMK", + b_layout: str = "LNK", + ) -> None: + a_mode = gemm_a_mode(a_layout) + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(c_layout) + + self.compute_capability = self._resolve_compute_capability( + None, + SUPPORTED_COMPUTE_CAPABILITIES, + "GemmSwigluSm100", + ) + self.a_desc = self._to_tensor_desc(sample_a, "sample_a", mode=a_mode) + self.b_desc = self._to_tensor_desc(sample_b, "sample_b", mode=b_mode) + self.acc_dtype = normalize_jax_dtype(acc_dtype, jnp.float32, "acc_dtype") + + self.is_block_scaled = any( + sample is not None + for sample in ( + sample_sfa, + sample_sfb, + sample_amax, + ) + ) + if self.is_block_scaled and (sample_sfa is None or sample_sfb is None): + raise ValueError("sample_sfa and sample_sfb are required for block-scaled GEMM + SwiGLU") + if not self.is_block_scaled: + if sf_vec_size != 16: + raise ValueError("sf_vec_size only applies to block-scaled GEMM + SwiGLU") + if vector_f32: + raise ValueError("vector_f32 only applies to block-scaled GEMM + SwiGLU") + if ab12_stages != 4: + raise ValueError("ab12_stages only applies to block-scaled GEMM + SwiGLU") + self.sfa_desc = None if sample_sfa is None else self._to_tensor_desc(sample_sfa, "sample_sfa", mode=BLOCK_SCALE_MODE) + self.sfb_desc = None if sample_sfb is None else self._to_tensor_desc(sample_sfb, "sample_sfb", mode=BLOCK_SCALE_MODE) + + if (sample_ab12 is None) != (sample_c is None): + raise ValueError("sample_ab12 and sample_c must be provided together") + if sample_ab12 is None: + resolved_ab12_dtype = normalize_jax_dtype(ab12_dtype, jnp.float32, "ab12_dtype") + resolved_c_dtype = normalize_jax_dtype(c_dtype, jnp.float16, "c_dtype") + self.ab12_desc, self.c_desc = self._default_output_descs( + resolved_ab12_dtype, + resolved_c_dtype, + output_mode, + ) + else: + if ab12_dtype is not None or c_dtype is not None: + raise ValueError("ab12_dtype and c_dtype cannot be specified with sample_ab12 and sample_c") + self.ab12_desc = self._to_tensor_desc( + sample_ab12, + "sample_ab12", + mode=output_mode, + ) + self.c_desc = self._to_tensor_desc( + sample_c, + "sample_c", + mode=output_mode, + ) + + self.amax_desc = None + if self.is_block_scaled: + needs_amax = self.a_desc.cudnn_dtype == data_type.FP4_E2M1 and self.c_desc.cudnn_dtype == data_type.BFLOAT16 + if sample_amax is not None: + self.amax_desc = self._to_tensor_desc(sample_amax, "sample_amax", init_value=float("-inf")) + elif needs_amax: + self.amax_desc = self._default_amax_desc() + + acc_cudnn_dtype = jax_to_cudnn_dtype(self.acc_dtype) + if acc_cudnn_dtype == data_type.NOT_SET: + raise ValueError(f"Unsupported JAX accumulator dtype {self.acc_dtype}") + + if self.is_block_scaled: + self._op = BlockScaledGemmSwigluSm100Op( + a=self.a_desc, + b=self.b_desc, + sfa=self.sfa_desc, + sfb=self.sfb_desc, + ab12=self.ab12_desc, + c=self.c_desc, + sfc=None, + amax=self.amax_desc, + norm_const=None, + alpha=alpha, + acc_dtype=acc_cudnn_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + ab12_stages=ab12_stages, + ) + else: + self._op = GemmSwigluSm100Op( + a=self.a_desc, + b=self.b_desc, + ab12=self.ab12_desc, + c=self.c_desc, + alpha=alpha, + acc_dtype=acc_cudnn_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + ) + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + + def _default_output_descs( + self, + ab12_dtype: Any, + c_dtype: Any, + output_mode: tuple[int, ...], + ) -> tuple[JaxTensorDesc, JaxTensorDesc]: + m, n, _, batch = require_gemm_inputs(self.a_desc, self.b_desc) + if n % 2: + raise ValueError(f"SwiGLU requires an even N dimension, got {n}") + + return ( + JaxTensorDesc.from_shape( + to_public_axes((m, n, batch), output_mode), + ab12_dtype, + name="sample_ab12", + mode=output_mode, + ), + JaxTensorDesc.from_shape( + to_public_axes((m, n // 2, batch), output_mode), + c_dtype, + name="sample_c", + mode=output_mode, + ), + ) + + def _default_amax_desc(self) -> JaxTensorDesc: + return JaxTensorDesc.from_shape( + (1,), + jnp.float32, + name="sample_amax", + init_value=float("-inf"), + ) + + def check_support(self) -> bool: + return self._op.check_support() + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any | None = None, + sfb_tensor: Any | None = None, + ) -> TupleDict: + self.check_support() + if self.is_block_scaled: + if sfa_tensor is None or sfb_tensor is None: + raise ValueError("sfa_tensor and sfb_tensor are required for this block-scaled callable") + elif sfa_tensor is not None or sfb_tensor is not None: + raise ValueError("Scale-factor tensors are not part of this standard GEMM + SwiGLU callable") + + max_active_clusters = self._get_max_active_clusters( + self._op.cluster_shape_mn[0] * self._op.cluster_shape_mn[1], + overlap_margin=self.num_cluster_overlap_margin, + ) + + if not self.is_block_scaled: + + def launch( + stream: Any, + a: Any, + b: Any, + ab12: Any, + c: Any, + ) -> None: + import cutlass + from cutlass.jax import jax_to_cutlass_dtype + + from .dense_gemm_persistent_swiglu import PersistentDenseGemmKernel + + kernel = PersistentDenseGemmKernel( + acc_dtype=jax_to_cutlass_dtype(self.acc_dtype), + use_2cta_instrs=self._op.mma_tiler_mn[0] == 256, + mma_tiler_mn=self._op.mma_tiler_mn, + cluster_shape_mn=self._op.cluster_shape_mn, + ) + kernel( + a, + b, + ab12, + c, + cutlass.Float32(self._op.alpha), + max_active_clusters, + stream, + ) + + ab12_tensor, c_tensor = self._call_kernel( + (a_tensor, b_tensor), + launch=launch, + input_descs=(self.a_desc, self.b_desc), + output_descs=(self.ab12_desc, self.c_desc), + compile_options=compile_options_for_target(self.compute_capability), + ) + return TupleDict( + ab12_tensor=ab12_tensor, + c_tensor=c_tensor, + sfc_tensor=None, + amax_tensor=None, + ) + + kernel_inputs = (a_tensor, b_tensor, sfa_tensor, sfb_tensor) + input_descs = (self.a_desc, self.b_desc, self.sfa_desc, self.sfb_desc) + + output_descs = (self.ab12_desc, self.c_desc) + if self.amax_desc is not None: + output_descs += (self.amax_desc,) + + def launch_block_scaled( + stream: Any, + a: Any, + b: Any, + sfa: Any, + sfb: Any, + ab12: Any, + c: Any, + *optional_outputs: Any, + ) -> None: + expected_optional_outputs = int(self.amax_desc is not None) + if len(optional_outputs) != expected_optional_outputs: + raise RuntimeError(f"GemmSwigluSm100 received {len(optional_outputs)} optional output buffers; expected {expected_optional_outputs}") + amax = optional_outputs[0] if self.amax_desc is not None else None + + import cutlass + + from .dense_blockscaled_gemm_persistent_swiglu_interleaved_quant import Sm100BlockScaledPersistentDenseGemmKernel + + kernel = Sm100BlockScaledPersistentDenseGemmKernel( + sf_vec_size=self._op.sf_vec_size, + mma_tiler_mn=self._op.mma_tiler_mn, + cluster_shape_mn=self._op.cluster_shape_mn, + vector_f32=self._op.vector_f32, + ab12_stages=self._op.ab12_stages, + ) + kernel( + a, + b, + sfa, + sfb, + c, + ab12, + amax, + None, + None, + cutlass.Float32(self._op.alpha), + max_active_clusters, + stream, + ) + + results = self._call_kernel( + kernel_inputs, + launch=launch_block_scaled, + input_descs=input_descs, + output_descs=output_descs, + compile_options=compile_options_for_target(self.compute_capability), + ) + ab12_tensor, c_tensor, *optional_results = results + amax_tensor = optional_results[0] if self.amax_desc is not None else None + return TupleDict( + ab12_tensor=ab12_tensor, + c_tensor=c_tensor, + sfc_tensor=None, + amax_tensor=amax_tensor, + ) + + +@partial( + jax.jit, + static_argnames=( + "alpha", + "c_layout", + "ab12_dtype", + "c_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "ab12_stages", + "a_layout", + "b_layout", + ), +) +def gemm_swiglu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + alpha: float = 1.0, + c_layout: str = "LMN", + ab12_dtype: Any | None = None, + c_dtype: Any | None = None, + acc_dtype: Any | None = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: tuple[int, int] | None = None, + sfa_tensor: Any | None = None, + sfb_tensor: Any | None = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + ab12_stages: int = 4, + *, + a_layout: str = "LMK", + b_layout: str = "LNK", +) -> TupleDict: + """Compute standard or block-scaled GEMM + SwiGLU on a local SM100-family GPU.""" + + return GemmSwigluSm100( + a_tensor, + b_tensor, + sample_sfa=sfa_tensor, + sample_sfb=sfb_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, + 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) + + +__all__ = [ + "GemmSwigluSm100", + "SUPPORTED_COMPUTE_CAPABILITIES", + "gemm_swiglu_wrapper_sm100", +] diff --git a/python/cudnn/gemm_swiglu/op.py b/python/cudnn/gemm_swiglu/op.py new file mode 100644 index 000000000..3dfde75a1 --- /dev/null +++ b/python/cudnn/gemm_swiglu/op.py @@ -0,0 +1,371 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral dense GEMM + SwiGLU operation.""" + +from __future__ import annotations + +from typing import Any, Optional + +from .. import data_type +from ..gemm.helpers import ( + block_scale_shape, + require_16_byte_alignment, + require_block_scale_layout, + require_cluster_shape, + require_compact_major, + require_gemm_inputs, + require_mma_tiler, + require_tensor_shape, +) +from ..common.op import Op +from ..common.tensor_desc import TensorDesc + +_STANDARD_AB_DTYPES = frozenset( + { + data_type.HALF, + data_type.BFLOAT16, + data_type.FLOAT, + data_type.FP8_E4M3, + data_type.FP8_E5M2, + } +) +_STANDARD_OUTPUT_DTYPES = frozenset( + { + data_type.HALF, + data_type.BFLOAT16, + } +) +_STANDARD_AB12_DTYPES = _STANDARD_OUTPUT_DTYPES | { + data_type.FLOAT, + data_type.FP8_E4M3, + data_type.FP8_E5M2, +} +_FP8_DTYPES = frozenset({data_type.FP8_E4M3, data_type.FP8_E5M2}) +_BLOCK_SCALED_AB_DTYPES = frozenset({data_type.FP4_E2M1, *_FP8_DTYPES}) +_BLOCK_SCALED_OUTPUT_DTYPES = frozenset( + { + data_type.HALF, + data_type.BFLOAT16, + data_type.FLOAT, + *_FP8_DTYPES, + } +) +_BLOCK_SCALE_DTYPES = frozenset({data_type.FP8_E4M3, data_type.FP8_E8M0}) +_SWIGLU_BLOCK_COLUMNS = 32 +_SWIGLU_COLUMNS_PER_PAIR = 2 * _SWIGLU_BLOCK_COLUMNS +_SWIGLU_MMA_N = tuple(range(_SWIGLU_COLUMNS_PER_PAIR, 257, _SWIGLU_COLUMNS_PER_PAIR)) + + +class GemmSwigluSm100Op(Op): + """Logical signature and launch configuration for standard GEMM + SwiGLU.""" + + def __init__( + self, + *, + a: TensorDesc[Any], + b: TensorDesc[Any], + ab12: TensorDesc[Any], + c: TensorDesc[Any], + alpha: float = 1.0, + acc_dtype: data_type = data_type.FLOAT, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: Optional[tuple[int, int]] = None, + ) -> None: + for name, desc in (("a", a), ("b", b), ("ab12", ab12), ("c", c)): + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + if not isinstance(acc_dtype, data_type): + raise TypeError(f"acc_dtype must be a cudnn.data_type, got {type(acc_dtype).__name__}") + + self.a = a + self.b = b + self.ab12 = ab12 + self.c = c + + self.alpha = alpha + self.acc_dtype = acc_dtype + self.mma_tiler_mn = tuple(mma_tiler_mn) + if cluster_shape_mn is None: + use_2cta_default = len(self.mma_tiler_mn) == 2 and self.mma_tiler_mn[0] == 256 + self.cluster_shape_mn = (2, 2) if use_2cta_default else (1, 1) + else: + self.cluster_shape_mn = tuple(cluster_shape_mn) + + self.m: Optional[int] = None + self.n: Optional[int] = None + self.k: Optional[int] = None + self.l: Optional[int] = None + self.output_n: Optional[int] = None + self.a_major: Optional[str] = None + self.b_major: Optional[str] = None + self.output_major: Optional[str] = None + + def check_support(self) -> bool: + """Validate the logical signature and resolve its canonical modes.""" + + self.m = self.n = self.k = self.l = self.output_n = None + self.a_major = self.b_major = self.output_major = None + + m, n, k, l = require_gemm_inputs(self.a, self.b) + if n % 2 != 0: + raise ValueError(f"N must be even for SwiGLU input/gate pairs, got {n}") + if n % _SWIGLU_COLUMNS_PER_PAIR != 0: + raise ValueError( + f"N must be divisible by {_SWIGLU_COLUMNS_PER_PAIR} because SwiGLU pairs consecutive {_SWIGLU_BLOCK_COLUMNS}-column blocks, got {n}" + ) + output_n = n // 2 + + require_tensor_shape(self.ab12, (m, n, l), label="AB12") + require_tensor_shape(self.c, (m, output_n, l), label="C") + + a_major = require_compact_major(self.a, "m", "k") + b_major = require_compact_major(self.b, "n", "k") + ab12_major = require_compact_major(self.ab12, "m", "n") + c_major = require_compact_major(self.c, "m", "n") + if ab12_major != c_major: + raise ValueError(f"AB12 and C must use the same major mode, got {ab12_major}-major and {c_major}-major") + + self.m, self.n, self.k, self.l = m, n, k, l + self.output_n = output_n + self.a_major, self.b_major, self.output_major = a_major, b_major, ab12_major + + self._check_standard_dtypes() + self.mma_tiler_mn = require_mma_tiler( + self.mma_tiler_mn, + allowed_n=_SWIGLU_MMA_N, + ) + cta_group_size = 2 if self.mma_tiler_mn[0] == 256 else 1 + if cta_group_size == 2 and m % self.mma_tiler_mn[0] != 0: + raise ValueError(f"M must be divisible by {self.mma_tiler_mn[0]} for 2-CTA MMA, got {m}") + self.cluster_shape_mn = require_cluster_shape( + self.cluster_shape_mn, + cta_group_size=cta_group_size, + ) + if cta_group_size == 1 and self.cluster_shape_mn != (1, 1): + raise ValueError("cluster_shape_mn must be (1, 1) with a 128-wide M tile") + + for tensor in (self.a, self.b, self.ab12, self.c): + require_16_byte_alignment(tensor) + return True + + def _check_standard_dtypes(self) -> None: + ab_dtype = self.a.cudnn_dtype + if ab_dtype not in _STANDARD_AB_DTYPES: + raise ValueError(f"A dtype must be one of the supported dense GEMM dtypes, got {self.a.dtype}") + if self.b.cudnn_dtype != ab_dtype: + raise ValueError(f"A and B must have the same dtype, got {self.a.dtype} and {self.b.dtype}") + + ab12_dtype = self.ab12.cudnn_dtype + if self.acc_dtype == data_type.FLOAT: + if ab12_dtype not in _STANDARD_AB12_DTYPES: + raise ValueError(f"AB12 has unsupported dtype {self.ab12.dtype} for float32 accumulation") + if ab12_dtype in _FP8_DTYPES: + raise NotImplementedError("FP8 AB12 output is currently disabled") + elif self.acc_dtype == data_type.HALF: + if ab12_dtype not in _STANDARD_OUTPUT_DTYPES: + raise ValueError(f"AB12 has unsupported dtype {self.ab12.dtype} for float16 accumulation") + if ab_dtype not in _FP8_DTYPES | {data_type.HALF}: + raise ValueError(f"A and B dtype {self.a.dtype} is unsupported for float16 accumulation") + else: + raise ValueError(f"Accumulator dtype must be float32 or float16, got {self.acc_dtype}") + + if self.c.cudnn_dtype not in _STANDARD_OUTPUT_DTYPES: + raise ValueError(f"C dtype must be float16 or bfloat16, got {self.c.dtype}") + + +class BlockScaledGemmSwigluSm100Op(Op): + """Logical signature for block-scaled GEMM + SwiGLU and quantized outputs.""" + + def __init__( + self, + *, + a: TensorDesc[Any], + b: TensorDesc[Any], + sfa: TensorDesc[Any], + sfb: TensorDesc[Any], + ab12: TensorDesc[Any], + c: TensorDesc[Any], + sfc: TensorDesc[Any] | None = None, + amax: TensorDesc[Any] | None = None, + norm_const: TensorDesc[Any] | None = None, + alpha: float = 1.0, + acc_dtype: data_type = data_type.FLOAT, + mma_tiler_mn: tuple[int, int] = (128, 128), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + ab12_stages: int = 4, + ) -> None: + for name, desc in ( + ("a", a), + ("b", b), + ("sfa", sfa), + ("sfb", sfb), + ("ab12", ab12), + ("c", c), + ): + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + for name, desc in (("sfc", sfc), ("amax", amax), ("norm_const", norm_const)): + if desc is not None and not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc or None, got {type(desc).__name__}") + if not isinstance(acc_dtype, data_type): + raise TypeError(f"acc_dtype must be a cudnn.data_type, got {type(acc_dtype).__name__}") + if isinstance(sf_vec_size, bool) or not isinstance(sf_vec_size, int): + raise TypeError(f"sf_vec_size must be an int, got {type(sf_vec_size).__name__}") + if not isinstance(vector_f32, bool): + raise TypeError(f"vector_f32 must be a bool, got {type(vector_f32).__name__}") + if isinstance(ab12_stages, bool) or not isinstance(ab12_stages, int): + raise TypeError(f"ab12_stages must be an int, got {type(ab12_stages).__name__}") + + self.a = a + self.b = b + self.sfa = sfa + self.sfb = sfb + self.ab12 = ab12 + self.c = c + self.sfc = sfc + self.amax = amax + self.norm_const = norm_const + self.alpha = alpha + self.acc_dtype = acc_dtype + self.requested_mma_tiler_mn = mma_tiler_mn + self.requested_cluster_shape_mn = cluster_shape_mn + self.sf_vec_size = sf_vec_size + self.vector_f32 = vector_f32 + self.ab12_stages = ab12_stages + + self.m: Optional[int] = None + self.n: Optional[int] = None + self.k: Optional[int] = None + self.l: Optional[int] = None + self.output_n: Optional[int] = None + self.a_major: Optional[str] = None + self.b_major: Optional[str] = None + self.output_major: Optional[str] = None + self.ab_dtype: Optional[data_type] = None + self.sf_dtype: Optional[data_type] = None + self.ab12_dtype: Optional[data_type] = None + self.c_dtype: Optional[data_type] = None + self.mma_tiler_mn: Optional[tuple[int, int]] = None + self.cluster_shape_mn: Optional[tuple[int, int]] = None + + def check_support(self) -> bool: + """Validate shapes, packed scale layouts, dtypes, and static tiling.""" + + self.m = self.n = self.k = self.l = self.output_n = None + self.a_major = self.b_major = self.output_major = None + self.ab_dtype = self.sf_dtype = self.ab12_dtype = self.c_dtype = None + self.mma_tiler_mn = self.cluster_shape_mn = None + + m, n, k, batch = require_gemm_inputs(self.a, self.b) + if n % _SWIGLU_COLUMNS_PER_PAIR: + raise ValueError(f"N must be divisible by {_SWIGLU_COLUMNS_PER_PAIR} for SwiGLU block pairs, got {n}") + output_n = n // 2 + require_tensor_shape(self.ab12, (m, n, batch), label="AB12") + require_tensor_shape(self.c, (m, output_n, batch), label="C") + + if self.sf_vec_size not in (16, 32): + raise ValueError(f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}") + require_tensor_shape(self.sfa, block_scale_shape(m, k, batch, self.sf_vec_size), label="SFA") + require_tensor_shape(self.sfb, block_scale_shape(n, k, batch, self.sf_vec_size), label="SFB") + require_block_scale_layout(self.sfa, "SFA") + require_block_scale_layout(self.sfb, "SFB") + + ab_dtype = self.a.cudnn_dtype + if ab_dtype not in _BLOCK_SCALED_AB_DTYPES: + raise ValueError(f"A has unsupported block-scaled dtype {self.a.dtype}") + if self.b.cudnn_dtype != ab_dtype: + raise ValueError(f"A and B must have the same dtype, got {self.a.dtype} and {self.b.dtype}") + + sf_dtype = self.sfa.cudnn_dtype + if sf_dtype not in _BLOCK_SCALE_DTYPES: + raise ValueError(f"SFA has unsupported scale-factor dtype {self.sfa.dtype}") + if self.sfb.cudnn_dtype != sf_dtype: + raise ValueError(f"SFA and SFB must have the same dtype, got {self.sfa.dtype} and {self.sfb.dtype}") + if ab_dtype in _FP8_DTYPES and (sf_dtype, self.sf_vec_size) != (data_type.FP8_E8M0, 32): + raise ValueError("FP8 A and B require FP8_E8M0 scales with sf_vec_size=32") + if ab_dtype == data_type.FP4_E2M1 and (sf_dtype, self.sf_vec_size) == (data_type.FP8_E4M3, 32): + raise ValueError("FP4 A and B do not support FP8_E4M3 scales with sf_vec_size=32") + + ab12_dtype = self.ab12.cudnn_dtype + c_dtype = self.c.cudnn_dtype + if ab12_dtype not in _BLOCK_SCALED_OUTPUT_DTYPES: + raise ValueError(f"AB12 has unsupported dtype {self.ab12.dtype}") + if c_dtype not in _BLOCK_SCALED_OUTPUT_DTYPES: + raise ValueError(f"C has unsupported dtype {self.c.dtype}") + if self.acc_dtype != data_type.FLOAT: + raise ValueError(f"Accumulator dtype must be float32, got {self.acc_dtype}") + if ab_dtype == data_type.FP4_E2M1 and c_dtype in _FP8_DTYPES: + raise ValueError("FP4 A and B are not compatible with FP8 C") + if c_dtype == data_type.FLOAT and ab12_dtype == data_type.FLOAT: + raise NotImplementedError("float32 C and float32 AB12 are disabled because the kernel fails to launch") + if ab_dtype in _FP8_DTYPES and (c_dtype in _FP8_DTYPES or ab12_dtype in _FP8_DTYPES or ab12_dtype == data_type.FLOAT): + raise NotImplementedError("MXFP8 inputs require float16 or bfloat16 AB12 and non-FP8 C") + + if c_dtype in _FP8_DTYPES: + if self.sfc is None or self.norm_const is None: + raise ValueError("sfc and norm_const are required when C is FP8") + if self.sfc is not None: + require_tensor_shape(self.sfc, block_scale_shape(m, output_n, batch, self.sf_vec_size), label="SFC") + require_block_scale_layout(self.sfc, "SFC") + if self.sfc.cudnn_dtype != sf_dtype: + raise ValueError(f"SFC must have the same dtype as SFA, got {self.sfc.dtype} and {self.sfa.dtype}") + if self.norm_const is not None: + require_tensor_shape(self.norm_const, (1,), label="norm_const") + if self.norm_const.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"norm_const must have float32 dtype, got {self.norm_const.dtype}") + if ab_dtype == data_type.FP4_E2M1 and c_dtype == data_type.BFLOAT16 and self.amax is None: + raise ValueError("amax is required when A and B are FP4 and C is bfloat16") + if self.amax is not None: + require_tensor_shape(self.amax, (1,), label="amax") + if self.amax.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"amax must have float32 dtype, got {self.amax.dtype}") + + a_major = require_compact_major(self.a, "m", "k") + b_major = require_compact_major(self.b, "n", "k") + ab12_major = require_compact_major(self.ab12, "m", "n") + c_major = require_compact_major(self.c, "m", "n") + if ab12_major != c_major: + raise ValueError(f"AB12 and C must use the same major mode, got {ab12_major}-major and {c_major}-major") + if ab_dtype == data_type.FP4_E2M1 and (a_major != "k" or b_major != "k"): + raise ValueError("FP4 A and B must use K-major layouts") + if ab_dtype == data_type.FP4_E2M1 and ab12_major != "n": + raise ValueError("FP4 AB12 and C must use N-major layouts") + + mma_tiler_mn = require_mma_tiler( + self.requested_mma_tiler_mn, + allowed_n=_SWIGLU_MMA_N, + ) + cta_group_size = 2 if mma_tiler_mn[0] == 256 else 1 + default_cluster = (2, 2) if cta_group_size == 2 else (1, 1) + cluster_shape_mn = require_cluster_shape( + default_cluster if self.requested_cluster_shape_mn is None else self.requested_cluster_shape_mn, + cta_group_size=cta_group_size, + ) + if cluster_shape_mn[0] > 4 or cluster_shape_mn[1] > 4: + raise ValueError(f"cluster_shape_mn entries must be at most 4 for scale-factor multicast, got {cluster_shape_mn}") + if m % mma_tiler_mn[0] or n % mma_tiler_mn[1]: + raise ValueError(f"M and N must be divisible by mma_tiler_mn, got M={m}, N={n}, tile={mma_tiler_mn}") + if self.ab12_stages <= 0: + raise ValueError(f"ab12_stages must be positive, got {self.ab12_stages}") + + for tensor in (self.a, self.b, self.ab12, self.c): + require_16_byte_alignment(tensor) + + self.m, self.n, self.k, self.l = m, n, k, batch + self.output_n = output_n + self.a_major, self.b_major, self.output_major = a_major, b_major, ab12_major + self.ab_dtype, self.sf_dtype = ab_dtype, sf_dtype + self.ab12_dtype, self.c_dtype = ab12_dtype, c_dtype + self.mma_tiler_mn = mma_tiler_mn + self.cluster_shape_mn = cluster_shape_mn + return True + + +__all__ = [ + "BlockScaledGemmSwigluSm100Op", + "GemmSwigluSm100Op", + "block_scale_shape", +] diff --git a/python/cudnn/grouped_gemm/__init__.py b/python/cudnn/grouped_gemm/__init__.py index 818a1ab99..9830fbd84 100644 --- a/python/cudnn/grouped_gemm/__init__.py +++ b/python/cudnn/grouped_gemm/__init__.py @@ -1,68 +1,30 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .grouped_gemm_swiglu.api import ( - GroupedGemmSwigluSm100, - grouped_gemm_swiglu_wrapper_sm100, +"""Lazy Torch API exports for the grouped GEMM operation family.""" + +from ..common.operation_api import make_operation_api + +_OPERATIONS = ( + "swiglu", + "dswiglu", + "quant", + "srelu", + "dsrelu", + "glu", + "glu_hadamard", + "dglu", + "wgrad", +) + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + f"grouped_gemm_{operation}.api": ( + f"GroupedGemm{''.join(part.capitalize() for part in operation.split('_'))}Sm100", + f"grouped_gemm_{operation}_wrapper_sm100", + ) + for operation in _OPERATIONS + }, + submodules=tuple(f"grouped_gemm_{operation}" for operation in _OPERATIONS), ) - -from .grouped_gemm_dswiglu.api import ( - GroupedGemmDswigluSm100, - grouped_gemm_dswiglu_wrapper_sm100, -) - -from .grouped_gemm_quant.api import ( - GroupedGemmQuantSm100, - grouped_gemm_quant_wrapper_sm100, -) - -from .grouped_gemm_srelu.api import ( - GroupedGemmSreluSm100, - grouped_gemm_srelu_wrapper_sm100, -) - -from .grouped_gemm_dsrelu.api import ( - GroupedGemmDsreluSm100, - grouped_gemm_dsrelu_wrapper_sm100, -) - -from .grouped_gemm_glu.api import ( - GroupedGemmGluSm100, - grouped_gemm_glu_wrapper_sm100, -) - -from .grouped_gemm_glu_hadamard.api import ( - GroupedGemmGluHadamardSm100, - grouped_gemm_glu_hadamard_wrapper_sm100, -) - -from .grouped_gemm_dglu.api import ( - GroupedGemmDgluSm100, - grouped_gemm_dglu_wrapper_sm100, -) - -from .grouped_gemm_wgrad.api import ( - GroupedGemmWgradSm100, - grouped_gemm_wgrad_wrapper_sm100, -) - -__all__ = [ - "GroupedGemmSwigluSm100", - "grouped_gemm_swiglu_wrapper_sm100", - "GroupedGemmDswigluSm100", - "grouped_gemm_dswiglu_wrapper_sm100", - "GroupedGemmQuantSm100", - "grouped_gemm_quant_wrapper_sm100", - "GroupedGemmSreluSm100", - "grouped_gemm_srelu_wrapper_sm100", - "GroupedGemmDsreluSm100", - "grouped_gemm_dsrelu_wrapper_sm100", - "GroupedGemmGluSm100", - "grouped_gemm_glu_wrapper_sm100", - "GroupedGemmGluHadamardSm100", - "grouped_gemm_glu_hadamard_wrapper_sm100", - "GroupedGemmDgluSm100", - "grouped_gemm_dglu_wrapper_sm100", - "GroupedGemmWgradSm100", - "grouped_gemm_wgrad_wrapper_sm100", -] diff --git a/python/cudnn/grouped_gemm/_jax_api.py b/python/cudnn/grouped_gemm/_jax_api.py new file mode 100644 index 000000000..82db2a9ee --- /dev/null +++ b/python/cudnn/grouped_gemm/_jax_api.py @@ -0,0 +1,693 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared JAX metadata and lowering helpers for contiguous grouped GEMMs.""" + +from __future__ import annotations + +from collections.abc import Iterable, Mapping, Sequence +from typing import Any + +import jax.numpy as jnp + +from .._jax.compiler import compile_options_for_target +from ..gemm.helpers import ( + block_scale_shape as _canonical_block_scale_shape, + data_type_bits, +) +from .._jax import JaxApiBase, JaxTensorDesc, TupleDict +from .._jax.gemm import BLOCK_SCALE_MODE, gemm_a_mode, gemm_b_mode, gemm_output_mode +from .._jax.layout import normalize_mode, to_public_axes + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) +MAX_EXPERTS = 1024 +FP8_SF_VEC_SIZE = 32 +TWO_CTA_MMA_TILER_M = 256 +FIX_PAD_SIZE = 256 +HADAMARD_SIZE = 16 +SF_VEC_SIZES = (16, 32) + +_NO_DEFAULT = object() + + +def as_dtype(value: Any) -> Any: + """Return a normalized JAX dtype without retaining an array value.""" + + if not isinstance(value, type) and hasattr(value, "dtype"): + value = value.dtype + return jnp.dtype(value) + + +def is_fp4_dtype(value: Any) -> bool: + """Return whether ``value`` is JAX's native logical FP4 dtype.""" + + return as_dtype(value) == jnp.dtype(jnp.float4_e2m1fn) + + +def is_fp8_dtype(value: Any) -> bool: + """Return whether ``value`` is one of the supported JAX FP8 data dtypes.""" + + return as_dtype(value) in { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + } + + +def is_low_precision_output_dtype(value: Any) -> bool: + """Return whether an output uses a native FP4 or FP8 dtype.""" + + return is_fp4_dtype(value) or is_fp8_dtype(value) + + +def require_dtype( + value: Any, + valid_dtypes: Iterable[Any], + *, + name: str = "dtype", + default: Any = _NO_DEFAULT, +) -> Any: + """Normalize and validate a JAX dtype argument.""" + + if value is None: + if default is _NO_DEFAULT: + raise ValueError(f"{name} must not be None") + value = default + dtype = as_dtype(value) + valid = tuple(as_dtype(item) for item in valid_dtypes) + if dtype not in valid: + supported = ", ".join(item.name for item in valid) + raise ValueError(f"{name} must be one of {{{supported}}}, got {dtype}") + return dtype + + +def require_array( + value: Any, + *, + name: str, + rank: int | Iterable[int] | None = None, + shape: Sequence[int] | None = None, + dtype: Any | Iterable[Any] | None = None, +) -> tuple[int, ...]: + """Validate abstract array metadata and return its public shape.""" + + if not hasattr(value, "shape") or not hasattr(value, "dtype"): + raise TypeError(f"{name} must have shape and dtype metadata") + metadata_shape = tuple(value.shape) + # Sample descriptors store canonical kernel axes. Exact shape checks in + # this adapter describe the public JAX value, so compare against the + # descriptor's inverse-mapped array shape when one is available. + actual_shape = ( + tuple(value.array_shape) + if shape is not None and hasattr(value, "array_shape") + else metadata_shape + ) + if rank is not None: + ranks = (rank,) if isinstance(rank, int) else tuple(rank) + if len(actual_shape) not in ranks: + raise ValueError( + f"{name} must have rank in {ranks}, got shape {actual_shape}" + ) + if shape is not None and actual_shape != tuple(shape): + raise ValueError(f"{name} must have shape {tuple(shape)}, got {actual_shape}") + if dtype is not None: + if ( + isinstance(dtype, Iterable) + and not isinstance(dtype, (str, bytes, type)) + and not hasattr(dtype, "dtype") + ): + valid_dtypes = tuple(dtype) + else: + valid_dtypes = (dtype,) + require_dtype(value, valid_dtypes, name=f"{name}.dtype") + return actual_shape + + +def require_layout(name: str, value: str, supported: tuple[str, ...]) -> str: + """Validate an explicit public axis-order string.""" + + if not isinstance(value, str): + raise TypeError(f"{name} must be a string, got {type(value).__name__}") + if value not in supported: + choices = ", ".join(repr(item) for item in supported) + raise ValueError(f"{name} must be one of ({choices}), got {value!r}") + return value + + +PROBABILITY_MODE = (2, 1, 0) +GROUPED_BIAS_MODE = (1, 0) +GROUPED_WORKSPACE_ALIGNMENT = 128 + + +def make_buffer_desc( + name: str, + shape: Sequence[int], + dtype: Any, + *, + mode: tuple[int, ...] | None = None, + public_stride_order: tuple[int, ...] | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, + init_value: bool | int | float | None = None, +) -> JaxTensorDesc: + """Describe an inferred JAX output or workspace buffer.""" + + if not name: + raise ValueError("buffer descriptor name must not be empty") + return JaxTensorDesc.from_shape( + tuple(shape), + dtype, + name=name, + mode=mode, + public_stride_order=public_stride_order, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + init_value=init_value, + ) + + +def as_gemm_tensor_desc( + name: str, + value: Any, + *, + mode: tuple[int, ...] | None = None, + public_stride_order: tuple[int, ...] | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, +) -> JaxTensorDesc: + if isinstance(value, JaxTensorDesc): + if value.mode != normalize_mode(value.ndim, mode): + raise ValueError( + f"{name} descriptor mode does not match the requested layout" + ) + expected = JaxTensorDesc.from_shape( + value.array_shape, + value.dtype, + name=name, + mode=mode, + public_stride_order=public_stride_order, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + ) + if (value.stride, value.stride_order) != ( + expected.stride, + expected.stride_order, + ): + raise ValueError( + f"{name} descriptor stride order does not match the requested layout" + ) + if value.divisibility != expected.divisibility: + raise ValueError( + f"{name} descriptor divisibility does not match the requested layout" + ) + if value.ptr_assumed_align != expected.ptr_assumed_align: + raise ValueError( + f"{name} descriptor pointer alignment does not match the requested layout" + ) + return value + return JaxTensorDesc.from_array( + value, + name=name, + mode=mode, + public_stride_order=public_stride_order, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + ) + + +class ApiBaseJax(JaxApiBase): + """Sample-signature-bound grouped GEMM callable using ``JaxApiBase``.""" + + def __init__(self) -> None: + self._is_supported = False + + def check_support(self) -> bool: + if not self._is_supported: + self._check_support() + self._is_supported = True + return True + + def _check_support(self) -> None: + raise NotImplementedError + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + self.check_support() + return self._call_impl(*args, **kwargs) + + def _call_impl(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError + + def make_tensor_desc( + self, + value: Any, + *, + mode: tuple[int, ...] | None = None, + public_stride_order: tuple[int, ...] | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, + name: str = "", + ) -> JaxTensorDesc: + return JaxTensorDesc.from_array( + value, + name=name or "value", + mode=mode, + public_stride_order=public_stride_order, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + ) + + def make_optional_tensor_desc( + self, + value: Any | None, + *, + mode: tuple[int, ...] | None = None, + public_stride_order: tuple[int, ...] | None = None, + divisibility: tuple[int | None, ...] | None = None, + ptr_assumed_align: int | None = None, + name: str = "", + ) -> JaxTensorDesc | None: + return ( + None + if value is None + else self.make_tensor_desc( + value, + mode=mode, + public_stride_order=public_stride_order, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + name=name, + ) + ) + + @staticmethod + def as_optional_dtype(value: Any | None) -> Any | None: + return None if value is None else as_dtype(value) + + @staticmethod + def freeze_mapping(values: Mapping[str, Any]) -> Mapping[str, Any]: + # JAX captures the current configuration while tracing. Keep ordinary + # Python state so advanced users may reconfigure an un-jitted class + # instance and let the surrounding JIT establish its own cache key. + return dict(values) + + def check_tensor_signatures( + self, + expected: Mapping[str, JaxTensorDesc | None], + values: Mapping[str, Any], + ) -> None: + for name, expected_desc in expected.items(): + value = values[name] + if value is None or expected_desc is None: + if value is not None or expected_desc is not None: + raise ValueError( + f"{name} presence does not match the sample signature" + ) + continue + actual_shape = require_array(value, name=name) + if actual_shape != expected_desc.array_shape: + raise ValueError( + f"{name} tensor shape mismatch: expected {expected_desc.array_shape}, got {actual_shape}" + ) + if as_dtype(value) != as_dtype(expected_desc.dtype): + raise ValueError( + f"{name} tensor dtype mismatch: expected {expected_desc.dtype}, got {value.dtype}" + ) + + +class _GroupedKernelCaller(JaxApiBase): + def check_support(self) -> bool: + return True + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + raise TypeError("_GroupedKernelCaller is an internal lowering helper") + + +_CALLER = _GroupedKernelCaller() + + +def call_cutedsl( + fn: Any, + inputs: Sequence[Any], + *, + input_descs: Sequence[JaxTensorDesc], + outputs: Sequence[JaxTensorDesc], + workspaces: Sequence[JaxTensorDesc] = (), + static_args: Mapping[str, Any] | None = None, + allow_cuda_graph: bool = True, + use_static_tensors: bool = True, +) -> tuple[Any, ...]: + """Lower a grouped kernel with canonical stream/input/output/workspace order. + + Tensor descriptors retain framework metadata and all CUTLASS lowering + constraints. + """ + + static = dict(static_args or {}) + cluster_shape = static.get("cluster_shape_mn") + margin = int(static.pop("cluster_overlap_margin", 0)) + + input_descs = tuple(input_descs) + output_descs = tuple(outputs) + workspace_descs = tuple(workspaces) + if len(input_descs) != len(inputs): + raise ValueError( + f"Expected {len(inputs)} input descriptors, got {len(input_descs)}" + ) + + if any(0 in desc.shape for desc in output_descs): + results = [] + for desc in output_descs: + metadata = _CALLER._to_shape_dtype_struct(desc) + if desc.init_value is None: + results.append(jnp.empty(metadata.shape, dtype=metadata.dtype)) + else: + results.append( + jnp.full( + metadata.shape, + desc.init_value, + dtype=metadata.dtype, + ) + ) + return tuple(results) + + if cluster_shape is not None: + static["max_active_clusters"] = _CALLER._get_max_active_clusters( + int(cluster_shape[0]) * int(cluster_shape[1]), + overlap_margin=margin, + ) + + def launch(stream: Any, *args: Any) -> None: + fn(stream, *args, **static) + + compute_capability = _CALLER._resolve_compute_capability( + None, + SUPPORTED_COMPUTE_CAPABILITIES, + "grouped GEMM", + ) + return _CALLER._call_kernel( + tuple(inputs), + launch=launch, + input_descs=input_descs, + output_descs=output_descs, + workspace_descs=workspace_descs, + allow_cuda_graph=allow_cuda_graph, + compile_options=compile_options_for_target(compute_capability), + use_static_tensors=use_static_tensors, + ) + + +def block_scale_shape( + rows: int, k: int, batch: int, sf_vec_size: int +) -> tuple[int, ...]: + """Return the public row-major scale shape used by grouped JAX APIs.""" + + return to_public_axes( + _canonical_block_scale_shape(rows, k, batch, sf_vec_size), + BLOCK_SCALE_MODE, + ) + + +def require_grouped_gemm_inputs( + a_tensor: Any, + b_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + *, + max_experts: int = MAX_EXPERTS, + valid_ab_dtypes: Iterable[Any] | None = None, +) -> tuple[int, int, int, int, Any]: + """Validate canonical grouped GEMM matrix metadata.""" + + if valid_ab_dtypes is None: + valid_ab_dtypes = ( + jnp.float4_e2m1fn, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + ) + a_shape = require_array(a_tensor, name="a_tensor", rank=3, dtype=valid_ab_dtypes) + b_shape = require_array(b_tensor, name="b_tensor", rank=3, dtype=as_dtype(a_tensor)) + m, k, a_batch = a_shape + n, b_k, experts = b_shape + if a_batch != 1 or b_k != k: + raise ValueError( + f"Grouped GEMM expects A=(M,K,1) and B=(N,K,L), got {a_shape} and {b_shape}" + ) + if m < 0 or any(value <= 0 for value in (n, k, experts)): + raise ValueError( + "Grouped GEMM requires M >= 0 and positive N, K, and L; " + f"got M={m}, N={n}, K={k}, L={experts}" + ) + if experts > max_experts: + raise ValueError( + f"The number of experts must be at most {max_experts}, got {experts}" + ) + require_array( + padded_offsets, name="padded_offsets", shape=(experts,), dtype=jnp.int32 + ) + require_array( + alpha_tensor, name="alpha_tensor", shape=(experts,), dtype=jnp.float32 + ) + return m, n, k, experts, as_dtype(a_tensor) + + +def require_grouped_block_scales( + sfa_tensor: Any, + sfb_tensor: Any, + *, + m: int, + n: int, + k: int, + experts: int, + sf_vec_size: int, + valid_dtypes: Iterable[Any], +) -> Any: + require_array( + sfa_tensor, + name="sfa_tensor", + shape=block_scale_shape(m, k, 1, sf_vec_size), + dtype=valid_dtypes, + ) + sf_dtype = as_dtype(sfa_tensor) + require_array( + sfb_tensor, + name="sfb_tensor", + shape=block_scale_shape(n, k, experts, sf_vec_size), + dtype=sf_dtype, + ) + return sf_dtype + + +def require_grouped_fp8_scales( + sfa_tensor: Any, + sfb_tensor: Any, + *, + m: int, + n: int, + k: int, + experts: int, + sf_vec_size: int, +) -> Any: + return require_grouped_block_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + valid_dtypes=(jnp.float8_e8m0fnu,), + ) + + +def require_grouped_input_scales( + sfa_tensor: Any, + sfb_tensor: Any, + *, + m: int, + n: int, + k: int, + experts: int, + sf_vec_size: int, + ab_dtype: Any, +) -> Any: + """Validate MXFP8, MXFP4, or NVFP4 scale metadata. + + Native JAX FP4 arrays use logical element shapes. This helper deliberately + does not accept raw ``uint8`` payloads or emulate Torch's packed-storage + reinterpretation. + """ + + if is_fp4_dtype(ab_dtype): + if sf_vec_size not in SF_VEC_SIZES: + raise ValueError( + f"FP4 grouped GEMM requires sf_vec_size in {SF_VEC_SIZES}, got {sf_vec_size}" + ) + sf_dtype = require_grouped_block_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + valid_dtypes=(jnp.float8_e8m0fnu, jnp.float8_e4m3fn), + ) + if sf_dtype == jnp.dtype(jnp.float8_e4m3fn) and sf_vec_size != 16: + raise ValueError("float8_e4m3fn scales require sf_vec_size=16") + return sf_dtype + + if is_fp8_dtype(ab_dtype): + if sf_vec_size != FP8_SF_VEC_SIZE: + raise ValueError( + f"FP8 grouped GEMM requires sf_vec_size={FP8_SF_VEC_SIZE}, got {sf_vec_size}" + ) + return require_grouped_fp8_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ) + + raise ValueError(f"Unsupported grouped GEMM input dtype {as_dtype(ab_dtype)}") + + +def require_grouped_vector( + name: str, tensor: Any, *, length: int, dtype: Any = None +) -> Any: + require_array( + tensor, + name=name, + shape=(length,), + dtype=jnp.float32 if dtype is None else dtype, + ) + return as_dtype(tensor) + + +def require_grouped_probability(name: str, tensor: Any, *, m: int) -> None: + require_array(tensor, name=name, shape=(1, 1, m), dtype=jnp.float32) + + +def require_16_byte_extent(name: str, elements: int, dtype: Any) -> None: + from .._jax.datatypes import jax_to_cudnn_dtype + + require_contiguous_alignment( + name, elements, data_type_bits(jax_to_cudnn_dtype(dtype)) + ) + + +def require_contiguous_alignment(name: str, elements: int, element_bits: int) -> None: + if elements * element_bits % 128: + raise ValueError(f"{name} contiguous extent must span a multiple of 16 bytes") + + +def require_grouped_mma_tiler( + value: Sequence[int], + *, + allowed_m: tuple[int, ...] = (64, 128, 256), + allowed_n: tuple[int, ...] = (128, 256), +) -> tuple[int, int]: + value = tuple(value) + if len(value) != 2 or value[0] not in allowed_m or value[1] not in allowed_n: + raise ValueError(f"Unsupported grouped GEMM mma_tiler_mn {value}") + return value + + +def require_grouped_cluster_shape( + value: Sequence[int], *, mma_tiler_mn: tuple[int, int] +) -> tuple[int, int]: + value = tuple(value) + if len(value) != 2 or any(item <= 0 or item & (item - 1) for item in value): + raise ValueError( + f"cluster_shape_mn entries must be positive powers of two, got {value}" + ) + if any(item > 4 for item in value) or value[0] * value[1] > 16: + raise ValueError(f"cluster_shape_mn product must be at most 16, got {value}") + cta_group = 2 if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else 1 + if value[0] % cta_group: + raise ValueError( + f"cluster_shape_mn[0] must be divisible by {cta_group}, got {value[0]}" + ) + cluster_m = value[0] // cta_group * mma_tiler_mn[0] + if cluster_m not in (128, 256): + raise ValueError( + f"Grouped GEMM cluster M tile must be 128 or 256, got {cluster_m}" + ) + return value + + +def dense_workspace_bytes(use_dynamic_sched: bool) -> int: + return 4 if use_dynamic_sched else 0 + + +def grouped_wgrad_workspace_bytes(expert_cnt: int, input_order: str) -> int: + from .moe_utils import MoEWeightMode, WGradInputOrder, WgradSfTensormapConstructor + + return WgradSfTensormapConstructor.get_workspace_size( + WGradInputOrder(input_order), + MoEWeightMode.DENSE, + expert_cnt, + ) + + +def normalize_wgrad_input_order(input_order: Any) -> Any: + from .moe_utils import WGradInputOrder + + if isinstance(input_order, WGradInputOrder): + return input_order + try: + return WGradInputOrder(input_order) + except ValueError: + choices = ", ".join(order.value for order in WGradInputOrder) + raise ValueError( + f"input_order must be one of {{{choices}}}, got {input_order!r}" + ) from None + + +def ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +__all__ = [ + "ApiBaseJax", + "BLOCK_SCALE_MODE", + "FIX_PAD_SIZE", + "FP8_SF_VEC_SIZE", + "HADAMARD_SIZE", + "GROUPED_BIAS_MODE", + "GROUPED_WORKSPACE_ALIGNMENT", + "MAX_EXPERTS", + "SF_VEC_SIZES", + "TWO_CTA_MMA_TILER_M", + "TupleDict", + "as_dtype", + "as_gemm_tensor_desc", + "block_scale_shape", + "call_cutedsl", + "ceil_div", + "dense_workspace_bytes", + "gemm_a_mode", + "gemm_b_mode", + "gemm_output_mode", + "grouped_wgrad_workspace_bytes", + "is_fp4_dtype", + "is_fp8_dtype", + "is_low_precision_output_dtype", + "make_buffer_desc", + "normalize_wgrad_input_order", + "PROBABILITY_MODE", + "require_16_byte_extent", + "require_array", + "require_contiguous_alignment", + "require_dtype", + "require_grouped_block_scales", + "require_grouped_cluster_shape", + "require_grouped_fp8_scales", + "require_grouped_gemm_inputs", + "require_grouped_input_scales", + "require_grouped_mma_tiler", + "require_grouped_probability", + "require_grouped_vector", + "require_layout", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py index 748246694..b6c91eddd 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dglu/__init__.py @@ -1,12 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmDgluSm100, - grouped_gemm_dglu_wrapper_sm100, -) +from ...common.operation_api import make_operation_api -__all__ = [ - "GroupedGemmDgluSm100", - "grouped_gemm_dglu_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmDgluSm100", "grouped_gemm_dglu_wrapper_sm100")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dglu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_dglu/jax.py new file mode 100644 index 000000000..811888516 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_dglu/jax.py @@ -0,0 +1,750 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for dense-weight grouped GEMM + dGLU on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + BLOCK_SCALE_MODE, + make_buffer_desc, + FIX_PAD_SIZE, + MAX_EXPERTS, + TWO_CTA_MMA_TILER_M, + TupleDict, + GROUPED_WORKSPACE_ALIGNMENT, + PROBABILITY_MODE, + as_dtype, + as_gemm_tensor_desc, + block_scale_shape, + call_cutedsl, + dense_workspace_bytes, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, + is_fp4_dtype, + is_fp8_dtype, + require_16_byte_extent, + require_dtype, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_input_scales, + require_grouped_mma_tiler, + require_grouped_probability, + require_grouped_vector, + require_layout, +) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + discrete_col_sfd: bool, + expert_cnt: int, + generate_dbias: bool, + generate_sfd: bool, + has_amax: bool, + use_dynamic_sched: bool, + act_func: str, + epilogue_op: str, + linear_offset: float, + geglu_alpha: float, + glu_clamp_max: float, + glu_clamp_min: float, + max_active_clusters: int, +): + import cutlass + import cutlass.cute as cute + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from ..moe_utils import MoEWeightMode + from .moe_blockscaled_grouped_gemm_dglu_dbias import ( + BlockScaledMoEGroupedGemmDgluDbiasKernel, + ) + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a = take() + b = take() + c = take() + sfa = take() + sfb = take() + padded_offsets = take() + alpha = take() + beta = take() + prob = take() + norm_const = take() if generate_sfd else None + + d_row = take() + d_col = take() + dprob = take() + dbias = take() if generate_dbias else None + amax = take() if has_amax else None + sfd_row = take() if generate_sfd else None + sfd_col = take() if generate_sfd else None + workspace = take() + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped GEMM argument count: consumed {arg_idx}, received {len(args)}" + ) + + if epilogue_op == "relu": + + def epilogue(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) + + elif epilogue_op == "srelu": + + def epilogue(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) ** 2 + + else: + + def epilogue(x): + return x + + kernel = BlockScaledMoEGroupedGemmDgluDbiasKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + discrete_col_sfd=discrete_col_sfd, + expert_cnt=expert_cnt, + weight_mode=MoEWeightMode.DENSE, + use_dynamic_sched=use_dynamic_sched, + act_func=act_func, + ) + kernel( + a, + b, + sfb, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int64(0), + OperandMajorMode.K, + workspace.iterator, + c, + d_row, + d_col, + sfa, + sfd_row, + sfd_col, + amax, + norm_const, + padded_offsets, + alpha, + beta, + prob, + dprob, + dbias, + max_active_clusters, + stream, + epilogue_op=epilogue, + linear_offset=linear_offset, + geglu_alpha=geglu_alpha, + glu_clamp_max=glu_clamp_max, + glu_clamp_min=glu_clamp_min, + ) + + +def _grouped_gemm_dglu_impl( + a_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Any, + prob_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + generate_dbias: bool = False, + norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "dswiglu", + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + epilogue_op: Optional[str] = None, + use_dynamic_sched: bool = False, + cluster_overlap_margin: int = 0, + *, + b_layout: str = "LNK", + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute a dense-weight block-scaled grouped GEMM with fused dGLU. + + ``dprob_tensor`` and optional ``dbias_tensor`` are fresh, zero-initialized + JAX results. Native FP4 inputs use logical array shapes; raw-byte + reinterpretation remains outside the JAX API. + """ + + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) + c_desc = as_gemm_tensor_desc("c_tensor", c_tensor, mode=output_mode) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + ) + if n % 32: + raise ValueError(f"b_tensor N must be divisible by 32 for dGLU, got {n}") + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_input_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ab_dtype=ab_dtype, + ) + + output_n = 2 * n + expected_c_shape = (m, output_n, 1) + if c_desc.shape != expected_c_shape: + raise ValueError( + f"c_tensor must have canonical shape {expected_c_shape}, got {c_desc.shape}" + ) + require_dtype( + c_desc, + ( + jnp.float32, + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + ), + name="c_tensor.dtype", + ) + c_dtype = as_dtype(c_desc) + fp8_dtypes = { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + } + if vector_f32 and c_dtype in fp8_dtypes: + raise ValueError("vector_f32 does not support an FP8 c_tensor") + require_grouped_probability("prob_tensor", prob_tensor, m=m) + require_grouped_vector("beta_tensor", beta_tensor, length=experts) + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + if is_fp4_dtype(ab_dtype): + valid_d_dtypes = (jnp.float16, jnp.bfloat16, jnp.float32) + default_d_dtype = jnp.bfloat16 + else: + valid_d_dtypes = (jnp.float8_e4m3fn, jnp.float8_e5m2) + default_d_dtype = ab_dtype + d_dtype = require_dtype( + d_dtype, valid_d_dtypes, name="d_dtype", default=default_d_dtype + ) + generate_sfd = is_fp8_dtype(ab_dtype) + has_amax = d_dtype in {jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)} + if generate_sfd: + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required for FP8 inputs") + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + discrete_col_sfd = False + if is_fp4_dtype(ab_dtype) and b_layout != "LNK": + raise ValueError("Native FP4 B must use the K-major LNK layout") + if ( + is_fp4_dtype(ab_dtype) + and sf_vec_size == 16 + and d_dtype == jnp.dtype(jnp.float32) + and not generate_dbias + ): + raise NotImplementedError( + "FP4 with sf_vec_size=16 and float32 D requires generate_dbias=True" + ) + if act_func not in ("dswiglu", "dgeglu"): + raise ValueError(f"act_func must be 'dswiglu' or 'dgeglu', got {act_func!r}") + if linear_offset is None: + linear_offset = 1.0 if act_func == "dgeglu" else 0.0 + normalized_epilogue = ( + "identity" if epilogue_op in (None, "none", "identity") else epilogue_op + ) + if normalized_epilogue not in ("identity", "relu", "srelu"): + raise ValueError( + f"epilogue_op must be None, 'none', 'identity', 'relu', or 'srelu', got {epilogue_op!r}" + ) + + mma_tiler_mn = require_grouped_mma_tiler( + mma_tiler_mn, allowed_m=(128, 256), allowed_n=(256,) + ) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + + require_16_byte_extent("a_tensor", k, ab_dtype) + require_16_byte_extent("b_tensor", n if b_layout == "LKN" else k, ab_dtype) + require_16_byte_extent("c_tensor", output_n, c_dtype) + require_16_byte_extent("d_tensor", output_n, d_dtype) + + if _validate_only: + return { + "acc_dtype": acc_dtype, + "d_dtype": d_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "linear_offset": linear_offset, + "epilogue_op": normalized_epilogue, + } + + outputs = [ + make_buffer_desc( + "d_row_tensor", (1, m, output_n), d_dtype, mode=output_mode + ), + make_buffer_desc( + "d_col_tensor", (1, m, output_n), d_dtype, mode=output_mode + ), + make_buffer_desc( + "dprob_tensor", + (1, 1, m), + jnp.float32, + mode=PROBABILITY_MODE, + init_value=0.0, + ), + ] + if generate_dbias: + outputs.append( + make_buffer_desc( + "dbias_tensor", + (experts, output_n, 1), + jnp.bfloat16, + init_value=0.0, + ) + ) + if has_amax: + outputs.append( + make_buffer_desc( + "amax_tensor", + (experts, 2, 1), + jnp.float32, + init_value=-float("inf"), + ) + ) + if generate_sfd: + outputs.extend( + ( + make_buffer_desc( + "sfd_row_tensor", + block_scale_shape(m, output_n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + mode=BLOCK_SCALE_MODE, + ), + make_buffer_desc( + "sfd_col_tensor", + block_scale_shape(output_n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + mode=BLOCK_SCALE_MODE, + ), + ) + ) + + workspace_bytes = max(dense_workspace_bytes(bool(use_dynamic_sched)), 1) + inputs = [ + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + ] + input_descs = [ + a_desc, + b_desc, + c_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + as_gemm_tensor_desc("beta_tensor", beta_tensor), + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE), + ] + if generate_sfd: + inputs.append(norm_const_tensor) + input_descs.append(as_gemm_tensor_desc("norm_const_tensor", norm_const_tensor)) + results = call_cutedsl( + _launch, + inputs, + input_descs=input_descs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "discrete_col_sfd": bool(discrete_col_sfd), + "expert_cnt": experts, + "generate_dbias": bool(generate_dbias), + "generate_sfd": generate_sfd, + "has_amax": has_amax, + "use_dynamic_sched": bool(use_dynamic_sched), + "act_func": act_func, + "epilogue_op": normalized_epilogue, + "linear_offset": float(linear_offset), + "geglu_alpha": float(geglu_alpha), + "glu_clamp_max": float(glu_clamp_max), + "glu_clamp_min": float(glu_clamp_min), + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + workspaces=( + make_buffer_desc( + "workspace", + (workspace_bytes,), + jnp.uint8, + ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT, + ), + ), + ) + result_idx = 0 + d_row_tensor = results[result_idx] + result_idx += 1 + d_col_tensor = results[result_idx] + result_idx += 1 + dprob_tensor = results[result_idx] + result_idx += 1 + dbias_tensor = results[result_idx] if generate_dbias else None + result_idx += int(bool(generate_dbias)) + amax_tensor = results[result_idx] if has_amax else None + result_idx += int(has_amax) + sfd_row_tensor = results[result_idx] if generate_sfd else None + sfd_col_tensor = results[result_idx + 1] if generate_sfd else None + return TupleDict( + d_row_tensor=d_row_tensor, + d_col_tensor=d_col_tensor, + dprob_tensor=dprob_tensor, + dbias_tensor=dbias_tensor, + amax_tensor=amax_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + ) + + +class GroupedGemmDgluSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM + dGLU.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_c_tensor: Any, + sample_sfa_tensor: Any, + sample_padded_offsets: Any, + sample_alpha_tensor: Any, + sample_beta_tensor: Any, + sample_prob_tensor: Any, + sample_b_tensor: Any, + sample_sfb_tensor: Any, + generate_dbias: bool = False, + sample_norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "dswiglu", + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + epilogue_op: Optional[str] = None, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", + ) -> None: + super().__init__() + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, mode=a_mode, name="sample_a_tensor" + ), + "c_tensor": self.make_tensor_desc( + sample_c_tensor, mode=output_mode, name="sample_c_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, 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, + mode=PROBABILITY_MODE, + name="sample_prob_tensor", + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, mode=b_mode, name="sample_b_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfb_tensor" + ), + "norm_const_tensor": self.make_optional_tensor_desc( + sample_norm_const_tensor, name="sample_norm_const_tensor" + ), + } + self._config = { + "generate_dbias": generate_dbias, + "acc_dtype": self.as_optional_dtype(acc_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "discrete_col_sfd": discrete_col_sfd, + "act_func": act_func, + "linear_offset": linear_offset, + "geglu_alpha": geglu_alpha, + "glu_clamp_max": glu_clamp_max, + "glu_clamp_min": glu_clamp_min, + "epilogue_op": epilogue_op, + "use_dynamic_sched": use_dynamic_sched, + "b_layout": b_layout, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_dglu_impl( + self._sample_descs["a_tensor"], + self._sample_descs["c_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["beta_tensor"], + self._sample_descs["prob_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["sfb_tensor"], + norm_const_tensor=self._sample_descs["norm_const_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Any, + prob_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + norm_const_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + c_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + b_tensor, + sfb_tensor, + norm_const_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Any, + prob_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + norm_const_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "c_tensor": c_tensor, + "sfa_tensor": sfa_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "beta_tensor": beta_tensor, + "prob_tensor": prob_tensor, + "b_tensor": b_tensor, + "sfb_tensor": sfb_tensor, + "norm_const_tensor": norm_const_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_dglu_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "generate_dbias", + "acc_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "act_func", + "linear_offset", + "geglu_alpha", + "glu_clamp_max", + "glu_clamp_min", + "epilogue_op", + "use_dynamic_sched", + "b_layout", + ), +) +def grouped_gemm_dglu_wrapper_sm100( + a_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Any, + prob_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + generate_dbias: bool = False, + norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "dswiglu", + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + epilogue_op: Optional[str] = None, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", +) -> TupleDict: + """Compute an MXFP8 dense-weight grouped GEMM with fused dGLU.""" + + op = GroupedGemmDgluSm100( + a_tensor, + c_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + b_tensor, + sfb_tensor, + generate_dbias=generate_dbias, + sample_norm_const_tensor=norm_const_tensor, + acc_dtype=acc_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + act_func=act_func, + linear_offset=linear_offset, + geglu_alpha=geglu_alpha, + glu_clamp_max=glu_clamp_max, + glu_clamp_min=glu_clamp_min, + epilogue_op=epilogue_op, + use_dynamic_sched=use_dynamic_sched, + b_layout=b_layout, + ) + return op( + a_tensor, + c_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + b_tensor, + sfb_tensor, + norm_const_tensor, + ) + + +__all__ = [ + "GroupedGemmDgluSm100", + "grouped_gemm_dglu_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py index e15137794..af74a57ab 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/__init__.py @@ -1,12 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmDsreluSm100, - grouped_gemm_dsrelu_wrapper_sm100, -) +from ...common.operation_api import make_operation_api -__all__ = [ - "GroupedGemmDsreluSm100", - "grouped_gemm_dsrelu_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmDsreluSm100", "grouped_gemm_dsrelu_wrapper_sm100")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/jax.py new file mode 100644 index 000000000..9594646d9 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_dsrelu/jax.py @@ -0,0 +1,668 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for dense-weight grouped GEMM + dSReLU on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + BLOCK_SCALE_MODE, + make_buffer_desc, + FIX_PAD_SIZE, + MAX_EXPERTS, + TWO_CTA_MMA_TILER_M, + TupleDict, + GROUPED_WORKSPACE_ALIGNMENT, + PROBABILITY_MODE, + as_dtype, + as_gemm_tensor_desc, + block_scale_shape, + call_cutedsl, + dense_workspace_bytes, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, + is_fp4_dtype, + is_fp8_dtype, + require_16_byte_extent, + require_dtype, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_input_scales, + require_grouped_mma_tiler, + require_grouped_probability, + require_grouped_vector, + require_layout, +) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + discrete_col_sfd: bool, + expert_cnt: int, + generate_dbias: bool, + generate_sfd: bool, + has_amax: bool, + use_dynamic_sched: bool, + use_dsrelu_reuse: bool, + max_active_clusters: int, +): + import cutlass + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from ..moe_utils import MoEWeightMode + from .moe_blockscaled_grouped_gemm_dsrelu_quant import ( + BlockScaledMoEGroupedGemmQuantBwdKernel, + EpilogueType, + ) + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a = take() + b = take() + c = take() + sfa = take() + sfb = take() + padded_offsets = take() + alpha = take() + prob = take() + norm_const = take() if generate_sfd else None + + d_row = take() + d_col = take() + d_srelu = take() + dprob = take() + dbias = take() if generate_dbias else None + amax = take() if has_amax else None + sfd_row = take() if generate_sfd else None + sfd_col = take() if generate_sfd else None + sfd_col_d_srelu = take() if generate_sfd else None + workspace = take() + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped GEMM argument count: consumed {arg_idx}, received {len(args)}" + ) + + kernel = BlockScaledMoEGroupedGemmQuantBwdKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + generate_sfd=generate_sfd, + discrete_col_sfd=discrete_col_sfd, + expert_cnt=expert_cnt, + weight_mode=MoEWeightMode.DENSE, + use_dynamic_sched=use_dynamic_sched, + epilogue_type=EpilogueType.DSRELU.value, + generate_dbias=generate_dbias, + generate_d_srelu=True, + use_dsrelu_reuse=use_dsrelu_reuse, + ) + kernel( + a, + b, + sfb, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int64(0), + OperandMajorMode.K, + workspace.iterator, + c, + d_row, + d_col, + sfa, + sfd_row, + sfd_col, + amax, + norm_const, + padded_offsets, + alpha, + prob, + dprob, + dbias, + d_srelu, + sfd_col_d_srelu, + max_active_clusters, + stream, + ) + + +def _grouped_gemm_dsrelu_impl( + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + generate_dbias: bool = False, + norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + use_dsrelu_reuse: bool = False, + cluster_overlap_margin: int = 0, + *, + b_layout: str = "LNK", + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute a dense-weight block-scaled grouped GEMM with fused dSReLU. + + The mutable Torch outputs are represented as functional JAX results. + ``dprob_tensor`` and optional ``dbias_tensor`` start at zero before the + kernel's atomic reductions. Native FP4 arrays use logical JAX shapes. + """ + + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) + c_desc = as_gemm_tensor_desc("c_tensor", c_tensor, mode=output_mode) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + ) + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_input_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ab_dtype=ab_dtype, + ) + + expected_c_shape = (m, n, 1) + if c_desc.shape != expected_c_shape: + raise ValueError( + f"c_tensor must have canonical shape {expected_c_shape}, got {c_desc.shape}" + ) + require_dtype( + c_desc, + ( + jnp.float32, + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + ), + name="c_tensor.dtype", + ) + c_dtype = as_dtype(c_desc) + fp8_dtypes = { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + } + if vector_f32 and c_dtype in fp8_dtypes: + raise ValueError("vector_f32 does not support an FP8 c_tensor") + require_grouped_probability("prob_tensor", prob_tensor, m=m) + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + if is_fp4_dtype(ab_dtype): + valid_d_dtypes = (jnp.float16, jnp.bfloat16, jnp.float32) + default_d_dtype = jnp.bfloat16 + else: + valid_d_dtypes = (jnp.float8_e4m3fn, jnp.float8_e5m2) + default_d_dtype = ab_dtype + d_dtype = require_dtype( + d_dtype, valid_d_dtypes, name="d_dtype", default=default_d_dtype + ) + generate_sfd = is_fp8_dtype(ab_dtype) + has_amax = d_dtype in {jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)} + if generate_sfd: + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required for FP8 inputs") + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + discrete_col_sfd = False + if is_fp4_dtype(ab_dtype) and b_layout != "LNK": + raise ValueError("Native FP4 B must use the K-major LNK layout") + if ( + is_fp4_dtype(ab_dtype) + and sf_vec_size == 16 + and d_dtype == jnp.dtype(jnp.float32) + and not generate_dbias + ): + raise NotImplementedError( + "FP4 with sf_vec_size=16 and float32 D requires generate_dbias=True" + ) + mma_tiler_mn = require_grouped_mma_tiler( + mma_tiler_mn, allowed_m=(128, 256), allowed_n=(256,) + ) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + + require_16_byte_extent("a_tensor", k, ab_dtype) + require_16_byte_extent("b_tensor", n if b_layout == "LKN" else k, ab_dtype) + require_16_byte_extent("c_tensor", n, c_dtype) + require_16_byte_extent("d_tensor", n, d_dtype) + + if _validate_only: + return { + "acc_dtype": acc_dtype, + "d_dtype": d_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + } + + outputs = [ + make_buffer_desc("d_row_tensor", (1, m, n), d_dtype, mode=output_mode), + make_buffer_desc("d_col_tensor", (1, m, n), d_dtype, mode=output_mode), + make_buffer_desc("d_srelu_tensor", (1, m, n), d_dtype, mode=output_mode), + make_buffer_desc( + "dprob_tensor", + (1, 1, m), + jnp.float32, + mode=PROBABILITY_MODE, + init_value=0.0, + ), + ] + if generate_dbias: + outputs.append( + make_buffer_desc( + "dbias_tensor", + (experts, n, 1), + jnp.bfloat16, + init_value=0.0, + ) + ) + if has_amax: + outputs.append( + make_buffer_desc( + "amax_tensor", (experts, 1), jnp.float32, init_value=-float("inf") + ) + ) + if generate_sfd: + outputs.extend( + ( + make_buffer_desc( + "sfd_row_tensor", + block_scale_shape(m, n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + mode=BLOCK_SCALE_MODE, + ), + make_buffer_desc( + "sfd_col_tensor", + block_scale_shape(n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + mode=BLOCK_SCALE_MODE, + ), + make_buffer_desc( + "sfd_col_d_srelu_tensor", + block_scale_shape(n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + mode=BLOCK_SCALE_MODE, + ), + ) + ) + + workspace_bytes = max(dense_workspace_bytes(bool(use_dynamic_sched)), 1) + inputs = [ + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + ] + input_descs = [ + a_desc, + b_desc, + c_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE), + ] + if generate_sfd: + inputs.append(norm_const_tensor) + input_descs.append(as_gemm_tensor_desc("norm_const_tensor", norm_const_tensor)) + results = call_cutedsl( + _launch, + inputs, + input_descs=input_descs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "discrete_col_sfd": bool(discrete_col_sfd), + "expert_cnt": experts, + "generate_dbias": bool(generate_dbias), + "generate_sfd": generate_sfd, + "has_amax": has_amax, + "use_dynamic_sched": bool(use_dynamic_sched), + "use_dsrelu_reuse": bool(use_dsrelu_reuse), + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + workspaces=( + make_buffer_desc( + "workspace", + (workspace_bytes,), + jnp.uint8, + ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT, + ), + ), + ) + result_idx = 0 + d_row_tensor = results[result_idx] + result_idx += 1 + d_col_tensor = results[result_idx] + result_idx += 1 + d_srelu_tensor = results[result_idx] + result_idx += 1 + dprob_tensor = results[result_idx] + result_idx += 1 + dbias_tensor = results[result_idx] if generate_dbias else None + result_idx += int(bool(generate_dbias)) + amax_tensor = results[result_idx] if has_amax else None + result_idx += int(has_amax) + sfd_row_tensor = results[result_idx] if generate_sfd else None + sfd_col_tensor = results[result_idx + 1] if generate_sfd else None + sfd_col_d_srelu_tensor = results[result_idx + 2] if generate_sfd else None + return TupleDict( + d_row_tensor=d_row_tensor, + d_col_tensor=d_col_tensor, + d_srelu_tensor=d_srelu_tensor, + dprob_tensor=dprob_tensor, + dbias_tensor=dbias_tensor, + amax_tensor=amax_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + sfd_col_d_srelu_tensor=sfd_col_d_srelu_tensor, + ) + + +class GroupedGemmDsreluSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM + dSReLU.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_b_tensor: Any, + sample_c_tensor: Any, + sample_sfa_tensor: Any, + sample_sfb_tensor: Any, + sample_padded_offsets: Any, + sample_alpha_tensor: Any, + sample_prob_tensor: Any, + generate_dbias: bool = False, + sample_norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + use_dsrelu_reuse: bool = False, + *, + b_layout: str = "LNK", + ) -> None: + super().__init__() + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, mode=a_mode, name="sample_a_tensor" + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, mode=b_mode, name="sample_b_tensor" + ), + "c_tensor": self.make_tensor_desc( + sample_c_tensor, mode=output_mode, name="sample_c_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfa_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, 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, + mode=PROBABILITY_MODE, + name="sample_prob_tensor", + ), + "norm_const_tensor": self.make_optional_tensor_desc( + sample_norm_const_tensor, name="sample_norm_const_tensor" + ), + } + self._config = { + "generate_dbias": generate_dbias, + "acc_dtype": self.as_optional_dtype(acc_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "discrete_col_sfd": discrete_col_sfd, + "use_dynamic_sched": use_dynamic_sched, + "use_dsrelu_reuse": use_dsrelu_reuse, + "b_layout": b_layout, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_dsrelu_impl( + self._sample_descs["a_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["c_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["sfb_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["prob_tensor"], + norm_const_tensor=self._sample_descs["norm_const_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + norm_const_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + norm_const_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + norm_const_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "b_tensor": b_tensor, + "c_tensor": c_tensor, + "sfa_tensor": sfa_tensor, + "sfb_tensor": sfb_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "prob_tensor": prob_tensor, + "norm_const_tensor": norm_const_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_dsrelu_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "generate_dbias", + "acc_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "use_dynamic_sched", + "use_dsrelu_reuse", + "b_layout", + ), +) +def grouped_gemm_dsrelu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + generate_dbias: bool = False, + norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + use_dsrelu_reuse: bool = False, + *, + b_layout: str = "LNK", +) -> TupleDict: + """Compute an MXFP8 dense-weight grouped GEMM with fused dSReLU.""" + + op = GroupedGemmDsreluSm100( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + generate_dbias=generate_dbias, + sample_norm_const_tensor=norm_const_tensor, + acc_dtype=acc_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + use_dynamic_sched=use_dynamic_sched, + use_dsrelu_reuse=use_dsrelu_reuse, + b_layout=b_layout, + ) + return op( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + norm_const_tensor, + ) + + +__all__ = [ + "GroupedGemmDsreluSm100", + "grouped_gemm_dsrelu_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py index f0261951a..541bb1986 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/__init__.py @@ -1,12 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmDswigluSm100, - grouped_gemm_dswiglu_wrapper_sm100, -) +from ...common.operation_api import make_operation_api -__all__ = [ - "GroupedGemmDswigluSm100", - "grouped_gemm_dswiglu_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmDswigluSm100", "grouped_gemm_dswiglu_wrapper_sm100")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/jax.py new file mode 100644 index 000000000..92d2cacf3 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_dswiglu/jax.py @@ -0,0 +1,623 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for contiguous grouped GEMM + dSwiGLU on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + BLOCK_SCALE_MODE, + make_buffer_desc, + FIX_PAD_SIZE, + MAX_EXPERTS, + TWO_CTA_MMA_TILER_M, + TupleDict, + PROBABILITY_MODE, + as_dtype, + as_gemm_tensor_desc, + block_scale_shape, + call_cutedsl, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, + is_fp4_dtype, + is_fp8_dtype, + require_16_byte_extent, + require_dtype, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_input_scales, + require_grouped_mma_tiler, + require_grouped_probability, + require_grouped_vector, + require_layout, +) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + discrete_col_sfd: bool, + expert_cnt: int, + epilogue_op: str, + generate_sfd: bool, + has_amax: bool, + max_active_clusters: int, +): + import cutlass.cute as cute + from cutlass.jax import jax_to_cutlass_dtype + + from .grouped_gemm_dswiglu_quant import BlockScaledContiguousGroupedGemmKernel + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a, b, c, sfa, sfb, padded_offsets, alpha, beta, prob = (take() for _ in range(9)) + norm_const = take() if generate_sfd else None + d_row, d_col, dprob = (take() for _ in range(3)) + amax = take() if has_amax else None + sfd_row = take() if generate_sfd else None + sfd_col = take() if generate_sfd else None + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped dSwiGLU argument count: consumed {arg_idx}, received {len(args)}" + ) + + if epilogue_op == "relu": + + def epilogue(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) + + elif epilogue_op == "srelu": + + def epilogue(x): + return cute.where(x > 0, x, cute.full_like(x, 0)) ** 2 + + else: + + def epilogue(x): + return x + + kernel = BlockScaledContiguousGroupedGemmKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + discrete_col_sfd=discrete_col_sfd, + expert_cnt=expert_cnt, + use_mono_increase_expert_idx=True, + ) + kernel( + a, + b, + c, + d_row, + d_col, + sfa, + sfb, + sfd_row, + sfd_col, + amax, + norm_const, + padded_offsets, + alpha, + beta, + prob, + dprob, + max_active_clusters, + stream, + epilogue_op=epilogue, + ) + + +def _grouped_gemm_dswiglu_impl( + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Optional[Any], + prob_tensor: Any, + norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + epilogue_op: Optional[str] = None, + b_layout: str = "LNK", + cluster_overlap_margin: int = 0, + *, + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute the contiguous block-scaled grouped dSwiGLU fusion. + + ``dprob_tensor`` is modeled as a fresh zero-initialized JAX result instead + of a caller-owned mutable buffer. If ``beta_tensor`` is ``None``, a vector + of ones is supplied, matching the convenience behavior of the Torch API. + Native FP4 inputs use logical JAX shapes; raw-byte reinterpretation is not + part of this adapter. Configuration values are static under JAX tracing. + """ + + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + a_mode = gemm_a_mode("LMK") + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) + c_desc = as_gemm_tensor_desc("c_tensor", c_tensor, mode=output_mode) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + ) + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_input_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ab_dtype=ab_dtype, + ) + + output_n = 2 * n + expected_c_shape = (m, output_n, 1) + if c_desc.shape != expected_c_shape: + raise ValueError( + f"c_tensor must have canonical shape {expected_c_shape}, got {c_desc.shape}" + ) + require_dtype( + c_desc, + ( + jnp.float32, + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + ), + name="c_tensor.dtype", + ) + c_dtype = as_dtype(c_desc) + if vector_f32 and c_dtype in { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + }: + raise ValueError("vector_f32 does not support an FP8 c_tensor") + require_grouped_probability("prob_tensor", prob_tensor, m=m) + if beta_tensor is None: + if not _validate_only: + beta_tensor = jnp.ones((experts,), dtype=jnp.float32) + else: + require_grouped_vector("beta_tensor", beta_tensor, length=experts) + + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + if is_fp4_dtype(ab_dtype): + valid_d_dtypes = (jnp.float16, jnp.bfloat16, jnp.float32) + default_d_dtype = jnp.bfloat16 + else: + valid_d_dtypes = (jnp.float8_e4m3fn, jnp.float8_e5m2) + default_d_dtype = ab_dtype + d_dtype = require_dtype( + d_dtype, valid_d_dtypes, name="d_dtype", default=default_d_dtype + ) + generate_sfd = is_fp8_dtype(ab_dtype) + has_amax = d_dtype in {jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)} + if generate_sfd: + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required for FP8 inputs") + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + discrete_col_sfd = False + if is_fp4_dtype(ab_dtype) and b_layout != "LNK": + raise ValueError("Native FP4 B must use the K-major LNK layout") + normalized_epilogue = ( + "identity" if epilogue_op in (None, "none", "identity") else epilogue_op + ) + if normalized_epilogue not in ("identity", "relu", "srelu"): + raise ValueError( + f"epilogue_op must be None, 'none', 'identity', 'relu', or 'srelu', got {epilogue_op!r}" + ) + + mma_tiler_mn = require_grouped_mma_tiler(mma_tiler_mn) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + + require_16_byte_extent("a_tensor", k, ab_dtype) + require_16_byte_extent("b_tensor", k, ab_dtype) + require_16_byte_extent("c_tensor", output_n, c_dtype) + require_16_byte_extent("d_tensor", output_n, d_dtype) + + if _validate_only: + return { + "acc_dtype": acc_dtype, + "d_dtype": d_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "epilogue_op": normalized_epilogue, + } + + inputs = [ + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + ] + input_descs = [ + a_desc, + b_desc, + c_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + as_gemm_tensor_desc("beta_tensor", beta_tensor), + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE), + ] + if generate_sfd: + inputs.append(norm_const_tensor) + input_descs.append(as_gemm_tensor_desc("norm_const_tensor", norm_const_tensor)) + + outputs = [ + make_buffer_desc( + "d_row_tensor", (1, m, output_n), d_dtype, mode=output_mode + ), + make_buffer_desc( + "d_col_tensor", (1, m, output_n), d_dtype, mode=output_mode + ), + make_buffer_desc( + "dprob_tensor", + (1, 1, m), + jnp.float32, + mode=PROBABILITY_MODE, + init_value=0.0, + ), + ] + if has_amax: + outputs.append( + make_buffer_desc( + "amax_tensor", + (experts, 2, 1), + jnp.float32, + init_value=-float("inf"), + ) + ) + if generate_sfd: + outputs.extend( + ( + make_buffer_desc( + "sfd_row_tensor", + block_scale_shape(m, output_n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + mode=BLOCK_SCALE_MODE, + ), + make_buffer_desc( + "sfd_col_tensor", + block_scale_shape(output_n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + mode=BLOCK_SCALE_MODE, + ), + ) + ) + + results = call_cutedsl( + _launch, + inputs, + input_descs=input_descs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "discrete_col_sfd": bool(discrete_col_sfd), + "expert_cnt": experts, + "epilogue_op": normalized_epilogue, + "generate_sfd": generate_sfd, + "has_amax": has_amax, + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + ) + result_idx = 3 + d_row_tensor, d_col_tensor, dprob_tensor = results[:3] + amax_tensor = results[result_idx] if has_amax else None + result_idx += int(has_amax) + sfd_row_tensor = results[result_idx] if generate_sfd else None + sfd_col_tensor = results[result_idx + 1] if generate_sfd else None + return TupleDict( + d_row_tensor=d_row_tensor, + d_col_tensor=d_col_tensor, + dprob_tensor=dprob_tensor, + amax_tensor=amax_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + ) + + +class GroupedGemmDswigluSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM + dSwiGLU.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_b_tensor: Any, + sample_c_tensor: Any, + sample_sfa_tensor: Any, + sample_sfb_tensor: Any, + sample_padded_offsets: Any, + sample_alpha_tensor: Any, + sample_beta_tensor: Optional[Any], + sample_prob_tensor: Any, + sample_norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + epilogue_op: Optional[str] = None, + b_layout: str = "LNK", + ) -> None: + super().__init__() + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + a_mode = gemm_a_mode("LMK") + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, mode=a_mode, name="sample_a_tensor" + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, mode=b_mode, name="sample_b_tensor" + ), + "c_tensor": self.make_tensor_desc( + sample_c_tensor, mode=output_mode, name="sample_c_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfa_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, 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, + mode=PROBABILITY_MODE, + name="sample_prob_tensor", + ), + "norm_const_tensor": self.make_optional_tensor_desc( + sample_norm_const_tensor, name="sample_norm_const_tensor" + ), + } + self._config = { + "acc_dtype": self.as_optional_dtype(acc_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "discrete_col_sfd": discrete_col_sfd, + "epilogue_op": epilogue_op, + "b_layout": b_layout, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_dswiglu_impl( + self._sample_descs["a_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["c_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["sfb_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["beta_tensor"], + self._sample_descs["prob_tensor"], + self._sample_descs["norm_const_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Optional[Any], + prob_tensor: Any, + norm_const_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + norm_const_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Optional[Any], + prob_tensor: Any, + norm_const_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "b_tensor": b_tensor, + "c_tensor": c_tensor, + "sfa_tensor": sfa_tensor, + "sfb_tensor": sfb_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "beta_tensor": beta_tensor, + "prob_tensor": prob_tensor, + "norm_const_tensor": norm_const_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_dswiglu_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "acc_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "epilogue_op", + "b_layout", + ), +) +def grouped_gemm_dswiglu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + c_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + beta_tensor: Optional[Any], + prob_tensor: Any, + norm_const_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + epilogue_op: Optional[str] = None, + b_layout: str = "LNK", +) -> TupleDict: + """Compute the contiguous block-scaled grouped dSwiGLU fusion.""" + + op = GroupedGemmDswigluSm100( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + norm_const_tensor, + acc_dtype=acc_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + epilogue_op=epilogue_op, + b_layout=b_layout, + ) + return op( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + norm_const_tensor, + ) + + +__all__ = [ + "GroupedGemmDswigluSm100", + "grouped_gemm_dswiglu_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_glu/__init__.py index 177db7f07..46a787dc2 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu/__init__.py @@ -1,12 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmGluSm100, - grouped_gemm_glu_wrapper_sm100, -) +from ...common.operation_api import make_operation_api -__all__ = [ - "GroupedGemmGluSm100", - "grouped_gemm_glu_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmGluSm100", "grouped_gemm_glu_wrapper_sm100")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_glu/jax.py new file mode 100644 index 000000000..45e9c3e46 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu/jax.py @@ -0,0 +1,696 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for dense-weight grouped GEMM + GLU on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + BLOCK_SCALE_MODE, + make_buffer_desc, + FIX_PAD_SIZE, + MAX_EXPERTS, + TWO_CTA_MMA_TILER_M, + TupleDict, + GROUPED_BIAS_MODE, + GROUPED_WORKSPACE_ALIGNMENT, + PROBABILITY_MODE, + as_gemm_tensor_desc, + block_scale_shape, + call_cutedsl, + dense_workspace_bytes, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, + is_fp4_dtype, + is_fp8_dtype, + require_16_byte_extent, + require_array, + require_dtype, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_input_scales, + require_grouped_mma_tiler, + require_grouped_probability, + require_grouped_vector, + require_layout, +) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + discrete_col_sfd: bool, + expert_cnt: int, + has_prob: bool, + has_bias: bool, + generate_sfd: bool, + has_amax: bool, + use_dynamic_sched: bool, + act_func: str, + linear_offset: float, + geglu_alpha: float, + glu_clamp_max: float, + glu_clamp_min: float, + max_active_clusters: int, +): + import cutlass + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from ..moe_utils import MoEWeightMode + from .moe_blockscaled_grouped_gemm_glu_bias import ( + BlockScaledMoEGroupedGemmGluBiasKernel, + ) + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a = take() + b = take() + sfa = take() + sfb = take() + padded_offsets = take() + alpha = take() + prob = take() if has_prob else None + bias = take() if has_bias else None + norm_const = take() if generate_sfd else None + + c = take() + d = take() + d_col = take() + amax = take() if has_amax else None + sfd_row = take() if generate_sfd else None + sfd_col = take() if generate_sfd else None + workspace = take() + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped GEMM argument count: consumed {arg_idx}, received {len(args)}" + ) + + kernel = BlockScaledMoEGroupedGemmGluBiasKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + generate_sfd=generate_sfd, + discrete_col_sfd=discrete_col_sfd, + expert_cnt=expert_cnt, + weight_mode=MoEWeightMode.DENSE, + use_dynamic_sched=use_dynamic_sched, + act_func=act_func, + enable_bias=has_bias, + ) + kernel( + a, + b, + sfb, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int64(0), + OperandMajorMode.K, + workspace.iterator, + c, + d, + d_col, + sfa, + sfd_row, + sfd_col, + amax, + norm_const, + padded_offsets, + alpha, + prob, + bias, + max_active_clusters, + stream, + linear_offset=cutlass.Float32(linear_offset), + geglu_alpha=cutlass.Float32(geglu_alpha), + glu_clamp_max=cutlass.Float32(glu_clamp_max), + glu_clamp_min=cutlass.Float32(glu_clamp_min), + ) + + +def _grouped_gemm_glu_impl( + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "swiglu", + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + use_dynamic_sched: bool = False, + cluster_overlap_margin: int = 0, + *, + b_layout: str = "LNK", + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute a dense-weight block-scaled grouped GEMM with fused GLU. + + Dense expert weights use public shape ``(L, N, K)``. ``swiglu`` and ``geglu`` + produce ``N / 2`` output columns. Native FP4 arrays cover MXFP4/NVFP4; + raw-byte reinterpretation is intentionally not accepted. Low-precision + outputs from FP8 inputs return scale factors. + """ + + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK",)) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + ) + if n % 64: + raise ValueError(f"b_tensor N must be divisible by 64 for GLU, got {n}") + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_input_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ab_dtype=ab_dtype, + ) + if prob_tensor is not None: + require_grouped_probability("prob_tensor", prob_tensor, m=m) + if bias_tensor is not None: + require_array( + bias_tensor, + name="bias_tensor", + shape=(experts, n), + dtype=(jnp.float16, jnp.bfloat16, jnp.float32), + ) + + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + c_dtype = require_dtype( + c_dtype, + ( + jnp.float32, + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + jnp.float4_e2m1fn, + ), + name="c_dtype", + default=jnp.bfloat16, + ) + if is_fp4_dtype(ab_dtype): + valid_d_dtypes = (jnp.float16, jnp.bfloat16, jnp.float32) + else: + valid_d_dtypes = ( + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + jnp.float4_e2m1fn, + ) + d_dtype = require_dtype( + d_dtype, valid_d_dtypes, name="d_dtype", default=jnp.bfloat16 + ) + if ( + is_fp4_dtype(ab_dtype) + and sf_vec_size == 16 + and d_dtype == jnp.dtype(jnp.float32) + and bias_tensor is None + ): + raise NotImplementedError( + "FP4 with sf_vec_size=16 and float32 D requires a bias tensor" + ) + fp8_dtypes = { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + } + # The Torch GLU wrapper requests row/column scales for every FP8-input + # configuration, including a high-precision D output. + generate_sfd = is_fp8_dtype(ab_dtype) + has_amax = d_dtype in {jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)} + if generate_sfd: + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required for FP8 inputs") + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + discrete_col_sfd = False + if vector_f32 and c_dtype in fp8_dtypes: + raise ValueError("vector_f32 does not support an FP8 c_dtype") + if act_func not in ("swiglu", "geglu"): + raise ValueError(f"act_func must be 'swiglu' or 'geglu', got {act_func!r}") + if linear_offset is None: + linear_offset = 1.0 if act_func == "geglu" else 0.0 + + mma_tiler_mn = require_grouped_mma_tiler( + mma_tiler_mn, allowed_m=(128, 256), allowed_n=(256,) + ) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + + output_n = n // 2 + 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_descs = [ + a_desc, + b_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + ] + if prob_tensor is not None: + inputs.append(prob_tensor) + input_descs.append( + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE) + ) + if bias_tensor is not None: + inputs.append(bias_tensor) + input_descs.append( + as_gemm_tensor_desc("bias_tensor", bias_tensor, mode=GROUPED_BIAS_MODE) + ) + if generate_sfd: + inputs.append(norm_const_tensor) + input_descs.append(as_gemm_tensor_desc("norm_const_tensor", norm_const_tensor)) + + outputs = [ + make_buffer_desc("c_tensor", (1, m, n), c_dtype, mode=output_mode), + make_buffer_desc( + "d_tensor", (1, m, output_n), d_dtype, mode=output_mode + ), + ] + outputs.append( + make_buffer_desc( + "d_col_tensor", (1, m, output_n), d_dtype, mode=output_mode + ) + ) + if has_amax: + outputs.append( + make_buffer_desc( + "amax_tensor", (experts, 1), jnp.float32, init_value=-float("inf") + ) + ) + if generate_sfd: + outputs.extend( + ( + make_buffer_desc( + "sfd_row_tensor", + block_scale_shape(m, output_n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + mode=BLOCK_SCALE_MODE, + ), + make_buffer_desc( + "sfd_col_tensor", + block_scale_shape(output_n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + mode=BLOCK_SCALE_MODE, + ), + ) + ) + + workspace_bytes = max(dense_workspace_bytes(bool(use_dynamic_sched)), 1) + workspaces = [ + make_buffer_desc( + "workspace", + (workspace_bytes,), + jnp.uint8, + ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT, + ) + ] + results = call_cutedsl( + _launch, + inputs, + input_descs=input_descs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "discrete_col_sfd": bool(discrete_col_sfd), + "expert_cnt": experts, + "has_prob": prob_tensor is not None, + "has_bias": bias_tensor is not None, + "generate_sfd": generate_sfd, + "has_amax": has_amax, + "use_dynamic_sched": bool(use_dynamic_sched), + "act_func": act_func, + "linear_offset": float(linear_offset), + "geglu_alpha": float(geglu_alpha), + "glu_clamp_max": float(glu_clamp_max), + "glu_clamp_min": float(glu_clamp_min), + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + workspaces=workspaces, + ) + c_tensor, d_tensor, d_col_tensor = results[:3] + result_idx = 3 + amax_tensor = results[result_idx] if has_amax else None + result_idx += int(has_amax) + sfd_row_tensor = results[result_idx] if generate_sfd else None + sfd_col_tensor = results[result_idx + 1] if generate_sfd else None + return TupleDict( + c_tensor=c_tensor, + d_tensor=d_tensor, + d_col_tensor=d_col_tensor, + amax_tensor=amax_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + ) + + +class GroupedGemmGluSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM + GLU.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_sfa_tensor: Any, + sample_padded_offsets: Any, + sample_alpha_tensor: Any, + sample_b_tensor: Any, + sample_sfb_tensor: Any, + sample_bias_tensor: Optional[Any] = None, + sample_norm_const_tensor: Optional[Any] = None, + sample_prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "swiglu", + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", + ) -> None: + super().__init__() + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK",)) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, mode=a_mode, name="sample_a_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, 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, mode=b_mode, name="sample_b_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfb_tensor" + ), + "bias_tensor": self.make_optional_tensor_desc( + sample_bias_tensor, + mode=GROUPED_BIAS_MODE, + 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, + mode=PROBABILITY_MODE, + name="sample_prob_tensor", + ), + } + self._config = { + "acc_dtype": self.as_optional_dtype(acc_dtype), + "c_dtype": self.as_optional_dtype(c_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "discrete_col_sfd": discrete_col_sfd, + "act_func": act_func, + "linear_offset": linear_offset, + "geglu_alpha": geglu_alpha, + "glu_clamp_max": glu_clamp_max, + "glu_clamp_min": glu_clamp_min, + "use_dynamic_sched": use_dynamic_sched, + "b_layout": b_layout, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_glu_impl( + self._sample_descs["a_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["sfb_tensor"], + self._sample_descs["bias_tensor"], + self._sample_descs["norm_const_tensor"], + self._sample_descs["prob_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + b_tensor, + sfb_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "sfa_tensor": sfa_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "b_tensor": b_tensor, + "sfb_tensor": sfb_tensor, + "bias_tensor": bias_tensor, + "norm_const_tensor": norm_const_tensor, + "prob_tensor": prob_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_glu_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "acc_dtype", + "c_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "act_func", + "linear_offset", + "geglu_alpha", + "glu_clamp_max", + "glu_clamp_min", + "use_dynamic_sched", + "b_layout", + ), +) +def grouped_gemm_glu_wrapper_sm100( + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "swiglu", + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", +) -> TupleDict: + """Compute an MXFP8 dense-weight grouped GEMM with fused GLU.""" + + op = GroupedGemmGluSm100( + a_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + b_tensor, + sfb_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + acc_dtype=acc_dtype, + c_dtype=c_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + act_func=act_func, + linear_offset=linear_offset, + geglu_alpha=geglu_alpha, + glu_clamp_max=glu_clamp_max, + glu_clamp_min=glu_clamp_min, + use_dynamic_sched=use_dynamic_sched, + b_layout=b_layout, + ) + return op( + a_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + b_tensor, + sfb_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + ) + + +__all__ = [ + "GroupedGemmGluSm100", + "grouped_gemm_glu_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/__init__.py index 2ef06aa02..89974fb93 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/__init__.py @@ -1,12 +1,15 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmGluHadamardSm100, - grouped_gemm_glu_hadamard_wrapper_sm100, -) +from ...common.operation_api import make_operation_api -__all__ = [ - "GroupedGemmGluHadamardSm100", - "grouped_gemm_glu_hadamard_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "api": ( + "GroupedGemmGluHadamardSm100", + "grouped_gemm_glu_hadamard_wrapper_sm100", + ) + }, + submodules=("api", "jax"), +) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/hadamard_utils.py b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/hadamard_utils.py index ec55574bc..08b7b056d 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/hadamard_utils.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/hadamard_utils.py @@ -10,7 +10,6 @@ from cutlass.cute.nvgpu import tcgen05 from cutlass.cute.nvgpu import OperandMajorMode from cutlass.cute.nvgpu.tcgen05 import OperandSource -import torch HADAMARD_SIZE = 16 TMEM_ROW_STRIDE = 1 << 16 @@ -235,6 +234,8 @@ def hadamard_smem_transpose_fwht_amax(sD, d_buffer, feature_offset: cutlass.Cons def hadamard_matrix(n, dtype=None, device=None): + import torch + if dtype is None: dtype = torch.float32 if n < 1: diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/jax.py new file mode 100644 index 000000000..552e78f6b --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard/jax.py @@ -0,0 +1,648 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for dense grouped GEMM + GLU + Hadamard on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + BLOCK_SCALE_MODE, + make_buffer_desc, + FIX_PAD_SIZE, + HADAMARD_SIZE, + MAX_EXPERTS, + SF_VEC_SIZES, + TupleDict, + GROUPED_BIAS_MODE, + GROUPED_WORKSPACE_ALIGNMENT, + PROBABILITY_MODE, + as_gemm_tensor_desc, + call_cutedsl, + dense_workspace_bytes, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, + require_array, + require_contiguous_alignment, + require_dtype, + require_grouped_block_scales, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_mma_tiler, + require_grouped_probability, + require_layout, +) + + +def hadamard_values(size: int) -> tuple[tuple[int, ...], ...]: + """Return an unnormalized Sylvester Hadamard matrix as Python values.""" + + if size < 1 or size & (size - 1): + raise ValueError(f"Hadamard size must be a positive power of two, got {size}") + matrix = ((1,),) + while len(matrix) < size: + matrix = tuple(row + row for row in matrix) + tuple( + row + tuple(-value for value in row) for row in matrix + ) + return matrix + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + expert_cnt: int, + act_func: str, + has_bias: bool, + has_hadamard: bool, + use_dynamic_sched: bool, + use_tmem_post_rht_amax: bool, + max_active_clusters: int, +): + import cutlass + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from ..moe_utils import MoEWeightMode + from .moe_blockscaled_grouped_gemm_glu_hadamard import ( + BlockScaledMoEGroupedGemmGluHadamardKernel, + ) + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a = take() + b = take() + sfa = take() + sfb = take() + padded_offsets = take() + alpha = take() + prob = take() + hadamard = take() if has_hadamard else None + bias = take() if has_bias else None + c = take() + d = take() + amax = take() + post_rht_amax = take() + workspace = take() + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped GEMM argument count: consumed {arg_idx}, received {len(args)}" + ) + + kernel = BlockScaledMoEGroupedGemmGluHadamardKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=True, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + expert_cnt=expert_cnt, + weight_mode=MoEWeightMode.DENSE, + use_dynamic_sched=use_dynamic_sched, + act_func=act_func, + enable_bias=has_bias, + use_tmem_post_rht_amax=use_tmem_post_rht_amax, + ) + kernel( + a, + b, + sfa, + sfb, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int64(0), + OperandMajorMode.K, + workspace.iterator, + c, + d, + amax, + post_rht_amax, + padded_offsets, + alpha, + prob, + hadamard, + bias, + max_active_clusters, + stream, + linear_offset=cutlass.Float32(1.0 if act_func == "geglu" else 0.0), + ) + + +def _grouped_gemm_glu_hadamard_impl( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + bias_tensor: Optional[Any] = None, + hadamard_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + m_aligned: int = 256, + act_func: str = "swiglu", + use_dynamic_sched: bool = False, + use_tmem_post_rht_amax: bool = False, + cluster_overlap_margin: int = 0, + *, + b_layout: str = "LNK", + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute dense native-FP4 grouped GEMM with GLU and Hadamard amax. + + A and B use JAX's native ``float4_e2m1fn`` dtype with public shapes + ``(1, M, K)`` and ``(L, N, K)``. Raw ``uint8`` FP4 payloads and discrete + per-expert pointer arrays are intentionally not part of this API. + + Runtime ``padded_offsets`` values must be nondecreasing multiples of 256, + must not exceed M, and must end at the number of rows whose outputs are + meaningful. Rows beyond the final offset are returned as zero. Empty + experts produce zero amax values. These value constraints are trusted + while tracing with :func:`jax.jit`. + """ + + from cutlass import Float32 + from cutlass.jax import jax_to_cutlass_dtype + + from .moe_blockscaled_grouped_gemm_glu_hadamard import ( + BlockScaledMoEGroupedGemmGluHadamardKernel, + ) + + kernel = BlockScaledMoEGroupedGemmGluHadamardKernel + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK",)) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + valid_ab_dtypes=(jnp.float4_e2m1fn,), + ) + if m % FIX_PAD_SIZE: + raise ValueError(f"M must be divisible by {FIX_PAD_SIZE}, got {m}") + if n % 64: + raise ValueError(f"N must be divisible by 64, got {n}") + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + if sf_vec_size not in SF_VEC_SIZES: + raise ValueError( + f"sf_vec_size must be one of {SF_VEC_SIZES}, got {sf_vec_size}" + ) + if act_func not in ("swiglu", "geglu", "srelu"): + raise ValueError( + f"act_func must be 'swiglu', 'geglu', or 'srelu', got {act_func!r}" + ) + sf_dtype = require_grouped_block_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + valid_dtypes=(jnp.float8_e8m0fnu, jnp.float8_e4m3fn), + ) + if sf_dtype == jnp.dtype(jnp.float8_e4m3fn) and sf_vec_size == 32: + raise ValueError("float8_e4m3fn scales require sf_vec_size=16") + require_grouped_probability("prob_tensor", prob_tensor, m=m) + + if bias_tensor is not None: + require_array( + bias_tensor, + name="bias_tensor", + shape=(experts, n), + dtype=(jnp.float16, jnp.bfloat16, jnp.float32), + ) + + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + c_dtype = require_dtype( + c_dtype, (jnp.float16, jnp.bfloat16), name="c_dtype", default=jnp.bfloat16 + ) + d_dtype = require_dtype( + d_dtype, (jnp.float16, jnp.bfloat16), name="d_dtype", default=jnp.bfloat16 + ) + + mma_tiler_mn = require_grouped_mma_tiler( + mma_tiler_mn, allowed_m=(256,), allowed_n=(256,) + ) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + + output_n = n if act_func == "srelu" else n // 2 + if output_n % HADAMARD_SIZE: + raise ValueError( + f"D's N dimension must be divisible by {HADAMARD_SIZE}, got {output_n}" + ) + require_contiguous_alignment("a_tensor", k, 4) + require_contiguous_alignment("b_tensor", k, 4) + require_contiguous_alignment("c_tensor", n, c_dtype.itemsize * 8) + require_contiguous_alignment("d_tensor", output_n, d_dtype.itemsize * 8) + + if not kernel.can_implement( + jax_to_cutlass_dtype(ab_dtype), + jax_to_cutlass_dtype(sf_dtype), + sf_vec_size, + Float32, + jax_to_cutlass_dtype(d_dtype), + True, + mma_tiler_mn, + cluster_shape_mn, + m, + n, + k, + experts, + "k", + "k", + "n", + m_aligned, + ): + raise ValueError("Unsupported grouped GEMM GLU Hadamard configuration") + + has_hadamard = bool(use_tmem_post_rht_amax) + if has_hadamard and hadamard_tensor is not None: + require_array( + hadamard_tensor, + name="hadamard_tensor", + shape=(HADAMARD_SIZE, HADAMARD_SIZE), + dtype=jnp.bfloat16, + ) + elif not has_hadamard and hadamard_tensor is not None: + raise ValueError( + "hadamard_tensor is used only when use_tmem_post_rht_amax=True" + ) + + if _validate_only: + return { + "acc_dtype": acc_dtype, + "c_dtype": c_dtype, + "d_dtype": d_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + } + + inputs = [ + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + ] + input_descs = [ + a_desc, + b_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE), + ] + + if has_hadamard: + if hadamard_tensor is None: + hadamard_tensor = jnp.asarray( + hadamard_values(HADAMARD_SIZE), dtype=jnp.bfloat16 + ) + inputs.append(hadamard_tensor) + input_descs.append(as_gemm_tensor_desc("hadamard_tensor", hadamard_tensor)) + + if bias_tensor is not None: + inputs.append(bias_tensor) + input_descs.append( + as_gemm_tensor_desc("bias_tensor", bias_tensor, mode=GROUPED_BIAS_MODE) + ) + + workspace_bytes = max(dense_workspace_bytes(bool(use_dynamic_sched)), 1) + c_tensor, d_tensor, amax_tensor, post_rht_amax_tensor = call_cutedsl( + _launch, + inputs, + input_descs=input_descs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "expert_cnt": experts, + "act_func": act_func, + "has_bias": bias_tensor is not None, + "has_hadamard": has_hadamard, + "use_dynamic_sched": bool(use_dynamic_sched), + "use_tmem_post_rht_amax": bool(use_tmem_post_rht_amax), + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=( + make_buffer_desc( + "c_tensor", (1, m, n), c_dtype, mode=output_mode, init_value=0 + ), + make_buffer_desc( + "d_tensor", + (1, m, output_n), + d_dtype, + mode=output_mode, + init_value=0, + ), + make_buffer_desc("amax_tensor", (experts, 1), jnp.float32, init_value=0.0), + make_buffer_desc( + "post_rht_amax_tensor", (experts, 1), jnp.float32, init_value=0.0 + ), + ), + workspaces=( + make_buffer_desc( + "workspace", + (workspace_bytes,), + jnp.uint8, + ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT, + init_value=0, + ), + ), + ) + 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_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, mode=a_mode, name="sample_a_tensor" + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, mode=b_mode, name="sample_b_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfa_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, 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, + mode=PROBABILITY_MODE, + name="sample_prob_tensor", + ), + "bias_tensor": self.make_optional_tensor_desc( + sample_bias_tensor, + mode=GROUPED_BIAS_MODE, + name="sample_bias_tensor", + ), + "hadamard_tensor": self.make_optional_tensor_desc( + sample_hadamard_tensor, name="sample_hadamard_tensor" + ), + } + self._config = { + "acc_dtype": self.as_optional_dtype(acc_dtype), + "c_dtype": self.as_optional_dtype(c_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "act_func": act_func, + "use_dynamic_sched": use_dynamic_sched, + "use_tmem_post_rht_amax": use_tmem_post_rht_amax, + "b_layout": b_layout, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_glu_hadamard_impl( + self._sample_descs["a_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["sfb_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["prob_tensor"], + self._sample_descs["bias_tensor"], + self._sample_descs["hadamard_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + bias_tensor: Optional[Any] = None, + hadamard_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + bias_tensor, + hadamard_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + bias_tensor: Optional[Any] = None, + hadamard_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "b_tensor": b_tensor, + "sfa_tensor": sfa_tensor, + "sfb_tensor": sfb_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "prob_tensor": prob_tensor, + "bias_tensor": bias_tensor, + "hadamard_tensor": hadamard_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_glu_hadamard_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "acc_dtype", + "c_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "act_func", + "use_dynamic_sched", + "use_tmem_post_rht_amax", + "b_layout", + ), +) +def grouped_gemm_glu_hadamard_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + prob_tensor: Any, + bias_tensor: Optional[Any] = None, + hadamard_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + m_aligned: int = 256, + act_func: str = "swiglu", + use_dynamic_sched: bool = False, + use_tmem_post_rht_amax: bool = False, + *, + b_layout: str = "LNK", +) -> TupleDict: + """Compute dense native-FP4 grouped GEMM with GLU and Hadamard amax.""" + + op = GroupedGemmGluHadamardSm100( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + bias_tensor, + hadamard_tensor, + acc_dtype=acc_dtype, + c_dtype=c_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + act_func=act_func, + use_dynamic_sched=use_dynamic_sched, + use_tmem_post_rht_amax=use_tmem_post_rht_amax, + b_layout=b_layout, + ) + return op( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + bias_tensor, + hadamard_tensor, + ) + + +__all__ = [ + "GroupedGemmGluHadamardSm100", + "grouped_gemm_glu_hadamard_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py index 9f88a8d00..2806ae24b 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_quant/__init__.py @@ -1,20 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -""" -Grouped GEMM Quant Kernel Module +from ...common.operation_api import make_operation_api -This module provides the contiguous grouped GEMM with output quantization -for MoE (Mixture of Experts) workloads on SM100+ GPUs. -Used for FC2 (forward down-projection) and dFC1 (backward FC1 GEMMs). -""" - -from .api import ( - GroupedGemmQuantSm100, - grouped_gemm_quant_wrapper_sm100, +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmQuantSm100", "grouped_gemm_quant_wrapper_sm100")}, + submodules=("api", "jax"), ) - -__all__ = [ - "GroupedGemmQuantSm100", - "grouped_gemm_quant_wrapper_sm100", -] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py b/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py index eaf7d7d98..8bc17b253 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_quant/api.py @@ -168,6 +168,7 @@ def __init__( stride_order=self.d_desc.stride_order, device=self.d_desc.device, name="sample_d_col", + init_value=self.d_desc.init_value, ) self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row") self.sfd_col_desc = self._make_tensor_desc(sample_sfd_col, name="sample_sfd_col") diff --git a/python/cudnn/grouped_gemm/grouped_gemm_quant/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_quant/jax.py new file mode 100644 index 000000000..7c0d62798 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_quant/jax.py @@ -0,0 +1,626 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for dense-weight grouped GEMM quantization on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + BLOCK_SCALE_MODE, + make_buffer_desc, + FIX_PAD_SIZE, + MAX_EXPERTS, + TWO_CTA_MMA_TILER_M, + TupleDict, + GROUPED_BIAS_MODE, + GROUPED_WORKSPACE_ALIGNMENT, + PROBABILITY_MODE, + as_gemm_tensor_desc, + block_scale_shape, + call_cutedsl, + dense_workspace_bytes, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, + is_fp4_dtype, + is_fp8_dtype, + is_low_precision_output_dtype, + require_16_byte_extent, + require_array, + require_dtype, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_input_scales, + require_grouped_mma_tiler, + require_grouped_probability, + require_grouped_vector, + require_layout, +) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + discrete_col_sfd: bool, + expert_cnt: int, + has_bias: bool, + has_row_scale: bool, + generate_sfd: bool, + has_amax: bool, + use_dynamic_sched: bool, + max_active_clusters: int, +): + import cutlass + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from .grouped_gemm_quant import BlockScaledMoEGroupedGemmQuantKernel + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a = take() + b = take() + sfa = take() + sfb = take() + padded_offsets = take() + alpha = take() + prob = take() + row_scale = take() if has_row_scale else None + bias = take() if has_bias else None + norm_const = take() if generate_sfd else None + d = take() + d_col = take() if generate_sfd else None + amax = take() if has_amax else None + sfd_row = take() if generate_sfd else None + sfd_col = take() if generate_sfd else None + workspace = take() + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped GEMM argument count: consumed {arg_idx}, received {len(args)}" + ) + + kernel = BlockScaledMoEGroupedGemmQuantKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + generate_sfd=generate_sfd, + discrete_col_sfd=discrete_col_sfd, + enable_bias=has_bias, + expert_cnt=expert_cnt, + use_dynamic_sched=use_dynamic_sched, + ) + kernel( + a, + b, + sfb, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int64(0), + OperandMajorMode.K, + workspace.iterator, + d, + d_col, + sfa, + sfd_row, + sfd_col, + amax, + norm_const, + padded_offsets, + alpha, + row_scale, + bias, + prob, + max_active_clusters, + stream, + ) + + +def _grouped_gemm_quant_impl( + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + row_scale_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + cluster_overlap_margin: int = 0, + *, + b_layout: str = "LNK", + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute a dense-weight block-scaled grouped GEMM with quantization. + + The grouped dimension is represented by B's leading ``L`` dimension and by the + runtime ``padded_offsets`` tensor. FP8 outputs return row/column E8M0 scale + factors; native FP4 inputs use logical JAX shapes and never raw-byte + reinterpretation. FP16/BF16 outputs return an initialized amax reduction. + """ + + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + ) + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_input_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ab_dtype=ab_dtype, + ) + if prob_tensor is None: + raise ValueError("prob_tensor is required; pass ones when no gating is needed") + require_grouped_probability("prob_tensor", prob_tensor, m=m) + if row_scale_tensor is not None: + require_grouped_vector("row_scale_tensor", row_scale_tensor, length=m) + if bias_tensor is not None: + require_array( + bias_tensor, + name="bias_tensor", + shape=(experts, n), + dtype=(jnp.float16, jnp.bfloat16, jnp.float32), + ) + + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + if is_fp4_dtype(ab_dtype): + valid_d_dtypes = (jnp.float16, jnp.bfloat16, jnp.float32) + else: + valid_d_dtypes = ( + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + jnp.float4_e2m1fn, + ) + d_dtype = require_dtype( + d_dtype, valid_d_dtypes, name="d_dtype", default=jnp.bfloat16 + ) + if ( + is_fp4_dtype(ab_dtype) + and sf_vec_size == 16 + and d_dtype == jnp.dtype(jnp.float32) + ): + raise NotImplementedError( + "FP4 with sf_vec_size=16 does not support a float32 D output" + ) + generate_sfd = is_fp8_dtype(ab_dtype) and is_low_precision_output_dtype(d_dtype) + has_amax = d_dtype in {jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)} + if generate_sfd: + if norm_const_tensor is None: + raise ValueError( + "norm_const_tensor is required for an FP8/FP4 output from FP8 inputs" + ) + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + discrete_col_sfd = False + if is_fp4_dtype(ab_dtype) and b_layout != "LNK": + raise ValueError("Native FP4 B must use the K-major LNK layout") + mma_tiler_mn = require_grouped_mma_tiler( + mma_tiler_mn, allowed_m=(128, 256), allowed_n=(256,) + ) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + + 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_descs = [ + a_desc, + b_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE), + ] + if row_scale_tensor is not None: + inputs.append(row_scale_tensor) + input_descs.append(as_gemm_tensor_desc("row_scale_tensor", row_scale_tensor)) + if bias_tensor is not None: + inputs.append(bias_tensor) + input_descs.append( + as_gemm_tensor_desc("bias_tensor", bias_tensor, mode=GROUPED_BIAS_MODE) + ) + if generate_sfd: + inputs.append(norm_const_tensor) + input_descs.append(as_gemm_tensor_desc("norm_const_tensor", norm_const_tensor)) + + outputs = [ + make_buffer_desc("d_tensor", (1, m, n), d_dtype, mode=output_mode) + ] + if generate_sfd: + outputs.append( + make_buffer_desc( + "d_col_tensor", (1, m, n), d_dtype, mode=output_mode + ) + ) + if has_amax: + outputs.append( + make_buffer_desc( + "amax_tensor", (experts, 1), jnp.float32, init_value=-float("inf") + ) + ) + if generate_sfd: + outputs.extend( + ( + make_buffer_desc( + "sfd_row_tensor", + block_scale_shape(m, n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + mode=BLOCK_SCALE_MODE, + ), + make_buffer_desc( + "sfd_col_tensor", + block_scale_shape(n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + mode=BLOCK_SCALE_MODE, + ), + ) + ) + + workspace_bytes = max(dense_workspace_bytes(bool(use_dynamic_sched)), 1) + results = call_cutedsl( + _launch, + inputs, + input_descs=input_descs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "discrete_col_sfd": bool(discrete_col_sfd), + "expert_cnt": experts, + "has_bias": bias_tensor is not None, + "has_row_scale": row_scale_tensor is not None, + "generate_sfd": generate_sfd, + "has_amax": has_amax, + "use_dynamic_sched": bool(use_dynamic_sched), + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + workspaces=( + make_buffer_desc( + "workspace", + (workspace_bytes,), + jnp.uint8, + ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT, + ), + ), + ) + result_idx = 1 + d_tensor = results[0] + d_col_tensor = results[result_idx] if generate_sfd else None + result_idx += int(generate_sfd) + amax_tensor = results[result_idx] if has_amax else None + result_idx += int(has_amax) + sfd_row_tensor = results[result_idx] if generate_sfd else None + sfd_col_tensor = results[result_idx + 1] if generate_sfd else None + return TupleDict( + d_tensor=d_tensor, + d_col_tensor=d_col_tensor, + amax_tensor=amax_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + ) + + +class GroupedGemmQuantSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM quantization.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_sfa_tensor: Any, + sample_padded_offsets: Any, + sample_alpha_tensor: Any, + sample_b_tensor: Any, + sample_sfb_tensor: Any, + sample_bias_tensor: Optional[Any] = None, + sample_norm_const_tensor: Optional[Any] = None, + sample_prob_tensor: Optional[Any] = None, + sample_row_scale_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", + ) -> None: + super().__init__() + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, mode=a_mode, name="sample_a_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, 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, mode=b_mode, name="sample_b_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfb_tensor" + ), + "bias_tensor": self.make_optional_tensor_desc( + sample_bias_tensor, + mode=GROUPED_BIAS_MODE, + 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, + mode=PROBABILITY_MODE, + name="sample_prob_tensor", + ), + "row_scale_tensor": self.make_optional_tensor_desc( + sample_row_scale_tensor, name="sample_row_scale_tensor" + ), + } + self._config = { + "acc_dtype": self.as_optional_dtype(acc_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "discrete_col_sfd": discrete_col_sfd, + "use_dynamic_sched": use_dynamic_sched, + "b_layout": b_layout, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_quant_impl( + self._sample_descs["a_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["sfb_tensor"], + self._sample_descs["bias_tensor"], + self._sample_descs["norm_const_tensor"], + self._sample_descs["prob_tensor"], + self._sample_descs["row_scale_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + row_scale_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + b_tensor, + sfb_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + row_scale_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + row_scale_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "sfa_tensor": sfa_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "b_tensor": b_tensor, + "sfb_tensor": sfb_tensor, + "bias_tensor": bias_tensor, + "norm_const_tensor": norm_const_tensor, + "prob_tensor": prob_tensor, + "row_scale_tensor": row_scale_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_quant_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "acc_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "use_dynamic_sched", + "b_layout", + ), +) +def grouped_gemm_quant_wrapper_sm100( + a_tensor: Any, + sfa_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + b_tensor: Any, + sfb_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + row_scale_tensor: Optional[Any] = None, + acc_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", +) -> TupleDict: + """Compute an MXFP8 dense-weight grouped GEMM with optional quantization.""" + + op = GroupedGemmQuantSm100( + a_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + b_tensor, + sfb_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + row_scale_tensor, + acc_dtype=acc_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + use_dynamic_sched=use_dynamic_sched, + b_layout=b_layout, + ) + return op( + a_tensor, + sfa_tensor, + padded_offsets, + alpha_tensor, + b_tensor, + sfb_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + row_scale_tensor, + ) + + +__all__ = [ + "GroupedGemmQuantSm100", + "grouped_gemm_quant_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_srelu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_srelu/__init__.py index 70b2bfe7d..b13e70906 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_srelu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_srelu/__init__.py @@ -1,19 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -""" -Grouped GEMM SReLU Kernel Module +from ...common.operation_api import make_operation_api -This module provides the forward grouped GEMM with SReLU activation -for MoE (Mixture of Experts) workloads on SM100+ GPUs. -""" - -from .api import ( - GroupedGemmSreluSm100, - grouped_gemm_srelu_wrapper_sm100, +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmSreluSm100", "grouped_gemm_srelu_wrapper_sm100")}, + submodules=("api", "jax"), ) - -__all__ = [ - "GroupedGemmSreluSm100", - "grouped_gemm_srelu_wrapper_sm100", -] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_srelu/api.py b/python/cudnn/grouped_gemm/grouped_gemm_srelu/api.py index 2ff670e08..97d7d26da 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_srelu/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_srelu/api.py @@ -179,6 +179,7 @@ def __init__( stride_order=self.d_desc.stride_order, device=self.d_desc.device, name="sample_d_col", + init_value=self.d_desc.init_value, ) self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row") self.sfd_col_desc = self._make_tensor_desc(sample_sfd_col, name="sample_sfd_col") diff --git a/python/cudnn/grouped_gemm/grouped_gemm_srelu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_srelu/jax.py new file mode 100644 index 000000000..854d2e904 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_srelu/jax.py @@ -0,0 +1,627 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for dense-weight grouped GEMM + squared ReLU on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + BLOCK_SCALE_MODE, + make_buffer_desc, + FIX_PAD_SIZE, + MAX_EXPERTS, + TWO_CTA_MMA_TILER_M, + TupleDict, + GROUPED_BIAS_MODE, + GROUPED_WORKSPACE_ALIGNMENT, + PROBABILITY_MODE, + as_gemm_tensor_desc, + block_scale_shape, + call_cutedsl, + dense_workspace_bytes, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, + is_fp4_dtype, + is_fp8_dtype, + is_low_precision_output_dtype, + require_16_byte_extent, + require_array, + require_dtype, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_input_scales, + require_grouped_mma_tiler, + require_grouped_probability, + require_grouped_vector, + require_layout, +) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + discrete_col_sfd: bool, + expert_cnt: int, + has_bias: bool, + generate_sfd: bool, + has_amax: bool, + use_dynamic_sched: bool, + max_active_clusters: int, +): + import cutlass + from cutlass.cute.nvgpu import OperandMajorMode + from cutlass.jax import jax_to_cutlass_dtype + + from .moe_blockscaled_grouped_gemm_srelu_quant import ( + BlockScaledMoEGroupedGemmQuantKernel, + EpilogueType, + ) + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a = take() + b = take() + sfa = take() + sfb = take() + padded_offsets = take() + alpha = take() + prob = take() + bias = take() if has_bias else None + norm_const = take() if generate_sfd else None + c = take() + d = take() + d_col = take() if generate_sfd else None + amax = take() if has_amax else None + sfd_row = take() if generate_sfd else None + sfd_col = take() if generate_sfd else None + workspace = take() + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped GEMM argument count: consumed {arg_idx}, received {len(args)}" + ) + + kernel = BlockScaledMoEGroupedGemmQuantKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=vector_f32, + generate_sfd=generate_sfd, + discrete_col_sfd=discrete_col_sfd, + generate_c=True, + enable_bias=has_bias, + expert_cnt=expert_cnt, + use_dynamic_sched=use_dynamic_sched, + epilogue_type=EpilogueType.SRELU.value, + ) + kernel( + a, + b, + sfb, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int64(0), + OperandMajorMode.K, + workspace.iterator, + c, + d, + d_col, + sfa, + sfd_row, + sfd_col, + amax, + norm_const, + padded_offsets, + alpha, + bias, + prob, + max_active_clusters, + stream, + ) + + +def _grouped_gemm_srelu_impl( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + cluster_overlap_margin: int = 0, + *, + b_layout: str = "LNK", + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute a dense-weight block-scaled grouped GEMM and squared-ReLU fusion. + + Native FP4 inputs cover MXFP4/NVFP4 without raw-byte reinterpretation. + Low-precision outputs from FP8 inputs return row/column scale factors; + FP16/BF16 outputs return an initialized per-expert amax reduction. + """ + + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + ) + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_input_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ab_dtype=ab_dtype, + ) + if prob_tensor is None: + raise ValueError("prob_tensor is required; pass ones when no gating is needed") + require_grouped_probability("prob_tensor", prob_tensor, m=m) + if bias_tensor is not None: + require_array( + bias_tensor, + name="bias_tensor", + shape=(experts, n), + dtype=(jnp.float16, jnp.bfloat16, jnp.float32), + ) + + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + c_dtype = require_dtype( + c_dtype, + (jnp.float32, jnp.float16, jnp.bfloat16, jnp.float8_e4m3fn, jnp.float8_e5m2), + name="c_dtype", + default=jnp.bfloat16, + ) + if is_fp4_dtype(ab_dtype): + valid_d_dtypes = (jnp.float16, jnp.bfloat16, jnp.float32) + else: + valid_d_dtypes = ( + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + jnp.float4_e2m1fn, + ) + d_dtype = require_dtype( + d_dtype, valid_d_dtypes, name="d_dtype", default=jnp.bfloat16 + ) + if ( + is_fp4_dtype(ab_dtype) + and sf_vec_size == 16 + and d_dtype == jnp.dtype(jnp.float32) + ): + raise NotImplementedError( + "FP4 with sf_vec_size=16 does not support a float32 D output" + ) + generate_sfd = is_fp8_dtype(ab_dtype) and is_low_precision_output_dtype(d_dtype) + has_amax = d_dtype in {jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)} + if generate_sfd: + if norm_const_tensor is None: + raise ValueError( + "norm_const_tensor is required for an FP8/FP4 output from FP8 inputs" + ) + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + discrete_col_sfd = False + if is_fp4_dtype(ab_dtype) and b_layout != "LNK": + raise ValueError("Native FP4 B must use the K-major LNK layout") + mma_tiler_mn = require_grouped_mma_tiler( + mma_tiler_mn, allowed_m=(128, 256), allowed_n=(256,) + ) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + + 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_descs = [ + a_desc, + b_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE), + ] + if bias_tensor is not None: + inputs.append(bias_tensor) + input_descs.append( + as_gemm_tensor_desc("bias_tensor", bias_tensor, mode=GROUPED_BIAS_MODE) + ) + if generate_sfd: + inputs.append(norm_const_tensor) + input_descs.append(as_gemm_tensor_desc("norm_const_tensor", norm_const_tensor)) + + outputs = [ + make_buffer_desc("c_tensor", (1, m, n), c_dtype, mode=output_mode), + make_buffer_desc("d_tensor", (1, m, n), d_dtype, mode=output_mode), + ] + if generate_sfd: + outputs.append( + make_buffer_desc( + "d_col_tensor", (1, m, n), d_dtype, mode=output_mode + ) + ) + if has_amax: + outputs.append( + make_buffer_desc( + "amax_tensor", (experts, 1), jnp.float32, init_value=-float("inf") + ) + ) + if generate_sfd: + outputs.extend( + ( + make_buffer_desc( + "sfd_row_tensor", + block_scale_shape(m, n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + mode=BLOCK_SCALE_MODE, + ), + make_buffer_desc( + "sfd_col_tensor", + block_scale_shape(n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + mode=BLOCK_SCALE_MODE, + ), + ) + ) + + workspace_bytes = max(dense_workspace_bytes(bool(use_dynamic_sched)), 1) + results = call_cutedsl( + _launch, + inputs, + input_descs=input_descs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "discrete_col_sfd": bool(discrete_col_sfd), + "expert_cnt": experts, + "has_bias": bias_tensor is not None, + "generate_sfd": generate_sfd, + "has_amax": has_amax, + "use_dynamic_sched": bool(use_dynamic_sched), + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + workspaces=( + make_buffer_desc( + "workspace", + (workspace_bytes,), + jnp.uint8, + ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT, + ), + ), + ) + result_idx = 0 + c_tensor, d_tensor = results[:2] + result_idx = 2 + d_col_tensor = results[result_idx] if generate_sfd else None + result_idx += int(generate_sfd) + amax_tensor = results[result_idx] if has_amax else None + result_idx += int(has_amax) + sfd_row_tensor = results[result_idx] if generate_sfd else None + sfd_col_tensor = results[result_idx + 1] if generate_sfd else None + return TupleDict( + c_tensor=c_tensor, + d_tensor=d_tensor, + d_col_tensor=d_col_tensor, + amax_tensor=amax_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + ) + + +class GroupedGemmSreluSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM + SReLU.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_b_tensor: Any, + sample_sfa_tensor: Any, + sample_sfb_tensor: Any, + sample_padded_offsets: Any, + sample_alpha_tensor: Any, + sample_bias_tensor: Optional[Any] = None, + sample_norm_const_tensor: Optional[Any] = None, + sample_prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", + ) -> None: + super().__init__() + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + b_layout = require_layout("b_layout", b_layout, ("LNK", "LKN")) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode(b_layout) + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, mode=a_mode, name="sample_a_tensor" + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, mode=b_mode, name="sample_b_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfa_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, 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, + mode=GROUPED_BIAS_MODE, + 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, + mode=PROBABILITY_MODE, + name="sample_prob_tensor", + ), + } + self._config = { + "acc_dtype": self.as_optional_dtype(acc_dtype), + "c_dtype": self.as_optional_dtype(c_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "discrete_col_sfd": discrete_col_sfd, + "use_dynamic_sched": use_dynamic_sched, + "b_layout": b_layout, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_srelu_impl( + self._sample_descs["a_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["sfb_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["bias_tensor"], + self._sample_descs["norm_const_tensor"], + self._sample_descs["prob_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "b_tensor": b_tensor, + "sfa_tensor": sfa_tensor, + "sfb_tensor": sfb_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "bias_tensor": bias_tensor, + "norm_const_tensor": norm_const_tensor, + "prob_tensor": prob_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_srelu_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "acc_dtype", + "c_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + "use_dynamic_sched", + "b_layout", + ), +) +def grouped_gemm_srelu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + bias_tensor: Optional[Any] = None, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + use_dynamic_sched: bool = False, + *, + b_layout: str = "LNK", +) -> TupleDict: + """Compute an MXFP8 dense-weight grouped GEMM and squared-ReLU fusion.""" + + op = GroupedGemmSreluSm100( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + acc_dtype=acc_dtype, + c_dtype=c_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + use_dynamic_sched=use_dynamic_sched, + b_layout=b_layout, + ) + return op( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + bias_tensor, + norm_const_tensor, + prob_tensor, + ) + + +__all__ = [ + "GroupedGemmSreluSm100", + "grouped_gemm_srelu_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py index 6bcd935f2..c62d1240f 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/__init__.py @@ -1,19 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -""" -Grouped GEMM SwiGLU Kernel Module +from ...common.operation_api import make_operation_api -This module provides the forward grouped GEMM with SwiGLU activation -for MoE (Mixture of Experts) workloads on SM100+ GPUs. -""" - -from .api import ( - GroupedGemmSwigluSm100, - grouped_gemm_swiglu_wrapper_sm100, +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmSwigluSm100", "grouped_gemm_swiglu_wrapper_sm100")}, + submodules=("api", "jax"), ) - -__all__ = [ - "GroupedGemmSwigluSm100", - "grouped_gemm_swiglu_wrapper_sm100", -] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_swiglu/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/jax.py new file mode 100644 index 000000000..901489cbc --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_swiglu/jax.py @@ -0,0 +1,571 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for contiguous grouped GEMM + SwiGLU on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + BLOCK_SCALE_MODE, + make_buffer_desc, + FIX_PAD_SIZE, + MAX_EXPERTS, + TWO_CTA_MMA_TILER_M, + TupleDict, + PROBABILITY_MODE, + as_gemm_tensor_desc, + block_scale_shape, + call_cutedsl, + gemm_a_mode, + gemm_b_mode, + gemm_output_mode, + is_fp4_dtype, + is_fp8_dtype, + require_16_byte_extent, + require_dtype, + require_grouped_cluster_shape, + require_grouped_gemm_inputs, + require_grouped_input_scales, + require_grouped_mma_tiler, + require_grouped_probability, + require_grouped_vector, + require_layout, +) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + vector_f32: bool, + discrete_col_sfd: bool, + expert_cnt: int, + has_prob: bool, + generate_sfd: bool, + has_amax: bool, + max_active_clusters: int, +): + from cutlass.jax import jax_to_cutlass_dtype + + from .grouped_gemm_swiglu_quant import BlockScaledContiguousGroupedGemmKernel + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a = take() + b = take() + sfa = take() + sfb = take() + padded_offsets = take() + alpha = take() + norm_const = take() if generate_sfd else None + prob = take() if has_prob else None + c = take() + d = take() + d_col = take() + amax = take() if has_amax else None + sfd_row = take() if generate_sfd else None + sfd_col = take() if generate_sfd else None + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped GEMM argument count: consumed {arg_idx}, received {len(args)}" + ) + + kernel = BlockScaledContiguousGroupedGemmKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vector_f32=vector_f32, + generate_sfd=generate_sfd, + discrete_col_sfd=discrete_col_sfd, + expert_cnt=expert_cnt, + use_mono_increase_expert_idx=True, + ) + kernel( + a, + b, + c, + d, + d_col, + sfa, + sfb, + sfd_row, + sfd_col, + amax, + norm_const, + padded_offsets, + alpha, + prob, + max_active_clusters, + stream, + ) + + +def _grouped_gemm_swiglu_impl( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + cluster_overlap_margin: int = 0, + *, + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute a contiguous block-scaled grouped GEMM and fused SwiGLU. + + ``a_tensor`` has public shape ``(1, M, K)`` and contains all padded expert rows. + ``b_tensor`` has public shape ``(L, N, K)``, where ``L`` is the expert count. + ``padded_offsets`` supplies each expert's cumulative padded row boundary at + runtime. Matrix outputs use public ``LMN`` order. + + FP4 inputs use JAX's native logical dtype; raw packed ``uint8`` storage is + deliberately excluded. Configuration values are static under JAX tracing. + """ + + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode("LNK") + output_mode = gemm_output_mode(output_layout, name="output_layout") + a_desc = as_gemm_tensor_desc("a_tensor", a_tensor, mode=a_mode) + b_desc = as_gemm_tensor_desc("b_tensor", b_tensor, mode=b_mode) + m, n, k, experts, ab_dtype = require_grouped_gemm_inputs( + a_desc, + b_desc, + padded_offsets, + alpha_tensor, + max_experts=MAX_EXPERTS, + ) + if n % 64: + raise ValueError(f"b_tensor N must be divisible by 64 for SwiGLU, got {n}") + if m_aligned != FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be {FIX_PAD_SIZE}, got {m_aligned}") + require_grouped_input_scales( + sfa_tensor, + sfb_tensor, + m=m, + n=n, + k=k, + experts=experts, + sf_vec_size=sf_vec_size, + ab_dtype=ab_dtype, + ) + if prob_tensor is not None: + require_grouped_probability("prob_tensor", prob_tensor, m=m) + + acc_dtype = require_dtype( + acc_dtype, (jnp.float32,), name="acc_dtype", default=jnp.float32 + ) + c_dtype = require_dtype( + c_dtype, + ( + jnp.float32, + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + jnp.float4_e2m1fn, + ), + name="c_dtype", + default=jnp.bfloat16, + ) + if is_fp4_dtype(ab_dtype): + valid_d_dtypes = (jnp.float16, jnp.bfloat16, jnp.float32) + else: + valid_d_dtypes = ( + jnp.float16, + jnp.bfloat16, + jnp.float8_e4m3fn, + jnp.float8_e5m2, + jnp.float4_e2m1fn, + ) + d_dtype = require_dtype( + d_dtype, valid_d_dtypes, name="d_dtype", default=jnp.bfloat16 + ) + if ( + is_fp4_dtype(ab_dtype) + and sf_vec_size == 16 + and d_dtype == jnp.dtype(jnp.float32) + ): + raise NotImplementedError( + "FP4 with sf_vec_size=16 does not support a float32 D output" + ) + generate_sfd = is_fp8_dtype(ab_dtype) + if generate_sfd: + if norm_const_tensor is None: + raise ValueError("norm_const_tensor is required for FP8 inputs") + require_grouped_vector("norm_const_tensor", norm_const_tensor, length=1) + else: + norm_const_tensor = None + discrete_col_sfd = False + if vector_f32 and c_dtype in { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + }: + raise ValueError("vector_f32 does not support an FP8 c_dtype") + mma_tiler_mn = require_grouped_mma_tiler(mma_tiler_mn) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, mma_tiler_mn=mma_tiler_mn + ) + if mma_tiler_mn[1] == 128 and d_dtype in { + jnp.dtype(jnp.float8_e4m3fn), + jnp.dtype(jnp.float8_e5m2), + }: + raise NotImplementedError("FP8 output requires mma_tiler_mn[1] == 256") + + output_n = n // 2 + require_16_byte_extent("a_tensor", k, ab_dtype) + require_16_byte_extent("b_tensor", k, ab_dtype) + require_16_byte_extent("c_tensor", n, c_dtype) + require_16_byte_extent("d_tensor", output_n, d_dtype) + + if _validate_only: + return { + "acc_dtype": acc_dtype, + "c_dtype": c_dtype, + "d_dtype": d_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + } + + inputs = [ + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + ] + input_descs = [ + a_desc, + b_desc, + as_gemm_tensor_desc("sfa_tensor", sfa_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("sfb_tensor", sfb_tensor, mode=BLOCK_SCALE_MODE), + as_gemm_tensor_desc("padded_offsets", padded_offsets), + as_gemm_tensor_desc("alpha_tensor", alpha_tensor), + ] + if generate_sfd: + inputs.append(norm_const_tensor) + input_descs.append(as_gemm_tensor_desc("norm_const_tensor", norm_const_tensor)) + if prob_tensor is not None: + inputs.append(prob_tensor) + input_descs.append( + as_gemm_tensor_desc("prob_tensor", prob_tensor, mode=PROBABILITY_MODE) + ) + + has_amax = d_dtype in {jnp.dtype(jnp.float16), jnp.dtype(jnp.bfloat16)} + outputs = [ + make_buffer_desc("c_tensor", (1, m, n), c_dtype, mode=output_mode), + make_buffer_desc( + "d_tensor", (1, m, output_n), d_dtype, mode=output_mode + ), + make_buffer_desc( + "d_col_tensor", (1, m, output_n), d_dtype, mode=output_mode + ), + ] + if has_amax: + outputs.append( + make_buffer_desc( + "amax_tensor", (experts, 1), jnp.float32, init_value=-float("inf") + ) + ) + if generate_sfd: + outputs.extend( + ( + make_buffer_desc( + "sfd_row_tensor", + block_scale_shape(m, output_n, 1, sf_vec_size), + jnp.float8_e8m0fnu, + mode=BLOCK_SCALE_MODE, + ), + make_buffer_desc( + "sfd_col_tensor", + block_scale_shape(output_n, m, 1, sf_vec_size), + jnp.float8_e8m0fnu, + mode=BLOCK_SCALE_MODE, + ), + ) + ) + results = call_cutedsl( + _launch, + inputs, + input_descs=input_descs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "vector_f32": bool(vector_f32), + "discrete_col_sfd": bool(discrete_col_sfd), + "expert_cnt": experts, + "has_prob": prob_tensor is not None, + "generate_sfd": generate_sfd, + "has_amax": has_amax, + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=outputs, + ) + result_idx = 0 + c_tensor = results[result_idx] + result_idx += 1 + d_tensor = results[result_idx] + result_idx += 1 + d_col_tensor = results[result_idx] + result_idx += 1 + amax_tensor = results[result_idx] if has_amax else None + result_idx += int(has_amax) + sfd_row_tensor = results[result_idx] if generate_sfd else None + sfd_col_tensor = results[result_idx + 1] if generate_sfd else None + return TupleDict( + c_tensor=c_tensor, + d_tensor=d_tensor, + d_col_tensor=d_col_tensor, + amax_tensor=amax_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + ) + + +class GroupedGemmSwigluSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM + SwiGLU.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_b_tensor: Any, + sample_sfa_tensor: Any, + sample_sfb_tensor: Any, + sample_padded_offsets: Any, + sample_alpha_tensor: Any, + sample_norm_const_tensor: Optional[Any] = None, + sample_prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + ) -> None: + super().__init__() + output_layout = require_layout("output_layout", output_layout, ("LMN",)) + a_mode = gemm_a_mode("LMK") + b_mode = gemm_b_mode("LNK") + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, mode=a_mode, name="sample_a_tensor" + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, mode=b_mode, name="sample_b_tensor" + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfa_tensor" + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, mode=BLOCK_SCALE_MODE, name="sample_sfb_tensor" + ), + "padded_offsets": self.make_tensor_desc( + sample_padded_offsets, name="sample_padded_offsets" + ), + "alpha_tensor": self.make_tensor_desc( + sample_alpha_tensor, name="sample_alpha_tensor" + ), + "norm_const_tensor": self.make_optional_tensor_desc( + sample_norm_const_tensor, name="sample_norm_const_tensor" + ), + "prob_tensor": self.make_optional_tensor_desc( + sample_prob_tensor, + mode=PROBABILITY_MODE, + name="sample_prob_tensor", + ), + } + self._config = { + "acc_dtype": self.as_optional_dtype(acc_dtype), + "c_dtype": self.as_optional_dtype(c_dtype), + "d_dtype": self.as_optional_dtype(d_dtype), + "output_layout": output_layout, + "mma_tiler_mn": tuple(mma_tiler_mn), + "cluster_shape_mn": ( + None if cluster_shape_mn is None else tuple(cluster_shape_mn) + ), + "sf_vec_size": sf_vec_size, + "vector_f32": vector_f32, + "m_aligned": m_aligned, + "discrete_col_sfd": discrete_col_sfd, + "cluster_overlap_margin": int( + os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0") + ), + } + + self._sample_descs = self.freeze_mapping(self._sample_descs) + self._config = self.freeze_mapping(self._config) + + def _check_support(self) -> None: + resolved = _grouped_gemm_swiglu_impl( + self._sample_descs["a_tensor"], + self._sample_descs["b_tensor"], + self._sample_descs["sfa_tensor"], + self._sample_descs["sfb_tensor"], + self._sample_descs["padded_offsets"], + self._sample_descs["alpha_tensor"], + self._sample_descs["norm_const_tensor"], + self._sample_descs["prob_tensor"], + **self._config, + _validate_only=True, + ) + self._config = self.freeze_mapping({**self._config, **resolved}) + + def __call__( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + ) -> TupleDict: + return super().__call__( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + norm_const_tensor, + prob_tensor, + ) + + def _call_impl( + self, + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + ) -> TupleDict: + values = { + "a_tensor": a_tensor, + "b_tensor": b_tensor, + "sfa_tensor": sfa_tensor, + "sfb_tensor": sfb_tensor, + "padded_offsets": padded_offsets, + "alpha_tensor": alpha_tensor, + "norm_const_tensor": norm_const_tensor, + "prob_tensor": prob_tensor, + } + self.check_tensor_signatures(self._sample_descs, values) + return _grouped_gemm_swiglu_impl(**values, **self._config) + + +@partial( + jax.jit, + static_argnames=( + "acc_dtype", + "c_dtype", + "d_dtype", + "output_layout", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "m_aligned", + "discrete_col_sfd", + ), +) +def grouped_gemm_swiglu_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + padded_offsets: Any, + alpha_tensor: Any, + norm_const_tensor: Optional[Any] = None, + prob_tensor: Optional[Any] = None, + acc_dtype: Any = None, + c_dtype: Any = None, + d_dtype: Any = None, + output_layout: str = "LMN", + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, +) -> TupleDict: + """Compute a contiguous block-scaled grouped GEMM and fused SwiGLU.""" + + return GroupedGemmSwigluSm100( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + norm_const_tensor, + prob_tensor, + acc_dtype=acc_dtype, + c_dtype=c_dtype, + d_dtype=d_dtype, + output_layout=output_layout, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + )( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + padded_offsets, + alpha_tensor, + norm_const_tensor, + prob_tensor, + ) + + +__all__ = [ + "GroupedGemmSwigluSm100", + "grouped_gemm_swiglu_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/__init__.py index 17a7ed70e..01f13bc67 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/__init__.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/__init__.py @@ -1,12 +1,10 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import ( - GroupedGemmWgradSm100, - grouped_gemm_wgrad_wrapper_sm100, -) +from ...common.operation_api import make_operation_api -__all__ = [ - "GroupedGemmWgradSm100", - "grouped_gemm_wgrad_wrapper_sm100", -] +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("GroupedGemmWgradSm100", "grouped_gemm_wgrad_wrapper_sm100")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py index 29b3d7328..d7fd32071 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py @@ -101,6 +101,7 @@ def __init__( stride_order=tuple(i for i, s in sorted(enumerate(self.wgrad_desc.stride[1:]), key=lambda x: x[1])), device=self.wgrad_desc.device, name="single_expert_wgrad", + init_value=self.wgrad_desc.init_value, ) else: # MoEWeightMode.DISCRETE self.expert_cnt = num_experts diff --git a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/jax.py b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/jax.py new file mode 100644 index 000000000..e6b8b9a6c --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/jax.py @@ -0,0 +1,550 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for dense-output grouped GEMM weight gradients on SM100.""" + +from __future__ import annotations + +from functools import partial +import os +from typing import Any, Optional + +import jax +import jax.numpy as jnp + +from .._jax_api import ( + ApiBaseJax, + GROUPED_WORKSPACE_ALIGNMENT, + make_buffer_desc, + TWO_CTA_MMA_TILER_M, + TupleDict, + as_dtype, + as_gemm_tensor_desc, + call_cutedsl, + ceil_div, + grouped_wgrad_workspace_bytes, + normalize_wgrad_input_order, + require_16_byte_extent, + require_array, + require_dtype, + require_grouped_cluster_shape, + require_grouped_mma_tiler, +) + + +def wgrad_scale_shape( + rows: int, + tokens: int, + sf_vec_size: int, +) -> tuple[int, int]: + """Return the packed 2Dx2D scale-buffer shape.""" + + return ( + ceil_div(rows, 128) * 128, + ceil_div(ceil_div(tokens, sf_vec_size), 4) * 4, + ) + + +WGRAD_ALIGNMENT = 16 +WGRAD_A_STRIDE_ORDER = (1, 0) +WGRAD_B_STRIDE_ORDER = (0, 1) +WGRAD_OUTPUT_STRIDE_ORDER = (2, 1, 0) + + +def _launch( + stream, + *args, + acc_dtype: Any, + mma_tiler_mn: tuple[int, int], + cluster_shape_mn: tuple[int, int], + sf_vec_size: int, + accumulate_on_output: bool, + expert_cnt: int, + input_order: str, + has_global_scale: bool, + max_active_clusters: int, +): + from cutlass.jax import jax_to_cutlass_dtype + + from ..moe_utils import MoEWeightMode, WGradInputOrder + from .moe_blockscaled_grouped_gemm_wgrad import ( + BlockScaledMoEGroupedGemmWgradKernel, + ) + + arg_idx = 0 + + def take(): + nonlocal arg_idx + value = args[arg_idx] + arg_idx += 1 + return value + + a_tensor = take() + b_tensor = take() + sfa_tensor = take() + sfb_tensor = take() + offsets_tensor = take() + global_scale_a = take() if has_global_scale else None + global_scale_b = take() if has_global_scale else None + wgrad_tensor = take() + workspace = take() + if arg_idx != len(args): + raise RuntimeError( + f"Unexpected grouped wgrad argument count: consumed {arg_idx}, " + f"received {len(args)}" + ) + + kernel = BlockScaledMoEGroupedGemmWgradKernel( + sf_vec_size=sf_vec_size, + acc_dtype=jax_to_cutlass_dtype(acc_dtype), + use_2cta_instrs=(mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M), + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + accumulate_on_output=accumulate_on_output, + expert_cnt=expert_cnt, + weight_mode=MoEWeightMode.DENSE, + input_order=WGradInputOrder(input_order), + ) + kernel( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + wgrad_tensor, + offsets_tensor, + workspace, + max_active_clusters, + stream, + global_scale_a, + global_scale_b, + None, + ) + + +def _grouped_gemm_wgrad_impl( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + offsets_tensor: Any, + global_scale_a: Optional[Any] = None, + global_scale_b: Optional[Any] = None, + acc_dtype: Any = None, + wgrad_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + accumulate_on_output: bool = False, + input_order: str = "tensor2d", + cluster_overlap_margin: int = 0, + *, + _validate_only: bool = False, +) -> TupleDict | dict[str, Any]: + """Compute dense expert weight gradients using the SM100 grouped kernel. + + ``a_tensor`` has shape ``(hidden,tokens)`` and ``b_tensor`` has shape + ``(tokens,intermediate)``. ``offsets_tensor`` contains each expert's + cumulative token end offset and determines grouping at runtime; its static + length determines the output shape + ``(experts,hidden,intermediate)``. Offset values must be non-decreasing, + aligned to ``sf_vec_size``, and no greater than ``tokens``. + + The JAX surface exposes only the dense output mode; raw per-expert pointer + outputs are not accepted. The output is inferred from the inputs and + allocated by JAX. When ``accumulate_on_output=True``, the reducing epilogue + starts from a zero-initialized output. + + ``input_order='tensor_ragged'`` is supported when A and B have already been + packed expert-by-expert according to the runtime offsets. JAX cannot infer + or perform that value-dependent packing from abstract shapes. + """ + + a_shape = require_array( + a_tensor, + name="a_tensor", + rank=2, + dtype=(jnp.float8_e4m3fn, jnp.float8_e5m2), + ) + ab_dtype = as_dtype(a_tensor) + b_shape = require_array(b_tensor, name="b_tensor", rank=2, dtype=ab_dtype) + hidden, tokens = a_shape + b_tokens, intermediate = b_shape + if b_tokens != tokens: + raise ValueError( + "a_tensor and b_tensor token dimensions must match, got " + f"{a_shape} and {b_shape}" + ) + dimensions = { + "hidden": hidden, + "tokens": tokens, + "intermediate": intermediate, + } + nonpositive = [ + f"{name}={value}" for name, value in dimensions.items() if value <= 0 + ] + if nonpositive: + raise ValueError( + "Grouped-wgrad dimensions must be positive, got " + ", ".join(nonpositive) + ) + + if sf_vec_size != 32: + raise ValueError( + f"The JAX FP8 grouped-wgrad path requires sf_vec_size=32, got {sf_vec_size}" + ) + + require_array( + sfa_tensor, + name="sfa_tensor", + shape=wgrad_scale_shape(hidden, tokens, sf_vec_size), + dtype=jnp.float8_e8m0fnu, + ) + sf_dtype = as_dtype(sfa_tensor) + require_array( + sfb_tensor, + name="sfb_tensor", + shape=wgrad_scale_shape(intermediate, tokens, sf_vec_size), + dtype=sf_dtype, + ) + + offsets_shape = require_array( + offsets_tensor, + name="offsets_tensor", + rank=1, + dtype=jnp.int32, + ) + (expert_cnt,) = offsets_shape + if expert_cnt <= 0: + raise ValueError( + f"offsets_tensor must contain at least one expert, got {expert_cnt}" + ) + output_shape = (expert_cnt, hidden, intermediate) + + has_global_scale = global_scale_a is not None or global_scale_b is not None + if has_global_scale: + if global_scale_a is None or global_scale_b is None: + raise ValueError( + "global_scale_a and global_scale_b must be provided together" + ) + require_array( + global_scale_a, + name="global_scale_a", + shape=(expert_cnt,), + dtype=jnp.float32, + ) + require_array( + global_scale_b, + name="global_scale_b", + shape=(expert_cnt,), + dtype=jnp.float32, + ) + + acc_dtype = require_dtype( + acc_dtype, + (jnp.float32,), + name="acc_dtype", + default=jnp.float32, + ) + valid_wgrad_dtypes = (jnp.bfloat16, jnp.float16, jnp.float32) + wgrad_dtype = require_dtype( + wgrad_dtype, + valid_wgrad_dtypes, + name="wgrad_dtype", + default=jnp.bfloat16, + ) + mma_tiler_mn = require_grouped_mma_tiler(mma_tiler_mn, allowed_m=(128, 256)) + if cluster_shape_mn is None: + cluster_shape_mn = (2, 1) if mma_tiler_mn[0] == TWO_CTA_MMA_TILER_M else (1, 1) + cluster_shape_mn = require_grouped_cluster_shape( + cluster_shape_mn, + mma_tiler_mn=mma_tiler_mn, + ) + input_order_value = normalize_wgrad_input_order(input_order).value + + require_16_byte_extent("a_tensor", tokens, ab_dtype) + require_16_byte_extent("b_tensor", tokens, ab_dtype) + require_16_byte_extent("wgrad_tensor", intermediate, wgrad_dtype) + + if _validate_only: + return { + "acc_dtype": acc_dtype, + "wgrad_dtype": wgrad_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + } + + inputs = [a_tensor, b_tensor, sfa_tensor, sfb_tensor, offsets_tensor] + input_descs = [ + as_gemm_tensor_desc( + "a_tensor", + a_tensor, + public_stride_order=WGRAD_A_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + ), + as_gemm_tensor_desc( + "b_tensor", + b_tensor, + public_stride_order=WGRAD_B_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + ), + as_gemm_tensor_desc( + "sfa_tensor", + sfa_tensor, + public_stride_order=WGRAD_A_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + ), + as_gemm_tensor_desc( + "sfb_tensor", + sfb_tensor, + public_stride_order=WGRAD_A_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + ), + as_gemm_tensor_desc("offsets_tensor", offsets_tensor), + ] + if has_global_scale: + inputs.extend((global_scale_a, global_scale_b)) + input_descs.extend( + ( + as_gemm_tensor_desc("global_scale_a", global_scale_a), + as_gemm_tensor_desc("global_scale_b", global_scale_b), + ) + ) + + (result_wgrad_tensor,) = call_cutedsl( + _launch, + inputs, + input_descs=input_descs, + static_args={ + "acc_dtype": acc_dtype, + "mma_tiler_mn": mma_tiler_mn, + "cluster_shape_mn": cluster_shape_mn, + "sf_vec_size": sf_vec_size, + "accumulate_on_output": bool(accumulate_on_output), + "expert_cnt": expert_cnt, + "input_order": input_order_value, + "has_global_scale": has_global_scale, + "cluster_overlap_margin": int(cluster_overlap_margin), + }, + outputs=( + make_buffer_desc( + "wgrad_tensor", + output_shape, + wgrad_dtype, + public_stride_order=WGRAD_OUTPUT_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + init_value=0.0 if accumulate_on_output else None, + ), + ), + workspaces=( + make_buffer_desc( + "workspace", + ( + grouped_wgrad_workspace_bytes( + expert_cnt, + input_order_value, + ), + ), + jnp.uint8, + ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT, + ), + ), + ) + return TupleDict(wgrad_tensor=result_wgrad_tensor) + + +class GroupedGemmWgradSm100(ApiBaseJax): + """Sample-signature-bound JAX callable for grouped GEMM weight gradients.""" + + def __init__( + self, + sample_a_tensor: Any, + sample_b_tensor: Any, + sample_sfa_tensor: Any, + sample_sfb_tensor: Any, + sample_offsets_tensor: Any, + sample_global_scale_a: Optional[Any] = None, + sample_global_scale_b: Optional[Any] = None, + acc_dtype: Any = None, + wgrad_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + accumulate_on_output: bool = False, + input_order: str = "tensor2d", + ) -> None: + super().__init__() + self._sample_descs = { + "a_tensor": self.make_tensor_desc( + sample_a_tensor, + public_stride_order=WGRAD_A_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + name="sample_a_tensor", + ), + "b_tensor": self.make_tensor_desc( + sample_b_tensor, + public_stride_order=WGRAD_B_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + name="sample_b_tensor", + ), + "sfa_tensor": self.make_tensor_desc( + sample_sfa_tensor, + public_stride_order=WGRAD_A_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + name="sample_sfa_tensor", + ), + "sfb_tensor": self.make_tensor_desc( + sample_sfb_tensor, + public_stride_order=WGRAD_A_STRIDE_ORDER, + ptr_assumed_align=WGRAD_ALIGNMENT, + name="sample_sfb_tensor", + ), + "offsets_tensor": self.make_tensor_desc( + sample_offsets_tensor, name="sample_offsets_tensor" + ), + "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) + + +@partial( + jax.jit, + static_argnames=( + "acc_dtype", + "wgrad_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "accumulate_on_output", + "input_order", + ), +) +def grouped_gemm_wgrad_wrapper_sm100( + a_tensor: Any, + b_tensor: Any, + sfa_tensor: Any, + sfb_tensor: Any, + offsets_tensor: Any, + global_scale_a: Optional[Any] = None, + global_scale_b: Optional[Any] = None, + acc_dtype: Any = None, + wgrad_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[tuple[int, int]] = None, + sf_vec_size: int = 32, + accumulate_on_output: bool = False, + input_order: str = "tensor2d", +) -> TupleDict: + """Compute dense expert weight gradients using the SM100 grouped kernel. + + The output is inferred and allocated by JAX. Set + ``accumulate_on_output=True`` to use the reducing epilogue with a + zero-initialized output. Pointer-table outputs are not supported. + """ + + op = GroupedGemmWgradSm100( + 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/moe_kernel_helpers.py b/python/cudnn/grouped_gemm/moe_kernel_helpers.py index ab41414da..bcf4adbf4 100644 --- a/python/cudnn/grouped_gemm/moe_kernel_helpers.py +++ b/python/cudnn/grouped_gemm/moe_kernel_helpers.py @@ -38,9 +38,12 @@ - Kernel helper functions that don't depend on kernel instance state """ -from typing import Type, Tuple, Union +from __future__ import annotations -import torch +from typing import TYPE_CHECKING, Type, Tuple, Union + +if TYPE_CHECKING: + import torch import cutlass import cutlass.cute as cute @@ -455,6 +458,8 @@ def silu_f32_geglu_scaled(a: Union[float, Float32], fastmath: bool = False) -> U def sigmoid(x): """PyTorch reference sigmoid using exp2 for numerical consistency.""" + import torch + LOG2_E = 1.4426950408889634 exp_x = torch.exp2(x * (-LOG2_E)) ret = 1.0 / (exp_x + 1.0) @@ -471,6 +476,8 @@ def compute_reference_amax(output_tensor: torch.Tensor) -> float: Returns: float: reference amax value """ + import torch + if output_tensor.dtype != torch.float32: output_fp32 = output_tensor.float() else: @@ -500,6 +507,8 @@ def compare_and_report_mismatches( rtol: Relative tolerance max_mismatches: Maximum number of mismatches to report """ + import torch + if gpu_tensor.is_cuda: gpu_data = gpu_tensor.cpu() else: diff --git a/python/cudnn/grouped_gemm/utils.py b/python/cudnn/grouped_gemm/utils.py index 1e58680ef..9302da871 100644 --- a/python/cudnn/grouped_gemm/utils.py +++ b/python/cudnn/grouped_gemm/utils.py @@ -1,5 +1,6 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT + # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions are met: @@ -33,7 +34,12 @@ the forward (grouped_gemm_swiglu) and backward (grouped_gemm_dswiglu) kernels. """ -from typing import Tuple, Union +from __future__ import annotations + +from typing import TYPE_CHECKING, Tuple, Union + +if TYPE_CHECKING: + import torch from cutlass.cutlass_dsl import ( Boolean, @@ -52,7 +58,6 @@ from cutlass.cute.typing import Float32, Int32 import cutlass.cute as cute import cutlass -import torch import cutlass.pipeline as pipeline from cutlass.pipeline import ( Agent, @@ -228,6 +233,8 @@ def warp_redux_sync_fmax( def logical_shape_fp4x2_aware(tensor: torch.Tensor) -> Tuple[int, ...]: """Return correct shapes for NVFP4 tensor.""" + import torch + if tensor.dtype == torch.float4_e2m1fn_x2: innermost_dim_index = next((i for i, s in enumerate(tensor.stride()) if s == 1), None) if innermost_dim_index is None: diff --git a/python/cudnn/jax/__init__.py b/python/cudnn/jax/__init__.py new file mode 100644 index 000000000..c4dc9e379 --- /dev/null +++ b/python/cudnn/jax/__init__.py @@ -0,0 +1,336 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Public optional-dependency boundary for JAX operation APIs. + +JAX and CUTLASS are validated eagerly at this boundary. Individual operation +adapters remain lazy so importing :mod:`cudnn.jax` does not import every CuTe +kernel or any Torch API module. +""" + +from importlib import import_module +from typing import Any + +_INSTALL_HINT = "pip install 'nvidia-cudnn-frontend[jax]'" +_OPTIONAL_DEPENDENCY_PREFIXES = ("cuda", "cutlass", "jax") + +try: + import cutlass.jax as _cutlass_jax + import jax as _jax +except ModuleNotFoundError as error: + missing_name = error.name + if missing_name is None or not any( + missing_name == prefix or missing_name.startswith(f"{prefix}.") + for prefix in _OPTIONAL_DEPENDENCY_PREFIXES + ): + raise + raise ImportError( + f"cuDNN JAX APIs require the optional JAX dependencies. Install them with `{_INSTALL_HINT}`." + ) from error + +if not _cutlass_jax.is_available(): + minimum_version = ".".join( + str(part) for part in _cutlass_jax.CUTE_DSL_MIN_SUPPORTED_JAX_VERSION + ) + installed_version = getattr(_jax, "__version__", "unknown") + raise ImportError( + f"CUTLASS JAX support is unavailable with JAX {installed_version}; " + f"the minimum supported JAX version is {minimum_version}. " + f"Install the optional JAX dependencies with `{_INSTALL_HINT}`." + ) + +from .._jax import JaxApiBase, JaxTensorDesc, TupleDict # noqa: E402 + +_OPERATION_EXPORTS = { + "RmsNormRhtAmaxSm100": ("..rmsnorm_rht_amax.jax", "RmsNormRhtAmaxSm100"), + "rmsnorm_rht_amax_sm100": ("..rmsnorm_rht_amax.jax", "rmsnorm_rht_amax_sm100"), + "GemmSwigluSm100": ("..gemm_swiglu.jax", "GemmSwigluSm100"), + "gemm_swiglu_wrapper_sm100": ("..gemm_swiglu.jax", "gemm_swiglu_wrapper_sm100"), + "GemmSreluSm100": ("..gemm_srelu.jax", "GemmSreluSm100"), + "gemm_srelu_wrapper_sm100": ("..gemm_srelu.jax", "gemm_srelu_wrapper_sm100"), + "GemmDsreluSm100": ("..gemm_dsrelu.jax", "GemmDsreluSm100"), + "gemm_dsrelu_wrapper_sm100": ("..gemm_dsrelu.jax", "gemm_dsrelu_wrapper_sm100"), + "GemmAmaxSm100": ("..gemm_amax.jax", "GemmAmaxSm100"), + "gemm_amax_wrapper_sm100": ("..gemm_amax.jax", "gemm_amax_wrapper_sm100"), + "SparseAttentionBackward": ( + "..deepseek_sparse_attention.sparse_attention_backward.jax", + "SparseAttentionBackward", + ), + "sparse_attention_backward_wrapper": ( + "..deepseek_sparse_attention.sparse_attention_backward.jax", + "sparse_attention_backward_wrapper", + ), + "IndexerForward": ( + "..deepseek_sparse_attention.indexer_forward.jax", + "IndexerForward", + ), + "indexer_forward_wrapper": ( + "..deepseek_sparse_attention.indexer_forward.jax", + "indexer_forward_wrapper", + ), + "IndexerTopK": ("..deepseek_sparse_attention.indexer_top_k.jax", "IndexerTopK"), + "indexer_top_k_wrapper": ( + "..deepseek_sparse_attention.indexer_top_k.jax", + "indexer_top_k_wrapper", + ), + "local_to_global_wrapper": ( + "..deepseek_sparse_attention.indexer_top_k.jax", + "local_to_global_wrapper", + ), + "compactify_wrapper": ( + "..deepseek_sparse_attention.indexer_top_k.jax", + "compactify_wrapper", + ), + "SparseIndexerScoreRecompute": ( + "..deepseek_sparse_attention.score_recompute.jax", + "SparseIndexerScoreRecompute", + ), + "sparse_indexer_score_recompute_wrapper": ( + "..deepseek_sparse_attention.score_recompute.jax", + "sparse_indexer_score_recompute_wrapper", + ), + "SparseAttnScoreRecompute": ( + "..deepseek_sparse_attention.score_recompute.jax", + "SparseAttnScoreRecompute", + ), + "sparse_attn_score_recompute_wrapper": ( + "..deepseek_sparse_attention.score_recompute.jax", + "sparse_attn_score_recompute_wrapper", + ), + "DenseIndexerScoreRecompute": ( + "..deepseek_sparse_attention.score_recompute.jax", + "DenseIndexerScoreRecompute", + ), + "dense_indexer_score_recompute_wrapper": ( + "..deepseek_sparse_attention.score_recompute.jax", + "dense_indexer_score_recompute_wrapper", + ), + "DenseAttnScoreRecompute": ( + "..deepseek_sparse_attention.score_recompute.jax", + "DenseAttnScoreRecompute", + ), + "dense_attn_score_recompute_wrapper": ( + "..deepseek_sparse_attention.score_recompute.jax", + "dense_attn_score_recompute_wrapper", + ), + "IndexerBackward": ( + "..deepseek_sparse_attention.indexer_backward.jax", + "IndexerBackward", + ), + "indexer_backward_wrapper": ( + "..deepseek_sparse_attention.indexer_backward.jax", + "indexer_backward_wrapper", + ), + "DenseIndexerBackward": ( + "..deepseek_sparse_attention.indexer_backward.jax", + "DenseIndexerBackward", + ), + "dense_indexer_backward_wrapper": ( + "..deepseek_sparse_attention.indexer_backward.jax", + "dense_indexer_backward_wrapper", + ), + "GroupedGemmSwigluSm100": ( + "..grouped_gemm.grouped_gemm_swiglu.jax", + "GroupedGemmSwigluSm100", + ), + "grouped_gemm_swiglu_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_swiglu.jax", + "grouped_gemm_swiglu_wrapper_sm100", + ), + "GroupedGemmDswigluSm100": ( + "..grouped_gemm.grouped_gemm_dswiglu.jax", + "GroupedGemmDswigluSm100", + ), + "grouped_gemm_dswiglu_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_dswiglu.jax", + "grouped_gemm_dswiglu_wrapper_sm100", + ), + "GroupedGemmQuantSm100": ( + "..grouped_gemm.grouped_gemm_quant.jax", + "GroupedGemmQuantSm100", + ), + "grouped_gemm_quant_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_quant.jax", + "grouped_gemm_quant_wrapper_sm100", + ), + "GroupedGemmSreluSm100": ( + "..grouped_gemm.grouped_gemm_srelu.jax", + "GroupedGemmSreluSm100", + ), + "grouped_gemm_srelu_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_srelu.jax", + "grouped_gemm_srelu_wrapper_sm100", + ), + "GroupedGemmDsreluSm100": ( + "..grouped_gemm.grouped_gemm_dsrelu.jax", + "GroupedGemmDsreluSm100", + ), + "grouped_gemm_dsrelu_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_dsrelu.jax", + "grouped_gemm_dsrelu_wrapper_sm100", + ), + "GroupedGemmGluSm100": ( + "..grouped_gemm.grouped_gemm_glu.jax", + "GroupedGemmGluSm100", + ), + "grouped_gemm_glu_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_glu.jax", + "grouped_gemm_glu_wrapper_sm100", + ), + "GroupedGemmGluHadamardSm100": ( + "..grouped_gemm.grouped_gemm_glu_hadamard.jax", + "GroupedGemmGluHadamardSm100", + ), + "grouped_gemm_glu_hadamard_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_glu_hadamard.jax", + "grouped_gemm_glu_hadamard_wrapper_sm100", + ), + "GroupedGemmDgluSm100": ( + "..grouped_gemm.grouped_gemm_dglu.jax", + "GroupedGemmDgluSm100", + ), + "grouped_gemm_dglu_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_dglu.jax", + "grouped_gemm_dglu_wrapper_sm100", + ), + "GroupedGemmWgradSm100": ( + "..grouped_gemm.grouped_gemm_wgrad.jax", + "GroupedGemmWgradSm100", + ), + "grouped_gemm_wgrad_wrapper_sm100": ( + "..grouped_gemm.grouped_gemm_wgrad.jax", + "grouped_gemm_wgrad_wrapper_sm100", + ), + "DiscreteGroupedGemmSwigluSm100": ( + "..discrete_grouped_gemm.discrete_grouped_gemm_swiglu.jax", + "DiscreteGroupedGemmSwigluSm100", + ), + "discrete_grouped_gemm_swiglu_wrapper_sm100": ( + "..discrete_grouped_gemm.discrete_grouped_gemm_swiglu.jax", + "discrete_grouped_gemm_swiglu_wrapper_sm100", + ), + "DiscreteGroupedGemmDswigluSm100": ( + "..discrete_grouped_gemm.discrete_grouped_gemm_dswiglu.jax", + "DiscreteGroupedGemmDswigluSm100", + ), + "discrete_grouped_gemm_dswiglu_wrapper_sm100": ( + "..discrete_grouped_gemm.discrete_grouped_gemm_dswiglu.jax", + "discrete_grouped_gemm_dswiglu_wrapper_sm100", + ), + "SelectionAttention": ( + "..native_sparse_attention.selection.jax", + "SelectionAttention", + ), + "selection_attention_wrapper": ( + "..native_sparse_attention.selection.jax", + "selection_attention_wrapper", + ), + "CompressionAttention": ( + "..native_sparse_attention.compression.jax", + "CompressionAttention", + ), + "compression_attention_wrapper": ( + "..native_sparse_attention.compression.jax", + "compression_attention_wrapper", + ), + "SlidingWindowAttention": ( + "..native_sparse_attention.sliding_window_attention.jax", + "SlidingWindowAttention", + ), + "sliding_window_attention_wrapper": ( + "..native_sparse_attention.sliding_window_attention.jax", + "sliding_window_attention_wrapper", + ), + "TopKReduction": ("..native_sparse_attention.top_k.jax", "TopKReduction"), + "topk_reduction_wrapper": ( + "..native_sparse_attention.top_k.jax", + "topk_reduction_wrapper", + ), + "SdpafwdSm100D256": ("..sdpa.fwd.jax", "SdpafwdSm100D256"), + "sdpa_fwd_wrapper_sm100_d256": ("..sdpa.fwd.jax", "sdpa_fwd_wrapper_sm100_d256"), + "SdpabwdSm100D256": ("..sdpa.bwd.jax", "SdpabwdSm100D256"), + "sdpa_bwd_wrapper_sm100_d256": ("..sdpa.bwd.jax", "sdpa_bwd_wrapper_sm100_d256"), + "BlockSparseAttentionForward": ( + "..block_sparse_attention.jax", + "BlockSparseAttentionForward", + ), + "block_sparse_attention_forward": ( + "..block_sparse_attention.jax", + "block_sparse_attention_forward", + ), + "BlockSparseAttentionBackward": ( + "..block_sparse_attention.jax", + "BlockSparseAttentionBackward", + ), + "block_sparse_attention_backward": ( + "..block_sparse_attention.jax", + "block_sparse_attention_backward", + ), +} + +_DSA_EXPORTS = frozenset( + name + for name, (module_name, _) in _OPERATION_EXPORTS.items() + if module_name.startswith("..deepseek_sparse_attention.") +) +_NSA_EXPORTS = frozenset( + name + for name, (module_name, _) in _OPERATION_EXPORTS.items() + if module_name.startswith("..native_sparse_attention.") +) +_BSA_EXPORTS = frozenset( + name + for name, (module_name, _) in _OPERATION_EXPORTS.items() + if module_name.startswith("..block_sparse_attention.") +) + + +def _load_operation(name: str) -> Any: + module_name, symbol_name = _OPERATION_EXPORTS[name] + value = getattr(import_module(module_name, __name__), symbol_name) + globals()[name] = value + return value + + +class _OperationNamespace: + """Lazy view of one JAX operation family.""" + + def __init__(self, qualified_name: str, exports: frozenset[str]) -> None: + self._qualified_name = qualified_name + self._exports = exports + + def __getattr__(self, name: str) -> Any: + if name not in self._exports: + raise AttributeError(f"{self._qualified_name} has no attribute {name!r}") + value = _load_operation(name) + setattr(self, name, value) + return value + + def __dir__(self) -> list[str]: + return sorted((*vars(self), *self._exports)) + + +DSA = _OperationNamespace("cudnn.jax.DSA", _DSA_EXPORTS) +NSA = _OperationNamespace("cudnn.jax.NSA", _NSA_EXPORTS) +BSA = _OperationNamespace("cudnn.jax.BSA", _BSA_EXPORTS) + + +def __getattr__(name: str) -> Any: + if name in _OPERATION_EXPORTS: + return _load_operation(name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__() -> list[str]: + return sorted((*globals(), *_OPERATION_EXPORTS)) + + +__all__ = [ + "BSA", + "DSA", + "JaxApiBase", + "JaxTensorDesc", + "NSA", + "TupleDict", + *_OPERATION_EXPORTS, +] diff --git a/python/cudnn/native_sparse_attention/__init__.py b/python/cudnn/native_sparse_attention/__init__.py index 6d4f8220b..fcb30fdf3 100644 --- a/python/cudnn/native_sparse_attention/__init__.py +++ b/python/cudnn/native_sparse_attention/__init__.py @@ -1,28 +1,37 @@ -from .selection import SelectionAttention, selection_attention_wrapper -from .compression import CompressionAttention, compression_attention_wrapper -from .sliding_window_attention import ( - SlidingWindowAttention, - sliding_window_attention_wrapper, +"""Lazy Torch namespace for native sparse-attention operations.""" + +from typing import Any + +from ..common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "selection": ("SelectionAttention", "selection_attention_wrapper"), + "compression": ("CompressionAttention", "compression_attention_wrapper"), + "sliding_window_attention": ( + "SlidingWindowAttention", + "sliding_window_attention_wrapper", + ), + "top_k": ("TopKReduction", "topk_reduction_wrapper"), + }, ) -from .top_k import TopKReduction, topk_reduction_wrapper +_OPERATION_EXPORTS = frozenset(__all__) class NSANamespace: - SelectionAttention = staticmethod(SelectionAttention) - selection_attention_wrapper = staticmethod(selection_attention_wrapper) + def __getattr__(self, name: str) -> Any: + if name not in _OPERATION_EXPORTS: + raise AttributeError(f"NSA has no attribute {name!r}") + value = __getattr__(name) + setattr(self, name, value) + return value - SlidingWindowAttention = staticmethod(SlidingWindowAttention) - sliding_window_attention_wrapper = staticmethod(sliding_window_attention_wrapper) - - CompressionAttention = staticmethod(CompressionAttention) - compression_attention_wrapper = staticmethod(compression_attention_wrapper) - - TopKReduction = staticmethod(TopKReduction) - topk_reduction_wrapper = staticmethod(topk_reduction_wrapper) + def __dir__(self) -> list[str]: + return sorted((*vars(self), *_OPERATION_EXPORTS)) NSA = NSANamespace() -__all__ = [ - "NSA", -] +# Preserve the historical root contract: operations are accessed through NSA. +__all__ = ["NSA"] diff --git a/python/cudnn/native_sparse_attention/compression/__init__.py b/python/cudnn/native_sparse_attention/compression/__init__.py index b44780956..75c7e5c34 100644 --- a/python/cudnn/native_sparse_attention/compression/__init__.py +++ b/python/cudnn/native_sparse_attention/compression/__init__.py @@ -1,6 +1,9 @@ -from .api import CompressionAttention, compression_attention_wrapper +"""Lazy Torch compression-attention API exports.""" -__all__ = [ - "CompressionAttention", - "compression_attention_wrapper", -] +from ...common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("CompressionAttention", "compression_attention_wrapper")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/native_sparse_attention/compression/fmha.py b/python/cudnn/native_sparse_attention/compression/fmha.py index a28eede15..20878b5da 100644 --- a/python/cudnn/native_sparse_attention/compression/fmha.py +++ b/python/cudnn/native_sparse_attention/compression/fmha.py @@ -104,6 +104,7 @@ def __init__( mma_tiler: Tuple[int, int, int], is_persistent: bool, mask_type: fmha_utils.MaskType, + persistent_sm_count: Optional[int] = None, ): """Initializes the configuration for a Blackwell Fused Multi-Head Attention (FMHA) kernel. @@ -155,6 +156,7 @@ def __init__( self.cluster_shape_mn = (1, 1) self.is_persistent = is_persistent self.mask_type = mask_type + self.persistent_sm_count = persistent_sm_count self.softmax0_warp_ids = (0, 1, 2, 3) self.softmax1_warp_ids = (4, 5, 6, 7) @@ -355,6 +357,7 @@ def __call__( cute.shape((s_q, d, ((h_r, h_k), b))), self.cta_tiler, self.is_persistent, + self.persistent_sm_count, ) self.q_major_mode = utils.LayoutEnum.from_tensor(q).mma_major_mode() diff --git a/python/cudnn/native_sparse_attention/compression/fmha_helpers.py b/python/cudnn/native_sparse_attention/compression/fmha_helpers.py index 51627cb80..cb6753cce 100644 --- a/python/cudnn/native_sparse_attention/compression/fmha_helpers.py +++ b/python/cudnn/native_sparse_attention/compression/fmha_helpers.py @@ -137,6 +137,7 @@ def __init__( def get_grid_shape( params: FmhaStaticTileSchedulerParams, *, + persistent_sm_count: Optional[int] = None, loc=None, ip=None, ) -> cute.Shape: @@ -154,10 +155,13 @@ def get_grid_shape( :rtype: cute.Shape """ if params.is_persistent: - hardware_info = HardwareInfo() - sm_count = hardware_info.get_device_multiprocessor_count() + if persistent_sm_count is None: + hardware_info = HardwareInfo() + persistent_sm_count = hardware_info.get_device_multiprocessor_count() + if persistent_sm_count <= 0: + raise ValueError(f"persistent_sm_count must be positive, got {persistent_sm_count}") return ( - min(sm_count, cute.size(params.problem_shape_mbh, loc=loc, ip=ip)), + min(persistent_sm_count, cute.size(params.problem_shape_mbh, loc=loc, ip=ip)), 1, 1, ) @@ -297,6 +301,7 @@ def compute_grid( o_shape: cute.Shape, cta_tiler: Tuple[int, int, int], is_persistent: bool, + persistent_sm_count: Optional[int] = None, ) -> Tuple[FmhaStaticTileSchedulerParams, Tuple[int, int, int]]: """ Compute grid parameters for FMHA operation. @@ -330,7 +335,10 @@ def compute_grid( cute.size(o_shape[2][1]), ), ) - grid = FmhaStaticTileScheduler.get_grid_shape(tile_sched_params) + grid = FmhaStaticTileScheduler.get_grid_shape( + tile_sched_params, + persistent_sm_count=persistent_sm_count, + ) return tile_sched_params, grid diff --git a/python/cudnn/native_sparse_attention/compression/jax.py b/python/cudnn/native_sparse_attention/compression/jax.py new file mode 100644 index 000000000..5215e515d --- /dev/null +++ b/python/cudnn/native_sparse_attention/compression/jax.py @@ -0,0 +1,536 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for fixed BHSD/BSHD and packed THD NSA compression attention.""" + +from __future__ import annotations + +from functools import partial +import math +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._jax.compiler import compile_options_for_target +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..jax_utils import ( + BHS_TO_BSH_MODE, + FIXED_LAYOUTS, + describe_fixed_data, + fixed_data_mode, + make_fixed_output, + normalize_attention_layout, + normalize_supported_dtype, + require_fixed_qkv, +) + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) + + +class CompressionAttention(JaxApiBase): + """JAX callable specialized from fixed BHSD/BSHD or packed THD metadata.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_cum_seqlen_q: Any | None = None, + sample_cum_seqlen_k: Any | None = None, + enable_lse: bool = False, + o_dtype: Any = None, + qk_acc_dtype: Any = None, + pv_acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + is_persistent: bool = False, + scale_q: float = 1.0, + scale_k: float = 1.0, + scale_v: float = 1.0, + inv_scale_o: float = 1.0, + scale_softmax: float | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + input_dtypes = ( + data_type.HALF, + data_type.BFLOAT16, + data_type.FP8_E4M3, + ) + ranks = tuple(len(sample.shape) for sample in (sample_q, sample_k, sample_v)) + if len(set(ranks)) != 1: + raise ValueError(f"Q, K, and V must all use the same rank, got {ranks}") + self.input_layout = normalize_attention_layout(layout, ranks[0]) + if self.input_layout in FIXED_LAYOUTS: + self.data_mode = fixed_data_mode(self.input_layout, kernel_axes="BSHD") + if sample_cum_seqlen_q is not None or sample_cum_seqlen_k is not None: + raise ValueError( + "cumulative sequence lengths are only valid for packed THD inputs" + ) + if max_s_q is not None or max_s_k is not None: + raise ValueError("max_s_q and max_s_k are only valid for packed THD") + self.q_desc = describe_fixed_data( + sample_q, + "sample_q", + layout=self.input_layout, + kernel_axes="BSHD", + ) + self.k_desc = describe_fixed_data( + sample_k, + "sample_k", + layout=self.input_layout, + kernel_axes="BSHD", + ) + self.v_desc = describe_fixed_data( + sample_v, + "sample_v", + layout=self.input_layout, + kernel_axes="BSHD", + ) + ( + self.batch, + self.num_query_heads, + self.num_kv_heads, + self.seqlen_q, + self.seqlen_k, + self.head_dim, + ) = require_fixed_qkv( + self.q_desc, + self.k_desc, + self.v_desc, + operation_name="CompressionAttention", + kernel_axes="BSHD", + input_dtypes=input_dtypes, + ) + self.cum_q_desc = self.cum_k_desc = None + self.lse_extent = self.seqlen_q + self.q_kernel_desc = self.q_desc + self.k_kernel_desc = self.k_desc + self.v_kernel_desc = self.v_desc + elif self.input_layout == "THD": + self.data_mode = None + if len(sample_k.shape) != 3 or len(sample_v.shape) != 3: + raise ValueError("Q, K, and V must all use the same rank") + if sample_cum_seqlen_q is None or sample_cum_seqlen_k is None: + raise ValueError( + "packed THD inputs require cumulative Q and K sequence lengths" + ) + if max_s_q is None or max_s_k is None: + raise ValueError("packed THD inputs require max_s_q and max_s_k") + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.k_desc = self._to_tensor_desc(sample_k, "sample_k") + self.v_desc = self._to_tensor_desc(sample_v, "sample_v") + total_q, self.num_query_heads, self.head_dim = self.q_desc.shape + total_k, self.num_kv_heads, k_head_dim = self.k_desc.shape + v_total, value_heads, value_dim = self.v_desc.shape + if min(total_q, total_k, self.num_query_heads, self.num_kv_heads) <= 0: + raise ValueError("packed THD dimensions must be positive") + if (v_total, value_heads, value_dim) != ( + total_k, + self.num_kv_heads, + self.head_dim, + ) or k_head_dim != self.head_dim: + raise ValueError("packed K and V shapes must match Q/K head metadata") + if self.q_desc.cudnn_dtype not in input_dtypes: + raise ValueError(f"unsupported Q dtype {self.q_desc.dtype}") + if ( + self.k_desc.cudnn_dtype != self.q_desc.cudnn_dtype + or self.v_desc.cudnn_dtype != self.q_desc.cudnn_dtype + ): + raise ValueError("Q, K, and V must have the same dtype") + if self.head_dim not in (32, 64, 128): + raise ValueError("head dimension must be 32, 64, or 128") + if self.num_query_heads % self.num_kv_heads: + raise ValueError("H_q must be divisible by H_kv") + self.cum_q_desc = self._to_tensor_desc( + sample_cum_seqlen_q, "sample_cum_seqlen_q" + ) + self.cum_k_desc = self._to_tensor_desc( + sample_cum_seqlen_k, "sample_cum_seqlen_k" + ) + if ( + self.cum_q_desc.ndim != 1 + or self.cum_q_desc.shape != self.cum_k_desc.shape + or self.cum_q_desc.shape[0] < 2 + ): + raise ValueError( + "cumulative Q and K sequence lengths must have matching " + "(B + 1,) shapes" + ) + if ( + self.cum_q_desc.cudnn_dtype not in (data_type.INT32, data_type.INT64) + or self.cum_k_desc.cudnn_dtype != self.cum_q_desc.cudnn_dtype + ): + raise ValueError( + "cumulative sequence lengths must share int32 or int64 dtype" + ) + self.batch = self.cum_q_desc.shape[0] - 1 + self.seqlen_q = int(max_s_q) + self.seqlen_k = int(max_s_k) + self.lse_extent = total_q + if ( + self.seqlen_q <= 0 + or self.seqlen_k <= 0 + or self.seqlen_q > total_q + or self.seqlen_k > total_k + ): + raise ValueError( + "max_s_q and max_s_k must be positive and no larger than " + "their packed token counts" + ) + self.q_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_q.shape), + sample_q.dtype, + name="q_tensor", + public_stride_order=(3, 2, 0, 1), + ) + self.k_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_k.shape), + sample_k.dtype, + name="k_tensor", + public_stride_order=(3, 2, 0, 1), + ) + self.v_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_v.shape), + sample_v.dtype, + name="v_tensor", + public_stride_order=(3, 2, 0, 1), + ) + if self.seqlen_q < self.seqlen_k or self.seqlen_q % self.seqlen_k: + raise ValueError( + "Compression attention requires S_q to be an integer multiple of " + f"S_k, got S_q={self.seqlen_q} and S_k={self.seqlen_k}" + ) + if tuple(mma_tiler_mn) != (128, 128): + raise ValueError(f"mma_tiler_mn must be (128, 128), got {mma_tiler_mn}") + + self.enable_lse = bool(enable_lse) + self.is_persistent = bool(is_persistent) + self.output_dtype = normalize_supported_dtype( + o_dtype, + sample_q.dtype, + "o_dtype", + (jnp.float16, jnp.bfloat16, jnp.float8_e4m3fn), + ) + normalize_supported_dtype( + qk_acc_dtype, jnp.float32, "qk_acc_dtype", (jnp.float32,) + ) + normalize_supported_dtype( + pv_acc_dtype, jnp.float32, "pv_acc_dtype", (jnp.float32,) + ) + base_scale = ( + 1.0 / math.sqrt(self.head_dim) + if scale_softmax is None + else float(scale_softmax) + ) + self.scale_softmax = float(scale_q) * float(scale_k) * base_scale + self.scale_output = float(scale_v) * float(inv_scale_o) + self.target_compute_capability = target_compute_capability + self.compute_capability: int | None = None + self.persistent_sm_count: int | None = None + + if self.input_layout in FIXED_LAYOUTS: + self.o_desc = make_fixed_output( + tuple(sample_q.shape), + self.output_dtype, + "o_tensor", + layout=self.input_layout, + kernel_axes="BSHD", + ) + self.o_kernel_desc = self.o_desc + self.lse_desc = ( + JaxTensorDesc.from_shape( + (self.batch, self.num_query_heads, self.seqlen_q), + jnp.float32, + name="lse_tensor", + mode=BHS_TO_BSH_MODE, + ) + if self.enable_lse + else None + ) + self.lse_kernel_desc = self.lse_desc + else: + output_shape = ( + self.q_desc.shape[0], + self.num_query_heads, + self.head_dim, + ) + self.o_desc = JaxTensorDesc.from_shape( + output_shape, self.output_dtype, name="o_tensor" + ) + self.o_kernel_desc = JaxTensorDesc.from_shape( + (1, *output_shape), + self.output_dtype, + name="o_kernel_tensor", + public_stride_order=(3, 2, 0, 1), + ) + self.lse_desc = ( + JaxTensorDesc.from_shape( + output_shape[:2], + jnp.float32, + name="lse_tensor", + ) + if self.enable_lse + else None + ) + self.lse_kernel_desc = ( + JaxTensorDesc.from_shape( + (1, *output_shape[:2]), + jnp.float32, + name="lse_kernel_tensor", + ) + if self.enable_lse + else None + ) + + def check_support(self) -> bool: + self.compute_capability = self._resolve_compute_capability( + self.target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "CompressionAttention", + ) + if self.is_persistent: + self.persistent_sm_count = self._get_device_multiprocessor_count() + return True + + def __call__( + self, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + ) -> TupleDict: + self.check_support() + if self.input_layout in FIXED_LAYOUTS: + if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: + raise ValueError( + "cumulative sequence lengths must be omitted for fixed inputs" + ) + inputs = (q_tensor, k_tensor, v_tensor) + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.v_kernel_desc, + ) + launch = self._launch_kernel + else: + if cum_seqlen_q_tensor is None or cum_seqlen_k_tensor is None: + raise ValueError( + "packed THD inputs require cumulative Q and K sequence lengths" + ) + for value, desc in ( + (q_tensor, self.q_desc), + (k_tensor, self.k_desc), + (v_tensor, self.v_desc), + ): + self._check_tensor_signature(value, desc) + q_tensor = jnp.reshape(q_tensor, self.q_kernel_desc.shape) + k_tensor = jnp.reshape(k_tensor, self.k_kernel_desc.shape) + v_tensor = jnp.reshape(v_tensor, self.v_kernel_desc.shape) + inputs = ( + q_tensor, + k_tensor, + v_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.v_kernel_desc, + self.cum_q_desc, + self.cum_k_desc, + ) + launch = self._launch_varlen_kernel + + results = self._call_kernel( + inputs, + launch=launch, + output_descs=(self.o_kernel_desc,) + if self.lse_kernel_desc is None + else (self.o_kernel_desc, self.lse_kernel_desc), + input_descs=input_descs, + compile_options=compile_options_for_target(self.compute_capability), + ) + output = results[0] + lse = results[1] if self.enable_lse else None + if self.input_layout == "THD": + output = jnp.reshape(output, self.o_desc.shape) + if lse is not None: + lse = jnp.reshape(lse, self.lse_desc.shape) + return TupleDict( + o_tensor=output, + lse_tensor=lse, + ) + + def _launch_kernel( + self, stream: Any, q: Any, k: Any, v: Any, output: Any, *optional: Any + ) -> None: + lse = optional[0] if optional else None + self._run_kernel(stream, q, k, v, output, lse, None, None) + + def _launch_varlen_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + cum_seqlen_q: Any, + cum_seqlen_k: Any, + output: Any, + *optional: Any, + ) -> None: + lse = optional[0] if optional else None + self._run_kernel( + stream, + q, + k, + v, + output, + lse, + cum_seqlen_q, + cum_seqlen_k, + ) + + def _run_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + output: Any, + lse: Any, + cum_seqlen_q: Any, + cum_seqlen_k: Any, + ) -> None: + from cutlass import Float32, Int32 + + from .fmha import BlackwellFusedMultiHeadAttentionForward + from .fmha_helpers import MaskType + + kernel = BlackwellFusedMultiHeadAttentionForward( + qk_acc_dtype=Float32, + pv_acc_dtype=Float32, + mma_tiler=(128, 128, self.head_dim), + is_persistent=self.is_persistent, + mask_type=MaskType.COMPRESSED_CAUSAL_MASK, + persistent_sm_count=self.persistent_sm_count, + ) + problem_size = tuple( + Int32(value) + for value in ( + self.batch, + self.seqlen_q, + self.lse_extent, + self.seqlen_k, + self.num_query_heads, + self.num_kv_heads, + self.head_dim, + ) + ) + kernel( + q, + k, + v, + output, + problem_size, + cum_seqlen_q, + cum_seqlen_k, + lse, + Float32(self.scale_softmax * math.log2(math.e)), + Float32(self.scale_softmax), + Float32(self.scale_output), + None, + Int32(0), + stream, + ) + + +@partial( + jax.jit, + static_argnames=( + "enable_lse", + "o_dtype", + "qk_acc_dtype", + "pv_acc_dtype", + "mma_tiler_mn", + "is_persistent", + "scale_q", + "scale_k", + "scale_v", + "inv_scale_o", + "scale_softmax", + "max_s_q", + "max_s_k", + "layout", + "target_compute_capability", + ), +) +def compression_attention_wrapper( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + enable_lse: bool = False, + o_dtype: Any = None, + qk_acc_dtype: Any = None, + pv_acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + is_persistent: bool = False, + scale_q: float = 1.0, + scale_k: float = 1.0, + scale_v: float = 1.0, + inv_scale_o: float = 1.0, + scale_softmax: float | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + layout: str | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute fixed BHSD/BSHD or packed THD compression attention. + + Packed THD calls must provide both cumulative sequence-length arrays and + static ``max_s_q``/``max_s_k`` bounds because JAX traces array values. + Fixed outputs follow ``layout``; LSE remains ``(B, H, S)``. + """ + + return CompressionAttention( + q_tensor, + k_tensor, + v_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + enable_lse=enable_lse, + o_dtype=o_dtype, + qk_acc_dtype=qk_acc_dtype, + pv_acc_dtype=pv_acc_dtype, + mma_tiler_mn=mma_tiler_mn, + is_persistent=is_persistent, + scale_q=scale_q, + scale_k=scale_k, + scale_v=scale_v, + inv_scale_o=inv_scale_o, + scale_softmax=scale_softmax, + max_s_q=max_s_q, + max_s_k=max_s_k, + layout=layout, + target_compute_capability=target_compute_capability, + )( + q_tensor, + k_tensor, + v_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + + +__all__ = [ + "CompressionAttention", + "SUPPORTED_COMPUTE_CAPABILITIES", + "compression_attention_wrapper", +] diff --git a/python/cudnn/native_sparse_attention/jax_utils.py b/python/cudnn/native_sparse_attention/jax_utils.py new file mode 100644 index 000000000..e4dd4c367 --- /dev/null +++ b/python/cudnn/native_sparse_attention/jax_utils.py @@ -0,0 +1,212 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared tensor metadata and validation for JAX NSA adapters.""" + +from __future__ import annotations + +from typing import Any + +import jax.numpy as jnp + +from .. import data_type +from .._jax import JaxApiBase, JaxTensorDesc +from .._jax.datatypes import normalize_jax_dtype +from .._jax.layout import mode_from_layout, stride_order_to_public, to_public_axes + +FIXED_LAYOUTS = ("BHSD", "BSHD") +BHS_TO_BSH_MODE = mode_from_layout("BHS", kernel_axes="BSH") + + +def normalize_attention_layout( + layout: str | None, + rank: int, + *, + allow_packed: bool = True, +) -> str: + """Normalize a fixed BHSD/BSHD or packed THD public layout.""" + + if layout is None: + if rank == 4: + return "BHSD" + if allow_packed and rank == 3: + return "THD" + raise ValueError(f"attention layout cannot be inferred from rank {rank}") + if not isinstance(layout, str): + raise TypeError(f"layout must be a string or None, got {layout!r}") + normalized = "".join( + character for character in layout.upper() if character.isalpha() + ) + allowed = (*FIXED_LAYOUTS, "THD") if allow_packed else FIXED_LAYOUTS + if normalized not in allowed: + choices = ", ".join(allowed) + raise ValueError(f"layout must be one of {choices}, got {layout!r}") + expected_rank = 3 if normalized == "THD" else 4 + if rank != expected_rank: + raise ValueError( + f"{normalized} layout requires rank-{expected_rank} tensors, got rank {rank}" + ) + return normalized + + +def fixed_data_mode(layout: str, *, kernel_axes: str) -> tuple[int, ...]: + if layout not in FIXED_LAYOUTS: + raise ValueError(f"fixed layout must be one of {FIXED_LAYOUTS}, got {layout!r}") + return mode_from_layout(layout, kernel_axes=kernel_axes) + + +def describe_fixed_data( + sample: Any, + name: str, + *, + layout: str, + kernel_axes: str, + kernel_stride_order: tuple[int, ...] | None = None, + init_value: bool | int | float | None = None, +) -> JaxTensorDesc: + """Describe public fixed-shape metadata in canonical kernel axes.""" + + mode = fixed_data_mode(layout, kernel_axes=kernel_axes) + public_stride_order = ( + None + if kernel_stride_order is None + else stride_order_to_public(kernel_stride_order, mode) + ) + return JaxApiBase._to_tensor_desc( + sample, + name, + mode=mode, + public_stride_order=public_stride_order, + init_value=init_value, + ) + + +def describe_bhs_as_bsh( + sample: Any, + name: str, + *, + init_value: bool | int | float | None = None, +) -> JaxTensorDesc: + """Describe public BHS metadata as the BSH modes consumed by a kernel.""" + + return JaxApiBase._to_tensor_desc( + sample, + name, + mode=BHS_TO_BSH_MODE, + init_value=init_value, + ) + + +def make_fixed_output( + public_shape: tuple[int, ...], + dtype: Any, + name: str, + *, + layout: str, + kernel_axes: str, + kernel_stride_order: tuple[int, ...] | None = None, + init_value: bool | int | float | None = None, +) -> JaxTensorDesc: + mode = fixed_data_mode(layout, kernel_axes=kernel_axes) + public_stride_order = ( + None + if kernel_stride_order is None + else stride_order_to_public(kernel_stride_order, mode) + ) + return JaxTensorDesc.from_shape( + public_shape, + dtype, + name=name, + mode=mode, + public_stride_order=public_stride_order, + init_value=init_value, + ) + + +def require_fixed_qkv( + q_desc: JaxTensorDesc, + k_desc: JaxTensorDesc, + v_desc: JaxTensorDesc | None = None, + *, + operation_name: str, + head_dims: tuple[int, ...] = (32, 64, 128), + kernel_axes: str = "BHSD", + input_dtypes: tuple[data_type, ...] = (data_type.HALF, data_type.BFLOAT16), +) -> tuple[int, int, int, int, int, int]: + """Validate fixed-shape public BHSD Q/K/V signatures.""" + + for desc in (q_desc, k_desc, v_desc): + if desc is not None and desc.ndim != 4: + raise ValueError( + f"{desc.name} must have rank 4 (B, H, S, D), got {desc.shape}" + ) + if q_desc.cudnn_dtype not in input_dtypes: + expected = ", ".join(str(dtype) for dtype in input_dtypes) + raise ValueError( + f"{operation_name} requires one of {{{expected}}}, got {q_desc.dtype}" + ) + for desc in (k_desc, v_desc): + if desc is not None and desc.cudnn_dtype != q_desc.cudnn_dtype: + raise ValueError(f"{desc.name} must have the same dtype as {q_desc.name}") + + logical_bhsd_mode = mode_from_layout("BHSD", kernel_axes=kernel_axes) + q_shape = to_public_axes(q_desc.shape, logical_bhsd_mode) + k_shape = to_public_axes(k_desc.shape, logical_bhsd_mode) + batch, num_query_heads, seqlen_q, head_dim = q_shape + k_batch, num_kv_heads, seqlen_k, k_head_dim = k_shape + dimensions = (batch, num_query_heads, num_kv_heads, seqlen_q, seqlen_k, head_dim) + if any(value <= 0 for value in dimensions): + raise ValueError( + f"{operation_name} dimensions must be positive, got {dimensions}" + ) + if k_batch != batch: + raise ValueError( + f"Q and K batch dimensions must match, got {batch} and {k_batch}" + ) + if k_head_dim != head_dim: + raise ValueError( + f"Q and K head dimensions must match, got {head_dim} and {k_head_dim}" + ) + if head_dim not in head_dims: + expected = ", ".join(str(value) for value in head_dims) + raise ValueError(f"head dimension must be one of {expected}, got {head_dim}") + if num_query_heads % num_kv_heads: + raise ValueError( + f"H_q ({num_query_heads}) must be divisible by H_kv ({num_kv_heads})" + ) + + if v_desc is not None: + v_batch, num_value_heads, v_seqlen, value_dim = to_public_axes( + v_desc.shape, logical_bhsd_mode + ) + if (v_batch, num_value_heads, v_seqlen) != (batch, num_kv_heads, seqlen_k): + raise ValueError("K and V batch, head, and sequence dimensions must match") + if value_dim != head_dim: + raise ValueError( + f"V head dimension must match Q/K ({head_dim}), got {value_dim}" + ) + return dimensions + + +def normalize_supported_dtype( + value: Any | None, default: Any, name: str, allowed: tuple[Any, ...] +) -> Any: + dtype = normalize_jax_dtype(value, default, name) + allowed_dtypes = tuple(jnp.dtype(item) for item in allowed) + if dtype not in allowed_dtypes: + expected = ", ".join(str(item) for item in allowed_dtypes) + raise ValueError(f"{name} must be one of {expected}, got {dtype}") + return dtype + + +__all__ = [ + "BHS_TO_BSH_MODE", + "FIXED_LAYOUTS", + "describe_bhs_as_bsh", + "describe_fixed_data", + "fixed_data_mode", + "make_fixed_output", + "normalize_attention_layout", + "normalize_supported_dtype", + "require_fixed_qkv", +] diff --git a/python/cudnn/native_sparse_attention/selection/__init__.py b/python/cudnn/native_sparse_attention/selection/__init__.py index 576723053..e61451af3 100644 --- a/python/cudnn/native_sparse_attention/selection/__init__.py +++ b/python/cudnn/native_sparse_attention/selection/__init__.py @@ -1,6 +1,9 @@ -from .api import SelectionAttention, selection_attention_wrapper +"""Lazy Torch selection-attention API exports.""" -__all__ = [ - "SelectionAttention", - "selection_attention_wrapper", -] +from ...common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("SelectionAttention", "selection_attention_wrapper")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/native_sparse_attention/selection/jax.py b/python/cudnn/native_sparse_attention/selection/jax.py new file mode 100644 index 000000000..381a50961 --- /dev/null +++ b/python/cudnn/native_sparse_attention/selection/jax.py @@ -0,0 +1,338 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for packed selection attention.""" + +from __future__ import annotations + +from functools import partial +import math +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._jax.compiler import compile_options_for_target +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..jax_utils import normalize_supported_dtype + +SUPPORTED_COMPUTE_CAPABILITIES = (90, 100, 103, 107) + + +class SelectionAttention(JaxApiBase): + """JAX callable specialized from packed THD selection metadata.""" + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_block_indices: Any, + sample_block_counts: Any, + sample_cum_seqlen: Any, + *, + max_s_q: int, + max_s_k: int, + block_size: int = 64, + scale_softmax: float | None = None, + o_dtype: Any = None, + acc_dtype: Any = None, + target_compute_capability: int | None = None, + ) -> None: + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.k_desc = self._to_tensor_desc(sample_k, "sample_k") + self.v_desc = self._to_tensor_desc(sample_v, "sample_v") + self.block_indices_desc = self._to_tensor_desc( + sample_block_indices, "sample_block_indices" + ) + self.block_counts_desc = self._to_tensor_desc( + sample_block_counts, "sample_block_counts" + ) + self.cum_seqlen_desc = self._to_tensor_desc( + sample_cum_seqlen, "sample_cum_seqlen" + ) + + if self.q_desc.ndim != 3 or self.k_desc.ndim != 3 or self.v_desc.ndim != 3: + raise ValueError( + "SelectionAttention requires packed rank-3 (T, H, D) Q, K, and V" + ) + if self.q_desc.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError( + f"Q must have dtype float16 or bfloat16, got {self.q_desc.dtype}" + ) + if ( + self.k_desc.cudnn_dtype != self.q_desc.cudnn_dtype + or self.v_desc.cudnn_dtype != self.q_desc.cudnn_dtype + ): + raise ValueError("Q, K, and V must have the same dtype") + + self.total_tokens, self.num_query_heads, self.head_dim = self.q_desc.shape + k_tokens, self.num_kv_heads, k_head_dim = self.k_desc.shape + v_tokens, num_value_heads, self.value_dim = self.v_desc.shape + dimensions = ( + self.total_tokens, + self.num_query_heads, + self.num_kv_heads, + self.head_dim, + self.value_dim, + ) + if any(value <= 0 for value in dimensions): + raise ValueError( + f"SelectionAttention dimensions must be positive, got {dimensions}" + ) + if (k_tokens, v_tokens) != (self.total_tokens, self.total_tokens): + raise ValueError("Q, K, and V must have the same packed token count") + if k_head_dim != self.head_dim: + raise ValueError("Q and K head dimensions must match") + if num_value_heads != self.num_kv_heads: + raise ValueError("K and V head counts must match") + if self.head_dim % 16 or self.value_dim % 16: + raise ValueError("Q/K and V head dimensions must be multiples of 16") + if self.num_query_heads % self.num_kv_heads: + raise ValueError("H_q must be divisible by H_kv") + self.gqa_group_size = self.num_query_heads // self.num_kv_heads + if self.gqa_group_size not in (1, 2, 4, 8, 16): + raise ValueError( + f"H_q / H_kv must be one of {{1, 2, 4, 8, 16}}, got {self.gqa_group_size}" + ) + + if self.block_indices_desc.ndim != 3 or self.block_indices_desc.shape[:2] != ( + self.total_tokens, + self.num_kv_heads, + ): + raise ValueError("block_indices must have shape (T, H_kv, K)") + if self.block_indices_desc.shape[2] <= 0: + raise ValueError("block_indices K dimension must be positive") + if self.block_counts_desc.shape != (self.total_tokens, self.num_kv_heads): + raise ValueError("block_counts must have shape (T, H_kv)") + if ( + self.block_indices_desc.cudnn_dtype != data_type.INT32 + or self.block_counts_desc.cudnn_dtype != data_type.INT32 + ): + raise ValueError("block_indices and block_counts must have dtype int32") + if self.cum_seqlen_desc.ndim != 1 or self.cum_seqlen_desc.shape[0] < 2: + raise ValueError("cum_seqlen must have shape (B + 1,) with B > 0") + if self.cum_seqlen_desc.cudnn_dtype not in ( + data_type.INT32, + data_type.INT64, + ): + raise ValueError( + "cumulative sequence lengths must have dtype int32 or int64" + ) + + self.max_s_q = int(max_s_q) + self.max_s_k = int(max_s_k) + if self.max_s_q <= 0 or self.max_s_k <= 0: + raise ValueError("max_s_q and max_s_k must be positive") + if self.max_s_q != self.max_s_k: + raise ValueError( + "SelectionAttention requires max_s_q and max_s_k to be identical" + ) + self.block_size = int(block_size) + if self.block_size not in (16, 32, 64): + raise ValueError("block_size must be 16, 32, or 64") + + self.output_dtype = normalize_supported_dtype( + o_dtype, + sample_q.dtype, + "o_dtype", + (sample_q.dtype,), + ) + normalize_supported_dtype(acc_dtype, jnp.float32, "acc_dtype", (jnp.float32,)) + self.scale_softmax = ( + 1.0 / math.sqrt(self.head_dim) + if scale_softmax is None + else float(scale_softmax) + ) + self.target_compute_capability = target_compute_capability + self.compute_capability: int | None = None + + self.q_kernel_desc = JaxTensorDesc.from_shape( + (1, *self.q_desc.shape), + self.q_desc.dtype, + name="q_tensor", + ) + self.k_kernel_desc = JaxTensorDesc.from_shape( + (1, *self.k_desc.shape), + self.k_desc.dtype, + name="k_tensor", + ) + self.v_kernel_desc = JaxTensorDesc.from_shape( + (1, *self.v_desc.shape), + self.v_desc.dtype, + name="v_tensor", + ) + self.o_desc = JaxTensorDesc.from_shape( + (1, self.total_tokens, self.num_query_heads, self.value_dim), + self.output_dtype, + name="o_tensor", + init_value=0, + ) + self.l_desc = JaxTensorDesc.from_shape( + (1, self.total_tokens, self.num_query_heads), + jnp.float32, + name="l_tensor", + init_value=0.0, + ) + self.m_desc = JaxTensorDesc.from_shape( + (1, self.total_tokens, self.num_query_heads), + jnp.float32, + name="m_tensor", + init_value=float("-inf"), + ) + + def check_support(self) -> bool: + self.compute_capability = self._resolve_compute_capability( + self.target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "SelectionAttention", + ) + return True + + def __call__( + self, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + block_indices_tensor: Any, + block_counts_tensor: Any, + cum_seqlen_tensor: Any, + ) -> TupleDict: + self.check_support() + for value, desc in ( + (q_tensor, self.q_desc), + (k_tensor, self.k_desc), + (v_tensor, self.v_desc), + ): + self._check_tensor_signature(value, desc) + + q_storage = jnp.reshape(q_tensor, self.q_kernel_desc.shape) + k_storage = jnp.reshape(k_tensor, self.k_kernel_desc.shape) + v_storage = jnp.reshape(v_tensor, self.v_kernel_desc.shape) + output, lse_sum, row_max = self._call_kernel( + ( + q_storage, + k_storage, + v_storage, + block_indices_tensor, + block_counts_tensor, + cum_seqlen_tensor, + ), + launch=self._launch_kernel, + output_descs=(self.o_desc, self.l_desc, self.m_desc), + input_descs=( + self.q_kernel_desc, + self.k_kernel_desc, + self.v_kernel_desc, + self.block_indices_desc, + self.block_counts_desc, + self.cum_seqlen_desc, + ), + compile_options=compile_options_for_target(self.compute_capability), + ) + return TupleDict( + o_tensor=jnp.reshape( + output, (self.total_tokens, self.num_query_heads, self.value_dim) + ), + l_tensor=jnp.reshape(lse_sum, (self.total_tokens, self.num_query_heads, 1)), + m_tensor=jnp.reshape(row_max, (self.total_tokens, self.num_query_heads, 1)), + ) + + def _launch_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + block_indices: Any, + block_counts: Any, + cum_seqlen: Any, + output: Any, + lse_sum: Any, + row_max: Any, + ) -> None: + from cutlass import Float32 + from cutlass.jax import jax_to_cutlass_dtype + + from .NSA_select_attn_fwd_hmma import HopperSelectAttentionFwd + + kernel = HopperSelectAttentionFwd( + head_dim=self.head_dim, + value_dim=self.value_dim, + GQA_group_size=self.gqa_group_size, + block_size=self.block_size, + dtype=jax_to_cutlass_dtype(self.q_desc.dtype), + acc_dtype=Float32, + ) + kernel( + q, + k, + v, + output, + lse_sum, + row_max, + block_indices, + block_counts, + self.max_s_q, + cum_seqlen, + Float32(self.scale_softmax), + stream, + ) + + +@partial( + jax.jit, + static_argnames=( + "max_s_q", + "max_s_k", + "block_size", + "scale_softmax", + "o_dtype", + "acc_dtype", + "target_compute_capability", + ), +) +def selection_attention_wrapper( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + block_indices_tensor: Any, + block_counts_tensor: Any, + cum_seqlen_tensor: Any, + *, + max_s_q: int, + max_s_k: int, + block_size: int = 64, + scale_softmax: float | None = None, + o_dtype: Any = None, + acc_dtype: Any = None, + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute packed THD selection attention.""" + + values = ( + q_tensor, + k_tensor, + v_tensor, + block_indices_tensor, + block_counts_tensor, + cum_seqlen_tensor, + ) + return SelectionAttention( + *values, + max_s_q=max_s_q, + max_s_k=max_s_k, + block_size=block_size, + scale_softmax=scale_softmax, + o_dtype=o_dtype, + acc_dtype=acc_dtype, + target_compute_capability=target_compute_capability, + )(*values) + + +__all__ = [ + "SUPPORTED_COMPUTE_CAPABILITIES", + "SelectionAttention", + "selection_attention_wrapper", +] diff --git a/python/cudnn/native_sparse_attention/sliding_window_attention/__init__.py b/python/cudnn/native_sparse_attention/sliding_window_attention/__init__.py index 7c0dc8fc7..980cc3fcb 100644 --- a/python/cudnn/native_sparse_attention/sliding_window_attention/__init__.py +++ b/python/cudnn/native_sparse_attention/sliding_window_attention/__init__.py @@ -1,6 +1,9 @@ -from .api import SlidingWindowAttention, sliding_window_attention_wrapper +"""Lazy Torch sliding-window-attention API exports.""" -__all__ = [ - "SlidingWindowAttention", - "sliding_window_attention_wrapper", -] +from ...common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("SlidingWindowAttention", "sliding_window_attention_wrapper")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/native_sparse_attention/sliding_window_attention/jax.py b/python/cudnn/native_sparse_attention/sliding_window_attention/jax.py new file mode 100644 index 000000000..8b24bf21f --- /dev/null +++ b/python/cudnn/native_sparse_attention/sliding_window_attention/jax.py @@ -0,0 +1,357 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for fixed and fully packed NSA sliding-window attention.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._jax import JaxApiBase, TupleDict +from ..jax_utils import ( + FIXED_LAYOUTS, + describe_fixed_data, + fixed_data_mode, + normalize_attention_layout, + normalize_supported_dtype, + require_fixed_qkv, +) + + +class SlidingWindowAttention(JaxApiBase): + """JAX callable for fixed BHSD/BSHD or fully packed THD attention. + + Packed inputs use cumulative sequence lengths and static maxima. Their + dynamic values must start at zero, be monotonic, end at the packed token + count, and not exceed the supplied maxima. + """ + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_cum_seqlen_q: Any | None = None, + sample_cum_seqlen_k: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + left_bound: int = 1, + right_bound: int = 0, + is_infer: bool = False, + attn_scale: float | None = None, + o_dtype: Any = None, + layout: str | None = None, + ) -> None: + ranks = tuple(len(sample.shape) for sample in (sample_q, sample_k, sample_v)) + if len(set(ranks)) != 1: + raise ValueError(f"Q, K, and V must use the same rank, got {ranks}") + self.input_layout = normalize_attention_layout(layout, ranks[0]) + if self.input_layout in FIXED_LAYOUTS: + if sample_cum_seqlen_q is not None or sample_cum_seqlen_k is not None: + raise ValueError( + "cumulative sequence lengths must be omitted for fixed inputs" + ) + if max_s_q is not None or max_s_k is not None: + raise ValueError("max_s_q and max_s_k are only valid for THD layout") + self._init_fixed(sample_q, sample_k, sample_v) + else: + self._init_packed( + sample_q, + sample_k, + sample_v, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + max_s_q, + max_s_k, + ) + self.left_bound = int(left_bound) + self.right_bound = int(right_bound) + self.is_infer = bool(is_infer) + self.attn_scale = None if attn_scale is None else float(attn_scale) + self.output_dtype = normalize_supported_dtype( + o_dtype, + sample_q.dtype, + "o_dtype", + (jnp.float16, jnp.bfloat16), + ) + if self.left_bound < 1: + raise ValueError(f"left_bound must be at least 1, got {self.left_bound}") + if self.right_bound < 0: + raise ValueError( + f"right_bound must be non-negative, got {self.right_bound}" + ) + + def _init_fixed(self, sample_q: Any, sample_k: Any, sample_v: Any) -> None: + self.data_mode = fixed_data_mode(self.input_layout, kernel_axes="BHSD") + self.q_desc = describe_fixed_data( + sample_q, + "sample_q", + layout=self.input_layout, + kernel_axes="BHSD", + ) + self.k_desc = describe_fixed_data( + sample_k, + "sample_k", + layout=self.input_layout, + kernel_axes="BHSD", + ) + self.v_desc = describe_fixed_data( + sample_v, + "sample_v", + layout=self.input_layout, + kernel_axes="BHSD", + ) + require_fixed_qkv( + self.q_desc, + self.k_desc, + self.v_desc, + operation_name="SlidingWindowAttention", + kernel_axes="BHSD", + ) + self.cum_q_desc = self.cum_k_desc = None + self.max_s_q = self.max_s_k = None + + def _init_packed( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_cum_seqlen_q: Any | None, + sample_cum_seqlen_k: Any | None, + max_s_q: int | None, + max_s_k: int | None, + ) -> None: + if sample_cum_seqlen_q is None or sample_cum_seqlen_k is None: + raise ValueError( + "packed THD inputs require cumulative Q and K sequence lengths" + ) + if max_s_q is None or max_s_k is None: + raise ValueError("packed THD inputs require max_s_q and max_s_k") + + self.data_mode = None + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.k_desc = self._to_tensor_desc(sample_k, "sample_k") + self.v_desc = self._to_tensor_desc(sample_v, "sample_v") + self.cum_q_desc = self._to_tensor_desc( + sample_cum_seqlen_q, "sample_cum_seqlen_q" + ) + self.cum_k_desc = self._to_tensor_desc( + sample_cum_seqlen_k, "sample_cum_seqlen_k" + ) + self.total_q, self.num_query_heads, self.head_dim = self.q_desc.shape + self.total_k, self.num_kv_heads, k_head_dim = self.k_desc.shape + v_total, v_heads, value_dim = self.v_desc.shape + dimensions = ( + self.total_q, + self.total_k, + self.num_query_heads, + self.num_kv_heads, + self.head_dim, + ) + if any(value <= 0 for value in dimensions): + raise ValueError( + f"SlidingWindowAttention dimensions must be positive, got {dimensions}" + ) + if (v_total, v_heads, value_dim) != ( + self.total_k, + self.num_kv_heads, + self.head_dim, + ) or k_head_dim != self.head_dim: + raise ValueError("packed K and V metadata must match Q/K head dimensions") + if self.q_desc.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError("packed Q, K, and V must use float16 or bfloat16") + if ( + self.k_desc.cudnn_dtype != self.q_desc.cudnn_dtype + or self.v_desc.cudnn_dtype != self.q_desc.cudnn_dtype + ): + raise ValueError("packed Q, K, and V must have the same dtype") + if self.num_query_heads % self.num_kv_heads: + raise ValueError("H_q must be divisible by H_kv") + if ( + self.cum_q_desc.ndim != 1 + or self.cum_q_desc.shape != self.cum_k_desc.shape + or self.cum_q_desc.shape[0] < 2 + ): + raise ValueError( + "cumulative Q and K sequence lengths must have matching (B + 1,) shapes" + ) + if ( + self.cum_q_desc.cudnn_dtype != data_type.INT32 + or self.cum_k_desc.cudnn_dtype != data_type.INT32 + ): + raise ValueError("cumulative sequence lengths must use int32") + self.batch = self.cum_q_desc.shape[0] - 1 + self.max_s_q = int(max_s_q) + self.max_s_k = int(max_s_k) + if ( + self.max_s_q <= 0 + or self.max_s_k <= 0 + or self.max_s_q > self.total_q + or self.max_s_k > self.total_k + ): + raise ValueError( + "max_s_q and max_s_k must be positive and no larger than packed token counts" + ) + + def check_support(self) -> bool: + return True + + def __call__( + self, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + ) -> TupleDict: + self.check_support() + for value, desc in ( + (q_tensor, self.q_desc), + (k_tensor, self.k_desc), + (v_tensor, self.v_desc), + ): + self._check_tensor_signature(value, desc) + + if self.input_layout == "BHSD": + if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: + raise ValueError( + "cumulative sequence lengths must be omitted for fixed inputs" + ) + q_btnh = jnp.transpose(q_tensor, (0, 2, 1, 3)) + k_btnh = jnp.transpose(k_tensor, (0, 2, 1, 3)) + v_btnh = jnp.transpose(v_tensor, (0, 2, 1, 3)) + query_lengths = key_value_lengths = None + elif self.input_layout == "BSHD": + if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: + raise ValueError( + "cumulative sequence lengths must be omitted for fixed inputs" + ) + q_btnh, k_btnh, v_btnh = q_tensor, k_tensor, v_tensor + query_lengths = key_value_lengths = None + else: + if cum_seqlen_q_tensor is None or cum_seqlen_k_tensor is None: + raise ValueError( + "packed THD inputs require cumulative Q and K sequence lengths" + ) + self._check_tensor_signature(cum_seqlen_q_tensor, self.cum_q_desc) + self._check_tensor_signature(cum_seqlen_k_tensor, self.cum_k_desc) + query_lengths = jnp.diff(cum_seqlen_q_tensor) + key_value_lengths = jnp.diff(cum_seqlen_k_tensor) + q_btnh = self._pad_packed(q_tensor, cum_seqlen_q_tensor[:-1], self.max_s_q) + k_btnh = self._pad_packed(k_tensor, cum_seqlen_k_tensor[:-1], self.max_s_k) + v_btnh = self._pad_packed(v_tensor, cum_seqlen_k_tensor[:-1], self.max_s_k) + attention_result = jax.nn.dot_product_attention( + q_btnh, + k_btnh, + v_btnh, + scale=self.attn_scale, + query_seq_lengths=query_lengths, + key_value_seq_lengths=key_value_lengths, + local_window_size=(self.left_bound - 1, self.right_bound), + implementation="cudnn", + return_residual=not self.is_infer, + ) + if self.is_infer: + output_btnh = attention_result + stats = None + else: + output_btnh, residual_btn = attention_result + stats = residual_btn + if self.input_layout == "THD": + output = self._unpad_packed( + output_btnh, cum_seqlen_q_tensor, self.total_q + ).astype(self.output_dtype) + if stats is not None: + stats = self._unpad_packed( + stats[..., None], cum_seqlen_q_tensor, self.total_q + ) + else: + output = ( + jnp.transpose(output_btnh, (0, 2, 1, 3)) + if self.input_layout == "BHSD" + else output_btnh + ).astype(self.output_dtype) + if stats is not None: + stats = jnp.transpose(stats, (0, 2, 1))[..., None] + return TupleDict(o_tensor=output, stats_tensor=stats) + + @staticmethod + def _pad_packed(values: Any, starts: Any, max_seqlen: int) -> Any: + padded = jnp.pad(values, ((0, max_seqlen), (0, 0), (0, 0))) + slice_shape = (max_seqlen, values.shape[1], values.shape[2]) + return jax.vmap( + lambda start: jax.lax.dynamic_slice(padded, (start, 0, 0), slice_shape) + )(starts) + + @staticmethod + def _unpad_packed(values: Any, cumulative: Any, total_tokens: int) -> Any: + token = jnp.arange(total_tokens, dtype=cumulative.dtype) + batch = jnp.searchsorted(cumulative[1:], token, side="right") + position = token - cumulative[batch] + return values[batch, position] + + +@partial( + jax.jit, + static_argnames=( + "left_bound", + "right_bound", + "is_infer", + "attn_scale", + "o_dtype", + "max_s_q", + "max_s_k", + "layout", + ), +) +def sliding_window_attention_wrapper( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + left_bound: int = 1, + right_bound: int = 0, + is_infer: bool = False, + attn_scale: float | None = None, + o_dtype: Any = None, + layout: str | None = None, +) -> TupleDict: + """Compute fixed BHSD/BSHD or fully packed THD attention. + + Packed THD calls require cumulative sequence lengths and static + ``max_s_q``/``max_s_k`` bounds. Arbitrary ragged byte-offset tensors are + intentionally outside this functional JAX API. + """ + + return SlidingWindowAttention( + q_tensor, + k_tensor, + v_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + max_s_q=max_s_q, + max_s_k=max_s_k, + left_bound=left_bound, + right_bound=right_bound, + is_infer=is_infer, + attn_scale=attn_scale, + o_dtype=o_dtype, + layout=layout, + )( + q_tensor, + k_tensor, + v_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + + +__all__ = ["SlidingWindowAttention", "sliding_window_attention_wrapper"] diff --git a/python/cudnn/native_sparse_attention/top_k/__init__.py b/python/cudnn/native_sparse_attention/top_k/__init__.py index 06dc32184..705ddb1a5 100644 --- a/python/cudnn/native_sparse_attention/top_k/__init__.py +++ b/python/cudnn/native_sparse_attention/top_k/__init__.py @@ -1,6 +1,9 @@ -from .api import TopKReduction, topk_reduction_wrapper +"""Lazy Torch top-K API exports.""" -__all__ = [ - "TopKReduction", - "topk_reduction_wrapper", -] +from ...common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("TopKReduction", "topk_reduction_wrapper")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/native_sparse_attention/top_k/jax.py b/python/cudnn/native_sparse_attention/top_k/jax.py new file mode 100644 index 000000000..73efbec05 --- /dev/null +++ b/python/cudnn/native_sparse_attention/top_k/jax.py @@ -0,0 +1,610 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for SM100 NSA top-K reduction.""" + +from __future__ import annotations + +from functools import partial +import math +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._jax.compiler import compile_options_for_target +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..._jax.layout import to_public_axes +from ..jax_utils import ( + FIXED_LAYOUTS, + describe_fixed_data, + fixed_data_mode, + make_fixed_output, + normalize_attention_layout, + normalize_supported_dtype, + require_fixed_qkv, +) + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) +_PACKED_LAYOUT = "THD" +_KERNEL_DATA_AXES = "BHSD" +_KERNEL_DATA_STRIDE_ORDER = (3, 1, 2, 0) + + +class TopKReduction(JaxApiBase): + """JAX callable specialized from fixed BHSD/BSHD or packed THD metadata. + + Packed inputs use logical ``(T, H, D)`` Q/K arrays and require both + cumulative sequence-length arrays plus static ``max_s_q`` and ``max_s_k``. + The cumulative arrays remain explicit custom-call operands. Their contents, + including the final packed-token offsets, are therefore validated by the + caller rather than read while JAX traces the operation. + """ + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_lse: Any, + sample_cum_seqlen_q: Any | None = None, + sample_cum_seqlen_k: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + acc_dtype: Any = None, + k_value: int = 16, + selection_block_size: int = 64, + compress_stride: int = 32, + is_causal: bool = True, + mma_tiler_mn: tuple[int, int] = (128, 128), + scale_softmax: float | None = None, + layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + q_rank = len(tuple(sample_q.shape)) + k_rank = len(tuple(sample_k.shape)) + if q_rank != k_rank: + raise ValueError( + f"Q and K must use the same rank, got {q_rank} and {k_rank}" + ) + self.input_layout = normalize_attention_layout(layout, q_rank) + + if self.input_layout in FIXED_LAYOUTS: + self.data_mode = fixed_data_mode( + self.input_layout, kernel_axes=_KERNEL_DATA_AXES + ) + self._init_fixed_shape( + sample_q, + sample_k, + sample_lse, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + ) + else: + self.data_mode = None + self._init_packed_shape( + sample_q, + sample_k, + sample_lse, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + max_s_q, + max_s_k, + ) + + normalize_supported_dtype(acc_dtype, jnp.float32, "acc_dtype", (jnp.float32,)) + self.k_value = int(k_value) + self.selection_block_size = int(selection_block_size) + self.compress_stride = int(compress_stride) + self.is_causal = bool(is_causal) + self.mma_tiler_mn = tuple(mma_tiler_mn) + self._check_configuration() + self.scale_softmax = ( + 1.0 / math.sqrt(self.head_dim) + if scale_softmax is None + else float(scale_softmax) + ) + self.target_compute_capability = target_compute_capability + self.compute_capability: int | None = None + + if self.input_layout in FIXED_LAYOUTS: + canonical_output_shape = ( + self.batch, + self.num_kv_heads, + self.seqlen_q, + self.k_value, + ) + output_shape = to_public_axes(canonical_output_shape, self.data_mode) + self.scores_desc = make_fixed_output( + output_shape, + jnp.float32, + "topk_scores_tensor", + layout=self.input_layout, + kernel_axes=_KERNEL_DATA_AXES, + kernel_stride_order=_KERNEL_DATA_STRIDE_ORDER, + init_value=float("-inf"), + ) + self.indices_desc = make_fixed_output( + output_shape, + jnp.int32, + "topk_indices_tensor", + layout=self.input_layout, + kernel_axes=_KERNEL_DATA_AXES, + kernel_stride_order=_KERNEL_DATA_STRIDE_ORDER, + init_value=-1, + ) + else: + output_shape = (self.total_q_tokens, self.num_kv_heads, self.k_value) + self.scores_desc = JaxTensorDesc.from_shape( + output_shape, + jnp.float32, + name="topk_scores_tensor", + init_value=float("-inf"), + ) + self.indices_desc = JaxTensorDesc.from_shape( + output_shape, + jnp.int32, + name="topk_indices_tensor", + init_value=-1, + ) + + def _init_fixed_shape( + self, + sample_q: Any, + sample_k: Any, + sample_lse: Any, + sample_cum_seqlen_q: Any | None, + sample_cum_seqlen_k: Any | None, + ) -> None: + if sample_cum_seqlen_q is not None or sample_cum_seqlen_k is not None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must both be omitted for BHSD layout" + ) + + # The kernel consumes BHSD modes with the H-inside-S storage used by + # the Torch wrapper's transposed BSHD tensors. + self.q_desc = describe_fixed_data( + sample_q, + "sample_q", + layout=self.input_layout, + kernel_axes=_KERNEL_DATA_AXES, + kernel_stride_order=_KERNEL_DATA_STRIDE_ORDER, + ) + self.k_desc = describe_fixed_data( + sample_k, + "sample_k", + layout=self.input_layout, + kernel_axes=_KERNEL_DATA_AXES, + kernel_stride_order=_KERNEL_DATA_STRIDE_ORDER, + ) + ( + self.batch, + self.num_query_heads, + self.num_kv_heads, + self.seqlen_q, + self.seqlen_k, + self.head_dim, + ) = require_fixed_qkv( + self.q_desc, + self.k_desc, + operation_name="TopKReduction", + kernel_axes=_KERNEL_DATA_AXES, + ) + self.total_q_tokens = self.batch * self.seqlen_q + self.lse_desc = self._to_tensor_desc(sample_lse, "sample_lse") + expected_lse_shape = (self.batch, self.num_query_heads, self.seqlen_q) + if self.lse_desc.shape != expected_lse_shape: + raise ValueError( + f"sample_lse must have shape {expected_lse_shape}, got {self.lse_desc.shape}" + ) + if self.lse_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"sample_lse must have dtype float32, got {self.lse_desc.dtype}" + ) + + self.max_s_q = self.seqlen_q + self.max_s_k = self.seqlen_k + self.cum_seqlen_q_desc = None + self.cum_seqlen_k_desc = None + self.q_kernel_desc = self.q_desc + self.k_kernel_desc = self.k_desc + self.lse_kernel_desc = self.lse_desc + + def _init_packed_shape( + self, + sample_q: Any, + sample_k: Any, + sample_lse: Any, + sample_cum_seqlen_q: Any | None, + sample_cum_seqlen_k: Any | None, + max_s_q: int | None, + max_s_k: int | None, + ) -> None: + if sample_cum_seqlen_q is None or sample_cum_seqlen_k is None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k are both required for THD layout" + ) + if max_s_q is None or max_s_k is None: + raise ValueError("max_s_q and max_s_k are both required for THD layout") + + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.k_desc = self._to_tensor_desc(sample_k, "sample_k") + self.lse_desc = self._to_tensor_desc(sample_lse, "sample_lse") + self.cum_seqlen_q_desc = self._to_tensor_desc( + sample_cum_seqlen_q, "sample_cum_seqlen_q" + ) + self.cum_seqlen_k_desc = self._to_tensor_desc( + sample_cum_seqlen_k, "sample_cum_seqlen_k" + ) + + self.total_q_tokens, self.num_query_heads, self.head_dim = self.q_desc.shape + self.total_k_tokens, self.num_kv_heads, k_head_dim = self.k_desc.shape + dimensions = ( + self.total_q_tokens, + self.total_k_tokens, + self.num_query_heads, + self.num_kv_heads, + self.head_dim, + ) + if any(value <= 0 for value in dimensions): + raise ValueError( + f"TopKReduction dimensions must be positive, got {dimensions}" + ) + if self.q_desc.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError( + f"TopKReduction requires float16 or bfloat16 inputs, got {self.q_desc.dtype}" + ) + if self.k_desc.cudnn_dtype != self.q_desc.cudnn_dtype: + raise ValueError("sample_k must have the same dtype as sample_q") + if k_head_dim != self.head_dim: + raise ValueError( + f"Q and K head dimensions must match, got {self.head_dim} and {k_head_dim}" + ) + if self.head_dim not in (32, 64, 128): + raise ValueError( + f"head dimension must be one of 32, 64, 128, got {self.head_dim}" + ) + if self.num_query_heads % self.num_kv_heads: + raise ValueError( + f"H_q ({self.num_query_heads}) must be divisible by H_kv ({self.num_kv_heads})" + ) + + expected_lse_shapes = ( + (self.total_q_tokens, self.num_query_heads), + (self.total_q_tokens, self.num_query_heads, 1), + ) + if self.lse_desc.shape not in expected_lse_shapes: + raise ValueError( + "sample_lse must have shape " + f"{expected_lse_shapes[0]} or {expected_lse_shapes[1]}, got {self.lse_desc.shape}" + ) + if self.lse_desc.cudnn_dtype != data_type.FLOAT: + raise ValueError( + f"sample_lse must have dtype float32, got {self.lse_desc.dtype}" + ) + + if ( + self.cum_seqlen_q_desc.ndim != 1 + or self.cum_seqlen_k_desc.ndim != 1 + or self.cum_seqlen_q_desc.shape != self.cum_seqlen_k_desc.shape + or self.cum_seqlen_q_desc.shape[0] < 2 + ): + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must have the same shape (B + 1,) with B > 0" + ) + if ( + self.cum_seqlen_q_desc.cudnn_dtype != data_type.INT32 + or self.cum_seqlen_k_desc.cudnn_dtype != data_type.INT32 + ): + raise ValueError("cum_seqlen_q and cum_seqlen_k must have dtype int32") + + self.batch = self.cum_seqlen_q_desc.shape[0] - 1 + self.seqlen_q = self.total_q_tokens + self.seqlen_k = self.total_k_tokens + self.max_s_q = int(max_s_q) + self.max_s_k = int(max_s_k) + if ( + self.max_s_q <= 0 + or self.max_s_k <= 0 + or self.max_s_q > self.total_q_tokens + or self.max_s_k > self.total_k_tokens + ): + raise ValueError( + "max_s_q and max_s_k must be positive and no larger than " + "their packed token counts" + ) + + # These virtual rank-4 views preserve the packed THD byte order while + # exposing the B,H,T,D shapes read by the CuTe kernel. LSE is likewise + # presented as B,H,T, matching the Torch wrapper's transpose. + self.q_kernel_desc = JaxTensorDesc.from_shape( + (1, self.num_query_heads, self.total_q_tokens, self.head_dim), + self.q_desc.dtype, + name="q_tensor", + public_stride_order=_KERNEL_DATA_STRIDE_ORDER, + ) + self.k_kernel_desc = JaxTensorDesc.from_shape( + (1, self.num_kv_heads, self.total_k_tokens, self.head_dim), + self.k_desc.dtype, + name="k_tensor", + public_stride_order=_KERNEL_DATA_STRIDE_ORDER, + ) + self.lse_kernel_desc = JaxTensorDesc.from_shape( + (1, self.num_query_heads, self.total_q_tokens), + jnp.float32, + name="lse_tensor", + ) + + def _check_configuration(self) -> None: + if self.mma_tiler_mn != (128, 128): + raise ValueError( + f"mma_tiler_mn must be (128, 128), got {self.mma_tiler_mn}" + ) + if self.selection_block_size <= 0 or self.compress_stride <= 0: + raise ValueError( + "selection_block_size and compress_stride must be positive" + ) + if self.selection_block_size % self.compress_stride: + raise ValueError( + "selection_block_size must be divisible by compress_stride" + ) + reduction_width = self.selection_block_size // self.compress_stride + if ( + reduction_width > self.mma_tiler_mn[1] + or self.mma_tiler_mn[1] % reduction_width + ): + raise ValueError( + "selection_block_size / compress_stride must divide the MMA N tile" + ) + candidate_count = self.mma_tiler_mn[1] // reduction_width + if self.k_value <= 0 or self.k_value % 4 or self.k_value > candidate_count: + raise ValueError( + "k_value must be a positive multiple of 4 no greater than the " + f"per-tile candidate count ({candidate_count}), got {self.k_value}" + ) + + def check_support(self) -> bool: + self.compute_capability = self._resolve_compute_capability( + self.target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "TopKReduction", + ) + return True + + def __call__( + self, + q_tensor: Any, + k_tensor: Any, + lse_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + ) -> TupleDict: + self.check_support() + if self.input_layout in FIXED_LAYOUTS: + if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must be omitted for fixed layout" + ) + inputs = (q_tensor, k_tensor, lse_tensor) + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.lse_kernel_desc, + ) + launch = self._launch_kernel + else: + if cum_seqlen_q_tensor is None or cum_seqlen_k_tensor is None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k are both required for THD layout" + ) + for value, desc in ( + (q_tensor, self.q_desc), + (k_tensor, self.k_desc), + (lse_tensor, self.lse_desc), + ): + self._check_tensor_signature(value, desc) + q_storage = jnp.transpose( + jnp.reshape( + q_tensor, + (1, self.total_q_tokens, self.num_query_heads, self.head_dim), + ), + (0, 2, 1, 3), + ) + k_storage = jnp.transpose( + jnp.reshape( + k_tensor, + (1, self.total_k_tokens, self.num_kv_heads, self.head_dim), + ), + (0, 2, 1, 3), + ) + lse_storage = jnp.transpose( + jnp.reshape(lse_tensor, (self.total_q_tokens, self.num_query_heads)), + (1, 0), + )[None, ...] + inputs = ( + q_storage, + k_storage, + lse_storage, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.lse_kernel_desc, + self.cum_seqlen_q_desc, + self.cum_seqlen_k_desc, + ) + launch = self._launch_packed_kernel + + scores, indices = self._call_kernel( + inputs, + launch=launch, + output_descs=(self.scores_desc, self.indices_desc), + input_descs=input_descs, + compile_options=compile_options_for_target(self.compute_capability), + ) + return TupleDict(topk_scores_tensor=scores, topk_indices_tensor=indices) + + def _launch_kernel( + self, + stream: Any, + q: Any, + k: Any, + lse: Any, + topk_scores: Any, + topk_indices: Any, + ) -> None: + self._invoke_kernel( + stream, + q, + k, + lse, + topk_scores, + topk_indices, + None, + None, + ) + + def _launch_packed_kernel( + self, + stream: Any, + q: Any, + k: Any, + lse: Any, + cum_seqlen_q: Any, + cum_seqlen_k: Any, + topk_scores: Any, + topk_indices: Any, + ) -> None: + self._invoke_kernel( + stream, + q, + k, + lse, + topk_scores, + topk_indices, + cum_seqlen_q, + cum_seqlen_k, + ) + + def _invoke_kernel( + self, + stream: Any, + q: Any, + k: Any, + lse: Any, + topk_scores: Any, + topk_indices: Any, + cum_seqlen_q: Any | None, + cum_seqlen_k: Any | None, + ) -> None: + from cutlass import Float32, Int32 + from cutlass.jax import jax_to_cutlass_dtype + + from .nsa_top_k_reduction_fwd import FineGrainedReductionQK + + kernel = FineGrainedReductionQK( + element_dtype=jax_to_cutlass_dtype(self.q_desc.dtype), + acc_dtype=Float32, + k_value=self.k_value, + selection_block_size=self.selection_block_size, + compress_block_sliding_stride=self.compress_stride, + mma_tiler=(128, 128, self.head_dim), + is_causal=self.is_causal, + ) + problem_size = tuple( + Int32(value) + for value in ( + self.batch, + self.max_s_q, + self.max_s_k, + self.num_query_heads, + self.num_kv_heads, + self.head_dim, + ) + ) + kernel( + problem_size, + q, + k, + lse, + topk_scores, + topk_indices, + Float32(self.scale_softmax * math.log2(math.e)), + cum_seqlen_q, + cum_seqlen_k, + stream, + ) + + +@partial( + jax.jit, + static_argnames=( + "max_s_q", + "max_s_k", + "acc_dtype", + "k_value", + "selection_block_size", + "compress_stride", + "is_causal", + "mma_tiler_mn", + "scale_softmax", + "layout", + "target_compute_capability", + ), +) +def topk_reduction_wrapper( + q_tensor: Any, + k_tensor: Any, + lse_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + acc_dtype: Any = None, + k_value: int = 16, + selection_block_size: int = 64, + compress_stride: int = 32, + is_causal: bool = True, + mma_tiler_mn: tuple[int, int] = (128, 128), + scale_softmax: float | None = None, + layout: str | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + """Select blocks for fixed BHSD/BSHD or packed THD inputs. + + Fixed outputs follow ``layout``; fixed LSE remains ``(B, H, S)``. + """ + + return TopKReduction( + q_tensor, + k_tensor, + lse_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + max_s_q=max_s_q, + max_s_k=max_s_k, + acc_dtype=acc_dtype, + k_value=k_value, + selection_block_size=selection_block_size, + compress_stride=compress_stride, + is_causal=is_causal, + mma_tiler_mn=mma_tiler_mn, + scale_softmax=scale_softmax, + layout=layout, + target_compute_capability=target_compute_capability, + )( + q_tensor, + k_tensor, + lse_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + + +__all__ = [ + "SUPPORTED_COMPUTE_CAPABILITIES", + "TopKReduction", + "topk_reduction_wrapper", +] diff --git a/python/cudnn/rmsnorm_rht_amax/__init__.py b/python/cudnn/rmsnorm_rht_amax/__init__.py index 9da5cae15..de543a280 100644 --- a/python/cudnn/rmsnorm_rht_amax/__init__.py +++ b/python/cudnn/rmsnorm_rht_amax/__init__.py @@ -1,16 +1,26 @@ # 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 Torch API and framework-neutral operation/kernel exports. + +JAX APIs are exported through :mod:`cudnn.jax`. +""" -__all__ = [ - "RmsNormRhtAmaxSm100", - "best_num_threads", - "pick_rows_per_cta", - "rmsnorm_rht_amax_wrapper_sm100", -] +from ..common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "op": ( + "RmsNormRhtAmaxSm100Op", + "best_num_threads", + "pick_rows_per_cta", + ), + "kernel": ("RMSNormRHTAmaxKernel",), + "api": ( + "RmsNormRhtAmaxSm100", + "rmsnorm_rht_amax_wrapper_sm100", + ), + }, + submodules=("api", "kernel", "op"), +) diff --git a/python/cudnn/rmsnorm_rht_amax/api.py b/python/cudnn/rmsnorm_rht_amax/api.py index 0294a3a22..933b88364 100644 --- a/python/cudnn/rmsnorm_rht_amax/api.py +++ b/python/cudnn/rmsnorm_rht_amax/api.py @@ -7,46 +7,21 @@ from typing import Optional from cuda.bindings import driver as cuda -import cutlass import cutlass.cute as cute import torch -from cutlass import Float32 from cutlass.cute.runtime import make_fake_stream from cudnn.api_base import APIBase, TupleDict from .kernel import RMSNormRHTAmaxKernel +from .op import ( + DEFAULT_NUM_THREADS_BY_N, + RmsNormRhtAmaxSm100Op, + best_num_threads, + pick_rows_per_cta, +) -DEFAULT_NUM_THREADS_BY_N = { - 2048: 128, - 4096: 256, - 7168: 128, - 8192: 512, - 16384: 1024, - 32768: 512, -} -RPC_CANDIDATES = (2, 4, 8) -TARGET_MIN_CTAS = 148 - - -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] +_TENSOR_ALIGNMENT = RMSNormRHTAmaxKernel.COPY_BITS // 8 class RmsNormRhtAmaxSm100(APIBase): @@ -74,57 +49,20 @@ def __init__( self.eps = eps self.requested_num_threads = num_threads self.requested_rows_per_cta = rows_per_cta - self.num_threads = None - self.rows_per_cta = None - self.n = None - def check_support(self) -> bool: - m, n = self._tensor_shape(self.x_desc, name="sample_x") - w_n = self._tensor_shape(self.w_desc, name="sample_w")[0] - o_m, o_n = self._tensor_shape(self.o_desc, name="sample_o") - - self._check_tensor_shape(self.x_desc, (m, n), "X") - self._check_tensor_shape(self.w_desc, (n,), "W") - self._check_tensor_shape(self.o_desc, (m, n), "O") - self._value_error_if(w_n != n, f"W length must match X hidden dimension, got {w_n} and {n}") - self._value_error_if((n % 16) != 0, f"N must be divisible by 16 for the Hadamard block size, got {n}") - self._value_error_if(o_m != m or o_n != n, f"O shape must match X shape, got {(o_m, o_n)} and {(m, n)}") - - self._check_tensor_stride(self.x_desc, stride=(n, 1), name="X", extra_error_msg="X must be row-major contiguous") - self._check_tensor_stride(self.w_desc, stride=(1,), name="W", extra_error_msg="W must be contiguous") - self._check_tensor_stride(self.o_desc, stride=(n, 1), name="O", extra_error_msg="O must be row-major contiguous") - - self._check_dtype(self.x_desc, dtype=torch.bfloat16, name="X") - self._check_dtype(self.w_desc, dtype=torch.bfloat16, name="W") - self._check_dtype(self.o_desc, dtype=torch.bfloat16, name="O") - self._check_dtype(self.amax_desc, dtype=torch.float32, name="Amax") - - resolved_num_threads = self.requested_num_threads - if resolved_num_threads is None: - resolved_num_threads = DEFAULT_NUM_THREADS_BY_N.get(n, best_num_threads(n)) - self._value_error_if(resolved_num_threads is None, f"No valid num_threads found for N={n}") - self._value_error_if(resolved_num_threads <= 0, f"num_threads must be positive, got {resolved_num_threads}") - self._value_error_if( - (resolved_num_threads % 32) != 0, - f"num_threads must be warp-aligned, got {resolved_num_threads}", - ) - self._value_error_if( - resolved_num_threads > 1024, - f"num_threads must not exceed the CUDA block size limit, got {resolved_num_threads}", + self._op = RmsNormRhtAmaxSm100Op( + x=self.x_desc, + weight=self.w_desc, + output=self.o_desc, + amax=self.amax_desc, + eps=eps, + num_threads=num_threads, + rows_per_cta=rows_per_cta, ) - 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") + def check_support(self) -> bool: + self._is_supported = False + self._op.check_support() self._runtime_error_if(not torch.cuda.is_available(), "CUDA is not available") major, minor = torch.cuda.get_device_capability(self.x_desc.device) @@ -134,9 +72,6 @@ def check_support(self) -> bool: f"RmsNormRhtAmaxSm100 requires SM100+, found SM{compute_capability}", ) - self.num_threads = resolved_num_threads - self.rows_per_cta = resolved_rows_per_cta - self.n = n self._is_supported = True return True @@ -146,35 +81,37 @@ def compile(self) -> None: return kernel = RMSNormRHTAmaxKernel( - n=self.n, - num_threads=self.num_threads, - eps=self.eps, - rows_per_cta=self.rows_per_cta, + n=self._op.n, + num_threads=self._op.num_threads, + eps=self._op.eps, + rows_per_cta=self._op.rows_per_cta, ) - valid_m = cute.sym_int(divisibility=self.rows_per_cta) + valid_m = cute.sym_int(divisibility=self._op.rows_per_cta) fake_x_tensor = self._make_fake_cute_compact_tensor( dtype=self.x_desc.dtype, - shape=(valid_m, self.n), + shape=(valid_m, self._op.n), stride_order=self.x_desc.stride_order, + assumed_align=_TENSOR_ALIGNMENT, dynamic_mode=None, - divisibility=self.rows_per_cta, + divisibility=self._op.rows_per_cta, ) - fake_w_tensor = self._make_fake_cute_tensor_from_desc(self.w_desc, assumed_align=16) + fake_w_tensor = self._make_fake_cute_tensor_from_desc(self.w_desc, assumed_align=_TENSOR_ALIGNMENT) fake_o_tensor = self._make_fake_cute_compact_tensor( dtype=self.o_desc.dtype, - shape=(valid_m, self.n), + shape=(valid_m, self._op.n), stride_order=self.o_desc.stride_order, + assumed_align=_TENSOR_ALIGNMENT, dynamic_mode=None, - divisibility=self.rows_per_cta, + divisibility=self._op.rows_per_cta, ) fake_num_ctas = cute.sym_int() fake_amax_tensor = self._make_fake_cute_tensor( dtype=self.amax_desc.dtype, shape=(fake_num_ctas,), stride=self.amax_desc.stride, - assumed_align=16, + assumed_align=_TENSOR_ALIGNMENT, ) fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) @@ -184,7 +121,6 @@ def compile(self) -> None: fake_w_tensor, fake_o_tensor, fake_amax_tensor, - Float32(self.eps), fake_stream, options="--enable-tvm-ffi", ) @@ -201,7 +137,6 @@ def tensor_api( w_tensor, o_tensor, amax_tensor, - Float32(self.eps), stream, ) @@ -222,6 +157,11 @@ def execute( o_tensor = self._unpad_tensor_to_ndim(o_tensor, 2, "o_tensor") amax_tensor = self._unpad_tensor_to_ndim(amax_tensor, 1, "amax_tensor") + # TVM-FFI validates the compiled tensor ABI. The amax extent is an + # independent symbol, so its relationship to M is checked here. + if x_tensor.ndim == 2: + self._check_tensor_shape(amax_tensor, (x_tensor.shape[0] // self._op.rows_per_cta,), "Amax") + if current_stream is None: current_stream = cuda.CUstream(torch.cuda.current_stream(x_tensor.device).cuda_stream) @@ -256,6 +196,8 @@ def rmsnorm_rht_amax_wrapper_sm100( if resolved_num_threads is None: raise ValueError(f"No valid num_threads found for N={n}") resolved_rows_per_cta = rows_per_cta if rows_per_cta is not None else pick_rows_per_cta(m) + if resolved_rows_per_cta <= 0: + raise ValueError(f"rows_per_cta must be positive, got {resolved_rows_per_cta}") if m % resolved_rows_per_cta != 0: raise ValueError(f"M must be divisible by rows_per_cta, got M={m}, rows_per_cta={resolved_rows_per_cta}") @@ -270,14 +212,14 @@ def rmsnorm_rht_amax_wrapper_sm100( tuple(x_tensor.stride()), tuple(w_tensor.stride()), tuple(o_tensor.stride()), + x_tensor.device, eps, resolved_num_threads, resolved_rows_per_cta, ) - if cache_key in _cache_of_RmsNormRhtAmaxSm100Objects: - api = _cache_of_RmsNormRhtAmaxSm100Objects[cache_key] - else: + api = _cache_of_RmsNormRhtAmaxSm100Objects.get(cache_key) + if api is None: api = RmsNormRhtAmaxSm100( sample_x=x_tensor, sample_w=w_tensor, @@ -287,7 +229,7 @@ def rmsnorm_rht_amax_wrapper_sm100( num_threads=resolved_num_threads, rows_per_cta=resolved_rows_per_cta, ) - assert api.check_support(), "Unsupported configuration" + api.check_support() api.compile() _cache_of_RmsNormRhtAmaxSm100Objects[cache_key] = api diff --git a/python/cudnn/rmsnorm_rht_amax/jax.py b/python/cudnn/rmsnorm_rht_amax/jax.py new file mode 100644 index 000000000..0d94335cc --- /dev/null +++ b/python/cudnn/rmsnorm_rht_amax/jax.py @@ -0,0 +1,139 @@ +# 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 + +from .. import data_type +from .._jax import JaxApiBase, JaxTensorDesc +from .kernel import RMSNormRHTAmaxKernel +from .op import RmsNormRhtAmaxSm100Op, pick_rows_per_cta + + +class RmsNormRhtAmaxSm100(JaxApiBase): + """JAX callable specialized from sample shape and dtype metadata.""" + + def __init__( + self, + sample_x: Any, + sample_w: Any, + *, + sample_o: Any | None = None, + sample_amax: Any | None = None, + eps: float = 1e-5, + num_threads: Optional[int] = None, + rows_per_cta: Optional[int] = None, + ) -> None: + self.x_desc = self._to_tensor_desc(sample_x, "sample_x") + self.w_desc = self._to_tensor_desc(sample_w, "sample_w") + if (sample_o is None) != (sample_amax is None): + raise ValueError("sample_o and sample_amax must be provided together") + + if sample_o is None: + self.o_desc, self.amax_desc = self._default_output_descs(rows_per_cta) + else: + self.o_desc = self._to_tensor_desc(sample_o, "sample_o") + self.amax_desc = self._to_tensor_desc(sample_amax, "sample_amax") + + self._op = RmsNormRhtAmaxSm100Op( + x=self.x_desc, + weight=self.w_desc, + output=self.o_desc, + amax=self.amax_desc, + eps=eps, + num_threads=num_threads, + rows_per_cta=rows_per_cta, + ) + + def _default_output_descs(self, rows_per_cta: Optional[int]) -> tuple[JaxTensorDesc, JaxTensorDesc]: + if self.x_desc.ndim != 2: + raise ValueError(f"X must have rank 2, got shape {self.x_desc.shape}") + m, n = self.x_desc.shape + resolved_rows_per_cta = pick_rows_per_cta(m) if rows_per_cta is None else rows_per_cta + if resolved_rows_per_cta <= 0: + raise ValueError(f"rows_per_cta must be positive, got {resolved_rows_per_cta}") + + return ( + self.x_desc.compact_like( + cudnn_dtype=data_type.BFLOAT16, + shape=(m, n), + name="sample_o", + ), + self.x_desc.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(m // resolved_rows_per_cta,), + name="sample_amax", + ), + ) + + def check_support(self) -> bool: + self._op.check_support() + self._check_device_compatibility( + minimum_compute_capability=100, + operation_name="RmsNormRhtAmaxSm100", + ) + return True + + def __call__(self, x: Any, weight: Any) -> tuple[Any, Any]: + self.check_support() + + x_desc = self.x_desc.with_divisibility( + (self._op.rows_per_cta, 16) + ) + weight_desc = self.w_desc.with_divisibility((16,)) + output_desc = self.o_desc.with_divisibility( + (self._op.rows_per_cta, 16) + ) + + def launch( + stream: Any, + x: Any, + weight: Any, + output: Any, + amax: Any, + ) -> None: + kernel = RMSNormRHTAmaxKernel( + n=self._op.n, + num_threads=self._op.num_threads, + eps=self._op.eps, + rows_per_cta=self._op.rows_per_cta, + ) + kernel(x, weight, output, amax, stream) + + return self._call_kernel( + (x, weight), + launch=launch, + output_descs=(output_desc, self.amax_desc), + input_descs=(x_desc, weight_desc), + ) + + +@jax.jit(static_argnames=("eps", "num_threads", "rows_per_cta")) +def rmsnorm_rht_amax_sm100( + x: Any, + weight: Any, + *, + eps: float = 1e-5, + num_threads: Optional[int] = None, + rows_per_cta: Optional[int] = None, +) -> tuple[Any, Any]: + """Apply fused RMSNorm, 16-wide RHT, and per-CTA amax from JAX.""" + + return RmsNormRhtAmaxSm100( + x, + weight, + eps=eps, + num_threads=num_threads, + rows_per_cta=rows_per_cta, + )(x, weight) + + +__all__ = [ + "RmsNormRhtAmaxSm100", + "rmsnorm_rht_amax_sm100", +] diff --git a/python/cudnn/rmsnorm_rht_amax/kernel.py b/python/cudnn/rmsnorm_rht_amax/kernel.py index 2445cae94..01715d98b 100644 --- a/python/cudnn/rmsnorm_rht_amax/kernel.py +++ b/python/cudnn/rmsnorm_rht_amax/kernel.py @@ -3,6 +3,8 @@ """CUTE DSL kernel for fused RMSNorm + RHT + per-CTA amax.""" +from __future__ import annotations + import math import operator @@ -10,7 +12,7 @@ import cutlass import cutlass.cute as cute import cutlass.utils as utils -from cutlass import Float32, Int32 +from cutlass import Float32 from cutlass._mlir.dialects import llvm from cutlass.cute.arch import shuffle_sync_bfly from cutlass.cutlass_dsl import T, dsl_user_op @@ -68,35 +70,67 @@ def redux_sync_max_f32(val, *, loc=None, ip=None): return Float32(result) +@dsl_user_op +def _as_global_tensor(tensor: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor: + """Preserve a generic kernel argument's global address space for cp.async.""" + + if tensor.memspace == cute.AddressSpace.gmem: + return tensor + if tensor.memspace != cute.AddressSpace.generic: + raise ValueError(f"Expected a global or generic tensor, got {tensor.memspace}") + + pointer = tensor.iterator + global_pointer = llvm.addrspacecast( + llvm.PointerType.get(cute.AddressSpace.gmem), + pointer.llvm_ptr, + loc=loc, + ip=ip, + ) + return cute.make_tensor( + cute.make_ptr( + tensor.element_type, + global_pointer, + cute.AddressSpace.gmem, + assumed_align=pointer.alignment, + loc=loc, + ip=ip, + ), + tensor.layout, + ) + + class RMSNormRHTAmaxKernel: - """Fused RMSNorm + block-diagonal Hadamard + running per-CTA amax.""" + """CuTe implementation of fused RMSNorm + RHT + per-CTA amax.""" COPY_BITS = 128 HAD_BLOCK = 16 - def __init__(self, n, num_threads=256, eps=1e-5, rows_per_cta=8): + def __init__( + self, + *, + n: int, + num_threads: int, + eps: float, + rows_per_cta: int, + ) -> None: self.n = n self.num_threads = num_threads self.eps = eps self.rows_per_cta = rows_per_cta self.vec_size = self.COPY_BITS // 16 - self.ept = n // num_threads - - assert n % num_threads == 0, f"N={n} must be divisible by num_threads={num_threads}" - assert self.ept % self.vec_size == 0, f"EPT={self.ept} must be a multiple of vec_size={self.vec_size}" - assert self.ept >= self.vec_size, f"EPT={self.ept} must be >= vec_size={self.vec_size}" + self.ept = self.n // self.num_threads self.num_vec_blocks = self.ept // self.vec_size - self.warps_per_row = num_threads // 32 + self.warps_per_row = self.num_threads // 32 self.inv_sqrt_had = 1.0 / math.sqrt(self.HAD_BLOCK) self.num_intra_stages = int(math.log2(self.vec_size)) self.num_cross_stages = 1 - self.tv_shape = ((num_threads, 1), (self.vec_size, self.num_vec_blocks)) - self.tv_stride = ((self.vec_size, 1), (1, self.vec_size * num_threads)) - self.tiler_mn = (1, n) + self.tv_shape = ((self.num_threads, 1), (self.vec_size, self.num_vec_blocks)) + self.tv_stride = ((self.vec_size, 1), (1, self.vec_size * self.num_threads)) + self.tiler_mn = (1, self.n) - tile_bytes = n * 2 + tile_bytes = self.n * 2 reduce_bytes = self.warps_per_row * 4 amax_bytes = self.warps_per_row * 4 self.smem_bytes = tile_bytes + reduce_bytes + amax_bytes + 128 @@ -114,6 +148,7 @@ def __init__(self, n, num_threads=256, eps=1e-5, rows_per_cta=8): @cute.kernel def kernel(self, m_x: cute.Tensor, m_w: cute.Tensor, m_o: cute.Tensor, m_amax: cute.Tensor, eps: Float32, tv_layout: cute.Layout, tiler_mn: cute.Shape): cfg = self + m_x = _as_global_tensor(m_x) tid = cute.arch.thread_idx()[0] bid = cute.arch.block_idx()[0] inv_sqrt_had = cutlass.Float32(cfg.inv_sqrt_had) @@ -241,11 +276,26 @@ def kernel(self, m_x: cute.Tensor, m_w: cute.Tensor, m_o: cute.Tensor, m_amax: c m_amax[bid] = cta_max @cute.jit - def __call__(self, x_tensor: cute.Tensor, w_tensor: cute.Tensor, o_tensor: cute.Tensor, amax_tensor: cute.Tensor, eps: Float32, stream: cuda.CUstream): + def __call__( + self, + x_tensor: cute.Tensor, + w_tensor: cute.Tensor, + o_tensor: cute.Tensor, + amax_tensor: cute.Tensor, + stream: cuda.CUstream, + ): m = x_tensor.shape[0] num_ctas = m // self.rows_per_cta tv_layout = cute.make_layout(self.tv_shape, stride=self.tv_stride) - self.kernel(x_tensor, w_tensor, o_tensor, amax_tensor, eps, tv_layout, self.tiler_mn).launch( + self.kernel( + x_tensor, + w_tensor, + o_tensor, + amax_tensor, + Float32(self.eps), + tv_layout, + self.tiler_mn, + ).launch( grid=(num_ctas, 1, 1), block=(self.num_threads, 1, 1), smem=self.smem_bytes, diff --git a/python/cudnn/rmsnorm_rht_amax/op.py b/python/cudnn/rmsnorm_rht_amax/op.py new file mode 100644 index 000000000..b3b967a8b --- /dev/null +++ b/python/cudnn/rmsnorm_rht_amax/op.py @@ -0,0 +1,193 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Framework-neutral RMSNorm + RHT + per-CTA amax operation.""" + +from __future__ import annotations + +from typing import Any, Optional + +from .. import data_type +from ..common.op import Op +from ..common.tensor_desc import TensorDesc + +HAD_BLOCK = 16 + +DEFAULT_NUM_THREADS_BY_N = { + 2048: 128, + 4096: 256, + 7168: 128, + 8192: 512, + 16384: 1024, + 32768: 512, +} +RPC_CANDIDATES = (2, 4, 8) +TARGET_MIN_CTAS = 148 + + +def best_num_threads(n: int) -> Optional[int]: + for num_threads in (1024, 512, 256, 128, 64): + if n % num_threads != 0: + continue + ept = n // num_threads + if ept >= 8 and ept % 8 == 0: + return num_threads + return None + + +def pick_rows_per_cta(m: int) -> int: + for rows_per_cta in reversed(RPC_CANDIDATES): + if m % rows_per_cta != 0: + continue + if m // rows_per_cta >= TARGET_MIN_CTAS: + return rows_per_cta + return RPC_CANDIDATES[0] + + +class RmsNormRhtAmaxSm100Op(Op): + """Complete logical signature and launch configuration for the SM100 op.""" + + def __init__( + self, + *, + x: TensorDesc[Any], + weight: TensorDesc[Any], + output: TensorDesc[Any], + amax: TensorDesc[Any], + eps: float = 1e-5, + num_threads: Optional[int] = None, + rows_per_cta: Optional[int] = None, + ) -> None: + for name, desc in (("x", x), ("weight", weight), ("output", output), ("amax", amax)): + if not isinstance(desc, TensorDesc): + raise TypeError(f"{name} must be a TensorDesc, got {type(desc).__name__}") + + self.x = x + self.weight = weight + self.output = output + self.amax = amax + self.eps = eps + self.requested_num_threads = num_threads + self.requested_rows_per_cta = rows_per_cta + + self.m: Optional[int] = None + self.n: Optional[int] = None + self.num_threads: Optional[int] = None + self.rows_per_cta: Optional[int] = None + + def check_support(self) -> bool: + """Validate the complete signature and resolve the launch configuration.""" + + self.m = None + self.n = None + self.num_threads = None + self.rows_per_cta = None + + self._check_input_descriptors() + m, n = self.x.shape + + num_threads = self.requested_num_threads + if num_threads is None: + num_threads = DEFAULT_NUM_THREADS_BY_N.get(n, best_num_threads(n)) + if num_threads is None: + raise ValueError(f"No valid num_threads found for N={n}") + + rows_per_cta = self._resolve_rows_per_cta(m, self.requested_rows_per_cta) + self._validate_launch_configuration(n, num_threads, rows_per_cta, m=m) + self._check_output_descriptors(m, n, rows_per_cta) + + self.m = m + self.n = n + self.num_threads = num_threads + self.rows_per_cta = rows_per_cta + return True + + def _check_input_descriptors(self) -> None: + if self.x.ndim != 2: + raise ValueError(f"X must have rank 2, got shape {self.x.shape}") + if self.weight.ndim != 1: + raise ValueError(f"W must have rank 1, got shape {self.weight.shape}") + if self.x.cudnn_dtype != data_type.BFLOAT16: + raise ValueError(f"X must have dtype bfloat16, got {self.x.dtype}") + if self.weight.cudnn_dtype != data_type.BFLOAT16: + raise ValueError(f"W must have dtype bfloat16, got {self.weight.dtype}") + + m, n = self.x.shape + if m <= 0: + raise ValueError(f"M must be positive, got {m}") + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if self.weight.shape != (n,): + raise ValueError(f"W must have shape {(n,)}, got {self.weight.shape}") + if self.x.stride != (n, 1) or self.x.stride_order != (1, 0): + raise ValueError(f"X must be row-major contiguous, got stride {self.x.stride} and stride order {self.x.stride_order}") + if self.weight.stride != (1,) or self.weight.stride_order != (0,): + raise ValueError(f"W must be contiguous, got stride {self.weight.stride} and stride order {self.weight.stride_order}") + if n % HAD_BLOCK != 0: + raise ValueError(f"N must be divisible by {HAD_BLOCK} for the Hadamard block size, got {n}") + + def _check_output_descriptors(self, m: int, n: int, rows_per_cta: int) -> None: + if self.output.ndim != 2: + raise ValueError(f"O must have rank 2, got shape {self.output.shape}") + if self.amax.ndim != 1: + raise ValueError(f"Amax must have rank 1, got shape {self.amax.shape}") + if self.output.cudnn_dtype != data_type.BFLOAT16: + raise ValueError(f"O must have dtype bfloat16, got {self.output.dtype}") + if self.amax.cudnn_dtype != data_type.FLOAT: + raise ValueError(f"Amax must have dtype float32, got {self.amax.dtype}") + if self.output.shape != (m, n): + raise ValueError(f"O must have shape {(m, n)}, got {self.output.shape}") + expected_amax_shape = (m // rows_per_cta,) + if self.amax.shape != expected_amax_shape: + raise ValueError(f"Amax must have shape {expected_amax_shape}, got {self.amax.shape}") + if self.output.stride != (n, 1) or self.output.stride_order != (1, 0): + raise ValueError(f"O must be row-major contiguous, got stride {self.output.stride} and stride order {self.output.stride_order}") + if self.amax.stride == (0,): + raise ValueError("Amax stride must be positive") + + @staticmethod + def _resolve_rows_per_cta(m: int, rows_per_cta: Optional[int]) -> int: + if rows_per_cta is None: + rows_per_cta = pick_rows_per_cta(m) + if rows_per_cta <= 0: + raise ValueError(f"rows_per_cta must be positive, got {rows_per_cta}") + if m % rows_per_cta != 0: + raise ValueError(f"M must be divisible by rows_per_cta, got M={m}, rows_per_cta={rows_per_cta}") + return rows_per_cta + + @staticmethod + def _validate_launch_configuration( + n: int, + num_threads: int, + rows_per_cta: int, + *, + m: Optional[int] = None, + ) -> None: + if n <= 0: + raise ValueError(f"N must be positive, got {n}") + if num_threads <= 0: + raise ValueError(f"num_threads must be positive, got {num_threads}") + if num_threads % 32 != 0: + raise ValueError(f"num_threads must be warp-aligned, got {num_threads}") + if num_threads > 1024: + raise ValueError(f"num_threads must not exceed the CUDA block size limit, got {num_threads}") + if n % num_threads != 0: + raise ValueError(f"N={n} must be divisible by num_threads={num_threads}") + + ept = n // num_threads + if ept < 8 or ept % 8 != 0: + raise ValueError(f"EPT={ept} must be >= 8 and divisible by 8") + if rows_per_cta <= 0: + raise ValueError(f"rows_per_cta must be positive, got {rows_per_cta}") + if m is not None and m % rows_per_cta != 0: + raise ValueError(f"M must be divisible by rows_per_cta, got M={m}, rows_per_cta={rows_per_cta}") + + +__all__ = [ + "DEFAULT_NUM_THREADS_BY_N", + "RPC_CANDIDATES", + "TARGET_MIN_CTAS", + "RmsNormRhtAmaxSm100Op", + "best_num_threads", + "pick_rows_per_cta", +] diff --git a/python/cudnn/sdpa/__init__.py b/python/cudnn/sdpa/__init__.py index b200dcbda..a77765b5f 100644 --- a/python/cudnn/sdpa/__init__.py +++ b/python/cudnn/sdpa/__init__.py @@ -1,12 +1,14 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .bwd import SdpabwdSm100D256, sdpa_bwd_wrapper_sm100_d256 -from .fwd import SdpafwdSm100D256, sdpa_fwd_wrapper_sm100_d256 +"""Lazy Torch API exports for the CuTe SDPA operations.""" -__all__ = [ - "SdpafwdSm100D256", - "sdpa_fwd_wrapper_sm100_d256", - "SdpabwdSm100D256", - "sdpa_bwd_wrapper_sm100_d256", -] +from ..common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={ + "fwd": ("SdpafwdSm100D256", "sdpa_fwd_wrapper_sm100_d256"), + "bwd": ("SdpabwdSm100D256", "sdpa_bwd_wrapper_sm100_d256"), + }, +) diff --git a/python/cudnn/sdpa/bwd/__init__.py b/python/cudnn/sdpa/bwd/__init__.py index 0d0c368ed..55775acd8 100644 --- a/python/cudnn/sdpa/bwd/__init__.py +++ b/python/cudnn/sdpa/bwd/__init__.py @@ -1,9 +1,12 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import SdpabwdSm100D256, sdpa_bwd_wrapper_sm100_d256 +"""Lazy Torch SDPA-backward API exports.""" -__all__ = [ - "SdpabwdSm100D256", - "sdpa_bwd_wrapper_sm100_d256", -] +from ...common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("SdpabwdSm100D256", "sdpa_bwd_wrapper_sm100_d256")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/sdpa/bwd/jax.py b/python/cudnn/sdpa/bwd/jax.py new file mode 100644 index 000000000..3af81d368 --- /dev/null +++ b/python/cudnn/sdpa/bwd/jax.py @@ -0,0 +1,693 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for fixed BHSD/BSHD and packed THD SM100 d=256 SDPA backward.""" + +from __future__ import annotations + +from functools import partial +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._jax.compiler import compile_options_for_target +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..jax_utils import ( + FIXED_LAYOUTS, + describe_fixed_data, + fixed_data_mode, + make_fixed_output, + normalize_sdpa_layout, + require_fixed_qkv, + require_float32_dtype, + resolve_sdpa_config, +) + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) +_PACKED_LAYOUT = "THD" + + +class SdpabwdSm100D256(JaxApiBase): + """JAX callable specialized from fixed BHSD/BSHD or packed THD metadata. + + Packed THD inputs require cumulative sequence lengths and explicit static + ``max_s_q``/``max_s_k`` bounds. JAX samples expose array metadata rather + than cumulative-length values, so these bounds cannot be inferred while + tracing. + """ + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_o: Any, + sample_do: Any, + sample_lse: Any, + sample_cum_seqlen_q: Any | None = None, + sample_cum_seqlen_k: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + dkdv_mma_tiler_mn: tuple[int, int] = (128, 64), + is_causal: bool = False, + window_size: tuple[int, int] = (-1, -1), + scale_softmax: float | None = None, + layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + ranks = tuple( + len(tuple(sample.shape)) + for sample in (sample_q, sample_k, sample_v, sample_o, sample_do) + ) + if len(set(ranks)) != 1: + raise ValueError(f"Q, K, V, O, and dO must use the same rank, got {ranks}") + self.input_layout = normalize_sdpa_layout(layout, ranks[0]) + if self.input_layout in FIXED_LAYOUTS: + if max_s_q is not None or max_s_k is not None: + raise ValueError("max_s_q and max_s_k are only valid for THD layout") + self.data_mode = fixed_data_mode(self.input_layout) + self._init_fixed_shape( + sample_q, + sample_k, + sample_v, + sample_o, + sample_do, + sample_lse, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + ) + elif self.input_layout == _PACKED_LAYOUT: + self.data_mode = None + self._init_packed_shape( + sample_q, + sample_k, + sample_v, + sample_o, + sample_do, + sample_lse, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + max_s_q, + max_s_k, + ) + + require_float32_dtype(acc_dtype, "acc_dtype") + if tuple(mma_tiler_mn) != (128, 128): + raise ValueError(f"mma_tiler_mn must be (128, 128), got {mma_tiler_mn}") + if tuple(dkdv_mma_tiler_mn) != (128, 64): + raise ValueError( + f"dkdv_mma_tiler_mn must be (128, 64), got {dkdv_mma_tiler_mn}" + ) + + self.is_causal = bool(is_causal) + ( + self.scale_softmax, + self.window_size_left, + self.window_size_right, + self.mask_kind, + ) = resolve_sdpa_config( + seqlen_q=self.max_s_q, + seqlen_k=self.max_s_k, + # The Torch API always selects the residual mask for non-causal + # varlen inputs, independent of the packed token count. + tile_extent=self.seqlen_q if self.input_layout in FIXED_LAYOUTS else 1, + is_causal=self.is_causal, + window_size=tuple(window_size), + scale_softmax=scale_softmax, + ) + self.target_compute_capability = target_compute_capability + self.compute_capability: int | None = None + + if self.input_layout in FIXED_LAYOUTS: + self.dq_desc = make_fixed_output( + tuple(sample_q.shape), + sample_q.dtype, + "dq_tensor", + layout=self.input_layout, + ) + self.dk_desc = make_fixed_output( + tuple(sample_k.shape), + sample_k.dtype, + "dk_tensor", + layout=self.input_layout, + ) + self.dv_desc = make_fixed_output( + tuple(sample_v.shape), + sample_v.dtype, + "dv_tensor", + layout=self.input_layout, + ) + self.dq_kernel_desc = self.dq_desc + self.dk_kernel_desc = self.dk_desc + self.dv_kernel_desc = self.dv_desc + self.workspace_seqlen_q = self.seqlen_q + self.workspace_batch = self.batch + else: + self.dq_desc = JaxTensorDesc.from_shape( + tuple(sample_q.shape), sample_q.dtype, name="dq_tensor" + ) + self.dk_desc = JaxTensorDesc.from_shape( + tuple(sample_k.shape), sample_k.dtype, name="dk_tensor" + ) + self.dv_desc = JaxTensorDesc.from_shape( + tuple(sample_v.shape), sample_v.dtype, name="dv_tensor" + ) + self.dq_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_q.shape), sample_q.dtype, name="dq_kernel_tensor" + ) + self.dk_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_k.shape), sample_k.dtype, name="dk_kernel_tensor" + ) + self.dv_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_v.shape), sample_v.dtype, name="dv_kernel_tensor" + ) + # Workspace size follows the promoted (1, T, H, D) ABI, matching + # the Torch adapter's use of s_qo=T and b=1. + self.workspace_seqlen_q = self.total_q_tokens + self.workspace_batch = 1 + + def _init_fixed_shape( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_o: Any, + sample_do: Any, + sample_lse: Any, + sample_cum_seqlen_q: Any | None, + sample_cum_seqlen_k: Any | None, + ) -> None: + if sample_cum_seqlen_q is not None or sample_cum_seqlen_k is not None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must be omitted for fixed layout" + ) + self.q_desc = describe_fixed_data( + sample_q, "sample_q", layout=self.input_layout + ) + self.k_desc = describe_fixed_data( + sample_k, "sample_k", layout=self.input_layout + ) + self.v_desc = describe_fixed_data( + sample_v, "sample_v", layout=self.input_layout + ) + self.o_desc = describe_fixed_data( + sample_o, "sample_o", layout=self.input_layout + ) + self.do_desc = describe_fixed_data( + sample_do, "sample_do", layout=self.input_layout + ) + self.lse_desc = self._to_tensor_desc(sample_lse, "sample_lse") + ( + self.batch, + self.num_query_heads, + self.num_kv_heads, + self.seqlen_q, + self.seqlen_k, + self.head_dim, + ) = require_fixed_qkv(self.q_desc, self.k_desc, self.v_desc) + for desc in (self.o_desc, self.do_desc): + if ( + desc.shape != self.q_desc.shape + or desc.cudnn_dtype != self.q_desc.cudnn_dtype + ): + raise ValueError(f"{desc.name} must match sample_q's shape and dtype") + expected_lse = (self.batch, self.num_query_heads, self.seqlen_q) + if ( + self.lse_desc.shape != expected_lse + or self.lse_desc.cudnn_dtype != data_type.FLOAT + ): + raise ValueError( + f"sample_lse must have shape {expected_lse} and dtype float32" + ) + self.total_q_tokens = self.batch * self.seqlen_q + self.total_k_tokens = self.batch * self.seqlen_k + self.max_s_q = self.seqlen_q + self.max_s_k = self.seqlen_k + self.cum_q_desc = None + self.cum_k_desc = None + self.q_kernel_desc = self.q_desc + self.k_kernel_desc = self.k_desc + self.v_kernel_desc = self.v_desc + self.o_kernel_desc = self.o_desc + self.do_kernel_desc = self.do_desc + self.lse_kernel_desc = self.lse_desc + + def _init_packed_shape( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_o: Any, + sample_do: Any, + sample_lse: Any, + sample_cum_seqlen_q: Any | None, + sample_cum_seqlen_k: Any | None, + max_s_q: int | None, + max_s_k: int | None, + ) -> None: + if sample_cum_seqlen_q is None or sample_cum_seqlen_k is None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k are both required for THD layout" + ) + if max_s_q is None or max_s_k is None: + raise ValueError("max_s_q and max_s_k are both required for THD layout") + + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.k_desc = self._to_tensor_desc(sample_k, "sample_k") + self.v_desc = self._to_tensor_desc(sample_v, "sample_v") + self.o_desc = self._to_tensor_desc(sample_o, "sample_o") + self.do_desc = self._to_tensor_desc(sample_do, "sample_do") + self.lse_desc = self._to_tensor_desc(sample_lse, "sample_lse") + self.cum_q_desc = self._to_tensor_desc( + sample_cum_seqlen_q, "sample_cum_seqlen_q" + ) + self.cum_k_desc = self._to_tensor_desc( + sample_cum_seqlen_k, "sample_cum_seqlen_k" + ) + + self.total_q_tokens, self.num_query_heads, self.head_dim = self.q_desc.shape + self.total_k_tokens, self.num_kv_heads, k_head_dim = self.k_desc.shape + v_total, v_heads, value_dim = self.v_desc.shape + dimensions = ( + self.total_q_tokens, + self.total_k_tokens, + self.num_query_heads, + self.num_kv_heads, + self.head_dim, + ) + if any(value <= 0 for value in dimensions): + raise ValueError(f"SDPA dimensions must be positive, got {dimensions}") + if (v_total, v_heads, value_dim) != ( + self.total_k_tokens, + self.num_kv_heads, + self.head_dim, + ) or k_head_dim != self.head_dim: + raise ValueError( + "K and V must share packed token, head, and head-dim metadata" + ) + if self.q_desc.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError( + f"SDPA requires float16 or bfloat16 inputs, got {self.q_desc.dtype}" + ) + if ( + self.k_desc.cudnn_dtype != self.q_desc.cudnn_dtype + or self.v_desc.cudnn_dtype != self.q_desc.cudnn_dtype + ): + raise ValueError("Q, K, and V must have the same dtype") + for desc in (self.o_desc, self.do_desc): + if ( + desc.shape != self.q_desc.shape + or desc.cudnn_dtype != self.q_desc.cudnn_dtype + ): + raise ValueError(f"{desc.name} must match sample_q's shape and dtype") + if self.head_dim != 256: + raise ValueError(f"head dimension must be 256, got {self.head_dim}") + if self.num_query_heads % self.num_kv_heads: + raise ValueError( + f"H_q ({self.num_query_heads}) must be divisible by H_kv ({self.num_kv_heads})" + ) + expected_lse = (self.total_q_tokens, self.num_query_heads) + if ( + self.lse_desc.shape != expected_lse + or self.lse_desc.cudnn_dtype != data_type.FLOAT + ): + raise ValueError( + f"sample_lse must have shape {expected_lse} and dtype float32" + ) + if ( + self.cum_q_desc.ndim != 1 + or self.cum_q_desc.shape != self.cum_k_desc.shape + or self.cum_q_desc.shape[0] < 2 + ): + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must have the same shape (B + 1,) with B > 0" + ) + if ( + self.cum_q_desc.cudnn_dtype != data_type.INT32 + or self.cum_k_desc.cudnn_dtype != data_type.INT32 + ): + raise ValueError("cum_seqlen_q and cum_seqlen_k must have dtype int32") + + self.batch = self.cum_q_desc.shape[0] - 1 + self.seqlen_q = self.total_q_tokens + self.seqlen_k = self.total_k_tokens + self.max_s_q = int(max_s_q) + self.max_s_k = int(max_s_k) + if ( + self.max_s_q <= 0 + or self.max_s_k <= 0 + or self.max_s_q > self.total_q_tokens + or self.max_s_k > self.total_k_tokens + ): + raise ValueError( + "max_s_q and max_s_k must be positive and no larger than their packed token counts" + ) + + self.q_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_q.shape), sample_q.dtype, name="q_tensor" + ) + self.k_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_k.shape), sample_k.dtype, name="k_tensor" + ) + self.v_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_v.shape), sample_v.dtype, name="v_tensor" + ) + self.o_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_o.shape), sample_o.dtype, name="o_tensor" + ) + self.do_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_do.shape), sample_do.dtype, name="do_tensor" + ) + # The kernel addresses packed LSE as (1, H, T) with T contiguous. + self.lse_kernel_desc = JaxTensorDesc.from_shape( + (1, self.num_query_heads, self.total_q_tokens), + jnp.float32, + name="lse_tensor", + ) + + def check_support(self) -> bool: + self.compute_capability = self._resolve_compute_capability( + self.target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "SdpabwdSm100D256", + ) + return True + + def __call__( + self, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + o_tensor: Any, + do_tensor: Any, + lse_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + ) -> TupleDict: + self.check_support() + if self.input_layout in FIXED_LAYOUTS: + if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must be omitted for fixed layout" + ) + inputs = (q_tensor, k_tensor, v_tensor, o_tensor, do_tensor, lse_tensor) + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.v_kernel_desc, + self.o_kernel_desc, + self.do_kernel_desc, + self.lse_kernel_desc, + ) + launch = self._launch_kernel + else: + if cum_seqlen_q_tensor is None or cum_seqlen_k_tensor is None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k are both required for THD layout" + ) + for value, desc in ( + (q_tensor, self.q_desc), + (k_tensor, self.k_desc), + (v_tensor, self.v_desc), + (o_tensor, self.o_desc), + (do_tensor, self.do_desc), + (lse_tensor, self.lse_desc), + ): + self._check_tensor_signature(value, desc) + lse_storage = jnp.transpose( + jnp.reshape( + lse_tensor, + (1, self.total_q_tokens, self.num_query_heads), + ), + (0, 2, 1), + ) + inputs = ( + jnp.reshape(q_tensor, self.q_kernel_desc.shape), + jnp.reshape(k_tensor, self.k_kernel_desc.shape), + jnp.reshape(v_tensor, self.v_kernel_desc.shape), + jnp.reshape(o_tensor, self.o_kernel_desc.shape), + jnp.reshape(do_tensor, self.do_kernel_desc.shape), + lse_storage, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.v_kernel_desc, + self.o_kernel_desc, + self.do_kernel_desc, + self.lse_kernel_desc, + self.cum_q_desc, + self.cum_k_desc, + ) + launch = self._launch_varlen_kernel + + dq, dk, dv = self._call_kernel( + inputs, + launch=launch, + output_descs=( + self.dq_kernel_desc, + self.dk_kernel_desc, + self.dv_kernel_desc, + ), + input_descs=input_descs, + workspace_descs=(self._workspace_desc(),), + compile_options=compile_options_for_target(self.compute_capability), + ) + if self.input_layout == _PACKED_LAYOUT: + dq = jnp.reshape(dq, self.dq_desc.shape) + dk = jnp.reshape(dk, self.dk_desc.shape) + dv = jnp.reshape(dv, self.dv_desc.shape) + return TupleDict(dq_tensor=dq, dk_tensor=dk, dv_tensor=dv) + + def _workspace_desc(self) -> JaxTensorDesc: + from cutlass import Float32 + + from .fmha_backward_sm100_2kernel import ( + BlackwellFusedMultiHeadAttentionBackward, + ) + + shape = BlackwellFusedMultiHeadAttentionBackward.get_workspace_size( + self.workspace_seqlen_q, + self.head_dim, + self.num_query_heads, + self.workspace_batch, + Float32, + ) + return self.q_kernel_desc.compact_like( + cudnn_dtype=data_type.UINT8, + shape=shape, + name="workspace", + init_value=0, + ) + + def _launch_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + output: Any, + doutput: Any, + lse: Any, + dq: Any, + dk: Any, + dv: Any, + workspace: Any, + ) -> None: + self._invoke_kernel( + stream, + q, + k, + v, + output, + doutput, + lse, + dq, + dk, + dv, + workspace, + None, + None, + ) + + def _launch_varlen_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + output: Any, + doutput: Any, + lse: Any, + cum_seqlen_q: Any, + cum_seqlen_k: Any, + dq: Any, + dk: Any, + dv: Any, + workspace: Any, + ) -> None: + self._invoke_kernel( + stream, + q, + k, + v, + output, + doutput, + lse, + dq, + dk, + dv, + workspace, + cum_seqlen_q, + cum_seqlen_k, + ) + + def _invoke_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + output: Any, + doutput: Any, + lse: Any, + dq: Any, + dk: Any, + dv: Any, + workspace: Any, + cum_seqlen_q: Any | None, + cum_seqlen_k: Any | None, + ) -> None: + from cutlass import Float32, Int32 + from cutlass.jax import jax_to_cutlass_dtype + + from ..fmha_utils import MaskEnum + from .fmha_backward_sm100_2kernel import ( + BlackwellFusedMultiHeadAttentionBackward, + ) + + mask_type = { + "residual": MaskEnum.RESIDUAL_MASK, + "window": MaskEnum.WINDOW_MASK_INFERENCE, + }[self.mask_kind] + kernel = BlackwellFusedMultiHeadAttentionBackward( + element_dtype=jax_to_cutlass_dtype(self.q_desc.dtype), + acc_dtype=Float32, + mma_tiler=(128, 128, 256), + dkdv_mma_tiler=(128, 64, 256), + varlen=self.input_layout == _PACKED_LAYOUT, + is_causal=self.is_causal, + mask_type=mask_type, + window_size_left=self.window_size_left, + window_size_right=self.window_size_right, + ) + problem_shape = ( + Int32(self.max_s_q), + Int32(self.max_s_k), + Int32(256), + ( + ( + Int32(self.num_query_heads // self.num_kv_heads), + Int32(self.num_kv_heads), + ), + Int32(self.batch), + ), + ) + kernel( + problem_shape, + q, + k, + v, + output, + dq, + dk, + dv, + doutput, + lse, + cum_seqlen_q, + cum_seqlen_k, + Float32(self.scale_softmax), + workspace, + stream, + ) + + +@partial( + jax.jit, + static_argnames=( + "max_s_q", + "max_s_k", + "acc_dtype", + "mma_tiler_mn", + "dkdv_mma_tiler_mn", + "is_causal", + "window_size", + "scale_softmax", + "layout", + "target_compute_capability", + ), +) +def sdpa_bwd_wrapper_sm100_d256( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + o_tensor: Any, + do_tensor: Any, + lse_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + dkdv_mma_tiler_mn: tuple[int, int] = (128, 64), + is_causal: bool = False, + window_size: tuple[int, int] = (-1, -1), + scale_softmax: float | None = None, + layout: str | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute fixed BHSD/BSHD or packed THD SDPA gradients. + + Packed calls must provide both cumulative sequence-length arrays and static + ``max_s_q``/``max_s_k`` bounds. Fixed gradients follow ``layout``; LSE + remains ``(B, H, S)``. + """ + + return SdpabwdSm100D256( + q_tensor, + k_tensor, + v_tensor, + o_tensor, + do_tensor, + lse_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + max_s_q=max_s_q, + max_s_k=max_s_k, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + dkdv_mma_tiler_mn=dkdv_mma_tiler_mn, + is_causal=is_causal, + window_size=window_size, + scale_softmax=scale_softmax, + layout=layout, + target_compute_capability=target_compute_capability, + )( + *values, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + + +__all__ = [ + "SUPPORTED_COMPUTE_CAPABILITIES", + "SdpabwdSm100D256", + "sdpa_bwd_wrapper_sm100_d256", +] diff --git a/python/cudnn/sdpa/fwd/__init__.py b/python/cudnn/sdpa/fwd/__init__.py index bf95abac2..c2cfef51f 100644 --- a/python/cudnn/sdpa/fwd/__init__.py +++ b/python/cudnn/sdpa/fwd/__init__.py @@ -1,9 +1,12 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: MIT -from .api import SdpafwdSm100D256, sdpa_fwd_wrapper_sm100_d256 +"""Lazy Torch SDPA-forward API exports.""" -__all__ = [ - "SdpafwdSm100D256", - "sdpa_fwd_wrapper_sm100_d256", -] +from ...common.operation_api import make_operation_api + +__all__, __getattr__, __dir__ = make_operation_api( + globals(), + exports={"api": ("SdpafwdSm100D256", "sdpa_fwd_wrapper_sm100_d256")}, + submodules=("api", "jax"), +) diff --git a/python/cudnn/sdpa/fwd/jax.py b/python/cudnn/sdpa/fwd/jax.py new file mode 100644 index 000000000..cedfaab03 --- /dev/null +++ b/python/cudnn/sdpa/fwd/jax.py @@ -0,0 +1,522 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX API for fixed BHSD/BSHD and packed THD SM100 d=256 SDPA forward.""" + +from __future__ import annotations + +from functools import partial +import math +from typing import Any + +import jax +import jax.numpy as jnp + +from ... import data_type +from ..._jax.compiler import compile_options_for_target +from ..._jax import JaxApiBase, JaxTensorDesc, TupleDict +from ..jax_utils import ( + FIXED_LAYOUTS, + describe_fixed_data, + fixed_data_mode, + make_fixed_output, + normalize_sdpa_layout, + require_fixed_qkv, + require_float32_dtype, + resolve_sdpa_config, +) + +SUPPORTED_COMPUTE_CAPABILITIES = (100, 103, 107) +_PACKED_LAYOUT = "THD" + + +class SdpafwdSm100D256(JaxApiBase): + """JAX callable specialized from fixed BHSD/BSHD or packed THD metadata. + + Packed THD inputs require cumulative sequence lengths and explicit static + ``max_s_q``/``max_s_k`` bounds. JAX samples expose array metadata rather + than cumulative-length values, so these bounds cannot be inferred while + tracing. + """ + + def __init__( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_cum_seqlen_q: Any | None = None, + sample_cum_seqlen_k: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + qk_acc_dtype: Any = None, + pv_acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + is_causal: bool = False, + window_size: tuple[int, int] = (-1, -1), + scale_softmax: float | None = None, + scale_output: float = 1.0, + layout: str | None = None, + target_compute_capability: int | None = None, + ) -> None: + ranks = tuple( + len(tuple(sample.shape)) for sample in (sample_q, sample_k, sample_v) + ) + if len(set(ranks)) != 1: + raise ValueError(f"Q, K, and V must use the same rank, got {ranks}") + self.input_layout = normalize_sdpa_layout(layout, ranks[0]) + if self.input_layout in FIXED_LAYOUTS: + if max_s_q is not None or max_s_k is not None: + raise ValueError("max_s_q and max_s_k are only valid for THD layout") + self.data_mode = fixed_data_mode(self.input_layout) + self._init_fixed_shape( + sample_q, + sample_k, + sample_v, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + ) + elif self.input_layout == _PACKED_LAYOUT: + self.data_mode = None + self._init_packed_shape( + sample_q, + sample_k, + sample_v, + sample_cum_seqlen_q, + sample_cum_seqlen_k, + max_s_q, + max_s_k, + ) + + require_float32_dtype(qk_acc_dtype, "qk_acc_dtype") + require_float32_dtype(pv_acc_dtype, "pv_acc_dtype") + if tuple(mma_tiler_mn) != (128, 128): + raise ValueError(f"mma_tiler_mn must be (128, 128), got {mma_tiler_mn}") + self.is_causal = bool(is_causal) + ( + self.scale_softmax, + self.window_size_left, + self.window_size_right, + self.mask_kind, + ) = resolve_sdpa_config( + seqlen_q=self.max_s_q, + seqlen_k=self.max_s_k, + # The Torch API always selects the residual mask for non-causal + # varlen inputs, independent of the packed token count. + tile_extent=self.seqlen_k if self.input_layout in FIXED_LAYOUTS else 1, + is_causal=self.is_causal, + window_size=tuple(window_size), + scale_softmax=scale_softmax, + ) + self.scale_output = float(scale_output) + self.target_compute_capability = target_compute_capability + self.compute_capability: int | None = None + + if self.input_layout in FIXED_LAYOUTS: + self.o_desc = make_fixed_output( + tuple(sample_q.shape), + sample_q.dtype, + "o_tensor", + layout=self.input_layout, + ) + self.o_kernel_desc = self.o_desc + self.lse_desc = JaxTensorDesc.from_shape( + (self.batch, self.num_query_heads, self.seqlen_q), + jnp.float32, + name="lse_tensor", + ) + self.lse_kernel_desc = self.lse_desc + else: + self.o_desc = JaxTensorDesc.from_shape( + tuple(sample_q.shape), + sample_q.dtype, + name="o_tensor", + ) + self.o_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_q.shape), + sample_q.dtype, + name="o_kernel_tensor", + ) + self.lse_desc = JaxTensorDesc.from_shape( + (self.total_q_tokens, self.num_query_heads), + jnp.float32, + name="lse_tensor", + # The kernel writes sequence-contiguous LSE. The descriptor + # records that physical order while preserving public (T, H). + public_stride_order=(0, 1), + ) + # Forward consumes packed LSE through its iterator, just as the + # Torch wrapper passes the rank-2 (T, H) buffer directly. + self.lse_kernel_desc = self.lse_desc + + def _init_fixed_shape( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_cum_seqlen_q: Any | None, + sample_cum_seqlen_k: Any | None, + ) -> None: + if sample_cum_seqlen_q is not None or sample_cum_seqlen_k is not None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must be omitted for fixed layout" + ) + self.q_desc = describe_fixed_data( + sample_q, "sample_q", layout=self.input_layout + ) + self.k_desc = describe_fixed_data( + sample_k, "sample_k", layout=self.input_layout + ) + self.v_desc = describe_fixed_data( + sample_v, "sample_v", layout=self.input_layout + ) + ( + self.batch, + self.num_query_heads, + self.num_kv_heads, + self.seqlen_q, + self.seqlen_k, + self.head_dim, + ) = require_fixed_qkv(self.q_desc, self.k_desc, self.v_desc) + self.total_q_tokens = self.batch * self.seqlen_q + self.total_k_tokens = self.batch * self.seqlen_k + self.max_s_q = self.seqlen_q + self.max_s_k = self.seqlen_k + self.cum_q_desc = None + self.cum_k_desc = None + self.q_kernel_desc = self.q_desc + self.k_kernel_desc = self.k_desc + self.v_kernel_desc = self.v_desc + + def _init_packed_shape( + self, + sample_q: Any, + sample_k: Any, + sample_v: Any, + sample_cum_seqlen_q: Any | None, + sample_cum_seqlen_k: Any | None, + max_s_q: int | None, + max_s_k: int | None, + ) -> None: + if sample_cum_seqlen_q is None or sample_cum_seqlen_k is None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k are both required for THD layout" + ) + if max_s_q is None or max_s_k is None: + raise ValueError("max_s_q and max_s_k are both required for THD layout") + + self.q_desc = self._to_tensor_desc(sample_q, "sample_q") + self.k_desc = self._to_tensor_desc(sample_k, "sample_k") + self.v_desc = self._to_tensor_desc(sample_v, "sample_v") + self.cum_q_desc = self._to_tensor_desc( + sample_cum_seqlen_q, "sample_cum_seqlen_q" + ) + self.cum_k_desc = self._to_tensor_desc( + sample_cum_seqlen_k, "sample_cum_seqlen_k" + ) + + self.total_q_tokens, self.num_query_heads, self.head_dim = self.q_desc.shape + self.total_k_tokens, self.num_kv_heads, k_head_dim = self.k_desc.shape + v_total, v_heads, value_dim = self.v_desc.shape + dimensions = ( + self.total_q_tokens, + self.total_k_tokens, + self.num_query_heads, + self.num_kv_heads, + self.head_dim, + ) + if any(value <= 0 for value in dimensions): + raise ValueError(f"SDPA dimensions must be positive, got {dimensions}") + if (v_total, v_heads, value_dim) != ( + self.total_k_tokens, + self.num_kv_heads, + self.head_dim, + ) or k_head_dim != self.head_dim: + raise ValueError( + "K and V must share packed token, head, and head-dim metadata" + ) + if self.q_desc.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError( + f"SDPA requires float16 or bfloat16 inputs, got {self.q_desc.dtype}" + ) + if ( + self.k_desc.cudnn_dtype != self.q_desc.cudnn_dtype + or self.v_desc.cudnn_dtype != self.q_desc.cudnn_dtype + ): + raise ValueError("Q, K, and V must have the same dtype") + if self.head_dim != 256: + raise ValueError(f"head dimension must be 256, got {self.head_dim}") + if self.num_query_heads % self.num_kv_heads: + raise ValueError( + f"H_q ({self.num_query_heads}) must be divisible by H_kv ({self.num_kv_heads})" + ) + if ( + self.cum_q_desc.ndim != 1 + or self.cum_q_desc.shape != self.cum_k_desc.shape + or self.cum_q_desc.shape[0] < 2 + ): + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must have the same shape (B + 1,) with B > 0" + ) + if ( + self.cum_q_desc.cudnn_dtype != data_type.INT32 + or self.cum_k_desc.cudnn_dtype != data_type.INT32 + ): + raise ValueError("cum_seqlen_q and cum_seqlen_k must have dtype int32") + + self.batch = self.cum_q_desc.shape[0] - 1 + self.seqlen_q = self.total_q_tokens + self.seqlen_k = self.total_k_tokens + self.max_s_q = int(max_s_q) + self.max_s_k = int(max_s_k) + if ( + self.max_s_q <= 0 + or self.max_s_k <= 0 + or self.max_s_q > self.total_q_tokens + or self.max_s_k > self.total_k_tokens + ): + raise ValueError( + "max_s_q and max_s_k must be positive and no larger than their packed token counts" + ) + + # Public THD remains row-major; adding the singleton batch dimension + # yields the exact (1, T, H, D) ABI used by the Torch wrapper. + self.q_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_q.shape), sample_q.dtype, name="q_tensor" + ) + self.k_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_k.shape), sample_k.dtype, name="k_tensor" + ) + self.v_kernel_desc = JaxTensorDesc.from_shape( + (1, *sample_v.shape), sample_v.dtype, name="v_tensor" + ) + + def check_support(self) -> bool: + self.compute_capability = self._resolve_compute_capability( + self.target_compute_capability, + SUPPORTED_COMPUTE_CAPABILITIES, + "SdpafwdSm100D256", + ) + return True + + def __call__( + self, + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + ) -> TupleDict: + self.check_support() + if self.input_layout in FIXED_LAYOUTS: + if cum_seqlen_q_tensor is not None or cum_seqlen_k_tensor is not None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k must be omitted for fixed layout" + ) + inputs = (q_tensor, k_tensor, v_tensor) + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.v_kernel_desc, + ) + launch = self._launch_kernel + else: + if cum_seqlen_q_tensor is None or cum_seqlen_k_tensor is None: + raise ValueError( + "cum_seqlen_q and cum_seqlen_k are both required for THD layout" + ) + for value, desc in ( + (q_tensor, self.q_desc), + (k_tensor, self.k_desc), + (v_tensor, self.v_desc), + ): + self._check_tensor_signature(value, desc) + inputs = ( + jnp.reshape(q_tensor, self.q_kernel_desc.shape), + jnp.reshape(k_tensor, self.k_kernel_desc.shape), + jnp.reshape(v_tensor, self.v_kernel_desc.shape), + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + input_descs = ( + self.q_kernel_desc, + self.k_kernel_desc, + self.v_kernel_desc, + self.cum_q_desc, + self.cum_k_desc, + ) + launch = self._launch_varlen_kernel + + output, lse = self._call_kernel( + inputs, + launch=launch, + output_descs=(self.o_kernel_desc, self.lse_kernel_desc), + input_descs=input_descs, + compile_options=compile_options_for_target(self.compute_capability), + ) + if self.input_layout == _PACKED_LAYOUT: + output = jnp.reshape(output, self.o_desc.shape) + return TupleDict(o_tensor=output, lse_tensor=lse) + + def _launch_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + output: Any, + lse: Any, + ) -> None: + self._invoke_kernel(stream, q, k, v, output, lse, None, None) + + def _launch_varlen_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + cum_seqlen_q: Any, + cum_seqlen_k: Any, + output: Any, + lse: Any, + ) -> None: + self._invoke_kernel( + stream, + q, + k, + v, + output, + lse, + cum_seqlen_q, + cum_seqlen_k, + ) + + def _invoke_kernel( + self, + stream: Any, + q: Any, + k: Any, + v: Any, + output: Any, + lse: Any, + cum_seqlen_q: Any | None, + cum_seqlen_k: Any | None, + ) -> None: + from cutlass import Float32, Int32 + + from ..fmha_utils import MaskEnum + from .fmha_forward_sm100_d256 import BlackwellFusedMultiHeadAttentionForward + + mask_type = { + "residual": MaskEnum.RESIDUAL_MASK, + "window": MaskEnum.WINDOW_MASK_INFERENCE, + }[self.mask_kind] + kernel = BlackwellFusedMultiHeadAttentionForward( + qk_acc_dtype=Float32, + pv_acc_dtype=Float32, + mma_tiler=(128, 128, 256), + is_persistent=False, + mask_type=mask_type, + ) + problem_size = tuple( + Int32(value) + for value in ( + self.batch, + self.max_s_q, + self.max_s_k, + self.num_query_heads, + self.num_kv_heads, + 256, + ) + ) + left = None if self.window_size_left < 0 else Int32(self.window_size_left) + right = None if self.window_size_right < 0 else Int32(self.window_size_right) + kernel( + q, + k, + v, + output, + problem_size, + cum_seqlen_q, + cum_seqlen_k, + lse, + Float32(self.scale_softmax * math.log2(math.e)), + Float32(self.scale_softmax), + Float32(self.scale_output), + left, + right, + stream, + ) + + +@partial( + jax.jit, + static_argnames=( + "max_s_q", + "max_s_k", + "qk_acc_dtype", + "pv_acc_dtype", + "mma_tiler_mn", + "is_causal", + "window_size", + "scale_softmax", + "scale_output", + "layout", + "target_compute_capability", + ), +) +def sdpa_fwd_wrapper_sm100_d256( + q_tensor: Any, + k_tensor: Any, + v_tensor: Any, + cum_seqlen_q_tensor: Any | None = None, + cum_seqlen_k_tensor: Any | None = None, + max_s_q: int | None = None, + max_s_k: int | None = None, + qk_acc_dtype: Any = None, + pv_acc_dtype: Any = None, + mma_tiler_mn: tuple[int, int] = (128, 128), + is_causal: bool = False, + window_size: tuple[int, int] = (-1, -1), + scale_softmax: float | None = None, + scale_output: float = 1.0, + layout: str | None = None, + target_compute_capability: int | None = None, +) -> TupleDict: + """Compute fixed BHSD/BSHD or packed THD SDPA forward. + + Packed calls must provide both cumulative sequence-length arrays and static + ``max_s_q``/``max_s_k`` bounds. Fixed outputs follow ``layout``; LSE remains + ``(B, H, S)``. + """ + + return SdpafwdSm100D256( + q_tensor, + k_tensor, + v_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + max_s_q=max_s_q, + max_s_k=max_s_k, + qk_acc_dtype=qk_acc_dtype, + pv_acc_dtype=pv_acc_dtype, + mma_tiler_mn=mma_tiler_mn, + is_causal=is_causal, + window_size=window_size, + scale_softmax=scale_softmax, + scale_output=scale_output, + layout=layout, + target_compute_capability=target_compute_capability, + )( + q_tensor, + k_tensor, + v_tensor, + cum_seqlen_q_tensor, + cum_seqlen_k_tensor, + ) + + +__all__ = [ + "SUPPORTED_COMPUTE_CAPABILITIES", + "SdpafwdSm100D256", + "sdpa_fwd_wrapper_sm100_d256", +] diff --git a/python/cudnn/sdpa/jax_utils.py b/python/cudnn/sdpa/jax_utils.py new file mode 100644 index 000000000..64547de05 --- /dev/null +++ b/python/cudnn/sdpa/jax_utils.py @@ -0,0 +1,189 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Shared fixed-shape metadata and configuration for JAX SDPA.""" + +from __future__ import annotations + +import math +from typing import Any + +import jax.numpy as jnp + +from .. import data_type +from .._jax import JaxApiBase, JaxTensorDesc +from .._jax.datatypes import normalize_jax_dtype +from .._jax.layout import mode_from_layout, stride_order_to_public, to_public_axes + +FIXED_LAYOUTS = ("BHSD", "BSHD") +KERNEL_AXES = "BSHD" +KERNEL_STRIDE_ORDER = (3, 2, 1, 0) + + +def normalize_sdpa_layout(layout: str | None, rank: int) -> str: + """Normalize the public fixed/packed axis order for an SDPA call.""" + + if layout is None: + if rank == 4: + return "BHSD" + if rank == 3: + return "THD" + raise ValueError(f"SDPA layout cannot be inferred from rank {rank}") + if not isinstance(layout, str): + raise TypeError(f"layout must be a string or None, got {layout!r}") + normalized = "".join( + character for character in layout.upper() if character.isalpha() + ) + if normalized not in (*FIXED_LAYOUTS, "THD"): + raise ValueError(f"layout must be BHSD, BSHD, or THD, got {layout!r}") + expected_rank = 3 if normalized == "THD" else 4 + if rank != expected_rank: + raise ValueError( + f"{normalized} layout requires rank-{expected_rank} tensors, got rank {rank}" + ) + return normalized + + +def fixed_data_mode(layout: str) -> tuple[int, ...]: + if layout not in FIXED_LAYOUTS: + raise ValueError(f"fixed layout must be one of {FIXED_LAYOUTS}, got {layout!r}") + return mode_from_layout(layout, kernel_axes=KERNEL_AXES) + + +def describe_fixed_data( + sample: Any, + name: str, + *, + layout: str, + init_value: bool | int | float | None = None, +) -> JaxTensorDesc: + """Describe fixed data with the BSHD layout hard-coded by the kernel.""" + + mode = fixed_data_mode(layout) + + return JaxApiBase._to_tensor_desc( + sample, + name, + mode=mode, + public_stride_order=stride_order_to_public(KERNEL_STRIDE_ORDER, mode), + init_value=init_value, + ) + + +def make_fixed_output( + public_shape: tuple[int, ...], + dtype: Any, + name: str, + *, + layout: str, + init_value: bool | int | float | None = None, +) -> JaxTensorDesc: + mode = fixed_data_mode(layout) + return JaxTensorDesc.from_shape( + public_shape, + dtype, + name=name, + mode=mode, + public_stride_order=stride_order_to_public(KERNEL_STRIDE_ORDER, mode), + init_value=init_value, + ) + + +def require_fixed_qkv( + q_desc: JaxTensorDesc, + k_desc: JaxTensorDesc, + v_desc: JaxTensorDesc, +) -> tuple[int, int, int, int, int, int]: + for desc in (q_desc, k_desc, v_desc): + if desc.ndim != 4: + raise ValueError( + f"{desc.name} must have rank 4 (B, H, S, D), got {desc.shape}" + ) + if q_desc.cudnn_dtype not in (data_type.HALF, data_type.BFLOAT16): + raise ValueError( + f"SDPA requires float16 or bfloat16 inputs, got {q_desc.dtype}" + ) + if ( + k_desc.cudnn_dtype != q_desc.cudnn_dtype + or v_desc.cudnn_dtype != q_desc.cudnn_dtype + ): + raise ValueError("Q, K, and V must have the same dtype") + + logical_bhsd_mode = mode_from_layout("BHSD", kernel_axes=KERNEL_AXES) + q_shape = to_public_axes(q_desc.shape, logical_bhsd_mode) + k_shape = to_public_axes(k_desc.shape, logical_bhsd_mode) + v_shape = to_public_axes(v_desc.shape, logical_bhsd_mode) + batch, num_query_heads, seqlen_q, head_dim = q_shape + k_batch, num_kv_heads, seqlen_k, k_head_dim = k_shape + v_batch, num_value_heads, v_seqlen, value_dim = v_shape + dimensions = (batch, num_query_heads, num_kv_heads, seqlen_q, seqlen_k, head_dim) + if any(value <= 0 for value in dimensions): + raise ValueError(f"SDPA dimensions must be positive, got {dimensions}") + if (k_batch, v_batch) != (batch, batch): + raise ValueError("Q, K, and V batch dimensions must match") + if (num_value_heads, v_seqlen) != (num_kv_heads, seqlen_k): + raise ValueError("K and V head and sequence dimensions must match") + if (k_head_dim, value_dim) != (head_dim, head_dim): + raise ValueError("Q, K, and V head dimensions must match") + if head_dim != 256: + raise ValueError(f"head dimension must be 256, got {head_dim}") + if num_query_heads % num_kv_heads: + raise ValueError( + f"H_q ({num_query_heads}) must be divisible by H_kv ({num_kv_heads})" + ) + return dimensions + + +def require_float32_dtype(value: Any | None, name: str) -> Any: + dtype = normalize_jax_dtype(value, jnp.float32, name) + if dtype != jnp.dtype(jnp.float32): + raise ValueError(f"{name} must be float32, got {dtype}") + return dtype + + +def resolve_sdpa_config( + *, + seqlen_q: int, + seqlen_k: int, + tile_extent: int, + is_causal: bool, + window_size: tuple[int, int], + scale_softmax: float | None, +) -> tuple[float, int, int, str]: + if len(window_size) != 2: + raise ValueError(f"window_size must contain two values, got {window_size}") + window_size_left, window_size_right = window_size + if is_causal: + window_size_right = 0 + elif (window_size_left, window_size_right) != (-1, -1): + raise NotImplementedError( + f"window_size must be (-1, -1) for non-causal mode, got {window_size}" + ) + if window_size_left >= seqlen_k - 1: + raise ValueError( + f"window_size_left must be less than S_kv - 1 ({seqlen_k - 1}), got {window_size_left}" + ) + if window_size_right >= seqlen_q - 1: + raise ValueError( + f"window_size_right must be less than S_q - 1 ({seqlen_q - 1}), got {window_size_right}" + ) + + scale = ( + 1.0 / math.sqrt(256) + if scale_softmax is None or scale_softmax == 0.0 + else float(scale_softmax) + ) + mask_kind = "residual" if not is_causal and tile_extent % 128 else "window" + return scale, window_size_left, window_size_right, mask_kind + + +__all__ = [ + "FIXED_LAYOUTS", + "describe_fixed_data", + "fixed_data_mode", + "make_fixed_output", + "normalize_sdpa_layout", + "require_fixed_qkv", + "require_float32_dtype", + "resolve_sdpa_config", +] diff --git a/test/python/fe_api/dsa/dsa_reference.py b/test/python/fe_api/dsa/dsa_reference.py index fe8dcb334..0a3ceba8c 100644 --- a/test/python/fe_api/dsa/dsa_reference.py +++ b/test/python/fe_api/dsa/dsa_reference.py @@ -66,7 +66,10 @@ def ref_sparse_attention_forward( q_f = q.to(torch.float32) k_f = kv.to(torch.float32) - v_f = kv.to(torch.float32) + # MLA uses a 512-wide value projection for the 576-wide Q/K variant. + # K and V share storage, so V is the leading D_v slice of KV. + head_dim_v = 512 if d == 576 else d + v_f = kv[..., :head_dim_v].to(torch.float32) mask = _make_topk_mask(topk_idxs, topk_length, t_kv) # (T, T_kv) mask = mask.unsqueeze(1).expand(t, h, t_kv) # (T, H, T_kv) @@ -76,8 +79,11 @@ def ref_sparse_attention_forward( scores = scores.masked_fill(~mask, float("-inf")) lse = torch.logsumexp(scores, dim=-1) # (T, H), excludes sink - sink = attn_sink.view(1, h) - lse_with_sink = torch.logaddexp(lse, sink) + sink = attn_sink.view(1, h, 1).expand(t, h, 1) + # Include the finite sink directly in the reduction. Composing + # logaddexp(logsumexp(all -inf), sink) gives NaN score gradients in + # PyTorch even though the sink makes the mathematical denominator finite. + lse_with_sink = torch.logsumexp(torch.cat((scores, sink), dim=-1), dim=-1) weights = torch.exp(scores - lse_with_sink.unsqueeze(-1)) out = torch.einsum("thk,kd->thd", weights, v_f) diff --git a/test/python/fe_api/dsa/dsa_utils.py b/test/python/fe_api/dsa/dsa_utils.py index d62375ca0..84bd7d741 100644 --- a/test/python/fe_api/dsa/dsa_utils.py +++ b/test/python/fe_api/dsa/dsa_utils.py @@ -17,7 +17,7 @@ ] DSA_SPARSE_ATTENTION_BACKWARD_PARAM_MARKS = [ - pytest.mark.parametrize("head_dim", [512]), + pytest.mark.parametrize("head_dim", [512, 576]), pytest.mark.parametrize("head_dim_v", [512]), pytest.mark.parametrize("num_heads", [64]), pytest.mark.parametrize("topk", [512]), diff --git a/test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py b/test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py index c6d10c061..f1ddf70ed 100644 --- a/test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py +++ b/test/python/fe_api/dsa/test_DSA_sparse_attention_backward.py @@ -40,10 +40,44 @@ def _allocate(cfg, has_topk_length: bool): topk_length = None if has_topk_length: topk_length = torch.randint(1, topk_k + 1, (total_s_q,), dtype=torch.int32, device=device) + # Empty compactified rows are valid: attention falls back entirely to + # the sink and all Q/KV gradients for that row are zero. + topk_length[0] = 0 return q, kv, attn_sink, topk_idxs, topk_length +@pytest.mark.L0 +@pytest.mark.parametrize("head_dim", [512, 576]) +def test_sparse_attention_reference_supports_empty_rows(head_dim): + """Keep the D=576 value width and sink-only rows numerically well-defined.""" + + torch.manual_seed(0) + total_q, total_kv, num_heads, topk = 2, 3, 2, 2 + q = torch.randn(total_q, num_heads, head_dim, dtype=torch.float32, requires_grad=True) + kv = torch.randn(total_kv, head_dim, dtype=torch.float32, requires_grad=True) + attn_sink = torch.randn(num_heads, dtype=torch.float32, requires_grad=True) + topk_idxs = torch.tensor([[0, 1], [1, 2]], dtype=torch.int32) + topk_length = torch.tensor([0, topk], dtype=torch.int32) + + out, lse = ref_sparse_attention_forward( + q, + kv, + attn_sink, + topk_idxs, + topk_length=topk_length, + ) + out.sum().backward() + + head_dim_v = 512 if head_dim == 576 else head_dim + assert out.shape == (total_q, num_heads, head_dim_v) + assert torch.isneginf(lse[0]).all() + assert torch.isfinite(q.grad).all() + assert torch.isfinite(kv.grad).all() + assert torch.isfinite(attn_sink.grad).all() + assert torch.count_nonzero(q.grad[0]) == 0 + + @pytest.mark.L0 @torch_fork_set_rng(seed=0) @with_dsa_sparse_attention_backward_params @@ -93,21 +127,18 @@ def test_DSA_sparse_attention_backward_wrapper( ) dout = torch.randn_like(out) - try: - result = DSA.sparse_attention_backward_wrapper( - q, - kv, - out, - dout, - lse, - attn_sink, - topk_idxs, - softmax_scale=softmax_scale, - topk_length=topk_length, - stream=stream, - ) - except (ValueError, NotImplementedError, RuntimeError) as e: - pytest.skip(f"Unsupported testcase: {e}") + result = DSA.sparse_attention_backward_wrapper( + q, + kv, + out, + dout, + lse, + attn_sink, + topk_idxs, + softmax_scale=softmax_scale, + topk_length=topk_length, + stream=stream, + ) dq, dkv, d_sink = result["dq"], result["dkv"], result["d_sink"] diff --git a/test/python/fe_api/test_dense_gemm_helpers.py b/test/python/fe_api/test_dense_gemm_helpers.py new file mode 100644 index 000000000..81d66bd9c --- /dev/null +++ b/test/python/fe_api/test_dense_gemm_helpers.py @@ -0,0 +1,198 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for framework-neutral dense GEMM helpers.""" + +import ast +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_PACKAGE = "cudnn_dense_gemm_test" + + +class _DataType(Enum): + NOT_SET = auto() + FLOAT = auto() + HALF = auto() + BFLOAT16 = auto() + FP8_E4M3 = auto() + FP8_E5M2 = auto() + FP8_E8M0 = auto() + FP4_E2M1 = auto() + + +class DenseGemmHelperTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + package = types.ModuleType(_PACKAGE) + package.__path__ = [str(_CUDNN_ROOT)] + package.__package__ = _PACKAGE + package.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + package.data_type = _DataType + sys.modules[_PACKAGE] = package + + cls.tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + cls.helpers = importlib.import_module(f"{_PACKAGE}.gemm.helpers") + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def desc(self, shape, dtype=_DataType.BFLOAT16, *, stride_order=None, name=""): + return self.tensor_module.make_compact_tensor_desc( + dtype=dtype, + shape=tuple(shape), + stride_order=stride_order, + name=name, + ) + + def test_gemm_inputs_resolve_mnkl(self): + a = self.desc((17, 31, 3), name="A") + b = self.desc((19, 31, 3), name="B") + + self.assertEqual(self.helpers.require_gemm_inputs(a, b), (17, 19, 31, 3)) + + def test_gemm_inputs_reject_invalid_rank_extents_and_matching(self): + cases = ( + (self.desc((17, 31)), self.desc((19, 31, 3)), "A must have rank 3"), + (self.desc((17, 31, 3)), self.desc((19, 31)), "B must have rank 3"), + (self.desc((0, 31, 3)), self.desc((19, 31, 3)), "M must be positive"), + (self.desc((17, 31, 3)), self.desc((0, 31, 3)), "N must be positive"), + (self.desc((17, 0, 3)), self.desc((19, 0, 3)), "K must be positive"), + (self.desc((17, 31, 0)), self.desc((19, 31, 0)), "L must be positive"), + (self.desc((17, 31, 3)), self.desc((19, 29, 3)), "B shape mismatch"), + (self.desc((17, 31, 3)), self.desc((19, 31, 4)), "B shape mismatch"), + ) + + for a, b, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self.helpers.require_gemm_inputs(a, b) + + def test_exact_output_shape_uses_logical_descriptor_shape(self): + output = self.desc((17, 19, 3), name="AB12") + self.helpers.require_tensor_shape(output, (17, 19, 3)) + + with self.assertRaisesRegex(ValueError, "AB12 must have shape"): + self.helpers.require_tensor_shape(output, (17, 18, 3)) + + def test_compact_major_recognizes_mode_zero_and_mode_one(self): + m_major = self.desc((8, 16, 2), stride_order=(0, 1, 2), name="A") + k_major = self.desc((8, 16, 2), stride_order=(1, 0, 2), name="A") + + self.assertEqual(self.helpers.require_compact_major(m_major, "m", "k"), "m") + self.assertEqual(self.helpers.require_compact_major(k_major, "m", "k"), "k") + + def test_compact_major_rejects_batch_major_and_noncompact_layouts(self): + batch_major = self.desc((8, 16, 2), stride_order=(2, 1, 0), name="A") + noncompact = self.tensor_module.TensorDesc( + dtype=_DataType.BFLOAT16, + shape=(8, 16, 2), + stride=(1, 16, 256), + stride_order=(0, 1, 2), + name="A", + ) + + for tensor in (batch_major, noncompact): + with self.subTest(stride=tensor.stride): + with self.assertRaisesRegex(ValueError, "compact m-major or k-major"): + self.helpers.require_compact_major(tensor, "m", "k") + + def test_data_type_widths_and_16_byte_alignment(self): + cases = ( + (_DataType.FLOAT, 32, 4), + (_DataType.HALF, 16, 8), + (_DataType.BFLOAT16, 16, 8), + (_DataType.FP8_E4M3, 8, 16), + (_DataType.FP8_E5M2, 8, 16), + (_DataType.FP4_E2M1, 4, 32), + ) + + for dtype, bits, extent in cases: + with self.subTest(dtype=dtype): + tensor = self.desc((extent, 7, 2), dtype, stride_order=(0, 1, 2), name="A") + self.assertEqual(self.helpers.data_type_bits(dtype), bits) + self.helpers.require_16_byte_alignment(tensor) + + misaligned = self.desc((extent - 1, 7, 2), dtype, stride_order=(0, 1, 2), name="A") + with self.assertRaisesRegex(ValueError, "16-byte alignment"): + self.helpers.require_16_byte_alignment(misaligned) + + with self.assertRaisesRegex(ValueError, "Unsupported cuDNN data type"): + self.helpers.data_type_bits(_DataType.NOT_SET) + + def test_16_byte_alignment_uses_the_selected_major_extent(self): + tensor = self.desc((7, 8, 2), stride_order=(1, 0, 2), name="A") + self.helpers.require_16_byte_alignment(tensor) + + def test_mma_tiler_validation_is_configurable(self): + self.assertEqual(self.helpers.require_mma_tiler((128, 32)), (128, 32)) + self.assertEqual(self.helpers.require_mma_tiler((256, 256)), (256, 256)) + self.assertEqual( + self.helpers.require_mma_tiler((64, 96), allowed_m=(64,), allowed_n=(96,)), + (64, 96), + ) + + cases = ( + ((64, 128), "mma_tiler_mn\\[0\\]"), + ((128, 48), "mma_tiler_mn\\[1\\]"), + ((128,), "must contain two integers"), + ) + for tiler, message in cases: + with self.subTest(tiler=tiler): + with self.assertRaisesRegex(ValueError, message): + self.helpers.require_mma_tiler(tiler) + + with self.assertRaisesRegex(TypeError, "must contain integers"): + self.helpers.require_mma_tiler((True, 128)) + + def test_cluster_shape_validation_handles_cta_groups(self): + self.assertEqual(self.helpers.require_cluster_shape((1, 1)), (1, 1)) + self.assertEqual( + self.helpers.require_cluster_shape((2, 4), cta_group_size=2), + (2, 4), + ) + + cases = ( + ((3, 1), {}, "positive powers of two"), + ((4, 8), {}, "must not exceed 16"), + ((1, 2), {"cta_group_size": 2}, "divisible by CTA group size"), + ) + for shape, kwargs, message in cases: + with self.subTest(shape=shape): + with self.assertRaisesRegex(ValueError, message): + self.helpers.require_cluster_shape(shape, **kwargs) + + def test_module_has_no_framework_or_kernel_dependencies(self): + tree = ast.parse((_CUDNN_ROOT / "gemm" / "helpers.py").read_text()) + imports = [] + for node in tree.body: + if isinstance(node, ast.Import): + imports.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + imports.append("." * node.level + (node.module or "")) + + self.assertFalse( + any(imported == dependency or imported.startswith(f"{dependency}.") for imported in imports for dependency in ("torch", "jax", "cutlass", "cuda")), + imports, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_dsa_indexer_backward_op.py b/test/python/fe_api/test_dsa_indexer_backward_op.py new file mode 100644 index 000000000..17e8fba78 --- /dev/null +++ b/test/python/fe_api/test_dsa_indexer_backward_op.py @@ -0,0 +1,231 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for DSA indexer-backward metadata.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "deepseek_sparse_attention" / "indexer_backward" +_PACKAGE = "cudnn_dsa_indexer_backward_op_test" + + +class _DataType(Enum): + NOT_SET = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + + +class IndexerBackwardOpContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + dsa_name = f"{_PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_CUDNN_ROOT / "deepseek_sparse_attention")] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + sys.modules[dsa_name] = dsa + + operation_name = f"{dsa_name}.indexer_backward" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + try: + cls.tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}.common.op") + cls.op_module = importlib.import_module(f"{operation_name}.op") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def _desc(self, shape, dtype=_DataType.BFLOAT16, *, name=""): + return self.tensor_module.make_compact_tensor_desc( + dtype=dtype, + shape=tuple(shape), + name=name, + ) + + def _sparse(self, **overrides): + batch, seqlen_q, seqlen_k, heads, dim, topk = 2, 64, 256, 64, 128, 128 + arguments = { + "index_q": self._desc((batch, seqlen_q, heads, dim), name="index_q"), + "weights": self._desc((batch, seqlen_q, heads), name="weights"), + "index_k": self._desc((batch, seqlen_k, dim), name="index_k"), + "d_index_q": self._desc((batch, seqlen_q, heads, dim), name="d_index_q"), + "d_weights": self._desc((batch, seqlen_q, heads), name="d_weights"), + "d_index_k": self._desc((batch, seqlen_k, dim), name="d_index_k"), + "attn_score": self._desc((batch, seqlen_q, topk), _DataType.FLOAT, name="attn_score"), + "index_score": self._desc((batch, seqlen_q, topk), _DataType.FLOAT, name="index_score"), + "topk_indices": self._desc((batch, seqlen_q, topk), _DataType.INT32, name="topk_indices"), + } + arguments.update(overrides) + return self.op_module.IndexerBackwardOp(**arguments) + + def _sparse_thd(self, **overrides): + total_q, total_k, heads, dim, topk = 96, 384, 64, 128, 128 + arguments = { + "index_q": self._desc((total_q, heads, dim), name="index_q"), + "weights": self._desc((total_q, heads), name="weights"), + "index_k": self._desc((total_k, dim), name="index_k"), + "d_index_q": self._desc((total_q, heads, dim), name="d_index_q"), + "d_weights": self._desc((total_q, heads), name="d_weights"), + "d_index_k": self._desc((total_k, dim), name="d_index_k"), + "attn_score": self._desc((total_q, topk), _DataType.FLOAT, name="attn_score"), + "index_score": self._desc((total_q, topk), _DataType.FLOAT, name="index_score"), + "topk_indices": self._desc((total_q, topk), _DataType.INT32, name="topk_indices"), + "topk_indices_global": True, + } + arguments.update(overrides) + return self.op_module.IndexerBackwardOp(**arguments) + + def _dense_bshd(self, **overrides): + batch, seqlen_q, seqlen_k, heads, dim = 2, 64, 256, 64, 128 + arguments = { + "index_q": self._desc((batch, seqlen_q, heads, dim), name="index_q"), + "weights": self._desc((batch, seqlen_q, heads), name="weights"), + "index_k": self._desc((batch, seqlen_k, dim), name="index_k"), + "d_index_q": self._desc((batch, seqlen_q, heads, dim), name="d_index_q"), + "d_weights": self._desc((batch, seqlen_q, heads), name="d_weights"), + "d_index_k": self._desc((batch, seqlen_k, dim), name="d_index_k"), + "attn_score": self._desc((batch, seqlen_q, seqlen_k), _DataType.FLOAT, name="attn_score"), + "attn_l1norm": self._desc((batch, seqlen_q), _DataType.FLOAT, name="attn_l1norm"), + "index_score": self._desc((batch, seqlen_q, seqlen_k), _DataType.FLOAT, name="index_score"), + "index_lse": self._desc((batch, seqlen_q), _DataType.FLOAT, name="index_lse"), + } + arguments.update(overrides) + return self.op_module.DenseIndexerBackwardOp(**arguments) + + def _dense_thd(self, **overrides): + batch, total_q, total_k, max_k, heads, dim = 2, 96, 384, 256, 64, 128 + arguments = { + "index_q": self._desc((total_q, heads, dim), name="index_q"), + "weights": self._desc((total_q, heads), name="weights"), + "index_k": self._desc((total_k, dim), name="index_k"), + "d_index_q": self._desc((total_q, heads, dim), name="d_index_q"), + "d_weights": self._desc((total_q, heads), name="d_weights"), + "d_index_k": self._desc((total_k, dim), name="d_index_k"), + "attn_score": self._desc((total_q, max_k), _DataType.FLOAT, name="attn_score"), + "attn_l1norm": self._desc((total_q,), _DataType.FLOAT, name="attn_l1norm"), + "index_score": self._desc((total_q, max_k), _DataType.FLOAT, name="index_score"), + "index_lse": self._desc((total_q,), _DataType.FLOAT, name="index_lse"), + "cu_seqlens_q": self._desc((batch + 1,), _DataType.INT32, name="cu_seqlens_q"), + "cu_seqlens_k": self._desc((batch + 1,), _DataType.INT32, name="cu_seqlens_k"), + "q_causal_offsets": self._desc((batch,), _DataType.INT32, name="q_causal_offsets"), + "max_seqlen_q": 64, + "max_seqlen_k": max_k, + } + arguments.update(overrides) + return self.op_module.DenseIndexerBackwardOp(**arguments) + + def test_sparse_complete_signature(self): + operation = self._sparse() + self.assertIsInstance(operation, self.base_module.Op) + self.assertTrue(operation.check_support()) + self.assertFalse(operation.is_thd) + self.assertEqual( + (operation.batch, operation.seqlen_q, operation.seqlen_k, operation.heads, operation.head_dim, operation.topk), + (2, 64, 256, 64, 128, 128), + ) + + def test_sparse_packed_thd_complete_signature(self): + operation = self._sparse_thd() + self.assertTrue(operation.check_support()) + self.assertTrue(operation.is_thd) + self.assertEqual( + (operation.batch, operation.seqlen_q, operation.seqlen_k, operation.heads, operation.head_dim, operation.topk), + (1, 96, 384, 64, 128, 128), + ) + + def test_sparse_packed_thd_requires_global_indices(self): + with self.assertRaisesRegex(ValueError, "requires topk_indices_global=True"): + self._sparse_thd(topk_indices_global=False).check_support() + + def test_sparse_rejects_invalid_output_score_and_configuration(self): + cases = ( + ({"d_index_q": self._desc((2, 64, 32, 128))}, "d_index_q must have shape"), + ({"index_score": self._desc((2, 64, 64), _DataType.FLOAT)}, "index_score must have shape"), + ({"topk_indices": self._desc((2, 64, 64), _DataType.INT32)}, r"topk \(64\) must be divisible"), + ({"block_i": 64}, "block_i must be 128"), + ({"sm_scale": float("inf")}, "sm_scale must be finite"), + ) + for overrides, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self._sparse(**overrides).check_support() + + def test_scale_preserves_values_supported_by_forward(self): + for scale in (0.0, -0.5): + with self.subTest(scale=scale): + operation = self._sparse(sm_scale=scale) + self.assertTrue(operation.check_support()) + self.assertEqual(operation.sm_scale, scale) + + def test_dense_bshd_complete_signature(self): + operation = self._dense_bshd(q_causal_offsets=self._desc((2,), _DataType.INT32)) + self.assertTrue(operation.check_support()) + self.assertFalse(operation.is_thd) + self.assertEqual( + (operation.batch, operation.normalization_tokens, operation.total_k, operation.max_seqlen_q, operation.max_seqlen_k), + (2, 128, 512, 64, 256), + ) + + def test_dense_thd_complete_signature(self): + operation = self._dense_thd() + self.assertTrue(operation.check_support()) + self.assertTrue(operation.is_thd) + self.assertEqual( + (operation.batch, operation.normalization_tokens, operation.total_k, operation.max_seqlen_q, operation.max_seqlen_k), + (2, 96, 384, 64, 256), + ) + + def test_dense_rejects_partial_or_invalid_thd_signature(self): + cases = ( + ({"cu_seqlens_k": None}, "must be provided together"), + ({"max_seqlen_q": None}, "requires max_seqlen_q and max_seqlen_k"), + ({"q_causal_offsets": self._desc((3,), _DataType.INT32)}, "q_causal_offsets must have shape"), + ({"attn_score": self._desc((96, 128), _DataType.FLOAT)}, "attn_score must have shape"), + ) + for overrides, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self._dense_thd(**overrides).check_support() + + def test_operation_imports_no_framework_or_kernel(self): + source = (_OPERATION_ROOT / "op.py").read_text() + for forbidden in ("import torch", "import jax", "import cutlass", "import cuda"): + self.assertNotIn(forbidden, source) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_dsa_indexer_ops.py b/test/python/fe_api/test_dsa_indexer_ops.py new file mode 100644 index 000000000..b38ce52b1 --- /dev/null +++ b/test/python/fe_api/test_dsa_indexer_ops.py @@ -0,0 +1,603 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for the DSA indexer logical operations.""" + +import ast +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_DSA_ROOT = _CUDNN_ROOT / "deepseek_sparse_attention" +_PACKAGE = "cudnn_dsa_indexer_op_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + INT64 = auto() + + +class _DType: + def __init__(self, name: str, itemsize: int): + self.name = name + self.itemsize = itemsize + + def __repr__(self): + return self.name + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + @property + def ndim(self): + return len(self.shape) + + def __getitem__(self, key): + if isinstance(key, tuple) and key[:-1] == (Ellipsis,) and isinstance(key[-1], slice): + return _Array((*self.shape[:-1], key[-1].stop), self.dtype) + if isinstance(key, tuple) and len(key) == self.ndim and all(isinstance(value, slice) for value in key): + shape = [] + for extent, value in zip(self.shape, key): + start = 0 if value.start is None else value.start + stop = extent if value.stop is None else min(value.stop, extent) + step = 1 if value.step is None else value.step + shape.append(max(0, (stop - start + step - 1) // step)) + return _Array(shape, self.dtype) + raise AssertionError(f"Unexpected array slice {key!r}") + + +class _TensorSpec: + def __init__(self, *, layout=None, mode=None, divisibility=None, **_kwargs): + self.layout = layout + self.mode = mode + self.divisibility = divisibility + + +class DsaIndexerOpContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + dsa_name = f"{_PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_DSA_ROOT)] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + sys.modules[dsa_name] = dsa + + for operation in ("indexer_forward", "indexer_top_k"): + name = f"{dsa_name}.{operation}" + module = types.ModuleType(name) + module.__path__ = [str(_DSA_ROOT / operation)] + module.__package__ = name + module.__spec__ = ModuleSpec(name, loader=None, is_package=True) + sys.modules[name] = module + + cls.tensor = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + importlib.import_module(f"{_PACKAGE}.common.op") + cls.forward = importlib.import_module(f"{dsa_name}.indexer_forward.op") + cls.topk = importlib.import_module(f"{dsa_name}.indexer_top_k.op") + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def desc(self, shape, dtype, name="", stride_order=None): + return self.tensor.make_compact_tensor_desc( + dtype=dtype, + shape=tuple(shape), + stride_order=stride_order, + name=name, + ) + + def test_fixed_forward_complete_signature(self): + operation = self.forward.IndexerForwardOp( + q=self.desc((2, 16, 32, 128), _DataType.BFLOAT16), + k=self.desc((2, 17, 1, 128), _DataType.BFLOAT16), + weight=self.desc((2, 16, 32), _DataType.BFLOAT16), + output=self.desc((2, 16, 20), _DataType.FLOAT), + q_causal_offsets=self.desc((2,), _DataType.INT32), + ratio=4, + target_compute_capability=100, + ) + self.assertTrue(operation.check_support()) + self.assertFalse(operation.is_varlen) + self.assertEqual((operation.s_q, operation.s_k), (16, 17)) + self.assertEqual(operation.qhead_per_kv_head, 32) + + def test_varlen_forward_complete_signature(self): + operation = self.forward.IndexerForwardOp( + q=self.desc((31, 64, 128), _DataType.BFLOAT16), + k=self.desc((47, 1, 128), _DataType.BFLOAT16), + weight=self.desc((31, 64), _DataType.BFLOAT16), + output=self.desc((31, 20), _DataType.FLOAT), + cu_seqlens_q=self.desc((3,), _DataType.INT32), + cu_seqlens_k=self.desc((3,), _DataType.INT32), + q_causal_offsets=self.desc((2,), _DataType.INT32), + max_seqlen_q=16, + max_seqlen_k=17, + target_compute_capability=90, + ) + self.assertTrue(operation.check_support()) + self.assertTrue(operation.is_varlen) + self.assertEqual(operation.batch_size, 2) + + def test_forward_rejects_invalid_output_and_optional_signature(self): + with self.assertRaisesRegex(ValueError, "leading shape"): + self.forward.IndexerForwardOp( + q=self.desc((1, 8, 32, 128), _DataType.BFLOAT16), + k=self.desc((1, 5, 1, 128), _DataType.BFLOAT16), + weight=self.desc((1, 8, 32), _DataType.BFLOAT16), + output=self.desc((1, 8, 5), _DataType.FLOAT), + ).check_support() + + oversized_output = self.forward.IndexerForwardOp( + q=self.desc((1, 8, 32, 128), _DataType.BFLOAT16), + k=self.desc((1, 5, 1, 128), _DataType.BFLOAT16), + weight=self.desc((1, 8, 32), _DataType.BFLOAT16), + output=self.desc((1, 8, 12), _DataType.FLOAT), + ) + self.assertTrue(oversized_output.check_support()) + + with self.assertRaisesRegex(ValueError, "requires both"): + self.forward.IndexerForwardOp( + q=self.desc((8, 32, 128), _DataType.BFLOAT16), + k=self.desc((8, 1, 128), _DataType.BFLOAT16), + weight=self.desc((8, 32), _DataType.BFLOAT16), + output=self.desc((8, 8), _DataType.FLOAT), + cu_seqlens_q=self.desc((2,), _DataType.INT32), + max_seqlen_q=8, + max_seqlen_k=8, + ).check_support() + + def test_forward_requires_the_kernel_vector_axis_to_be_contiguous(self): + operation = self.forward.IndexerForwardOp( + q=self.desc( + (1, 8, 32, 128), + _DataType.BFLOAT16, + stride_order=(2, 3, 1, 0), + ), + k=self.desc((1, 5, 1, 128), _DataType.BFLOAT16), + weight=self.desc((1, 8, 32), _DataType.BFLOAT16), + output=self.desc((1, 8, 8), _DataType.FLOAT), + ratio=1, + ) + with self.assertRaisesRegex(ValueError, "canonical D axis contiguous"): + operation.check_support() + + def test_sm100_forward_requires_a_complete_packed_query_tile(self): + operation = self.forward.IndexerForwardOp( + q=self.desc((1, 4, 32, 128), _DataType.BFLOAT16), + k=self.desc((1, 5, 1, 128), _DataType.BFLOAT16), + weight=self.desc((1, 4, 32), _DataType.BFLOAT16), + output=self.desc((1, 4, 8), _DataType.FLOAT), + ratio=1, + target_compute_capability=100, + ) + with self.assertRaisesRegex(ValueError, "complete packed query tile"): + operation.check_support() + + operation.target_compute_capability = 90 + self.assertTrue(operation.check_support()) + + def make_topk(self, *, return_val=True, **overrides): + input_values = self.desc((4, 64), _DataType.FLOAT) + arguments = { + "input_values": input_values, + "seq_lens": self.desc((4,), _DataType.INT32), + "output_indices": self.desc((4, 8), _DataType.INT32), + "output_values": self.desc((4, 8), _DataType.FLOAT) if return_val else None, + "workspace": self.desc((4, 2, 64), _DataType.INT32), + "top_k": 8, + "return_val": return_val, + } + arguments.update(overrides) + return self.topk.IndexerTopKOp(**arguments) + + def test_topk_complete_signature_with_and_without_values(self): + for return_val in (True, False): + with self.subTest(return_val=return_val): + operation = self.make_topk(return_val=return_val) + self.assertTrue(operation.check_support()) + self.assertEqual(operation.buffer_count, 2) + self.assertEqual(operation.max_num_cols, 64) + + def test_topk_rejects_invalid_rows_workspace_and_copy_width(self): + cases = ( + ({"seq_lens": self.desc((2,), _DataType.INT32)}, "must equal"), + ({"workspace": self.desc((4, 1, 64), _DataType.INT32)}, "workspace must"), + ({"num_copy_bits": 24}, "power-of-two"), + ({"top_k": 65}, "num_cols=64"), + ) + for overrides, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self.make_topk(**overrides).check_support() + + +class DsaIndexerImportContractTest(unittest.TestCase): + @staticmethod + def top_level_imports(path: Path) -> tuple[str, ...]: + tree = ast.parse(path.read_text(), filename=str(path)) + imports = [] + for node in tree.body: + if isinstance(node, ast.Import): + imports.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + imports.append("." * node.level + (node.module or "")) + return tuple(imports) + + def test_operation_packages_and_ops_are_framework_lazy(self): + for operation in ("indexer_forward", "indexer_top_k"): + root = _DSA_ROOT / operation + for filename in ("__init__.py", "op.py"): + imports = self.top_level_imports(root / filename) + self.assertFalse( + any(name == framework or name.startswith(f"{framework}.") for name in imports for framework in ("torch", "jax", "cutlass", "cuda")), + f"{operation}/{filename}: {imports}", + ) + + def test_jax_reachable_kernels_do_not_require_torch(self): + files = ( + _DSA_ROOT / "indexer_forward" / "jax.py", + _DSA_ROOT / "indexer_forward" / "indexer_fwd_sm90.py", + _DSA_ROOT / "indexer_forward" / "indexer_fwd_sm100.py", + _DSA_ROOT / "indexer_top_k" / "jax.py", + _DSA_ROOT / "indexer_top_k" / "indexer_top_k_decode_varlen.py", + _DSA_ROOT / "indexer_top_k" / "indexer_top_k_varlen_util.py", + _DSA_ROOT / "indexer_top_k" / "local_to_global_dsl.py", + _DSA_ROOT / "indexer_top_k" / "compactify.py", + ) + for path in files: + imports = self.top_level_imports(path) + self.assertNotIn("torch", imports, str(path)) + + def test_torch_indexer_forward_passes_all_optional_kernel_arguments(self): + path = _DSA_ROOT / "indexer_forward" / "api.py" + tree = ast.parse(path.read_text(), filename=str(path)) + compile_method = next(node for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) and node.name == "compile" and node.lineno < 220) + compile_call = next( + node for node in ast.walk(compile_method) if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "compile" + ) + self.assertEqual(sum(isinstance(argument, ast.Constant) and argument.value is None for argument in compile_call.args), 3) + + tensor_api = next(node for node in ast.walk(compile_method) if isinstance(node, ast.FunctionDef) and node.name == "tensor_api") + launch_call = next(node for node in ast.walk(tensor_api) if isinstance(node, ast.Call)) + self.assertEqual(sum(isinstance(argument, ast.Constant) and argument.value is None for argument in launch_call.args), 3) + + +class DsaIndexerJaxAdapterContractTest(unittest.TestCase): + _PACKAGE = "cudnn_dsa_indexer_jax_test" + + @classmethod + def setUpClass(cls) -> None: + cls.float16 = _DType("float16", 2) + cls.bfloat16 = _DType("bfloat16", 2) + cls.float32 = _DType("float32", 4) + cls.int32 = _DType("int32", 4) + cls.int64 = _DType("int64", 8) + + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.float16 = cls.float16 + fake_jnp.bfloat16 = cls.bfloat16 + fake_jnp.float32 = cls.float32 + fake_jnp.int32 = cls.int32 + fake_jnp.int64 = cls.int64 + fake_jnp.dtype = lambda value: value.dtype if hasattr(value, "dtype") else value + fake_jnp.full = lambda shape, _value, dtype: _Array(shape, dtype) + fake_jnp.reshape = lambda value, shape: _Array(shape, value.dtype) + + fake_jax = types.ModuleType("jax") + fake_jax.__path__ = [] + fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + fake_jax.numpy = fake_jnp + fake_jax.ShapeDtypeStruct = _Array + + def fake_jit(function=None, **options): + def decorate(resolved): + resolved._test_static_argnames = tuple(options.get("static_argnames", ())) + return resolved + + return decorate if function is None else decorate(function) + + fake_jax.jit = fake_jit + fake_jax.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, + ) + + fake_cutlass = types.ModuleType("cutlass") + fake_cutlass.__path__ = [] + fake_cutlass_jax = types.ModuleType("cutlass.jax") + fake_cutlass_jax.TensorSpec = _TensorSpec + fake_cutlass_jax.cutlass_call = None + fake_cutlass.jax = fake_cutlass_jax + + cls.fake_modules = { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "cutlass": fake_cutlass, + "cutlass.jax": fake_cutlass_jax, + } + cls.module_patch = mock.patch.dict(sys.modules, cls.fake_modules) + cls.module_patch.start() + + root = types.ModuleType(cls._PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = cls._PACKAGE + root.__spec__ = ModuleSpec(cls._PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[cls._PACKAGE] = root + + dsa_name = f"{cls._PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_DSA_ROOT)] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + sys.modules[dsa_name] = dsa + for operation in ("indexer_forward", "indexer_top_k"): + name = f"{dsa_name}.{operation}" + package = types.ModuleType(name) + package.__path__ = [str(_DSA_ROOT / operation)] + package.__package__ = name + package.__spec__ = ModuleSpec(name, loader=None, is_package=True) + sys.modules[name] = package + + cls.jax_base = importlib.import_module(f"{cls._PACKAGE}._jax") + + def resolve_compute_capability( + target_compute_capability, + supported_compute_capabilities, + operation_name, + ): + del operation_name + resolved = 100 if target_compute_capability is None else target_compute_capability + if resolved not in supported_compute_capabilities: + raise ValueError(f"unsupported synthetic target SM{resolved}") + return resolved + + cls.target_resolution_patch = mock.patch.object( + cls.jax_base.JaxApiBase, + "_resolve_compute_capability", + staticmethod(resolve_compute_capability), + ) + cls.target_resolution_patch.start() + cls.forward = importlib.import_module(f"{dsa_name}.indexer_forward.jax") + cls.topk = importlib.import_module(f"{dsa_name}.indexer_top_k.jax") + + @classmethod + def tearDownClass(cls) -> None: + cls.target_resolution_patch.stop() + for name in tuple(sys.modules): + if name == cls._PACKAGE or name.startswith(f"{cls._PACKAGE}."): + sys.modules.pop(name, None) + cls.module_patch.stop() + + @staticmethod + def fake_call(captured): + def call(_inputs, *, output_descs, **options): + input_descs = options.get("input_descs") + if input_descs is not None: + derived_input_specs = tuple( + DsaIndexerJaxAdapterContractTest.jax_base.JaxApiBase._to_tensor_spec( + desc + ) + for desc in input_descs + ) + supplied_input_specs = options.get("input_spec") + options["input_spec"] = ( + derived_input_specs + if supplied_input_specs is None + else tuple( + derived if supplied is None else supplied + for derived, supplied in zip( + derived_input_specs, + supplied_input_specs, + ) + ) + ) + + derived_output_specs = tuple( + DsaIndexerJaxAdapterContractTest.jax_base.JaxApiBase._to_tensor_spec( + desc + ) + for desc in output_descs + ) + supplied_output_specs = options.get("output_spec") + options["output_spec"] = ( + derived_output_specs + if supplied_output_specs is None + else tuple( + derived if supplied is None else supplied + for derived, supplied in zip( + derived_output_specs, + supplied_output_specs, + ) + ) + ) + captured.update(output_descs=output_descs, **options) + output_specs = options["output_spec"] + + def public_shape(desc, spec): + if spec is None or spec.mode is None: + return desc.shape + shape = [None] * desc.ndim + for canonical_axis, public_axis in enumerate(spec.mode): + shape[public_axis] = desc.shape[canonical_axis] + return tuple(shape) + + return tuple(_Array(public_shape(desc, spec), desc.dtype) for desc, spec in zip(output_descs, output_specs)) + + return call + + def test_forward_declares_initialized_padded_output_for_fixed_and_thd(self): + fixed = self.forward.IndexerForward( + _Array((1, 8, 32, 128), self.bfloat16), + _Array((1, 5, 1, 128), self.bfloat16), + _Array((1, 8, 32), self.bfloat16), + ratio=1, + target_compute_capability=100, + ) + captured = {} + fixed._call_kernel = self.fake_call(captured) + result = fixed( + _Array((1, 8, 32, 128), self.bfloat16), + _Array((1, 5, 1, 128), self.bfloat16), + _Array((1, 8, 32), self.bfloat16), + ) + self.assertEqual(result["scores"].shape, (1, 8, 5)) + output = captured["output_descs"][0] + self.assertEqual(output.shape, (1, 8, 8)) + self.assertEqual(output.init_value, float("-inf")) + + thd = self.forward.IndexerForward( + _Array((8, 32, 128), self.bfloat16), + _Array((9, 1, 128), self.bfloat16), + _Array((8, 32), self.bfloat16), + sample_cu_seqlens_q=_Array((2,), self.int32), + sample_cu_seqlens_k=_Array((2,), self.int32), + max_seqlen_q=8, + max_seqlen_k=5, + target_compute_capability=90, + ) + captured = {} + thd._call_kernel = self.fake_call(captured) + result = thd( + _Array((8, 32, 128), self.bfloat16), + _Array((9, 1, 128), self.bfloat16), + _Array((8, 32), self.bfloat16), + cu_seqlens_q=_Array((2,), self.int32), + cu_seqlens_k=_Array((2,), self.int32), + ) + self.assertEqual(result["scores"].shape, (8, 5)) + self.assertEqual(captured["output_descs"][0].shape, (8, 8)) + + def test_forward_maps_independent_public_layouts_to_canonical_axes(self): + operation = self.forward.IndexerForward( + _Array((8, 2, 32, 128), self.bfloat16), + _Array((2, 5, 1, 128), self.bfloat16), + _Array((8, 2, 32), self.bfloat16), + ratio=1, + q_layout="SBHD", + k_layout="BSHD", + w_layout="SBH", + output_layout="SBK", + target_compute_capability=100, + ) + captured = {} + operation._call_kernel = self.fake_call(captured) + + result = operation( + _Array((8, 2, 32, 128), self.bfloat16), + _Array((2, 5, 1, 128), self.bfloat16), + _Array((8, 2, 32), self.bfloat16), + ) + + self.assertEqual(operation.q_desc.shape, (2, 8, 32, 128)) + self.assertEqual(operation.q_desc.stride_order, (3, 2, 0, 1)) + self.assertEqual(result["scores"].shape, (8, 2, 5)) + self.assertEqual( + tuple(spec.mode for spec in captured["input_spec"][:3]), + ((1, 0, 2, 3), (0, 1, 2, 3), (1, 0, 2)), + ) + self.assertEqual( + tuple(spec.layout for spec in captured["input_spec"][:3]), + ((3, 2, 1, 0), (3, 2, 1, 0), (2, 1, 0)), + ) + self.assertEqual(captured["output_spec"][0].mode, (1, 0, 2)) + self.assertEqual(captured["output_spec"][0].layout, (2, 1, 0)) + + explicit_output = self.forward.IndexerForward( + _Array((8, 2, 32, 128), self.bfloat16), + _Array((2, 5, 1, 128), self.bfloat16), + _Array((8, 2, 32), self.bfloat16), + sample_out=_Array((8, 2, 8), self.float32), + ratio=1, + q_layout="SBHD", + k_layout="BSHD", + w_layout="SBH", + output_layout="SBK", + target_compute_capability=100, + ) + self.assertEqual(explicit_output.o_desc.shape, (2, 8, 8)) + self.assertEqual(explicit_output.o_desc.stride_order, (2, 0, 1)) + + def test_forward_rejects_unimplemented_feature_axis_permutations(self): + with self.assertRaisesRegex(ValueError, "q_layout must be one of"): + self.forward.IndexerForward( + _Array((1, 32, 8, 128), self.bfloat16), + _Array((1, 5, 1, 128), self.bfloat16), + _Array((1, 8, 32), self.bfloat16), + q_layout="BHSD", + target_compute_capability=100, + ) + + def test_forward_layout_arguments_are_static(self): + static_argnames = set(self.forward.indexer_forward_wrapper._test_static_argnames) + self.assertTrue({"q_layout", "k_layout", "w_layout", "output_layout"}.issubset(static_argnames)) + + def test_topk_declares_outputs_and_hidden_workspace_for_both_modes(self): + for return_val in (True, False): + with self.subTest(return_val=return_val): + api = self.topk.IndexerTopK( + _Array((4, 64), self.float32), + _Array((4,), self.int32), + 8, + return_val=return_val, + target_compute_capability=100, + ) + captured = {} + api._call_kernel = self.fake_call(captured) + result = api(_Array((4, 64), self.float32), _Array((4,), self.int32)) + self.assertEqual(result["indices"].shape, (4, 8)) + self.assertEqual(result["values"] is not None, return_val) + self.assertEqual(captured["workspace_descs"][0].shape, (4, 2, 64)) + self.assertEqual(len(captured["output_descs"]), 2 if return_val else 1) + + def test_topk_constructor_reports_invalid_inputs_before_inference(self): + with self.assertRaisesRegex(ValueError, "input_values must have rank 2"): + self.topk.IndexerTopK( + _Array((64,), self.float32), + _Array((1,), self.int32), + 8, + target_compute_capability=100, + ) + + with self.assertRaisesRegex(ValueError, "input_values must have dtype"): + self.topk.IndexerTopK( + _Array((4, 64), self.int64), + _Array((4,), self.int32), + 8, + target_compute_capability=100, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_dsa_score_recompute_op.py b/test/python/fe_api/test_dsa_score_recompute_op.py new file mode 100644 index 000000000..31cd2a4b2 --- /dev/null +++ b/test/python/fe_api/test_dsa_score_recompute_op.py @@ -0,0 +1,231 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for DSA score-recompute operations.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_DSA_ROOT = _CUDNN_ROOT / "deepseek_sparse_attention" +_SCORE_ROOT = _DSA_ROOT / "score_recompute" +_PACKAGE = "cudnn_dsa_score_op_test" + + +class _DataType(Enum): + NOT_SET = auto() + FLOAT = auto() + BFLOAT16 = auto() + INT32 = auto() + + +class DsaScoreRecomputeOpTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + dsa_name = f"{_PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_DSA_ROOT)] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + sys.modules[dsa_name] = dsa + + score_name = f"{dsa_name}.score_recompute" + score = types.ModuleType(score_name) + score.__path__ = [str(_SCORE_ROOT)] + score.__package__ = score_name + score.__spec__ = ModuleSpec(score_name, loader=None, is_package=True) + sys.modules[score_name] = score + + try: + cls.tensor = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + cls.op = importlib.import_module(f"{score_name}.op") + cls.config = importlib.import_module(f"{score_name}.config") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def desc(self, shape, dtype=_DataType.BFLOAT16, *, stride_order=None, name=""): + return self.tensor.make_compact_tensor_desc( + dtype=dtype, + shape=tuple(shape), + stride_order=stride_order, + name=name, + ) + + def sparse( + self, *, target=100, score_type="indexer", with_length=False, **overrides + ): + aux_dtype = _DataType.BFLOAT16 if score_type == "indexer" else _DataType.FLOAT + arguments = { + "q": self.desc((2, 4, 32, 128), name="q"), + "k": self.desc((2, 256, 128), name="k"), + "per_head": self.desc((2, 4, 32), aux_dtype, name="per_head"), + "topk_indices": self.desc( + (2, 4, 128), _DataType.INT32, name="topk_indices" + ), + "output": self.desc((2, 4, 128), _DataType.FLOAT, name="output"), + "topk_length": self.desc((2, 4), _DataType.INT32, name="topk_length") + if with_length + else None, + "score_type": score_type, + "softmax_scale": 0.125 if score_type == "attention" else 1.0, + "target_compute_capability": target, + } + arguments.update(overrides) + return self.op.SparseScoreRecomputeOp(**arguments) + + def dense(self, *, target=100, score_type="indexer", thd=False, **overrides): + aux_dtype = _DataType.BFLOAT16 if score_type == "indexer" else _DataType.FLOAT + if thd: + q_shape, k_shape, aux_shape = (7, 32, 128), (11, 1, 128), (7, 32) + output_shape, denominator_shape = (7, 8), (7,) + cu_q = self.desc((3,), _DataType.INT32, name="cu_q") + cu_k = self.desc((3,), _DataType.INT32, name="cu_k") + max_q, max_k = 4, 8 + else: + q_shape, k_shape, aux_shape = (2, 4, 32, 128), (2, 8, 1, 128), (2, 4, 32) + output_shape, denominator_shape = (2, 4, 8), (2, 4) + cu_q = cu_k = None + max_q = max_k = None + arguments = { + "q": self.desc(q_shape, name="q"), + "k": self.desc(k_shape, name="k"), + "per_head": self.desc(aux_shape, aux_dtype, name="per_head"), + "output": self.desc(output_shape, _DataType.FLOAT, name="output"), + "denominator": self.desc( + denominator_shape, _DataType.FLOAT, name="denominator" + ), + "score_type": score_type, + "scale": 0.125 if score_type == "attention" else 1.0, + "ratio": 2, + "target_compute_capability": target, + "is_thd": thd, + "cu_seqlens_q": cu_q, + "cu_seqlens_k": cu_k, + "max_seqlen_q": max_q, + "max_seqlen_k": max_k, + } + arguments.update(overrides) + return self.op.DenseScoreRecomputeOp(**arguments) + + def test_sparse_resolves_sm100_and_preserves_explicit_lengths(self): + operation = self.sparse(with_length=True) + self.assertTrue(operation.check_support()) + self.assertEqual(operation.qhead_per_kv_head, 32) + self.assertIsInstance(operation.config, self.config.SparseScoreKernelConfig) + self.assertEqual( + (operation.config.m_block_size, operation.config.n_block_size), (32, 128) + ) + self.assertTrue(operation.config.have_topk_length) + + def test_sparse_resolves_sm90(self): + operation = self.sparse(target=90) + self.assertTrue(operation.check_support()) + self.assertIsInstance(operation.config, self.op.SparseScoreSm90Config) + self.assertEqual( + ( + operation.config.tile_m, + operation.config.tile_n, + operation.config.num_threads, + ), + (32, 64, 256), + ) + + def test_dense_bshd_and_thd_resolve_sm100(self): + bshd = self.dense() + self.assertTrue(bshd.check_support()) + self.assertEqual((bshd.max_seqlen_q, bshd.max_seqlen_k), (4, 8)) + self.assertIsInstance(bshd.config, self.config.DenseScoreKernelConfig) + + thd = self.dense(thd=True, score_type="attention") + self.assertTrue(thd.check_support()) + self.assertEqual((thd.max_seqlen_q, thd.max_seqlen_k), (4, 8)) + + def test_dense_sm90_thd_is_explicitly_unsupported(self): + with self.assertRaisesRegex( + NotImplementedError, "host-side sequence-length reads" + ): + self.dense(target=90, thd=True).check_support() + + def test_complete_signatures_are_cross_checked(self): + with self.assertRaisesRegex(ValueError, "output must have shape"): + self.sparse(output=self.desc((2, 4, 64), _DataType.FLOAT)).check_support() + with self.assertRaisesRegex(ValueError, "per-head tensor must have shape"): + self.dense(per_head=self.desc((2, 4, 16))).check_support() + with self.assertRaisesRegex(ValueError, "qhead_per_kv_head must equal"): + self.dense(qhead_per_kv_head=16).check_support() + with self.assertRaisesRegex(ValueError, "ratio must be >= 1"): + self.dense(ratio=0).check_support() + + def test_native_abi_feature_axes_must_be_contiguous(self): + unsupported_layout = self.tensor.TensorDesc( + dtype=_DataType.BFLOAT16, + shape=(2, 4, 32, 128), + stride=(16384, 128, 512, 1), + stride_order=(3, 1, 2, 0), + ) + with self.assertRaisesRegex( + ValueError, "Q must have its final 2 axes contiguous" + ): + self.sparse(q=unsupported_layout).check_support() + + def test_sm100_configuration_rejects_incompatible_topk(self): + with self.assertRaisesRegex( + ValueError, "multiple of the selected n_block_size" + ): + self.sparse( + topk_indices=self.desc((2, 4, 96), _DataType.INT32), + output=self.desc((2, 4, 96), _DataType.FLOAT), + ).check_support() + + def test_rejects_signatures_the_native_kernels_cannot_cover(self): + with self.assertRaisesRegex(ValueError, "SM90.*multiple of 128"): + self.sparse( + target=90, + topk_indices=self.desc((2, 4, 96), _DataType.INT32), + output=self.desc((2, 4, 96), _DataType.FLOAT), + ).check_support() + + with self.assertRaisesRegex(ValueError, "requires MQA with H_kv=1"): + self.dense( + q=self.desc((2, 4, 64, 128)), + k=self.desc((2, 8, 2, 128)), + per_head=self.desc((2, 4, 64)), + ).check_support() + + def test_common_modules_do_not_load_frameworks_or_kernels(self): + self.sparse().check_support() + prefix = f"{_PACKAGE}.deepseek_sparse_attention.score_recompute" + self.assertNotIn(f"{prefix}.jax", sys.modules) + self.assertNotIn(f"{prefix}.api", sys.modules) + self.assertNotIn(f"{prefix}.sparse_score_recompute_sm100", sys.modules) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_dsa_sparse_attention_backward_op.py b/test/python/fe_api/test_dsa_sparse_attention_backward_op.py new file mode 100644 index 000000000..f81acc026 --- /dev/null +++ b/test/python/fe_api/test_dsa_sparse_attention_backward_op.py @@ -0,0 +1,193 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for DSA sparse-attention backward metadata.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "deepseek_sparse_attention" / "sparse_attention_backward" +_PACKAGE = "cudnn_dsa_sparse_attention_backward_op_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + UINT8 = auto() + + +class SparseAttentionBackwardOpContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + dsa_name = f"{_PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_CUDNN_ROOT / "deepseek_sparse_attention")] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + sys.modules[dsa_name] = dsa + + operation_name = f"{dsa_name}.sparse_attention_backward" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + try: + cls.tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}.common.op") + cls.op_module = importlib.import_module(f"{operation_name}.op") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def _desc(self, shape, dtype=_DataType.BFLOAT16, *, stride=None, name=""): + shape = tuple(shape) + if stride is None: + stride_values = [0] * len(shape) + running = 1 + for dimension in reversed(range(len(shape))): + stride_values[dimension] = running + running *= max(shape[dimension], 1) + stride = tuple(stride_values) + else: + stride = tuple(stride) + stride_order = tuple(index for index, _ in sorted(enumerate(stride), key=lambda item: (item[1], shape[item[0]]))) + return self.tensor_module.TensorDesc( + dtype=dtype, + shape=shape, + stride=stride, + stride_order=stride_order, + name=name, + ) + + def _op(self, *, head_dim=512, dtype=_DataType.BFLOAT16, with_length=True, **overrides): + m, n, h, topk = 65, 130, 64, 32 + head_dim_v = 512 if head_dim == 576 else head_dim + arguments = { + "q": self._desc((m, h, head_dim), dtype, name="q"), + "kv": self._desc((n, head_dim), dtype, name="kv"), + "output": self._desc((m, h, head_dim_v), dtype, name="output"), + "doutput": self._desc((m, h, head_dim_v), dtype, name="doutput"), + "lse": self._desc((m, h), _DataType.FLOAT, name="lse"), + "attn_sink": self._desc((h,), _DataType.FLOAT, name="attn_sink"), + "topk_idxs": self._desc((m, topk), _DataType.INT32, name="topk_idxs"), + "topk_length": self._desc((m,), _DataType.INT32, name="topk_length") if with_length else None, + "dq": self._desc((m, h, head_dim), dtype, name="dq"), + "dkv": self._desc((n, head_dim), dtype, name="dkv"), + "d_sink": self._desc((h,), _DataType.FLOAT, name="d_sink"), + } + arguments.update(overrides) + return self.op_module.SparseAttentionBackwardOp(**arguments) + + def test_validates_both_supported_head_dimensions(self): + for head_dim, expected_v in ((512, 512), (576, 512)): + with self.subTest(head_dim=head_dim): + operation = self._op(head_dim=head_dim) + self.assertIsInstance(operation, self.base_module.Op) + self.assertTrue(operation.check_support()) + self.assertEqual(operation.head_dim_v, expected_v) + self.assertEqual(operation.softmax_scale, 1.0 / head_dim**0.5) + self.assertEqual( + (operation.total_seqlen_q, operation.total_seqlen_kv, operation.num_heads, operation.max_topk), + (65, 130, 64, 32), + ) + + def test_supports_fp16_and_optional_topk_length(self): + operation = self._op(dtype=_DataType.HALF, with_length=False, softmax_scale=0.125) + self.assertTrue(operation.check_support()) + self.assertIsNone(operation.topk_length) + self.assertEqual(operation.softmax_scale, 0.125) + + def test_rejects_incomplete_signatures(self): + cases = ( + ({"q": self._desc((65, 512))}, "Q must have rank 3"), + ({"kv": self._desc((130, 256))}, "KV head dimension must match Q"), + ({"output": self._desc((65, 64, 576))}, "O must have shape"), + ({"doutput": self._desc((65, 64, 576))}, "dO must have shape"), + ({"lse": self._desc((65, 32), _DataType.FLOAT)}, "LSE must have shape"), + ({"attn_sink": self._desc((32,), _DataType.FLOAT)}, "attn_sink must have shape"), + ({"topk_idxs": self._desc((64, 32), _DataType.INT32)}, "topk_idxs must have shape"), + ({"topk_length": self._desc((64,), _DataType.INT32)}, "topk_length must have shape"), + ({"dq": self._desc((64, 64, 512))}, "dQ must have shape"), + ({"dkv": self._desc((130, 256))}, "dKV must have shape"), + ({"d_sink": self._desc((32,), _DataType.FLOAT)}, "d_sink must have shape"), + ({"q": self._desc((65, 64, 512), _DataType.FLOAT)}, "Q must have dtype"), + ({"topk_idxs": self._desc((65, 32), _DataType.FLOAT)}, "topk_idxs must have dtype int32"), + ( + {"q": self._desc((65, 64, 256)), "kv": self._desc((130, 256)), "dq": self._desc((65, 64, 256)), "dkv": self._desc((130, 256))}, + "head_dim must be one of", + ), + ( + { + "q": self._desc((65, 32, 512)), + "output": self._desc((65, 32, 512)), + "doutput": self._desc((65, 32, 512)), + "lse": self._desc((65, 32), _DataType.FLOAT), + "attn_sink": self._desc((32,), _DataType.FLOAT), + "dq": self._desc((65, 32, 512)), + "d_sink": self._desc((32,), _DataType.FLOAT), + }, + "num_heads must be divisible", + ), + ({"q": self._desc((65, 64, 512), stride=(1, 65 * 512, 65))}, "must be row-major contiguous"), + ({"block_tile": 32}, "block_tile must be 64"), + ({"softmax_scale": float("inf")}, "softmax_scale must be finite"), + ) + for overrides, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self._op(**overrides).check_support() + + def test_explicit_scale_preserves_finite_values(self): + for scale in (0.0, -0.125): + with self.subTest(scale=scale): + operation = self._op(softmax_scale=scale) + self.assertTrue(operation.check_support()) + self.assertEqual(operation.softmax_scale, scale) + + def test_operation_imports_no_framework_or_kernel(self): + source = (_OPERATION_ROOT / "op.py").read_text() + for forbidden in ("import torch", "import jax", "import cutlass", "import cuda"): + self.assertNotIn(forbidden, source) + operation = self._op() + operation.check_support() + prefix = f"{_PACKAGE}.deepseek_sparse_attention.sparse_attention_backward" + self.assertNotIn(f"{prefix}.api", sys.modules) + self.assertNotIn(f"{prefix}.jax", sys.modules) + self.assertNotIn(f"{prefix}.dsa_bwd_sm90", sys.modules) + self.assertNotIn(f"{prefix}.dsa_bwd_sm100", sys.modules) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_gemm_amax_imports.py b/test/python/fe_api/test_gemm_amax_imports.py new file mode 100644 index 000000000..7b19102cc --- /dev/null +++ b/test/python/fe_api/test_gemm_amax_imports.py @@ -0,0 +1,148 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Import and Torch-adapter contracts for dense GEMM + amax.""" + +import ast +import importlib.util +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "gemm_amax" + + +def _load_operation_package(name: str): + parent_name, separator, _ = name.rpartition(".") + if not separator: + raise ValueError("Synthetic operation package must have a parent") + parent = types.ModuleType(parent_name) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = parent_name + sys.modules[parent_name] = parent + + spec = importlib.util.spec_from_file_location( + name, + _OPERATION_ROOT / "__init__.py", + submodule_search_locations=[str(_OPERATION_ROOT)], + ) + if spec is None or spec.loader is None: + raise RuntimeError("Unable to load GEMM + amax package") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _remove_package(name: str) -> None: + root_name = name.split(".", 1)[0] + for module_name in tuple(sys.modules): + if module_name == root_name or module_name.startswith(f"{root_name}."): + sys.modules.pop(module_name, None) + + +class GemmAmaxImportContractTest(unittest.TestCase): + def test_package_import_loads_no_framework_adapter(self): + name = "cudnn_gemm_amax_import_test.gemm_amax" + blocked = {module_name: None for module_name in ("torch", "jax", "cutlass", "cuda")} + try: + with mock.patch.dict(sys.modules, blocked): + package = _load_operation_package(name) + self.assertNotIn(f"{name}.api", sys.modules) + self.assertNotIn(f"{name}.op", sys.modules) + self.assertTrue( + { + "GemmAmaxSm100Op", + "GemmAmaxSm100", + "gemm_amax_wrapper_sm100", + "api", + "op", + }.issubset(dir(package)) + ) + finally: + _remove_package(name) + + def test_static_adapter_boundaries_and_public_torch_signatures(self): + def imports_framework(node, framework: str): + if isinstance(node, ast.Import): + return any(alias.name == framework or alias.name.startswith(f"{framework}.") for alias in node.names) + return node.level == 0 and (node.module == framework or (node.module or "").startswith(f"{framework}.")) + + op_tree = ast.parse((_OPERATION_ROOT / "op.py").read_text(), filename="op.py") + api_tree = ast.parse((_OPERATION_ROOT / "api.py").read_text(), filename="api.py") + op_imports = [node for node in ast.walk(op_tree) if isinstance(node, (ast.Import, ast.ImportFrom))] + api_imports = [node for node in ast.walk(api_tree) if isinstance(node, (ast.Import, ast.ImportFrom))] + self.assertFalse(any(imports_framework(node, framework) for node in op_imports for framework in ("torch", "jax", "cutlass", "cuda"))) + self.assertFalse(any(imports_framework(node, "jax") for node in api_imports)) + + adapter = next(node for node in api_tree.body if isinstance(node, ast.ClassDef) and node.name == "GemmAmaxSm100") + constructor = next(node for node in adapter.body if isinstance(node, ast.FunctionDef) and node.name == "__init__") + self.assertEqual( + [argument.arg for argument in constructor.args.args], + [ + "self", + "sample_a", + "sample_b", + "sample_sfa", + "sample_sfb", + "sample_c", + "sample_amax", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + ], + ) + op_constructors = [ + node for node in ast.walk(constructor) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "GemmAmaxSm100Op" + ] + self.assertEqual(len(op_constructors), 1) + + execute = next(node for node in adapter.body if isinstance(node, ast.FunctionDef) and node.name == "execute") + self.assertEqual( + [argument.arg for argument in execute.args.args], + [ + "self", + "a_tensor", + "b_tensor", + "sfa_tensor", + "sfb_tensor", + "c_tensor", + "amax_tensor", + "current_stream", + ], + ) + + wrapper = next(node for node in api_tree.body if isinstance(node, ast.FunctionDef) and node.name == "gemm_amax_wrapper_sm100") + self.assertEqual( + [argument.arg for argument in wrapper.args.args], + [ + "a_tensor", + "b_tensor", + "sfa_tensor", + "sfb_tensor", + "c_major", + "c_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "stream", + ], + ) + self.assertFalse(any(isinstance(node, ast.FunctionDef) and node.name.startswith("_check_torch_support") for node in adapter.body)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_gemm_amax_op.py b/test/python/fe_api/test_gemm_amax_op.py new file mode 100644 index 000000000..942a86b0e --- /dev/null +++ b/test/python/fe_api/test_gemm_amax_op.py @@ -0,0 +1,295 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for the dense GEMM + amax operation.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "gemm_amax" +_PACKAGE = "cudnn_gemm_amax_op_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT8 = auto() + UINT8 = auto() + FP4_E2M1 = auto() + FP8_E4M3 = auto() + FP8_E5M2 = auto() + FP8_E8M0 = auto() + + +class GemmAmaxOpContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + operation_name = f"{_PACKAGE}.gemm_amax" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + try: + cls.tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}.common.op") + cls.op_module = importlib.import_module(f"{operation_name}.op") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def _desc(self, shape, dtype, order, name): + shape = tuple(shape) + stride = [0] * len(shape) + running = 1 + for dimension in order: + stride[dimension] = running + running *= max(shape[dimension], 1) + return self.tensor_module.TensorDesc( + dtype=dtype, + shape=shape, + stride=tuple(stride), + stride_order=tuple(order), + name=name, + ) + + def _op(self, **overrides): + m, n, k, batch = 128, 256, 128, 2 + scale_order = (3, 1, 0, 4, 2, 5) + arguments = { + "a": self._desc((m, k, batch), _DataType.FP8_E4M3, (1, 0, 2), "A"), + "b": self._desc((n, k, batch), _DataType.FP8_E4M3, (1, 0, 2), "B"), + "sfa": self._desc((32, 4, 1, 4, 1, batch), _DataType.FP8_E8M0, scale_order, "SFA"), + "sfb": self._desc((32, 4, 2, 4, 1, batch), _DataType.FP8_E8M0, scale_order, "SFB"), + "c": self._desc((m, n, batch), _DataType.FLOAT, (1, 0, 2), "C"), + "amax": self._desc((1, 1, 1), _DataType.FLOAT, (2, 1, 0), "Amax"), + } + arguments.update(overrides) + return self.op_module.GemmAmaxSm100Op(**arguments) + + def test_validates_complete_signature_and_resolves_configuration(self): + operation = self._op() + + self.assertIsInstance(operation, self.base_module.Op) + self.assertTrue(operation.check_support()) + self.assertEqual((operation.m, operation.n, operation.k, operation.l), (128, 256, 128, 2)) + self.assertEqual((operation.a_major, operation.b_major, operation.c_major), ("k", "k", "n")) + self.assertEqual(operation.ab_dtype, _DataType.FP8_E4M3) + self.assertEqual(operation.scale_dtype, _DataType.FP8_E8M0) + self.assertEqual(operation.c_dtype, _DataType.FLOAT) + self.assertEqual(operation.mma_tiler_mn, (128, 128)) + self.assertEqual(operation.cluster_shape_mn, (1, 1)) + self.assertNotIn(f"{_PACKAGE}.gemm_amax.api", sys.modules) + self.assertNotIn( + f"{_PACKAGE}.gemm_amax.dense_blockscaled_gemm_persistent_amax", + sys.modules, + ) + + def test_accepts_supported_major_modes_and_storage_aliases(self): + scale_order = (3, 1, 0, 4, 2, 5) + operation = self._op( + a=self._desc((128, 128, 2), _DataType.UINT8, (1, 0, 2), "A"), + b=self._desc((256, 128, 2), _DataType.UINT8, (1, 0, 2), "B"), + sfa=self._desc((32, 4, 1, 4, 2, 2), _DataType.INT8, scale_order, "SFA"), + sfb=self._desc((32, 4, 2, 4, 2, 2), _DataType.INT8, scale_order, "SFB"), + c=self._desc((128, 256, 2), _DataType.UINT8, (1, 0, 2), "C"), + sf_vec_size=16, + ) + self.assertTrue(operation.check_support()) + self.assertEqual(operation.ab_dtype, _DataType.FP4_E2M1) + self.assertEqual(operation.scale_dtype, _DataType.FP8_E8M0) + self.assertEqual(operation.c_dtype, _DataType.FP4_E2M1) + + alternate = self._op( + a=self._desc((128, 128, 2), _DataType.FP8_E4M3, (0, 1, 2), "A"), + b=self._desc((256, 128, 2), _DataType.FP8_E4M3, (0, 1, 2), "B"), + c=self._desc((128, 256, 2), _DataType.BFLOAT16, (0, 1, 2), "C"), + ) + self.assertTrue(alternate.check_support()) + self.assertEqual((alternate.a_major, alternate.b_major, alternate.c_major), ("m", "n", "m")) + + def test_scale_layout_accepts_size_one_stride_order_ties(self): + scale = self.tensor_module.TensorDesc( + dtype=_DataType.FP8_E8M0, + shape=(32, 4, 1, 4, 1, 1), + stride=(16, 4, 512, 1, 512, 512), + stride_order=(3, 1, 0, 2, 4, 5), + name="scale", + ) + operation = self._op( + a=self._desc((128, 128, 1), _DataType.FP8_E4M3, (1, 0, 2), "A"), + b=self._desc((128, 128, 1), _DataType.FP8_E4M3, (1, 0, 2), "B"), + sfa=scale, + sfb=scale, + c=self._desc((128, 128, 1), _DataType.FLOAT, (1, 0, 2), "C"), + ) + self.assertTrue(operation.check_support()) + + def test_existing_torch_configuration_families_remain_accepted(self): + scale_order = (3, 1, 0, 4, 2, 5) + cases = ( + { + "a": self._desc((128, 256, 2), _DataType.FP8_E5M2, (0, 1, 2), "A"), + "b": self._desc((256, 256, 2), _DataType.FP8_E5M2, (0, 1, 2), "B"), + "sfa": self._desc((32, 4, 1, 4, 2, 2), _DataType.FP8_E8M0, scale_order, "SFA"), + "sfb": self._desc((32, 4, 2, 4, 2, 2), _DataType.FP8_E8M0, scale_order, "SFB"), + "c": self._desc((128, 256, 2), _DataType.BFLOAT16, (0, 1, 2), "C"), + "mma_tiler_mn": (128, 256), + "cluster_shape_mn": (2, 2), + }, + { + "a": self._desc((128, 128, 2), _DataType.FP4_E2M1, (1, 0, 2), "A"), + "b": self._desc((256, 128, 2), _DataType.FP4_E2M1, (1, 0, 2), "B"), + "sfa": self._desc((32, 4, 1, 4, 2, 2), _DataType.FP8_E4M3, scale_order, "SFA"), + "sfb": self._desc((32, 4, 2, 4, 2, 2), _DataType.FP8_E4M3, scale_order, "SFB"), + "c": self._desc((128, 256, 2), _DataType.FP8_E4M3, (1, 0, 2), "C"), + "sf_vec_size": 16, + "cluster_shape_mn": (2, 2), + }, + { + "a": self._desc((128, 128, 2), _DataType.FP4_E2M1, (1, 0, 2), "A"), + "b": self._desc((256, 128, 2), _DataType.FP4_E2M1, (1, 0, 2), "B"), + "c": self._desc((128, 256, 2), _DataType.FLOAT, (0, 1, 2), "C"), + "cluster_shape_mn": (4, 4), + }, + ) + for overrides in cases: + with self.subTest(overrides=tuple(sorted(overrides))): + self.assertTrue(self._op(**overrides).check_support()) + + def test_rejects_inconsistent_shapes_and_scale_layout(self): + scale_order = (3, 1, 0, 4, 2, 5) + cases = ( + ( + {"b": self._desc((256, 64, 2), _DataType.FP8_E4M3, (1, 0, 2), "B")}, + "B shape mismatch", + ), + ( + {"c": self._desc((128, 128, 2), _DataType.FLOAT, (1, 0, 2), "C")}, + "C must have shape", + ), + ( + { + "sfa": self._desc( + (32, 4, 1, 4, 2, 2), + _DataType.FP8_E8M0, + scale_order, + "SFA", + ) + }, + "SFA must have shape", + ), + ( + { + "sfa": self._desc( + (32, 4, 1, 4, 1, 2), + _DataType.FP8_E8M0, + (5, 4, 3, 2, 1, 0), + "SFA", + ) + }, + "packed block-scale layout", + ), + ( + {"amax": self._desc((1,), _DataType.FLOAT, (0,), "Amax")}, + "Amax must have shape", + ), + ) + for overrides, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self._op(**overrides).check_support() + + def test_rejects_invalid_dtype_combinations(self): + cases = ( + ( + {"b": self._desc((256, 128, 2), _DataType.FP8_E5M2, (1, 0, 2), "B")}, + ValueError, + "A and B must have the same dtype", + ), + ( + {"amax": self._desc((1, 1, 1), _DataType.HALF, (2, 1, 0), "Amax")}, + ValueError, + "Amax must have dtype float32", + ), + ( + {"acc_dtype": _DataType.HALF}, + ValueError, + "Accumulator dtype must be float32", + ), + ( + {"c": self._desc((128, 256, 2), _DataType.FP8_E5M2, (1, 0, 2), "C")}, + NotImplementedError, + "FP8 A and B with FP8 C", + ), + ( + {"sf_vec_size": 16}, + ValueError, + "FP8 A and B do not support sf_vec_size=16", + ), + ) + for overrides, error_type, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(error_type, message): + self._op(**overrides).check_support() + + def test_validates_tiler_cluster_and_alignment(self): + with self.assertRaisesRegex(ValueError, r"mma_tiler_mn\[1\]"): + self._op(mma_tiler_mn=(128, 64)).check_support() + with self.assertRaisesRegex(NotImplementedError, "currently hangs"): + self._op(mma_tiler_mn=(256, 128)).check_support() + with self.assertRaisesRegex(ValueError, "entries must be at most 4"): + self._op(cluster_shape_mn=(8, 1)).check_support() + with self.assertRaisesRegex(TypeError, "mma_tiler_mn must contain two integers"): + self._op(mma_tiler_mn=None).check_support() + + scale_order = (3, 1, 0, 4, 2, 5) + with self.assertRaisesRegex(ValueError, "16-byte alignment"): + self._op( + a=self._desc((128, 120, 2), _DataType.FP8_E4M3, (1, 0, 2), "A"), + b=self._desc((256, 120, 2), _DataType.FP8_E4M3, (1, 0, 2), "B"), + sfa=self._desc((32, 4, 1, 4, 1, 2), _DataType.FP8_E8M0, scale_order, "SFA"), + sfb=self._desc((32, 4, 2, 4, 1, 2), _DataType.FP8_E8M0, scale_order, "SFB"), + ).check_support() + + def test_constructor_requires_descriptors_and_canonical_accumulator(self): + for name in ("a", "b", "sfa", "sfb", "c", "amax"): + with self.subTest(name=name): + with self.assertRaisesRegex(TypeError, f"{name} must be a TensorDesc"): + self._op(**{name: types.SimpleNamespace(shape=(1,))}) + with self.assertRaisesRegex(TypeError, "acc_dtype must be a cudnn.data_type"): + self._op(acc_dtype="float32") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_gemm_relu_imports.py b/test/python/fe_api/test_gemm_relu_imports.py new file mode 100644 index 000000000..10dc35344 --- /dev/null +++ b/test/python/fe_api/test_gemm_relu_imports.py @@ -0,0 +1,123 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Import boundaries for dense GEMM + sReLU and dsReLU.""" + +import ast +import importlib.util +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" + + +def _load_operation_package(name: str, operation: str): + parent_name = name.rpartition(".")[0] + parent = types.ModuleType(parent_name) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = parent_name + sys.modules[parent_name] = parent + + operation_root = _CUDNN_ROOT / operation + spec = importlib.util.spec_from_file_location( + name, + operation_root / "__init__.py", + submodule_search_locations=[str(operation_root)], + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"Unable to load {operation} package") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _remove_package(name: str) -> None: + root_name = name.split(".", 1)[0] + for module_name in tuple(sys.modules): + if module_name == root_name or module_name.startswith(f"{root_name}."): + sys.modules.pop(module_name, None) + + +class GemmReluImportContractTest(unittest.TestCase): + def test_public_jax_facade_routes_both_operations(self): + tree = ast.parse((_CUDNN_ROOT / "jax" / "__init__.py").read_text(), filename="cudnn/jax/__init__.py") + assignment = next( + node + for node in tree.body + if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == "_OPERATION_EXPORTS" for target in node.targets) + ) + exports = ast.literal_eval(assignment.value) + self.assertEqual(exports["GemmSreluSm100"], ("..gemm_srelu.jax", "GemmSreluSm100")) + self.assertEqual(exports["gemm_srelu_wrapper_sm100"], ("..gemm_srelu.jax", "gemm_srelu_wrapper_sm100")) + self.assertEqual(exports["GemmDsreluSm100"], ("..gemm_dsrelu.jax", "GemmDsreluSm100")) + self.assertEqual(exports["gemm_dsrelu_wrapper_sm100"], ("..gemm_dsrelu.jax", "gemm_dsrelu_wrapper_sm100")) + + def test_package_imports_are_framework_free_and_lazy(self): + blocked = {module_name: None for module_name in ("torch", "jax", "cutlass", "cuda")} + for operation, op_class, api_class, wrapper in ( + ("gemm_srelu", "GemmSreluSm100Op", "GemmSreluSm100", "gemm_srelu_wrapper_sm100"), + ("gemm_dsrelu", "GemmDsreluSm100Op", "GemmDsreluSm100", "gemm_dsrelu_wrapper_sm100"), + ): + name = f"cudnn_{operation}_import_test.{operation}" + try: + with self.subTest(operation=operation), mock.patch.dict(sys.modules, blocked): + package = _load_operation_package(name, operation) + self.assertNotIn(f"{name}.api", sys.modules) + self.assertNotIn(f"{name}.op", sys.modules) + self.assertTrue({op_class, api_class, wrapper, "api", "op"}.issubset(dir(package))) + finally: + _remove_package(name) + + def test_exports_route_to_op_and_torch_modules_independently(self): + for operation, op_class, api_class, wrapper in ( + ("gemm_srelu", "GemmSreluSm100Op", "GemmSreluSm100", "gemm_srelu_wrapper_sm100"), + ("gemm_dsrelu", "GemmDsreluSm100Op", "GemmDsreluSm100", "gemm_dsrelu_wrapper_sm100"), + ): + name = f"cudnn_{operation}_route_test.{operation}" + op_sentinel = object() + api_sentinel = object() + op_module = types.ModuleType(f"{name}.op") + setattr(op_module, op_class, op_sentinel) + api_module = types.ModuleType(f"{name}.api") + setattr(api_module, api_class, api_sentinel) + setattr(api_module, wrapper, object()) + try: + with self.subTest(operation=operation): + sys.modules[op_module.__name__] = op_module + sys.modules[api_module.__name__] = api_module + package = _load_operation_package(name, operation) + self.assertIs(getattr(package, op_class), op_sentinel) + self.assertIs(getattr(package, api_class), api_sentinel) + finally: + _remove_package(name) + + def test_framework_neutral_op_modules_have_no_framework_imports(self): + paths = [ + _CUDNN_ROOT / "gemm" / "__init__.py", + _CUDNN_ROOT / "gemm" / "helpers.py", + _CUDNN_ROOT / "gemm" / "srelu.py", + _CUDNN_ROOT / "gemm_srelu" / "op.py", + _CUDNN_ROOT / "gemm_dsrelu" / "op.py", + ] + for path in paths: + imports = [node for node in ast.walk(ast.parse(path.read_text(), filename=str(path))) if isinstance(node, (ast.Import, ast.ImportFrom))] + imported = {alias.name.split(".", 1)[0] for node in imports if isinstance(node, ast.Import) for alias in node.names} + imported.update((node.module or "").split(".", 1)[0] for node in imports if isinstance(node, ast.ImportFrom) and node.level == 0) + with self.subTest(path=path.name): + self.assertTrue(imported.isdisjoint({"torch", "jax", "cutlass", "cuda"})) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_gemm_swiglu_imports.py b/test/python/fe_api/test_gemm_swiglu_imports.py new file mode 100644 index 000000000..db44c06aa --- /dev/null +++ b/test/python/fe_api/test_gemm_swiglu_imports.py @@ -0,0 +1,192 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Import and Torch-adapter contracts for dense GEMM + SwiGLU.""" + +import ast +import importlib.util +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "gemm_swiglu" + + +def _load_operation_package(name: str): + parent_name, separator, _ = name.rpartition(".") + if not separator: + raise ValueError("Synthetic operation package must have a parent") + parent = types.ModuleType(parent_name) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = parent_name + sys.modules[parent_name] = parent + + spec = importlib.util.spec_from_file_location( + name, + _OPERATION_ROOT / "__init__.py", + submodule_search_locations=[str(_OPERATION_ROOT)], + ) + if spec is None or spec.loader is None: + raise RuntimeError("Unable to load GEMM + SwiGLU package") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _remove_package(name: str) -> None: + root_name = name.split(".", 1)[0] + for module_name in tuple(sys.modules): + if module_name == root_name or module_name.startswith(f"{root_name}."): + sys.modules.pop(module_name, None) + + +class GemmSwigluImportContractTest(unittest.TestCase): + def test_package_import_loads_no_framework_adapter(self): + name = "cudnn_gemm_swiglu_import_test.gemm_swiglu" + blocked = {module_name: None for module_name in ("torch", "jax", "cutlass", "cuda")} + try: + with mock.patch.dict(sys.modules, blocked): + package = _load_operation_package(name) + self.assertNotIn(f"{name}.api", sys.modules) + self.assertNotIn(f"{name}.op", sys.modules) + self.assertTrue( + { + "BlockScaledGemmSwigluSm100Op", + "GemmSwigluSm100Op", + "GemmSwigluSm100", + "gemm_swiglu_wrapper_sm100", + "api", + "op", + }.issubset(dir(package)) + ) + finally: + _remove_package(name) + + def test_exports_route_without_loading_the_other_layer(self): + name = "cudnn_gemm_swiglu_route_test.gemm_swiglu" + api_sentinel = object() + op_sentinel = object() + api = types.ModuleType(f"{name}.api") + api.GemmSwigluSm100 = api_sentinel + api.gemm_swiglu_wrapper_sm100 = object() + op = types.ModuleType(f"{name}.op") + op.BlockScaledGemmSwigluSm100Op = object() + op.GemmSwigluSm100Op = op_sentinel + try: + sys.modules[api.__name__] = api + sys.modules[op.__name__] = op + package = _load_operation_package(name) + self.assertIs(package.GemmSwigluSm100Op, op_sentinel) + self.assertIs(package.GemmSwigluSm100, api_sentinel) + finally: + _remove_package(name) + + def test_static_adapter_boundaries_and_torch_constructor(self): + def imports_framework(node, framework: str): + if isinstance(node, ast.Import): + return any(alias.name == framework or alias.name.startswith(f"{framework}.") for alias in node.names) + return node.level == 0 and (node.module == framework or (node.module or "").startswith(f"{framework}.")) + + op_tree = ast.parse((_OPERATION_ROOT / "op.py").read_text(), filename="op.py") + api_tree = ast.parse((_OPERATION_ROOT / "api.py").read_text(), filename="api.py") + op_imports = [node for node in ast.walk(op_tree) if isinstance(node, (ast.Import, ast.ImportFrom))] + api_imports = [node for node in ast.walk(api_tree) if isinstance(node, (ast.Import, ast.ImportFrom))] + self.assertFalse(any(imports_framework(node, framework) for node in op_imports for framework in ("torch", "jax", "cutlass", "cuda"))) + self.assertFalse(any(imports_framework(node, "jax") for node in api_imports)) + + adapter = next(node for node in api_tree.body if isinstance(node, ast.ClassDef) and node.name == "GemmSwigluSm100") + self.assertEqual( + [base.id for base in adapter.bases if isinstance(base, ast.Name)], + ["APIBase"], + ) + constructor = next(node for node in adapter.body if isinstance(node, ast.FunctionDef) and node.name == "__init__") + self.assertEqual( + [argument.arg for argument in constructor.args.args], + [ + "self", + "sample_a", + "sample_b", + "sample_ab12", + "sample_c", + "alpha", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sample_sfa", + "sample_sfb", + "sample_amax", + "sample_sfc", + "sample_norm_const", + "sf_vec_size", + "vector_f32", + "ab12_stages", + ], + ) + required_count = len(constructor.args.args) - len(constructor.args.defaults) + self.assertEqual( + [argument.arg for argument in constructor.args.args[:required_count]], + ["self", "sample_a", "sample_b", "sample_ab12", "sample_c"], + ) + op_constructors = [ + node for node in ast.walk(constructor) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "GemmSwigluSm100Op" + ] + self.assertEqual(len(op_constructors), 1) + self.assertTrue({"a", "b", "ab12", "c"}.issubset({keyword.arg for keyword in op_constructors[0].keywords})) + + execute = next(node for node in adapter.body if isinstance(node, ast.FunctionDef) and node.name == "execute") + self.assertEqual( + [argument.arg for argument in execute.args.args], + [ + "self", + "a_tensor", + "b_tensor", + "ab12_tensor", + "c_tensor", + "sfa_tensor", + "sfb_tensor", + "amax_tensor", + "sfc_tensor", + "norm_const_tensor", + "alpha", + "current_stream", + ], + ) + + wrapper = next(node for node in api_tree.body if isinstance(node, ast.FunctionDef) and node.name == "gemm_swiglu_wrapper_sm100") + self.assertEqual( + [argument.arg for argument in wrapper.args.args], + [ + "a_tensor", + "b_tensor", + "alpha", + "c_major", + "ab12_dtype", + "c_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sfa_tensor", + "sfb_tensor", + "norm_const_tensor", + "sf_vec_size", + "vector_f32", + "ab12_stages", + "stream", + ], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_gemm_swiglu_op.py b/test/python/fe_api/test_gemm_swiglu_op.py new file mode 100644 index 000000000..a2b5bb611 --- /dev/null +++ b/test/python/fe_api/test_gemm_swiglu_op.py @@ -0,0 +1,332 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for the dense GEMM + SwiGLU operation.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "gemm_swiglu" +_PACKAGE = "cudnn_gemm_swiglu_op_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + FP8_E4M3 = auto() + FP8_E5M2 = auto() + FP8_E8M0 = auto() + FP4_E2M1 = auto() + + +class GemmSwigluOpContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + operation_name = f"{_PACKAGE}.gemm_swiglu" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + try: + cls.tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}.common.op") + cls.op_module = importlib.import_module(f"{operation_name}.op") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def _desc(self, shape, dtype, order, name): + shape = tuple(shape) + stride = [0] * len(shape) + running = 1 + for dimension in order: + stride[dimension] = running + running *= max(shape[dimension], 1) + return self.tensor_module.TensorDesc( + dtype=dtype, + shape=shape, + stride=tuple(stride), + stride_order=tuple(order), + name=name, + ) + + def _op(self, **overrides): + m, n, k, batch = 128, 128, 128, 2 + arguments = { + "a": self._desc((m, k, batch), _DataType.BFLOAT16, (1, 0, 2), "A"), + "b": self._desc((n, k, batch), _DataType.BFLOAT16, (1, 0, 2), "B"), + "ab12": self._desc((m, n, batch), _DataType.FLOAT, (1, 0, 2), "AB12"), + "c": self._desc((m, n // 2, batch), _DataType.HALF, (1, 0, 2), "C"), + } + arguments.update(overrides) + return self.op_module.GemmSwigluSm100Op(**arguments) + + def _block_op(self, **overrides): + m, n, k, batch = 128, 128, 128, 2 + scale_order = (3, 1, 0, 4, 2, 5) + arguments = { + "a": self._desc((m, k, batch), _DataType.FP8_E4M3, (1, 0, 2), "A"), + "b": self._desc((n, k, batch), _DataType.FP8_E4M3, (1, 0, 2), "B"), + "sfa": self._desc((32, 4, 1, 4, 1, batch), _DataType.FP8_E8M0, scale_order, "SFA"), + "sfb": self._desc((32, 4, 1, 4, 1, batch), _DataType.FP8_E8M0, scale_order, "SFB"), + "ab12": self._desc((m, n, batch), _DataType.BFLOAT16, (1, 0, 2), "AB12"), + "c": self._desc((m, n // 2, batch), _DataType.BFLOAT16, (1, 0, 2), "C"), + "sf_vec_size": 32, + } + arguments.update(overrides) + return self.op_module.BlockScaledGemmSwigluSm100Op(**arguments) + + def test_validates_complete_signature_and_resolves_configuration(self): + operation = self._op(alpha=0.5) + + self.assertIsInstance(operation, self.base_module.Op) + self.assertTrue(operation.check_support()) + self.assertEqual((operation.m, operation.n, operation.k, operation.l), (128, 128, 128, 2)) + self.assertEqual(operation.output_n, 64) + self.assertEqual( + (operation.a_major, operation.b_major, operation.output_major), + ("k", "k", "n"), + ) + self.assertEqual(operation.mma_tiler_mn, (128, 128)) + self.assertEqual(operation.cluster_shape_mn, (1, 1)) + self.assertEqual(operation.alpha, 0.5) + self.assertNotIn(f"{_PACKAGE}.gemm_swiglu.api", sys.modules) + self.assertNotIn(f"{_PACKAGE}.gemm_swiglu.dense_gemm_persistent_swiglu", sys.modules) + + def test_accepts_each_supported_compact_major_mode(self): + m, n, k, batch = 128, 128, 128, 2 + for a_order, b_order, output_order, expected in ( + ((0, 1, 2), (0, 1, 2), (0, 1, 2), ("m", "n", "m")), + ((1, 0, 2), (1, 0, 2), (1, 0, 2), ("k", "k", "n")), + ): + with self.subTest(expected=expected): + operation = self._op( + a=self._desc((m, k, batch), _DataType.BFLOAT16, a_order, "A"), + b=self._desc((n, k, batch), _DataType.BFLOAT16, b_order, "B"), + ab12=self._desc((m, n, batch), _DataType.FLOAT, output_order, "AB12"), + c=self._desc((m, n // 2, batch), _DataType.HALF, output_order, "C"), + ) + self.assertTrue(operation.check_support()) + self.assertEqual( + (operation.a_major, operation.b_major, operation.output_major), + expected, + ) + + def test_rejects_invalid_shapes_and_layouts(self): + cases = ( + ( + {"a": self._desc((128, 128), _DataType.BFLOAT16, (1, 0), "A")}, + "A must have rank 3", + ), + ( + {"b": self._desc((128, 64, 2), _DataType.BFLOAT16, (1, 0, 2), "B")}, + "B shape mismatch", + ), + ( + { + "b": self._desc((127, 128, 2), _DataType.BFLOAT16, (1, 0, 2), "B"), + "ab12": self._desc((128, 127, 2), _DataType.FLOAT, (1, 0, 2), "AB12"), + "c": self._desc((128, 63, 2), _DataType.HALF, (1, 0, 2), "C"), + }, + "N must be even", + ), + ( + { + "b": self._desc((96, 128, 2), _DataType.BFLOAT16, (1, 0, 2), "B"), + "ab12": self._desc((128, 96, 2), _DataType.FLOAT, (1, 0, 2), "AB12"), + "c": self._desc((128, 48, 2), _DataType.HALF, (1, 0, 2), "C"), + }, + "N must be divisible by 64", + ), + ( + {"c": self._desc((128, 128, 2), _DataType.HALF, (1, 0, 2), "C")}, + "C must have shape", + ), + ( + {"c": self._desc((128, 64, 2), _DataType.HALF, (0, 1, 2), "C")}, + "AB12 and C must use the same major mode", + ), + ) + for overrides, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self._op(**overrides).check_support() + + def test_rejects_invalid_standard_dtypes(self): + cases = ( + ( + {"b": self._desc((128, 128, 2), _DataType.HALF, (1, 0, 2), "B")}, + ValueError, + "A and B must have the same dtype", + ), + ( + {"ab12": self._desc((128, 128, 2), _DataType.FP8_E4M3, (1, 0, 2), "AB12")}, + NotImplementedError, + "FP8 AB12 output is currently disabled", + ), + ( + { + "acc_dtype": _DataType.HALF, + "ab12": self._desc((128, 128, 2), _DataType.HALF, (1, 0, 2), "AB12"), + }, + ValueError, + "unsupported for float16 accumulation", + ), + ( + {"c": self._desc((128, 64, 2), _DataType.FLOAT, (1, 0, 2), "C")}, + ValueError, + "C dtype must be float16 or bfloat16", + ), + ) + for overrides, error_type, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(error_type, message): + self._op(**overrides).check_support() + + def test_validates_alignment_tiler_and_cluster(self): + with self.assertRaisesRegex(ValueError, "A contiguous extent must be a multiple"): + self._op( + a=self._desc((128, 124, 2), _DataType.BFLOAT16, (1, 0, 2), "A"), + b=self._desc((128, 124, 2), _DataType.BFLOAT16, (1, 0, 2), "B"), + ).check_support() + + with self.assertRaisesRegex(ValueError, r"mma_tiler_mn\[1\]"): + self._op(mma_tiler_mn=(128, 32)).check_support() + + with self.assertRaisesRegex(ValueError, "must contain two integers"): + self._op(mma_tiler_mn=()).check_support() + + with self.assertRaisesRegex(ValueError, r"must be \(1, 1\)"): + self._op(cluster_shape_mn=(2, 1)).check_support() + + with self.assertRaisesRegex(ValueError, "M must be divisible by 256"): + self._op(mma_tiler_mn=(256, 128)).check_support() + + operation = self._op( + a=self._desc((256, 128, 2), _DataType.BFLOAT16, (1, 0, 2), "A"), + ab12=self._desc((256, 128, 2), _DataType.FLOAT, (1, 0, 2), "AB12"), + c=self._desc((256, 64, 2), _DataType.HALF, (1, 0, 2), "C"), + mma_tiler_mn=(256, 128), + ) + self.assertTrue(operation.check_support()) + self.assertEqual(operation.cluster_shape_mn, (2, 2)) + + def test_constructor_requires_canonical_descriptors_and_dtype(self): + for name in ("a", "b", "ab12", "c"): + with self.subTest(name=name): + with self.assertRaisesRegex(TypeError, f"{name} must be a TensorDesc"): + self._op(**{name: types.SimpleNamespace(shape=(1,))}) + + with self.assertRaisesRegex(TypeError, "acc_dtype must be a cudnn.data_type"): + self._op(acc_dtype="float32") + + def test_validates_block_scaled_mxfp8_signature(self): + operation = self._block_op(alpha=0.25, vector_f32=True, ab12_stages=3) + + self.assertTrue(operation.check_support()) + self.assertEqual((operation.m, operation.n, operation.k, operation.l), (128, 128, 128, 2)) + self.assertEqual(operation.output_n, 64) + self.assertEqual((operation.ab_dtype, operation.sf_dtype), (_DataType.FP8_E4M3, _DataType.FP8_E8M0)) + self.assertEqual((operation.ab12_dtype, operation.c_dtype), (_DataType.BFLOAT16, _DataType.BFLOAT16)) + self.assertEqual(operation.mma_tiler_mn, (128, 128)) + self.assertEqual(operation.cluster_shape_mn, (1, 1)) + self.assertTrue(operation.vector_f32) + self.assertEqual(operation.ab12_stages, 3) + + def test_validates_block_scaled_fp4_amax_signature(self): + m, n, k, batch = 256, 256, 128, 1 + scale_order = (3, 1, 0, 4, 2, 5) + operation = self._block_op( + a=self._desc((m, k, batch), _DataType.FP4_E2M1, (1, 0, 2), "A"), + b=self._desc((n, k, batch), _DataType.FP4_E2M1, (1, 0, 2), "B"), + sfa=self._desc((32, 4, 2, 4, 2, batch), _DataType.FP8_E4M3, scale_order, "SFA"), + sfb=self._desc((32, 4, 2, 4, 2, batch), _DataType.FP8_E4M3, scale_order, "SFB"), + ab12=self._desc((m, n, batch), _DataType.HALF, (1, 0, 2), "AB12"), + c=self._desc((m, n // 2, batch), _DataType.BFLOAT16, (1, 0, 2), "C"), + amax=self._desc((1,), _DataType.FLOAT, (0,), "amax"), + sf_vec_size=16, + mma_tiler_mn=(256, 256), + ) + + self.assertTrue(operation.check_support()) + self.assertEqual(operation.cluster_shape_mn, (2, 2)) + self.assertEqual(operation.output_major, "n") + + def test_rejects_invalid_block_scaled_signatures(self): + scale_order = (3, 1, 0, 4, 2, 5) + cases = ( + ( + {"sfa": self._desc((32, 4, 1, 4, 2, 2), _DataType.FP8_E8M0, scale_order, "SFA")}, + ValueError, + "SFA must have shape", + ), + ( + {"sfa": self._desc((32, 4, 1, 4, 1, 2), _DataType.FP8_E8M0, (5, 4, 3, 2, 1, 0), "SFA")}, + ValueError, + "packed block-scale layout", + ), + ( + { + "sfa": self._desc((32, 4, 1, 4, 1, 2), _DataType.FP8_E4M3, scale_order, "SFA"), + "sfb": self._desc((32, 4, 1, 4, 1, 2), _DataType.FP8_E4M3, scale_order, "SFB"), + }, + ValueError, + "FP8 A and B require FP8_E8M0 scales with sf_vec_size=32", + ), + ( + {"ab12": self._desc((128, 128, 2), _DataType.FLOAT, (1, 0, 2), "AB12")}, + NotImplementedError, + "MXFP8 inputs require", + ), + ( + {"mma_tiler_mn": (128, 192)}, + ValueError, + "M and N must be divisible by mma_tiler_mn", + ), + ) + for overrides, error_type, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(error_type, message): + self._block_op(**overrides).check_support() + + fp4 = _DataType.FP4_E2M1 + with self.assertRaisesRegex(ValueError, "amax is required"): + self._block_op( + a=self._desc((128, 128, 2), fp4, (1, 0, 2), "A"), + b=self._desc((128, 128, 2), fp4, (1, 0, 2), "B"), + ).check_support() + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_api_base.py b/test/python/fe_api/test_jax_api_base.py new file mode 100644 index 000000000..0a9d1f764 --- /dev/null +++ b/test/python/fe_api/test_jax_api_base.py @@ -0,0 +1,700 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for the JAX API base.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_PACKAGE = "cudnn_jax_api_base_test" + + +class _DataType(Enum): + NOT_SET = auto() + FLOAT = auto() + BFLOAT16 = auto() + + +class _ArrayMetadata: + def __init__(self, shape, dtype, label=None): + self.shape = shape + self.dtype = dtype + self.label = label + + +class _TensorSpec: + def __init__( + self, + *, + layout=None, + mode=None, + divisibility=None, + ptr_assumed_align=None, + ): + self.layout = layout + self.mode = mode + self.divisibility = divisibility + self.ptr_assumed_align = ptr_assumed_align + + +class _ShapeDtypeStruct: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + +class _Device: + platform = "gpu" + + def __init__(self, device_id, compute_capability): + self.id = device_id + self.compute_capability = compute_capability + + def __str__(self): + return f"CudaDevice(id={self.id})" + + +class JaxApiBaseTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + datatypes_name = f"{internal_name}.datatypes" + datatypes = types.ModuleType(datatypes_name) + datatypes.jax_to_cudnn_dtype = lambda dtype: { + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, + }.get(dtype, _DataType.NOT_SET) + datatypes.cudnn_to_jax_dtype = lambda dtype: { + _DataType.BFLOAT16: "bfloat16", + _DataType.FLOAT: "float32", + }[dtype] + sys.modules[datatypes_name] = datatypes + + try: + with mock.patch.dict( + sys.modules, + { + "jax": None, + "jax.numpy": None, + "cutlass": None, + "cutlass.jax": None, + }, + ): + cls.tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + cls.module = importlib.import_module(f"{internal_name}.api_base") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def test_checks_all_local_jax_gpu_compute_capabilities(self): + jax = types.ModuleType("jax") + + def local_devices(*, backend): + self.assertEqual(backend, "gpu") + return (_Device(0, "10.0"), _Device(1, "12.0")) + + jax.local_devices = local_devices + with mock.patch.dict(sys.modules, {"jax": jax}): + self.module.JaxApiBase._check_device_compatibility( + minimum_compute_capability=100, + operation_name="TestOp", + ) + + def test_rejects_an_incompatible_local_jax_gpu(self): + jax = types.ModuleType("jax") + jax.local_devices = lambda *, backend: (_Device(0, "10.0"), _Device(1, "9.0")) + + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaisesRegex(RuntimeError, r"TestOp requires SM100\+, found CudaDevice\(id=1\) \(SM90\)"): + self.module.JaxApiBase._check_device_compatibility( + minimum_compute_capability=100, + operation_name="TestOp", + ) + + def test_requires_a_local_jax_gpu(self): + jax = types.ModuleType("jax") + jax.local_devices = lambda *, backend: () + + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaisesRegex(RuntimeError, r"no local JAX GPU"): + self.module.JaxApiBase._check_device_compatibility( + minimum_compute_capability=100, + operation_name="TestOp", + ) + + def test_reports_device_discovery_and_capability_failures(self): + jax = types.ModuleType("jax") + + def fail_discovery(*, backend): + raise RuntimeError("GPU backend unavailable") + + jax.local_devices = fail_discovery + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaisesRegex(RuntimeError, r"could not discover a local GPU"): + self.module.JaxApiBase._check_device_compatibility( + minimum_compute_capability=100, + operation_name="TestOp", + ) + + jax.local_devices = lambda *, backend: (_Device(0, "unknown"),) + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaisesRegex(RuntimeError, r"invalid compute capability 'unknown'"): + self.module.JaxApiBase._check_device_compatibility( + minimum_compute_capability=100, + operation_name="TestOp", + ) + + def test_resolves_exact_local_target_for_multi_arch_operation(self): + jax = types.ModuleType("jax") + jax.local_devices = lambda *, backend: (_Device(0, "10.3"),) + + with mock.patch.dict(sys.modules, {"jax": jax}): + target = self.module.JaxApiBase._resolve_compute_capability( + None, + (90, 100, 103, 107), + "TestOp", + ) + + self.assertEqual(target, 103) + + def test_rejects_heterogeneous_implicit_targets(self): + jax = types.ModuleType("jax") + jax.local_devices = lambda *, backend: (_Device(0, "10.0"), _Device(1, "10.3")) + + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaisesRegex(RuntimeError, r"heterogeneous targets"): + self.module.JaxApiBase._resolve_compute_capability(None, (90, 100, 103, 107), "TestOp") + + def test_explicit_target_must_match_local_devices(self): + jax = types.ModuleType("jax") + jax.local_devices = lambda *, backend: (_Device(0, "10.0"),) + + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaisesRegex(RuntimeError, r"targets SM90.*found CudaDevice\(id=0\) \(SM100\)"): + self.module.JaxApiBase._resolve_compute_capability(90, (90, 100, 103, 107), "TestOp") + self.assertEqual( + self.module.JaxApiBase._resolve_compute_capability(100, (90, 100, 103, 107), "TestOp"), + 100, + ) + + def test_rejects_unknown_explicit_target_before_lowering(self): + with self.assertRaisesRegex(ValueError, r"no kernel for SM101.*supported targets"): + self.module.JaxApiBase._resolve_compute_capability( + 101, + (90, 100, 103, 107), + "TestOp", + ) + + def test_implicit_target_requires_a_local_device(self): + jax = types.ModuleType("jax") + jax.local_devices = lambda *, backend: () + + with mock.patch.dict(sys.modules, {"jax": jax}): + with self.assertRaisesRegex(RuntimeError, r"no local JAX GPU"): + self.module.JaxApiBase._resolve_compute_capability(None, (90, 100), "TestOp") + + def test_caches_raw_max_active_clusters_per_instance_and_cluster_size(self): + query_calls = [] + + class HardwareInfo: + def get_max_active_clusters(self, cluster_size): + query_calls.append(cluster_size) + return {2: 8, 4: 12}[cluster_size] + + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass_utils = types.ModuleType("cutlass.utils") + cutlass_utils.HardwareInfo = HardwareInfo + cutlass.utils = cutlass_utils + + class Adapter(self.module.JaxApiBase): + def check_support(self): + return True + + def __call__(self): + return self.check_support() + + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + "cutlass.utils": cutlass_utils, + }, + ): + api = Adapter() + self.assertEqual(api._get_max_active_clusters(4, overlap_margin=2), 10) + self.assertEqual(api._get_max_active_clusters(4, overlap_margin=3), 9) + self.assertEqual(api._get_max_active_clusters(2, overlap_margin=1), 7) + with self.assertRaisesRegex(ValueError, "max_active_clusters must be positive"): + api._get_max_active_clusters(4, overlap_margin=12) + + other_api = Adapter() + self.assertEqual(other_api._get_max_active_clusters(4), 12) + + self.assertEqual(query_calls, [4, 2, 4]) + + def test_converts_array_metadata_to_shared_tensor_desc(self): + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3, 4), "bfloat16"), + "sample", + ) + + self.assertIsInstance(desc, self.tensor_module.TensorDesc) + self.assertEqual(desc.dtype, "bfloat16") + self.assertEqual(desc.shape, (2, 3, 4)) + self.assertEqual(desc.stride, (12, 4, 1)) + self.assertEqual(desc.stride_order, (2, 1, 0)) + self.assertEqual(desc.name, "sample") + self.assertIsNone(desc.init_value) + self.assertEqual(desc.cudnn_dtype, _DataType.BFLOAT16) + + def test_converts_public_array_axes_to_canonical_descriptor_axes(self): + mode = (2, 0, 1) + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3, 4), "bfloat16"), + "sample", + mode=mode, + ) + + self.assertEqual(desc.shape, (4, 2, 3)) + self.assertEqual(desc.stride, (1, 12, 4)) + self.assertEqual(desc.stride_order, (0, 2, 1)) + self.assertEqual(desc.mode, mode) + self.assertEqual(desc.array_shape, (2, 3, 4)) + + def test_constructs_descriptor_directly_from_shape(self): + desc = self.module.JaxTensorDesc.from_shape( + (2, 3, 4), + "bfloat16", + name="sample", + mode=(2, 0, 1), + ) + + self.assertEqual(desc.shape, (4, 2, 3)) + self.assertEqual(desc.array_shape, (2, 3, 4)) + self.assertEqual(desc.mode, (2, 0, 1)) + self.assertEqual(desc.cudnn_dtype, _DataType.BFLOAT16) + + def test_converts_public_physical_layout_to_canonical_descriptor_axes(self): + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3, 4), "bfloat16"), + "sample", + mode=(1, 2, 0), + public_stride_order=(0, 2, 1), + ) + + self.assertEqual(desc.shape, (3, 4, 2)) + self.assertEqual(desc.stride, (8, 2, 1)) + self.assertEqual(desc.stride_order, (2, 1, 0)) + + def test_rejects_invalid_public_stride_orders(self): + value = _ArrayMetadata((2, 3, 4), "bfloat16") + + with self.assertRaisesRegex(ValueError, "public_stride_order rank mismatch"): + self.module.JaxApiBase._to_tensor_desc(value, "sample", public_stride_order=(1, 0)) + with self.assertRaisesRegex(ValueError, "public_stride_order must be a permutation"): + self.module.JaxApiBase._to_tensor_desc(value, "sample", public_stride_order=(2, 2, 0)) + with self.assertRaisesRegex(TypeError, "public_stride_order entries must be integers"): + self.module.JaxApiBase._to_tensor_desc(value, "sample", public_stride_order=(2, True, 0)) + + def test_jax_descriptor_derives_compact_output_metadata(self): + source = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3), "bfloat16"), + "input", + ) + + output = source.compact_like( + cudnn_dtype=_DataType.FLOAT, + shape=(5, 7), + stride_order=(0, 1), + name="output", + init_value=-2.0, + ) + + self.assertIsInstance(output, self.module.JaxTensorDesc) + self.assertEqual(output.dtype, "float32") + self.assertEqual(output.cudnn_dtype, _DataType.FLOAT) + self.assertEqual(output.shape, (5, 7)) + self.assertEqual(output.stride, (1, 5)) + self.assertEqual(output.stride_order, (0, 1)) + self.assertEqual(output.name, "output") + self.assertEqual(output.init_value, -2.0) + self.assertEqual(output.mode, (0, 1)) + + def test_checks_invocation_signature(self): + desc = self.module.JaxApiBase._to_tensor_desc(_ArrayMetadata((2, 3), "bfloat16"), "sample") + self.module.JaxApiBase._check_tensor_signature(_ArrayMetadata((2, 3), "bfloat16"), desc) + + with self.assertRaisesRegex(ValueError, "sample tensor shape mismatch"): + self.module.JaxApiBase._check_tensor_signature(_ArrayMetadata((1, 3), "bfloat16"), desc) + with self.assertRaisesRegex(ValueError, "sample tensor dtype mismatch"): + self.module.JaxApiBase._check_tensor_signature(_ArrayMetadata((2, 3), "float32"), desc) + + def test_checks_public_invocation_against_canonical_descriptor(self): + mode = (2, 0, 1) + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3, 4), "bfloat16"), + "sample", + mode=mode, + ) + + self.module.JaxApiBase._check_tensor_signature( + _ArrayMetadata((2, 3, 4), "bfloat16"), + desc, + ) + with self.assertRaisesRegex(ValueError, "sample tensor shape mismatch"): + self.module.JaxApiBase._check_tensor_signature( + _ArrayMetadata((1, 3, 4), "bfloat16"), + desc, + ) + + def test_builds_public_tensor_spec_from_canonical_metadata(self): + mode = (2, 0, 1) + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3, 4), "bfloat16"), + "sample", + mode=mode, + divisibility=(8, 2, 4), + ptr_assumed_align=64, + ) + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass_jax = types.ModuleType("cutlass.jax") + cutlass_jax.TensorSpec = _TensorSpec + cutlass.jax = cutlass_jax + + with mock.patch.dict(sys.modules, {"cutlass": cutlass, "cutlass.jax": cutlass_jax}): + spec = self.module.JaxApiBase._to_tensor_spec(desc) + + self.assertEqual(spec.layout, (2, 1, 0)) + self.assertEqual(spec.mode, mode) + self.assertEqual(spec.divisibility, (2, 4, 8)) + self.assertEqual(spec.ptr_assumed_align, 64) + + def test_tensor_spec_preserves_explicit_public_physical_layout(self): + mode = (1, 2, 0) + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3, 4), "bfloat16"), + "sample", + mode=mode, + public_stride_order=(0, 2, 1), + ) + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass_jax = types.ModuleType("cutlass.jax") + cutlass_jax.TensorSpec = _TensorSpec + cutlass.jax = cutlass_jax + + with mock.patch.dict(sys.modules, {"cutlass": cutlass, "cutlass.jax": cutlass_jax}): + spec = self.module.JaxApiBase._to_tensor_spec(desc) + + self.assertEqual(spec.layout, (0, 2, 1)) + self.assertEqual(spec.mode, mode) + + def test_default_descriptor_uses_cutlass_default_tensor_spec(self): + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3), "bfloat16"), + "sample", + ) + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass_jax = types.ModuleType("cutlass.jax") + cutlass_jax.TensorSpec = _TensorSpec + cutlass.jax = cutlass_jax + + with mock.patch.dict(sys.modules, {"cutlass": cutlass, "cutlass.jax": cutlass_jax}): + self.assertIsNone(self.module.JaxApiBase._to_tensor_spec(desc)) + + def test_tensor_spec_rejects_divisibility_with_the_wrong_rank(self): + desc = self.module.JaxApiBase._to_tensor_desc( + _ArrayMetadata((2, 3), "bfloat16"), + "sample", + ) + + with self.assertRaisesRegex(ValueError, "divisibility rank mismatch"): + self.module.JaxTensorDesc.from_array( + _ArrayMetadata((2, 3), "bfloat16"), + name="sample", + divisibility=(2,), + ) + + def test_builds_shape_dtype_struct_in_array_axis_order(self): + desc = self.module.JaxTensorDesc.from_shape( + (2, 3, 4), + "bfloat16", + mode=(2, 0, 1), + ) + jax = types.ModuleType("jax") + jax.ShapeDtypeStruct = _ShapeDtypeStruct + + with mock.patch.dict(sys.modules, {"jax": jax}): + output = self.module.JaxApiBase._to_shape_dtype_struct(desc) + + self.assertEqual(output.shape, (2, 3, 4)) + self.assertEqual(output.dtype, "bfloat16") + + def test_calls_explicit_launcher_with_initialized_and_fresh_outputs(self): + seen = {} + + output_descs = ( + self.module.JaxTensorDesc.from_shape( + (2,), + "float32", + name="output", + ), + self.module.JaxTensorDesc.from_shape( + (2, 3, 4), + "float32", + name="amax", + init_value=float("-inf"), + mode=(2, 0, 1), + ), + ) + workspace_descs = ( + self.module.JaxTensorDesc.from_shape( + (3,), + "float32", + name="counter", + init_value=0, + ), + self.module.JaxTensorDesc.from_shape( + (5,), + "bfloat16", + name="scratch", + ), + ) + + def launch(stream, *buffers): + seen["launch_args"] = (stream, *buffers) + + class Adapter(self.module.JaxApiBase): + def check_support(self): + return True + + def __call__(self, value): + return self._call_kernel( + (value,), + launch=launch, + input_descs=(input_desc,), + output_descs=output_descs, + workspace_descs=workspace_descs, + ) + + full_calls = [] + jax = types.ModuleType("jax") + jax.__path__ = [] + jax.ShapeDtypeStruct = _ShapeDtypeStruct + jnp = types.ModuleType("jax.numpy") + + def full(shape, value, dtype): + result = _ArrayMetadata(tuple(shape), dtype, label=f"full({value})") + full_calls.append((tuple(shape), value, dtype, result)) + return result + + jnp.full = full + jax.numpy = jnp + + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass_cute = types.ModuleType("cutlass.cute") + + def cute_jit(*, preprocess): + self.assertFalse(preprocess) + + def decorate(function): + seen["cute_launcher"] = function + return function + + return decorate + + cutlass_cute.jit = cute_jit + cutlass_jax = types.ModuleType("cutlass.jax") + cutlass_jax.TensorSpec = _TensorSpec + + def cutlass_call(launcher, **options): + seen["cutlass_launcher"] = launcher + seen["call_options"] = options + + def invoke(*args): + aliases = options["input_output_aliases"] + result_metadata = options["output_shape_dtype"] + results = [None] * len(result_metadata) + for input_index, result_index in aliases.items(): + results[result_index] = args[input_index] + for result_index, metadata in enumerate(result_metadata): + if results[result_index] is None: + results[result_index] = _ArrayMetadata( + metadata.shape, + metadata.dtype, + label=f"allocated({result_index})", + ) + + unaliased_results = [results[index] for index in range(len(results)) if index not in aliases.values()] + launcher("stream", *args, *unaliased_results) + return tuple(results) + + return invoke + + cutlass_jax.cutlass_call = cutlass_call + cutlass.cute = cutlass_cute + cutlass.jax = cutlass_jax + + input_value = _ArrayMetadata((2,), "float32", label="input") + input_desc = self.module.JaxTensorDesc.from_array( + input_value, + name="input", + ) + api = Adapter() + with mock.patch.dict( + sys.modules, + { + "jax": jax, + "jax.numpy": jnp, + "cutlass": cutlass, + "cutlass.cute": cutlass_cute, + "cutlass.jax": cutlass_jax, + }, + ): + results = api._call_kernel( + (input_value,), + launch=launch, + input_descs=(input_desc,), + output_descs=output_descs, + workspace_descs=workspace_descs, + ) + with self.assertRaisesRegex(TypeError, "input #0 must have shape and dtype metadata"): + api._call_kernel( + ([input_value],), + launch=launch, + input_descs=(input_desc,), + output_descs=output_descs, + workspace_descs=workspace_descs, + ) + with self.assertRaisesRegex(ValueError, "Expected 1 input descriptors, got 0"): + api._call_kernel( + (input_value,), + launch=launch, + input_descs=(), + output_descs=output_descs, + ) + with self.assertRaisesRegex(TypeError, "input_descs must contain JaxTensorDesc"): + api._call_kernel( + (input_value,), + launch=launch, + input_descs=( + self.tensor_module.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2,), + ), + ), + output_descs=output_descs, + ) + with self.assertRaisesRegex(TypeError, "output_descs must contain JaxTensorDesc"): + api._call_kernel( + (input_value,), + launch=launch, + input_descs=(input_desc,), + output_descs=( + self.tensor_module.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2,), + ), + ), + workspace_descs=workspace_descs, + ) + with self.assertRaisesRegex(TypeError, "launch must be callable"): + api._call_kernel( + (input_value,), + launch=None, + input_descs=(input_desc,), + output_descs=output_descs, + ) + + self.assertEqual( + [(shape, value, dtype) for shape, value, dtype, _ in full_calls], + [ + ((2, 3, 4), float("-inf"), "float32"), + ((3,), 0, "float32"), + ], + ) + self.assertEqual( + seen["call_options"]["input_output_aliases"], + {1: 1, 2: 2}, + ) + self.assertEqual(len(seen["call_options"]["input_spec"]), 3) + self.assertIsNone(seen["call_options"]["input_spec"][0]) + self.assertIs( + seen["call_options"]["input_spec"][1], + seen["call_options"]["output_spec"][1], + ) + self.assertIs( + seen["call_options"]["input_spec"][2], + seen["call_options"]["output_spec"][2], + ) + + stream, launch_input, output, amax, counter, scratch = seen["launch_args"] + self.assertEqual(stream, "stream") + self.assertIs(seen["cute_launcher"], seen["cutlass_launcher"]) + self.assertIs(launch_input, input_value) + self.assertEqual(output.label, "allocated(0)") + self.assertIs(amax, full_calls[0][3]) + self.assertIs(counter, full_calls[1][3]) + self.assertEqual(scratch.label, "allocated(3)") + self.assertEqual([result.label for result in results], ["allocated(0)", "full(-inf)"]) + self.assertNotIn(full_calls[0][3], vars(api).values()) + self.assertNotIn(full_calls[1][3], vars(api).values()) + + def test_returns_callable_without_owning_jit_policy(self): + base = self.module.JaxApiBase + + class Adapter(base): + def check_support(self): + return True + + def __call__(self): + return self.check_support() + + api = Adapter() + api.option = 1 + self.assertIs(api.get_jax_callable(), api) + self.assertTrue(api()) + api.option = 2 + self.assertEqual(api.option, 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_bsa_contract.py b/test/python/fe_api/test_jax_bsa_contract.py new file mode 100644 index 000000000..c648bf2ec --- /dev/null +++ b/test/python/fe_api/test_jax_bsa_contract.py @@ -0,0 +1,442 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for the JAX block-sparse attention adapters.""" + +from __future__ import annotations + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_BSA_ROOT = _CUDNN_ROOT / "block_sparse_attention" +_PACKAGE = "cudnn_jax_bsa_contract_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + + +_DTYPE_TO_CUDNN = { + "float16": _DataType.HALF, + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, + "int32": _DataType.INT32, +} +_CUDNN_TO_DTYPE = {value: key for key, value in _DTYPE_TO_CUDNN.items()} + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + +class _TensorSpec: + def __init__(self, *, layout, mode, divisibility=None): + self.layout = layout + self.mode = mode + self.divisibility = divisibility + + +class JaxBsaContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + bsa_name = f"{_PACKAGE}.block_sparse_attention" + bsa = types.ModuleType(bsa_name) + bsa.__path__ = [str(_BSA_ROOT)] + bsa.__package__ = bsa_name + bsa.__spec__ = ModuleSpec(bsa_name, loader=None, is_package=True) + sys.modules[bsa_name] = bsa + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + layout_module = importlib.import_module(f"{internal_name}.layout") + result_module = importlib.import_module(f"{_PACKAGE}.common.result") + + class JaxTensorDesc(tensor_module.TensorDesc): + @classmethod + def from_shape( + cls, + shape, + dtype, + *, + name="", + mode=None, + public_stride_order=None, + init_value=None, + ): + return JaxApiBase._to_tensor_desc( + _Array(shape, dtype), + name, + mode=mode, + public_stride_order=public_stride_order, + init_value=init_value, + ) + + @property + def cudnn_dtype(self): + return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) + + def compact_like( + self, + *, + cudnn_dtype, + shape, + stride_order=None, + name="", + init_value=None, + mode=None, + ): + if stride_order is None: + stride_order = tuple(reversed(range(len(shape)))) + stride = layout_module.compact_stride(tuple(shape), tuple(stride_order)) + desc = JaxTensorDesc( + dtype=_CUDNN_TO_DTYPE[cudnn_dtype], + shape=tuple(shape), + stride=stride, + stride_order=tuple(stride_order), + name=name, + init_value=init_value, + ) + object.__setattr__( + desc, + "mode", + layout_module.normalize_mode(len(shape), mode), + ) + return desc + + class JaxApiBase: + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del operation_name + resolved = 100 if target is None else target + if resolved not in supported: + raise ValueError(f"unsupported target {resolved}") + return resolved + + @staticmethod + def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): + public_shape = tuple(value.shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + public_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride(public_shape, public_order) + canonical_axis_by_public_axis = layout_module.to_public_axes( + tuple(range(len(public_shape))), mode + ) + desc = JaxTensorDesc( + dtype=value.dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple( + canonical_axis_by_public_axis[axis] for axis in public_order + ), + name=name, + init_value=init_value, + ) + object.__setattr__(desc, "mode", mode) + return desc + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + if mode is None: + mode = expected.mode + actual_shape = layout_module.to_canonical_axes(tuple(value.shape), mode) + if actual_shape != expected.shape: + raise ValueError(f"{expected.name} shape mismatch") + if ( + _DTYPE_TO_CUDNN.get(value.dtype, _DataType.NOT_SET) + != expected.cudnn_dtype + ): + raise ValueError(f"{expected.name} dtype mismatch") + + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + if mode is None: + mode = desc.mode + return _TensorSpec( + layout=layout_module.to_cutlass_layout( + desc.shape, + desc.stride, + desc.stride_order, + mode=mode, + name=desc.name, + ), + mode=mode, + divisibility=divisibility, + ) + + def _call_kernel( + self, + inputs, + *, + output_descs, + input_descs=None, + output_spec=None, + **options, + ): + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc) + if output_spec is None: + output_spec = tuple( + self._to_tensor_spec(desc) for desc in output_descs + ) + self.captured_call = { + "inputs": tuple(inputs), + "output_descs": tuple(output_descs), + "output_spec": tuple(output_spec), + **options, + } + return tuple( + _Array( + layout_module.to_public_axes(desc.shape, spec.mode), + _CUDNN_TO_DTYPE[desc.cudnn_dtype], + ) + for desc, spec in zip(output_descs, output_spec) + ) + + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = result_module.TupleDict + + fake_jax = types.ModuleType("jax") + fake_jax.__path__ = [] + fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + fake_jax.ShapeDtypeStruct = _Array + cls.static_argnames = {} + + def jit(function=None, *, static_argnames=()): + def decorate(target): + cls.static_argnames[target.__name__] = tuple(static_argnames) + return target + + return decorate if function is None else decorate(function) + + fake_jax.jit = jit + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.float32 = "float32" + fake_jnp.int32 = "int32" + fake_jnp.iinfo = lambda _dtype: types.SimpleNamespace(max=2**31 - 1) + fake_jax.numpy = fake_jnp + + try: + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "torch": None, + "cutlass": None, + }, + ): + cls.module = importlib.import_module(f"{bsa_name}.jax") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + @staticmethod + def _forward_samples(layout="bhsd"): + q_shape = (2, 4, 256, 128) if layout == "bhsd" else (2, 256, 4, 128) + kv_shape = (2, 2, 512, 128) if layout == "bhsd" else (2, 512, 2, 128) + return ( + _Array(q_shape, "bfloat16"), + _Array(kv_shape, "bfloat16"), + _Array(kv_shape, "bfloat16"), + _Array((2, 2, 4, 8), "int32"), + ) + + def test_forward_infers_public_output_layout_and_cutlass_buffers(self): + q, k, v, block_index = self._forward_samples("bshd") + api = self.module.BlockSparseAttentionForward( + q, + k, + v, + block_index, + block_sparse_num=8, + layout="bshd", + target_compute_capability=100, + ) + result = api(q, k, v, block_index) + + self.assertEqual(api.data_mode, (0, 2, 1, 3)) + self.assertEqual(api.q_desc.shape, (2, 4, 256, 128)) + self.assertTrue(api._op.pack_gqa_effective) + self.assertEqual(tuple(result.keys()), ("o_tensor", "lse_tensor")) + self.assertEqual(result["o_tensor"].shape, (2, 256, 4, 128)) + self.assertEqual(result["lse_tensor"].shape, (2, 4, 256)) + self.assertEqual(api.captured_call["workspace_descs"], ()) + self.assertIn("--gpu-arch sm_100a", api.captured_call["compile_options"]) + + def test_backward_declares_csr_and_accumulator_workspaces(self): + q = _Array((2, 4, 256, 128), "bfloat16") + k = _Array((2, 4, 512, 128), "bfloat16") + lse = _Array((2, 4, 256), "float32") + block_index = _Array((2, 4, 2, 8), "int32") + api = self.module.BlockSparseAttentionBackward( + q, + q, + k, + k, + q, + lse, + block_index, + block_sparse_num=8, + target_compute_capability=100, + ) + result = api(q, q, k, k, q, lse, block_index) + + self.assertEqual(tuple(result.keys()), ("dq_tensor", "dk_tensor", "dv_tensor")) + self.assertEqual(result["dq_tensor"].shape, q.shape) + self.assertEqual(result["dk_tensor"].shape, k.shape) + workspace_names = tuple( + desc.name for desc in api.captured_call["workspace_descs"] + ) + self.assertIn("counts_workspace", workspace_names) + self.assertIn("bucket_offsets_workspace", workspace_names) + self.assertIn("dq_accum_workspace", workspace_names) + counts = next( + desc + for desc in api.captured_call["workspace_descs"] + if desc.name == "counts_workspace" + ) + self.assertEqual(counts.init_value, 0) + + def test_invalid_metadata_and_unsupported_paths_fail_before_lowering(self): + q, k, v, _ = self._forward_samples() + wrong_index = _Array((2, 4, 3, 8), "int32") + with self.assertRaisesRegex(ValueError, "block_index shape prefix"): + self.module.BlockSparseAttentionForward( + q, + k, + v, + wrong_index, + block_sparse_num=8, + pack_gqa=False, + target_compute_capability=100, + ) + + with self.assertRaisesRegex(ValueError, "sparse_block_size=64"): + self.module.BlockSparseAttentionForward( + q, + k, + v, + _Array((2, 4, 4, 8), "int32"), + block_sparse_num=8, + sparse_block_size=128, + target_compute_capability=90, + ) + + def test_split_kv_disables_the_clc_scheduler(self): + resolve = self.module._resolve_sm100_blk64_use_clc + + self.assertFalse( + resolve( + kv_splits=2, + requested=None, + batch=1, + heads=4, + seqlen_q=256, + block_sparse_num=8, + has_variable_block_nums=True, + ) + ) + + q = _Array((2, 4, 256, 128), "bfloat16") + k = _Array((2, 4, 512, 128), "bfloat16") + v = _Array((2, 4, 512, 128), "bfloat16") + block_index = _Array((2, 4, 4, 8), "int32") + with self.assertRaisesRegex(ValueError, r"\[1, 256\]"): + self.module.BlockSparseAttentionForward( + q, + k, + v, + block_index, + block_sparse_num=8, + kv_splits=257, + sparse_block_size=64, + target_compute_capability=100, + ) + + q, k, v, block_index = self._forward_samples() + with self.assertRaisesRegex(ValueError, "SM100-family blk64"): + self.module.BlockSparseAttentionForward( + q, + k, + v, + block_index, + block_sparse_num=8, + use_clc=False, + sparse_block_size=128, + target_compute_capability=100, + ) + + def test_backward_auto_bucket_matches_kernel_tuning(self): + choose = self.module._default_bucket_size + + self.assertEqual(choose(100, 64, 2048, 4), 1152) + self.assertEqual(choose(100, 64, 3000, 4), 1024) + self.assertEqual(choose(100, 128, 4096, 1), 256) + + def test_forward_auto_kv_splits_uses_static_metadata(self): + choose = self.module._sm100_blk64_auto_kv_splits + + self.assertEqual(choose(255), 1) + self.assertEqual(choose(256), 2) + self.assertEqual(choose(450), 4) + self.assertEqual(choose(900), 8) + + def test_function_wrappers_make_configuration_static(self): + self.assertIn( + "block_sparse_num", self.static_argnames["block_sparse_attention_forward"] + ) + self.assertIn("layout", self.static_argnames["block_sparse_attention_forward"]) + self.assertIn( + "bucket_size_blocks", + self.static_argnames["block_sparse_attention_backward"], + ) + self.assertIn( + "target_compute_capability", + self.static_argnames["block_sparse_attention_backward"], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_compiler.py b/test/python/fe_api/test_jax_compiler.py new file mode 100644 index 000000000..7ba5efc0d --- /dev/null +++ b/test/python/fe_api/test_jax_compiler.py @@ -0,0 +1,74 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for CUTLASS JAX compilation targets.""" + +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_PACKAGE = "cudnn_jax_compiler_test" + + +class JaxCompilerTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + sys.modules[_PACKAGE] = root + + jax_package_name = f"{_PACKAGE}._jax" + jax_package = types.ModuleType(jax_package_name) + jax_package.__path__ = [str(_CUDNN_ROOT / "_jax")] + jax_package.__package__ = jax_package_name + jax_package.__spec__ = ModuleSpec(jax_package_name, loader=None, is_package=True) + sys.modules[jax_package_name] = jax_package + + cls.arch = importlib.import_module(f"{_PACKAGE}.common.cute_arch") + cls.compiler = importlib.import_module(f"{jax_package_name}.compiler") + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def test_explicit_targets_do_not_import_frameworks(self): + torch_before = sys.modules.get("torch") + jax_before = sys.modules.get("jax") + self.assertEqual(self.arch.gpu_arch_flag_for_compute_capability(90), "sm_90a") + self.assertEqual(self.arch.gpu_arch_flag_for_compute_capability(100), "sm_100a") + self.assertEqual(self.arch.gpu_arch_flag_for_compute_capability(103), "sm_103a") + self.assertEqual(self.arch.gpu_arch_flag_for_compute_capability(107), "sm_100f") + self.assertIs(sys.modules.get("torch"), torch_before) + self.assertIs(sys.modules.get("jax"), jax_before) + + def test_builds_jax_compile_options_from_explicit_target(self): + self.assertEqual( + self.compiler.compile_options_for_target(103, "--opt-level 3"), + "--gpu-arch sm_103a --opt-level 3", + ) + + def test_rejects_unknown_or_non_integer_targets(self): + with self.assertRaisesRegex(RuntimeError, "Unsupported GPU compute capability SM101"): + self.compiler.compile_options_for_target(101) + with self.assertRaisesRegex(TypeError, "must be an int"): + self.compiler.compile_options_for_target(True) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_datatypes.py b/test/python/fe_api/test_jax_datatypes.py new file mode 100644 index 000000000..d6f554274 --- /dev/null +++ b/test/python/fe_api/test_jax_datatypes.py @@ -0,0 +1,92 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Tests for common JAX/cuDNN dtype conversion.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_PACKAGE = "cudnn_jax_datatypes_test" + + +class _DataType(Enum): + NOT_SET = auto() + FLOAT = auto() + BFLOAT16 = auto() + + +class JaxDataTypesTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + jax = types.ModuleType("jax") + jax.__path__ = [] + jnp = types.ModuleType("jax.numpy") + jnp.bfloat16 = "bfloat16" + jnp.float32 = "float32" + jnp.dtype = lambda value: value + jax.numpy = jnp + + try: + with mock.patch.dict(sys.modules, {"jax": jax, "jax.numpy": jnp}): + cls.module = importlib.import_module(f"{internal_name}.datatypes") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def test_supported_types_round_trip(self): + self.assertEqual(self.module.jax_to_cudnn_dtype("bfloat16"), _DataType.BFLOAT16) + self.assertEqual(self.module.jax_to_cudnn_dtype("float32"), _DataType.FLOAT) + self.assertEqual(self.module.cudnn_to_jax_dtype(_DataType.BFLOAT16), "bfloat16") + self.assertEqual(self.module.cudnn_to_jax_dtype(_DataType.FLOAT), "float32") + + def test_normalizes_dtype_values_and_defaults(self): + self.assertEqual(self.module.normalize_jax_dtype("bfloat16", "float32", "output_dtype"), "bfloat16") + self.assertEqual(self.module.normalize_jax_dtype(None, "float32", "output_dtype"), "float32") + + with mock.patch.object(self.module.jnp, "dtype", side_effect=TypeError("invalid dtype")): + with self.assertRaisesRegex(TypeError, "output_dtype must be a JAX dtype, got 'invalid'"): + self.module.normalize_jax_dtype("invalid", "float32", "output_dtype") + + def test_unsupported_types_have_explicit_behavior(self): + self.assertEqual(self.module.jax_to_cudnn_dtype("unsupported"), _DataType.NOT_SET) + with self.assertRaisesRegex(ValueError, "Unsupported JAX data type"): + self.module.cudnn_to_jax_dtype(_DataType.NOT_SET) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_discrete_grouped_gemm_contract.py b/test/python/fe_api/test_jax_discrete_grouped_gemm_contract.py new file mode 100644 index 000000000..1c344c63b --- /dev/null +++ b/test/python/fe_api/test_jax_discrete_grouped_gemm_contract.py @@ -0,0 +1,517 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for the discrete grouped GEMM JAX adapters.""" + +from __future__ import annotations + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_DISCRETE_ROOT = _CUDNN_ROOT / "discrete_grouped_gemm" +_PACKAGE = "cudnn_jax_discrete_grouped_gemm_contract_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + INT64 = auto() + UINT8 = auto() + FP8_E4M3 = auto() + FP8_E5M2 = auto() + FP8_E8M0 = auto() + FP4_E2M1 = auto() + + +_DTYPE_TO_CUDNN = { + "float16": _DataType.HALF, + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, + "int32": _DataType.INT32, + "int64": _DataType.INT64, + "uint8": _DataType.UINT8, + "float4_e2m1fn": _DataType.FP4_E2M1, + "float8_e4m3fn": _DataType.FP8_E4M3, + "float8_e5m2": _DataType.FP8_E5M2, + "float8_e8m0fnu": _DataType.FP8_E8M0, +} + + +class _Array: + def __init__(self, shape, dtype, fill_value=None): + self.shape = tuple(shape) + self.dtype = dtype + self.fill_value = fill_value + self.iterator = object() + + +class _TensorSpec: + def __init__(self, mode=None, ptr_assumed_align=None): + self.mode = mode + self.ptr_assumed_align = ptr_assumed_align + + +class JaxDiscreteGroupedGemmContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + discrete_name = f"{_PACKAGE}.discrete_grouped_gemm" + discrete = types.ModuleType(discrete_name) + discrete.__path__ = [str(_DISCRETE_ROOT)] + discrete.__package__ = discrete_name + discrete.__spec__ = ModuleSpec(discrete_name, loader=None, is_package=True) + sys.modules[discrete_name] = discrete + + for child in ("discrete_grouped_gemm_swiglu", "discrete_grouped_gemm_dswiglu"): + child_name = f"{discrete_name}.{child}" + child_module = types.ModuleType(child_name) + child_module.__path__ = [str(_DISCRETE_ROOT / child)] + child_module.__package__ = child_name + child_module.__spec__ = ModuleSpec(child_name, loader=None, is_package=True) + sys.modules[child_name] = child_module + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + layout_module = importlib.import_module(f"{internal_name}.layout") + result_module = importlib.import_module(f"{_PACKAGE}.common.result") + + class JaxTensorDesc(tensor_module.TensorDesc): + @classmethod + def from_shape( + cls, + shape, + dtype, + *, + name="", + mode=None, + public_stride_order=None, + init_value=None, + divisibility=None, + ptr_assumed_align=None, + ): + return JaxApiBase._to_tensor_desc( + _Array(shape, dtype), + name, + mode=mode, + public_stride_order=public_stride_order, + init_value=init_value, + divisibility=divisibility, + ptr_assumed_align=ptr_assumed_align, + ) + + @property + def cudnn_dtype(self): + return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) + + class JaxApiBase: + @staticmethod + def _resolve_compute_capability(_target, _supported, _name): + return 100 + + @staticmethod + def _to_tensor_desc( + value, + name, + *, + mode=None, + init_value=None, + divisibility=None, + ptr_assumed_align=None, + **_unused, + ): + public_shape = tuple(value.shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + public_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride(public_shape, public_order) + canonical_axis_by_public_axis = layout_module.to_public_axes( + tuple(range(len(public_shape))), mode + ) + desc = JaxTensorDesc( + dtype=value.dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple( + canonical_axis_by_public_axis[axis] for axis in public_order + ), + name=name, + init_value=init_value, + ) + object.__setattr__(desc, "mode", mode) + object.__setattr__(desc, "divisibility", divisibility) + object.__setattr__(desc, "ptr_assumed_align", ptr_assumed_align) + return desc + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + if mode is None: + mode = expected.mode + actual = layout_module.to_canonical_axes(tuple(value.shape), mode) + if actual != expected.shape: + raise ValueError( + f"{expected.name} shape mismatch: expected {expected.shape}, got {actual}" + ) + if ( + _DTYPE_TO_CUDNN.get(value.dtype, _DataType.NOT_SET) + != expected.cudnn_dtype + ): + raise ValueError(f"{expected.name} dtype mismatch") + + @staticmethod + def _to_tensor_spec(_desc, *, mode=None, **_unused): + if mode is None: + mode = _desc.mode + return _TensorSpec( + mode, + ptr_assumed_align=getattr( + _desc, "ptr_assumed_align", None + ), + ) + + @staticmethod + def _to_shape_dtype_struct(desc, *, mode=None): + if mode is None: + mode = desc.mode + return _Array( + layout_module.to_public_axes(desc.shape, mode), desc.dtype + ) + + @staticmethod + def _get_max_active_clusters(_cluster_size, *, overlap_margin=0): + return 8 - overlap_margin + + def _call_kernel( + self, + inputs, + *, + launch, + output_descs, + input_descs=None, + workspace_descs=(), + output_spec=None, + workspace_spec=None, + **options, + ): + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc) + if output_spec is None: + output_spec = tuple( + self._to_tensor_spec(desc) for desc in output_descs + ) + if workspace_spec is None: + workspace_spec = tuple( + self._to_tensor_spec(desc) for desc in workspace_descs + ) + outputs = tuple( + _Array( + layout_module.to_public_axes(desc.shape, spec.mode), desc.dtype + ) + for desc, spec in zip(output_descs, output_spec) + ) + workspaces = tuple( + _Array(desc.shape, desc.dtype) for desc in workspace_descs + ) + launch("stream", *inputs, *outputs, *workspaces) + self.captured_kernel_call = ( + tuple(inputs), + tuple(output_descs), + tuple(workspace_descs), + options, + ) + self.captured_workspaces = workspaces + self.captured_workspace_specs = workspace_spec + return outputs + + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = result_module.TupleDict + + datatypes = types.ModuleType(f"{internal_name}.datatypes") + datatypes.jax_to_cudnn_dtype = lambda dtype: _DTYPE_TO_CUDNN.get( + dtype, _DataType.NOT_SET + ) + datatypes.normalize_jax_dtype = lambda value, default, _name: ( + default if value is None else value + ) + sys.modules[datatypes.__name__] = datatypes + + fake_jax = types.ModuleType("jax") + fake_jax.ShapeDtypeStruct = _Array + fake_jax.jit = lambda function=None, **_kwargs: ( + (lambda target: target) if function is None else function + ) + fake_jnp = types.ModuleType("jax.numpy") + for dtype in _DTYPE_TO_CUDNN: + setattr(fake_jnp, dtype, dtype) + fake_jnp.dtype = lambda dtype: dtype + fake_jnp.ones = lambda shape, dtype: _Array(shape, dtype) + fake_jnp.empty = lambda shape, dtype: _Array(shape, dtype) + fake_jnp.full = lambda shape, value, dtype: _Array(shape, dtype, value) + fake_jax.numpy = fake_jnp + + try: + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "torch": None, + "cutlass": None, + }, + ): + cls.common = importlib.import_module(f"{discrete_name}._jax_common") + cls.forward = importlib.import_module( + f"{discrete_name}.discrete_grouped_gemm_swiglu.jax" + ) + cls.backward = importlib.import_module( + f"{discrete_name}.discrete_grouped_gemm_dswiglu.jax" + ) + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + @staticmethod + def _scale_shape( + rows: int, k: int, experts: int, sf_vec_size: int = 32 + ) -> tuple[int, ...]: + canonical = ( + 32, + 4, + (rows + 127) // 128, + 4, + ((k + sf_vec_size - 1) // sf_vec_size + 3) // 4, + experts, + ) + return (experts, canonical[2], canonical[4], 32, 4, 4) + + def test_forward_layouts_map_stacked_experts_to_canonical_kernel_axes(self): + m, n, k, experts = 256, 256, 128, 4 + operation = self.forward.DiscreteGroupedGemmSwigluSm100( + _Array((1, m, k), "float8_e4m3fn"), + _Array((experts, n, k), "float8_e4m3fn"), + _Array(self._scale_shape(m, k, 1), "float8_e8m0fnu"), + _Array(self._scale_shape(n, k, experts), "float8_e8m0fnu"), + _Array((experts,), "int32"), + _Array((experts,), "float32"), + sample_norm_const=_Array((1,), "float32"), + sample_prob=_Array((1, 1, m), "float32"), + sample_bias=_Array((experts, n), "bfloat16"), + ) + + self.assertTrue(operation.check_support()) + self.assertEqual(operation.a_desc.shape, (m, k, 1)) + self.assertEqual(operation.b_desc.shape, (n, k, experts)) + self.assertEqual(operation.b_desc.stride_order, (1, 0, 2)) + self.assertEqual(operation.c_desc.shape, (m, n, 1)) + self.assertEqual(operation.d_desc.shape, (m, n // 2, 1)) + self.assertEqual(operation.d_desc.stride_order, (1, 0, 2)) + self.assertEqual(operation.sfd_row_desc.shape, (32, 4, 2, 4, 1, 1)) + self.assertEqual(operation.amax_desc.shape, (experts, 1)) + self.assertEqual(operation.amax_desc.init_value, float("-inf")) + + def test_backward_infers_initialized_reduction_outputs(self): + m, n, k, experts = 256, 256, 128, 4 + operation = self.backward.DiscreteGroupedGemmDswigluSm100( + _Array((1, m, k), "float8_e4m3fn"), + _Array((experts, n, k), "float8_e4m3fn"), + _Array((1, m, 2 * n), "bfloat16"), + _Array(self._scale_shape(m, k, 1), "float8_e8m0fnu"), + _Array(self._scale_shape(n, k, experts), "float8_e8m0fnu"), + _Array((experts,), "int32"), + _Array((experts,), "float32"), + _Array((experts,), "float32"), + _Array((1, 1, m), "float32"), + generate_dbias=True, + ) + + self.assertTrue(operation.check_support()) + self.assertEqual(operation.d_row_desc.shape, (m, 2 * n, 1)) + self.assertEqual(operation.dprob_desc.shape, (m, 1, 1)) + self.assertEqual(operation.dprob_desc.init_value, 0.0) + self.assertEqual(operation.dbias_desc.shape, (experts, 2 * n, 1)) + self.assertEqual(operation.dbias_desc.init_value, 0.0) + self.assertEqual(operation.amax_desc.shape, (experts, 2, 1)) + self.assertEqual(operation.amax_desc.init_value, float("-inf")) + + def test_non_k_major_stacked_weights_are_rejected(self): + m, n, k, experts = 256, 256, 128, 4 + operation = self.forward.DiscreteGroupedGemmSwigluSm100( + _Array((1, m, k), "float8_e4m3fn"), + _Array((experts, k, n), "float8_e4m3fn"), + _Array(self._scale_shape(m, k, 1), "float8_e8m0fnu"), + _Array(self._scale_shape(n, k, experts), "float8_e8m0fnu"), + _Array((experts,), "int32"), + _Array((experts,), "float32"), + b_layout="LKN", + ) + + with self.assertRaisesRegex(ValueError, "K-major A and B"): + operation.check_support() + + def test_launcher_passes_live_stacked_operands_before_outputs_and_workspace(self): + m, n, k, experts = 256, 256, 128, 4 + a = _Array((1, m, k), "float8_e4m3fn") + b = _Array((experts, n, k), "float8_e4m3fn") + sfa = _Array(self._scale_shape(m, k, 1), "float8_e8m0fnu") + sfb = _Array(self._scale_shape(n, k, experts), "float8_e8m0fnu") + offsets = _Array((experts,), "int32") + alpha = _Array((experts,), "float32") + operation = self.forward.DiscreteGroupedGemmSwigluSm100( + a, b, sfa, sfb, offsets, alpha + ) + + captured = {} + + class FakeKernel: + def __init__(self, **config): + captured["config"] = config + + @staticmethod + def get_workspace_bytes(): + return 512 + + def __call__(self, *args): + captured["args"] = args + + kernel_module_name = f"{_PACKAGE}.discrete_grouped_gemm.discrete_grouped_gemm_swiglu.discrete_B_blockscaled_grouped_gemm_glu_bias" + kernel_module = types.ModuleType(kernel_module_name) + kernel_module.BlockScaledDiscreteWeightGroupedGemmBiasKernel = FakeKernel + + cutlass = types.ModuleType("cutlass") + cutlass.Int32 = lambda value: ("i32", value) + cutlass.Int64 = lambda value: ("i64", value) + cutlass.Float32 = lambda value: ("f32", value) + cute = types.ModuleType("cutlass.cute") + nvgpu = types.ModuleType("cutlass.cute.nvgpu") + nvgpu.OperandMajorMode = types.SimpleNamespace(K="K") + cutlass_jax = types.ModuleType("cutlass.jax") + cutlass_jax.TensorSpec = _TensorSpec + cutlass_jax.jax_to_cutlass_dtype = lambda dtype: dtype + cutlass.cute = cute + + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + "cutlass.cute": cute, + "cutlass.cute.nvgpu": nvgpu, + "cutlass.jax": cutlass_jax, + kernel_module_name: kernel_module, + }, + ): + result = operation(a, b, sfa, sfb, offsets, alpha) + + kernel_args = captured["args"] + self.assertIs(kernel_args[0], a) + self.assertIs(kernel_args[1], b.iterator) + self.assertIs(kernel_args[2], sfb.iterator) + self.assertEqual(kernel_args[6], "K") + self.assertIs(kernel_args[7], operation.captured_workspaces[0].iterator) + self.assertEqual(operation.captured_workspaces[0].shape, (512,)) + self.assertEqual(operation.captured_workspace_specs[0].ptr_assumed_align, 128) + self.assertEqual(kernel_args[18].shape, (1, 1, m)) + self.assertEqual(kernel_args[18].dtype, "float32") + self.assertEqual( + tuple(result.keys()), + ( + "c_tensor", + "d_tensor", + "d_col_tensor", + "amax_tensor", + "sfd_row_tensor", + "sfd_col_tensor", + ), + ) + self.assertEqual(kernel_args[20], 8) + self.assertEqual(kernel_args[21], "stream") + self.assertTrue(captured["config"]["stacked_expert_inputs"]) + + def test_zero_m_materializes_results_without_launching(self): + m, n, k, experts = 0, 256, 128, 4 + a = _Array((1, m, k), "float8_e4m3fn") + b = _Array((experts, n, k), "float8_e4m3fn") + sfa = _Array(self._scale_shape(m, k, 1), "float8_e8m0fnu") + sfb = _Array(self._scale_shape(n, k, experts), "float8_e8m0fnu") + offsets = _Array((experts,), "int32") + alpha = _Array((experts,), "float32") + + forward = self.forward.DiscreteGroupedGemmSwigluSm100( + a, b, sfa, sfb, offsets, alpha + ) + forward_result = forward(a, b, sfa, sfb, offsets, alpha) + self.assertEqual(forward_result["c_tensor"].shape, (1, 0, n)) + self.assertEqual(forward_result["d_tensor"].shape, (1, 0, n // 2)) + self.assertEqual(forward_result["amax_tensor"].fill_value, float("-inf")) + + c = _Array((1, m, 2 * n), "bfloat16") + beta = _Array((experts,), "float32") + prob = _Array((1, 1, m), "float32") + backward = self.backward.DiscreteGroupedGemmDswigluSm100( + a, + b, + c, + sfa, + sfb, + offsets, + alpha, + beta, + prob, + generate_dbias=True, + ) + backward_result = backward(a, b, c, sfa, sfb, offsets, alpha, beta, prob) + self.assertEqual(backward_result["d_row_tensor"].shape, (1, 0, 2 * n)) + self.assertEqual(backward_result["dprob_tensor"].shape, (1, 1, 0)) + self.assertEqual(backward_result["dbias_tensor"].fill_value, 0.0) + self.assertEqual(backward_result["amax_tensor"].fill_value, float("-inf")) + self.assertEqual( + tuple(backward_result.keys()), + ( + "d_row_tensor", + "d_col_tensor", + "dprob_tensor", + "dbias_tensor", + "amax_tensor", + "sfd_row_tensor", + "sfd_col_tensor", + ), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_dsa_backward.py b/test/python/fe_api/test_jax_dsa_backward.py new file mode 100644 index 000000000..1947026a3 --- /dev/null +++ b/test/python/fe_api/test_jax_dsa_backward.py @@ -0,0 +1,360 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX integration coverage for DSA backward adapters.""" + +import pytest + + +def _jax_runtime(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + pytest.importorskip("cudnn") + return jax, jnp + + +def _supported_gpu(jax, minimum_compute_capability): + supported_targets = {90, 100, 103, 107} + for device in jax.local_devices(): + if device.platform != "gpu": + continue + reported = getattr(device, "compute_capability", None) + if isinstance(reported, (tuple, list)): + major, minor = int(reported[0]), int(reported[1]) + else: + text = str(reported) + if "." in text: + major_text, minor_text = text.split(".", 1) + major, minor = int(major_text), int(minor_text) + else: + try: + major, minor = divmod(int(text), 10) + except ValueError: + continue + capability = major * 10 + minor + if capability >= minimum_compute_capability and capability in supported_targets: + return device, capability + pytest.skip(f"A supported JAX SM{minimum_compute_capability}+ GPU is not available") + + +@pytest.mark.L0 +def test_jax_dsa_backward_abstract_wrappers(monkeypatch): + """Trace every backward signature on a CPU-only JAX installation.""" + + jax, jnp = _jax_runtime() + from cudnn._jax import JaxApiBase + from cudnn.deepseek_sparse_attention.indexer_backward.jax import ( + dense_indexer_backward_wrapper, + indexer_backward_wrapper, + ) + from cudnn.deepseek_sparse_attention.sparse_attention_backward.jax import sparse_attention_backward_wrapper + + def abstract_call(_self, _inputs, *, output_descs, **_options): + return tuple(jnp.empty(desc.shape, dtype=desc.dtype) for desc in output_descs) + + monkeypatch.setattr(JaxApiBase, "_call_kernel", abstract_call) + + def resolve_compute_capability( + target_compute_capability, + supported_compute_capabilities, + operation_name, + ): + del operation_name + resolved = 100 if target_compute_capability is None else target_compute_capability + assert resolved in supported_compute_capabilities + return resolved + + with monkeypatch.context() as target_resolution: + target_resolution.setattr( + JaxApiBase, + "_resolve_compute_capability", + staticmethod(resolve_compute_capability), + ) + q = jax.ShapeDtypeStruct((1, 2, 64, 128), jnp.bfloat16) + weights = jax.ShapeDtypeStruct((1, 2, 64), jnp.bfloat16) + k = jax.ShapeDtypeStruct((1, 128, 128), jnp.bfloat16) + sparse_score = jax.ShapeDtypeStruct((1, 2, 128), jnp.float32) + topk = jax.ShapeDtypeStruct((1, 2, 128), jnp.int32) + sparse = jax.eval_shape( + lambda q, w, k, target, predict, topk: indexer_backward_wrapper( + q, + w, + k, + target, + predict, + topk, + target_compute_capability=100, + ), + q, + weights, + k, + sparse_score, + sparse_score, + topk, + ) + assert sparse["d_index_q"].shape == q.shape + assert sparse["d_weights"].shape == weights.shape + assert sparse["d_index_k"].shape == k.shape + + q_sparse_thd = jax.ShapeDtypeStruct((2, 64, 128), jnp.bfloat16) + weights_sparse_thd = jax.ShapeDtypeStruct((2, 64), jnp.bfloat16) + k_sparse_thd = jax.ShapeDtypeStruct((128, 128), jnp.bfloat16) + score_sparse_thd = jax.ShapeDtypeStruct((2, 128), jnp.float32) + topk_sparse_thd = jax.ShapeDtypeStruct((2, 128), jnp.int32) + sparse_thd = jax.eval_shape( + lambda q, w, k, target, predict, topk: indexer_backward_wrapper( + q, + w, + k, + target, + predict, + topk, + topk_indices_global=True, + target_compute_capability=100, + ), + q_sparse_thd, + weights_sparse_thd, + k_sparse_thd, + score_sparse_thd, + score_sparse_thd, + topk_sparse_thd, + ) + assert sparse_thd["d_index_q"].shape == q_sparse_thd.shape + assert sparse_thd["d_weights"].shape == weights_sparse_thd.shape + assert sparse_thd["d_index_k"].shape == k_sparse_thd.shape + + dense_score = jax.ShapeDtypeStruct((1, 2, 128), jnp.float32) + dense_denom = jax.ShapeDtypeStruct((1, 2), jnp.float32) + dense = jax.eval_shape( + lambda q, w, k, target, denom, predict, lse: dense_indexer_backward_wrapper( + q, + w, + k, + target, + denom, + predict, + lse, + target_compute_capability=100, + ), + q, + weights, + k, + dense_score, + dense_denom, + dense_score, + dense_denom, + ) + assert dense["d_index_q"].shape == q.shape + assert dense["d_weights"].shape == weights.shape + assert dense["d_index_k"].shape == k.shape + + q_thd = jax.ShapeDtypeStruct((2, 64, 128), jnp.bfloat16) + weights_thd = jax.ShapeDtypeStruct((2, 64), jnp.bfloat16) + k_thd = jax.ShapeDtypeStruct((128, 128), jnp.bfloat16) + score_thd = jax.ShapeDtypeStruct((2, 128), jnp.float32) + denom_thd = jax.ShapeDtypeStruct((2,), jnp.float32) + cu = jax.ShapeDtypeStruct((2,), jnp.int32) + dense_thd = jax.eval_shape( + lambda q, w, k, target, denom, predict, lse, cu_q, cu_k: dense_indexer_backward_wrapper( + q, + w, + k, + target, + denom, + predict, + lse, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=2, + max_seqlen_k=128, + target_compute_capability=90, + ), + q_thd, + weights_thd, + k_thd, + score_thd, + denom_thd, + score_thd, + denom_thd, + cu, + cu, + ) + assert dense_thd["d_index_q"].shape == q_thd.shape + assert dense_thd["d_index_k"].shape == k_thd.shape + + attn_q = jax.ShapeDtypeStruct((2, 64, 512), jnp.bfloat16) + attn_kv = jax.ShapeDtypeStruct((128, 512), jnp.bfloat16) + attn_out = jax.ShapeDtypeStruct((2, 64, 512), jnp.bfloat16) + lse = jax.ShapeDtypeStruct((2, 64), jnp.float32) + sink = jax.ShapeDtypeStruct((64,), jnp.float32) + attn_topk = jax.ShapeDtypeStruct((2, 32), jnp.int32) + attention = jax.eval_shape( + lambda q, kv, out, dout, lse, sink, topk: sparse_attention_backward_wrapper( + q, + kv, + out, + dout, + lse, + sink, + topk, + target_compute_capability=100, + ), + attn_q, + attn_kv, + attn_out, + attn_out, + lse, + sink, + attn_topk, + ) + assert attention["dq"].shape == attn_q.shape + assert attention["dkv"].shape == attn_kv.shape + assert attention["d_sink"].shape == sink.shape + + +@pytest.mark.L0 +def test_jax_dense_indexer_backward_sm100_jit_and_numerics(): + jax, jnp = _jax_runtime() + cutlass_jax = pytest.importorskip("cutlass.jax") + if not cutlass_jax.is_available(): + pytest.skip("Installed JAX version is unsupported by CUTLASS JAX") + device, compute_capability = _supported_gpu(jax, 100) + + from cudnn.deepseek_sparse_attention.indexer_backward.jax import dense_indexer_backward_wrapper + + batch, seqlen_q, seqlen_k = 1, 128, 128 + heads, head_dim = 64, 128 + sm_scale = 0.5 + index_q = jax.device_put( + jax.random.normal( + jax.random.key(40), + (batch, seqlen_q, heads, head_dim), + dtype=jnp.bfloat16, + ), + device, + ) + weights = jax.device_put( + jax.random.normal( + jax.random.key(41), + (batch, seqlen_q, heads), + dtype=jnp.bfloat16, + ) + * jnp.bfloat16(1.0 / heads), + device, + ) + index_k = jax.device_put( + jax.random.normal( + jax.random.key(42), + (batch, seqlen_k, head_dim), + dtype=jnp.bfloat16, + ), + device, + ) + valid = jnp.arange(seqlen_k)[None, :] < jnp.arange(seqlen_q)[:, None] + 1 + attn_score = jax.device_put( + jax.random.uniform( + jax.random.key(43), + (batch, seqlen_q, seqlen_k), + dtype=jnp.float32, + minval=0.1, + maxval=1.0, + ), + device, + ) + attn_score = jnp.where(valid[None, :, :], attn_score, 0.0) + attn_l1norm = jnp.sum(attn_score, axis=-1) + scores_per_head = jnp.einsum( + "bqhd,bkd->bqhk", + index_q.astype(jnp.float32), + index_k.astype(jnp.float32), + ) + index_score = jnp.sum( + jnp.maximum(scores_per_head, 0.0) * (weights.astype(jnp.float32) * sm_scale)[..., None], + axis=2, + ) + index_lse = jax.nn.logsumexp( + jnp.where(valid[None, :, :], index_score, -jnp.inf), + axis=-1, + ) + + @jax.jit + def run(q, w, k, target, target_denom, predict, predict_lse, grad_loss): + return dense_indexer_backward_wrapper( + q, + w, + k, + target, + target_denom, + predict, + predict_lse, + grad_loss=grad_loss, + sm_scale=sm_scale, + target_compute_capability=compute_capability, + ) + + grad_loss = jax.device_put(jnp.ones((1,), dtype=jnp.float32), device) + arguments = ( + index_q, + weights, + index_k, + attn_score, + attn_l1norm, + index_score, + index_lse, + grad_loss, + ) + lowered = run.lower(*arguments) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + compiled = lowered.compile() + result = compiled(*arguments) + result["d_index_k"].block_until_ready() + assert result["d_index_q"].shape == index_q.shape + assert result["d_weights"].shape == weights.shape + assert result["d_index_k"].shape == index_k.shape + assert all(jnp.all(jnp.isfinite(value)) for value in result.values()) + + target = attn_score / attn_l1norm[..., None] + predict = jax.nn.softmax( + jnp.where(valid[None, :, :], index_score, -jnp.inf), + axis=-1, + ) + grad_scale = 1.0 / (batch * seqlen_q) + g = jnp.where( + valid[None, :, :], + -jnp.maximum(target, jnp.exp(-100.0)) * grad_scale, + 0.0, + ) + grad_signal = jnp.where( + valid[None, :, :], + g - predict * jnp.sum(g, axis=-1, keepdims=True), + 0.0, + ) + + def raw_scores(q, w, k): + per_head = jnp.einsum("bqhd,bkd->bqhk", q, k) + return jnp.sum( + jnp.maximum(per_head, 0.0) * (w * sm_scale)[..., None], + axis=2, + ) + + _, pullback = jax.vjp( + raw_scores, + index_q.astype(jnp.float32), + weights.astype(jnp.float32), + index_k.astype(jnp.float32), + ) + reference = pullback(grad_signal) + for name, expected in zip(("d_index_q", "d_weights", "d_index_k"), reference): + actual_f32 = result[name].astype(jnp.float32).ravel() + expected_f32 = expected.astype(jnp.float32).ravel() + cosine = jnp.vdot(actual_f32, expected_f32) / (jnp.linalg.norm(actual_f32) * jnp.linalg.norm(expected_f32)) + relative_rms = jnp.sqrt(jnp.mean(jnp.square(actual_f32 - expected_f32))) / jnp.sqrt(jnp.mean(jnp.square(expected_f32))) + assert cosine >= 0.97, f"{name} cosine similarity: {cosine}" + assert relative_rms <= 0.55, f"{name} relative RMS error: {relative_rms}" + + zero_result = compiled(*arguments[:-1], jnp.zeros_like(grad_loss)) + zero_result["d_index_k"].block_until_ready() + assert all(jnp.array_equal(value, jnp.zeros_like(value)) for value in zero_result.values()) diff --git a/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py b/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py new file mode 100644 index 000000000..57ec4bae7 --- /dev/null +++ b/test/python/fe_api/test_jax_dsa_indexer_backward_contract.py @@ -0,0 +1,800 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX contracts for DSA indexer backward.""" + +from __future__ import annotations + +import ast +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_DSA_ROOT = _CUDNN_ROOT / "deepseek_sparse_attention" +_OPERATION_ROOT = _DSA_ROOT / "indexer_backward" +_PACKAGE = "cudnn_jax_dsa_indexer_backward_test" + + +class _DataType(Enum): + NOT_SET = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + + +class _DType: + def __init__(self, name: str): + self.name = name + + def __repr__(self): + return self.name + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + def astype(self, dtype): + return _Array(self.shape, dtype) + + +class _TupleDict(dict): + def __iter__(self): + return iter(self.values()) + + +class _TensorSpec: + def __init__(self, *, mode=None, divisibility=None): + self.mode = mode + self.divisibility = divisibility + + +class _Scalar: + def __init__(self, value): + self.value = value + + +class _CuteTensor: + def __init__(self, shape, stride): + self.shape = tuple(shape) + self.stride = tuple(stride) + self.iterator = self + + +class _RecordedKernel: + events = [] + + def __init__(self, label, *configuration, **keywords): + self.label = label + self.configuration = (configuration, keywords) + + def __call__(self, *arguments): + self.events.append((self.label, self.configuration, arguments)) + + +class JaxIndexerBackwardContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.bfloat16 = _DType("bfloat16") + cls.float32 = _DType("float32") + cls.int32 = _DType("int32") + cls.dtype_to_cudnn = { + cls.bfloat16: _DataType.BFLOAT16, + cls.float32: _DataType.FLOAT, + cls.int32: _DataType.INT32, + } + cls.cudnn_to_dtype = {value: key for key, value in cls.dtype_to_cudnn.items()} + + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + dsa_name = f"{_PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_DSA_ROOT)] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + sys.modules[dsa_name] = dsa + + operation_name = f"{dsa_name}.indexer_backward" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + + class JaxTensorDesc(tensor_module.TensorDesc): + @property + def cudnn_dtype(self): + return cls.dtype_to_cudnn.get(self.dtype, _DataType.NOT_SET) + + def compact_like( + self, + *, + cudnn_dtype, + shape, + stride_order=None, + name="", + init_value=None, + mode=None, + ): + canonical = tensor_module.make_compact_tensor_desc( + dtype=cudnn_dtype, + shape=shape, + stride_order=stride_order, + name=name, + init_value=init_value, + ) + desc = JaxTensorDesc( + dtype=cls.cudnn_to_dtype[cudnn_dtype], + shape=canonical.shape, + stride=canonical.stride, + stride_order=canonical.stride_order, + name=name, + init_value=init_value, + ) + object.__setattr__( + desc, + "mode", + tuple(range(len(shape))) if mode is None else tuple(mode), + ) + return desc + + def with_divisibility(self, divisibility): + desc = JaxTensorDesc( + dtype=self.dtype, + shape=self.shape, + stride=self.stride, + stride_order=self.stride_order, + name=self.name, + init_value=self.init_value, + ) + object.__setattr__(desc, "mode", self.mode) + object.__setattr__(desc, "divisibility", tuple(divisibility)) + return desc + + class JaxApiBase: + captured_call = None + + @staticmethod + def _to_tensor_desc(value, name, *, mode=None, init_value=None): + public_shape = tuple(value.shape) + mode = tuple(range(len(public_shape))) if mode is None else tuple(mode) + public_stride = [0] * len(public_shape) + running = 1 + for dimension in reversed(range(len(public_shape))): + public_stride[dimension] = running + running *= max(public_shape[dimension], 1) + canonical_axis_by_public = [0] * len(mode) + for canonical_axis, public_axis in enumerate(mode): + canonical_axis_by_public[public_axis] = canonical_axis + desc = JaxTensorDesc( + dtype=value.dtype, + shape=tuple(public_shape[axis] for axis in mode), + stride=tuple(public_stride[axis] for axis in mode), + stride_order=tuple( + canonical_axis_by_public[axis] + for axis in reversed(range(len(public_shape))) + ), + name=name, + init_value=init_value, + ) + object.__setattr__(desc, "mode", mode) + return desc + + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del supported, operation_name + if target is None: + raise RuntimeError("target required by test") + return target + + @staticmethod + def _compute_capability_family(target, supported): + return max( + (value for value in supported if value <= target), default=None + ) + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + mode = expected.mode if mode is None else tuple(mode) + actual_dtype = cls.dtype_to_cudnn.get(value.dtype, _DataType.NOT_SET) + actual_shape = tuple(value.shape[axis] for axis in mode) + if ( + actual_shape != expected.shape + or actual_dtype != expected.cudnn_dtype + ): + raise ValueError(f"{expected.name} tensor signature mismatch") + + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + if mode is None: + mode = desc.mode + if divisibility is None: + divisibility = getattr(desc, "divisibility", None) + return _TensorSpec(mode=mode, divisibility=divisibility) + + def _call_kernel(self, inputs, *, launch, **options): + options["launch"] = launch + input_descs = options.get("input_descs") + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc) + derived_specs = tuple( + self._to_tensor_spec(desc) for desc in input_descs + ) + supplied_specs = options.get("input_spec") + options["input_spec"] = ( + derived_specs + if supplied_specs is None + else tuple( + derived if supplied is None else supplied + for derived, supplied in zip( + derived_specs, supplied_specs + ) + ) + ) + output_descs = options["output_descs"] + derived_output_specs = tuple( + self._to_tensor_spec(desc) for desc in output_descs + ) + supplied_output_specs = options.get("output_spec") + options["output_spec"] = ( + derived_output_specs + if supplied_output_specs is None + else tuple( + derived if supplied is None else supplied + for derived, supplied in zip( + derived_output_specs, supplied_output_specs + ) + ) + ) + options["workspace_spec"] = tuple( + self._to_tensor_spec(desc) + for desc in options.get("workspace_descs", ()) + ) + JaxApiBase.captured_call = (inputs, options) + specs = options.get("output_spec") or (None,) * len( + options["output_descs"] + ) + results = [] + for desc, spec in zip(options["output_descs"], specs): + mode = ( + desc.mode + if spec is None or spec.mode is None + else tuple(spec.mode) + ) + canonical_axis_by_public = [0] * len(mode) + for canonical_axis, public_axis in enumerate(mode): + canonical_axis_by_public[public_axis] = canonical_axis + public_shape = tuple( + desc.shape[canonical_axis_by_public[public_axis]] + for public_axis in range(desc.ndim) + ) + results.append( + _Array( + public_shape, + cls.cudnn_to_dtype[desc.cudnn_dtype], + ) + ) + return tuple(results) + + internal = types.ModuleType(f"{_PACKAGE}._jax") + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal.__name__ + internal.__spec__ = ModuleSpec(internal.__name__, loader=None, is_package=True) + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = _TupleDict + sys.modules[internal.__name__] = internal + + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.bfloat16 = cls.bfloat16 + fake_jnp.float32 = cls.float32 + fake_jnp.int32 = cls.int32 + fake_jnp.dtype = lambda value: value.dtype if hasattr(value, "dtype") else value + fake_jnp.asarray = lambda value, dtype: _Array( + getattr(value, "shape", ()), dtype + ) + fake_jnp.reshape = lambda value, shape: _Array(shape, value.dtype) + + fake_jax = types.ModuleType("jax") + fake_jax.__path__ = [] + fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + fake_jax.numpy = fake_jnp + fake_jax.ShapeDtypeStruct = _Array + fake_jax.jit = lambda function=None, **_kwargs: ( + (lambda fn: fn) if function is None else function + ) + + try: + with mock.patch.dict(sys.modules, {"jax": fake_jax, "jax.numpy": fake_jnp}): + cls.module = importlib.import_module(f"{operation_name}.jax") + except Exception: + cls.tearDownClass() + raise + cls.JaxApiBase = JaxApiBase + cls.fake_jax_modules = {"jax": fake_jax, "jax.numpy": fake_jnp} + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def setUp(self): + self.JaxApiBase.captured_call = None + _RecordedKernel.events.clear() + + @classmethod + def _sparse_inputs(cls, *, topk=128): + score_shape = (2, 64, topk) + return ( + _Array((2, 64, 64, 128), cls.bfloat16), + _Array((2, 64, 64), cls.bfloat16), + _Array((2, 256, 128), cls.bfloat16), + _Array(score_shape, cls.float32), + _Array(score_shape, cls.float32), + _Array(score_shape, cls.int32), + ) + + @classmethod + def _sparse_thd_inputs(cls, *, topk=128): + score_shape = (96, topk) + return ( + _Array((96, 64, 128), cls.bfloat16), + _Array((96, 64), cls.bfloat16), + _Array((384, 128), cls.bfloat16), + _Array(score_shape, cls.float32), + _Array(score_shape, cls.float32), + _Array(score_shape, cls.int32), + ) + + @classmethod + def _sparse_sequence_major_inputs(cls, *, topk=128): + score_shape = (64, 2, topk) + return ( + _Array((64, 2, 64, 128), cls.bfloat16), + _Array((64, 2, 64), cls.bfloat16), + _Array((256, 2, 128), cls.bfloat16), + _Array(score_shape, cls.float32), + _Array(score_shape, cls.float32), + _Array(score_shape, cls.int32), + ) + + @classmethod + def _dense_bshd_inputs(cls): + return ( + _Array((2, 64, 64, 128), cls.bfloat16), + _Array((2, 64, 64), cls.bfloat16), + _Array((2, 256, 128), cls.bfloat16), + _Array((2, 64, 256), cls.float32), + _Array((2, 64), cls.float32), + _Array((2, 64, 256), cls.float32), + _Array((2, 64), cls.float32), + ) + + @classmethod + def _dense_thd_inputs(cls): + return ( + _Array((96, 64, 128), cls.bfloat16), + _Array((96, 64), cls.bfloat16), + _Array((384, 128), cls.bfloat16), + _Array((96, 256), cls.float32), + _Array((96,), cls.float32), + _Array((96, 256), cls.float32), + _Array((96,), cls.float32), + ) + + @classmethod + def _dense_sequence_major_inputs(cls): + return ( + _Array((64, 2, 64, 128), cls.bfloat16), + _Array((64, 2, 64), cls.bfloat16), + _Array((256, 2, 128), cls.bfloat16), + _Array((64, 2, 256), cls.float32), + _Array((64, 2), cls.float32), + _Array((64, 2, 256), cls.float32), + _Array((64, 2), cls.float32), + ) + + def test_sparse_wrapper_declares_functional_outputs_and_zeroed_dk(self): + inputs = self._sparse_inputs() + with mock.patch.dict(sys.modules, self.fake_jax_modules): + result = self.module.indexer_backward_wrapper( + *inputs, target_compute_capability=103 + ) + + self.assertEqual( + tuple(result), + (result["d_index_q"], result["d_weights"], result["d_index_k"]), + ) + self.assertEqual(result["d_index_q"].shape, inputs[0].shape) + self.assertEqual(result["d_weights"].shape, inputs[1].shape) + self.assertEqual(result["d_index_k"].shape, inputs[2].shape) + kernel_inputs, options = self.JaxApiBase.captured_call + self.assertEqual(kernel_inputs[:-1], inputs) + self.assertEqual( + (kernel_inputs[-1].shape, kernel_inputs[-1].dtype), ((1,), self.float32) + ) + self.assertEqual(options["output_descs"][2].init_value, 0.0) + self.assertEqual(options["workspace_descs"][0].shape, inputs[3].shape) + self.assertIn("--gpu-arch sm_103a", options["compile_options"]) + self.assertEqual(options["input_spec"][0].divisibility, (None, None, None, 128)) + self.assertEqual(options["input_spec"][2].divisibility, (None, None, 128)) + + def test_sparse_packed_thd_signature_requires_global_indices(self): + inputs = self._sparse_thd_inputs() + with mock.patch.dict(sys.modules, self.fake_jax_modules): + with self.assertRaisesRegex( + ValueError, "requires topk_indices_global=True" + ): + self.module.indexer_backward_wrapper( + *inputs, + target_compute_capability=100, + ) + result = self.module.indexer_backward_wrapper( + *inputs, + topk_indices_global=True, + target_compute_capability=100, + ) + + self.assertEqual(result["d_index_q"].shape, inputs[0].shape) + self.assertEqual(result["d_weights"].shape, inputs[1].shape) + self.assertEqual(result["d_index_k"].shape, inputs[2].shape) + _, options = self.JaxApiBase.captured_call + self.assertEqual(options["input_spec"][0].divisibility, (None, None, 128)) + self.assertEqual(options["input_spec"][2].divisibility, (None, 128)) + + def test_sparse_sequence_major_layouts_preserve_public_shapes(self): + inputs = self._sparse_sequence_major_inputs() + with mock.patch.dict(sys.modules, self.fake_jax_modules): + result = self.module.indexer_backward_wrapper( + *inputs, + q_layout="SBHD", + w_layout="SBH", + k_layout="SBD", + score_layout="SBK", + target_compute_capability=100, + ) + + self.assertEqual(result["d_index_q"].shape, inputs[0].shape) + self.assertEqual(result["d_weights"].shape, inputs[1].shape) + self.assertEqual(result["d_index_k"].shape, inputs[2].shape) + _, options = self.JaxApiBase.captured_call + self.assertEqual( + tuple(spec.mode for spec in options["input_spec"][:6]), + ( + (1, 0, 2, 3), + (1, 0, 2), + (1, 0, 2), + (1, 0, 2), + (1, 0, 2), + (1, 0, 2), + ), + ) + self.assertEqual( + tuple(spec.mode for spec in options["output_spec"]), + ((1, 0, 2, 3), (1, 0, 2), (1, 0, 2)), + ) + self.assertEqual(options["workspace_spec"][0].mode, (1, 0, 2)) + + def test_grad_loss_uses_a_fixed_runtime_abi_descriptor(self): + inputs = self._sparse_inputs() + with mock.patch.dict(sys.modules, self.fake_jax_modules): + api = self.module.IndexerBackward( + *inputs, + target_compute_capability=100, + ) + self.assertEqual(api.grad_loss_desc.shape, (1,)) + self.assertIs(api.grad_loss_desc.dtype, self.float32) + + def test_dense_bshd_declares_functional_outputs_and_runtime_grad_loss(self): + inputs = self._dense_bshd_inputs() + grad_loss = _Array((1,), self.float32) + with mock.patch.dict(sys.modules, self.fake_jax_modules): + result = self.module.dense_indexer_backward_wrapper( + *inputs, + grad_loss=grad_loss, + target_compute_capability=100, + ) + + self.assertEqual(result["d_index_q"].shape, inputs[0].shape) + self.assertEqual(result["d_weights"].shape, inputs[1].shape) + self.assertEqual(result["d_index_k"].shape, inputs[2].shape) + kernel_inputs, options = self.JaxApiBase.captured_call + self.assertEqual( + (kernel_inputs[7].shape, kernel_inputs[7].dtype), ((1,), self.float32) + ) + self.assertEqual(options["output_descs"][2].init_value, 0.0) + self.assertEqual(options["workspace_descs"][0].shape, inputs[3].shape) + self.assertIn("--gpu-arch sm_100a", options["compile_options"]) + + def test_dense_thd_signature_and_optional_runtime_operands(self): + inputs = self._dense_thd_inputs() + cu_q = _Array((3,), self.int32) + cu_k = _Array((3,), self.int32) + offsets = _Array((2,), self.int32) + with mock.patch.dict(sys.modules, self.fake_jax_modules): + api = self.module.DenseIndexerBackward( + *inputs, + sample_cu_seqlens_q=cu_q, + sample_cu_seqlens_k=cu_k, + sample_q_causal_offsets=offsets, + max_seqlen_q=64, + max_seqlen_k=256, + target_compute_capability=90, + ) + result = api( + *inputs, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + q_causal_offsets=offsets, + ) + + self.assertTrue(api._op.is_thd) + self.assertEqual(result["d_index_q"].shape, inputs[0].shape) + kernel_inputs, options = self.JaxApiBase.captured_call + self.assertEqual(kernel_inputs[-3:], (cu_q, cu_k, offsets)) + self.assertEqual(options["input_spec"][0].divisibility, (None, None, 128)) + self.assertEqual(options["input_spec"][2].divisibility, (None, 128)) + self.assertIn("--gpu-arch sm_90a", options["compile_options"]) + + def test_dense_sequence_major_layouts_preserve_public_shapes(self): + inputs = self._dense_sequence_major_inputs() + with mock.patch.dict(sys.modules, self.fake_jax_modules): + result = self.module.dense_indexer_backward_wrapper( + *inputs, + q_layout="SBHD", + w_layout="SBH", + k_layout="SBD", + score_layout="SBK", + denom_layout="SB", + target_compute_capability=100, + ) + + self.assertEqual(result["d_index_q"].shape, inputs[0].shape) + self.assertEqual(result["d_weights"].shape, inputs[1].shape) + self.assertEqual(result["d_index_k"].shape, inputs[2].shape) + _, options = self.JaxApiBase.captured_call + self.assertEqual( + tuple(spec.mode for spec in options["input_spec"][:7]), + ( + (1, 0, 2, 3), + (1, 0, 2), + (1, 0, 2), + (1, 0, 2), + (1, 0), + (1, 0, 2), + (1, 0), + ), + ) + self.assertEqual( + tuple(spec.mode for spec in options["output_spec"]), + ((1, 0, 2, 3), (1, 0, 2), (1, 0, 2)), + ) + self.assertEqual(options["workspace_spec"][0].mode, (1, 0, 2)) + + def _kernel_modules(self): + package = self.module.__package__ + + def kernel(label): + return lambda *args, **kwargs: _RecordedKernel(label, *args, **kwargs) + + sparse_100 = types.ModuleType(f"{package}.indexer_backward_sm100") + sparse_100.ScoreGradSm100 = kernel("sparse_score_100") + sparse_100.IndexerBackwardSm100 = kernel("sparse_gemm_100") + sparse_90 = types.ModuleType(f"{package}.indexer_backward_sm90") + sparse_90.ScoreGradSm90 = kernel("sparse_score_90") + sparse_90.IndexerBackwardSm90 = kernel("sparse_gemm_90") + dense_100 = types.ModuleType(f"{package}.dense_indexer_backward_sm100") + dense_100.ScoreGradDense = kernel("dense_score_100") + dense_100.DenseIndexerBackward2QGemmSm100 = kernel("dense_gemm_100") + dense_90 = types.ModuleType(f"{package}.dense_indexer_backward_sm90") + dense_90.ScoreGradDenseSm90 = kernel("dense_score_90") + return { + module.__name__: module + for module in (sparse_100, sparse_90, dense_100, dense_90) + } + + @staticmethod + def _fake_cutlass(): + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass.Float32 = _Scalar + cutlass.Int32 = _Scalar + return cutlass + + @staticmethod + def _fake_cute(): + cute = types.ModuleType("cutlass.cute") + cute.make_layout = lambda shape, stride: types.SimpleNamespace( + shape=tuple(shape), + stride=tuple(stride), + ) + cute.make_tensor = lambda _iterator, layout: _CuteTensor( + layout.shape, layout.stride + ) + return cute + + @staticmethod + def _compact_cute(shape): + stride = [0] * len(shape) + running = 1 + for dimension in reversed(range(len(shape))): + stride[dimension] = running + running *= shape[dimension] + return _CuteTensor(shape, stride) + + def test_sparse_sm100_and_sm90_launch_functionally(self): + inputs = self._sparse_inputs() + kernel_inputs = tuple(object() for _ in range(7)) + outputs = tuple(object() for _ in range(3)) + workspace = (object(),) + modules = { + **self.fake_jax_modules, + "cutlass": self._fake_cutlass(), + **self._kernel_modules(), + } + + for target, labels in ( + (100, ("sparse_score_100", "sparse_gemm_100")), + (90, ("sparse_score_90", "sparse_gemm_90")), + ): + with self.subTest(target=target), mock.patch.dict(sys.modules, modules): + _RecordedKernel.events.clear() + api = self.module.IndexerBackward( + *inputs, target_compute_capability=target + ) + api.check_support() + api._launch_kernel(object(), *kernel_inputs, *outputs, *workspace) + self.assertEqual( + tuple(event[0] for event in _RecordedKernel.events), labels + ) + score_args = _RecordedKernel.events[0][2] + self.assertIs(score_args[-2 if target == 100 else -1], workspace[0]) + gemm_args = _RecordedKernel.events[1][2] + self.assertIs(gemm_args[6], workspace[0]) + + def test_sparse_packed_thd_launch_adds_a_synthetic_batch_view(self): + inputs = self._sparse_thd_inputs() + cute = self._fake_cute() + cutlass = self._fake_cutlass() + cutlass.cute = cute + modules = { + **self.fake_jax_modules, + "cutlass": cutlass, + "cutlass.cute": cute, + **self._kernel_modules(), + } + kernel_inputs = tuple(self._compact_cute(value.shape) for value in inputs) + ( + self._compact_cute((1,)), + ) + outputs = ( + self._compact_cute(inputs[0].shape), + self._compact_cute(inputs[1].shape), + self._compact_cute(inputs[2].shape), + ) + workspace = (self._compact_cute(inputs[3].shape),) + + for target, labels in ( + (100, ("sparse_score_100", "sparse_gemm_100")), + (90, ("sparse_score_90", "sparse_gemm_90")), + ): + with self.subTest(target=target), mock.patch.dict(sys.modules, modules): + _RecordedKernel.events.clear() + api = self.module.IndexerBackward( + *inputs, + topk_indices_global=True, + target_compute_capability=target, + ) + api.check_support() + api._launch_kernel(object(), *kernel_inputs, *outputs, *workspace) + + self.assertEqual( + tuple(event[0] for event in _RecordedKernel.events), labels + ) + score_args = _RecordedKernel.events[0][2] + self.assertEqual(score_args[0].shape, (1, 96, 128)) + self.assertEqual(score_args[1].shape, (1, 96, 128)) + self.assertEqual(score_args[5].shape, (1, 96, 128)) + gemm_args = _RecordedKernel.events[1][2] + self.assertEqual(gemm_args[0].shape, (1, 96, 64, 128)) + self.assertEqual(gemm_args[1].shape, (1, 96, 64)) + self.assertEqual(gemm_args[2].shape, (1, 384, 128)) + self.assertEqual(gemm_args[7].shape, (1, 96, 128)) + + def test_dense_sm100_bshd_and_sm90_thd_launch_functionally(self): + modules = { + **self.fake_jax_modules, + "cutlass": self._fake_cutlass(), + **self._kernel_modules(), + } + outputs = tuple(object() for _ in range(3)) + workspace = (object(),) + + bshd = self._dense_bshd_inputs() + with mock.patch.dict(sys.modules, modules): + api = self.module.DenseIndexerBackward(*bshd, target_compute_capability=100) + api.check_support() + kernel_inputs = tuple(object() for _ in range(8)) + api._launch_kernel(object(), *kernel_inputs, *outputs, *workspace) + self.assertEqual( + tuple(event[0] for event in _RecordedKernel.events), + ("dense_score_100", "dense_gemm_100"), + ) + score_args = _RecordedKernel.events[0][2] + self.assertIs(score_args[11], workspace[0]) + self.assertIs(score_args[12], kernel_inputs[7]) + + _RecordedKernel.events.clear() + thd = self._dense_thd_inputs() + cu_q = _Array((3,), self.int32) + cu_k = _Array((3,), self.int32) + offsets = _Array((2,), self.int32) + with mock.patch.dict(sys.modules, modules): + api = self.module.DenseIndexerBackward( + *thd, + sample_cu_seqlens_q=cu_q, + sample_cu_seqlens_k=cu_k, + sample_q_causal_offsets=offsets, + max_seqlen_q=64, + max_seqlen_k=256, + target_compute_capability=90, + ) + api.check_support() + kernel_inputs = tuple(object() for _ in range(11)) + api._launch_kernel(object(), *kernel_inputs, *outputs, *workspace) + self.assertEqual( + tuple(event[0] for event in _RecordedKernel.events), + ("dense_score_90", "sparse_gemm_90"), + ) + score_args = _RecordedKernel.events[0][2] + self.assertEqual(score_args[4:7], kernel_inputs[8:11]) + self.assertIs(score_args[11], workspace[0]) + self.assertIs(score_args[12], kernel_inputs[7]) + gemm_args = _RecordedKernel.events[1][2] + self.assertEqual(gemm_args[10:12], kernel_inputs[8:10]) + self.assertIs(gemm_args[14], kernel_inputs[10]) + + def test_jax_reachable_modules_do_not_import_torch_at_module_scope(self): + for filename in ( + "jax.py", + "indexer_backward_sm90.py", + "indexer_backward_sm100.py", + "dense_indexer_backward_sm90.py", + "dense_indexer_backward_sm100.py", + ): + path = _OPERATION_ROOT / filename + tree = ast.parse(path.read_text(), filename=str(path)) + imports = [] + for node in tree.body: + if isinstance(node, ast.Import): + imports.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + imports.append(node.module or "") + self.assertNotIn("torch", imports, filename) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_dsa_indexers.py b/test/python/fe_api/test_jax_dsa_indexers.py new file mode 100644 index 000000000..d41bf4e83 --- /dev/null +++ b/test/python/fe_api/test_jax_dsa_indexers.py @@ -0,0 +1,279 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX integration coverage for DSA indexer forward and top-K.""" + +import pytest + + +def _jax_runtime(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + cutlass_jax = pytest.importorskip("cutlass.jax") + if not cutlass_jax.is_available(): + pytest.skip("Installed JAX version is unsupported by CUTLASS JAX") + return jax, jnp + + +def _supported_gpu(jax, minimum=90): + supported_targets = {90, 100, 103, 107} + try: + devices = tuple(jax.local_devices(backend="gpu")) + except RuntimeError as error: + pytest.skip(f"A JAX GPU is not available: {error}") + if not devices: + pytest.skip("A JAX GPU is not available") + + capabilities = [] + for device in devices: + reported = getattr(device, "compute_capability", None) + try: + if isinstance(reported, (tuple, list)): + major, minor = int(reported[0]), int(reported[1]) + else: + major_text, minor_text = str(reported).split(".", 1) + major, minor = int(major_text), int(minor_text) + except (TypeError, ValueError): + pytest.skip(f"JAX reported an unsupported compute capability {reported!r}") + capabilities.append(major * 10 + minor) + + if len(set(capabilities)) != 1: + pytest.skip(f"DSA indexers require homogeneous local GPU targets, found {capabilities}") + capability = capabilities[0] + if capability < minimum or capability not in supported_targets: + pytest.skip(f"A supported JAX SM{minimum}+ GPU is not available") + return devices[0], capability + + +@pytest.mark.L0 +def test_jax_dsa_indexer_abstract_shapes(monkeypatch): + jax, jnp = _jax_runtime() + from cudnn._jax import JaxApiBase + from cudnn.jax import ( + compactify_wrapper, + indexer_forward_wrapper, + indexer_top_k_wrapper, + local_to_global_wrapper, + ) + + monkeypatch.setattr(JaxApiBase, "_local_gpu_capabilities", staticmethod(lambda _operation_name: ((object(), 100),))) + + q = jax.ShapeDtypeStruct((1, 8, 32, 128), jnp.bfloat16) + k = jax.ShapeDtypeStruct((1, 5, 1, 128), jnp.bfloat16) + w = jax.ShapeDtypeStruct((1, 8, 32), jnp.bfloat16) + fixed = jax.eval_shape( + lambda q, k, w: indexer_forward_wrapper(q, k, w, ratio=1, target_compute_capability=100), + q, + k, + w, + )["scores"] + assert fixed.shape == (1, 8, 5) + assert fixed.dtype == jnp.float32 + + q_sequence_major = jax.ShapeDtypeStruct((8, 2, 32, 128), jnp.bfloat16) + k_sequence_major = jax.ShapeDtypeStruct((5, 2, 1, 128), jnp.bfloat16) + w_sequence_major = jax.ShapeDtypeStruct((8, 2, 32), jnp.bfloat16) + sequence_major = jax.eval_shape( + lambda q, k, w: indexer_forward_wrapper( + q, + k, + w, + ratio=1, + q_layout="SBHD", + k_layout="SBHD", + w_layout="SBH", + output_layout="SBK", + target_compute_capability=100, + ), + q_sequence_major, + k_sequence_major, + w_sequence_major, + )["scores"] + assert sequence_major.shape == (8, 2, 5) + assert sequence_major.dtype == jnp.float32 + + q_thd = jax.ShapeDtypeStruct((8, 32, 128), jnp.bfloat16) + k_thd = jax.ShapeDtypeStruct((5, 1, 128), jnp.bfloat16) + w_thd = jax.ShapeDtypeStruct((8, 32), jnp.bfloat16) + cu = jax.ShapeDtypeStruct((2,), jnp.int32) + packed = jax.eval_shape( + lambda q, k, w, cu_q, cu_k: indexer_forward_wrapper( + q, + k, + w, + ratio=1, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=8, + max_seqlen_k=5, + target_compute_capability=100, + ), + q_thd, + k_thd, + w_thd, + cu, + cu, + )["scores"] + assert packed.shape == (8, 5) + + values = jax.ShapeDtypeStruct((2, 64), jnp.float32) + lengths = jax.ShapeDtypeStruct((2,), jnp.int32) + with_values = jax.eval_shape( + lambda values, lengths: indexer_top_k_wrapper( + values, + lengths, + 8, + target_compute_capability=100, + ), + values, + lengths, + ) + assert with_values["indices"].shape == (2, 8) + assert with_values["indices"].dtype == jnp.int32 + assert with_values["values"].shape == (2, 8) + without_values = jax.eval_shape( + lambda values, lengths: indexer_top_k_wrapper( + values, + lengths, + 8, + return_val=False, + target_compute_capability=100, + ), + values, + lengths, + ) + assert without_values["indices"].shape == (2, 8) + assert without_values["values"] is None + + local = jax.ShapeDtypeStruct((1, 2, 8), jnp.int64) + global_indices = jax.eval_shape( + lambda indices: local_to_global_wrapper(indices, 64, target_compute_capability=100), + local, + )["indices"] + assert global_indices.shape == local.shape + assert global_indices.dtype == jnp.int32 + packed_global = jax.eval_shape( + lambda indices, cu_q, cu_k: local_to_global_wrapper( + indices, + 64, + cu_q, + cu_k, + target_compute_capability=100, + ), + jax.ShapeDtypeStruct((2, 8), jnp.int32), + jax.ShapeDtypeStruct((2,), jnp.int32), + jax.ShapeDtypeStruct((2,), jnp.int32), + )["indices"] + assert packed_global.shape == (2, 8) + + compact = jax.eval_shape( + lambda indices: compactify_wrapper(indices, target_compute_capability=100), + jax.ShapeDtypeStruct((1, 2, 8), jnp.int32), + ) + assert compact["indices"].shape == (2, 8) + assert compact["topk_length"].shape == (2,) + + +@pytest.mark.L0 +def test_jax_indexer_forward_jit_and_numerics(): + jax, jnp = _jax_runtime() + device, capability = _supported_gpu(jax, 90) + from cudnn.jax import indexer_forward_wrapper + + batch, seqlen_q, seqlen_k = 2, 8, 5 + num_query_heads, head_dim = 32, 128 + q = jax.device_put( + jax.random.normal(jax.random.key(0), (batch, seqlen_q, num_query_heads, head_dim), dtype=jnp.bfloat16), + device, + ) + k = jax.device_put( + jax.random.normal(jax.random.key(1), (batch, seqlen_k, 1, head_dim), dtype=jnp.bfloat16), + device, + ) + w = jax.device_put( + jax.random.normal(jax.random.key(2), (batch, seqlen_q, num_query_heads), dtype=jnp.bfloat16), + device, + ) + + lowered = indexer_forward_wrapper.lower( + q, + k, + w, + ratio=1, + target_compute_capability=capability, + ) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + scores = lowered.compile()(q, k, w)["scores"] + scores.block_until_ready() + + expanded_k = jnp.repeat(k.astype(jnp.float32), num_query_heads, axis=2) + per_head = jnp.einsum("bqhd,bkhd->bqhk", q.astype(jnp.float32), expanded_k) + expected = jnp.sum(jnp.maximum(per_head, 0.0) * w.astype(jnp.float32)[..., None], axis=2) + valid = jnp.arange(seqlen_k)[None, :] < (jnp.arange(seqlen_q)[:, None] + 1) + expected = jnp.where(valid[None, :, :], expected, -jnp.inf) + assert jnp.array_equal(jnp.isneginf(scores), jnp.isneginf(expected)) + finite = jnp.isfinite(expected) + assert jnp.allclose(scores[finite], expected[finite], atol=5e-2, rtol=2e-2) + + q_sequence_major = jnp.transpose(q, (1, 0, 2, 3)) + k_sequence_major = jnp.transpose(k, (1, 0, 2, 3)) + w_sequence_major = jnp.transpose(w, (1, 0, 2)) + sequence_major_lowered = indexer_forward_wrapper.lower( + q_sequence_major, + k_sequence_major, + w_sequence_major, + ratio=1, + q_layout="SBHD", + k_layout="SBHD", + w_layout="SBH", + output_layout="SBK", + target_compute_capability=capability, + ) + sequence_major_scores = sequence_major_lowered.compile()( + q_sequence_major, + k_sequence_major, + w_sequence_major, + )["scores"] + sequence_major_scores.block_until_ready() + expected_sequence_major = jnp.transpose(expected, (1, 0, 2)) + assert jnp.array_equal( + jnp.isneginf(sequence_major_scores), + jnp.isneginf(expected_sequence_major), + ) + finite = jnp.isfinite(expected_sequence_major) + assert jnp.allclose( + sequence_major_scores[finite], + expected_sequence_major[finite], + atol=5e-2, + rtol=2e-2, + ) + + +@pytest.mark.L0 +def test_jax_indexer_top_k_jit_and_numerics(): + jax, jnp = _jax_runtime() + device, capability = _supported_gpu(jax, 90) + from cudnn.jax import indexer_top_k_wrapper + + num_rows, num_cols, top_k = 4, 64, 8 + row = jnp.arange(num_cols, dtype=jnp.float32) + input_values = jax.device_put(jnp.broadcast_to(row, (num_rows, num_cols)), device) + seq_lens = jax.device_put(jnp.asarray((64, 63, 62, 61), dtype=jnp.int32), device) + lowered = indexer_top_k_wrapper.lower( + input_values, + seq_lens, + top_k, + target_compute_capability=capability, + ) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + result = lowered.compile()(input_values, seq_lens) + result["indices"].block_until_ready() + + valid = jnp.arange(num_cols)[None, :] < seq_lens[:, None] + reference_values, reference_indices = jax.lax.top_k(jnp.where(valid, input_values, -jnp.inf), top_k) + assert jnp.array_equal(jnp.sort(result["indices"], axis=-1), jnp.sort(reference_indices, axis=-1)) + assert jnp.array_equal(jnp.sort(result["values"], axis=-1), jnp.sort(reference_values, axis=-1)) diff --git a/test/python/fe_api/test_jax_dsa_score_recompute.py b/test/python/fe_api/test_jax_dsa_score_recompute.py new file mode 100644 index 000000000..3b9ee967f --- /dev/null +++ b/test/python/fe_api/test_jax_dsa_score_recompute.py @@ -0,0 +1,349 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX integration tests for DSA score-recompute kernels.""" + +import pytest + + +def _jax_runtime(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + cutlass_jax = pytest.importorskip("cutlass.jax") + if not cutlass_jax.is_available(): + pytest.skip("Installed JAX version is unsupported by CUTLASS JAX") + return jax, jnp + + +def _sm100_device(jax): + supported = {100, 103, 107} + try: + devices = tuple(jax.local_devices(backend="gpu")) + except RuntimeError as error: + pytest.skip(f"A JAX GPU is not available: {error}") + if not devices: + pytest.skip("A JAX GPU is not available") + capabilities = [] + for device in devices: + reported = getattr(device, "compute_capability", None) + try: + if isinstance(reported, (tuple, list)): + major, minor = (int(value) for value in reported[:2]) + capability = major * 10 + minor + elif "." in str(reported): + major, minor = (int(value) for value in str(reported).split(".", 1)) + capability = major * 10 + minor + else: + capability = int(reported) + if capability < 10: + capability *= 10 + except (TypeError, ValueError): + pytest.skip(f"JAX reported an unsupported compute capability {reported!r}") + capabilities.append(capability) + if len(set(capabilities)) != 1 or capabilities[0] not in supported: + pytest.skip(f"Score recompute requires homogeneous SM100/SM103/SM107 devices, found {capabilities}") + return devices[0], capabilities[0] + + +@pytest.mark.L0 +def test_jax_score_recompute_eval_shapes(monkeypatch): + jax, jnp = _jax_runtime() + from cudnn._jax import JaxApiBase + from cudnn.jax import ( + dense_attn_score_recompute_wrapper, + dense_indexer_score_recompute_wrapper, + sparse_attn_score_recompute_wrapper, + sparse_indexer_score_recompute_wrapper, + ) + + monkeypatch.setattr(JaxApiBase, "_local_gpu_capabilities", staticmethod(lambda _operation_name: ((object(), 100),))) + + sparse_q = jax.ShapeDtypeStruct((1, 4, 32, 128), jnp.bfloat16) + sparse_k = jax.ShapeDtypeStruct((1, 128, 128), jnp.bfloat16) + sparse_weights = jax.ShapeDtypeStruct((1, 4, 32), jnp.bfloat16) + sparse_lse = jax.ShapeDtypeStruct((1, 4, 32), jnp.float32) + sparse_indices = jax.ShapeDtypeStruct((1, 4, 128), jnp.int32) + predict = jax.eval_shape( + lambda q, k, weights, indices: sparse_indexer_score_recompute_wrapper( + q, + k, + weights, + indices, + target_compute_capability=100, + ), + sparse_q, + sparse_k, + sparse_weights, + sparse_indices, + )["predict"] + target = jax.eval_shape( + lambda q, k, lse, indices: sparse_attn_score_recompute_wrapper( + q, + k, + lse, + indices, + softmax_scale=0.125, + target_compute_capability=100, + ), + sparse_q, + sparse_k, + sparse_lse, + sparse_indices, + )["target"] + assert (predict.shape, predict.dtype) == ((1, 4, 128), jnp.float32) + assert (target.shape, target.dtype) == ((1, 4, 128), jnp.float32) + + dense_q = jax.ShapeDtypeStruct((1, 4, 32, 128), jnp.bfloat16) + dense_k = jax.ShapeDtypeStruct((1, 8, 1, 128), jnp.bfloat16) + dense_weights = jax.ShapeDtypeStruct((1, 4, 32), jnp.bfloat16) + dense_lse = jax.ShapeDtypeStruct((1, 4, 32), jnp.float32) + dense_indexer = jax.eval_shape( + lambda q, k, weights: dense_indexer_score_recompute_wrapper(q, k, weights, target_compute_capability=100), + dense_q, + dense_k, + dense_weights, + ) + dense_attention = jax.eval_shape( + lambda q, k, lse: dense_attn_score_recompute_wrapper( + q, + k, + lse, + softmax_scale=0.125, + target_compute_capability=100, + ), + dense_q, + dense_k, + dense_lse, + ) + for result in (dense_indexer, dense_attention): + assert (result["out"].shape, result["out"].dtype) == ((1, 4, 8), jnp.float32) + assert (result["denom"].shape, result["denom"].dtype) == ((1, 4), jnp.float32) + + +@pytest.mark.L0 +def test_jax_dense_score_recompute_jit_and_numerics(): + jax, jnp = _jax_runtime() + device, target_compute_capability = _sm100_device(jax) + from cudnn.jax import dense_attn_score_recompute_wrapper, dense_indexer_score_recompute_wrapper + + batch, seqlen_q, seqlen_k = 1, 2, 128 + num_query_heads, head_dim = 32, 128 + ratio = 1 + softmax_scale = head_dim**-0.5 + q = jax.device_put(jax.random.normal(jax.random.key(40), (batch, seqlen_q, num_query_heads, head_dim), dtype=jnp.bfloat16), device) + k = jax.device_put(jax.random.normal(jax.random.key(41), (batch, seqlen_k, 1, head_dim), dtype=jnp.bfloat16), device) + weights = jax.device_put( + jax.random.normal(jax.random.key(42), (batch, seqlen_q, num_query_heads), dtype=jnp.bfloat16) * jnp.bfloat16(1.0 / num_query_heads), + device, + ) + q_causal_offsets = jax.device_put(jnp.array([63], dtype=jnp.int32), device) + + qk = jnp.einsum("bqhd,bkd->bqhk", q.astype(jnp.float32), k[:, :, 0, :].astype(jnp.float32)) + scaled_qk = qk * softmax_scale + row_max = jnp.max(scaled_qk, axis=-1) + lse = row_max + jnp.log(jnp.sum(jnp.exp(scaled_qk - row_max[..., None]), axis=-1)) + + @jax.jit + def run_indexer(q, k, weights, offsets): + return dense_indexer_score_recompute_wrapper( + q, + k, + weights, + ratio=ratio, + q_causal_offsets=offsets, + target_compute_capability=target_compute_capability, + ) + + @jax.jit + def run_attention(q, k, lse, offsets): + return dense_attn_score_recompute_wrapper( + q, + k, + lse, + softmax_scale, + ratio=ratio, + q_causal_offsets=offsets, + target_compute_capability=target_compute_capability, + ) + + indexer_lowered = run_indexer.lower(q, k, weights, q_causal_offsets) + attention_lowered = run_attention.lower(q, k, lse, q_causal_offsets) + for lowered in (indexer_lowered, attention_lowered): + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + indexer = indexer_lowered.compile()(q, k, weights, q_causal_offsets) + attention = attention_lowered.compile()(q, k, lse, q_causal_offsets) + attention["denom"].block_until_ready() + + q_positions = jnp.arange(seqlen_q, dtype=jnp.int32)[None, :, None] + col_limits = (q_causal_offsets[:, None, None] + q_positions + 1) // ratio + valid = jnp.arange(seqlen_k, dtype=jnp.int32)[None, None, :] < col_limits + + indexer_scores = jnp.sum(jnp.maximum(qk, 0.0) * weights.astype(jnp.float32)[..., None], axis=2) + expected_indexer_out = jnp.where(valid, indexer_scores, -jnp.inf) + indexer_max = jnp.max(expected_indexer_out, axis=-1) + expected_indexer_denom = indexer_max + jnp.log(jnp.sum(jnp.exp(expected_indexer_out - indexer_max[..., None]), axis=-1)) + + attention_scores = jnp.sum(jnp.exp(scaled_qk - lse[..., None]), axis=2) + expected_attention_out = jnp.where(valid, attention_scores, -jnp.inf) + expected_attention_denom = jnp.sum(jnp.where(valid, attention_scores, 0.0), axis=-1) + + assert jnp.allclose(indexer["out"], expected_indexer_out, atol=6e-3, rtol=2e-2) + assert jnp.allclose(indexer["denom"], expected_indexer_denom, atol=6e-3, rtol=2e-2) + assert jnp.allclose(attention["out"], expected_attention_out, atol=6e-3, rtol=2e-2) + assert jnp.allclose(attention["denom"], expected_attention_denom, atol=6e-3, rtol=2e-2) + assert jnp.all(jnp.isneginf(indexer["out"][~valid])) + assert jnp.all(jnp.isneginf(attention["out"][~valid])) + + +@pytest.mark.L0 +def test_jax_sparse_indexer_score_recompute_jit(): + jax, jnp = _jax_runtime() + device, target_compute_capability = _sm100_device(jax) + from cudnn.jax import 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, + ) + indices = jax.device_put( + jnp.broadcast_to(jnp.arange(seqlen_k - 1, -1, -1, dtype=jnp.int32), (batch, seqlen_q, topk)), + device, + ) + + lowered = sparse_indexer_score_recompute_wrapper.lower( + q, + k, + weights, + indices, + target_compute_capability=target_compute_capability, + ) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + predict = lowered.compile()(q, k, weights, indices)["predict"] + predict.block_until_ready() + + k_gather = jax.vmap(lambda batch_k, batch_indices: batch_k[batch_indices])(k.astype(jnp.float32), indices) + scores_per_head = jnp.einsum("bqhd,bqtd->bqht", q.astype(jnp.float32), k_gather) + scores = jnp.sum(jnp.maximum(scores_per_head, 0.0) * weights.astype(jnp.float32)[..., None], axis=2) + expected = jax.nn.softmax(scores, axis=-1) + assert jnp.allclose(predict, expected, atol=3e-3, rtol=2e-2) + + +@pytest.mark.L0 +def test_jax_sparse_attn_score_recompute_jit(): + jax, jnp = _jax_runtime() + device, target_compute_capability = _sm100_device(jax) + from cudnn.jax import 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) + indices = jax.device_put( + jnp.broadcast_to(jnp.arange(seqlen_k - 1, -1, -1, dtype=jnp.int32), (batch, seqlen_q, topk)), + device, + ) + full_scores = jnp.einsum("bqhd,bkd->bqhk", q.astype(jnp.float32), k.astype(jnp.float32)) * softmax_scale + row_max = jnp.max(full_scores, axis=-1) + lse = row_max + jnp.log(jnp.sum(jnp.exp(full_scores - row_max[..., None]), axis=-1)) + + lowered = sparse_attn_score_recompute_wrapper.lower( + q, + k, + lse, + indices, + softmax_scale=softmax_scale, + target_compute_capability=target_compute_capability, + ) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + target = lowered.compile()(q, k, lse, indices)["target"] + target.block_until_ready() + + k_gather = jax.vmap(lambda batch_k, batch_indices: batch_k[batch_indices])(k.astype(jnp.float32), indices) + sparse_scores = jnp.einsum("bqhd,bqtd->bqht", q.astype(jnp.float32), k_gather) * softmax_scale + expected = jnp.sum(jnp.exp(sparse_scores - lse[..., None]), axis=2) + expected /= jnp.sum(expected, axis=-1, keepdims=True) + assert jnp.allclose(target, expected, atol=3e-3, rtol=2e-2) + + +@pytest.mark.L0 +def test_jax_sparse_score_recompute_global_indices_and_partial_lengths(): + jax, jnp = _jax_runtime() + device, target_compute_capability = _sm100_device(jax) + from cudnn.jax import sparse_attn_score_recompute_wrapper, sparse_indexer_score_recompute_wrapper + + batch, seqlen_q, seqlen_k = 2, 1, 128 + num_query_heads, head_dim, topk = 32, 128, 128 + softmax_scale = head_dim**-0.5 + q = jax.device_put(jax.random.normal(jax.random.key(30), (batch, seqlen_q, num_query_heads, head_dim), dtype=jnp.bfloat16), device) + k = jax.device_put(jax.random.normal(jax.random.key(31), (batch, seqlen_k, head_dim), dtype=jnp.bfloat16), device) + weights = jax.device_put( + jax.random.normal(jax.random.key(32), (batch, seqlen_q, num_query_heads), dtype=jnp.bfloat16) * jnp.bfloat16(1.0 / num_query_heads), + device, + ) + local_indices = jnp.broadcast_to(jnp.arange(topk, dtype=jnp.int32), (batch, seqlen_q, topk)) + global_indices = jax.device_put(local_indices + jnp.arange(batch, dtype=jnp.int32)[:, None, None] * seqlen_k, device) + topk_length = jax.device_put(jnp.array([[0], [73]], dtype=jnp.int32), device) + + full_scores = jnp.einsum("bqhd,bkd->bqhk", q.astype(jnp.float32), k.astype(jnp.float32)) * softmax_scale + row_max = jnp.max(full_scores, axis=-1) + lse = row_max + jnp.log(jnp.sum(jnp.exp(full_scores - row_max[..., None]), axis=-1)) + + @jax.jit + def run(q, k, weights, lse, indices, lengths): + predict = sparse_indexer_score_recompute_wrapper( + q, + k, + weights, + indices, + topk_length=lengths, + topk_indices_global=True, + target_compute_capability=target_compute_capability, + )["predict"] + target = sparse_attn_score_recompute_wrapper( + q, + k, + lse, + indices, + softmax_scale, + topk_length=lengths, + topk_indices_global=True, + target_compute_capability=target_compute_capability, + )["target"] + return predict, target + + lowered = run.lower(q, k, weights, lse, global_indices, topk_length) + assert lowered.as_text("stablehlo").count("stablehlo.custom_call") == 2 + predict, target = lowered.compile()(q, k, weights, lse, global_indices, topk_length) + target.block_until_ready() + + k_gather = jax.vmap(lambda batch_k, batch_indices: batch_k[batch_indices])(k.astype(jnp.float32), local_indices) + qk = jnp.einsum("bqhd,bqtd->bqht", q.astype(jnp.float32), k_gather) + valid = jnp.arange(topk)[None, None, :] < topk_length[..., None] + + indexer_scores = jnp.sum(jnp.maximum(qk, 0.0) * weights.astype(jnp.float32)[..., None], axis=2) + expected_predict = jax.nn.softmax(jnp.where(valid, indexer_scores, -jnp.inf), axis=-1) + expected_predict = jnp.where(topk_length[..., None] > 0, expected_predict, 0.0) + + attention_scores = jnp.sum(jnp.exp(qk * softmax_scale - lse[..., None]), axis=2) + attention_scores = jnp.where(valid, attention_scores, 0.0) + attention_denom = jnp.sum(attention_scores, axis=-1, keepdims=True) + expected_target = jnp.where(attention_denom > 0, attention_scores / jnp.maximum(attention_denom, 1e-12), 0.0) + + assert jnp.allclose(predict, expected_predict, atol=3e-3, rtol=2e-2) + assert jnp.allclose(target, expected_target, atol=3e-3, rtol=2e-2) + assert jnp.array_equal(predict[0], jnp.zeros_like(predict[0])) + assert jnp.array_equal(target[0], jnp.zeros_like(target[0])) + assert jnp.array_equal(predict[1, :, 73:], jnp.zeros_like(predict[1, :, 73:])) + assert jnp.array_equal(target[1, :, 73:], jnp.zeros_like(target[1, :, 73:])) diff --git a/test/python/fe_api/test_jax_dsa_score_recompute_contract.py b/test/python/fe_api/test_jax_dsa_score_recompute_contract.py new file mode 100644 index 000000000..e2b8a743a --- /dev/null +++ b/test/python/fe_api/test_jax_dsa_score_recompute_contract.py @@ -0,0 +1,706 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for JAX DSA score-recompute adapters.""" + +from enum import Enum, auto +import ast +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_DSA_ROOT = _CUDNN_ROOT / "deepseek_sparse_attention" +_SCORE_ROOT = _DSA_ROOT / "score_recompute" +_PACKAGE = "cudnn_jax_dsa_score_contract_test" + + +class _DataType(Enum): + NOT_SET = auto() + FLOAT = auto() + BFLOAT16 = auto() + INT32 = auto() + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + @property + def ndim(self): + return len(self.shape) + + +class _TupleDict(dict): + pass + + +class _TensorSpec: + def __init__(self, *, mode=None, divisibility=None): + self.mode = mode + self.divisibility = divisibility + + +def _identity_jit(fn=None, **_kwargs): + return (lambda decorated: decorated) if fn is None else fn + + +class JaxDsaScoreRecomputeContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + cls.tensor = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.bfloat16 = _DataType.BFLOAT16 + fake_jnp.float32 = _DataType.FLOAT + fake_jnp.int32 = _DataType.INT32 + fake_jnp.expand_dims = lambda value, axis: _Array( + value.shape[:axis] + (1,) + value.shape[axis:], value.dtype + ) + fake_jnp.transpose = lambda value, axes: _Array( + tuple(value.shape[axis] for axis in axes), value.dtype + ) + + fake_jax = types.ModuleType("jax") + fake_jax.__path__ = [] + fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + fake_jax.numpy = fake_jnp + fake_jax.jit = _identity_jit + fake_jax.ShapeDtypeStruct = _Array + + tensor_module = cls.tensor + + class JaxTensorDesc(tensor_module.TensorDesc): + @classmethod + def from_shape( + cls, + shape, + dtype, + *, + name="", + mode=None, + public_stride_order=None, + init_value=None, + ): + return FakeJaxApiBase._to_tensor_desc( + _Array(shape, dtype), + name, + mode=mode, + init_value=init_value, + ) + + @property + def cudnn_dtype(self): + return self.dtype + + class FakeJaxApiBase: + @staticmethod + def _to_tensor_desc(value, name, *, mode=None, init_value=None): + public_shape = tuple(value.shape) + mode = tuple(range(len(public_shape))) if mode is None else tuple(mode) + public_stride = [0] * len(public_shape) + running = 1 + for dimension in reversed(range(len(public_shape))): + public_stride[dimension] = running + running *= max(public_shape[dimension], 1) + canonical_axis_by_public = [0] * len(mode) + for canonical_axis, public_axis in enumerate(mode): + canonical_axis_by_public[public_axis] = canonical_axis + desc = JaxTensorDesc( + dtype=value.dtype, + shape=tuple(public_shape[axis] for axis in mode), + stride=tuple(public_stride[axis] for axis in mode), + stride_order=tuple( + canonical_axis_by_public[axis] + for axis in reversed(range(len(public_shape))) + ), + name=name, + init_value=init_value, + ) + object.__setattr__(desc, "mode", mode) + return desc + + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del operation_name + resolved = 100 if target is None else target + if resolved not in supported: + raise ValueError(f"unsupported target {resolved}") + return resolved + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + mode = expected.mode if mode is None else tuple(mode) + if tuple(value.shape[axis] for axis in mode) != expected.shape: + raise ValueError(f"{expected.name} tensor shape mismatch") + if value.dtype != expected.cudnn_dtype: + raise ValueError(f"{expected.name} tensor dtype mismatch") + + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + if mode is None: + mode = desc.mode + return _TensorSpec(mode=mode, divisibility=divisibility) + + def _call_kernel( + self, inputs, *, launch, output_descs, workspace_descs=(), **options + ): + input_descs = options.get("input_descs") + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc) + options["input_spec"] = tuple( + self._to_tensor_spec(desc) for desc in input_descs + ) + options["output_spec"] = tuple( + self._to_tensor_spec(desc) for desc in output_descs + ) + self.captured_call = { + "inputs": tuple(inputs), + "outputs": tuple(output_descs), + "workspaces": tuple(workspace_descs), + "launch": launch, + **options, + } + specs = options.get("output_spec") or (None,) * len(output_descs) + outputs = [] + for desc, spec in zip(output_descs, specs): + mode = ( + desc.mode + if spec is None or spec.mode is None + else tuple(spec.mode) + ) + canonical_axis_by_public = [0] * len(mode) + for canonical_axis, public_axis in enumerate(mode): + canonical_axis_by_public[public_axis] = canonical_axis + outputs.append( + _Array( + tuple( + desc.shape[canonical_axis_by_public[public_axis]] + for public_axis in range(desc.ndim) + ), + desc.cudnn_dtype, + ) + ) + return tuple(outputs) + + fake_internal_jax = types.ModuleType(f"{_PACKAGE}._jax") + fake_internal_jax.__path__ = [str(_CUDNN_ROOT / "_jax")] + fake_internal_jax.__package__ = fake_internal_jax.__name__ + fake_internal_jax.__spec__ = ModuleSpec( + fake_internal_jax.__name__, loader=None, is_package=True + ) + fake_internal_jax.JaxApiBase = FakeJaxApiBase + fake_internal_jax.JaxTensorDesc = JaxTensorDesc + fake_internal_jax.TupleDict = _TupleDict + + dsa_name = f"{_PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_DSA_ROOT)] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + + score_name = f"{dsa_name}.score_recompute" + score = types.ModuleType(score_name) + score.__path__ = [str(_SCORE_ROOT)] + score.__package__ = score_name + score.__spec__ = ModuleSpec(score_name, loader=None, is_package=True) + + utils_name = f"{dsa_name}.utils" + utils = types.ModuleType(utils_name) + utils.__path__ = [] + cls.optional_modules = { + "jax": fake_jax, + "jax.numpy": fake_jnp, + f"{_PACKAGE}._jax": fake_internal_jax, + dsa_name: dsa, + score_name: score, + utils_name: utils, + } + sys.modules.update(cls.optional_modules) + try: + cls.adapter = importlib.import_module(f"{score_name}.jax") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + for name in ("jax.numpy", "jax"): + module = sys.modules.get(name) + if module is cls.optional_modules.get(name): + sys.modules.pop(name, None) + + def array(self, shape, dtype=_DataType.BFLOAT16): + return _Array(shape, dtype) + + def sparse_samples(self): + return ( + self.array((2, 4, 32, 128)), + self.array((2, 256, 128)), + self.array((2, 4, 32)), + self.array((2, 4, 128), _DataType.INT32), + ) + + def sequence_major_sparse_samples(self): + return ( + self.array((4, 2, 32, 128)), + self.array((256, 2, 128)), + self.array((4, 2, 32)), + self.array((4, 2, 128), _DataType.INT32), + ) + + def dense_samples(self, *, attention=False): + return ( + self.array((2, 4, 32, 128)), + self.array((2, 8, 1, 128)), + self.array( + (2, 4, 32), _DataType.FLOAT if attention else _DataType.BFLOAT16 + ), + ) + + def sequence_major_dense_samples(self): + return ( + self.array((4, 2, 32, 128)), + self.array((8, 2, 1, 128)), + self.array((4, 2, 32)), + ) + + def test_sm100_sparse_infers_output_and_hidden_length_workspace(self): + q, k, weights, indices = self.sparse_samples() + api = self.adapter.SparseIndexerScoreRecompute( + q, k, weights, indices, target_compute_capability=100 + ) + result = api(q, k, weights, indices) + + self.assertIsInstance(result, _TupleDict) + self.assertEqual( + (result["predict"].shape, result["predict"].dtype), + ((2, 4, 128), _DataType.FLOAT), + ) + self.assertEqual(len(api.captured_call["workspaces"]), 1) + workspace = api.captured_call["workspaces"][0] + self.assertEqual( + (workspace.shape, workspace.cudnn_dtype), ((1, 1), _DataType.INT32) + ) + self.assertEqual(api.captured_call["compile_options"], "--gpu-arch sm_100a") + + def test_explicit_topk_length_remains_a_real_input(self): + q, k, _, indices = self.sparse_samples() + lse = self.array((2, 4, 32), _DataType.FLOAT) + lengths = self.array((2, 4), _DataType.INT32) + api = self.adapter.SparseAttnScoreRecompute( + q, + k, + lse, + indices, + 0.125, + sample_topk_length=lengths, + target_compute_capability=100, + ) + result = api(q, k, lse, indices, lengths) + + self.assertEqual(result["target"].shape, (2, 4, 128)) + self.assertEqual(api.captured_call["workspaces"], ()) + self.assertIs(api.captured_call["inputs"][-1], lengths) + self.assertTrue(api._op.config.have_topk_length) + + def test_sm90_sparse_adapts_k_and_per_head_layouts(self): + q, k, weights, indices = self.sparse_samples() + api = self.adapter.SparseIndexerScoreRecompute( + q, k, weights, indices, target_compute_capability=90 + ) + api(q, k, weights, indices) + + _, kernel_k, kernel_weights, _ = api.captured_call["inputs"] + self.assertEqual(kernel_k.shape, (2, 256, 1, 128)) + self.assertEqual(kernel_weights.shape, (2, 32, 4)) + self.assertEqual(api.captured_call["workspaces"], ()) + + def test_sparse_sequence_major_layouts_preserve_public_shapes(self): + q, k, weights, indices = self.sequence_major_sparse_samples() + api = self.adapter.SparseIndexerScoreRecompute( + q, + k, + weights, + indices, + q_layout="SBHD", + k_layout="SBD", + per_head_layout="SBH", + score_layout="SBK", + output_layout="SBK", + target_compute_capability=90, + ) + result = api(q, k, weights, indices) + + self.assertEqual(result["predict"].shape, indices.shape) + _, kernel_k, kernel_weights, _ = api.captured_call["inputs"] + self.assertEqual(kernel_k.shape, (256, 2, 1, 128)) + self.assertEqual(kernel_weights.shape, (2, 32, 4)) + self.assertEqual( + tuple(spec.mode for spec in api.captured_call["input_spec"]), + ((1, 0, 2, 3), (1, 0, 2, 3), (0, 1, 2), (1, 0, 2)), + ) + self.assertEqual( + api.captured_call["output_spec"][0].mode, + (1, 0, 2), + ) + + def test_dense_bshd_outputs_are_functional_and_initialized(self): + q, k, weights = self.dense_samples() + offsets = self.array((2,), _DataType.INT32) + api = self.adapter.DenseIndexerScoreRecompute( + q, + k, + weights, + ratio=2, + sample_q_causal_offsets=offsets, + target_compute_capability=100, + ) + result = api(q, k, weights, q_causal_offsets=offsets) + + self.assertEqual(result["out"].shape, (2, 4, 8)) + self.assertEqual(result["denom"].shape, (2, 4)) + self.assertEqual(api.out_desc.init_value, float("-inf")) + self.assertIsNone(api.denom_desc.init_value) + self.assertIs(api.captured_call["inputs"][-1], offsets) + + explicit = self.adapter.DenseIndexerScoreRecompute( + q, + k, + weights, + sample_out=self.array((2, 4, 8), _DataType.FLOAT), + sample_denom_out=self.array((2, 4), _DataType.FLOAT), + target_compute_capability=100, + ) + self.assertEqual(explicit.out_desc.init_value, float("-inf")) + self.assertIsNone(explicit.denom_desc.init_value) + + def test_dense_sequence_major_layouts_preserve_public_shapes(self): + q, k, weights = self.sequence_major_dense_samples() + api = self.adapter.DenseIndexerScoreRecompute( + q, + k, + weights, + q_layout="SBHD", + k_layout="SBHD", + per_head_layout="SBH", + output_layout="SBK", + denom_layout="SB", + target_compute_capability=100, + ) + result = api(q, k, weights) + + self.assertEqual(result["out"].shape, (4, 2, 8)) + self.assertEqual(result["denom"].shape, (4, 2)) + self.assertEqual( + tuple(spec.mode for spec in api.captured_call["input_spec"]), + ((1, 0, 2, 3), (1, 0, 2, 3), (1, 0, 2)), + ) + self.assertEqual( + tuple(spec.mode for spec in api.captured_call["output_spec"]), + ((1, 0, 2), (1, 0)), + ) + + def test_dense_thd_sm100_and_sm90_rejection(self): + q = self.array((7, 32, 128)) + k = self.array((11, 1, 128)) + lse = self.array((7, 32), _DataType.FLOAT) + cu_q = self.array((3,), _DataType.INT32) + cu_k = self.array((3,), _DataType.INT32) + api = self.adapter.DenseAttnScoreRecompute( + q, + k, + lse, + 0.125, + sample_cu_seqlens_q=cu_q, + sample_cu_seqlens_k=cu_k, + max_seqlen_q=4, + max_seqlen_k=8, + target_compute_capability=100, + ) + result = api(q, k, lse, cu_q, cu_k) + self.assertEqual((result["out"].shape, result["denom"].shape), ((7, 8), (7,))) + self.assertEqual(api.captured_call["inputs"][-2:], (cu_q, cu_k)) + + sm90 = self.adapter.DenseAttnScoreRecompute( + q, + k, + lse, + 0.125, + sample_cu_seqlens_q=cu_q, + sample_cu_seqlens_k=cu_k, + max_seqlen_q=4, + max_seqlen_k=8, + target_compute_capability=90, + ) + with self.assertRaisesRegex(NotImplementedError, "cannot be traced by JAX"): + sm90(q, k, lse, cu_q, cu_k) + + def test_explicit_output_samples_are_cross_checked(self): + q, k, weights = self.dense_samples() + bad_output = self.array((2, 4, 4), _DataType.FLOAT) + denominator = self.array((2, 4), _DataType.FLOAT) + api = self.adapter.DenseIndexerScoreRecompute( + q, + k, + weights, + sample_out=bad_output, + sample_denom_out=denominator, + target_compute_capability=100, + ) + with self.assertRaisesRegex(ValueError, "output must have shape"): + api(q, k, weights) + + def test_high_level_wrapper_infers_shapes(self): + q, k, weights, indices = self.sparse_samples() + with mock.patch.dict(sys.modules, self.optional_modules): + result = self.adapter.sparse_indexer_score_recompute_wrapper( + q, + k, + weights, + indices, + target_compute_capability=100, + ) + self.assertEqual(result["predict"].shape, indices.shape) + + def test_sparse_launchers_preserve_native_argument_order(self): + calls = [] + + class FakeSm90Kernel: + def __init__(self, dtype, **options): + calls.append(("sm90_init", dtype, options)) + + def __call__(self, *args): + calls.append(("sm90_call", args)) + + class FakeSm100Kernel: + def __init__(self, **options): + calls.append(("sm100_init", options)) + + def __call__(self, *args): + calls.append(("sm100_call", args)) + + cutlass = types.ModuleType("cutlass") + cutlass.BFloat16 = "BFloat16" + cutlass.Float32 = lambda value: ("Float32", value) + package = self.adapter.__package__ + sm90_module = types.ModuleType(f"{package}.sparse_score_recompute_sm90") + sm90_module.SparseScoreRecomputeSm90 = FakeSm90Kernel + sm100_module = types.ModuleType(f"{package}.sparse_score_recompute_sm100") + sm100_module.SparseScoreRecomputeSm100 = FakeSm100Kernel + + q, k, weights, indices = self.sparse_samples() + sm90 = self.adapter.SparseIndexerScoreRecompute( + q, k, weights, indices, target_compute_capability=90 + ) + sm90.check_support() + sm100 = self.adapter.SparseIndexerScoreRecompute( + q, k, weights, indices, target_compute_capability=100 + ) + sm100.check_support() + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + sm90_module.__name__: sm90_module, + sm100_module.__name__: sm100_module, + }, + ): + sm90._launch_kernel("stream90", "q90", "k90", "aux90", "indices90", "out90") + sm100._launch_kernel( + "stream100", "q100", "k100", "aux100", "indices100", "out100", "dummy" + ) + + self.assertEqual( + [entry for entry in calls if entry[0].endswith("call")], + [ + ( + "sm90_call", + ( + "q90", + "k90", + "indices90", + "stream90", + "out90", + "aux90", + None, + None, + ), + ), + ( + "sm100_call", + ( + "q100", + "k100", + "aux100", + "indices100", + "out100", + "dummy", + ("Float32", 1.0), + "stream100", + ), + ), + ], + ) + + def test_dense_launchers_preserve_native_argument_order(self): + calls = [] + + class FakeSm90Kernel: + def __init__(self, dtype, **options): + calls.append(("sm90_init", dtype, options)) + + def __call__(self, *args): + calls.append(("sm90_call", args)) + + class FakeSm100Kernel: + def __init__(self, **options): + calls.append(("sm100_init", options)) + + def __call__(self, *args): + calls.append(("sm100_call", args)) + + cutlass = types.ModuleType("cutlass") + cutlass.BFloat16 = "BFloat16" + cutlass.Float32 = lambda value: ("Float32", value) + cutlass.Int32 = lambda value: ("Int32", value) + package = self.adapter.__package__ + sm90_module = types.ModuleType(f"{package}.dense_score_recompute_sm90") + sm90_module.DenseScoreRecomputeSm90 = FakeSm90Kernel + sm100_module = types.ModuleType(f"{package}.dense_score_recompute_sm100") + sm100_module.DenseScoreRecomputeSm100 = FakeSm100Kernel + + q, k, weights = self.dense_samples() + offsets = self.array((2,), _DataType.INT32) + sm90 = self.adapter.DenseIndexerScoreRecompute( + q, + k, + weights, + sample_q_causal_offsets=offsets, + target_compute_capability=90, + ) + sm90.check_support() + sm100 = self.adapter.DenseIndexerScoreRecompute( + q, + k, + weights, + sample_q_causal_offsets=offsets, + target_compute_capability=100, + ) + sm100.check_support() + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + sm90_module.__name__: sm90_module, + sm100_module.__name__: sm100_module, + }, + ): + sm90._launch_kernel( + "stream90", "q90", "k90", "aux90", "offset90", "out90", "denom90" + ) + sm100._launch_kernel( + "stream100", "q100", "k100", "aux100", "offset100", "out100", "denom100" + ) + + self.assertEqual( + [entry for entry in calls if entry[0].endswith("call")], + [ + ( + "sm90_call", + ( + "q90", + "k90", + None, + "stream90", + "out90", + "aux90", + None, + "denom90", + "offset90", + ), + ), + ( + "sm100_call", + ( + "q100", + "k100", + "aux100", + "out100", + "denom100", + ("Float32", 1.0), + ("Int32", 4), + ("Int32", 8), + None, + None, + "offset100", + "stream100", + ), + ), + ], + ) + + def test_common_and_jax_modules_keep_framework_import_boundaries(self): + def imports_framework(path, framework): + tree = ast.parse(path.read_text(), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import) and any( + alias.name == framework or alias.name.startswith(f"{framework}.") + for alias in node.names + ): + return True + if ( + isinstance(node, ast.ImportFrom) + and node.level == 0 + and ( + node.module == framework + or (node.module or "").startswith(f"{framework}.") + ) + ): + return True + return False + + for filename in ("config.py", "op.py"): + path = _SCORE_ROOT / filename + for framework in ("torch", "jax", "cutlass", "cuda"): + self.assertFalse( + imports_framework(path, framework), + f"{filename} imports {framework}", + ) + self.assertFalse(imports_framework(_SCORE_ROOT / "jax.py", "torch")) + + init_tree = ast.parse((_SCORE_ROOT / "__init__.py").read_text()) + imported_modules = { + node.module + for node in ast.walk(init_tree) + if isinstance(node, ast.ImportFrom) + } + self.assertNotIn("api", imported_modules) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py b/test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py new file mode 100644 index 000000000..8ba4f7beb --- /dev/null +++ b/test/python/fe_api/test_jax_dsa_sparse_attention_backward_contract.py @@ -0,0 +1,401 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX contracts for DSA sparse-attention backward.""" + +from __future__ import annotations + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_DSA_ROOT = _CUDNN_ROOT / "deepseek_sparse_attention" +_OPERATION_ROOT = _DSA_ROOT / "sparse_attention_backward" +_PACKAGE = "cudnn_jax_dsa_sparse_attention_backward_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + UINT8 = auto() + + +_DTYPE_TO_CUDNN = { + "float16": _DataType.HALF, + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, + "int32": _DataType.INT32, + "uint8": _DataType.UINT8, +} +_CUDNN_TO_DTYPE = {value: key for key, value in _DTYPE_TO_CUDNN.items()} + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + +class _TupleDict(dict): + def __iter__(self): + return iter(self.values()) + + +class _TensorSpec: + def __init__(self, *, mode=None, divisibility=None): + self.mode = mode + self.divisibility = divisibility + + +class JaxSparseAttentionBackwardContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + dsa_name = f"{_PACKAGE}.deepseek_sparse_attention" + dsa = types.ModuleType(dsa_name) + dsa.__path__ = [str(_DSA_ROOT)] + dsa.__package__ = dsa_name + dsa.__spec__ = ModuleSpec(dsa_name, loader=None, is_package=True) + sys.modules[dsa_name] = dsa + + operation_name = f"{dsa_name}.sparse_attention_backward" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + + class JaxTensorDesc(tensor_module.TensorDesc): + @property + def cudnn_dtype(self): + return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) + + def compact_like(self, *, cudnn_dtype, shape, stride_order=None, name="", init_value=None): + canonical = tensor_module.make_compact_tensor_desc( + dtype=cudnn_dtype, + shape=shape, + stride_order=stride_order, + name=name, + init_value=init_value, + ) + return JaxTensorDesc( + dtype=_CUDNN_TO_DTYPE[cudnn_dtype], + shape=canonical.shape, + stride=canonical.stride, + stride_order=canonical.stride_order, + name=name, + init_value=init_value, + ) + + def with_divisibility(self, divisibility): + desc = JaxTensorDesc( + dtype=self.dtype, + shape=self.shape, + stride=self.stride, + stride_order=self.stride_order, + name=self.name, + init_value=self.init_value, + ) + object.__setattr__(desc, "divisibility", tuple(divisibility)) + return desc + + class JaxApiBase: + captured_call = None + + @staticmethod + def _to_tensor_desc(value, name, *, mode=None, init_value=None): + del mode + shape = tuple(value.shape) + stride = [] + running = 1 + for size in reversed(shape): + stride.append(running) + running *= max(size, 1) + stride = tuple(reversed(stride)) + return JaxTensorDesc( + dtype=value.dtype, + shape=shape, + stride=stride, + stride_order=tuple(reversed(range(len(shape)))), + name=name, + init_value=init_value, + ) + + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del supported, operation_name + if target is None: + raise RuntimeError("target required by test") + return target + + @staticmethod + def _compute_capability_family(target, supported): + return max((value for value in supported if value <= target), default=None) + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + del mode + if tuple(value.shape) != expected.shape or _DTYPE_TO_CUDNN[value.dtype] != expected.cudnn_dtype: + raise ValueError(f"{expected.name} tensor signature mismatch") + + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + del desc + return _TensorSpec(mode=mode, divisibility=divisibility) + + def _call_kernel(self, inputs, *, launch, **options): + options["launch"] = launch + JaxApiBase.captured_call = (inputs, options) + return tuple(_Array(desc.shape, _CUDNN_TO_DTYPE[desc.cudnn_dtype]) for desc in options["output_descs"]) + + internal = types.ModuleType(f"{_PACKAGE}._jax") + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal.__name__ + internal.__spec__ = ModuleSpec(internal.__name__, loader=None, is_package=True) + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = _TupleDict + sys.modules[internal.__name__] = internal + + fake_jax = types.ModuleType("jax") + fake_jax.jit = lambda function=None, **_kwargs: (lambda fn: fn) if function is None else function + fake_jax.ShapeDtypeStruct = _Array + + try: + with mock.patch.dict(sys.modules, {"jax": fake_jax}): + cls.module = importlib.import_module(f"{operation_name}.jax") + except Exception: + cls.tearDownClass() + raise + cls.JaxApiBase = JaxApiBase + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + @staticmethod + def _inputs(*, head_dim=512, dtype="bfloat16", with_length=False): + m, n, h, topk = 65, 130, 64, 32 + head_dim_v = 512 if head_dim == 576 else head_dim + values = ( + _Array((m, h, head_dim), dtype), + _Array((n, head_dim), dtype), + _Array((m, h, head_dim_v), dtype), + _Array((m, h, head_dim_v), dtype), + _Array((m, h), "float32"), + _Array((h,), "float32"), + _Array((m, topk), "int32"), + ) + return (*values, _Array((m,), "int32") if with_length else None) + + def test_sm100_declares_functional_outputs_and_zeroed_workspaces(self): + inputs = self._inputs(with_length=True) + api = self.module.SparseAttentionBackward( + *inputs[:-1], + sample_topk_length=inputs[-1], + target_compute_capability=103, + ) + result = api(*inputs[:-1], inputs[-1]) + self.assertEqual(tuple(result), (result["dq"], result["dkv"], result["d_sink"])) + self.assertEqual(result["dq"].shape, inputs[0].shape) + self.assertEqual(result["dkv"].shape, inputs[1].shape) + self.assertEqual(result["d_sink"].shape, inputs[5].shape) + + _, options = self.JaxApiBase.captured_call + outputs = options["output_descs"] + self.assertIsNone(outputs[0].init_value) + self.assertIsNone(outputs[1].init_value) + self.assertEqual(outputs[2].init_value, 0.0) + workspaces = options["workspace_descs"] + self.assertEqual( + [(desc.shape, desc.cudnn_dtype, desc.init_value) for desc in workspaces], + [ + ((1, 64, 72, 8), _DataType.UINT8, 0), + ((1, 1, 136, 2048), _DataType.UINT8, 0), + ], + ) + self.assertIn("--gpu-arch sm_103a", options["compile_options"]) + + def test_sm90_declares_pipeline_workspaces_and_optional_length_dummy(self): + inputs = self._inputs(dtype="float16") + api = self.module.SparseAttentionBackward(*inputs[:-1], target_compute_capability=90) + api(*inputs[:-1]) + _, options = self.JaxApiBase.captured_call + self.assertEqual( + [(desc.shape, desc.cudnn_dtype, desc.init_value) for desc in options["workspace_descs"]], + [ + ((1, 128, 64), _DataType.FLOAT, None), + ((1, 128, 64), _DataType.FLOAT, None), + ((1, 1, 192 * 512), _DataType.FLOAT, 0.0), + ((1,), _DataType.INT32, None), + ], + ) + self.assertIn("--gpu-arch sm_90a", options["compile_options"]) + + def test_target_specific_dtype_and_576_shape_validation(self): + fp16 = self._inputs(dtype="float16") + with self.assertRaisesRegex(ValueError, "SM100.*bfloat16"): + self.module.SparseAttentionBackward(*fp16[:-1], target_compute_capability=100)(*fp16[:-1]) + + inputs_576 = self._inputs(head_dim=576) + result = self.module.SparseAttentionBackward(*inputs_576[:-1], target_compute_capability=100)(*inputs_576[:-1]) + self.assertEqual(result["dq"].shape, (65, 64, 576)) + self.assertEqual(result["dkv"].shape, (130, 576)) + + def test_sm100_launcher_passes_max_topk_to_kernel(self): + inputs = self._inputs() + api = self.module.SparseAttentionBackward(*inputs[:-1], target_compute_capability=100) + api.check_support() + seen = {} + + class Kernel: + def __init__(self, **configuration): + seen["configuration"] = configuration + + def __call__(self, *arguments): + seen["arguments"] = arguments + + kernel_module = types.ModuleType(f"{self.module.__package__}.dsa_bwd_sm100") + kernel_module.FlashAttentionDSABackwardSm100 = Kernel + cutlass = types.ModuleType("cutlass") + cutlass.Int32 = int + cutlass.Float32 = float + placeholders = tuple(object() for _ in range(12)) + with mock.patch.dict(sys.modules, {"cutlass": cutlass, kernel_module.__name__: kernel_module}): + api._launch_sm100(placeholders[:7], placeholders[7:10], placeholders[10:12], object()) + + self.assertEqual(seen["configuration"]["max_topk"], 32) + self.assertEqual(seen["configuration"]["head_dim"], 512) + self.assertIsNone(seen["arguments"][8]) + + def test_sm90_launcher_builds_flat_mqa_views_and_orders_pipeline(self): + inputs = self._inputs() + api = self.module.SparseAttentionBackward(*inputs[:-1], target_compute_capability=90) + api.check_support() + seen = {} + + class CuteTensor: + def __init__(self, shape, stride, *, element_type="BFloat16", label=""): + self.shape = tuple(shape) + self.stride = tuple(stride) + self.element_type = element_type + self.iterator = self + self.label = label + + def compact(shape, *, element_type="BFloat16", label=""): + stride = [] + running = 1 + for size in reversed(shape): + stride.append(running) + running *= max(size, 1) + return CuteTensor(shape, tuple(reversed(stride)), element_type=element_type, label=label) + + class Stage: + def __init__(self, label, *configuration, **keywords): + seen[f"{label}_configuration"] = (configuration, keywords) + self.label = label + + def __call__(self, *arguments): + seen[f"{self.label}_arguments"] = arguments + + kernel_module = types.ModuleType(f"{self.module.__package__}.dsa_bwd_sm90") + kernel_module._FlashAttentionDSABackwardPreprocessSm90 = lambda *args, **kwargs: Stage("pre", *args, **kwargs) + kernel_module.FlashAttentionDSABackwardSm90 = lambda *args, **kwargs: Stage("main", *args, **kwargs) + kernel_module._FlashAttentionDSABackwardPostprocessSm90 = lambda *args, **kwargs: Stage("post", *args, **kwargs) + + cute = types.ModuleType("cutlass.cute") + cute.make_layout = lambda shape, stride: types.SimpleNamespace(shape=tuple(shape), stride=tuple(stride)) + cute.make_tensor = lambda iterator, layout: CuteTensor( + layout.shape, + layout.stride, + element_type=iterator.element_type, + label=iterator.label, + ) + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass.Int32 = int + cutlass.Float32 = float + cutlass.cute = cute + + kernel_inputs = ( + compact((65, 64, 512), label="q"), + compact((130, 512), label="kv"), + compact((65, 64, 512), label="out"), + compact((65, 64, 512), label="dout"), + compact((65, 64), element_type="Float32", label="lse"), + compact((64,), element_type="Float32", label="sink"), + compact((65, 32), element_type="Int32", label="topk"), + ) + kernel_outputs = ( + compact((65, 64, 512), label="dq"), + compact((130, 512), label="dkv"), + compact((64,), element_type="Float32", label="dsink"), + ) + kernel_workspaces = ( + compact((1, 128, 64), element_type="Float32", label="dpsum"), + compact((1, 128, 64), element_type="Float32", label="lse_log2"), + compact((1, 1, 192 * 512), element_type="Float32", label="dkv_accum"), + compact((1,), element_type="Int32", label="dummy_length"), + ) + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + "cutlass.cute": cute, + kernel_module.__name__: kernel_module, + }, + ): + api._launch_sm90(kernel_inputs, kernel_outputs, kernel_workspaces, object()) + + pre_args = seen["pre_arguments"] + self.assertEqual(pre_args[0].shape, (1, 65, 64, 512)) + self.assertEqual(pre_args[1].shape, (1, 65, 64, 512)) + main_args = seen["main_arguments"] + self.assertEqual(main_args[0].shape, (1, 65, 64, 512)) + self.assertEqual(main_args[1].shape, (1, 130, 1, 512)) + self.assertEqual(main_args[7].shape, (1, 65, 32)) + self.assertEqual(main_args[8].shape, (1, 1)) + main_config, main_keywords = seen["main_configuration"] + self.assertEqual(main_config[:4], ("BFloat16", 512, 512, 64)) + self.assertEqual(main_keywords["max_topk"], 32) + self.assertFalse(main_keywords["have_topk_length"]) + post_args = seen["post_arguments"] + self.assertEqual(post_args[1].shape, (1, 130, 1, 512)) + self.assertEqual(post_args[2], 130) + + def test_jax_module_has_no_torch_import(self): + source = (_OPERATION_ROOT / "jax.py").read_text() + self.assertNotIn("import torch", source) + for kernel_name in ("dsa_bwd_sm90.py", "dsa_bwd_sm100.py"): + self.assertNotIn("import torch", (_OPERATION_ROOT / kernel_name).read_text()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_gemm_amax.py b/test/python/fe_api/test_jax_gemm_amax.py new file mode 100644 index 000000000..bb21cc7be --- /dev/null +++ b/test/python/fe_api/test_jax_gemm_amax.py @@ -0,0 +1,196 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX integration tests for block-scaled dense GEMM + amax.""" + +import pytest + + +def _jax_runtime(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + cutlass_jax = pytest.importorskip("cutlass.jax") + if not cutlass_jax.is_available(): + pytest.skip("Installed JAX version is unsupported by CUTLASS JAX") + if not hasattr(jnp, "float8_e8m0fnu"): + pytest.skip("JAX does not provide the E8M0 scale-factor dtype") + return jax, jnp + + +def _compute_capability(device): + reported = getattr(device, "compute_capability", None) + if isinstance(reported, (tuple, list)): + major, minor = int(reported[0]), int(reported[1]) + else: + text = str(reported) + if "." in text: + major_text, minor_text = text.split(".", 1) + major, minor = int(major_text), int(minor_text) + else: + major, minor = divmod(int(text), 10) + return major * 10 + minor + + +def _supported_gpu(jax): + try: + devices = tuple(jax.local_devices(backend="gpu")) + capabilities = {_compute_capability(device) for device in devices} + except RuntimeError: + pytest.skip("A JAX GPU backend is not available") + except (TypeError, ValueError): + pytest.skip("JAX GPU compute capability is unavailable") + if not devices or len(capabilities) != 1: + pytest.skip("A homogeneous JAX SM100-family GPU configuration is required") + capability = capabilities.pop() + if capability not in {100, 103, 107}: + pytest.skip("A supported JAX SM100-family GPU is not available") + return devices[0] + + +def _scale_shape(batch, rows, k, sf_vec_size): + return ( + batch, + (rows + 127) // 128, + ((k + sf_vec_size - 1) // sf_vec_size + 3) // 4, + 32, + 4, + 4, + ) + + +@pytest.mark.L0 +def test_jax_gemm_amax_abstract_contract(monkeypatch): + jax, jnp = _jax_runtime() + + from cudnn._jax import JaxApiBase + from cudnn.jax import GemmAmaxSm100, gemm_amax_wrapper_sm100 + + def abstract_call(self, _inputs, *, output_descs, output_spec=None, **_options): + if output_spec is None: + output_spec = tuple(self._to_tensor_spec(desc) for desc in output_descs) + values = [] + for desc, spec in zip(output_descs, output_spec): + shape = self._to_shape_dtype_struct(desc, mode=spec.mode).shape + fill = 0 if desc.init_value is None else desc.init_value + values.append(jnp.full(shape, fill, dtype=desc.dtype)) + return tuple(values) + + monkeypatch.setattr(JaxApiBase, "_call_kernel", abstract_call) + monkeypatch.setattr( + JaxApiBase, + "_resolve_compute_capability", + staticmethod(lambda _target, _supported, _operation, **_options: 100), + ) + monkeypatch.setattr( + JaxApiBase, + "_get_max_active_clusters", + lambda self, _cluster_size, **_options: 8, + ) + + batch, m, n, k = 2, 128, 256, 128 + sample_a = jax.ShapeDtypeStruct((batch, m, k), jnp.float8_e4m3fn) + sample_b = jax.ShapeDtypeStruct((batch, n, k), jnp.float8_e4m3fn) + sample_sfa = jax.ShapeDtypeStruct( + _scale_shape(batch, m, k, 32), + jnp.float8_e8m0fnu, + ) + sample_sfb = jax.ShapeDtypeStruct( + _scale_shape(batch, n, k, 32), + jnp.float8_e8m0fnu, + ) + operation = GemmAmaxSm100(sample_a, sample_b, sample_sfa, sample_sfb) + result = jax.eval_shape( + operation, + sample_a, + sample_b, + sample_sfa, + sample_sfb, + ) + assert tuple(result.keys()) == ("c_tensor", "amax_tensor") + assert result["c_tensor"].shape == (batch, m, n) + assert result["c_tensor"].dtype == jnp.float32 + assert result["amax_tensor"].shape == (1, 1, 1) + assert result["amax_tensor"].dtype == jnp.float32 + + alternate_a = jax.ShapeDtypeStruct((batch, k, m), jnp.float8_e4m3fn) + alternate_b = jax.ShapeDtypeStruct((batch, k, n), jnp.float8_e4m3fn) + alternate = jax.eval_shape( + lambda a, b, sfa, sfb: gemm_amax_wrapper_sm100( + a, + b, + sfa, + sfb, + a_layout="LKM", + b_layout="LKN", + c_layout="LNM", + c_dtype=jnp.bfloat16, + ), + alternate_a, + alternate_b, + sample_sfa, + sample_sfb, + ) + assert alternate["c_tensor"].shape == (batch, n, m) + assert alternate["c_tensor"].dtype == jnp.bfloat16 + + +@pytest.mark.L0 +def test_jax_gemm_amax_sm100_jit_numerics_and_amax_reset(): + jax, jnp = _jax_runtime() + device = _supported_gpu(jax) + + from cudnn.jax import gemm_amax_wrapper_sm100 + + batch, m, n, k = 1, 128, 128, 128 + a_base = jax.random.uniform( + jax.random.key(0), + (batch, m, k), + minval=-0.5, + maxval=0.5, + dtype=jnp.float32, + ).astype(jnp.float8_e4m3fn) + b = jax.random.uniform( + jax.random.key(1), + (batch, n, k), + minval=-0.5, + maxval=0.5, + dtype=jnp.float32, + ).astype(jnp.float8_e4m3fn) + a_large = (a_base.astype(jnp.float32) * 2).astype(jnp.float8_e4m3fn) + sfa = jnp.ones( + _scale_shape(batch, m, k, 32), + dtype=jnp.float8_e8m0fnu, + ) + sfb = jnp.ones( + _scale_shape(batch, n, k, 32), + dtype=jnp.float8_e8m0fnu, + ) + a_base, a_large, b, sfa, sfb = (jax.device_put(value, device) for value in (a_base, a_large, b, sfa, sfb)) + + lowered = gemm_amax_wrapper_sm100.lower(a_large, b, sfa, sfb) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + compiled = lowered.compile() + + def reference(a_value): + c = jnp.einsum( + "lmk,lnk->lmn", + a_value.astype(jnp.float32), + b.astype(jnp.float32), + ) + return c, jnp.max(jnp.abs(c)).reshape(1, 1, 1) + + # Run the larger input first. The smaller second result proves that the + # aliased amax seed is reset to -inf for each invocation. + for a_value in (a_large, a_base): + result = compiled(a_value, b, sfa, sfb) + result["amax_tensor"].block_until_ready() + reference_c, reference_amax = reference(a_value) + assert jnp.allclose(result["c_tensor"], reference_c, atol=2e-1, rtol=5e-2) + assert jnp.allclose( + result["amax_tensor"], + reference_amax, + atol=2e-1, + rtol=5e-2, + ) diff --git a/test/python/fe_api/test_jax_gemm_amax_contract.py b/test/python/fe_api/test_jax_gemm_amax_contract.py new file mode 100644 index 000000000..9cdd1cf30 --- /dev/null +++ b/test/python/fe_api/test_jax_gemm_amax_contract.py @@ -0,0 +1,449 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for the JAX GEMM + amax adapter.""" + +from __future__ import annotations + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +import os +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "gemm_amax" +_PACKAGE = "cudnn_jax_gemm_amax_contract_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT8 = auto() + UINT8 = auto() + FP4_E2M1 = auto() + FP8_E4M3 = auto() + FP8_E5M2 = auto() + FP8_E8M0 = auto() + + +_DTYPE_TO_CUDNN = { + "float16": _DataType.HALF, + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, + "float4_e2m1fn": _DataType.FP4_E2M1, + "float8_e4m3fn": _DataType.FP8_E4M3, + "float8_e5m2": _DataType.FP8_E5M2, + "float8_e8m0fnu": _DataType.FP8_E8M0, +} +_CUDNN_TO_DTYPE = {value: key for key, value in _DTYPE_TO_CUDNN.items()} + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + +class _TensorSpec: + def __init__(self, *, layout, mode, divisibility=None): + self.layout = layout + self.mode = mode + self.divisibility = divisibility + + +class JaxGemmAmaxContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + operation_name = f"{_PACKAGE}.gemm_amax" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + layout_module = importlib.import_module(f"{internal_name}.layout") + result_module = importlib.import_module(f"{_PACKAGE}.common.result") + + class JaxTensorDesc(tensor_module.TensorDesc): + @classmethod + def from_shape(cls, shape, dtype, *, name, mode=None, init_value=None): + public_shape = tuple(shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + public_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride(public_shape, public_order) + canonical_axis_by_public_axis = layout_module.to_public_axes(tuple(range(len(public_shape))), mode) + desc = cls( + dtype=dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple(canonical_axis_by_public_axis[axis] for axis in public_order), + name=name, + init_value=init_value, + ) + object.__setattr__(desc, "mode", mode) + return desc + + @property + def cudnn_dtype(self): + return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) + + class JaxApiBase: + @staticmethod + def _to_tensor_desc( + value, + name, + *, + mode=None, + public_stride_order=None, + init_value=None, + ): + public_shape = tuple(value.shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + if public_stride_order is None: + public_stride_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride( + public_shape, + public_stride_order, + ) + canonical_axis_by_public_axis = layout_module.to_public_axes( + tuple(range(len(public_shape))), + mode, + ) + desc = JaxTensorDesc( + dtype=value.dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple(canonical_axis_by_public_axis[axis] for axis in public_stride_order), + name=name, + init_value=init_value, + ) + object.__setattr__(desc, "mode", mode) + return desc + + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del operation_name + resolved = 100 if target is None else target + if resolved not in supported: + raise ValueError(f"unsupported target {resolved}") + return resolved + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + actual_shape = layout_module.to_canonical_axes(tuple(value.shape), mode) + if actual_shape != expected.shape: + raise ValueError(f"{expected.name} tensor shape mismatch: expected " f"{expected.shape}, got {actual_shape}") + actual_dtype = _DTYPE_TO_CUDNN.get(value.dtype, _DataType.NOT_SET) + if actual_dtype != expected.cudnn_dtype: + raise ValueError(f"{expected.name} tensor dtype mismatch: expected " f"{expected.cudnn_dtype}, got {actual_dtype}") + + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + mode = layout_module.normalize_mode(desc.ndim, mode) + return _TensorSpec( + layout=layout_module.to_cutlass_layout( + desc.shape, + desc.stride, + desc.stride_order, + mode=mode, + name=desc.name, + ), + mode=mode, + divisibility=divisibility, + ) + + def _call_kernel(self, inputs, **options): + input_descs = options.get("input_descs") + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc, mode=desc.mode) + options["input_spec"] = tuple(self._to_tensor_spec(desc, mode=desc.mode) for desc in input_descs) + if "output_spec" not in options: + options["output_spec"] = tuple( + self._to_tensor_spec(desc, mode=getattr(desc, "mode", None)) for desc in options["output_descs"] + ) + self.captured_call = (tuple(inputs), options) + return tuple( + _Array( + layout_module.to_public_axes(desc.shape, spec.mode), + _CUDNN_TO_DTYPE[desc.cudnn_dtype], + ) + for desc, spec in zip( + options["output_descs"], + options["output_spec"], + ) + ) + + def _get_max_active_clusters(self, cluster_size, *, overlap_margin=0): + self.captured_active_cluster_queries = getattr( + self, + "captured_active_cluster_queries", + [], + ) + self.captured_active_cluster_queries.append((cluster_size, overlap_margin)) + return 12 - overlap_margin + + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = result_module.TupleDict + + datatypes = types.ModuleType(f"{internal_name}.datatypes") + datatypes.jax_to_cudnn_dtype = lambda dtype: _DTYPE_TO_CUDNN.get(dtype, _DataType.NOT_SET) + datatypes.normalize_jax_dtype = lambda value, default, name: (default if value is None else value) + datatypes.cudnn_to_jax_dtype = lambda dtype: _CUDNN_TO_DTYPE[dtype] + sys.modules[datatypes.__name__] = datatypes + + cls.jit_static_argnames = {} + fake_jax = types.ModuleType("jax") + + def jit(function=None, *, static_argnames=()): + def decorate(target): + cls.jit_static_argnames[target.__name__] = tuple(static_argnames) + return target + + return decorate if function is None else decorate(function) + + fake_jax.jit = jit + fake_jax.ShapeDtypeStruct = _Array + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.dtype = lambda dtype: dtype + for dtype_name in _DTYPE_TO_CUDNN: + setattr(fake_jnp, dtype_name, dtype_name) + fake_jax.numpy = fake_jnp + + try: + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "torch": None, + "cutlass": None, + }, + ): + cls.module = importlib.import_module(f"{operation_name}.jax") + except Exception: + cls.tearDownClass() + raise + + cls.operation_name = operation_name + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + @staticmethod + def _samples( + *, + a_layout="LMK", + b_layout="LNK", + ab_dtype="float8_e4m3fn", + sf_dtype="float8_e8m0fnu", + sf_vec_size=32, + ): + m, n, k, batch = 128, 256, 128, 2 + a_shape = (batch, m, k) if a_layout == "LMK" else (batch, k, m) + b_shape = (batch, n, k) if b_layout == "LNK" else (batch, k, n) + k_tiles = ((k + sf_vec_size - 1) // sf_vec_size + 3) // 4 + sfa_shape = (batch, 1, k_tiles, 32, 4, 4) + sfb_shape = (batch, 2, k_tiles, 32, 4, 4) + return ( + _Array(a_shape, ab_dtype), + _Array(b_shape, ab_dtype), + _Array(sfa_shape, sf_dtype), + _Array(sfb_shape, sf_dtype), + ) + + def test_default_layouts_bind_row_major_scales_and_infer_outputs(self): + samples = self._samples() + api = self.module.GemmAmaxSm100(*samples) + + self.assertEqual(api.a_desc.mode, (1, 2, 0)) + self.assertEqual(api.b_desc.mode, (1, 2, 0)) + self.assertEqual(api.c_desc.mode, (1, 2, 0)) + self.assertEqual(api.sfa_desc.mode, (3, 4, 1, 5, 2, 0)) + self.assertEqual(api.sfa_desc.shape, (32, 4, 1, 4, 1, 2)) + self.assertEqual(api.sfa_desc.stride_order, (3, 1, 0, 4, 2, 5)) + self.assertEqual(api.amax_desc.init_value, float("-inf")) + + result = api(*samples) + self.assertEqual(tuple(result.keys()), ("c_tensor", "amax_tensor")) + self.assertEqual(result["c_tensor"].shape, (2, 128, 256)) + self.assertEqual(result["c_tensor"].dtype, "float32") + self.assertEqual(result["amax_tensor"].shape, (1, 1, 1)) + + inputs, options = api.captured_call + self.assertEqual(inputs, samples) + self.assertTrue(callable(options["launch"])) + self.assertEqual( + tuple(spec.mode for spec in options["input_spec"]), + (api.a_desc.mode, api.b_desc.mode, api.sfa_desc.mode, api.sfb_desc.mode), + ) + self.assertEqual(options["input_spec"][2].layout, (5, 4, 3, 2, 1, 0)) + self.assertEqual(options["input_spec"][3].layout, (5, 4, 3, 2, 1, 0)) + self.assertEqual(options["output_descs"], (api.c_desc, api.amax_desc)) + self.assertIn("--gpu-arch sm_100a", options["compile_options"]) + + def test_alternate_gemm_layouts_and_native_fp4_are_supported(self): + alternate_samples = self._samples( + a_layout="LKM", + b_layout="LKN", + ) + alternate = self.module.GemmAmaxSm100( + *alternate_samples, + a_layout="LKM", + b_layout="LKN", + c_layout="LNM", + c_dtype="bfloat16", + ) + alternate_result = alternate(*alternate_samples) + self.assertEqual(alternate.a_desc.stride_order, (0, 1, 2)) + self.assertEqual(alternate.b_desc.stride_order, (0, 1, 2)) + self.assertEqual(alternate_result["c_tensor"].shape, (2, 256, 128)) + + fp4_samples = self._samples( + ab_dtype="float4_e2m1fn", + sf_dtype="float8_e4m3fn", + sf_vec_size=16, + ) + fp4 = self.module.GemmAmaxSm100( + *fp4_samples, + c_dtype="bfloat16", + sf_vec_size=16, + ) + fp4_result = fp4(*fp4_samples) + self.assertEqual(fp4_result["c_tensor"].shape, (2, 128, 256)) + self.assertEqual(fp4_result["c_tensor"].dtype, "bfloat16") + + for c_dtype in ("float8_e4m3fn", "float8_e5m2", "float4_e2m1fn"): + with self.subTest(c_dtype=c_dtype): + quantized_output = self.module.GemmAmaxSm100( + *fp4_samples, + c_dtype=c_dtype, + sf_vec_size=16, + )(*fp4_samples) + self.assertEqual(quantized_output["c_tensor"].dtype, c_dtype) + + def test_explicit_outputs_and_runtime_signatures_are_checked(self): + samples = self._samples() + api = self.module.GemmAmaxSm100( + *samples, + sample_c=_Array((2, 128, 256), "bfloat16"), + sample_amax=_Array((1, 1, 1), "float32"), + ) + result = api(*samples) + self.assertEqual(result["c_tensor"].dtype, "bfloat16") + + with self.assertRaisesRegex( + ValueError, + "sample_c and sample_amax must be provided together", + ): + self.module.GemmAmaxSm100( + *samples, + sample_c=_Array((2, 128, 256), "float32"), + ) + with self.assertRaisesRegex(ValueError, "c_dtype cannot be specified"): + self.module.GemmAmaxSm100( + *samples, + sample_c=_Array((2, 128, 256), "bfloat16"), + sample_amax=_Array((1, 1, 1), "float32"), + c_dtype="bfloat16", + ) + with self.assertRaisesRegex(ValueError, "sample_a tensor shape mismatch"): + api(_Array((2, 64, 128), samples[0].dtype), *samples[1:]) + + def test_wrapper_configuration_is_static(self): + self.assertEqual( + self.jit_static_argnames["gemm_amax_wrapper_sm100"], + ( + "c_layout", + "c_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "a_layout", + "b_layout", + ), + ) + samples = self._samples() + result = self.module.gemm_amax_wrapper_sm100(*samples) + self.assertEqual(result["c_tensor"].shape, (2, 128, 256)) + + def test_inline_launch_constructs_kernel_and_preserves_abi_order(self): + samples = self._samples() + with mock.patch.dict(os.environ, {"CUDNNFE_CLUSTER_OVERLAP_MARGIN": "2"}): + api = self.module.GemmAmaxSm100(*samples, cluster_shape_mn=(2, 2)) + api(*samples) + self.assertEqual(api.captured_active_cluster_queries, [(4, 2)]) + _, options = api.captured_call + seen = {} + + class Kernel: + def __init__(self, **configuration): + seen["configuration"] = configuration + + def __call__(self, *arguments): + seen["arguments"] = arguments + + kernel_module = types.ModuleType(f"{self.operation_name}.dense_blockscaled_gemm_persistent_amax") + kernel_module.Sm100BlockScaledPersistentDenseGemmKernel = Kernel + stream = object() + with mock.patch.dict(sys.modules, {kernel_module.__name__: kernel_module}): + options["launch"](stream, "A", "B", "SFA", "SFB", "C", "Amax") + + self.assertEqual( + seen["configuration"], + { + "sf_vec_size": 32, + "mma_tiler_mn": (128, 128), + "cluster_shape_mn": (2, 2), + }, + ) + self.assertEqual( + seen["arguments"], + ("A", "B", "SFA", "SFB", "C", "Amax", 10, stream), + ) + + def test_adapter_import_does_not_load_torch_or_kernel(self): + self.assertNotIn(f"{self.operation_name}.api", sys.modules) + self.assertNotIn( + f"{self.operation_name}.dense_blockscaled_gemm_persistent_amax", + sys.modules, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_gemm_layout.py b/test/python/fe_api/test_jax_gemm_layout.py new file mode 100644 index 000000000..fab222835 --- /dev/null +++ b/test/python/fe_api/test_jax_gemm_layout.py @@ -0,0 +1,129 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for dense JAX GEMM axis bindings.""" + +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_PACKAGE = "cudnn_jax_gemm_layout_test" + + +class JaxGemmLayoutTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + sys.modules[_PACKAGE] = root + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + try: + with mock.patch.dict( + sys.modules, + { + "jax": None, + "jax.numpy": None, + "cutlass": None, + "cutlass.jax": None, + "torch": None, + }, + ): + cls.module = importlib.import_module(f"{internal_name}.gemm") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls): + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def test_maps_supported_public_layouts_to_canonical_modes(self): + cases = ( + (self.module.gemm_a_mode, "LMK", (1, 2, 0)), + (self.module.gemm_a_mode, "LKM", (2, 1, 0)), + (self.module.gemm_b_mode, "LNK", (1, 2, 0)), + (self.module.gemm_b_mode, "LKN", (2, 1, 0)), + (self.module.gemm_output_mode, "LMN", (1, 2, 0)), + (self.module.gemm_output_mode, "LNM", (2, 1, 0)), + ) + + for helper, layout, expected in cases: + with self.subTest(layout=layout): + self.assertEqual(helper(layout), expected) + + def test_rejects_unsupported_axis_orders(self): + with self.assertRaisesRegex(ValueError, r"a_layout must be one of \('LMK', 'LKM'\)"): + self.module.gemm_a_mode("MKL") + with self.assertRaisesRegex(ValueError, r"b_layout must be one of \('LNK', 'LKN'\)"): + self.module.gemm_b_mode("NKL") + with self.assertRaisesRegex(ValueError, r"c_layout must be one of \('LMN', 'LNM'\)"): + self.module.gemm_output_mode("MNL") + + with self.assertRaisesRegex(ValueError, r"d_layout must be one of \('LMN', 'LNM'\)"): + self.module.gemm_output_mode("MNL", name="d_layout") + + def test_layout_validation_is_case_sensitive(self): + with self.assertRaisesRegex(ValueError, "got 'lmk'"): + self.module.gemm_a_mode("lmk") + + def test_layout_validation_requires_a_string(self): + with self.assertRaisesRegex(TypeError, "a_layout must be a string"): + self.module.gemm_a_mode(("L", "M", "K")) + + def test_fixed_auxiliary_stride_orders_match_kernel_abis(self): + self.assertEqual(self.module.ROW_MAJOR_STRIDE_ORDER_3D, (2, 1, 0)) + self.assertEqual(self.module.BLOCK_SCALE_MODE, (3, 4, 1, 5, 2, 0)) + self.assertEqual(self.module.BLOCK_SCALE_STRIDE_ORDER, (3, 1, 0, 4, 2, 5)) + self.assertEqual(self.module.PROBABILITY_MODE, (2, 1, 0)) + self.assertEqual(self.module.PROBABILITY_STRIDE_ORDER, (0, 1, 2)) + + layout = importlib.import_module(f"{self.module.__package__}.layout") + block_scale_shape = (32, 4, 2, 4, 3, 5) + block_scale_stride = layout.compact_stride(block_scale_shape, self.module.BLOCK_SCALE_STRIDE_ORDER) + self.assertEqual( + layout.to_cutlass_layout( + block_scale_shape, + block_scale_stride, + self.module.BLOCK_SCALE_STRIDE_ORDER, + ), + (2, 1, 4, 0, 3, 5), + ) + + probability_shape = (7, 1, 5) + probability_stride = layout.compact_stride(probability_shape, self.module.PROBABILITY_STRIDE_ORDER) + self.assertEqual( + layout.to_cutlass_layout( + probability_shape, + probability_stride, + self.module.PROBABILITY_STRIDE_ORDER, + ), + (0, 1, 2), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_gemm_relu.py b/test/python/fe_api/test_jax_gemm_relu.py new file mode 100644 index 000000000..94cdf8469 --- /dev/null +++ b/test/python/fe_api/test_jax_gemm_relu.py @@ -0,0 +1,276 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX integration tests for dense block-scaled sReLU fusions.""" + +import pytest + + +def _jax_runtime(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + cutlass_jax = pytest.importorskip("cutlass.jax") + if not cutlass_jax.is_available(): + pytest.skip("Installed JAX version is unsupported by CUTLASS JAX") + return jax, jnp + + +def _compute_capability(device): + reported = getattr(device, "compute_capability", None) + if isinstance(reported, (tuple, list)): + major, minor = int(reported[0]), int(reported[1]) + else: + text = str(reported) + if "." in text: + major_text, minor_text = text.split(".", 1) + major, minor = int(major_text), int(minor_text) + else: + major, minor = divmod(int(text), 10) + return major * 10 + minor + + +def _supported_gpu(jax): + devices = tuple(device for device in jax.local_devices() if device.platform == "gpu") + try: + capabilities = {_compute_capability(device) for device in devices} + except (TypeError, ValueError): + pytest.skip("JAX GPU compute capability is unavailable") + if not devices or len(capabilities) != 1: + pytest.skip("A homogeneous JAX SM100-family GPU configuration is required") + capability = capabilities.pop() + if capability not in {100, 103, 107}: + pytest.skip("A supported JAX SM100-family GPU is not available") + return devices[0] + + +@pytest.mark.L0 +def test_jax_gemm_relu_abstract_contract(monkeypatch): + jax, jnp = _jax_runtime() + + from cudnn._jax import JaxApiBase + from cudnn.jax import GemmDsreluSm100, GemmSreluSm100, gemm_dsrelu_wrapper_sm100, gemm_srelu_wrapper_sm100 + + def abstract_call(self, _inputs, *, output_descs, output_spec=None, **_options): + if output_spec is None: + output_spec = tuple(self._to_tensor_spec(desc) for desc in output_descs) + return tuple( + jnp.empty( + self._to_shape_dtype_struct(desc, mode=spec.mode).shape, + dtype=desc.dtype, + ) + for desc, spec in zip(output_descs, output_spec) + ) + + monkeypatch.setattr(JaxApiBase, "_call_kernel", abstract_call) + monkeypatch.setattr( + JaxApiBase, + "_resolve_compute_capability", + staticmethod(lambda _target, _supported, _operation: 100), + ) + monkeypatch.setattr(JaxApiBase, "_get_max_active_clusters", lambda self, _size, **_options: 8) + + batch, m, n, k = 2, 256, 256, 512 + a = jax.ShapeDtypeStruct((batch, m, k), jnp.float8_e4m3fn) + b = jax.ShapeDtypeStruct((batch, n, k), jnp.float8_e4m3fn) + scales = jax.ShapeDtypeStruct((batch, m // 128, k // 128, 32, 4, 4), jnp.float8_e8m0fnu) + prob = jax.ShapeDtypeStruct((batch, 1, m), jnp.float32) + c = jax.ShapeDtypeStruct((batch, m, n), jnp.bfloat16) + + forward = jax.eval_shape( + GemmSreluSm100(a, b, scales, scales, prob, sf_vec_size=32), + a, + b, + scales, + scales, + prob, + ) + backward = jax.eval_shape( + GemmDsreluSm100(a, b, c, scales, scales, prob, sf_vec_size=32), + a, + b, + c, + scales, + scales, + prob, + ) + assert tuple(forward.keys()) == ("c_tensor", "d_tensor", "amax_tensor", "sfd_tensor") + assert forward["c_tensor"].shape == (batch, m, n) + assert forward["d_tensor"].shape == (batch, m, n) + assert forward["amax_tensor"] is None + assert backward["d_tensor"].shape == (batch, m, n) + assert backward["dprob_tensor"].shape == (batch, 1, m) + + alternate = jax.eval_shape( + lambda a_value, b_value, scale_a, scale_b, probability: gemm_srelu_wrapper_sm100( + a_value, + b_value, + scale_a, + scale_b, + probability, + c_layout="LNM", + sf_vec_size=32, + a_layout="LKM", + b_layout="LKN", + ), + jax.ShapeDtypeStruct((batch, k, m), jnp.float8_e4m3fn), + jax.ShapeDtypeStruct((batch, k, n), jnp.float8_e4m3fn), + scales, + scales, + prob, + ) + assert alternate["c_tensor"].shape == (batch, n, m) + + with pytest.raises(NotImplementedError, match="FP8 D output is unavailable"): + jax.eval_shape( + lambda *args: gemm_dsrelu_wrapper_sm100( + *args, + d_dtype=jnp.float8_e4m3fn, + sf_vec_size=32, + ), + a, + b, + c, + scales, + scales, + prob, + ) + + fp4_dtype = getattr(jnp, "float4_e2m1fn", None) + if fp4_dtype is not None: + fp4_a = jax.ShapeDtypeStruct((1, 128, 128), fp4_dtype) + fp4_b = jax.ShapeDtypeStruct((1, 128, 128), fp4_dtype) + fp4_scales = jax.ShapeDtypeStruct((1, 1, 2, 32, 4, 4), jnp.float8_e8m0fnu) + fp4_prob = jax.ShapeDtypeStruct((1, 1, 128), jnp.float32) + fp4_result = jax.eval_shape( + lambda *args: gemm_srelu_wrapper_sm100( + *args, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ), + fp4_a, + fp4_b, + fp4_scales, + fp4_scales, + fp4_prob, + ) + assert fp4_result["amax_tensor"].shape == (1,) + + +@pytest.mark.L0 +def test_jax_gemm_relu_sm100_jit_and_numerics(): + jax, jnp = _jax_runtime() + device = _supported_gpu(jax) + + from cudnn.jax import gemm_dsrelu_wrapper_sm100, gemm_srelu_wrapper_sm100 + + batch = 1 + m = n = k = 128 + alpha = 0.5 + fp8_dtype = jnp.float8_e4m3fn + a = jax.device_put( + jax.random.uniform(jax.random.key(0), (batch, m, k), minval=-0.5, maxval=0.5).astype(fp8_dtype), + device, + ) + b = jax.device_put( + jax.random.uniform(jax.random.key(1), (batch, n, k), minval=-0.5, maxval=0.5).astype(fp8_dtype), + device, + ) + c = jax.device_put( + jax.random.uniform(jax.random.key(2), (batch, m, n), minval=-0.5, maxval=0.5).astype(jnp.bfloat16), + device, + ) + prob = jax.device_put( + jax.random.uniform(jax.random.key(3), (batch, 1, m), minval=0.25, maxval=1.0), + device, + ) + scales = jax.device_put( + jnp.ones((batch, 1, 1, 32, 4, 4), dtype=jnp.float32).astype(jnp.float8_e8m0fnu), + device, + ) + + forward_lowered = gemm_srelu_wrapper_sm100.lower( + a, + b, + scales, + scales, + prob, + alpha=alpha, + sf_vec_size=32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + backward_lowered = gemm_dsrelu_wrapper_sm100.lower( + a, + b, + c, + scales, + scales, + prob, + alpha=alpha, + sf_vec_size=32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + for lowered in (forward_lowered, backward_lowered): + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + gemm = alpha * jnp.einsum("lmk,lnk->lmn", a.astype(jnp.float32), b.astype(jnp.float32)) + relu_gemm = jnp.maximum(gemm, 0.0) + probability = prob.transpose(0, 2, 1) + + forward_result = forward_lowered.compile()(a, b, scales, scales, prob) + forward_result["d_tensor"].block_until_ready() + assert jnp.allclose(forward_result["c_tensor"].astype(jnp.float32), gemm, atol=2e-1, rtol=3e-2) + assert jnp.allclose( + forward_result["d_tensor"].astype(jnp.float32), + jnp.square(relu_gemm) * probability, + atol=4e-1, + rtol=5e-2, + ) + + backward_compiled = backward_lowered.compile() + backward_result = backward_compiled(a, b, c, scales, scales, prob) + backward_result["dprob_tensor"].block_until_ready() + expected_d = c.astype(jnp.float32) * probability * 2 * relu_gemm + expected_dprob = jnp.sum(c.astype(jnp.float32) * jnp.square(relu_gemm), axis=2, keepdims=True).transpose(0, 2, 1) + assert jnp.allclose(backward_result["d_tensor"].astype(jnp.float32), expected_d, atol=4e-1, rtol=5e-2) + assert jnp.allclose(backward_result["dprob_tensor"], expected_dprob, atol=2.0, rtol=5e-2) + + zero_result = backward_compiled(a, b, jnp.zeros_like(c), scales, scales, prob) + zero_result["dprob_tensor"].block_until_ready() + assert jnp.array_equal(zero_result["dprob_tensor"], jnp.zeros_like(zero_result["dprob_tensor"])) + + +@pytest.mark.L0 +def test_jax_gemm_srelu_native_fp4_amax_smoke(): + jax, jnp = _jax_runtime() + device = _supported_gpu(jax) + fp4_dtype = getattr(jnp, "float4_e2m1fn", None) + if fp4_dtype is None: + pytest.skip("JAX does not provide native float4_e2m1fn") + + from cudnn.jax import gemm_srelu_wrapper_sm100 + + m = n = k = 128 + a = jax.device_put(jnp.ones((1, m, k), dtype=jnp.float32).astype(fp4_dtype), device) + b = jax.device_put(jnp.ones((1, n, k), dtype=jnp.float32).astype(fp4_dtype), device) + scales = jax.device_put( + jnp.ones((1, 1, 2, 32, 4, 4), dtype=jnp.float32).astype(jnp.float8_e8m0fnu), + device, + ) + prob = jax.device_put(jnp.ones((1, 1, m), dtype=jnp.float32), device) + lowered = gemm_srelu_wrapper_sm100.lower( + a, + b, + scales, + scales, + prob, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + result = lowered.compile()(a, b, scales, scales, prob) + result["amax_tensor"].block_until_ready() + assert result["amax_tensor"].shape == (1,) + assert jnp.all(jnp.isfinite(result["amax_tensor"])) diff --git a/test/python/fe_api/test_jax_gemm_relu_contract.py b/test/python/fe_api/test_jax_gemm_relu_contract.py new file mode 100644 index 000000000..fc690e3fb --- /dev/null +++ b/test/python/fe_api/test_jax_gemm_relu_contract.py @@ -0,0 +1,684 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for the JAX sReLU and dsReLU adapters.""" + +from __future__ import annotations + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +import os +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_PACKAGE = "cudnn_jax_gemm_relu_contract_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + UINT8 = auto() + FP4_E2M1 = auto() + FP8_E4M3 = auto() + FP8_E5M2 = auto() + FP8_E8M0 = auto() + + +_DTYPE_TO_CUDNN = { + "float4_e2m1fn": _DataType.FP4_E2M1, + "float8_e4m3fn": _DataType.FP8_E4M3, + "float8_e5m2": _DataType.FP8_E5M2, + "float8_e8m0fnu": _DataType.FP8_E8M0, + "float16": _DataType.HALF, + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, +} +_CUDNN_TO_DTYPE = {value: key for key, value in _DTYPE_TO_CUDNN.items()} + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + +class _TensorSpec: + def __init__(self, *, layout, mode, divisibility=None): + self.layout = layout + self.mode = mode + self.divisibility = divisibility + + +class JaxGemmReluContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + for operation in ("gemm_srelu", "gemm_dsrelu"): + operation_name = f"{_PACKAGE}.{operation}" + operation_module = types.ModuleType(operation_name) + operation_module.__path__ = [str(_CUDNN_ROOT / operation)] + operation_module.__package__ = operation_name + operation_module.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation_module + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + layout_module = importlib.import_module(f"{internal_name}.layout") + result_module = importlib.import_module(f"{_PACKAGE}.common.result") + + class JaxTensorDesc(tensor_module.TensorDesc): + @classmethod + def from_shape(cls, shape, dtype, *, name, mode=None, init_value=None): + public_shape = tuple(shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + public_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride(public_shape, public_order) + canonical_axis_by_public_axis = layout_module.to_public_axes(tuple(range(len(public_shape))), mode) + desc = cls( + dtype=dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple(canonical_axis_by_public_axis[axis] for axis in public_order), + name=name, + init_value=init_value, + ) + object.__setattr__(desc, "mode", mode) + return desc + + @property + def cudnn_dtype(self): + return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) + + def compact_like(self, *, cudnn_dtype, shape, stride_order=None, name="", init_value=None): + canonical = tensor_module.make_compact_tensor_desc( + dtype=cudnn_dtype, + shape=shape, + stride_order=stride_order, + name=name, + init_value=init_value, + ) + return JaxTensorDesc( + dtype=_CUDNN_TO_DTYPE[cudnn_dtype], + shape=canonical.shape, + stride=canonical.stride, + stride_order=canonical.stride_order, + name=name, + init_value=init_value, + ) + + class JaxApiBase: + @staticmethod + def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): + public_shape = tuple(value.shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + public_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride(public_shape, public_order) + canonical_axis_by_public_axis = layout_module.to_public_axes(tuple(range(len(public_shape))), mode) + desc = JaxTensorDesc( + dtype=value.dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple(canonical_axis_by_public_axis[axis] for axis in public_order), + name=name, + init_value=init_value, + ) + object.__setattr__(desc, "mode", mode) + return desc + + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del operation_name + resolved = 100 if target is None else target + if resolved not in supported: + raise ValueError(f"unsupported target {resolved}") + return resolved + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + actual_shape = layout_module.to_canonical_axes(tuple(value.shape), mode) + if actual_shape != expected.shape: + raise ValueError(f"{expected.name} tensor shape mismatch: expected {expected.shape}, got {actual_shape}") + actual_dtype = _DTYPE_TO_CUDNN.get(value.dtype, _DataType.NOT_SET) + if actual_dtype != expected.cudnn_dtype: + raise ValueError(f"{expected.name} tensor dtype mismatch: expected {expected.cudnn_dtype}, got {actual_dtype}") + + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + return _TensorSpec( + layout=layout_module.to_cutlass_layout( + desc.shape, + desc.stride, + desc.stride_order, + mode=mode, + name=desc.name, + ), + mode=mode, + divisibility=divisibility, + ) + + def _call_kernel(self, inputs, **options): + input_descs = options.get("input_descs") + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc, mode=desc.mode) + options["input_spec"] = tuple(self._to_tensor_spec(desc, mode=desc.mode) for desc in input_descs) + if "output_spec" not in options: + options["output_spec"] = tuple( + self._to_tensor_spec(desc, mode=getattr(desc, "mode", None)) for desc in options["output_descs"] + ) + self.captured_call = (tuple(inputs), options) + return tuple( + _Array( + layout_module.to_public_axes(desc.shape, spec.mode), + _CUDNN_TO_DTYPE[desc.cudnn_dtype], + ) + for desc, spec in zip(options["output_descs"], options["output_spec"]) + ) + + def _get_max_active_clusters(self, cluster_size, *, overlap_margin=0): + self.captured_active_cluster_query = (cluster_size, overlap_margin) + return 12 - overlap_margin + + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = result_module.TupleDict + + datatypes = types.ModuleType(f"{internal_name}.datatypes") + datatypes.jax_to_cudnn_dtype = lambda dtype: _DTYPE_TO_CUDNN.get(dtype, _DataType.NOT_SET) + datatypes.normalize_jax_dtype = lambda value, default, name: default if value is None else value + datatypes.cudnn_to_jax_dtype = lambda dtype: _CUDNN_TO_DTYPE[dtype] + sys.modules[datatypes.__name__] = datatypes + + cls.jit_static_argnames = {} + fake_jax = types.ModuleType("jax") + + def jit(function=None, *, static_argnames=()): + def decorate(target): + cls.jit_static_argnames[target.__name__] = tuple(static_argnames) + return target + + return decorate if function is None else decorate(function) + + fake_jax.jit = jit + fake_jax.ShapeDtypeStruct = _Array + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.dtype = lambda dtype: dtype + for dtype_name in _DTYPE_TO_CUDNN: + setattr(fake_jnp, dtype_name, dtype_name) + fake_jax.numpy = fake_jnp + + try: + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "torch": None, + "cutlass": None, + }, + ): + cls.srelu = importlib.import_module(f"{_PACKAGE}.gemm_srelu.jax") + cls.dsrelu = importlib.import_module(f"{_PACKAGE}.gemm_dsrelu.jax") + except Exception: + cls.tearDownClass() + raise + + cls.operation_names = { + cls.srelu: f"{_PACKAGE}.gemm_srelu", + cls.dsrelu: f"{_PACKAGE}.gemm_dsrelu", + } + cls.tensor_module = tensor_module + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + @staticmethod + def _samples(*, dtype="float8_e4m3fn", sf_vec_size=32, batch=2, m=256, n=256, k=512): + rest_k = ((k + sf_vec_size - 1) // sf_vec_size + 3) // 4 + a = _Array((batch, m, k), dtype) + b = _Array((batch, n, k), dtype) + sfa = _Array((batch, (m + 127) // 128, rest_k, 32, 4, 4), "float8_e8m0fnu") + sfb = _Array((batch, (n + 127) // 128, rest_k, 32, 4, 4), "float8_e8m0fnu") + prob = _Array((batch, 1, m), "float32") + c = _Array((batch, m, n), "bfloat16") + return a, b, c, sfa, sfb, prob + + def test_forward_and_backward_use_row_major_public_auxiliary_layouts(self): + a, b, c, sfa, sfb, prob = self._samples() + forward = self.srelu.GemmSreluSm100(a, b, sfa, sfb, prob, sf_vec_size=32) + backward = self.dsrelu.GemmDsreluSm100(a, b, c, sfa, sfb, prob, sf_vec_size=32) + + forward_result = forward(a, b, sfa, sfb, prob) + backward_result = backward(a, b, c, sfa, sfb, prob) + + self.assertEqual(forward.sfa_desc.shape, (32, 4, 2, 4, 4, 2)) + self.assertEqual(forward.sfa_desc.stride_order, (3, 1, 0, 4, 2, 5)) + self.assertEqual(forward.prob_desc.shape, (256, 1, 2)) + self.assertEqual(forward.prob_desc.stride_order, (0, 1, 2)) + self.assertIs(forward._op.sfa, forward.sfa_desc) + self.assertIs(backward._op.dprob, backward.dprob_desc) + + self.assertEqual(forward_result["c_tensor"].shape, (2, 256, 256)) + self.assertEqual(forward_result["d_tensor"].shape, (2, 256, 256)) + self.assertIsNone(forward_result["amax_tensor"]) + self.assertEqual(backward_result["d_tensor"].shape, (2, 256, 256)) + self.assertEqual(backward_result["dprob_tensor"].shape, (2, 1, 256)) + self.assertEqual(backward.dprob_desc.init_value, 0.0) + + forward_inputs, forward_options = forward.captured_call + self.assertEqual(forward_inputs, (a, b, sfa, sfb, prob)) + self.assertTrue(callable(forward_options["launch"])) + self.assertEqual( + tuple(spec.mode for spec in forward_options["input_spec"]), + ( + forward.a_desc.mode, + forward.b_desc.mode, + forward.sfa_desc.mode, + forward.sfb_desc.mode, + forward.prob_desc.mode, + ), + ) + self.assertEqual(forward_options["input_spec"][2].layout, (5, 4, 3, 2, 1, 0)) + self.assertEqual(forward_options["input_spec"][4].layout, (2, 1, 0)) + self.assertIn("--gpu-arch sm_100a", forward_options["compile_options"]) + + def test_native_fp4_returns_initialized_amax(self): + a, b, c, sfa, sfb, prob = self._samples(dtype="float4_e2m1fn", sf_vec_size=16) + forward = self.srelu.GemmSreluSm100(a, b, sfa, sfb, prob, sf_vec_size=16) + backward = self.dsrelu.GemmDsreluSm100(a, b, c, sfa, sfb, prob, sf_vec_size=16) + + forward_result = forward(a, b, sfa, sfb, prob) + backward_result = backward(a, b, c, sfa, sfb, prob) + + for api, result in ((forward, forward_result), (backward, backward_result)): + self.assertIsNotNone(api.amax_desc) + self.assertEqual(api.amax_desc.shape, (1,)) + self.assertEqual(api.amax_desc.init_value, float("-inf")) + self.assertEqual(result["amax_tensor"].shape, (1,)) + self.assertIs(api.captured_call[1]["output_descs"][-1], api.amax_desc) + + def test_fp8_d_output_is_rejected_until_sfd_generation_is_implemented(self): + a, b, c, sfa, sfb, prob = self._samples() + constructors = ( + lambda: self.srelu.GemmSreluSm100( + a, + b, + sfa, + sfb, + prob, + d_dtype="float8_e4m3fn", + sf_vec_size=32, + ), + lambda: self.dsrelu.GemmDsreluSm100( + a, + b, + c, + sfa, + sfb, + prob, + d_dtype="float8_e4m3fn", + sf_vec_size=32, + ), + ) + for construct in constructors: + with self.subTest(construct=construct), self.assertRaisesRegex(NotImplementedError, "FP8 D output is unavailable"): + construct().check_support() + + def test_explicit_output_samples_are_alternative_to_dtype_arguments(self): + a, b, c, sfa, sfb, prob = self._samples() + sample_d = _Array(c.shape, "bfloat16") + sample_dprob = _Array(prob.shape, "float32") + + forward = self.srelu.GemmSreluSm100( + a, + b, + sfa, + sfb, + prob, + sample_c=c, + sample_d=sample_d, + sf_vec_size=32, + ) + backward = self.dsrelu.GemmDsreluSm100( + a, + b, + c, + sfa, + sfb, + prob, + sample_d=sample_d, + sample_dprob=sample_dprob, + sf_vec_size=32, + ) + self.assertTrue(forward.check_support()) + self.assertTrue(backward.check_support()) + + with self.assertRaisesRegex(ValueError, "c_dtype and d_dtype cannot be specified"): + self.srelu.GemmSreluSm100( + a, + b, + sfa, + sfb, + prob, + sample_c=c, + sample_d=sample_d, + c_dtype="bfloat16", + sf_vec_size=32, + ) + with self.assertRaisesRegex(ValueError, "d_dtype cannot be specified"): + self.dsrelu.GemmDsreluSm100( + a, + b, + c, + sfa, + sfb, + prob, + sample_d=sample_d, + sample_dprob=sample_dprob, + d_dtype="bfloat16", + sf_vec_size=32, + ) + + def test_unit_scale_tiles_and_batch_preserve_packed_layout(self): + a, b, c, sfa, sfb, prob = self._samples(batch=1, m=128, n=128, k=128) + forward = self.srelu.GemmSreluSm100( + a, + b, + sfa, + sfb, + prob, + sf_vec_size=32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + backward = self.dsrelu.GemmDsreluSm100( + a, + b, + c, + sfa, + sfb, + prob, + sf_vec_size=32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + + self.assertTrue(forward.check_support()) + self.assertTrue(backward.check_support()) + self.assertEqual(forward.sfa_desc.shape, (32, 4, 1, 4, 1, 1)) + self.assertTrue(forward.sfa_desc.is_compact((3, 1, 0, 4, 2, 5))) + + def test_launch_callbacks_preserve_native_kernel_abis(self): + a, b, c, sfa, sfb, prob = self._samples(dtype="float4_e2m1fn", sf_vec_size=16) + with mock.patch.dict(os.environ, {"CUDNNFE_CLUSTER_OVERLAP_MARGIN": "2"}): + forward = self.srelu.GemmSreluSm100(a, b, sfa, sfb, prob, alpha=0.25) + backward = self.dsrelu.GemmDsreluSm100(a, b, c, sfa, sfb, prob, alpha=0.5) + forward(a, b, sfa, sfb, prob) + backward(a, b, c, sfa, sfb, prob) + self.assertEqual(forward.captured_active_cluster_query, (2, 2)) + self.assertEqual(backward.captured_active_cluster_query, (2, 2)) + + calls = {} + + def kernel_type(label): + class Kernel: + def __init__(self, **configuration): + calls[f"{label}_configuration"] = configuration + + def __call__(self, *arguments, **keywords): + calls[f"{label}_arguments"] = arguments + calls[f"{label}_keywords"] = keywords + + return Kernel + + cutlass = types.ModuleType("cutlass") + cutlass.Float32 = lambda value: ("Float32", value) + cute = types.ModuleType("cutlass.cute") + cute.where = lambda condition, yes, no: yes if condition else no + cute.full_like = lambda value, fill: fill + cutlass.cute = cute + + forward_kernel_name = f"{self.operation_names[self.srelu]}.dense_blockscaled_gemm_persistent_srelu_quant" + forward_kernel = types.ModuleType(forward_kernel_name) + forward_kernel.Sm100BlockScaledPersistentDenseGemmKernel = kernel_type("forward") + backward_kernel_name = f"{self.operation_names[self.dsrelu]}.dense_blockscaled_gemm_persistent_dsrelu_quant" + backward_kernel = types.ModuleType(backward_kernel_name) + backward_kernel.Sm100BlockScaledPersistentDenseGemmKernel = kernel_type("backward") + + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + "cutlass.cute": cute, + forward_kernel_name: forward_kernel, + backward_kernel_name: backward_kernel, + }, + ): + forward.captured_call[1]["launch"]("stream", "A", "B", "SFA", "SFB", "PROB", "C", "D", "AMAX") + backward.captured_call[1]["launch"]( + "stream", + "A", + "B", + "C", + "SFA", + "SFB", + "PROB", + "D", + "DPROB", + "AMAX", + ) + + self.assertEqual( + calls["forward_arguments"], + ("A", "B", "SFA", "SFB", "C", "D", "PROB", "AMAX", None, None, ("Float32", 0.25), 10, "stream"), + ) + self.assertEqual( + calls["backward_arguments"], + ("A", "B", "SFA", "SFB", "C", "D", "PROB", "DPROB", "AMAX", None, None, ("Float32", 0.5), 10, "stream"), + ) + self.assertEqual(set(calls["forward_keywords"]), {"epilogue_op"}) + self.assertEqual(set(calls["backward_keywords"]), {"epilogue_op"}) + + def test_wrappers_are_jitted_with_static_configuration(self): + for wrapper_name in ("gemm_srelu_wrapper_sm100", "gemm_dsrelu_wrapper_sm100"): + static = self.jit_static_argnames[wrapper_name] + self.assertIn("mma_tiler_mn", static) + self.assertIn("cluster_shape_mn", static) + self.assertIn("sf_vec_size", static) + self.assertIn("a_layout", static) + self.assertIn("b_layout", static) + + def test_validation_rejects_bad_scale_shape_and_partial_mma_rows(self): + a, b, _, sfa, sfb, prob = self._samples() + bad_sfa = _Array((2, 2, 3, 32, 4, 4), "float8_e8m0fnu") + with self.assertRaisesRegex(ValueError, "SFA must have shape"): + self.srelu.GemmSreluSm100(a, b, bad_sfa, sfb, prob, sf_vec_size=32)(a, b, bad_sfa, sfb, prob) + + partial_a, partial_b, _, partial_sfa, partial_sfb, partial_prob = self._samples(m=129) + with self.assertRaisesRegex(ValueError, "CTA_TILE_M=128"): + self.srelu.GemmSreluSm100( + partial_a, + partial_b, + partial_sfa, + partial_sfb, + partial_prob, + sf_vec_size=32, + )(partial_a, partial_b, partial_sfa, partial_sfb, partial_prob) + + def test_shared_op_normalizes_legacy_uint8_storage_to_logical_fp4(self): + make_desc = self.tensor_module.make_compact_tensor_desc + a = make_desc(dtype=_DataType.UINT8, shape=(128, 128, 1), stride_order=(1, 0, 2), name="A") + b = make_desc(dtype=_DataType.UINT8, shape=(128, 128, 1), stride_order=(1, 0, 2), name="B") + c = make_desc(dtype=_DataType.BFLOAT16, shape=(128, 128, 1), stride_order=(1, 0, 2), name="C") + d = make_desc(dtype=_DataType.BFLOAT16, shape=(128, 128, 1), stride_order=(1, 0, 2), name="D") + sfa = make_desc(dtype=_DataType.FP8_E8M0, shape=(32, 4, 1, 4, 2, 1), stride_order=(3, 1, 0, 4, 2, 5), name="SFA") + sfb = make_desc(dtype=_DataType.FP8_E8M0, shape=(32, 4, 1, 4, 2, 1), stride_order=(3, 1, 0, 4, 2, 5), name="SFB") + prob = make_desc(dtype=_DataType.FLOAT, shape=(128, 1, 1), stride_order=(0, 1, 2), name="prob") + op = self.srelu.GemmSreluSm100Op( + a=a, + b=b, + c=c, + d=d, + sfa=sfa, + sfb=sfb, + prob=prob, + sf_vec_size=16, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + + self.assertTrue(op.check_support()) + self.assertEqual(op.ab_dtype, _DataType.FP4_E2M1) + + with self.assertRaisesRegex(ValueError, "at most 4"): + self.srelu.GemmSreluSm100Op( + a=a, + b=b, + c=c, + d=d, + sfa=sfa, + sfb=sfb, + prob=prob, + sf_vec_size=16, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(8, 1), + ).check_support() + + def test_relu_ops_enforce_forward_and_backward_specific_signatures(self): + make_desc = self.tensor_module.make_compact_tensor_desc + a = make_desc(dtype=_DataType.FP4_E2M1, shape=(128, 128, 1), stride_order=(1, 0, 2), name="A") + b = make_desc(dtype=_DataType.FP4_E2M1, shape=(128, 128, 1), stride_order=(1, 0, 2), name="B") + c = make_desc(dtype=_DataType.BFLOAT16, shape=(128, 128, 1), stride_order=(1, 0, 2), name="C") + d = make_desc(dtype=_DataType.BFLOAT16, shape=(128, 128, 1), stride_order=(1, 0, 2), name="D") + dprob = make_desc(dtype=_DataType.FLOAT, shape=(128, 1, 1), stride_order=(0, 1, 2), name="dprob") + sfa = make_desc(dtype=_DataType.FP8_E8M0, shape=(32, 4, 1, 4, 2, 1), stride_order=(3, 1, 0, 4, 2, 5), name="SFA") + sfb = make_desc(dtype=_DataType.FP8_E8M0, shape=(32, 4, 1, 4, 2, 1), stride_order=(3, 1, 0, 4, 2, 5), name="SFB") + prob = make_desc(dtype=_DataType.FLOAT, shape=(128, 1, 1), stride_order=(0, 1, 2), name="prob") + arguments = dict( + a=a, + b=b, + c=c, + d=d, + sfa=sfa, + sfb=sfb, + prob=prob, + sf_vec_size=16, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + + with self.assertRaisesRegex(ValueError, "dprob is only part of"): + self.srelu.GemmSreluSm100Op(**arguments, dprob=dprob).check_support() + with self.assertRaisesRegex(ValueError, "dprob is required"): + self.dsrelu.GemmDsreluSm100Op(**arguments).check_support() + self.assertTrue(self.dsrelu.GemmDsreluSm100Op(**arguments, dprob=dprob).check_support()) + + fp8_c = make_desc(dtype=_DataType.FP8_E4M3, shape=(128, 128, 1), stride_order=(1, 0, 2), name="C") + with self.assertRaisesRegex(ValueError, "dsReLU C must use"): + self.dsrelu.GemmDsreluSm100Op(**{**arguments, "c": fp8_c}, dprob=dprob).check_support() + + def test_shared_op_rejects_unimplemented_fp8_output(self): + make_desc = self.tensor_module.make_compact_tensor_desc + a = make_desc(dtype=_DataType.FP8_E4M3, shape=(128, 128, 1), stride_order=(1, 0, 2), name="A") + b = make_desc(dtype=_DataType.FP8_E4M3, shape=(128, 128, 1), stride_order=(1, 0, 2), name="B") + c = make_desc(dtype=_DataType.BFLOAT16, shape=(128, 128, 1), stride_order=(1, 0, 2), name="C") + d = make_desc(dtype=_DataType.FP8_E4M3, shape=(128, 128, 1), stride_order=(1, 0, 2), name="D") + sfa = make_desc(dtype=_DataType.FP8_E8M0, shape=(32, 4, 1, 4, 1, 1), stride_order=(3, 1, 0, 4, 2, 5), name="SFA") + sfb = make_desc(dtype=_DataType.FP8_E8M0, shape=(32, 4, 1, 4, 1, 1), stride_order=(3, 1, 0, 4, 2, 5), name="SFB") + prob = make_desc(dtype=_DataType.FLOAT, shape=(128, 1, 1), stride_order=(0, 1, 2), name="prob") + arguments = dict( + a=a, + b=b, + c=c, + d=d, + sfa=sfa, + sfb=sfb, + prob=prob, + sf_vec_size=32, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ) + with self.assertRaisesRegex(NotImplementedError, "SFD generation is not implemented"): + self.srelu.GemmSreluSm100Op(**arguments).check_support() + + sfd = make_desc(dtype=_DataType.FP8_E8M0, shape=(32, 4, 1, 4, 1, 1), stride_order=(3, 1, 0, 4, 2, 5), name="SFD") + norm_const = make_desc(dtype=_DataType.FLOAT, shape=(1,), name="norm_const") + with self.assertRaisesRegex(NotImplementedError, "SFD generation is not implemented"): + self.srelu.GemmSreluSm100Op( + **arguments, + sfd=sfd, + norm_const=norm_const, + ).check_support() + + def test_shared_op_rejects_scale_formats_that_do_not_match_the_mma(self): + make_desc = self.tensor_module.make_compact_tensor_desc + + def descriptor(dtype, shape, stride_order, name): + return make_desc(dtype=dtype, shape=shape, stride_order=stride_order, name=name) + + c = descriptor(_DataType.BFLOAT16, (128, 128, 1), (1, 0, 2), "C") + d = descriptor(_DataType.BFLOAT16, (128, 128, 1), (1, 0, 2), "D") + prob = descriptor(_DataType.FLOAT, (128, 1, 1), (0, 1, 2), "prob") + + cases = ( + (_DataType.FP8_E4M3, _DataType.FP8_E4M3, 16, 2, "FP8 inputs require FP8_E8M0"), + (_DataType.FP4_E2M1, _DataType.FP8_E4M3, 32, 1, "FP4 inputs with sf_vec_size=32 require FP8_E8M0"), + ) + for ab_dtype, sf_dtype, sf_vec_size, rest_k, message in cases: + a = descriptor(ab_dtype, (128, 128, 1), (1, 0, 2), "A") + b = descriptor(ab_dtype, (128, 128, 1), (1, 0, 2), "B") + sfa = descriptor(sf_dtype, (32, 4, 1, 4, rest_k, 1), (3, 1, 0, 4, 2, 5), "SFA") + sfb = descriptor(sf_dtype, (32, 4, 1, 4, rest_k, 1), (3, 1, 0, 4, 2, 5), "SFB") + with self.subTest(ab_dtype=ab_dtype, sf_dtype=sf_dtype, sf_vec_size=sf_vec_size), self.assertRaisesRegex(ValueError, message): + self.srelu.GemmSreluSm100Op( + a=a, + b=b, + c=c, + d=d, + sfa=sfa, + sfb=sfb, + prob=prob, + sf_vec_size=sf_vec_size, + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + ).check_support() + + def test_jax_modules_do_not_load_torch_or_kernels(self): + self.assertNotIn(f"{_PACKAGE}.gemm_srelu.api", sys.modules) + self.assertNotIn(f"{_PACKAGE}.gemm_dsrelu.api", sys.modules) + self.assertNotIn(f"{_PACKAGE}.gemm_srelu.dense_blockscaled_gemm_persistent_srelu_quant", sys.modules) + self.assertNotIn(f"{_PACKAGE}.gemm_dsrelu.dense_blockscaled_gemm_persistent_dsrelu_quant", sys.modules) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_gemm_swiglu.py b/test/python/fe_api/test_jax_gemm_swiglu.py new file mode 100644 index 000000000..28d3a18d2 --- /dev/null +++ b/test/python/fe_api/test_jax_gemm_swiglu.py @@ -0,0 +1,242 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX integration tests for the standard dense GEMM + SwiGLU adapter.""" + +import pytest + + +def _jax_runtime(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + cutlass_jax = pytest.importorskip("cutlass.jax") + if not cutlass_jax.is_available(): + pytest.skip("Installed JAX version is unsupported by CUTLASS JAX") + return jax, jnp + + +def _compute_capability(device): + reported = getattr(device, "compute_capability", None) + if isinstance(reported, (tuple, list)): + major, minor = int(reported[0]), int(reported[1]) + else: + text = str(reported) + if "." in text: + major_text, minor_text = text.split(".", 1) + major, minor = int(major_text), int(minor_text) + else: + major, minor = divmod(int(text), 10) + return major * 10 + minor + + +def _supported_gpu(jax): + devices = tuple(device for device in jax.local_devices() if device.platform == "gpu") + try: + capabilities = {_compute_capability(device) for device in devices} + except (TypeError, ValueError): + pytest.skip("JAX GPU compute capability is unavailable") + if not devices or len(capabilities) != 1: + pytest.skip("A homogeneous JAX SM100-family GPU configuration is required") + capability = capabilities.pop() + if capability not in {100, 103, 107}: + pytest.skip("A supported JAX SM100-family GPU is not available") + return devices[0] + + +@pytest.mark.L0 +def test_jax_gemm_swiglu_abstract_contract( + monkeypatch, +): + jax, jnp = _jax_runtime() + + from cudnn._jax import JaxApiBase + from cudnn.jax import GemmSwigluSm100, gemm_swiglu_wrapper_sm100 + + def abstract_call(self, _inputs, *, output_descs, output_spec=None, **_options): + if output_spec is None: + output_spec = tuple(self._to_tensor_spec(desc) for desc in output_descs) + return tuple( + jnp.empty( + self._to_shape_dtype_struct(desc, mode=spec.mode).shape, + dtype=desc.dtype, + ) + for desc, spec in zip(output_descs, output_spec) + ) + + monkeypatch.setattr(JaxApiBase, "_call_kernel", abstract_call) + monkeypatch.setattr( + JaxApiBase, + "_resolve_compute_capability", + staticmethod(lambda _target, _supported, _operation, **_options: 100), + ) + + sample_a = jax.ShapeDtypeStruct((3, 128, 64), jnp.bfloat16) + sample_b = jax.ShapeDtypeStruct((3, 192, 64), jnp.bfloat16) + api = GemmSwigluSm100( + sample_a, + sample_b, + ) + result = jax.eval_shape(api, sample_a, sample_b) + assert tuple(result.keys()) == ( + "ab12_tensor", + "c_tensor", + "sfc_tensor", + "amax_tensor", + ) + assert result["ab12_tensor"].shape == (3, 128, 192) + assert result["ab12_tensor"].dtype == jnp.float32 + assert result["c_tensor"].shape == (3, 128, 96) + assert result["c_tensor"].dtype == jnp.float16 + assert result["sfc_tensor"] is None + assert result["amax_tensor"] is None + + alternate_a = jax.ShapeDtypeStruct((3, 64, 128), jnp.bfloat16) + alternate_b = jax.ShapeDtypeStruct((3, 64, 192), jnp.bfloat16) + alternate = jax.eval_shape( + lambda a, b: gemm_swiglu_wrapper_sm100( + a, + b, + a_layout="LKM", + b_layout="LKN", + c_layout="LNM", + c_dtype=jnp.bfloat16, + ), + alternate_a, + alternate_b, + ) + assert alternate["ab12_tensor"].shape == (3, 192, 128) + assert alternate["c_tensor"].shape == (3, 96, 128) + assert alternate["c_tensor"].dtype == jnp.bfloat16 + + explicit = GemmSwigluSm100( + sample_a, + sample_b, + sample_ab12=jax.ShapeDtypeStruct((3, 128, 192), jnp.bfloat16), + sample_c=jax.ShapeDtypeStruct((3, 128, 96), jnp.bfloat16), + ) + explicit_result = jax.eval_shape(explicit, sample_a, sample_b) + assert explicit_result["ab12_tensor"].dtype == jnp.bfloat16 + assert explicit_result["c_tensor"].dtype == jnp.bfloat16 + + +@pytest.mark.L0 +def test_jax_gemm_swiglu_block_scaled_sm100_jit_and_numerics(): + jax, jnp = _jax_runtime() + device = _supported_gpu(jax) + if not hasattr(jnp, "float8_e4m3fn") or not hasattr(jnp, "float8_e8m0fnu"): + pytest.skip("The installed JAX runtime does not expose MXFP8 dtypes") + + from cudnn.jax import gemm_swiglu_wrapper_sm100 + + batch, m, n, k = 1, 128, 128, 128 + alpha = 0.5 + a = jax.device_put( + (jax.random.normal(jax.random.key(10), (batch, m, k), dtype=jnp.float32) * 0.1).astype(jnp.float8_e4m3fn), + device, + ) + b = jax.device_put( + (jax.random.normal(jax.random.key(11), (batch, n, k), dtype=jnp.float32) * 0.1).astype(jnp.float8_e4m3fn), + device, + ) + scale_shape = (batch, 1, 1, 32, 4, 4) + sfa = jax.device_put(jnp.ones(scale_shape, dtype=jnp.float32).astype(jnp.float8_e8m0fnu), device) + sfb = jax.device_put(jnp.ones(scale_shape, dtype=jnp.float32).astype(jnp.float8_e8m0fnu), device) + + lowered = gemm_swiglu_wrapper_sm100.lower( + a, + b, + alpha=alpha, + ab12_dtype=jnp.bfloat16, + c_dtype=jnp.bfloat16, + sfa_tensor=sfa, + sfb_tensor=sfb, + sf_vec_size=32, + ) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + result = lowered.compile()(a, b, sfa_tensor=sfa, sfb_tensor=sfb) + result["c_tensor"].block_until_ready() + reference_ab12 = alpha * jnp.einsum("lmk,lnk->lmn", a.astype(jnp.float32), b.astype(jnp.float32)) + blocks = reference_ab12.reshape(batch, m, n // 32, 32) + input_blocks = blocks[:, :, 0::2].reshape(batch, m, n // 2) + gate_blocks = blocks[:, :, 1::2].reshape(batch, m, n // 2) + reference_c = input_blocks * jax.nn.silu(gate_blocks) + + assert jnp.allclose(result["ab12_tensor"].astype(jnp.float32), reference_ab12, atol=8e-2, rtol=4e-2) + assert jnp.allclose(result["c_tensor"].astype(jnp.float32), reference_c, atol=1e-1, rtol=5e-2) + assert result["sfc_tensor"] is None + assert result["amax_tensor"] is None + + +@pytest.mark.L0 +def test_jax_gemm_swiglu_sm100_jit_and_numerics(): + jax, jnp = _jax_runtime() + device = _supported_gpu(jax) + + from cudnn.jax import gemm_swiglu_wrapper_sm100 + + batch, m, n, k = 1, 128, 128, 128 + alpha = 0.5 + a = jax.device_put( + jax.random.normal( + jax.random.key(0), + (batch, k, m), + dtype=jnp.bfloat16, + ) + * jnp.bfloat16(0.1), + device, + ) + b = jax.device_put( + jax.random.normal( + jax.random.key(1), + (batch, k, n), + dtype=jnp.bfloat16, + ) + * jnp.bfloat16(0.1), + device, + ) + + lowered = gemm_swiglu_wrapper_sm100.lower( + a, + b, + alpha=alpha, + a_layout="LKM", + b_layout="LKN", + c_layout="LNM", + ) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + compiled = lowered.compile() + + def reference(a_value): + ab12_lmn = alpha * jnp.einsum( + "lmk,lnk->lmn", + a_value.transpose(0, 2, 1).astype(jnp.float32), + b.transpose(0, 2, 1).astype(jnp.float32), + ) + blocks = ab12_lmn.reshape(batch, m, n // 32, 32) + input_blocks = blocks[:, :, 0::2].reshape(batch, m, n // 2) + gate_blocks = blocks[:, :, 1::2].reshape(batch, m, n // 2) + c_lmn = input_blocks * jax.nn.silu(gate_blocks) + return ab12_lmn.transpose(0, 2, 1), c_lmn.transpose(0, 2, 1) + + for a_value in (a, -a): + result = compiled(a_value, b) + result["c_tensor"].block_until_ready() + reference_ab12, reference_c = reference(a_value) + assert jnp.allclose( + result["ab12_tensor"], + reference_ab12, + atol=5e-2, + rtol=2e-2, + ) + assert jnp.allclose( + result["c_tensor"].astype(jnp.float32), + reference_c, + atol=8e-2, + rtol=3e-2, + ) diff --git a/test/python/fe_api/test_jax_gemm_swiglu_contract.py b/test/python/fe_api/test_jax_gemm_swiglu_contract.py new file mode 100644 index 000000000..8d3e8611f --- /dev/null +++ b/test/python/fe_api/test_jax_gemm_swiglu_contract.py @@ -0,0 +1,616 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for the standard JAX GEMM + SwiGLU adapter.""" + +from __future__ import annotations + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +import os +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "gemm_swiglu" +_PACKAGE = "cudnn_jax_gemm_swiglu_contract_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + FP8_E4M3 = auto() + FP8_E5M2 = auto() + FP8_E8M0 = auto() + FP4_E2M1 = auto() + + +_DTYPE_TO_CUDNN = { + "float16": _DataType.HALF, + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, + "float4_e2m1fn": _DataType.FP4_E2M1, + "float8_e4m3fn": _DataType.FP8_E4M3, + "float8_e5m2": _DataType.FP8_E5M2, + "float8_e8m0fnu": _DataType.FP8_E8M0, +} +_CUDNN_TO_DTYPE = {value: key for key, value in _DTYPE_TO_CUDNN.items()} + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + +class _TensorSpec: + def __init__(self, *, layout, mode, divisibility=None): + self.layout = layout + self.mode = mode + self.divisibility = divisibility + + +class JaxGemmSwigluContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + operation_name = f"{_PACKAGE}.gemm_swiglu" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + layout_module = importlib.import_module(f"{internal_name}.layout") + result_module = importlib.import_module(f"{_PACKAGE}.common.result") + + class JaxTensorDesc(tensor_module.TensorDesc): + @classmethod + def from_shape(cls, shape, dtype, *, name, mode=None, init_value=None): + public_shape = tuple(shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + public_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride(public_shape, public_order) + canonical_axis_by_public_axis = layout_module.to_public_axes(tuple(range(len(public_shape))), mode) + desc = cls( + dtype=dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple(canonical_axis_by_public_axis[axis] for axis in public_order), + name=name, + init_value=init_value, + ) + object.__setattr__(desc, "mode", mode) + return desc + + @property + def cudnn_dtype(self): + return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) + + class JaxApiBase: + @staticmethod + def _to_tensor_desc(value, name, *, mode=None, init_value=None, **_unused): + public_shape = tuple(value.shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + public_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride(public_shape, public_order) + canonical_axis_by_public_axis = layout_module.to_public_axes(tuple(range(len(public_shape))), mode) + desc = JaxTensorDesc( + dtype=value.dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple(canonical_axis_by_public_axis[axis] for axis in public_order), + name=name, + init_value=init_value, + ) + object.__setattr__(desc, "mode", mode) + return desc + + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del operation_name + resolved = 100 if target is None else target + if resolved not in supported: + raise ValueError(f"unsupported target {resolved}") + return resolved + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + actual_shape = layout_module.to_canonical_axes(tuple(value.shape), mode) + if actual_shape != expected.shape: + raise ValueError(f"{expected.name} tensor shape mismatch: " f"expected {expected.shape}, got {actual_shape}") + actual_dtype = _DTYPE_TO_CUDNN.get(value.dtype, _DataType.NOT_SET) + if actual_dtype != expected.cudnn_dtype: + raise ValueError(f"{expected.name} tensor dtype mismatch: " f"expected {expected.cudnn_dtype}, got {actual_dtype}") + + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + return _TensorSpec( + layout=layout_module.to_cutlass_layout( + desc.shape, + desc.stride, + desc.stride_order, + mode=mode, + name=desc.name, + ), + mode=mode, + divisibility=divisibility, + ) + + def _call_kernel(self, inputs, *, launch, **options): + input_descs = options.get("input_descs") + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc, mode=desc.mode) + options["input_spec"] = tuple(self._to_tensor_spec(desc, mode=desc.mode) for desc in input_descs) + if "output_spec" not in options: + options["output_spec"] = tuple( + self._to_tensor_spec(desc, mode=getattr(desc, "mode", None)) for desc in options["output_descs"] + ) + options["launch"] = launch + self.captured_call = (tuple(inputs), options) + return tuple( + _Array( + layout_module.to_public_axes(desc.shape, spec.mode), + _CUDNN_TO_DTYPE[desc.cudnn_dtype], + ) + for desc, spec in zip(options["output_descs"], options["output_spec"]) + ) + + def _get_max_active_clusters(self, cluster_size, *, overlap_margin=0): + if not hasattr(self, "captured_active_cluster_queries"): + self.captured_active_cluster_queries = [] + self.captured_active_cluster_queries.append((cluster_size, overlap_margin)) + return 12 - overlap_margin + + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = result_module.TupleDict + + datatypes = types.ModuleType(f"{internal_name}.datatypes") + datatypes.jax_to_cudnn_dtype = lambda dtype: _DTYPE_TO_CUDNN.get(dtype, _DataType.NOT_SET) + datatypes.normalize_jax_dtype = lambda value, default, name: default if value is None else value + datatypes.cudnn_to_jax_dtype = lambda dtype: _CUDNN_TO_DTYPE[dtype] + sys.modules[datatypes.__name__] = datatypes + + cls.jit_static_argnames = {} + fake_jax = types.ModuleType("jax") + + def jit(function=None, *, static_argnames=()): + def decorate(target): + cls.jit_static_argnames[target.__name__] = tuple(static_argnames) + return target + + return decorate if function is None else decorate(function) + + fake_jax.jit = jit + fake_jax.ShapeDtypeStruct = _Array + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.dtype = lambda dtype: dtype + fake_jnp.float16 = "float16" + fake_jnp.bfloat16 = "bfloat16" + fake_jnp.float32 = "float32" + fake_jnp.float4_e2m1fn = "float4_e2m1fn" + fake_jnp.float8_e4m3fn = "float8_e4m3fn" + fake_jnp.float8_e5m2 = "float8_e5m2" + fake_jnp.float8_e8m0fnu = "float8_e8m0fnu" + fake_jax.numpy = fake_jnp + + try: + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "torch": None, + "cutlass": None, + }, + ): + cls.module = importlib.import_module(f"{operation_name}.jax") + except Exception: + cls.tearDownClass() + raise + + cls.JaxApiBase = JaxApiBase + cls.operation_name = operation_name + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def setUp(self) -> None: + self.JaxApiBase.captured_call = None + + @staticmethod + def _samples(a_layout="LMK", b_layout="LNK"): + m, n, k, batch = 128, 192, 64, 3 + a_shape = (batch, m, k) if a_layout == "LMK" else (batch, k, m) + b_shape = (batch, n, k) if b_layout == "LNK" else (batch, k, n) + return _Array(a_shape, "bfloat16"), _Array(b_shape, "bfloat16") + + def test_default_layouts_infer_outputs_and_declare_cutlass_call(self): + sample_a, sample_b = self._samples() + api = self.module.GemmSwigluSm100( + sample_a, + sample_b, + ) + + self.assertEqual(api.a_desc.mode, (1, 2, 0)) + self.assertEqual(api.b_desc.mode, (1, 2, 0)) + self.assertEqual(api.c_desc.mode, (1, 2, 0)) + self.assertEqual(api.a_desc.shape, (128, 64, 3)) + self.assertEqual(api.b_desc.shape, (192, 64, 3)) + self.assertEqual(api.ab12_desc.shape, (128, 192, 3)) + self.assertEqual(api.c_desc.shape, (128, 96, 3)) + self.assertIs(api._op.a, api.a_desc) + self.assertIs(api._op.b, api.b_desc) + self.assertIs(api._op.ab12, api.ab12_desc) + self.assertIs(api._op.c, api.c_desc) + + result = api(sample_a, sample_b) + self.assertEqual( + tuple(result.keys()), + ("ab12_tensor", "c_tensor", "sfc_tensor", "amax_tensor"), + ) + self.assertEqual(result["ab12_tensor"].shape, (3, 128, 192)) + self.assertEqual(result["ab12_tensor"].dtype, "float32") + self.assertEqual(result["c_tensor"].shape, (3, 128, 96)) + self.assertEqual(result["c_tensor"].dtype, "float16") + self.assertIsNone(result["sfc_tensor"]) + self.assertIsNone(result["amax_tensor"]) + + inputs, options = api.captured_call + self.assertEqual(inputs, (sample_a, sample_b)) + self.assertTrue(callable(options["launch"])) + self.assertEqual(options["output_descs"], (api.ab12_desc, api.c_desc)) + self.assertEqual( + tuple(spec.mode for spec in options["input_spec"]), + (api.a_desc.mode, api.b_desc.mode), + ) + self.assertEqual( + tuple(spec.mode for spec in options["output_spec"]), + (api.ab12_desc.mode, api.c_desc.mode), + ) + self.assertEqual( + tuple(spec.layout for spec in options["input_spec"]), + ((2, 1, 0), (2, 1, 0)), + ) + self.assertEqual( + tuple(spec.layout for spec in options["output_spec"]), + ((2, 1, 0), (2, 1, 0)), + ) + self.assertIn("--gpu-arch sm_100a", options["compile_options"]) + + def test_alternate_layouts_map_public_shapes_to_the_same_canonical_op(self): + sample_a, sample_b = self._samples(a_layout="LKM", b_layout="LKN") + api = self.module.GemmSwigluSm100( + sample_a, + sample_b, + a_layout="LKM", + b_layout="LKN", + c_layout="LNM", + ) + result = api(sample_a, sample_b) + + self.assertEqual(api.a_desc.mode, (2, 1, 0)) + self.assertEqual(api.b_desc.mode, (2, 1, 0)) + self.assertEqual(api.c_desc.mode, (2, 1, 0)) + self.assertEqual(api.a_desc.shape, (128, 64, 3)) + self.assertEqual(api.b_desc.shape, (192, 64, 3)) + self.assertEqual(api.a_desc.stride_order, (0, 1, 2)) + self.assertEqual(api.b_desc.stride_order, (0, 1, 2)) + self.assertEqual(api.ab12_desc.stride_order, (0, 1, 2)) + self.assertEqual(api.c_desc.stride_order, (0, 1, 2)) + self.assertEqual(result["ab12_tensor"].shape, (3, 192, 128)) + self.assertEqual(result["c_tensor"].shape, (3, 96, 128)) + self.assertIn("--gpu-arch sm_100a", api.captured_call[1]["compile_options"]) + + def test_explicit_output_exemplars_are_checked_and_preserved(self): + sample_a, sample_b = self._samples() + sample_ab12 = _Array((3, 128, 192), "bfloat16") + sample_c = _Array((3, 128, 96), "bfloat16") + api = self.module.GemmSwigluSm100( + sample_a, + sample_b, + sample_ab12=sample_ab12, + sample_c=sample_c, + ) + result = api(sample_a, sample_b) + self.assertEqual(result["ab12_tensor"].dtype, "bfloat16") + self.assertEqual(result["c_tensor"].dtype, "bfloat16") + + with self.assertRaisesRegex(ValueError, "sample_ab12 and sample_c must be provided together"): + self.module.GemmSwigluSm100( + sample_a, + sample_b, + sample_ab12=sample_ab12, + ) + with self.assertRaisesRegex(ValueError, "ab12_dtype and c_dtype cannot be specified"): + self.module.GemmSwigluSm100( + sample_a, + sample_b, + sample_ab12=sample_ab12, + sample_c=sample_c, + c_dtype="float16", + ) + + def test_runtime_inputs_must_match_the_specialized_signature(self): + sample_a, sample_b = self._samples() + api = self.module.GemmSwigluSm100( + sample_a, + sample_b, + ) + with self.assertRaisesRegex(ValueError, "sample_a tensor shape mismatch"): + api(_Array((3, 64, 64), "bfloat16"), sample_b) + with self.assertRaisesRegex(ValueError, "sample_b tensor dtype mismatch"): + api(sample_a, _Array(sample_b.shape, "float16")) + + def test_rejects_fp8_c_and_block_only_standard_options(self): + sample_a, sample_b = self._samples() + sample_ab12 = _Array((3, 128, 192), "bfloat16") + sample_c_fp8 = _Array((3, 128, 96), "float8_e4m3fn") + with self.assertRaisesRegex(ValueError, "C dtype must be float16 or bfloat16"): + self.module.GemmSwigluSm100( + sample_a, + sample_b, + sample_ab12=sample_ab12, + sample_c=sample_c_fp8, + ).check_support() + + for option, value in (("sf_vec_size", 32), ("vector_f32", True), ("ab12_stages", 3)): + with self.subTest(option=option), self.assertRaisesRegex(ValueError, "only applies to block-scaled"): + self.module.GemmSwigluSm100(sample_a, sample_b, **{option: value}) + + def test_wrapper_marks_configuration_as_static_and_remains_functional(self): + self.assertIn("gemm_swiglu_wrapper_sm100", self.module.__all__) + self.assertNotIn("gemm_swiglu_wrapper", self.module.__all__) + static = self.jit_static_argnames["gemm_swiglu_wrapper_sm100"] + self.assertEqual( + static, + ( + "alpha", + "c_layout", + "ab12_dtype", + "c_dtype", + "acc_dtype", + "mma_tiler_mn", + "cluster_shape_mn", + "sf_vec_size", + "vector_f32", + "ab12_stages", + "a_layout", + "b_layout", + ), + ) + + sample_a, sample_b = self._samples(a_layout="LKM", b_layout="LKN") + result = self.module.gemm_swiglu_wrapper_sm100( + sample_a, + sample_b, + c_layout="LNM", + a_layout="LKM", + b_layout="LKN", + ) + self.assertEqual(result["ab12_tensor"].shape, (3, 192, 128)) + self.assertEqual(result["c_tensor"].shape, (3, 96, 128)) + + def test_block_scaled_mxfp8_uses_packed_scale_modes_and_kernel_abi(self): + batch, m, n, k = 2, 128, 128, 128 + sample_a = _Array((batch, m, k), "float8_e4m3fn") + sample_b = _Array((batch, n, k), "float8_e4m3fn") + scale_shape = (batch, 1, 1, 32, 4, 4) + sample_sfa = _Array(scale_shape, "float8_e8m0fnu") + sample_sfb = _Array(scale_shape, "float8_e8m0fnu") + api = self.module.GemmSwigluSm100( + sample_a, + sample_b, + sample_sfa=sample_sfa, + sample_sfb=sample_sfb, + alpha=0.5, + ab12_dtype="bfloat16", + c_dtype="bfloat16", + sf_vec_size=32, + vector_f32=True, + ab12_stages=3, + ) + + self.assertTrue(api.is_block_scaled) + self.assertEqual(api.sfa_desc.mode, (3, 4, 1, 5, 2, 0)) + self.assertEqual(api.sfa_desc.shape, (32, 4, 1, 4, 1, batch)) + self.assertEqual(api.sfa_desc.stride_order, (3, 1, 0, 4, 2, 5)) + result = api(sample_a, sample_b, sample_sfa, sample_sfb) + self.assertEqual(result["ab12_tensor"].shape, (batch, m, n)) + self.assertEqual(result["ab12_tensor"].dtype, "bfloat16") + self.assertEqual(result["c_tensor"].shape, (batch, m, n // 2)) + self.assertEqual(result["c_tensor"].dtype, "bfloat16") + self.assertIsNone(result["sfc_tensor"]) + self.assertIsNone(result["amax_tensor"]) + + inputs, options = api.captured_call + self.assertEqual(inputs, (sample_a, sample_b, sample_sfa, sample_sfb)) + self.assertEqual( + tuple(spec.mode for spec in options["input_spec"]), + (api.a_desc.mode, api.b_desc.mode, api.sfa_desc.mode, api.sfb_desc.mode), + ) + self.assertEqual( + tuple(spec.mode for spec in options["output_spec"]), + (api.ab12_desc.mode, api.c_desc.mode), + ) + self.assertEqual(api.captured_active_cluster_queries, [(1, 0)]) + + seen = {} + + class Kernel: + def __init__(self, **configuration): + seen["configuration"] = configuration + + def __call__(self, *arguments): + seen["arguments"] = arguments + + cutlass = types.ModuleType("cutlass") + cutlass.Float32 = lambda value: ("Float32", value) + kernel_module = types.ModuleType(f"{self.operation_name}.dense_blockscaled_gemm_persistent_swiglu_interleaved_quant") + kernel_module.Sm100BlockScaledPersistentDenseGemmKernel = Kernel + stream = object() + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + kernel_module.__name__: kernel_module, + }, + ): + options["launch"](stream, "A", "B", "SFA", "SFB", "AB12", "C") + + self.assertEqual( + seen["configuration"], + { + "sf_vec_size": 32, + "mma_tiler_mn": (128, 128), + "cluster_shape_mn": (1, 1), + "vector_f32": True, + "ab12_stages": 3, + }, + ) + self.assertEqual( + seen["arguments"], + ("A", "B", "SFA", "SFB", "C", "AB12", None, None, None, ("Float32", 0.5), 12, stream), + ) + + wrapped = self.module.gemm_swiglu_wrapper_sm100( + sample_a, + sample_b, + alpha=0.5, + ab12_dtype="bfloat16", + c_dtype="bfloat16", + sfa_tensor=sample_sfa, + sfb_tensor=sample_sfb, + sf_vec_size=32, + ) + self.assertEqual(wrapped["ab12_tensor"].shape, (batch, m, n)) + self.assertEqual(wrapped["c_tensor"].shape, (batch, m, n // 2)) + + def test_block_scaled_fp4_infers_initialized_amax(self): + sample_a = _Array((1, 128, 128), "float4_e2m1fn") + sample_b = _Array((1, 128, 128), "float4_e2m1fn") + scale_shape = (1, 1, 2, 32, 4, 4) + sample_sfa = _Array(scale_shape, "float8_e8m0fnu") + sample_sfb = _Array(scale_shape, "float8_e8m0fnu") + api = self.module.GemmSwigluSm100( + sample_a, + sample_b, + sample_sfa=sample_sfa, + sample_sfb=sample_sfb, + c_dtype="bfloat16", + sf_vec_size=16, + ) + + result = api(sample_a, sample_b, sample_sfa, sample_sfb) + self.assertEqual(result["amax_tensor"].shape, (1,)) + self.assertEqual(result["amax_tensor"].dtype, "float32") + self.assertEqual(api.amax_desc.init_value, float("-inf")) + self.assertEqual(api.captured_call[1]["output_descs"], (api.ab12_desc, api.c_desc, api.amax_desc)) + + with self.assertRaisesRegex(ValueError, "sample_sfa and sample_sfb are required"): + self.module.GemmSwigluSm100(sample_a, sample_b, sample_sfa=sample_sfa) + + def test_launch_constructs_the_kernel_internally_and_preserves_abi_order(self): + _, sample_b = self._samples() + sample_a = _Array((3, 256, 64), "bfloat16") + with mock.patch.dict(os.environ, {"CUDNNFE_CLUSTER_OVERLAP_MARGIN": "2"}): + api = self.module.GemmSwigluSm100( + sample_a, + sample_b, + alpha=0.25, + mma_tiler_mn=(256, 128), + ) + self.assertFalse(hasattr(api, "captured_active_cluster_queries")) + api(sample_a, sample_b) + self.assertEqual(api.captured_active_cluster_queries, [(4, 2)]) + launch = api.captured_call[1]["launch"] + seen = {} + + class HardwareInfo: + def get_max_active_clusters(self, cluster_size): + raise AssertionError(f"HardwareInfo queried from launch callback for cluster size {cluster_size}") + + class Kernel: + def __init__(self, **configuration): + seen["configuration"] = configuration + + def __call__(self, *arguments): + seen["arguments"] = arguments + + cutlass = types.ModuleType("cutlass") + cutlass.Float32 = lambda value: ("Float32", value) + cutlass.utils = types.SimpleNamespace(HardwareInfo=HardwareInfo) + cutlass_jax = types.ModuleType("cutlass.jax") + cutlass_jax.jax_to_cutlass_dtype = lambda dtype: ("cutlass_dtype", dtype) + cutlass.jax = cutlass_jax + kernel_module = types.ModuleType(f"{self.operation_name}.dense_gemm_persistent_swiglu") + kernel_module.PersistentDenseGemmKernel = Kernel + stream = object() + + with mock.patch.dict( + sys.modules, + { + "cutlass": cutlass, + "cutlass.jax": cutlass_jax, + kernel_module.__name__: kernel_module, + }, + ): + launch(stream, "A", "B", "AB12", "C") + + self.assertEqual( + seen["configuration"], + { + "acc_dtype": ("cutlass_dtype", "float32"), + "use_2cta_instrs": True, + "mma_tiler_mn": (256, 128), + "cluster_shape_mn": (2, 2), + }, + ) + self.assertEqual( + seen["arguments"], + ("A", "B", "AB12", "C", ("Float32", 0.25), 10, stream), + ) + + with self.assertRaises(TypeError): + launch(stream, "A", "B", "AB12", "C", object()) + + def test_adapter_import_does_not_load_torch_or_the_cute_kernel(self): + self.assertNotIn(f"{self.operation_name}.api", sys.modules) + self.assertNotIn(f"{self.operation_name}.dense_gemm_persistent_swiglu", sys.modules) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_grouped_gemm_contract.py b/test/python/fe_api/test_jax_grouped_gemm_contract.py new file mode 100644 index 000000000..97cd8c8b1 --- /dev/null +++ b/test/python/fe_api/test_jax_grouped_gemm_contract.py @@ -0,0 +1,624 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Source contracts for JAX grouped GEMM adapters.""" + +from __future__ import annotations + +import ast +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +from unittest import mock + +import pytest + +pytestmark = pytest.mark.L0 + +_ROOT = Path(__file__).resolve().parents[3] +_GROUPED_ROOT = _ROOT / "python" / "cudnn" / "grouped_gemm" +_FAMILIES = ( + "swiglu", + "dswiglu", + "quant", + "srelu", + "dsrelu", + "glu", + "glu_hadamard", + "dglu", + "wgrad", +) + +_NATIVE_FP4_FAMILIES = ( + "swiglu", + "dswiglu", + "quant", + "srelu", + "dsrelu", + "glu", + "dglu", +) + + +def _class_stem(family: str) -> str: + return "".join(part.capitalize() for part in family.split("_")) + + +def _jax_path(family: str) -> Path: + return _GROUPED_ROOT / f"grouped_gemm_{family}" / "jax.py" + + +@pytest.mark.parametrize("family", _FAMILIES) +def test_each_torch_family_has_a_jax_class_and_wrapper(family: str): + tree = ast.parse(_jax_path(family).read_text()) + definitions = { + node.name: node + for node in tree.body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) + } + + assert f"GroupedGemm{_class_stem(family)}Sm100" in definitions + wrapper = definitions[f"grouped_gemm_{family}_wrapper_sm100"] + assert isinstance(wrapper, ast.FunctionDef) + assert wrapper.decorator_list + assert "jax.jit" in ast.unparse(wrapper.decorator_list[0]) + + +@pytest.mark.parametrize("family", _FAMILIES) +def test_launchers_receive_precomputed_occupancy_and_preserve_stream_first(family: str): + source = _jax_path(family).read_text() + tree = ast.parse(source) + launcher = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "_launch" + ) + + assert launcher.args.args[0].arg == "stream" + assert "max_active_clusters" in { + argument.arg for argument in launcher.args.kwonlyargs + } + assert "HardwareInfo" not in source + assert "resolve_max_active_clusters" not in source + + +@pytest.mark.parametrize("family", _FAMILIES) +def test_jax_module_has_no_top_level_torch_or_kernel_import(family: str): + tree = ast.parse(_jax_path(family).read_text()) + imported_modules = [] + for node in tree.body: + if isinstance(node, ast.Import): + imported_modules.extend(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + imported_modules.append(node.module or "") + + assert not any( + module == "torch" or module.startswith("torch.") for module in imported_modules + ) + assert not any( + "grouped_gemm_" in module and module != "_jax_api" + for module in imported_modules + ) + + +def test_shared_adapter_uses_current_jax_api_base_and_call_kernel(): + source = (_GROUPED_ROOT / "_jax_api.py").read_text() + tree = ast.parse(source) + classes = {node.name: node for node in tree.body if isinstance(node, ast.ClassDef)} + + assert ast.unparse(classes["ApiBaseJax"].bases[0]) == "JaxApiBase" + assert "_CALLER._call_kernel(" in source + assert "_CALLER._get_max_active_clusters(" in source + assert source.index("_CALLER._get_max_active_clusters(") < source.index( + "def launch(stream:" + ) + + +def test_grouped_layouts_are_explicit_row_major_public_mappings(): + source = (_GROUPED_ROOT / "_jax_api.py").read_text() + + assert "PROBABILITY_MODE = (2, 1, 0)" in source + assert "GROUPED_BIAS_MODE = (1, 0)" in source + assert "GROUPED_WORKSPACE_ALIGNMENT = 128" in source + assert "_canonical_block_scale_shape(rows, k, batch, sf_vec_size)" in source + assert "BLOCK_SCALE_MODE" in source + + +@pytest.mark.parametrize("family", ("dswiglu", "dglu", "dsrelu")) +def test_dprob_descriptors_use_public_probability_shape(family: str): + tree = ast.parse(_jax_path(family).read_text()) + dprob_calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "make_buffer_desc" + and node.args + and isinstance(node.args[0], ast.Constant) + and node.args[0].value == "dprob_tensor" + ] + + assert dprob_calls + assert all(ast.unparse(call.args[1]) == "(1, 1, m)" for call in dprob_calls) + + +def test_grouped_buffers_use_tensor_descriptors_directly(): + source = (_GROUPED_ROOT / "_jax_api.py").read_text() + + assert "class BufferSpec" not in source + assert "class _SampleDesc" not in source + assert "def make_buffer_desc(" in source + assert "JaxTensorDesc.from_shape(" in source + assert "JaxTensorDesc.from_array(" in source + assert "init_value=init_value" in source + + +@pytest.mark.parametrize("family", _FAMILIES) +def test_grouped_lowerings_pass_descriptors_without_parallel_specs(family: str): + tree = ast.parse(_jax_path(family).read_text()) + calls = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "call_cutedsl" + ] + + assert calls + for call in calls: + keywords = {keyword.arg for keyword in call.keywords} + assert "input_descs" in keywords + assert "outputs" in keywords + assert not { + "input_specs", + "output_specs", + "workspace_specs", + } & keywords + + +@pytest.mark.parametrize("family", _NATIVE_FP4_FAMILIES) +def test_dense_grouped_adapters_use_native_fp4_recipe_validation(family: str): + source = _jax_path(family).read_text() + + assert "require_grouped_input_scales(" in source + assert "jnp.float4_e2m1fn" in source or "is_fp4_dtype(ab_dtype)" in source + assert "torch.uint8" not in source + + +@pytest.mark.parametrize("family", ("swiglu", "srelu", "quant")) +def test_native_fp4_float32_output_preserves_torch_recipe_guard(family: str): + source = _jax_path(family).read_text() + + assert "sf_vec_size == 16" in source + assert "d_dtype == jnp.dtype(jnp.float32)" in source + + +def test_native_fp4_glu_float32_output_requires_bias_for_sf_vec_16(): + source = _jax_path("glu").read_text() + + assert "sf_vec_size == 16" in source + assert "d_dtype == jnp.dtype(jnp.float32)" in source + assert "bias_tensor is None" in source + + +def test_glu_preserves_torch_sfd_generation_for_all_fp8_inputs(): + source = _jax_path("glu").read_text() + + assert "generate_sfd = is_fp8_dtype(ab_dtype)" in source + + +def test_glu_accepts_native_fp4_intermediate_output(): + tree = ast.parse(_jax_path("glu").read_text()) + c_dtype_call = next( + node.value + for node in ast.walk(tree) + if isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == "c_dtype" + for target in node.targets + ) + and isinstance(node.value, ast.Call) + and isinstance(node.value.func, ast.Name) + and node.value.func.id == "require_dtype" + ) + + assert "jnp.float4_e2m1fn" in ast.unparse(c_dtype_call.args[1]) + + +@pytest.mark.parametrize("family", ("dswiglu", "dglu", "dsrelu")) +def test_native_fp4_backward_paths_return_amax_without_sfd(family: str): + source = _jax_path(family).read_text() + + assert "generate_sfd = is_fp8_dtype(ab_dtype)" in source + assert '"has_amax": has_amax' in source + assert '"amax_tensor"' in source + + +@pytest.mark.parametrize("family", ("dswiglu", "dglu", "dsrelu")) +def test_backward_inputs_validate_canonical_descriptor_shapes(family: str): + tree = ast.parse(_jax_path(family).read_text()) + invalid_checks = [ + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "require_array" + and node.args + and isinstance(node.args[0], ast.Name) + and node.args[0].id == "c_desc" + ] + + assert not invalid_checks + assert "if c_desc.shape != expected_c_shape:" in _jax_path(family).read_text() + + +def test_grouped_packages_keep_torch_exports_lazy(): + root_source = (_GROUPED_ROOT / "__init__.py").read_text() + assert "make_operation_api" in root_source + assert ".api import" not in root_source + + for family in _FAMILIES: + source = (_GROUPED_ROOT / f"grouped_gemm_{family}" / "__init__.py").read_text() + assert "make_operation_api" in source + assert 'submodules=("api", "jax")' in source + assert ".api import" not in source + + +def test_grouped_package_import_does_not_load_torch_apis(): + package_name = "cudnn_frontend_grouped_lazy_contract" + package = types.ModuleType(package_name) + package.__path__ = [str(_ROOT / "python" / "cudnn")] + package.__package__ = package_name + + with mock.patch.dict(sys.modules, {package_name: package}): + grouped = importlib.import_module(f"{package_name}.grouped_gemm") + expected = { + f"GroupedGemm{_class_stem(family)}Sm100" for family in _FAMILIES + } | {f"grouped_gemm_{family}_wrapper_sm100" for family in _FAMILIES} + assert set(grouped.__all__) == expected + assert not any( + module_name.startswith(f"{package_name}.grouped_gemm.") + for module_name in sys.modules + ) + + for family in _FAMILIES: + subpackage_name = f"{package_name}.grouped_gemm.grouped_gemm_{family}" + subpackage = importlib.import_module(subpackage_name) + assert "jax" in dir(subpackage) + assert f"{subpackage_name}.api" not in sys.modules + + for module_name in tuple(sys.modules): + if module_name == package_name or module_name.startswith(f"{package_name}."): + sys.modules.pop(module_name, None) + + +def test_shared_layout_helpers_map_row_major_arrays_to_kernel_axes(): + class DataType(Enum): + NOT_SET = auto() + + class TensorSpec: + def __init__(self, *, layout=None, mode=None, **kwargs): + self.layout = layout + self.mode = mode + self.__dict__.update(kwargs) + + class ShapeDtypeStruct: + def __init__(self, shape, dtype, fill_value=None): + self.shape = tuple(shape) + self.dtype = dtype + self.fill_value = fill_value + + fake_jax = types.ModuleType("jax") + fake_jax.__path__ = [] + fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + fake_jax.ShapeDtypeStruct = ShapeDtypeStruct + fake_jax.local_devices = lambda **_kwargs: () + fake_jax.tree_util = types.SimpleNamespace( + DictKey=lambda key: key, + register_pytree_with_keys=lambda *_args: None, + ) + + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.dtype = lambda value: value + for dtype_name in ( + "float4_e2m1fn", + "float8_e4m3fn", + "float8_e5m2", + "float8_e8m0fnu", + ): + setattr(fake_jnp, dtype_name, dtype_name) + fake_jnp.empty = lambda shape, dtype: ShapeDtypeStruct(shape, dtype) + fake_jnp.full = lambda shape, fill_value, dtype: ShapeDtypeStruct( + shape, dtype, fill_value + ) + fake_jax.numpy = fake_jnp + + fake_cutlass = types.ModuleType("cutlass") + fake_cutlass.__path__ = [] + fake_cutlass.__spec__ = ModuleSpec("cutlass", loader=None, is_package=True) + fake_cutlass_jax = types.ModuleType("cutlass.jax") + fake_cutlass_jax.TensorSpec = TensorSpec + fake_cutlass.jax = fake_cutlass_jax + + package_name = "cudnn_frontend_grouped_layout_contract" + package = types.ModuleType(package_name) + package.__path__ = [str(_ROOT / "python" / "cudnn")] + package.__package__ = package_name + package.data_type = DataType + + grouped_package = types.ModuleType(f"{package_name}.grouped_gemm") + grouped_package.__path__ = [str(_GROUPED_ROOT)] + grouped_package.__package__ = grouped_package.__name__ + + modules = { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "cutlass": fake_cutlass, + "cutlass.jax": fake_cutlass_jax, + package_name: package, + grouped_package.__name__: grouped_package, + } + with mock.patch.dict(sys.modules, modules): + adapter = importlib.import_module(f"{package_name}.grouped_gemm._jax_api") + + assert adapter.gemm_a_mode("LMK") == (1, 2, 0) + assert adapter.gemm_b_mode("LKN") == (2, 1, 0) + assert adapter.PROBABILITY_MODE == (2, 1, 0) + assert adapter.GROUPED_BIAS_MODE == (1, 0) + assert adapter.block_scale_shape(256, 64, 4, 32) == (4, 2, 1, 32, 4, 4) + + sample = ShapeDtypeStruct((1, 256, 64), "float8") + desc = adapter.ApiBaseJax().make_tensor_desc( + sample, + mode=adapter.gemm_a_mode("LMK"), + name="sample_a", + ) + assert type(desc) is adapter.JaxTensorDesc + assert desc.shape == (256, 64, 1) + assert desc.stride == (64, 1, 16384) + assert desc.mode == (1, 2, 0) + assert desc.array_shape == sample.shape + + fp4_sfa = ShapeDtypeStruct( + adapter.block_scale_shape(256, 256, 1, 16), + "float8_e4m3fn", + ) + fp4_sfb = ShapeDtypeStruct( + adapter.block_scale_shape(128, 256, 2, 16), + "float8_e4m3fn", + ) + assert ( + adapter.require_grouped_input_scales( + fp4_sfa, + fp4_sfb, + m=256, + n=128, + k=256, + experts=2, + sf_vec_size=16, + ab_dtype="float4_e2m1fn", + ) + == "float8_e4m3fn" + ) + + invalid_fp4_sfa = ShapeDtypeStruct( + adapter.block_scale_shape(256, 256, 1, 32), + "float8_e4m3fn", + ) + invalid_fp4_sfb = ShapeDtypeStruct( + adapter.block_scale_shape(128, 256, 2, 32), + "float8_e4m3fn", + ) + with pytest.raises(ValueError, match="sf_vec_size=16"): + adapter.require_grouped_input_scales( + invalid_fp4_sfa, + invalid_fp4_sfb, + m=256, + n=128, + k=256, + experts=2, + sf_vec_size=32, + ab_dtype="float4_e2m1fn", + ) + + mxfp4_sfa = ShapeDtypeStruct( + adapter.block_scale_shape(256, 256, 1, 32), + "float8_e8m0fnu", + ) + mxfp4_sfb = ShapeDtypeStruct( + adapter.block_scale_shape(128, 256, 2, 32), + "float8_e8m0fnu", + ) + assert ( + adapter.require_grouped_input_scales( + mxfp4_sfa, + mxfp4_sfb, + m=256, + n=128, + k=256, + experts=2, + sf_vec_size=32, + ab_dtype="float4_e2m1fn", + ) + == "float8_e8m0fnu" + ) + with pytest.raises(ValueError, match="Unsupported grouped GEMM input dtype"): + adapter.require_grouped_input_scales( + mxfp4_sfa, + mxfp4_sfb, + m=256, + n=128, + k=256, + experts=2, + sf_vec_size=32, + ab_dtype="uint8", + ) + + fp8_sfa = ShapeDtypeStruct( + adapter.block_scale_shape(256, 256, 1, 32), + "float8_e8m0fnu", + ) + fp8_sfb = ShapeDtypeStruct( + adapter.block_scale_shape(128, 256, 2, 32), + "float8_e8m0fnu", + ) + assert ( + adapter.require_grouped_input_scales( + fp8_sfa, + fp8_sfb, + m=256, + n=128, + k=256, + experts=2, + sf_vec_size=32, + ab_dtype="float8_e4m3fn", + ) + == "float8_e8m0fnu" + ) + + launched = [] + lowered = {} + + def launch(stream, input_value, output_value, workspace_value, **static): + launched.append( + (stream, input_value, output_value, workspace_value, static) + ) + + def fake_call_kernel(inputs, **options): + lowered.update(inputs=inputs, **options) + options["launch"]("stream", *inputs, "output", "workspace") + return ("result",) + + input_desc = adapter.as_gemm_tensor_desc( + "input", + sample, + mode=adapter.gemm_a_mode("LMK"), + ) + output_mode = adapter.gemm_output_mode("LMN", name="output_layout") + with ( + mock.patch.object( + adapter._CALLER, "_get_max_active_clusters", return_value=7 + ) as occupancy, + mock.patch.object( + adapter._CALLER, "_resolve_compute_capability", return_value=100 + ), + mock.patch.object( + adapter._CALLER, "_call_kernel", side_effect=fake_call_kernel + ), + ): + result = adapter.call_cutedsl( + launch, + (sample,), + input_descs=(input_desc,), + outputs=( + adapter.make_buffer_desc( + "output", + (1, 256, 128), + "bfloat16", + mode=output_mode, + ), + ), + workspaces=( + adapter.make_buffer_desc( + "workspace", + (4,), + "uint8", + ptr_assumed_align=adapter.GROUPED_WORKSPACE_ALIGNMENT, + ), + ), + static_args={ + "cluster_shape_mn": (2, 1), + "cluster_overlap_margin": 1, + "expert_cnt": 4, + }, + ) + + assert result == ("result",) + occupancy.assert_called_once_with(2, overlap_margin=1) + assert launched == [ + ( + "stream", + sample, + "output", + "workspace", + { + "cluster_shape_mn": (2, 1), + "expert_cnt": 4, + "max_active_clusters": 7, + }, + ) + ] + assert lowered["output_descs"][0].shape == (256, 128, 1) + assert type(lowered["output_descs"][0]) is adapter.JaxTensorDesc + assert lowered["output_descs"][0].mode == output_mode + assert lowered["input_descs"] == (input_desc,) + assert lowered["workspace_descs"][0].shape == (4,) + assert ( + lowered["workspace_descs"][0].ptr_assumed_align + == adapter.GROUPED_WORKSPACE_ALIGNMENT + ) + + with pytest.raises(ValueError, match="Expected 1 input descriptors, got 0"): + adapter.call_cutedsl( + launch, + (sample,), + input_descs=(), + outputs=(adapter.make_buffer_desc("output", (1, 0, 128), "bfloat16"),), + ) + + def materialize(desc): + return ShapeDtypeStruct( + desc.array_shape, + desc.dtype, + ) + + with ( + mock.patch.object( + adapter._CALLER, + "_get_max_active_clusters", + side_effect=AssertionError( + "empty grouped GEMM must not query occupancy" + ), + ), + mock.patch.object( + adapter._CALLER, + "_call_kernel", + side_effect=AssertionError("empty grouped GEMM must not lower"), + ), + mock.patch.object( + adapter._CALLER, + "_to_shape_dtype_struct", + side_effect=materialize, + ), + ): + empty, initialized = adapter.call_cutedsl( + launch, + (sample,), + input_descs=(input_desc,), + outputs=( + adapter.make_buffer_desc( + "empty", + (1, 0, 128), + "bfloat16", + mode=output_mode, + ), + adapter.make_buffer_desc( + "amax", + (4, 1), + "float32", + init_value=float("-inf"), + ), + ), + static_args={"cluster_shape_mn": (2, 1)}, + ) + + assert empty.shape == (1, 0, 128) + assert initialized.shape == (4, 1) + assert initialized.fill_value == float("-inf") + + for module_name in tuple(sys.modules): + if module_name == package_name or module_name.startswith(f"{package_name}."): + sys.modules.pop(module_name, None) diff --git a/test/python/fe_api/test_jax_grouped_wgrad_accumulation_contract.py b/test/python/fe_api/test_jax_grouped_wgrad_accumulation_contract.py new file mode 100644 index 000000000..ed88fc600 --- /dev/null +++ b/test/python/fe_api/test_jax_grouped_wgrad_accumulation_contract.py @@ -0,0 +1,138 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for JAX-owned grouped-wgrad outputs.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.L0 + +_ROOT = Path(__file__).resolve().parents[3] +_WGRAD_PATH = ( + _ROOT / "python" / "cudnn" / "grouped_gemm" / "grouped_gemm_wgrad" / "jax.py" +) +_SHARED_PATH = _ROOT / "python" / "cudnn" / "grouped_gemm" / "_jax_api.py" + + +def _definitions() -> dict[str, ast.AST]: + tree = ast.parse(_WGRAD_PATH.read_text()) + return { + node.name: node + for node in tree.body + if isinstance(node, (ast.ClassDef, ast.FunctionDef)) + } + + +def _method(node: ast.ClassDef, name: str) -> ast.FunctionDef: + return next( + child + for child in node.body + if isinstance(child, ast.FunctionDef) and child.name == name + ) + + +def _argument_names(function: ast.FunctionDef) -> tuple[str, ...]: + return tuple( + argument.arg for argument in (*function.args.args, *function.args.kwonlyargs) + ) + + +def test_wgrad_output_is_not_part_of_the_class_or_wrapper_signatures(): + definitions = _definitions() + operation = definitions["GroupedGemmWgradSm100"] + assert isinstance(operation, ast.ClassDef) + wrapper = definitions["grouped_gemm_wgrad_wrapper_sm100"] + assert isinstance(wrapper, ast.FunctionDef) + implementation = definitions["_grouped_gemm_wgrad_impl"] + assert isinstance(implementation, ast.FunctionDef) + + assert "sample_wgrad_tensor" not in _argument_names( + _method(operation, "__init__") + ) + assert "wgrad_tensor" not in _argument_names(_method(operation, "__call__")) + assert "wgrad_tensor" not in _argument_names(_method(operation, "_call_impl")) + assert "wgrad_tensor" not in _argument_names(implementation) + assert "wgrad_tensor" not in _argument_names(wrapper) + + +def test_accumulating_wgrad_is_zero_initialized_by_its_descriptor(): + implementation = _definitions()["_grouped_gemm_wgrad_impl"] + assert isinstance(implementation, ast.FunctionDef) + calls = [ + node + for node in ast.walk(implementation) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "call_cutedsl" + ] + assert len(calls) == 1 + + keywords = {keyword.arg: keyword.value for keyword in calls[0].keywords} + assert "output_seeds" not in keywords + + output_descs = [ + node + for node in ast.walk(implementation) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "make_buffer_desc" + and node.args + and isinstance(node.args[0], ast.Constant) + and node.args[0].value == "wgrad_tensor" + ] + assert len(output_descs) == 1 + output_keywords = { + keyword.arg: keyword.value for keyword in output_descs[0].keywords + } + assert ast.unparse(output_keywords["init_value"]) == ( + "0.0 if accumulate_on_output else None" + ) + + source = ast.get_source_segment(_WGRAD_PATH.read_text(), implementation) + assert source is not None + assert "default=jnp.bfloat16" in source + + +def test_grouped_lowering_does_not_accept_preallocated_outputs(): + tree = ast.parse(_SHARED_PATH.read_text()) + call_cutedsl = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "call_cutedsl" + ) + + assert "output_seeds" not in _argument_names(call_cutedsl) + assert "output_seeds" not in ast.unparse(call_cutedsl) + + +def test_wgrad_layout_and_alignment_constraints_live_on_descriptors(): + source = _WGRAD_PATH.read_text() + implementation = _definitions()["_grouped_gemm_wgrad_impl"] + assert isinstance(implementation, ast.FunctionDef) + call = next( + node + for node in ast.walk(implementation) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "call_cutedsl" + ) + keywords = {keyword.arg for keyword in call.keywords} + + assert "WGRAD_B_STRIDE_ORDER = (0, 1)" in source + assert "WGRAD_ALIGNMENT = 16" in source + assert "ptr_assumed_align=GROUPED_WORKSPACE_ALIGNMENT" in source + assert "input_descs" in keywords + assert not {"input_specs", "output_specs", "workspace_specs"} & keywords + + +def test_jax_owned_output_and_zero_initialization_are_documented(): + source = _WGRAD_PATH.read_text() + + assert "The output is inferred and allocated by JAX" in source + assert "zero-initialized output" in source + assert "Pointer-table outputs are not supported" in source diff --git a/test/python/fe_api/test_jax_import_contract.py b/test/python/fe_api/test_jax_import_contract.py new file mode 100644 index 000000000..f001f206d --- /dev/null +++ b/test/python/fe_api/test_jax_import_contract.py @@ -0,0 +1,324 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Automatically discovered import-boundary contracts for JAX APIs.""" + +from __future__ import annotations + +import ast +from collections import deque +import importlib.util +from pathlib import Path +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_CUDNN_ROOT = _REPO_ROOT / "python" / "cudnn" + + +def _module_name(path): + relative = path.relative_to(_CUDNN_ROOT) + parts = ( + relative.parts[:-1] + if path.name == "__init__.py" + else relative.with_suffix("").parts + ) + return ".".join(("cudnn", *parts)) + + +def _static_condition(node): + if isinstance(node, ast.Constant) and isinstance(node.value, bool): + return node.value + if isinstance(node, ast.Name) and node.id == "TYPE_CHECKING": + return False + if isinstance(node, ast.Attribute) and node.attr == "TYPE_CHECKING": + return False + if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.Not): + value = _static_condition(node.operand) + return None if value is None else not value + if ( + isinstance(node, ast.Compare) + and isinstance(node.left, ast.Name) + and node.left.id == "__name__" + and len(node.ops) == len(node.comparators) == 1 + and isinstance(node.comparators[0], ast.Constant) + and node.comparators[0].value == "__main__" + ): + if isinstance(node.ops[0], ast.Eq): + return False + if isinstance(node.ops[0], ast.NotEq): + return True + return None + + +def _catches_import_error(node): + if isinstance(node, ast.Name): + return node.id in {"ImportError", "ModuleNotFoundError"} + if isinstance(node, ast.Attribute): + return node.attr in {"ImportError", "ModuleNotFoundError"} + if isinstance(node, ast.Tuple): + return any(_catches_import_error(item) for item in node.elts) + return False + + +def _imports(path, *, all_scopes): + """Collect imports and whether a direct Torch import is optional.""" + + imports = [] + + def visit(node, torch_optional=False): + if isinstance(node, (ast.Import, ast.ImportFrom)): + imports.append((node, torch_optional)) + return + if isinstance(node, ast.If): + condition = _static_condition(node.test) + if condition is True: + branches = node.body + elif condition is False: + branches = node.orelse + else: + branches = (*node.body, *node.orelse) + for child in branches: + visit(child, torch_optional) + return + if isinstance(node, ast.Try): + catches_import = any( + handler.type is not None and _catches_import_error(handler.type) + for handler in node.handlers + ) + for child in node.body: + visit(child, torch_optional or catches_import) + for handler in node.handlers: + for child in handler.body: + visit(child, torch_optional) + for child in (*node.orelse, *node.finalbody): + visit(child, torch_optional) + return + if not all_scopes and isinstance( + node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda) + ): + return + for child in ast.iter_child_nodes(node): + visit(child, torch_optional) + + tree = ast.parse(path.read_text(), filename=str(path)) + for node in tree.body: + visit(node) + return imports + + +def _import_targets(path, node, modules): + if isinstance(node, ast.Import): + return [alias.name for alias in node.names] + + if node.level: + current_module = _module_name(path) + package = ( + current_module + if path.name == "__init__.py" + else current_module.rpartition(".")[0] + ) + base = importlib.util.resolve_name( + "." * node.level + (node.module or ""), package + ) + else: + base = node.module or "" + + targets = [base] if base else [] + for alias in node.names: + candidate = f"{base}.{alias.name}" if base else alias.name + if alias.name != "*" and candidate in modules: + targets.append(candidate) + return targets + + +def _package_initializers(path): + initializers = [_CUDNN_ROOT / "__init__.py"] + current = _CUDNN_ROOT + for part in path.parent.relative_to(_CUDNN_ROOT).parts: + current /= part + initializer = current / "__init__.py" + if initializer != path and initializer.is_file(): + initializers.append(initializer) + return initializers + + +def _jax_import_violations(root, modules): + # Scan every scope in JAX roots because kernel imports are deferred until + # tracing. In reached modules, scan only import-time code so shared kernel + # files may retain function-local Torch wrappers. + queue = deque([(root, True, (root,))]) + queue.extend((path, False, (root, path)) for path in _package_initializers(root)) + scanned_all_scopes = set() + scanned_module_scope = set() + visited = set() + violations = [] + + while queue: + path, all_scopes, chain = queue.popleft() + if path in scanned_all_scopes or ( + not all_scopes and path in scanned_module_scope + ): + continue + (scanned_all_scopes if all_scopes else scanned_module_scope).add(path) + visited.add(path) + + for node, torch_optional in _imports(path, all_scopes=all_scopes): + for target in _import_targets(path, node, modules): + chain_text = " -> ".join( + str(item.relative_to(_REPO_ROOT)) for item in chain + ) + if target == "torch" or target.startswith("torch."): + if not torch_optional: + violations.append( + f"{chain_text}:{node.lineno} imports {target}" + ) + continue + + dependency = modules.get(target) + if dependency is None: + continue + if dependency.name == "api.py": + violations.append( + f"{chain_text}:{node.lineno} imports Torch API module {dependency.relative_to(_REPO_ROOT)}" + ) + continue + queue.extend( + (initializer, False, (*chain, initializer)) + for initializer in _package_initializers(dependency) + ) + queue.append((dependency, False, (*chain, dependency))) + + return visited, violations + + +class JaxImportContractTest(unittest.TestCase): + @staticmethod + def _literal_assignment(path, name): + tree = ast.parse(path.read_text(), filename=str(path)) + assignment = next( + node + for node in tree.body + if isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == name + for target in node.targets + ) + ) + return ast.literal_eval(assignment.value) + + @staticmethod + def _operation_exports(path): + tree = ast.parse(path.read_text(), filename=str(path)) + call = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "make_operation_api" + ) + exports = ast.literal_eval( + next(keyword.value for keyword in call.keywords if keyword.arg == "exports") + ) + return {name for module_exports in exports.values() for name in module_exports} + + def test_jax_dsa_facade_matches_the_torch_operation_surface(self): + torch_exports = self._literal_assignment( + _CUDNN_ROOT / "deepseek_sparse_attention" / "__init__.py", + "_SYMBOLS", + ) + jax_exports = self._literal_assignment( + _CUDNN_ROOT / "jax" / "__init__.py", "_OPERATION_EXPORTS" + ) + jax_dsa_exports = { + name + for name, (module_name, _) in jax_exports.items() + if module_name.startswith("..deepseek_sparse_attention.") + } + + self.assertEqual(jax_dsa_exports, set(torch_exports)) + + def test_jax_facade_covers_the_remaining_cute_operation_packages(self): + jax_exports = self._literal_assignment( + _CUDNN_ROOT / "jax" / "__init__.py", "_OPERATION_EXPORTS" + ) + grouped_families = ( + "swiglu", + "dswiglu", + "quant", + "srelu", + "dsrelu", + "glu", + "glu_hadamard", + "dglu", + "wgrad", + ) + grouped_names = { + name + for family in grouped_families + for name in ( + f"GroupedGemm{''.join(part.capitalize() for part in family.split('_'))}Sm100", + f"grouped_gemm_{family}_wrapper_sm100", + ) + } + discrete_names = { + "DiscreteGroupedGemmSwigluSm100", + "discrete_grouped_gemm_swiglu_wrapper_sm100", + "DiscreteGroupedGemmDswigluSm100", + "discrete_grouped_gemm_dswiglu_wrapper_sm100", + } + nsa_names = self._operation_exports( + _CUDNN_ROOT / "native_sparse_attention" / "__init__.py" + ) + sdpa_names = self._operation_exports(_CUDNN_ROOT / "sdpa" / "__init__.py") + bsa_names = { + name + for name in self._literal_assignment( + _CUDNN_ROOT / "block_sparse_attention" / "__init__.py", + "_SYMBOLS", + ) + if not name.endswith("Op") + } + + self.assertLessEqual( + grouped_names | discrete_names | nsa_names | sdpa_names | bsa_names, + set(jax_exports), + ) + + def test_jax_import_graph_has_no_required_torch_dependency(self): + modules = {_module_name(path): path for path in _CUDNN_ROOT.rglob("*.py")} + roots = set(_CUDNN_ROOT.rglob("jax.py")) + roots.update((_CUDNN_ROOT / "jax").rglob("*.py")) + self.assertTrue(roots, "No JAX API modules were discovered") + + visited = set() + violations = [] + for root in sorted(roots): + root_visited, root_violations = _jax_import_violations(root, modules) + visited.update(root_visited) + violations.extend(root_violations) + + self.assertTrue( + any(path.name == "__init__.py" and path not in roots for path in visited), + "JAX import graph did not discover package initializers", + ) + self.assertTrue( + any(path.name != "__init__.py" and path not in roots for path in visited), + "JAX import graph did not discover deferred kernel dependencies", + ) + self.assertEqual( + violations, + [], + "JAX import paths must not require Torch:\n" + "\n".join(violations), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_layout.py b/test/python/fe_api/test_jax_layout.py new file mode 100644 index 000000000..068af5cf9 --- /dev/null +++ b/test/python/fe_api/test_jax_layout.py @@ -0,0 +1,157 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Tests for CUTLASS JAX layout conversion.""" + +import importlib.util +from pathlib import Path +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_LAYOUT_FILE = ( + Path(__file__).resolve().parents[3] / "python" / "cudnn" / "_jax" / "layout.py" +) + + +def _load_layout_module(): + spec = importlib.util.spec_from_file_location("cudnn_jax_layout_test", _LAYOUT_FILE) + if spec is None or spec.loader is None: + raise RuntimeError("Unable to load CUTLASS JAX layout helpers") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class JaxLayoutTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.layout = _load_layout_module() + + def test_non_self_inverse_stride_order_maps_to_cutlass_axis_ranks(self): + shape = (7, 3, 5) + stride_order = (1, 2, 0) + stride = self.layout.compact_stride(shape, stride_order) + + self.assertEqual(stride, (15, 1, 3)) + self.assertEqual( + self.layout.to_cutlass_layout(shape, stride, stride_order), + (2, 0, 1), + ) + + def test_named_layout_maps_kernel_axes_to_public_axes(self): + cases = ( + ("LMK", "MKL", (1, 2, 0)), + ("LKM", "MKL", (2, 1, 0)), + ("LNK", "NKL", (1, 2, 0)), + ("LKN", "NKL", (2, 1, 0)), + ("LMN", "MNL", (1, 2, 0)), + ("LNM", "MNL", (2, 1, 0)), + ) + for public_layout, kernel_axes, expected in cases: + with self.subTest(public_layout=public_layout, kernel_axes=kernel_axes): + self.assertEqual( + self.layout.mode_from_layout( + public_layout, + kernel_axes=kernel_axes, + ), + expected, + ) + + def test_named_layout_requires_strings(self): + with self.assertRaisesRegex(TypeError, "layout must be a string"): + self.layout.mode_from_layout(("L", "M", "K"), kernel_axes="MKL") + with self.assertRaisesRegex(TypeError, "kernel_axes must be a string"): + self.layout.mode_from_layout("LMK", kernel_axes=("M", "K", "L")) + + def test_named_layout_rejects_rank_mismatch(self): + with self.assertRaisesRegex(ValueError, "layout rank must match"): + self.layout.mode_from_layout("MK", kernel_axes="MKL") + + def test_named_layout_rejects_duplicate_axes(self): + with self.assertRaisesRegex(ValueError, "layout axes must be unique"): + self.layout.mode_from_layout("MML", kernel_axes="MKL") + with self.assertRaisesRegex(ValueError, "kernel_axes must be unique"): + self.layout.mode_from_layout("MKL", kernel_axes="MML") + + def test_named_layout_requires_the_exact_axis_set(self): + with self.assertRaisesRegex(ValueError, "layout must contain exactly"): + self.layout.mode_from_layout("MKN", kernel_axes="MKL") + + def test_noncompact_stride_is_rejected(self): + with self.assertRaisesRegex(ValueError, "cannot represent non-compact stride"): + self.layout.to_cutlass_layout( + shape=(7, 3, 5), + stride=(16, 1, 3), + stride_order=(1, 2, 0), + name="sample", + ) + + def test_compact_stride_preserves_order_for_zero_extents(self): + self.assertEqual( + self.layout.compact_stride((2, 0, 3), (2, 1, 0)), + (3, 3, 1), + ) + + def test_identity_mode_is_the_default(self): + self.assertEqual(self.layout.normalize_mode(3), (0, 1, 2)) + self.assertEqual(self.layout.to_canonical_axes((2, 3, 4)), (2, 3, 4)) + self.assertEqual(self.layout.to_public_axes((2, 3, 4)), (2, 3, 4)) + + def test_mode_maps_kernel_axes_to_public_axes(self): + mode = (2, 0, 1) + + self.assertEqual(self.layout.to_canonical_axes((2, 3, 4), mode), (4, 2, 3)) + self.assertEqual(self.layout.to_public_axes((4, 2, 3), mode), (2, 3, 4)) + + def test_cutlass_layout_is_indexed_by_public_axes(self): + self.assertEqual( + self.layout.to_cutlass_layout( + shape=(4, 2, 3), + stride=(1, 12, 4), + stride_order=(0, 2, 1), + mode=(2, 0, 1), + ), + (2, 1, 0), + ) + + def test_canonical_stride_order_maps_to_named_public_layout(self): + bhsd_to_bshd = self.layout.mode_from_layout("BHSD", kernel_axes="BSHD") + bshd_to_bshd = self.layout.mode_from_layout("BSHD", kernel_axes="BSHD") + + self.assertEqual( + self.layout.stride_order_to_public((3, 2, 1, 0), bhsd_to_bshd), + (3, 1, 2, 0), + ) + self.assertEqual( + self.layout.stride_order_to_public((3, 2, 1, 0), bshd_to_bshd), + (3, 2, 1, 0), + ) + + bshd_to_bhsd = self.layout.mode_from_layout("BSHD", kernel_axes="BHSD") + self.assertEqual( + self.layout.to_public_axes((2, 8, 32, 16), bshd_to_bhsd), + (2, 32, 8, 16), + ) + self.assertEqual( + self.layout.stride_order_to_public((3, 1, 2, 0), bshd_to_bhsd), + (3, 2, 1, 0), + ) + + def test_canonical_stride_order_requires_a_permutation(self): + with self.assertRaisesRegex(ValueError, "must be a permutation"): + self.layout.stride_order_to_public((2, 2, 0)) + + def test_invalid_mode_is_rejected(self): + with self.assertRaisesRegex(ValueError, "mode must be a permutation"): + self.layout.normalize_mode(3, (0, 0, 2)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_nsa_sdpa_source_contract.py b/test/python/fe_api/test_jax_nsa_sdpa_source_contract.py new file mode 100644 index 000000000..eb2eb495f --- /dev/null +++ b/test/python/fe_api/test_jax_nsa_sdpa_source_contract.py @@ -0,0 +1,355 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Source contracts shared by the JAX NSA and d=256 SDPA adapters.""" + +from __future__ import annotations + +import ast +from pathlib import Path +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_ADAPTERS = { + "selection": _CUDNN_ROOT / "native_sparse_attention" / "selection" / "jax.py", + "compression": _CUDNN_ROOT / "native_sparse_attention" / "compression" / "jax.py", + "sliding": _CUDNN_ROOT + / "native_sparse_attention" + / "sliding_window_attention" + / "jax.py", + "top_k": _CUDNN_ROOT / "native_sparse_attention" / "top_k" / "jax.py", + "sdpa_fwd": _CUDNN_ROOT / "sdpa" / "fwd" / "jax.py", + "sdpa_bwd": _CUDNN_ROOT / "sdpa" / "bwd" / "jax.py", +} + + +def _tree(path: Path) -> ast.Module: + return ast.parse(path.read_text(), filename=str(path)) + + +def _class(tree: ast.Module, name: str) -> ast.ClassDef: + return next( + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == name + ) + + +def _method(class_node: ast.ClassDef, name: str) -> ast.FunctionDef: + return next( + node + for node in class_node.body + if isinstance(node, ast.FunctionDef) and node.name == name + ) + + +class JaxNsaSdpaSourceContractTest(unittest.TestCase): + def test_layout_helpers_separate_logical_modes_from_physical_storage(self): + nsa_helpers = ( + _CUDNN_ROOT / "native_sparse_attention" / "jax_utils.py" + ).read_text() + sdpa_helpers = (_CUDNN_ROOT / "sdpa" / "jax_utils.py").read_text() + + self.assertNotIn("BHSD_STORAGE_ORDER", nsa_helpers) + self.assertNotIn("BHSD_STORAGE_ORDER", sdpa_helpers) + self.assertIn("mode_from_layout(layout, kernel_axes=kernel_axes)", nsa_helpers) + self.assertIn("kernel_stride_order is None", nsa_helpers) + self.assertIn('KERNEL_AXES = "BSHD"', sdpa_helpers) + self.assertIn("KERNEL_STRIDE_ORDER = (3, 2, 1, 0)", sdpa_helpers) + + def test_all_adapters_are_torch_independent_jax_api_classes(self): + expected_classes = { + "selection": "SelectionAttention", + "compression": "CompressionAttention", + "sliding": "SlidingWindowAttention", + "top_k": "TopKReduction", + "sdpa_fwd": "SdpafwdSm100D256", + "sdpa_bwd": "SdpabwdSm100D256", + } + for operation, path in _ADAPTERS.items(): + with self.subTest(operation=operation): + tree = _tree(path) + imports = [ + node + for node in ast.walk(tree) + if isinstance(node, (ast.Import, ast.ImportFrom)) + ] + self.assertFalse( + any( + ( + isinstance(node, ast.Import) + and any(alias.name == "torch" for alias in node.names) + ) + or (isinstance(node, ast.ImportFrom) and node.module == "torch") + for node in imports + ) + ) + adapter = _class(tree, expected_classes[operation]) + self.assertEqual( + [base.id for base in adapter.bases if isinstance(base, ast.Name)], + ["JaxApiBase"], + ) + self.assertTrue( + any( + isinstance(node, ast.FunctionDef) + and node.name == "check_support" + for node in adapter.body + ) + ) + self.assertTrue( + any( + isinstance(node, ast.FunctionDef) and node.name == "__call__" + for node in adapter.body + ) + ) + + def test_cute_launchers_follow_stream_inputs_outputs_workspaces_order(self): + expected = { + "selection": ( + "stream", + "q", + "k", + "v", + "block_indices", + "block_counts", + "cum_seqlen", + "output", + "lse_sum", + "row_max", + ), + "compression": ("stream", "q", "k", "v", "output"), + "top_k": ("stream", "q", "k", "lse", "topk_scores", "topk_indices"), + "sdpa_fwd": ("stream", "q", "k", "v", "output", "lse"), + "sdpa_bwd": ( + "stream", + "q", + "k", + "v", + "output", + "doutput", + "lse", + "dq", + "dk", + "dv", + "workspace", + ), + } + class_names = { + "selection": "SelectionAttention", + "compression": "CompressionAttention", + "top_k": "TopKReduction", + "sdpa_fwd": "SdpafwdSm100D256", + "sdpa_bwd": "SdpabwdSm100D256", + } + for operation, argument_names in expected.items(): + with self.subTest(operation=operation): + launcher = _method( + _class(_tree(_ADAPTERS[operation]), class_names[operation]), + "_launch_kernel", + ) + actual = tuple(argument.arg for argument in launcher.args.args[1:]) + self.assertEqual(actual, argument_names) + + def test_functional_buffers_preserve_required_initial_values(self): + expected_literals = { + "selection": {"0", "0.0", "float('-inf')"}, + "top_k": {"float('-inf')", "-1"}, + "sdpa_bwd": {"0"}, + } + for operation, expected in expected_literals.items(): + with self.subTest(operation=operation): + values = set() + for call in ( + node + for node in ast.walk(_tree(_ADAPTERS[operation])) + if isinstance(node, ast.Call) + ): + for keyword in call.keywords: + if keyword.arg == "init_value": + values.add(ast.unparse(keyword.value)) + self.assertTrue( + expected.issubset(values), + f"{operation} initialized buffers changed: {values}", + ) + + def test_selection_uses_one_self_attention_offset_operand(self): + tree = _tree(_ADAPTERS["selection"]) + wrapper = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) + and node.name == "selection_attention_wrapper" + ) + argument_names = tuple(argument.arg for argument in wrapper.args.args) + + self.assertIn("cum_seqlen_tensor", argument_names) + self.assertNotIn("cum_seqlen_q_tensor", argument_names) + self.assertNotIn("cum_seqlen_k_tensor", argument_names) + self.assertIn("(1, 2, 4, 8, 16)", _ADAPTERS["selection"].read_text()) + + def test_persistent_compression_resolves_hardware_before_lowering(self): + tree = _tree(_ADAPTERS["compression"]) + adapter = _class(tree, "CompressionAttention") + check_support = _method(adapter, "check_support") + runner = _method(adapter, "_run_kernel") + + self.assertIn( + "self._get_device_multiprocessor_count()", + ast.unparse(check_support), + ) + self.assertIn("persistent_sm_count", ast.unparse(runner)) + self.assertNotIn("HardwareInfo", _ADAPTERS["compression"].read_text()) + + def test_compression_supports_fixed_bhsd_and_packed_thd(self): + tree = _tree(_ADAPTERS["compression"]) + wrapper = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) + and node.name == "compression_attention_wrapper" + ) + argument_names = { + argument.arg for argument in (*wrapper.args.args, *wrapper.args.kwonlyargs) + } + varlen_launcher = _method( + _class(tree, "CompressionAttention"), "_launch_varlen_kernel" + ) + + self.assertTrue( + { + "cum_seqlen_q_tensor", + "cum_seqlen_k_tensor", + "max_s_q", + "max_s_k", + }.issubset(argument_names) + ) + self.assertEqual( + tuple(argument.arg for argument in varlen_launcher.args.args[1:]), + ( + "stream", + "q", + "k", + "v", + "cum_seqlen_q", + "cum_seqlen_k", + "output", + ), + ) + source = _ADAPTERS["compression"].read_text() + self.assertIn("normalize_attention_layout(layout, ranks[0])", source) + self.assertIn("self.input_layout in FIXED_LAYOUTS", source) + self.assertIn("jnp.float8_e4m3fn", source) + self.assertIn("self.lse_extent = total_q", source) + + def test_fixed_adapters_expose_named_bhsd_and_bshd_layouts(self): + for operation in ("compression", "sliding", "top_k", "sdpa_fwd", "sdpa_bwd"): + source = _ADAPTERS[operation].read_text() + with self.subTest(operation=operation): + self.assertIn("layout: str | None", source) + self.assertIn('"layout",', source) + + def test_compression_packed_views_duplicate_the_token_stride(self): + tree = _tree(_ADAPTERS["compression"]) + packed_orders = [ + ast.literal_eval(keyword.value) + for call in ast.walk(tree) + if isinstance(call, ast.Call) + for keyword in call.keywords + if keyword.arg == "public_stride_order" + and ast.unparse(keyword.value) == "(3, 2, 0, 1)" + ] + + self.assertEqual(len(packed_orders), 4) + token_count, heads, head_dim = 20, 8, 128 + shape = (1, token_count, heads, head_dim) + running = 1 + strides = [0] * len(shape) + for axis in packed_orders[0]: + strides[axis] = running + running *= max(shape[axis], 1) + self.assertEqual( + tuple(strides), + (heads * head_dim, heads * head_dim, head_dim, 1), + ) + + def test_sliding_window_supports_stats_two_sided_windows_and_packed_thd(self): + source = _ADAPTERS["sliding"].read_text() + + self.assertIn("return_residual=not self.is_infer", source) + self.assertIn("(self.left_bound - 1, self.right_bound)", source) + self.assertIn("query_seq_lengths=query_lengths", source) + self.assertIn("key_value_seq_lengths=key_value_lengths", source) + self.assertIn("jnp.diff(cum_seqlen_q_tensor)", source) + self.assertIn("self._pad_packed(", source) + self.assertIn("self._unpad_packed(", source) + self.assertNotIn("currently supports inference only", source) + self.assertNotIn("currently requires S_q == S_k", source) + + def test_cutlass_kernels_are_deferred_until_adapter_calls(self): + kernel_modules = { + "selection": "NSA_select_attn_fwd_hmma", + "compression": "fmha", + "top_k": "nsa_top_k_reduction_fwd", + "sdpa_fwd": "fmha_forward_sm100_d256", + "sdpa_bwd": "fmha_backward_sm100_2kernel", + } + for operation, module_name in kernel_modules.items(): + with self.subTest(operation=operation): + tree = _tree(_ADAPTERS[operation]) + top_level_imports = [ + node + for node in tree.body + if isinstance(node, (ast.Import, ast.ImportFrom)) + ] + self.assertFalse( + any( + isinstance(node, ast.ImportFrom) + and (node.module or "").endswith(module_name) + for node in top_level_imports + ) + ) + self.assertTrue( + any( + isinstance(node, ast.ImportFrom) + and (node.module or "").endswith(module_name) + for node in ast.walk(tree) + ) + ) + + def test_package_initializers_keep_torch_apis_lazy(self): + initializers = [ + _CUDNN_ROOT / "native_sparse_attention" / "__init__.py", + *( + _CUDNN_ROOT / "native_sparse_attention" / name / "__init__.py" + for name in ( + "selection", + "compression", + "sliding_window_attention", + "top_k", + ) + ), + _CUDNN_ROOT / "sdpa" / "__init__.py", + _CUDNN_ROOT / "sdpa" / "fwd" / "__init__.py", + _CUDNN_ROOT / "sdpa" / "bwd" / "__init__.py", + ] + for path in initializers: + with self.subTest(path=path.relative_to(_CUDNN_ROOT)): + direct_api_imports = [ + node + for node in _tree(path).body + if isinstance(node, ast.ImportFrom) + and node.module in {"api", ".api"} + ] + self.assertEqual(direct_api_imports, []) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_nsa_topk_thd_contract.py b/test/python/fe_api/test_jax_nsa_topk_thd_contract.py new file mode 100644 index 000000000..db9e98ea4 --- /dev/null +++ b/test/python/fe_api/test_jax_nsa_topk_thd_contract.py @@ -0,0 +1,440 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for packed JAX NSA top-K reduction.""" + +from __future__ import annotations + +import ast +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_NSA_ROOT = _CUDNN_ROOT / "native_sparse_attention" +_TOPK_ROOT = _NSA_ROOT / "top_k" +_PACKAGE = "cudnn_jax_nsa_topk_thd_contract_test" + + +class _DataType(Enum): + NOT_SET = auto() + HALF = auto() + BFLOAT16 = auto() + FLOAT = auto() + INT32 = auto() + INT64 = auto() + + +_DTYPE_TO_CUDNN = { + "float16": _DataType.HALF, + "bfloat16": _DataType.BFLOAT16, + "float32": _DataType.FLOAT, + "int32": _DataType.INT32, + "int64": _DataType.INT64, +} +_CUDNN_TO_DTYPE = {value: key for key, value in _DTYPE_TO_CUDNN.items()} + + +class _Array: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.dtype = dtype + + def __getitem__(self, key): + if key == (None, Ellipsis): + return _Array((1, *self.shape), self.dtype) + raise TypeError(f"unsupported test index {key!r}") + + +class _TensorSpec: + def __init__(self, *, layout, mode, divisibility=None): + self.layout = layout + self.mode = mode + self.divisibility = divisibility + + +class JaxNsaTopkThdContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + nsa_name = f"{_PACKAGE}.native_sparse_attention" + nsa = types.ModuleType(nsa_name) + nsa.__path__ = [str(_NSA_ROOT)] + nsa.__package__ = nsa_name + nsa.__spec__ = ModuleSpec(nsa_name, loader=None, is_package=True) + sys.modules[nsa_name] = nsa + + topk_name = f"{nsa_name}.top_k" + topk = types.ModuleType(topk_name) + topk.__path__ = [str(_TOPK_ROOT)] + topk.__package__ = topk_name + topk.__spec__ = ModuleSpec(topk_name, loader=None, is_package=True) + sys.modules[topk_name] = topk + + internal_name = f"{_PACKAGE}._jax" + internal = types.ModuleType(internal_name) + internal.__path__ = [str(_CUDNN_ROOT / "_jax")] + internal.__package__ = internal_name + internal.__spec__ = ModuleSpec(internal_name, loader=None, is_package=True) + sys.modules[internal_name] = internal + + tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + layout_module = importlib.import_module(f"{internal_name}.layout") + result_module = importlib.import_module(f"{_PACKAGE}.common.result") + + class JaxTensorDesc(tensor_module.TensorDesc): + @classmethod + def from_shape( + cls, + shape, + dtype, + *, + name="", + mode=None, + public_stride_order=None, + init_value=None, + ): + return JaxApiBase._to_tensor_desc( + _Array(shape, dtype), + name, + mode=mode, + public_stride_order=public_stride_order, + init_value=init_value, + ) + + @property + def cudnn_dtype(self): + return _DTYPE_TO_CUDNN.get(self.dtype, _DataType.NOT_SET) + + class JaxApiBase: + @staticmethod + def _resolve_compute_capability(target, supported, operation_name): + del operation_name + resolved = 100 if target is None else target + if resolved not in supported: + raise ValueError(f"unsupported target {resolved}") + return resolved + + @staticmethod + def _to_tensor_desc( + value, + name, + *, + mode=None, + public_stride_order=None, + init_value=None, + ): + public_shape = tuple(value.shape) + mode = layout_module.normalize_mode(len(public_shape), mode) + if public_stride_order is None: + public_stride_order = tuple(reversed(range(len(public_shape)))) + public_stride = layout_module.compact_stride( + public_shape, tuple(public_stride_order) + ) + canonical_axis_by_public_axis = layout_module.to_public_axes( + tuple(range(len(public_shape))), mode + ) + desc = JaxTensorDesc( + dtype=value.dtype, + shape=layout_module.to_canonical_axes(public_shape, mode), + stride=layout_module.to_canonical_axes(public_stride, mode), + stride_order=tuple( + canonical_axis_by_public_axis[axis] + for axis in public_stride_order + ), + name=name, + init_value=init_value, + ) + object.__setattr__(desc, "mode", mode) + return desc + + @staticmethod + def _check_tensor_signature(value, expected, *, mode=None): + if mode is None: + mode = expected.mode + actual_shape = layout_module.to_canonical_axes(tuple(value.shape), mode) + if actual_shape != expected.shape: + raise ValueError(f"{expected.name} shape mismatch") + if ( + _DTYPE_TO_CUDNN.get(value.dtype, _DataType.NOT_SET) + != expected.cudnn_dtype + ): + raise ValueError(f"{expected.name} dtype mismatch") + + @staticmethod + def _to_tensor_spec(desc, *, mode=None, divisibility=None): + if mode is None: + mode = desc.mode + mode = layout_module.normalize_mode(desc.ndim, mode) + return _TensorSpec( + layout=layout_module.to_cutlass_layout( + desc.shape, + desc.stride, + desc.stride_order, + mode=mode, + name=desc.name, + ), + mode=mode, + divisibility=divisibility, + ) + + def _call_kernel( + self, + inputs, + *, + launch, + output_descs, + input_descs=None, + input_spec=None, + output_spec=None, + **options, + ): + if input_descs is not None: + for value, desc in zip(inputs, input_descs): + self._check_tensor_signature(value, desc) + if input_spec is None: + input_spec = tuple( + self._to_tensor_spec(desc) for desc in input_descs or () + ) + if output_spec is None: + output_spec = tuple( + self._to_tensor_spec(desc) for desc in output_descs + ) + self.captured_call = { + "inputs": tuple(inputs), + "launch": launch, + "output_descs": tuple(output_descs), + "input_spec": tuple(input_spec), + "output_spec": tuple(output_spec), + **options, + } + return tuple( + _Array( + layout_module.to_public_axes(desc.shape, spec.mode), + _CUDNN_TO_DTYPE[desc.cudnn_dtype], + ) + for desc, spec in zip(output_descs, output_spec) + ) + + internal.JaxApiBase = JaxApiBase + internal.JaxTensorDesc = JaxTensorDesc + internal.TupleDict = result_module.TupleDict + + datatypes = types.ModuleType(f"{internal_name}.datatypes") + datatypes.normalize_jax_dtype = lambda value, default, _name: ( + default if value is None else value + ) + sys.modules[datatypes.__name__] = datatypes + + fake_jax = types.ModuleType("jax") + fake_jax.__path__ = [] + fake_jax.__spec__ = ModuleSpec("jax", loader=None, is_package=True) + fake_jax.ShapeDtypeStruct = _Array + cls.static_argnames = {} + + def jit(function=None, *, static_argnames=()): + def decorate(target): + cls.static_argnames[target.__name__] = tuple(static_argnames) + return target + + return decorate if function is None else decorate(function) + + fake_jax.jit = jit + fake_jnp = types.ModuleType("jax.numpy") + fake_jnp.float16 = "float16" + fake_jnp.bfloat16 = "bfloat16" + fake_jnp.float32 = "float32" + fake_jnp.int32 = "int32" + fake_jnp.dtype = lambda value: value + fake_jnp.reshape = lambda value, shape: _Array(shape, value.dtype) + fake_jnp.transpose = lambda value, axes: _Array( + tuple(value.shape[axis] for axis in axes), value.dtype + ) + fake_jax.numpy = fake_jnp + + try: + with mock.patch.dict( + sys.modules, + { + "jax": fake_jax, + "jax.numpy": fake_jnp, + "torch": None, + "cutlass": None, + }, + ): + cls.module = importlib.import_module(f"{topk_name}.jax") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + @staticmethod + def _packed_samples(cum_dtype="int32"): + return ( + _Array((20, 8, 128), "bfloat16"), + _Array((18, 2, 128), "bfloat16"), + _Array((20, 8, 1), "float32"), + _Array((3,), cum_dtype), + _Array((3,), cum_dtype), + ) + + def test_packed_thd_uses_explicit_offsets_and_virtual_kernel_views(self): + q, k, lse, cum_q, cum_k = self._packed_samples() + api = self.module.TopKReduction( + q, + k, + lse, + cum_q, + cum_k, + max_s_q=12, + max_s_k=10, + layout="T,H,D", + target_compute_capability=100, + ) + result = api(q, k, lse, cum_q, cum_k) + + self.assertEqual(api.input_layout, "THD") + self.assertEqual(api.batch, 2) + self.assertEqual(api.q_kernel_desc.shape, (1, 8, 20, 128)) + self.assertEqual(api.k_kernel_desc.shape, (1, 2, 18, 128)) + self.assertEqual(api.q_kernel_desc.stride_order, (3, 1, 2, 0)) + self.assertEqual(api.lse_kernel_desc.shape, (1, 8, 20)) + self.assertEqual( + tuple(value.shape for value in api.captured_call["inputs"]), + ((1, 8, 20, 128), (1, 2, 18, 128), (1, 8, 20), (3,), (3,)), + ) + self.assertEqual( + tuple(spec.layout for spec in api.captured_call["input_spec"]), + ((3, 1, 2, 0), (3, 1, 2, 0), (2, 1, 0), (0,), (0,)), + ) + self.assertEqual(api.captured_call["launch"].__name__, "_launch_packed_kernel") + self.assertEqual(result["topk_scores_tensor"].shape, (20, 2, 16)) + self.assertEqual(result["topk_indices_tensor"].shape, (20, 2, 16)) + self.assertEqual(api.scores_desc.init_value, float("-inf")) + self.assertEqual(api.indices_desc.init_value, -1) + + def test_fixed_bhsd_contract_is_preserved(self): + q = _Array((2, 8, 20, 128), "float16") + k = _Array((2, 2, 18, 128), "float16") + lse = _Array((2, 8, 20), "float32") + api = self.module.TopKReduction( + q, + k, + lse, + target_compute_capability=100, + ) + result = api(q, k, lse) + + self.assertEqual(api.input_layout, "BHSD") + self.assertEqual(len(api.captured_call["inputs"]), 3) + self.assertEqual(api.captured_call["launch"].__name__, "_launch_kernel") + self.assertEqual(api.captured_call["input_spec"][0].layout, (3, 1, 2, 0)) + self.assertEqual(result["topk_scores_tensor"].shape, (2, 2, 20, 16)) + self.assertEqual(result["topk_indices_tensor"].shape, (2, 2, 20, 16)) + + def test_packed_metadata_is_rejected_before_lowering(self): + q, k, lse, cum_q, cum_k = self._packed_samples() + with self.assertRaisesRegex(ValueError, "max_s_q and max_s_k"): + self.module.TopKReduction(q, k, lse, cum_q, cum_k) + with self.assertRaisesRegex(ValueError, "both required"): + self.module.TopKReduction( + q, + k, + lse, + cum_q, + None, + max_s_q=12, + max_s_k=10, + ) + + _, _, _, bad_cum_q, bad_cum_k = self._packed_samples("int64") + with self.assertRaisesRegex(ValueError, "dtype int32"): + self.module.TopKReduction( + q, + k, + lse, + bad_cum_q, + bad_cum_k, + max_s_q=12, + max_s_k=10, + ) + with self.assertRaisesRegex(ValueError, "BHSD layout requires rank-4"): + self.module.TopKReduction(q, k, lse, layout="BHSD") + + def test_wrapper_marks_packed_problem_extents_and_layout_static(self): + static = self.static_argnames["topk_reduction_wrapper"] + self.assertIn("max_s_q", static) + self.assertIn("max_s_k", static) + self.assertIn("layout", static) + + def test_packed_launcher_follows_stream_inputs_outputs_order(self): + path = _TOPK_ROOT / "jax.py" + tree = ast.parse(path.read_text(), filename=str(path)) + adapter = next( + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == "TopKReduction" + ) + launcher = next( + node + for node in adapter.body + if isinstance(node, ast.FunctionDef) + and node.name == "_launch_packed_kernel" + ) + arguments = tuple(argument.arg for argument in launcher.args.args[1:]) + self.assertEqual( + arguments, + ( + "stream", + "q", + "k", + "lse", + "cum_seqlen_q", + "cum_seqlen_k", + "topk_scores", + "topk_indices", + ), + ) + imports = [ + node + for node in ast.walk(tree) + if isinstance(node, (ast.Import, ast.ImportFrom)) + ] + self.assertFalse( + any( + ( + isinstance(node, ast.Import) + and any(alias.name == "torch" for alias in node.names) + ) + or (isinstance(node, ast.ImportFrom) and node.module == "torch") + for node in imports + ) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_rmsnorm_rht_amax.py b/test/python/fe_api/test_jax_rmsnorm_rht_amax.py new file mode 100644 index 000000000..0a1e1ba29 --- /dev/null +++ b/test/python/fe_api/test_jax_rmsnorm_rht_amax.py @@ -0,0 +1,158 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""JAX integration tests for RMSNorm + RHT + amax.""" + +import math + +import pytest + + +def _jax_runtime(): + jax = pytest.importorskip("jax") + jnp = pytest.importorskip("jax.numpy") + cutlass_jax = pytest.importorskip("cutlass.jax") + if not cutlass_jax.is_available(): + pytest.skip("Installed JAX version is unsupported by CUTLASS JAX") + return jax, jnp + + +def _hadamard_16(jnp): + values = [[1.0 if ((row & column).bit_count() % 2 == 0) else -1.0 for column in range(16)] for row in range(16)] + return jnp.asarray(values, dtype=jnp.float32) / math.sqrt(16) + + +@pytest.mark.L0 +def test_jax_rmsnorm_rht_amax_abstract_contract(monkeypatch): + jax, jnp = _jax_runtime() + + from cudnn import Op, TensorDesc + from cudnn.jax import JaxApiBase, RmsNormRhtAmaxSm100, rmsnorm_rht_amax_sm100 + from cudnn.rmsnorm_rht_amax.op import RmsNormRhtAmaxSm100Op + + monkeypatch.setattr(JaxApiBase, "_local_gpu_capabilities", staticmethod(lambda _operation_name: ((object(), 100),))) + + sample_x = jax.ShapeDtypeStruct((256, 2048), jnp.bfloat16) + sample_weight = jax.ShapeDtypeStruct((2048,), jnp.bfloat16) + api = RmsNormRhtAmaxSm100(sample_x, sample_weight) + + assert isinstance(api, JaxApiBase) + assert isinstance(api._op, Op) + assert isinstance(api._op, RmsNormRhtAmaxSm100Op) + assert isinstance(api.x_desc, TensorDesc) + assert isinstance(api.w_desc, TensorDesc) + assert isinstance(api.o_desc, TensorDesc) + assert isinstance(api.amax_desc, TensorDesc) + assert api._op.x is api.x_desc + assert api._op.weight is api.w_desc + assert api._op.output is api.o_desc + assert api._op.amax is api.amax_desc + assert api.x_desc.dtype == jnp.dtype(jnp.bfloat16) + assert api.x_desc.stride == (2048, 1) + assert api.x_desc.stride_order == (1, 0) + assert api.x_desc.name == "sample_x" + assert all(value is not sample_x and value is not sample_weight for value in vars(api).values()) + output, amax = jax.eval_shape(api, sample_x, sample_weight) + assert output.shape == (256, 2048) + assert output.dtype == jnp.bfloat16 + assert amax.shape == (128,) + assert amax.dtype == jnp.float32 + + explicit_api = RmsNormRhtAmaxSm100( + sample_x, + sample_weight, + sample_o=jax.ShapeDtypeStruct((256, 2048), jnp.bfloat16), + sample_amax=jax.ShapeDtypeStruct((128,), jnp.float32), + ) + explicit_output, explicit_amax = jax.eval_shape(explicit_api, sample_x, sample_weight) + assert explicit_output.shape == (256, 2048) + assert explicit_amax.shape == (128,) + + array_api = RmsNormRhtAmaxSm100( + jnp.empty((2, 2048), jnp.bfloat16), + jnp.empty((2048,), jnp.bfloat16), + ) + assert array_api.x_desc.shape == (2, 2048) + assert array_api.w_desc.shape == (2048,) + + with pytest.raises(ValueError, match="sample_o and sample_amax must be provided together"): + RmsNormRhtAmaxSm100( + sample_x, + sample_weight, + sample_o=jax.ShapeDtypeStruct((256, 2048), jnp.bfloat16), + ) + + functional_output, functional_amax = jax.eval_shape( + lambda x, weight: rmsnorm_rht_amax_sm100( + x, + weight, + eps=1e-4, + num_threads=128, + rows_per_cta=2, + ), + sample_x, + sample_weight, + ) + assert functional_output.shape == (256, 2048) + assert functional_output.dtype == jnp.bfloat16 + assert functional_amax.shape == (128,) + assert functional_amax.dtype == jnp.float32 + + wrong_x = jax.ShapeDtypeStruct((128, 2048), jnp.bfloat16) + with pytest.raises(ValueError, match="sample_x tensor shape mismatch"): + jax.eval_shape(api, wrong_x, sample_weight) + + invalid_api = RmsNormRhtAmaxSm100(sample_x, sample_weight, rows_per_cta=3) + with pytest.raises(ValueError, match="M must be divisible by rows_per_cta"): + jax.eval_shape(invalid_api, sample_x, sample_weight) + + +@pytest.mark.L0 +def test_jax_rmsnorm_rht_amax_jit(): + jax, jnp = _jax_runtime() + gpu_devices = [device for device in jax.local_devices() if device.platform == "gpu"] + all_local_devices_are_supported = bool(gpu_devices) + for device in gpu_devices: + capability = getattr(device, "compute_capability", None) + if capability is None: + all_local_devices_are_supported = False + break + major, minor = (int(value) for value in str(capability).split(".", 1)) + if (major, minor) < (10, 0): + all_local_devices_are_supported = False + break + if not all_local_devices_are_supported: + pytest.skip("RMSNorm + RHT requires every local JAX GPU to be SM100+") + + from cudnn.jax import rmsnorm_rht_amax_sm100 + + device = gpu_devices[0] + m, n = 256, 2048 + rows_per_cta = 2 + eps = 1e-4 + x = jax.device_put(jax.random.normal(jax.random.key(0), (m, n), dtype=jnp.bfloat16), device) + weight = jax.device_put(jax.random.normal(jax.random.key(1), (n,), dtype=jnp.bfloat16), device) + lowered = rmsnorm_rht_amax_sm100.lower( + x, + weight, + eps=eps, + num_threads=128, + rows_per_cta=rows_per_cta, + ) + stablehlo = lowered.as_text("stablehlo") + assert stablehlo.count("stablehlo.custom_call") == 1 + assert "CuteDSLRT_NvJaxCutlassCall" in stablehlo + + output, amax = lowered.compile()(x, weight) + output.block_until_ready() + x_f32 = x.astype(jnp.float32) + normalized = x_f32 * jax.lax.rsqrt(jnp.mean(jnp.square(x_f32), axis=-1, keepdims=True) + eps) + normalized *= weight.astype(jnp.float32)[None, :] + reference = (normalized.reshape(m, n // 16, 16) @ _hadamard_16(jnp)).reshape(m, n) + amax_reference = jnp.max( + jnp.abs(reference).reshape(m // rows_per_cta, rows_per_cta, n), + axis=(1, 2), + ) + + assert jnp.allclose(output.astype(jnp.float32), reference, atol=4e-2, rtol=1e-2) + assert jnp.allclose(amax, amax_reference, atol=2e-3, rtol=1e-3) diff --git a/test/python/fe_api/test_jax_sdpa_thd_contract.py b/test/python/fe_api/test_jax_sdpa_thd_contract.py new file mode 100644 index 000000000..050edce56 --- /dev/null +++ b/test/python/fe_api/test_jax_sdpa_thd_contract.py @@ -0,0 +1,185 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Source contracts for packed-THD JAX SDPA.""" + +from __future__ import annotations + +import ast +from pathlib import Path +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_FORWARD = _CUDNN_ROOT / "sdpa" / "fwd" / "jax.py" +_BACKWARD = _CUDNN_ROOT / "sdpa" / "bwd" / "jax.py" + + +def _tree(path: Path) -> ast.Module: + return ast.parse(path.read_text(), filename=str(path)) + + +def _class(tree: ast.Module, name: str) -> ast.ClassDef: + return next( + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == name + ) + + +def _function(nodes: list[ast.stmt], name: str) -> ast.FunctionDef: + return next( + node + for node in nodes + if isinstance(node, ast.FunctionDef) and node.name == name + ) + + +def _argument_names(function: ast.FunctionDef) -> tuple[str, ...]: + return tuple(argument.arg for argument in function.args.args) + + +def _jit_static_argnames(function: ast.FunctionDef) -> tuple[str, ...]: + decorator = next( + node + for node in function.decorator_list + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "partial" + ) + keyword = next( + keyword for keyword in decorator.keywords if keyword.arg == "static_argnames" + ) + return ast.literal_eval(keyword.value) + + +class JaxSdpaThdContractTest(unittest.TestCase): + def test_constructors_bind_cumulative_metadata_and_static_bounds(self): + for path, class_name in ( + (_FORWARD, "SdpafwdSm100D256"), + (_BACKWARD, "SdpabwdSm100D256"), + ): + constructor = _function(_class(_tree(path), class_name).body, "__init__") + arguments = _argument_names(constructor) + with self.subTest(path=path.relative_to(_CUDNN_ROOT)): + self.assertIn("sample_cum_seqlen_q", arguments) + self.assertIn("sample_cum_seqlen_k", arguments) + self.assertIn("max_s_q", arguments) + self.assertIn("max_s_k", arguments) + + def test_functional_wrappers_keep_offsets_dynamic_and_bounds_static(self): + for path, wrapper_name in ( + (_FORWARD, "sdpa_fwd_wrapper_sm100_d256"), + (_BACKWARD, "sdpa_bwd_wrapper_sm100_d256"), + ): + wrapper = _function(_tree(path).body, wrapper_name) + arguments = _argument_names(wrapper) + static_arguments = _jit_static_argnames(wrapper) + with self.subTest(path=path.relative_to(_CUDNN_ROOT)): + self.assertIn("cum_seqlen_q_tensor", arguments) + self.assertIn("cum_seqlen_k_tensor", arguments) + self.assertNotIn("cum_seqlen_q_tensor", static_arguments) + self.assertNotIn("cum_seqlen_k_tensor", static_arguments) + self.assertIn("max_s_q", static_arguments) + self.assertIn("max_s_k", static_arguments) + + def test_varlen_launchers_follow_stream_inputs_outputs_workspaces_order(self): + expected = { + _FORWARD: ( + "stream", + "q", + "k", + "v", + "cum_seqlen_q", + "cum_seqlen_k", + "output", + "lse", + ), + _BACKWARD: ( + "stream", + "q", + "k", + "v", + "output", + "doutput", + "lse", + "cum_seqlen_q", + "cum_seqlen_k", + "dq", + "dk", + "dv", + "workspace", + ), + } + class_names = { + _FORWARD: "SdpafwdSm100D256", + _BACKWARD: "SdpabwdSm100D256", + } + for path, arguments in expected.items(): + launcher = _function( + _class(_tree(path), class_names[path]).body, + "_launch_varlen_kernel", + ) + with self.subTest(path=path.relative_to(_CUDNN_ROOT)): + self.assertEqual(_argument_names(launcher)[1:], arguments) + + def test_thd_is_promoted_without_changing_public_result_ranks(self): + forward = _FORWARD.read_text() + backward = _BACKWARD.read_text() + + for source in (forward, backward): + self.assertIn("JaxTensorDesc.from_shape(", source) + self.assertIn("(1, *sample_q.shape)", source) + self.assertIn("jnp.reshape(q_tensor, self.q_kernel_desc.shape)", source) + self.assertIn("max_s_q and max_s_k are both required for THD", source) + + self.assertIn("jnp.reshape(output, self.o_desc.shape)", forward) + self.assertIn("(self.total_q_tokens, self.num_query_heads)", forward) + self.assertIn("public_stride_order=(0, 1)", forward) + for gradient in ("dq", "dk", "dv"): + self.assertIn( + f"{gradient} = jnp.reshape({gradient}, self.{gradient}_desc.shape)", + backward, + ) + self.assertIn("(1, self.num_query_heads, self.total_q_tokens)", backward) + self.assertNotIn("public_stride_order=(1, 2, 0)", backward) + + def test_backward_workspace_is_initialized_and_adapters_do_not_import_torch(self): + backward_tree = _tree(_BACKWARD) + init_values = { + ast.unparse(keyword.value) + for call in ast.walk(backward_tree) + if isinstance(call, ast.Call) + for keyword in call.keywords + if keyword.arg == "init_value" + } + self.assertIn("0", init_values) + + for path in (_FORWARD, _BACKWARD): + imports = [ + node + for node in ast.walk(_tree(path)) + if isinstance(node, (ast.Import, ast.ImportFrom)) + ] + with self.subTest(path=path.relative_to(_CUDNN_ROOT)): + self.assertFalse( + any( + ( + isinstance(node, ast.Import) + and any(alias.name == "torch" for alias in node.names) + ) + or (isinstance(node, ast.ImportFrom) and node.module == "torch") + for node in imports + ) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_jax_sliding_packed_contract.py b/test/python/fe_api/test_jax_sliding_packed_contract.py new file mode 100644 index 000000000..781053a0d --- /dev/null +++ b/test/python/fe_api/test_jax_sliding_packed_contract.py @@ -0,0 +1,64 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Executable contracts for packed-THD sliding-window repacking.""" + +from __future__ import annotations + +import ast +from pathlib import Path +from typing import Any + +import pytest + +pytestmark = pytest.mark.L0 + +jax = pytest.importorskip("jax") +jnp = pytest.importorskip("jax.numpy") + +_ADAPTER = ( + Path(__file__).resolve().parents[3] + / "python" + / "cudnn" + / "native_sparse_attention" + / "sliding_window_attention" + / "jax.py" +) + + +def _load_static_method(name: str): + tree = ast.parse(_ADAPTER.read_text(), filename=str(_ADAPTER)) + class_node = next( + node + for node in tree.body + if isinstance(node, ast.ClassDef) and node.name == "SlidingWindowAttention" + ) + method = next( + node + for node in class_node.body + if isinstance(node, ast.FunctionDef) and node.name == name + ) + method.decorator_list = [] + module = ast.fix_missing_locations(ast.Module(body=[method], type_ignores=[])) + namespace = {"Any": Any, "jax": jax, "jnp": jnp} + exec(compile(module, str(_ADAPTER), "exec"), namespace) + return namespace[name] + + +def test_packed_padding_and_gather_round_trip_under_jit(): + pad_packed = _load_static_method("_pad_packed") + unpad_packed = _load_static_method("_unpad_packed") + values = jnp.arange(5, dtype=jnp.float32).reshape(5, 1, 1) + cumulative = jnp.array([0, 2, 5], dtype=jnp.int32) + + padded = jax.jit(lambda x, offsets: pad_packed(x, offsets[:-1], 3))( + values, cumulative + ) + assert padded.shape == (2, 3, 1, 1) + assert padded[:, :, 0, 0].tolist() == [[0.0, 1.0, 2.0], [2.0, 3.0, 4.0]] + + restored = jax.jit(lambda x, offsets: unpad_packed(x, offsets, 5))( + padded, cumulative + ) + assert restored.shape == values.shape + assert restored[:, 0, 0].tolist() == values[:, 0, 0].tolist() diff --git a/test/python/fe_api/test_result.py b/test/python/fe_api/test_result.py new file mode 100644 index 000000000..8a7376fcc --- /dev/null +++ b/test/python/fe_api/test_result.py @@ -0,0 +1,66 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for framework-neutral frontend results.""" + +import importlib.util +from pathlib import Path +import sys +import types + +import pytest + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_COMMON_SPEC = importlib.util.spec_from_file_location( + "_cudnn_result_contract", + _CUDNN_ROOT / "common" / "result.py", +) +assert _COMMON_SPEC is not None and _COMMON_SPEC.loader is not None +_COMMON_MODULE = importlib.util.module_from_spec(_COMMON_SPEC) +_COMMON_SPEC.loader.exec_module(_COMMON_MODULE) +TupleDict = _COMMON_MODULE.TupleDict + + +pytestmark = pytest.mark.L0 + + +def test_tuple_dict_preserves_named_and_positional_access(): + result = TupleDict(output=1, amax=2) + + assert tuple(result) == (1, 2) + assert result[0] == result["output"] == 1 + assert result[1] == result["amax"] == 2 + with pytest.raises(IndexError, match="index -1 out of range"): + _ = result[-1] + + +def test_tuple_dict_is_a_jax_pytree(monkeypatch): + jax = pytest.importorskip("jax") + + package_name = "_cudnn_result_jax_contract" + package = types.ModuleType(package_name) + package.__path__ = [str(_CUDNN_ROOT)] + common_package = types.ModuleType(f"{package_name}.common") + common_package.__path__ = [str(_CUDNN_ROOT / "common")] + jax_package = types.ModuleType(f"{package_name}._jax") + jax_package.__path__ = [str(_CUDNN_ROOT / "_jax")] + monkeypatch.setitem(sys.modules, package_name, package) + monkeypatch.setitem(sys.modules, f"{package_name}.common", common_package) + monkeypatch.setitem(sys.modules, f"{package_name}._jax", jax_package) + monkeypatch.setitem(sys.modules, f"{package_name}.common.result", _COMMON_MODULE) + + spec = importlib.util.spec_from_file_location( + f"{package_name}._jax.result", + _CUDNN_ROOT / "_jax" / "result.py", + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + JaxTupleDict = module.TupleDict + + result = JaxTupleDict(output=1, amax=2) + leaves, treedef = jax.tree_util.tree_flatten(result) + + assert leaves == [1, 2] + assert jax.tree_util.tree_unflatten(treedef, leaves) == result diff --git a/test/python/fe_api/test_rmsnorm_rht_amax.py b/test/python/fe_api/test_rmsnorm_rht_amax.py index c00d04b02..e4bb15864 100644 --- a/test/python/fe_api/test_rmsnorm_rht_amax.py +++ b/test/python/fe_api/test_rmsnorm_rht_amax.py @@ -82,6 +82,9 @@ def test_rmsnorm_rht_amax_compile_execute(n, num_threads, request): num_threads=num_threads, rows_per_cta=rows_per_cta, ) + assert api.eps == eps + assert api.requested_num_threads == num_threads + assert api.requested_rows_per_cta == rows_per_cta try: assert api.check_support(), "Unsupported testcase" @@ -89,10 +92,76 @@ def test_rmsnorm_rht_amax_compile_execute(n, num_threads, request): pytest.skip(f"Unsupported testcase: {exc}") api.compile() - api.execute(x_tensor=x, w_tensor=w, o_tensor=o, amax_tensor=amax) + result = api.execute(x_tensor=x, w_tensor=w, o_tensor=o, amax_tensor=amax) + assert result is None _assert_ref_close(x, w, o, amax, eps=eps, rows_per_cta=rows_per_cta, skip_ref=skip_ref) +@pytest.mark.L0 +def test_rmsnorm_rht_amax_class_rejects_invalid_output_exemplars(monkeypatch): + try: + from cudnn import RmsNormRhtAmaxSm100 + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + sample_x = torch.empty((4, 256), dtype=torch.bfloat16) + sample_w = torch.empty((256,), dtype=torch.bfloat16) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (10, 0)) + + cases = ( + (torch.empty((2, 256), dtype=torch.bfloat16), torch.empty((2,), dtype=torch.float32), "O must have shape"), + (torch.empty((4, 256), dtype=torch.float32), torch.empty((2,), dtype=torch.float32), "O must have dtype"), + (torch.empty_strided((4, 256), (1, 4), dtype=torch.bfloat16), torch.empty((2,), dtype=torch.float32), "O must be row-major contiguous"), + (torch.empty((4, 256), dtype=torch.bfloat16), torch.empty((1,), dtype=torch.float32), "Amax must have shape"), + (torch.empty((4, 256), dtype=torch.bfloat16), torch.empty((2,), dtype=torch.bfloat16), "Amax must have dtype"), + ) + for sample_o, sample_amax, message in cases: + with pytest.raises(ValueError, match=message): + RmsNormRhtAmaxSm100( + sample_x, + sample_w, + sample_o, + sample_amax, + num_threads=32, + rows_per_cta=2, + ).check_support() + + strided_amax = torch.empty_strided((2,), (2,), dtype=torch.float32) + assert RmsNormRhtAmaxSm100( + sample_x, + sample_w, + torch.empty_like(sample_x), + strided_amax, + num_threads=32, + rows_per_cta=2, + ).check_support() + + +@pytest.mark.L0 +def test_rmsnorm_rht_amax_execute_checks_dynamic_amax_shape(): + try: + from cudnn import RmsNormRhtAmaxSm100 + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + x = torch.empty((4, 256), dtype=torch.bfloat16) + w = torch.empty((256,), dtype=torch.bfloat16) + o = torch.empty_like(x) + amax = torch.empty((2,), dtype=torch.float32) + api = RmsNormRhtAmaxSm100(x, w, o, amax, num_threads=32, rows_per_cta=2) + api._op.check_support() + api._compiled_kernel = lambda **_kwargs: None + + with pytest.raises(ValueError, match="Amax tensor shape mismatch"): + api.execute(x, w, o, torch.empty((1,), dtype=torch.float32)) + + dynamic_x = torch.empty((8, 256), dtype=torch.bfloat16) + dynamic_o = torch.empty_like(dynamic_x) + dynamic_amax = torch.empty((4,), dtype=torch.float32) + assert api.execute(dynamic_x, w, dynamic_o, dynamic_amax, current_stream=object()) is None + + @pytest.mark.L0 @torch_fork_set_rng(seed=0) @pytest.mark.parametrize("n,num_threads", SUPPORTED_N_NUM_THREADS) @@ -122,3 +191,89 @@ def test_rmsnorm_rht_amax_wrapper(n, num_threads, rows_per_cta, request): assert outputs["o_tensor"].shape == (m, n) assert outputs["amax_tensor"].shape == (m // rows_per_cta,) _assert_ref_close(x, w, outputs["o_tensor"], outputs["amax_tensor"], eps=eps, rows_per_cta=rows_per_cta, skip_ref=skip_ref) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_rmsnorm_rht_amax_wrapper_reuses_compiled_kernel(monkeypatch): + try: + from cudnn.rmsnorm_rht_amax import api as rmsnorm_api + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + construction_count = 0 + original_init = rmsnorm_api.RMSNormRHTAmaxKernel.__init__ + + def counted_init(self, *args, **kwargs): + nonlocal construction_count + construction_count += 1 + original_init(self, *args, **kwargs) + + monkeypatch.setattr(rmsnorm_api.RMSNormRHTAmaxKernel, "__init__", counted_init) + rmsnorm_api._cache_of_RmsNormRhtAmaxSm100Objects.clear() + + first_m, second_m, n = 256, 512, 2048 + rows_per_cta = 2 + first_x, w = _make_inputs(m=first_m, n=n) + second_x, _ = _make_inputs(m=second_m, n=n) + try: + try: + first = rmsnorm_api.rmsnorm_rht_amax_wrapper_sm100(first_x, w, num_threads=128, rows_per_cta=rows_per_cta) + except (ValueError, RuntimeError) as exc: + pytest.skip(f"Unsupported testcase: {exc}") + + second = rmsnorm_api.rmsnorm_rht_amax_wrapper_sm100(second_x, w, num_threads=128, rows_per_cta=rows_per_cta) + _assert_ref_close(second_x, w, second["o_tensor"], second["amax_tensor"], eps=1e-5, rows_per_cta=rows_per_cta) + + assert construction_count == 1 + assert first["o_tensor"].shape == (first_m, n) + assert second["o_tensor"].shape == (second_m, n) + assert first["amax_tensor"].shape == (first_m // rows_per_cta,) + assert second["amax_tensor"].shape == (second_m // rows_per_cta,) + finally: + rmsnorm_api._cache_of_RmsNormRhtAmaxSm100Objects.clear() + + +@pytest.mark.L0 +def test_rmsnorm_rht_amax_wrapper_passes_allocated_outputs_to_class(monkeypatch): + try: + from cudnn.rmsnorm_rht_amax import api as rmsnorm_api + except ImportError: + pytest.skip("Environment not supported: cudnn optional dependencies not installed") + + constructed = {} + + class FakeApi: + def __init__(self, **kwargs): + constructed.update(kwargs) + + def check_support(self): + return True + + def compile(self): + pass + + def execute(self, **kwargs): + constructed["execute"] = kwargs + + monkeypatch.setattr(rmsnorm_api, "RmsNormRhtAmaxSm100", FakeApi) + rmsnorm_api._cache_of_RmsNormRhtAmaxSm100Objects.clear() + + x = torch.empty((4, 256), dtype=torch.bfloat16) + w = torch.empty((256,), dtype=torch.bfloat16) + try: + result = rmsnorm_api.rmsnorm_rht_amax_wrapper_sm100(x, w, num_threads=32, rows_per_cta=4) + finally: + rmsnorm_api._cache_of_RmsNormRhtAmaxSm100Objects.clear() + + assert set(constructed) == {"sample_x", "sample_w", "sample_o", "sample_amax", "eps", "num_threads", "rows_per_cta", "execute"} + assert constructed["sample_o"] is result["o_tensor"] + assert constructed["sample_amax"] is result["amax_tensor"] + assert torch.isneginf(result["amax_tensor"]).all() + assert constructed["execute"] == { + "x_tensor": x, + "w_tensor": w, + "o_tensor": result["o_tensor"], + "amax_tensor": result["amax_tensor"], + "current_stream": None, + } diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_imports.py b/test/python/fe_api/test_rmsnorm_rht_amax_imports.py new file mode 100644 index 000000000..a4d13faa2 --- /dev/null +++ b/test/python/fe_api/test_rmsnorm_rht_amax_imports.py @@ -0,0 +1,377 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Import-boundary contracts for the RMSNorm framework adapters.""" + +import ast +import importlib +import importlib.util +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "rmsnorm_rht_amax" + + +def _load_cudnn_submodule(name: str, submodule: str): + package = types.ModuleType(name) + package.__path__ = [str(_CUDNN_ROOT)] + package.__package__ = name + sys.modules[name] = package + return importlib.import_module(f"{name}.{submodule}") + + +def _load_operation_package(name: str): + parent_name, separator, _ = name.rpartition(".") + if not separator: + raise ValueError("Synthetic operation package must have a parent") + parent = types.ModuleType(parent_name) + parent.__path__ = [str(_CUDNN_ROOT)] + parent.__package__ = parent_name + sys.modules[parent_name] = parent + + spec = importlib.util.spec_from_file_location( + name, + _OPERATION_ROOT / "__init__.py", + submodule_search_locations=[str(_OPERATION_ROOT)], + ) + if spec is None or spec.loader is None: + raise RuntimeError("Unable to load RMSNorm package") + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _remove_package(name: str) -> None: + root_name = name.split(".", 1)[0] + for module_name in tuple(sys.modules): + if module_name == root_name or module_name.startswith(f"{root_name}."): + sys.modules.pop(module_name, None) + + +class RmsNormImportContractTest(unittest.TestCase): + def test_jax_facade_reports_missing_jax_dependency(self): + name = "cudnn_jax_missing_dependency_test" + try: + with mock.patch.dict(sys.modules, {"jax": None}): + with self.assertRaisesRegex(ImportError, r"nvidia-cudnn-frontend\[jax\]"): + _load_cudnn_submodule(name, "jax") + finally: + _remove_package(name) + + def test_jax_facade_reports_missing_cutlass_dependency(self): + name = "cudnn_jax_missing_cutlass_test" + jax = types.ModuleType("jax") + try: + with mock.patch.dict(sys.modules, {"jax": jax, "cutlass": None, "cutlass.jax": None}): + with self.assertRaisesRegex(ImportError, r"nvidia-cudnn-frontend\[jax\]"): + _load_cudnn_submodule(name, "jax") + finally: + _remove_package(name) + + def test_jax_facade_reports_incompatible_jax(self): + name = "cudnn_jax_incompatible_test" + jax = types.ModuleType("jax") + jax.__version__ = "0.8.0" + cutlass = types.ModuleType("cutlass") + cutlass.__path__ = [] + cutlass_jax = types.ModuleType("cutlass.jax") + cutlass_jax.is_available = lambda: False + cutlass_jax.CUTE_DSL_MIN_SUPPORTED_JAX_VERSION = (0, 9, 1) + cutlass.jax = cutlass_jax + try: + with mock.patch.dict(sys.modules, {"jax": jax, "cutlass": cutlass, "cutlass.jax": cutlass_jax}): + with self.assertRaisesRegex(ImportError, r"JAX 0\.8\.0.*minimum supported JAX version is 0\.9\.1.*nvidia-cudnn-frontend\[jax\]"): + _load_cudnn_submodule(name, "jax") + finally: + _remove_package(name) + + def test_operation_package_does_not_advertise_jax_submodule(self): + name = "cudnn_jax_operator_route_test.rmsnorm_rht_amax" + try: + package = _load_operation_package(name) + self.assertNotIn("jax", dir(package)) + finally: + _remove_package(name) + + def test_package_import_loads_no_framework_adapter(self): + name = "cudnn_rmsnorm_import_test.rmsnorm_rht_amax" + blocked = {module_name: None for module_name in ("torch", "jax", "cutlass", "cuda")} + try: + with mock.patch.dict(sys.modules, blocked): + package = _load_operation_package(name) + self.assertNotIn(f"{name}.api", sys.modules) + self.assertNotIn(f"{name}.jax", sys.modules) + self.assertNotIn(f"{name}.kernel", sys.modules) + self.assertTrue( + { + "RMSNormRHTAmaxKernel", + "RmsNormRhtAmaxSm100Op", + "RmsNormRhtAmaxSm100", + "api", + "kernel", + "op", + }.issubset(dir(package)) + ) + finally: + _remove_package(name) + + def test_unqualified_exports_route_only_to_torch_api(self): + name = "cudnn_rmsnorm_torch_route_test.rmsnorm_rht_amax" + sentinel = object() + api = types.ModuleType(f"{name}.api") + api.RmsNormRhtAmaxSm100 = sentinel + try: + sys.modules[api.__name__] = api + package = _load_operation_package(name) + self.assertIs(package.RmsNormRhtAmaxSm100, sentinel) + self.assertNotIn(f"{name}.jax", sys.modules) + finally: + _remove_package(name) + + def test_kernel_exports_do_not_load_an_adapter(self): + name = "cudnn_rmsnorm_kernel_route_test.rmsnorm_rht_amax" + sentinel = object() + kernel = types.ModuleType(f"{name}.kernel") + kernel.RMSNormRHTAmaxKernel = sentinel + try: + sys.modules[kernel.__name__] = kernel + package = _load_operation_package(name) + self.assertIs(package.RMSNormRHTAmaxKernel, sentinel) + self.assertNotIn(f"{name}.api", sys.modules) + self.assertNotIn(f"{name}.jax", sys.modules) + finally: + _remove_package(name) + + def test_operation_exports_do_not_load_a_framework_adapter(self): + name = "cudnn_rmsnorm_op_route_test.rmsnorm_rht_amax" + op_sentinel = object() + best_num_threads_sentinel = object() + pick_rows_per_cta_sentinel = object() + op = types.ModuleType(f"{name}.op") + op.RmsNormRhtAmaxSm100Op = op_sentinel + op.best_num_threads = best_num_threads_sentinel + op.pick_rows_per_cta = pick_rows_per_cta_sentinel + try: + sys.modules[op.__name__] = op + package = _load_operation_package(name) + self.assertIs(package.RmsNormRhtAmaxSm100Op, op_sentinel) + self.assertIs(package.best_num_threads, best_num_threads_sentinel) + self.assertIs(package.pick_rows_per_cta, pick_rows_per_cta_sentinel) + self.assertNotIn(f"{name}.api", sys.modules) + self.assertNotIn(f"{name}.jax", sys.modules) + self.assertNotIn(f"{name}.kernel", sys.modules) + finally: + _remove_package(name) + + def test_static_import_directions(self): + def all_imports(filename): + path = Path(filename) + if not path.is_absolute(): + path = _OPERATION_ROOT / path + tree = ast.parse(path.read_text(), filename=str(path)) + return [node for node in ast.walk(tree) if isinstance(node, (ast.Import, ast.ImportFrom))] + + def imports_framework(node, framework: str): + if isinstance(node, ast.Import): + return any(alias.name == framework or alias.name.startswith(f"{framework}.") for alias in node.names) + return node.level == 0 and (node.module == framework or (node.module or "").startswith(f"{framework}.")) + + base_op_imports = all_imports(_CUDNN_ROOT / "common" / "op.py") + self.assertFalse(any(imports_framework(node, framework) for node in base_op_imports for framework in ("torch", "jax", "cutlass", "cuda"))) + + operation_imports = all_imports("op.py") + self.assertFalse(any(imports_framework(node, framework) for node in operation_imports for framework in ("torch", "jax", "cutlass", "cuda"))) + + kernel_imports = all_imports("kernel.py") + self.assertFalse(any(imports_framework(node, framework) for node in kernel_imports for framework in ("torch", "jax"))) + + api_imports = all_imports("api.py") + self.assertFalse(any(imports_framework(node, "jax") for node in api_imports)) + + jax_imports = all_imports("jax.py") + self.assertFalse(any(imports_framework(node, "torch") for node in jax_imports)) + + api_tree = ast.parse((_OPERATION_ROOT / "api.py").read_text(), filename="api.py") + jax_tree = ast.parse((_OPERATION_ROOT / "jax.py").read_text(), filename="jax.py") + jax_package_tree = ast.parse((_CUDNN_ROOT / "_jax" / "__init__.py").read_text(), filename="_jax/__init__.py") + jax_facade_tree = ast.parse((_CUDNN_ROOT / "jax" / "__init__.py").read_text(), filename="jax/__init__.py") + jax_base_tree = ast.parse((_CUDNN_ROOT / "_jax" / "api_base.py").read_text(), filename="_jax/api_base.py") + op_tree = ast.parse((_OPERATION_ROOT / "op.py").read_text(), filename="op.py") + kernel_tree = ast.parse((_OPERATION_ROOT / "kernel.py").read_text(), filename="kernel.py") + torch_adapter = next(node for node in api_tree.body if isinstance(node, ast.ClassDef) and node.name == "RmsNormRhtAmaxSm100") + self.assertEqual([base.id for base in torch_adapter.bases if isinstance(base, ast.Name)], ["APIBase"]) + constructor = next(node for node in torch_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "__init__") + required_argument_count = len(constructor.args.args) - len(constructor.args.defaults) + self.assertEqual( + [argument.arg for argument in constructor.args.args[:required_argument_count]], + ["self", "sample_x", "sample_w", "sample_o", "sample_amax"], + ) + self.assertNotIn( + "_kernel", + [argument.arg for argument in (*constructor.args.args, *constructor.args.kwonlyargs)], + ) + op_constructors = [ + node for node in ast.walk(api_tree) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RmsNormRhtAmaxSm100Op" + ] + self.assertEqual(len(op_constructors), 1) + self.assertIn(op_constructors[0], ast.walk(constructor)) + self.assertTrue({"x", "weight", "output", "amax"}.issubset({keyword.arg for keyword in op_constructors[0].keywords})) + self.assertFalse( + any( + isinstance(node, ast.FunctionDef) and any(isinstance(decorator, ast.Name) and decorator.id == "property" for decorator in node.decorator_list) + for node in torch_adapter.body + ) + ) + torch_check_support = next(node for node in torch_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "check_support") + self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "check_support" for node in ast.walk(torch_check_support))) + self.assertFalse(any(isinstance(node, ast.Attribute) and node.attr == "check_output" for node in ast.walk(torch_check_support))) + self.assertFalse( + any( + isinstance(node, ast.Attribute) + and isinstance(node.ctx, ast.Store) + and isinstance(node.value, ast.Attribute) + and isinstance(node.value.value, ast.Name) + and node.value.value.id == "self" + and node.value.attr == "_op" + for node in ast.walk(torch_check_support) + ) + ) + + jax_adapter = next(node for node in jax_tree.body if isinstance(node, ast.ClassDef) and node.name == "RmsNormRhtAmaxSm100") + self.assertEqual([base.id for base in jax_adapter.bases if isinstance(base, ast.Name)], ["JaxApiBase"]) + self.assertFalse(jax_adapter.decorator_list) + self.assertTrue( + any(isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RmsNormRhtAmaxSm100Op" for node in ast.walk(jax_adapter)) + ) + jax_check_support = next(node for node in jax_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "check_support") + self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "_check_device_compatibility" for node in ast.walk(jax_check_support))) + jax_call = next(node for node in jax_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "__call__") + jax_launch = next(node for node in ast.walk(jax_call) if isinstance(node, ast.FunctionDef) and node.name == "launch") + self.assertTrue( + any(isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RMSNormRHTAmaxKernel" for node in ast.walk(jax_launch)) + ) + jax_kernel_call = next( + node for node in ast.walk(jax_call) if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "_call_kernel" + ) + launch_keyword = next(keyword.value for keyword in jax_kernel_call.keywords if keyword.arg == "launch") + self.assertIsInstance(launch_keyword, ast.Name) + self.assertEqual(launch_keyword.id, "launch") + + jax_wrapper = next(node for node in jax_tree.body if isinstance(node, ast.FunctionDef) and node.name == "rmsnorm_rht_amax_sm100") + jit_decorator = next( + decorator + for decorator in jax_wrapper.decorator_list + if isinstance(decorator, ast.Call) and isinstance(decorator.func, ast.Attribute) and decorator.func.attr == "jit" + ) + static_argnames = next(keyword.value for keyword in jit_decorator.keywords if keyword.arg == "static_argnames") + self.assertEqual( + tuple(element.value for element in static_argnames.elts), + ("eps", "num_threads", "rows_per_cta"), + ) + + torch_wrapper = next(node for node in api_tree.body if isinstance(node, ast.FunctionDef) and node.name == "rmsnorm_rht_amax_wrapper_sm100") + self.assertFalse( + any(isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RMSNormRHTAmaxKernel" for node in ast.walk(torch_wrapper)) + ) + self.assertFalse(any(isinstance(node, ast.Attribute) and node.attr == "_materialize_outputs" for node in ast.walk(torch_wrapper))) + self.assertFalse( + any( + isinstance(node, ast.Attribute) + and node.attr in {"infer_output_from", "from_tensor", "materialize", "_to_tensor_desc", "_materialize_tensor_desc"} + for node in ast.walk(torch_wrapper) + ) + ) + torch_api_construction = next( + node for node in ast.walk(torch_wrapper) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RmsNormRhtAmaxSm100" + ) + self.assertTrue({"sample_o", "sample_amax"}.issubset({keyword.arg for keyword in torch_api_construction.keywords})) + + torch_execute = next(node for node in torch_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "execute") + self.assertEqual( + [argument.arg for argument in torch_execute.args.args], + ["self", "x_tensor", "w_tensor", "o_tensor", "amax_tensor", "current_stream"], + ) + self.assertFalse(any(isinstance(node, ast.Attribute) and node.attr in {"infer_output", "materialize"} for node in ast.walk(torch_execute))) + + kernel_class = next(node for node in kernel_tree.body if isinstance(node, ast.ClassDef) and node.name == "RMSNormRHTAmaxKernel") + self.assertFalse(kernel_class.bases) + launcher = next(node for node in kernel_class.body if isinstance(node, ast.FunctionDef) and node.name == "__call__") + self.assertEqual( + [argument.arg for argument in launcher.args.args], + ["self", "x_tensor", "w_tensor", "o_tensor", "amax_tensor", "stream"], + ) + + compile_method = next(node for node in torch_adapter.body if isinstance(node, ast.FunctionDef) and node.name == "compile") + kernel_constructors = [ + node for node in ast.walk(api_tree) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "RMSNormRHTAmaxKernel" + ] + self.assertEqual(len(kernel_constructors), 1) + self.assertIn(kernel_constructors[0], ast.walk(compile_method)) + self.assertTrue( + {"n", "rows_per_cta"}.issubset( + { + node.attr + for node in ast.walk(compile_method) + if isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Attribute) + and isinstance(node.value.value, ast.Name) + and node.value.value.id == "self" + and node.value.attr == "_op" + } + ) + ) + tensor_api = next(node for node in ast.walk(compile_method) if isinstance(node, ast.FunctionDef) and node.name == "tensor_api") + compiled_call = next( + node for node in ast.walk(tensor_api) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "compiled_kernel" + ) + self.assertEqual( + [argument.id for argument in compiled_call.args if isinstance(argument, ast.Name)], + ["x_tensor", "w_tensor", "o_tensor", "amax_tensor", "stream"], + ) + + operation_class = next(node for node in op_tree.body if isinstance(node, ast.ClassDef) and node.name == "RmsNormRhtAmaxSm100Op") + self.assertEqual([base.id for base in operation_class.bases if isinstance(base, ast.Name)], ["Op"]) + self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "_call_kernel" for node in ast.walk(jax_adapter))) + self.assertFalse( + any(isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "cutlass_call" for node in ast.walk(jax_adapter)) + ) + self.assertIn("_INSTALL_HINT", {node.id for node in ast.walk(jax_facade_tree) if isinstance(node, ast.Name)}) + self.assertTrue(any(isinstance(node, ast.Attribute) and node.attr == "is_available" for node in ast.walk(jax_facade_tree))) + self.assertTrue(any(isinstance(node, ast.Try) for node in jax_facade_tree.body)) + self.assertTrue(any(isinstance(node, ast.Import) and any(alias.name == "jax" for alias in node.names) for node in jax_tree.body)) + self.assertFalse(any(isinstance(node, ast.ImportFrom) and any(alias.name == "jax" for alias in node.names) for node in jax_tree.body)) + for tree in (jax_package_tree, jax_tree): + self.assertNotIn("_INSTALL_HINT", {node.id for node in ast.walk(tree) if isinstance(node, ast.Name)}) + self.assertFalse(any(isinstance(node, ast.Try) for node in tree.body)) + self.assertFalse(any(isinstance(node, ast.Attribute) and node.attr == "is_available" for node in ast.walk(tree))) + jax_base = next(node for node in jax_base_tree.body if isinstance(node, ast.ClassDef) and node.name == "JaxApiBase") + base_call = next(node for node in jax_base.body if isinstance(node, ast.FunctionDef) and node.name == "_call_kernel") + cutlass_call = next( + node for node in ast.walk(base_call) if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "cutlass_call" + ) + self.assertTrue(cutlass_call.args) + self.assertNotIn("_call_cache", {node.attr for node in ast.walk(jax_adapter) if isinstance(node, ast.Attribute)}) + self.assertNotIn("_trace_lock", {node.attr for node in ast.walk(jax_adapter) if isinstance(node, ast.Attribute)}) + + root_source = (_CUDNN_ROOT / "__init__.py").read_text() + self.assertIn("from .common.op import Op", root_source) + self.assertIn("from .common.tensor_desc import TensorDesc", root_source) + self.assertIn('if name == "jax":', root_source) + self.assertIn('importlib.import_module(".jax", __name__)', root_source) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py b/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py new file mode 100644 index 000000000..bf8a36a4c --- /dev/null +++ b/test/python/fe_api/test_rmsnorm_rht_amax_kernel.py @@ -0,0 +1,179 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for the RMSNorm CuTe kernel.""" + +import ast +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest +from unittest import mock + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "rmsnorm_rht_amax" +_PACKAGE = "cudnn_rmsnorm_kernel_test" + + +def _identity_decorator(function=None, **_kwargs): + if function is None: + return lambda decorated: decorated + return function + + +def _module(name: str, *, package: bool = False, **attributes): + module = types.ModuleType(name) + if package: + module.__path__ = [] + for attribute, value in attributes.items(): + setattr(module, attribute, value) + return module + + +def _cutlass_stubs() -> dict[str, types.ModuleType]: + driver = _module("cuda.bindings.driver", CUstream=object) + bindings = _module("cuda.bindings", package=True, driver=driver) + cuda = _module("cuda", package=True, bindings=bindings) + + arch = _module("cutlass.cute.arch", shuffle_sync_bfly=lambda value, **_kwargs: value) + cute = _module( + "cutlass.cute", + package=True, + arch=arch, + kernel=_identity_decorator, + jit=_identity_decorator, + ) + utils = _module("cutlass.utils") + llvm = _module( + "cutlass._mlir.dialects.llvm", + AsmDialect=types.SimpleNamespace(AD_ATT=object()), + ) + dialects = _module("cutlass._mlir.dialects", package=True, llvm=llvm) + mlir = _module("cutlass._mlir", package=True, dialects=dialects) + cutlass_dsl = _module( + "cutlass.cutlass_dsl", + T=types.SimpleNamespace(f32=lambda: None), + dsl_user_op=_identity_decorator, + ) + cutlass = _module( + "cutlass", + package=True, + cute=cute, + utils=utils, + Float32=object, + Int32=object, + BFloat16=object, + ) + + return { + module.__name__: module + for module in ( + cuda, + bindings, + driver, + cutlass, + cute, + arch, + utils, + mlir, + dialects, + llvm, + cutlass_dsl, + ) + } + + +class RmsNormRhtAmaxKernelContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + sys.modules[_PACKAGE] = root + + operation_name = f"{_PACKAGE}.rmsnorm_rht_amax" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + try: + with mock.patch.dict(sys.modules, _cutlass_stubs()): + cls.kernel_module = importlib.import_module(f"{operation_name}.kernel") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def test_kernel_consumes_only_resolved_configuration(self): + kernel = self.kernel_module.RMSNormRHTAmaxKernel( + n=2048, + num_threads=128, + eps=1e-4, + rows_per_cta=2, + ) + + self.assertEqual(kernel.n, 2048) + self.assertEqual(kernel.num_threads, 128) + self.assertEqual(kernel.eps, 1e-4) + self.assertEqual(kernel.rows_per_cta, 2) + self.assertEqual(kernel.ept, 16) + self.assertEqual(kernel.vec_size, 8) + self.assertEqual(kernel.num_vec_blocks, 2) + self.assertEqual(kernel.warps_per_row, 4) + self.assertEqual(kernel.tiler_mn, (1, 2048)) + + def test_kernel_module_has_no_operation_or_framework_dependency(self): + tree = ast.parse((_OPERATION_ROOT / "kernel.py").read_text(), filename="kernel.py") + imports = [node for node in ast.walk(tree) if isinstance(node, (ast.Import, ast.ImportFrom))] + + imported_modules = set() + for node in imports: + if isinstance(node, ast.Import): + imported_modules.update(alias.name for alias in node.names) + else: + if node.module: + imported_modules.add(node.module) + imported_modules.update(alias.name for alias in node.names) + + self.assertNotIn("torch", imported_modules) + self.assertNotIn("jax", imported_modules) + self.assertNotIn("_op", imported_modules) + self.assertNotIn("_tensor_desc", imported_modules) + self.assertNotIn("data_type", imported_modules) + + kernel_class = next(node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == "RMSNormRHTAmaxKernel") + self.assertFalse(kernel_class.bases) + self.assertFalse(any(isinstance(node, ast.FunctionDef) and node.name == "check_support" for node in kernel_class.body)) + self.assertFalse(any(isinstance(node, ast.FunctionDef) and node.name.startswith("infer_") for node in kernel_class.body)) + + def test_hadamard_block_matches_operation_contract(self): + tree = ast.parse((_OPERATION_ROOT / "op.py").read_text(), filename="op.py") + assignment = next( + node + for node in tree.body + if isinstance(node, ast.Assign) and any(isinstance(target, ast.Name) and target.id == "HAD_BLOCK" for target in node.targets) + ) + + self.assertEqual(ast.literal_eval(assignment.value), self.kernel_module.RMSNormRHTAmaxKernel.HAD_BLOCK) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_rmsnorm_rht_amax_op.py b/test/python/fe_api/test_rmsnorm_rht_amax_op.py new file mode 100644 index 000000000..5db57fa78 --- /dev/null +++ b/test/python/fe_api/test_rmsnorm_rht_amax_op.py @@ -0,0 +1,186 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Contracts for the RMSNorm logical operation.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_OPERATION_ROOT = _CUDNN_ROOT / "rmsnorm_rht_amax" +_PACKAGE = "cudnn_rmsnorm_op_test" + + +class _DataType(Enum): + NOT_SET = auto() + FLOAT = auto() + BFLOAT16 = auto() + + +class RmsNormRhtAmaxOpContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + root = types.ModuleType(_PACKAGE) + root.__path__ = [str(_CUDNN_ROOT)] + root.__package__ = _PACKAGE + root.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + root.data_type = _DataType + sys.modules[_PACKAGE] = root + + operation_name = f"{_PACKAGE}.rmsnorm_rht_amax" + operation = types.ModuleType(operation_name) + operation.__path__ = [str(_OPERATION_ROOT)] + operation.__package__ = operation_name + operation.__spec__ = ModuleSpec(operation_name, loader=None, is_package=True) + sys.modules[operation_name] = operation + + try: + cls.tensor_module = importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + cls.base_module = importlib.import_module(f"{_PACKAGE}.common.op") + cls.op_module = importlib.import_module(f"{operation_name}.op") + except Exception: + cls.tearDownClass() + raise + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def _desc(self, shape, dtype=_DataType.BFLOAT16, *, stride=None, name=""): + shape = tuple(shape) + if stride is None: + stride_order = tuple(reversed(range(len(shape)))) + stride_values = [0] * len(shape) + running = 1 + for dimension in stride_order: + stride_values[dimension] = running + running *= max(shape[dimension], 1) + stride = tuple(stride_values) + else: + stride = tuple(stride) + stride_order = tuple(index for index, _ in sorted(enumerate(stride), key=lambda item: (item[1], shape[item[0]]))) + return self.tensor_module.TensorDesc( + dtype=dtype, + shape=shape, + stride=stride, + stride_order=stride_order, + name=name, + ) + + def _op(self, **overrides): + arguments = { + "x": self._desc((256, 2048), name="x"), + "weight": self._desc((2048,), name="weight"), + "output": self._desc((256, 2048), name="output"), + "amax": self._desc((128,), _DataType.FLOAT, name="amax"), + } + arguments.update(overrides) + return self.op_module.RmsNormRhtAmaxSm100Op(**arguments) + + def test_operation_validates_complete_signature_and_resolves_configuration(self): + operation = self._op() + + self.assertIsInstance(operation, self.base_module.Op) + self.assertTrue(operation.check_support()) + self.assertEqual((operation.m, operation.n), (256, 2048)) + self.assertEqual((operation.num_threads, operation.rows_per_cta), (128, 2)) + self.assertFalse(hasattr(operation, "infer_output")) + self.assertFalse(hasattr(operation, "check_output")) + self.assertNotIn(f"{_PACKAGE}.rmsnorm_rht_amax.kernel", sys.modules) + self.assertNotIn(f"{_PACKAGE}.rmsnorm_rht_amax.api", sys.modules) + self.assertNotIn(f"{_PACKAGE}.rmsnorm_rht_amax.jax", sys.modules) + + def test_positive_strided_amax_is_supported_but_broadcast_amax_is_not(self): + self.assertTrue(self._op(amax=self._desc((128,), _DataType.FLOAT, stride=(2,))).check_support()) + + with self.assertRaisesRegex(ValueError, "Amax stride must be positive"): + self._op(amax=self._desc((128,), _DataType.FLOAT, stride=(0,))).check_support() + + def test_framework_descriptors_are_compared_through_cudnn_dtype(self): + tensor_module = self.tensor_module + + class FrameworkTensorDesc(tensor_module.TensorDesc): + @property + def cudnn_dtype(self): + return { + "framework.bfloat16": _DataType.BFLOAT16, + "framework.float32": _DataType.FLOAT, + }.get(self.dtype, _DataType.NOT_SET) + + def desc(shape, dtype, *, stride=None): + canonical = self._desc(shape, stride=stride) + return FrameworkTensorDesc( + dtype=dtype, + shape=canonical.shape, + stride=canonical.stride, + stride_order=canonical.stride_order, + ) + + operation = self.op_module.RmsNormRhtAmaxSm100Op( + x=desc((256, 2048), "framework.bfloat16"), + weight=desc((2048,), "framework.bfloat16"), + output=desc((256, 2048), "framework.bfloat16"), + amax=desc((128,), "framework.float32"), + ) + self.assertTrue(operation.check_support()) + + def test_invalid_complete_signatures_and_configuration(self): + cases = ( + ({"x": self._desc((2048,))}, "X must have rank 2"), + ({"weight": self._desc((1, 2048))}, "W must have rank 1"), + ({"output": self._desc((256 * 2048,))}, "O must have rank 2"), + ({"amax": self._desc((128, 1), _DataType.FLOAT)}, "Amax must have rank 1"), + ({"x": self._desc((256, 2048), _DataType.FLOAT)}, "X must have dtype bfloat16"), + ({"weight": self._desc((2048,), _DataType.FLOAT)}, "W must have dtype bfloat16"), + ({"output": self._desc((256, 2048), _DataType.FLOAT)}, "O must have dtype bfloat16"), + ({"amax": self._desc((128,), _DataType.BFLOAT16)}, "Amax must have dtype float32"), + ({"x": self._desc((0, 2048))}, "M must be positive"), + ( + { + "x": self._desc((256, 2049)), + "weight": self._desc((2049,)), + "output": self._desc((256, 2049)), + }, + "N must be divisible by 16", + ), + ({"weight": self._desc((1024,))}, "W must have shape"), + ({"output": self._desc((128, 2048))}, "O must have shape"), + ({"amax": self._desc((64,), _DataType.FLOAT)}, "Amax must have shape"), + ({"x": self._desc((256, 2048), stride=(1, 256))}, "X must be row-major contiguous"), + ({"weight": self._desc((2048,), stride=(2,))}, "W must be contiguous"), + ({"output": self._desc((256, 2048), stride=(1, 256))}, "O must be row-major contiguous"), + ({"num_threads": 512}, "EPT=4 must be >= 8 and divisible by 8"), + ({"num_threads": 2048}, "must not exceed the CUDA block size limit"), + ({"rows_per_cta": 0}, "rows_per_cta must be positive"), + ({"rows_per_cta": 3}, "M must be divisible by rows_per_cta"), + ) + + for overrides, message in cases: + with self.subTest(message=message): + with self.assertRaisesRegex(ValueError, message): + self._op(**overrides).check_support() + + def test_constructor_requires_tensor_descriptors(self): + for name in ("x", "weight", "output", "amax"): + with self.subTest(name=name): + with self.assertRaisesRegex(TypeError, f"{name} must be a TensorDesc"): + self._op(**{name: types.SimpleNamespace(shape=(1,))}) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_shared_tensor_desc.py b/test/python/fe_api/test_shared_tensor_desc.py new file mode 100644 index 000000000..cd48a097d --- /dev/null +++ b/test/python/fe_api/test_shared_tensor_desc.py @@ -0,0 +1,205 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Tests for framework-neutral tensor descriptors.""" + +from enum import Enum, auto +import importlib +from importlib.machinery import ModuleSpec +from pathlib import Path +import sys +import types +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +_CUDNN_ROOT = Path(__file__).resolve().parents[3] / "python" / "cudnn" +_PACKAGE = "cudnn_tensor_desc_test" + + +class _DataType(Enum): + FLOAT = auto() + BFLOAT16 = auto() + + +def _load_tensor_desc_module(): + package = types.ModuleType(_PACKAGE) + package.__path__ = [str(_CUDNN_ROOT)] + package.__package__ = _PACKAGE + package.__spec__ = ModuleSpec(_PACKAGE, loader=None, is_package=True) + package.data_type = _DataType + sys.modules[_PACKAGE] = package + return importlib.import_module(f"{_PACKAGE}.common.tensor_desc") + + +class SharedTensorDescTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.tensor_desc = _load_tensor_desc_module() + + @classmethod + def tearDownClass(cls) -> None: + for name in tuple(sys.modules): + if name == _PACKAGE or name.startswith(f"{_PACKAGE}."): + sys.modules.pop(name, None) + + def test_compact_descriptor_defaults_to_row_major(self): + desc = self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.BFLOAT16, + shape=(2, 3, 4), + name="output", + ) + + self.assertIsInstance(desc, self.tensor_desc.TensorDesc) + self.assertEqual(desc.dtype, _DataType.BFLOAT16) + self.assertEqual(desc.cudnn_dtype, _DataType.BFLOAT16) + self.assertEqual(desc.shape, (2, 3, 4)) + self.assertEqual(desc.stride, (12, 4, 1)) + self.assertEqual(desc.stride_order, (2, 1, 0)) + self.assertEqual(desc.ndim, 3) + self.assertEqual(desc.name, "output") + self.assertIsNone(desc.init_value) + + def test_compact_descriptor_preserves_initial_value(self): + for init_value in (0, False, float("-inf")): + with self.subTest(init_value=init_value): + desc = self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2, 3), + init_value=init_value, + ) + + if init_value is False: + self.assertIs(desc.init_value, False) + else: + self.assertEqual(desc.init_value, init_value) + + with self.assertRaisesRegex(TypeError, "init_value must be a real scalar or None"): + self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2, 3), + init_value=[], + ) + + def test_compact_like_derives_a_canonical_descriptor(self): + source = self.tensor_desc.TensorDesc( + dtype=_DataType.BFLOAT16, + shape=(7,), + stride=(1,), + stride_order=(0,), + name="input", + ) + + derived = source.compact_like( + cudnn_dtype=_DataType.FLOAT, + shape=(2, 3), + stride_order=(0, 1), + name="workspace", + init_value=0, + ) + + self.assertIs(type(derived), self.tensor_desc.TensorDesc) + self.assertEqual(derived.dtype, _DataType.FLOAT) + self.assertEqual(derived.shape, (2, 3)) + self.assertEqual(derived.stride, (1, 2)) + self.assertEqual(derived.stride_order, (0, 1)) + self.assertEqual(derived.name, "workspace") + self.assertEqual(derived.init_value, 0) + + def test_compact_descriptor_supports_an_explicit_dimension_order(self): + desc = self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(7, 3, 5), + stride_order=(1, 2, 0), + ) + + self.assertEqual(desc.stride, (15, 1, 3)) + self.assertEqual(desc.stride_order, (1, 2, 0)) + + def test_compact_descriptor_supports_scalars(self): + desc = self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(), + ) + + self.assertEqual(desc.shape, ()) + self.assertEqual(desc.stride, ()) + self.assertEqual(desc.stride_order, ()) + self.assertEqual(desc.ndim, 0) + self.assertTrue(desc.is_compact()) + + def test_is_compact_uses_canonical_or_explicit_dimension_order(self): + row_major = self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2, 3, 4), + ) + self.assertTrue(row_major.is_compact()) + self.assertTrue(row_major.is_compact((2, 1, 0))) + self.assertFalse(row_major.is_compact((0, 1, 2))) + + size_one = self.tensor_desc.TensorDesc( + dtype=_DataType.FLOAT, + shape=(2, 1, 4), + stride=(4, 17, 1), + stride_order=(2, 0, 1), + ) + self.assertTrue(size_one.is_compact((2, 1, 0))) + + with self.assertRaisesRegex(ValueError, "must be a permutation"): + row_major.is_compact((2, 2, 0)) + + def test_compact_descriptor_preserves_order_for_zero_extents(self): + desc = self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2, 0, 3), + ) + + self.assertEqual(desc.stride, (3, 3, 1)) + self.assertEqual(desc.stride_order, (2, 1, 0)) + + def test_compact_descriptor_rejects_noncanonical_dtype(self): + with self.assertRaisesRegex(TypeError, "dtype must be a cudnn.data_type"): + self.tensor_desc.make_compact_tensor_desc(dtype="float32", shape=(4,)) + + def test_compact_descriptor_rejects_invalid_shapes(self): + for shape, message in ( + ((-1, 4), "must be non-negative"), + ((2.5, 4), "must be integers"), + ((True, 4), "must be integers"), + ): + with self.subTest(shape=shape): + with self.assertRaisesRegex((TypeError, ValueError), message): + self.tensor_desc.make_compact_tensor_desc(dtype=_DataType.FLOAT, shape=shape) + + def test_compact_descriptor_rejects_invalid_dimension_orders(self): + for stride_order, message in ( + ((1,), "rank mismatch"), + ((0, 0), "must be a permutation"), + ((0, True), "must be integers"), + ): + with self.subTest(stride_order=stride_order): + with self.assertRaisesRegex((TypeError, ValueError), message): + self.tensor_desc.make_compact_tensor_desc( + dtype=_DataType.FLOAT, + shape=(2, 3), + stride_order=stride_order, + ) + + def test_descriptor_rejects_stride_order_that_contradicts_strides(self): + with self.assertRaisesRegex(ValueError, "is inconsistent with stride"): + self.tensor_desc.TensorDesc( + dtype=_DataType.FLOAT, + shape=(2, 3), + stride=(3, 1), + stride_order=(0, 1), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/python/fe_api/test_tensor_desc.py b/test/python/fe_api/test_tensor_desc.py new file mode 100644 index 000000000..4d55bac01 --- /dev/null +++ b/test/python/fe_api/test_tensor_desc.py @@ -0,0 +1,252 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Compatibility tests for shared and Torch tensor descriptors.""" + +import unittest + +try: + import pytest +except ImportError: + pass +else: + pytestmark = pytest.mark.L0 + + +try: + import torch + + from cudnn import TensorDesc as SharedTensorDesc, data_type + from cudnn.api_base import APIBase, TensorDesc as TorchTensorDesc +except ImportError as error: + _IMPORT_ERROR = error +else: + _IMPORT_ERROR = None + + +@unittest.skipIf(_IMPORT_ERROR is not None, f"Optional FE dependencies unavailable: {_IMPORT_ERROR}") +class TensorDescTest(unittest.TestCase): + def _assert_torch_metadata(self, desc, *, name="sample", init_value=None): + self.assertIsInstance(desc, SharedTensorDesc) + self.assertEqual(desc.dtype, torch.bfloat16) + self.assertEqual(desc.device, torch.device("cpu")) + self.assertFalse(desc.interpret_uint8_as_fp4x2) + self.assertEqual(desc.name, name) + self.assertEqual(desc.init_value, init_value) + self.assertEqual(desc.ndim, len(desc.shape)) + self.assertEqual(desc.cudnn_dtype, data_type.BFLOAT16) + + def test_torch_descriptor_preserves_constructor_and_layout_operations(self): + desc = TorchTensorDesc( + torch.bfloat16, + (2, 3, 4), + (12, 4, 1), + (2, 1, 0), + "cpu", + False, + "sample", + ) + self._assert_torch_metadata(desc) + self.assertEqual(desc.size(), torch.Size((2, 3, 4))) + + transformed = ( + desc.transpose(0, 1), + desc.unsqueeze(-1), + desc.unsqueeze(-1).squeeze(-1), + desc.view(6, 4), + desc.as_strided((2, 3, 4), (12, 4, 1)), + desc.transpose(0, 1).contiguous(), + ) + for result in transformed: + with self.subTest(shape=result.shape, stride=result.stride): + self._assert_torch_metadata(result) + + def test_packed_uint8_maps_to_canonical_fp4(self): + desc = TorchTensorDesc( + dtype=torch.uint8, + shape=(16,), + stride=(1,), + stride_order=(0,), + device="cpu", + interpret_uint8_as_fp4x2=True, + name="packed_fp4", + ) + self.assertEqual(desc.cudnn_dtype, data_type.FP4_E2M1) + + def test_layout_operations_preserve_initial_value(self): + desc = TorchTensorDesc( + dtype=torch.float32, + shape=(2, 3, 4), + stride=(12, 4, 1), + stride_order=(2, 1, 0), + device="cpu", + name="workspace", + init_value=float("-inf"), + ) + + transformed = ( + desc.permute(1, 0, 2), + desc.transpose(0, 1), + desc.unsqueeze(-1), + desc.unsqueeze(-1).squeeze(-1), + desc.view(6, 4), + desc.as_strided((2, 3, 4), (12, 4, 1)), + desc.transpose(0, 1).contiguous(), + ) + for result in transformed: + with self.subTest(shape=result.shape, stride=result.stride): + self.assertEqual(result.init_value, float("-inf")) + + def test_api_base_converts_torch_tensor_metadata(self): + tensor = torch.empty_strided((2, 3, 4), (12, 1, 3), dtype=torch.bfloat16) + desc = APIBase._to_tensor_desc(tensor, "sample") + + self._assert_torch_metadata(desc) + self.assertEqual(desc.shape, (2, 3, 4)) + self.assertEqual(desc.stride, (12, 1, 3)) + self.assertEqual(desc.stride_order, (1, 2, 0)) + + def test_torch_descriptor_constructs_from_tensor_metadata(self): + tensor = torch.empty_strided((2, 3, 4), (12, 1, 3), dtype=torch.uint8) + desc = TorchTensorDesc.from_tensor( + tensor, + "logical_fp4", + shape=(2, 3, 8), + stride=(12, 1, 3), + interpret_uint8_as_fp4x2=True, + ) + + self.assertEqual(desc.dtype, torch.uint8) + self.assertEqual(desc.shape, (2, 3, 8)) + self.assertEqual(desc.stride, (12, 1, 3)) + self.assertEqual(desc.stride_order, (1, 2, 0)) + self.assertEqual(desc.device, tensor.device) + self.assertEqual(desc.name, "logical_fp4") + self.assertTrue(desc.interpret_uint8_as_fp4x2) + self.assertEqual(desc.cudnn_dtype, data_type.FP4_E2M1) + self.assertEqual( + APIBase._to_tensor_desc( + tensor, + "logical_fp4", + shape=(2, 3, 8), + stride=(12, 1, 3), + interpret_uint8_as_fp4x2=True, + ), + desc, + ) + + def test_torch_descriptor_derives_and_materializes_compact_output(self): + source = TorchTensorDesc.from_tensor( + torch.empty((4, 8), dtype=torch.bfloat16), + "input", + ) + output = source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(3, 5), + stride_order=(0, 1), + name="output", + init_value=float("-inf"), + ) + + self.assertIsInstance(output, TorchTensorDesc) + self.assertEqual(output.dtype, torch.float32) + self.assertEqual(output.device, source.device) + self.assertEqual(output.shape, (3, 5)) + self.assertEqual(output.stride, (1, 3)) + self.assertEqual(output.stride_order, (0, 1)) + self.assertEqual(output.name, "output") + self.assertEqual(output.init_value, float("-inf")) + self.assertFalse(output.interpret_uint8_as_fp4x2) + + tensor = output.materialize() + self.assertEqual(tensor.dtype, torch.float32) + self.assertEqual(tuple(tensor.shape), output.shape) + self.assertEqual(tuple(tensor.stride()), output.stride) + self.assertTrue(torch.isneginf(tensor).all()) + + def test_torch_descriptor_materializes_false_and_zero_initializers(self): + source = TorchTensorDesc.from_tensor(torch.empty((1,), dtype=torch.float32)) + for init_value in (False, 0): + with self.subTest(init_value=init_value): + output = source.compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(4,), + init_value=init_value, + ) + self.assertTrue(torch.equal(output.materialize(), torch.zeros(4))) + + def test_torch_descriptor_rejects_materializing_logical_packed_storage(self): + packed = TorchTensorDesc.from_tensor( + torch.empty((4,), dtype=torch.uint8), + interpret_uint8_as_fp4x2=True, + ) + + with self.assertRaisesRegex(ValueError, "logical FP4"): + packed.materialize() + + unpacked = TorchTensorDesc( + dtype=torch.float32, + shape=(4,), + stride=(1,), + stride_order=(0,), + device="cpu", + interpret_uint8_as_fp4x2=True, + init_value=1, + ) + self.assertTrue(torch.equal(unpacked.materialize(), torch.ones(4))) + + def test_api_base_materializes_a_canonical_descriptor(self): + desc = SharedTensorDesc( + dtype=data_type.FLOAT, + shape=(2, 3), + stride=(1, 2), + stride_order=(0, 1), + name="output", + init_value=float("-inf"), + ) + + tensor = APIBase._materialize_tensor_desc(desc, device=torch.device("cpu")) + + self.assertEqual(tensor.dtype, torch.float32) + self.assertEqual(tuple(tensor.shape), desc.shape) + self.assertEqual(tuple(tensor.stride()), desc.stride) + self.assertTrue(torch.isneginf(tensor).all()) + + def test_api_base_initializes_on_the_requested_cuda_stream(self): + if not torch.cuda.is_available(): + self.skipTest("CUDA is unavailable") + + from cuda.bindings import driver as cuda + + device = torch.device("cuda", torch.cuda.current_device()) + stream = torch.cuda.Stream(device=device) + desc = SharedTensorDesc( + dtype=data_type.FLOAT, + shape=(8,), + stride=(1,), + stride_order=(0,), + init_value=7, + ) + + tensor = APIBase._materialize_tensor_desc( + desc, + device=device, + stream=cuda.CUstream(stream.cuda_stream), + ) + stream.synchronize() + + self.assertTrue(torch.equal(tensor.cpu(), torch.full((8,), 7.0))) + + torch_desc = TorchTensorDesc.from_tensor(torch.empty((1,), dtype=torch.float32, device=device)).compact_like( + cudnn_dtype=data_type.FLOAT, + shape=(8,), + init_value=9, + ) + tensor = torch_desc.materialize(stream=cuda.CUstream(stream.cuda_stream)) + stream.synchronize() + + self.assertTrue(torch.equal(tensor.cpu(), torch.full((8,), 9.0))) + + +if __name__ == "__main__": + unittest.main()